Compare commits

..
Author SHA1 Message Date
Michaelandkhluu 2a69949bda [Bugfix]: Fix Gemma4ToolParser.__init__() missing tools parameter (#38847)
Signed-off-by: Michael Hospedales <hospedales@me.com>
(cherry picked from commit bb39382b2b)
2026-04-02 16:45:38 -07:00
Luciano Martinsandkhluu 8adcf8c40a feat(models): implement Google Gemma 4 architecture support (MoE, Multimodal, Reasoning, Tool-Use) (#38826)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Signed-off-by: Luciano Martins <lucianomartins@google.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
(cherry picked from commit 08ed2b9688)
2026-04-02 11:49:53 -07:00
khluu cfad6a509c Revert "[Bugfix] Restrict TRTLLM attention to SM100, fixing GB300 (SM103) hang (#38730)"
This reverts commit c284a6671c.
2026-04-01 15:14:58 -07:00
Stefano Castagnettaandkhluu c284a6671c [Bugfix] Restrict TRTLLM attention to SM100, fixing GB300 (SM103) hang (#38730)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
(cherry picked from commit 6183cae1bd)
2026-04-01 12:11:03 -07:00
Chaunceyandkhluu 3a30a1a6a8 [Misc] Rename think_start_str/think_end_str to reasoning_start_str/reasoning_end_str (#38242)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
(cherry picked from commit cbe7d18096)
2026-04-01 12:10:53 -07:00
Juan Pérez de Algabaandkhluu 29982d48b3 (security) Enforce frame limit in VideoMediaIO (#38636)
Signed-off-by: jperezde <jperezde@redhat.com>
(cherry picked from commit 58ee614221)
2026-04-01 12:10:40 -07:00
Yifan Qiaoandkhluu 1dbbafd3f3 [Feat][v1] Simple yet General CPU KV Cache Offloading (#37160)
Signed-off-by: Yifan Qiao <yifanqiao@berkeley.edu>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
(cherry picked from commit 91e4521f9f)
2026-04-01 01:03:14 -07:00
Lucas Wilkinsonandkhluu 0ee3b7fc3d [Bugfix][MLA] Add logits size budget to sparse indexer prefill chunking (#36178)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
(cherry picked from commit eb47454987)
2026-04-01 01:02:58 -07:00
Matthew Bonanniandkhluu 268bed9cf3 [Bugfix][Async] Fix async spec decoding with hybrid models (#38556)
Signed-off-by: SandishKumarHN <sandishkumarhn@gmail.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: SandishKumarHN <sandishkumarhn@gmail.com>
(cherry picked from commit 757068dc65)
2026-04-01 01:02:35 -07:00
Jiangyun Zhuandkhluu bcc0fdd0f3 [CI] fix LM Eval Qwen3.5 Models (B200) (#38632)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
(cherry picked from commit ea7bfde6e4)
2026-04-01 01:02:20 -07:00
wang.yuqiandkhluu 69b8bd4b33 [CI Failure] pin colmodernvbert revision (#38612)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
Signed-off-by: wang.yuqi <noooop@126.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
(cherry picked from commit 719735d6c5)
2026-04-01 01:02:04 -07:00
Li, Jiangandkhluu 12449f9492 [Bugfix][CPU] Skip set_num_threads after thread binding (#38535)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
(cherry picked from commit 6557f4937f)
2026-03-30 23:01:42 -07:00
haosdentandkhluu b92312dfd7 [CI] Fix SPLADE pooler test broken by #38139 (#38495)
Signed-off-by: haosdent <haosdent@gmail.com>
(cherry picked from commit a08b7733fd)
2026-03-30 21:52:13 -07:00
JaewonandGitHub d816834c1a [MoE] Add RoutingMethodType.Simulated to TRT-LLM FP8/NVFP4 kernel allowlists (#38329)
Signed-off-by: Jaewon Lee <jaewon@meta.com>
2026-03-29 22:53:43 -07:00
Roger WangandGitHub 92f0db57a8 [Misc] Always use forward_mulmat for Conv3d on newer versions of torch. (#38487) 2026-03-30 05:39:41 +00:00
Andreas KaratzasandGitHub bea23536f6 [CI] Add temperature=0.0, reduce max_tokens, and add debug prints to audio_in_video tests (#38492)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-30 05:36:45 +00:00
Jiangyun ZhuandGitHub c133f33746 Add @ZJY0516 to CODEOWNERS (#38497)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-03-29 21:10:00 -07:00
a6db99ba02 [Bugfix] Support multi-type params parsing for DeepSeek v3.2 (#33703)
Signed-off-by: Stanislav Kirillov <stas@nebius.com>
Co-authored-by: Stanislav Kirillov <stas@nebius.com>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
2026-03-30 04:07:28 +00:00
Andreas KaratzasandGitHub 4f2ed5fddb [ROCm][CI] Enable hybrid chunked prefill test (#38317)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-30 10:30:26 +08:00
Kyle SayersandGitHub d28d86e8a3 [QeRL] Fix online quantized reloading (#38442)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
2026-03-29 14:56:41 -06:00
Wentao YeandGitHub 995dea1354 [Perf] Remove redundant device copies for CPU-only pooling token IDs, 48.9% E2E throughput improvement (#38139)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-03-29 18:12:50 +00:00
allgatherandGitHub 8c0b6267d7 [Transformers v5] fix missing pixtral/voxtral multimodal dispatch (#38410)
Signed-off-by: allgather <all2allops@gmail.com>
2026-03-29 09:59:06 +00:00
Andreas KaratzasandGitHub 43cc5138e5 [ROCm][CI] Fix cross-attention dispatch for encoder-decoder models (#38450)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-28 22:08:03 -07:00
Shubhra PanditandGitHub 5b8c30d62b [Spec Decode, BugFix] Propagate norm_before_fc from Eagle3 speculator (#38111)
Signed-off-by: Shubhra Pandit <shubhra.pandit@gmail.com>
2026-03-29 00:42:06 +00:00
haosdentandGitHub d39b8daf5f [Feature] Add Qwen3-ForcedAligner support via token classification pooling (#35367)
Signed-off-by: haosdent <haosdent@gmail.com>
2026-03-29 00:27:52 +00:00
Walter Beller-MoralesandGitHub fafca38adc [BugFix][Frontend] apply task instruction as system prompt in cohere v2/embed (#38362)
Signed-off-by: walterbm <walter.beller.morales@gmail.com>
2026-03-28 18:30:54 +00:00
aa4eb0db78 [CI]revert initialize_model context manager (#38426)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-03-28 16:56:50 +00:00
Andreas KaratzasandGitHub af89140efc [ROCm][CI] Fix UV install in Dockerfile.rocm to detect curl failures and retry (#38415)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-29 00:47:42 +08:00
haosdentandGitHub b2bc736b12 [CI] Fix Ernie4.5-VL initialization test (#38429)
Signed-off-by: haosdent <haosdent@gmail.com>
2026-03-28 22:43:24 +08:00
whyiugandGitHub 58c959a767 [Misc]: clean up non-core lint issues (#37049)
Signed-off-by: whyiug <whyiug@hotmail.com>
2026-03-28 10:28:16 -04:00
BviciiandGitHub bda3eda82d [Bugfix] Disallow renderer_num_workers > 1 with mm processor cache (#38418)
Signed-off-by: Bvicii <yizhanhuang2002@gmail.com>
2026-03-28 06:32:52 -07:00
Michael GoinandGitHub 2bf5b70ae8 [CI Bugfix] Pre-download missing FlashInfer headers in Docker build (#38391)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-03-28 06:09:00 -07:00
yzong-rhandGitHub 6dad4c5722 [Test] Fix flaky race condition in test_abort_final_step (#38414)
Signed-off-by: Yifan <yzong@redhat.com>
2026-03-28 09:06:56 +00:00
171775f306 Fix Device Index for ROCm Ray Workers in MoE Benchmark (#38108)
Signed-off-by: Liwen <53441624+li-liwen@users.noreply.github.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-03-28 08:27:11 +00:00
TJianandGitHub 58a249bc61 [ROCm] [Release] Update ROCm variant from rocm700 to rocm721 (#38413)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-03-28 06:07:03 +00:00
IriKaandGitHub 148a5c1226 [Bugfix]fix output Nan/Inf in marlin if dtype=float16 (#33972)
Signed-off-by: IriKa Qiu <qiujie.jq@gmail.com>
2026-03-27 16:36:08 -07:00
Wei ZhaoandGitHub b69bf2f0b1 [Perf] Use torch compile to fuse pack topk in trtllm moe (#37695)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>
2026-03-27 17:30:46 -06:00
rongfu.lengandGitHub 88149b635e Add nvidia h800 moe config (#31201)
Signed-off-by: rongfu.leng <rongfu.leng@daocloud.io>
2026-03-27 16:28:48 -07:00
83a4df049d [ROCm][Documentation] update quickstart and installation to include rocm nightly docker tips (#38367)
Signed-off-by: Hongxia Yang <hongxiay.yang@amd.com>
Co-authored-by: Hongxia Yang <hongxiay.yang@amd.com>
2026-03-27 23:20:19 +00:00
Gregory ShtrasbergandGitHub 731285c939 [ROCm][CI/Build] ROCm 7.2.1 release version; torch 2.10; triton 3.6 (#38252)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
2026-03-27 18:03:12 -05:00
+3 97d19197bc [NVIDIA] Fix DGX Spark logic (#38126)
Signed-off-by: johnnynunez <johnnynuca14@gmail.com>
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Signed-off-by: Mark McLoughlin <markmc@redhat.com>
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Signed-off-by: Sathish Sanjeevi <sathish.krishnan.p.s@gmail.com>
Signed-off-by: guillaume_guy <guillaume.guy@airbnb.com>
Signed-off-by: Guillaume Guy <guillaume.c.guy@gmail.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Mark McLoughlin <markmc@redhat.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com>
Co-authored-by: Sathish Sanjeevi <SKPsanjeevi@users.noreply.github.com>
Co-authored-by: Guillaume Guy <guillaume.c.guy@gmail.com>
Co-authored-by: guillaume_guy <guillaume.guy@airbnb.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
2026-03-27 15:26:07 -07:00
Giancarlo DelfinandGitHub 384e4d5f48 [Model Runner V2] Rebuild attention metadata before eagle decode full… (#38311)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
2026-03-27 13:46:42 -07:00
Nicolò LucchesiandGitHub 44a6528028 [CI] Skip failing test (#38369)
Signed-off-by: NickLucche <nlucches@redhat.com>
2026-03-27 13:25:19 -07:00
Kyle SayersandGitHub 648edcf729 [QeRL] Compose online quantization with quantized reloading (#38032)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
2026-03-27 13:22:33 -07:00
7ba425e916 Add short flag -sc for --speculative-config argument (#38380)
Co-authored-by: Claude <noreply@anthropic.com>
2026-03-27 12:04:22 -07:00
Gregory ShtrasbergandGitHub b8665383df [ROCm] Fix GPT-OSS import for triton 3.6 (#37453)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
2026-03-27 18:00:57 +00:00
0e9358c11d {ROCm]: gpt-oss fusion/padding fixes (#38043)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-03-27 12:19:15 -04:00
Harry MellorGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
21d2b53f88 Remove need for explicit \n in docstring lists for --help formatting (#38350)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-03-27 08:38:00 -07:00
Jonas M. KüblerandGitHub 98e7f223b9 enable skipping of SW attention layers when using FP8 KV cache (#33695)
Signed-off-by: Jonas Kuebler <kuebj@amazon.com>
2026-03-27 07:25:02 -06:00
b111f8a61f fix(security): Add VLLM_MAX_N_SEQUENCES environment variable and enforce limit (#37952)
Signed-off-by: jperezde <jperezde@redhat.com>
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Co-authored-by: Russell Bryant <rbryant@redhat.com>
2026-03-27 09:02:10 -04:00
Sage MooreandGitHub 497e234d38 [EPLB] Cleanup the transfer logic for the various eplb maps (#34520)
Signed-off-by: Sage Moore <sagmoore@redhat.com>
Signed-off-by: Sage Moore <sage@neuralmagic.com>
2026-03-27 10:18:46 +01:00
dtcandGitHub 6287e7fa20 [P/D] Mooncake: Add unit tests and minor fixes for mooncake connector (#36946)
Signed-off-by: Tianchen Ding <dtcccc@linux.alibaba.com>
2026-03-27 09:26:40 +01:00
Shengqi ChenGitHubClaude Opus 4.6gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>Copilot
84e439a9cb [CI/Build] Move nightly wheel index generation to a single post-build step (#38322)
Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-03-27 07:44:18 +00:00
Yuichiro UtsumiGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
a1746ff9ec [Doc] Clarify Helm chart location in deployment guide (#38328)
Signed-off-by: Yuichiro Utsumi <utsumi.yuichiro@fujitsu.com>
Signed-off-by: Yuichiro Utsumi <81412151+utsumi-fj@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-03-27 15:43:02 +08:00
Flora FengandGitHub aee4c14689 [Bugfix] Fix Hermes tool parser when stream interval > 1 (#38168)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-03-27 14:42:26 +08:00
Bowen BaoandGitHub 0ae89f18fd [Refactor] Move FusedMoE hidden_size roundup to quant_method (#34285)
Signed-off-by: Bowen Bao <bowenbao@amd.com>
2026-03-26 23:38:26 -07:00
wenjun liuandGitHub c2b17d71af [CI] Add xpu auto-label rule for Intel GPU/XPU PRs (#38320)
Signed-off-by: wendyliu235 <wenjun.liu@intel.com>
2026-03-27 14:22:38 +08:00
Li, JiangandGitHub becaed6ec8 [CPU] Support CT W4A16 on CPU MP kernel (#38219)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
2026-03-27 14:15:28 +08:00
Xiaoshuang WangandGitHub a8eab8f30d [Model] Extract GatedDeltaNetAttention into shared layer for Qwen3Next and Qwen3.5 (#37975)
Signed-off-by: wxsIcey <1790571317@qq.com>
Signed-off-by: Icey <1790571317@qq.com>
2026-03-27 14:13:21 +08:00
cjackalandGitHub 2babac0bed [frontend] dump openai responses type by alias (#38262)
Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>
2026-03-27 05:58:20 +00:00
Or OzeriandGitHub 7cc302dd87 [kv_offload+HMA][7/N]: Support register_kv_caches for hybrid models (#37853)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
2026-03-27 08:38:33 +03:00
999dfc1622 [Bugfix] Offload blocking tokenizer ops to shared thread pool to unblock event loop (#34789)
Signed-off-by: Bvicii <yizhanhuang2002@gmail.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-03-26 22:17:00 -07:00
wenjun liuandGitHub d86060122a [CI/Build] enable Intel XPU test flow with prebuilt image (#37447)
Signed-off-by: wendyliu235 <wenjun.liu@intel.com>
2026-03-26 18:16:04 -07:00
Harry MellorandGitHub f73bcb1c51 Various Transformers v5 config fixes (#38247)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-03-26 23:06:59 +00:00
218 changed files with 16181 additions and 3765 deletions
+23
View File
@@ -0,0 +1,23 @@
name: vllm_intel_ci
job_dirs:
- ".buildkite/intel_jobs"
run_all_patterns:
- "docker/Dockerfile"
- "CMakeLists.txt"
- "requirements/common.txt"
- "requirements/xpu.txt"
- "requirements/build.txt"
- "requirements/test.txt"
- "setup.py"
- "csrc/"
- "cmake/"
run_all_exclude_patterns:
- "docker/Dockerfile."
- "csrc/cpu/"
- "csrc/rocm/"
- "cmake/hipify.py"
- "cmake/cpu_extension.cmake"
registries: public.ecr.aws/q9t5s3a7
repositories:
main: "vllm-ci-test-repo"
premerge: "vllm-ci-test-repo"
+34
View File
@@ -0,0 +1,34 @@
#!/bin/bash
set -e
if [[ $# -lt 3 ]]; then
echo "Usage: $0 <registry> <repo> <commit>"
exit 1
fi
REGISTRY=$1
REPO=$2
BUILDKITE_COMMIT=$3
# authenticate with AWS ECR
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
# skip build if image already exists
if ! docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu &> /dev/null; then
echo "Image not found, proceeding with build..."
else
echo "Image found"
exit 0
fi
# build
docker build \
--file docker/Dockerfile.xpu \
--build-arg max_jobs=16 \
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
--tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu \
--progress plain .
# push
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu
+64
View File
@@ -0,0 +1,64 @@
group: Intel
steps:
- label: ":docker: Build XPU image"
soft_fail: true
depends_on: []
key: image-build-xpu
commands:
- bash -lc '.buildkite/image_build/image_build_xpu.sh "public.ecr.aws/q9t5s3a7" "vllm-ci-test-repo" "$BUILDKITE_COMMIT"'
env:
DOCKER_BUILDKIT: "1"
retry:
automatic:
- exit_status: -1 # Agent was lost
limit: 2
- exit_status: -10 # Agent was lost
limit: 2
- label: "XPU example Test"
depends_on:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- vllm/
- .buildkite/intel_jobs/test-intel.yaml
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install tblib==3.1.0 &&
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager &&
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE &&
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp &&
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN &&
python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8 &&
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel'
- label: "XPU V1 test"
depends_on:
- image-build-xpu
timeout_in_minutes: 30
device: intel_gpu
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
source_file_dependencies:
- vllm/
- .buildkite/intel_jobs/test-intel.yaml
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py &&
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
pytest -v -s v1/structured_output &&
pytest -v -s v1/test_serial_utils.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py'
+10 -1
View File
@@ -90,6 +90,14 @@ steps:
env:
DOCKER_BUILDKIT: "1"
- label: "Generate and upload wheel indices"
depends_on: "build-wheels"
allow_dependency_failure: true
agents:
queue: cpu_queue_release
commands:
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
- group: "Build release Docker images"
key: "build-release-images"
steps:
@@ -603,7 +611,7 @@ steps:
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
env:
S3_BUCKET: "vllm-wheels"
VARIANT: "rocm700"
VARIANT: "rocm721"
# ROCm Job 6: Build ROCm Release Docker Image
- label: ":docker: Build release image - x86_64 - ROCm"
@@ -673,6 +681,7 @@ steps:
- label: "Publish nightly ROCm image to DockerHub"
depends_on:
- build-rocm-release-image
if: build.env("NIGHTLY") == "1"
agents:
queue: small_cpu_queue_release
commands:
+84
View File
@@ -0,0 +1,84 @@
#!/usr/bin/env bash
set -ex
# Generate and upload wheel indices for all wheels in the commit directory.
# This script should run once after all wheels have been built and uploaded.
# ======== setup ========
BUCKET="vllm-wheels"
INDICES_OUTPUT_DIR="indices"
DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
PYTHON="${PYTHON_PROG:-python3}" # try to read from env var, otherwise use python3
SUBPATH=$BUILDKITE_COMMIT
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# detect if python3.12+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
if [[ "$has_new_python" -eq 0 ]]; then
# use new python from docker
docker pull python:3-slim
PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
fi
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ======== generate and upload indices ========
# list all wheels in the commit directory
echo "Existing wheels on S3:"
aws s3 ls "$S3_COMMIT_PREFIX"
obj_json="objects.json"
aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
mkdir -p "$INDICES_OUTPUT_DIR"
# call script to generate indices for all existing wheels
# these indices have relative paths that work as long as they are next to the wheel directory in s3
# i.e., the wheels are always in s3://vllm-wheels/<commit>/
# and indices can be placed in /<commit>/, or /nightly/, or /<version>/
alias_args=()
if [[ -n "$DEFAULT_VARIANT_ALIAS" ]]; then
alias_args=(--alias-to-default "$DEFAULT_VARIANT_ALIAS")
fi
# HACK: we do not need regex module here, but it is required by pre-commit hook
# To avoid any external dependency, we simply replace it back to the stdlib re module
sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" "${alias_args[@]}"
# copy indices to /<commit>/ unconditionally
echo "Uploading indices to $S3_COMMIT_PREFIX"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
# copy to /nightly/ only if it is on the main branch and not a PR
if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
echo "Uploading indices to overwrite /nightly/"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
fi
# detect version from any wheel in the commit directory
# download the first wheel we find to extract version metadata
first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))")
if [[ -z "$first_wheel_key" ]]; then
echo "Error: No wheels found in $S3_COMMIT_PREFIX"
exit 1
fi
first_wheel=$(basename "$first_wheel_key")
aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}"
version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
rm -f "/tmp/${first_wheel}"
echo "Version in wheel: $version"
pure_version="${version%%+*}"
echo "Pure version (without variant): $pure_version"
# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
if [[ "$version" != *"dev"* ]]; then
echo "Re-generating indices for /$pure_version/"
rm -rf "${INDICES_OUTPUT_DIR:?}"
mkdir -p "$INDICES_OUTPUT_DIR"
# wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" "${alias_args[@]}"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/"
fi
@@ -1,9 +1,10 @@
#!/bin/bash
set -euox pipefail
export VLLM_CPU_CI_ENV=0
export VLLM_CPU_KVCACHE_SPACE=1 # avoid OOM
echo "--- PP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 &
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
@@ -23,7 +24,7 @@ if [ "$failed_req" -ne 0 ]; then
fi
echo "--- DP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 &
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
+276
View File
@@ -0,0 +1,276 @@
#!/bin/bash
# This script runs tests inside the Intel XPU docker container.
# It mirrors the structure of run-amd-test.sh while keeping Intel-specific
# container setup and allowing commands to be sourced from YAML or env.
#
# Command sources (in priority order):
# 1) VLLM_TEST_COMMANDS env var (preferred, preserves quoting)
# 2) Positional args (legacy)
# 3) One or more YAML files with a commands list (test-area style)
###############################################################################
set -o pipefail
DRY_RUN=${DRY_RUN:-0}
if [[ "${1:-}" == "--dry-run" ]]; then
DRY_RUN=1
shift
fi
# Export Python path
export PYTHONPATH=".."
###############################################################################
# Helper Functions
###############################################################################
cleanup_docker() {
docker_root=$(docker info -f '{{.DockerRootDir}}')
if [ -z "$docker_root" ]; then
echo "Failed to determine Docker root directory." >&2
exit 1
fi
echo "Docker root directory: $docker_root"
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
threshold=70
if [ "$disk_usage" -gt "$threshold" ]; then
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
docker image prune -f
docker volume prune -f && docker system prune --force --filter "until=72h" --all
echo "Docker images and volumes cleanup completed."
else
echo "Disk usage is below $threshold%. No cleanup needed."
fi
}
re_quote_pytest_markers() {
local input="$1"
local output=""
local collecting=false
local marker_buf=""
local flat="${input//$'\n'/ }"
local restore_glob
restore_glob="$(shopt -p -o noglob 2>/dev/null || true)"
set -o noglob
local -a words
read -ra words <<< "$flat"
eval "$restore_glob"
for word in "${words[@]}"; do
if $collecting; then
if [[ "$word" == *"'"* ]]; then
if [[ -n "$marker_buf" ]]; then
output+="${marker_buf} "
marker_buf=""
fi
output+="${word} "
collecting=false
continue
fi
local is_boundary=false
case "$word" in
"&&"|"||"|";"|"|")
is_boundary=true ;;
--*)
is_boundary=true ;;
-[a-zA-Z])
is_boundary=true ;;
*/*)
is_boundary=true ;;
*.py|*.py::*)
is_boundary=true ;;
*=*)
if [[ "$word" =~ ^[A-Z_][A-Z0-9_]*= ]]; then
is_boundary=true
fi
;;
esac
if $is_boundary; then
if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
output+="'${marker_buf}' "
else
output+="${marker_buf} "
fi
collecting=false
marker_buf=""
if [[ "$word" == "-m" || "$word" == "-k" ]]; then
output+="${word} "
collecting=true
else
output+="${word} "
fi
else
if [[ -n "$marker_buf" ]]; then
marker_buf+=" ${word}"
else
marker_buf="${word}"
fi
fi
elif [[ "$word" == "-m" || "$word" == "-k" ]]; then
output+="${word} "
collecting=true
marker_buf=""
else
output+="${word} "
fi
done
if $collecting && [[ -n "$marker_buf" ]]; then
if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
output+="'${marker_buf}'"
else
output+="${marker_buf}"
fi
fi
echo "${output% }"
}
apply_intel_test_overrides() {
local cmds="$1"
# Placeholder for Intel-specific exclusions/overrides.
echo "$cmds"
}
is_yaml_file() {
local p="$1"
[[ -f "$p" && "$p" == *.yaml ]]
}
extract_yaml_commands() {
local yaml_path="$1"
awk '
$1 == "commands:" { in_cmds=1; next }
in_cmds && $0 ~ /^[[:space:]]*-[[:space:]]/ {
sub(/^[[:space:]]*-[[:space:]]/, "");
print;
next
}
in_cmds && $0 ~ /^[^[:space:]]/ { exit }
' "$yaml_path"
}
###############################################################################
# Main
###############################################################################
default_image_name="${REGISTRY}/${REPO}:${BUILDKITE_COMMIT}-xpu"
#default_image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-xpu"
image_name="${IMAGE_TAG_XPU:-${default_image_name}}"
container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
# ---- Command source selection ----
commands=""
if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
commands="${VLLM_TEST_COMMANDS}"
echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
elif [[ $# -gt 0 ]]; then
all_yaml=true
for arg in "$@"; do
if ! is_yaml_file "$arg"; then
all_yaml=false
break
fi
done
if $all_yaml; then
for yaml in "$@"; do
mapfile -t COMMANDS < <(extract_yaml_commands "$yaml")
if [[ ${#COMMANDS[@]} -eq 0 ]]; then
echo "Error: No commands found in ${yaml}" >&2
exit 1
fi
for cmd in "${COMMANDS[@]}"; do
if [[ -z "$commands" ]]; then
commands="${cmd}"
else
commands+=" && ${cmd}"
fi
done
done
echo "Commands sourced from YAML files: $*"
else
commands="$*"
echo "Commands sourced from positional args (legacy mode)"
fi
else
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEFAULT_YAML="${SCRIPT_DIR}/intel-test.yaml"
if [[ ! -f "${DEFAULT_YAML}" ]]; then
echo "Error: YAML file not found: ${DEFAULT_YAML}" >&2
exit 1
fi
mapfile -t COMMANDS < <(extract_yaml_commands "${DEFAULT_YAML}")
if [[ ${#COMMANDS[@]} -eq 0 ]]; then
echo "Error: No commands found in ${DEFAULT_YAML}" >&2
exit 1
fi
for cmd in "${COMMANDS[@]}"; do
if [[ -z "$commands" ]]; then
commands="${cmd}"
else
commands+=" && ${cmd}"
fi
done
echo "Commands sourced from default YAML: ${DEFAULT_YAML}"
fi
if [[ -z "$commands" ]]; then
echo "Error: No test commands provided." >&2
exit 1
fi
echo "Raw commands: $commands"
commands=$(re_quote_pytest_markers "$commands")
echo "After re-quoting: $commands"
commands=$(apply_intel_test_overrides "$commands")
echo "Final commands: $commands"
# Dry-run mode prints final commands and exits before Docker.
if [[ "$DRY_RUN" == "1" ]]; then
echo "DRY_RUN=1 set, skipping Docker execution."
exit 0
fi
# --- Docker housekeeping ---
cleanup_docker
# --- Build or pull test image ---
if [[ -n "${IMAGE_TAG_XPU:-}" ]]; then
echo "Using prebuilt XPU image: ${IMAGE_TAG_XPU}"
docker pull "${IMAGE_TAG_XPU}"
else
echo "Using prebuilt XPU image: ${image_name}"
docker pull "${image_name}"
fi
remove_docker_container() {
docker rm -f "${container_name}" || true
docker image rm -f "${image_name}" || true
docker system prune -f || true
}
trap remove_docker_container EXIT
# --- Single-node job ---
if [[ -z "${ZE_AFFINITY_MASK:-}" ]]; then
echo "Warning: ZE_AFFINITY_MASK is not set. Proceeding without device affinity." >&2
fi
docker run \
--device /dev/dri:/dev/dri \
--net=host \
--ipc=host \
--privileged \
-v /dev/dri/by-path:/dev/dri/by-path \
--entrypoint="" \
-e "HF_TOKEN=${HF_TOKEN:-}" \
-e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \
-e "CMDS=${commands}" \
--name "${container_name}" \
"${image_name}" \
bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"'
+4 -65
View File
@@ -2,27 +2,14 @@
set -ex
# ======== part 0: setup ========
# Upload a single wheel to S3 (rename linux -> manylinux).
# Index generation is handled separately by generate-and-upload-nightly-index.sh.
BUCKET="vllm-wheels"
INDICES_OUTPUT_DIR="indices"
DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
PYTHON=${PYTHON_PROG:=python3} # try to read from env var, otherwise use python3
SUBPATH=$BUILDKITE_COMMIT
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# detect if python3.10+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
if [[ "$has_new_python" -eq 0 ]]; then
# use new python from docker
docker pull python:3-slim
PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
fi
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ========= part 1: collect, rename & upload the wheel ==========
# ========= collect, rename & upload the wheel ==========
# Assume wheels are in artifacts/dist/*.whl
wheel_files=(artifacts/dist/*.whl)
@@ -52,56 +39,8 @@ echo "Renamed wheel to: $wheel"
# Extract the version from the wheel
version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
echo "Version in wheel: $version"
pure_version="${version%%+*}"
echo "Pure version (without variant): $pure_version"
# copy wheel to its own bucket
aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
# ========= part 2: generate and upload indices ==========
# generate indices for all existing wheels in the commit directory
# this script might be run multiple times if there are multiple variants being built
# so we need to guarantee there is little chance for "TOCTOU" issues
# i.e., one process is generating indices while another is uploading a new wheel
# so we need to ensure no time-consuming operations happen below
# list all wheels in the commit directory
echo "Existing wheels on S3:"
aws s3 ls "$S3_COMMIT_PREFIX"
obj_json="objects.json"
aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
mkdir -p "$INDICES_OUTPUT_DIR"
# call script to generate indices for all existing wheels
# this indices have relative paths that could work as long as it is next to the wheel directory in s3
# i.e., the wheels are always in s3://vllm-wheels/<commit>/
# and indices can be placed in /<commit>/, or /nightly/, or /<version>/
alias_args=()
if [[ -n "$DEFAULT_VARIANT_ALIAS" ]]; then
alias_args=(--alias-to-default "$DEFAULT_VARIANT_ALIAS")
fi
# HACK: we do not need regex module here, but it is required by pre-commit hook
# To avoid any external dependency, we simply replace it back to the stdlib re module
sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" "${alias_args[@]}"
# copy indices to /<commit>/ unconditionally
echo "Uploading indices to $S3_COMMIT_PREFIX"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
# copy to /nightly/ only if it is on the main branch and not a PR
if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
echo "Uploading indices to overwrite /nightly/"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
fi
# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
if [[ "$version" != *"dev"* ]]; then
echo "Re-generating indices for /$pure_version/"
rm -rf "${INDICES_OUTPUT_DIR:?}/*"
mkdir -p "$INDICES_OUTPUT_DIR"
# wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" "${alias_args[@]}"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/"
fi
echo "Wheel uploaded. Index generation is handled by a separate step."
+16 -3
View File
@@ -812,7 +812,7 @@ steps:
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s model_executor
- pytest -v -s model_executor -m '(not slow_test)'
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
@@ -1242,7 +1242,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
@@ -1801,6 +1801,19 @@ steps:
- tests/v1/e2e
commands:
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy"
- label: V1 e2e (4xH100-4xMI325) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
agent_pool: mi325_4
optional: true
source_file_dependencies:
- vllm/v1/attention/backends/utils.py
- vllm/v1/worker/gpu_model_runner.py
- tests/v1/e2e/test_hybrid_chunked_prefill.py
commands:
- pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py
- label: V1 Spec Decode # TBD
@@ -2501,7 +2514,7 @@ steps:
- tests/models/
commands:
- TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
@@ -8,8 +8,10 @@ steps:
source_file_dependencies:
- vllm/distributed/eplb
- tests/distributed/test_eplb_algo.py
- tests/distributed/test_eplb_utils.py
commands:
- pytest -v -s distributed/test_eplb_algo.py
- pytest -v -s distributed/test_eplb_utils.py
- label: EPLB Execution
timeout_in_minutes: 20
+1 -1
View File
@@ -13,5 +13,5 @@ steps:
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s model_executor
- pytest -v -s model_executor -m '(not slow_test)'
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
@@ -14,7 +14,7 @@ steps:
- tests/models/
commands:
- TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
# Avoid importing model tests that cause CUDA reinitialization error
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
+3
View File
@@ -9,6 +9,7 @@
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
/vllm/model_executor/layers/mamba @tdoublep
/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516
/vllm/model_executor/model_loader @22quinn
/vllm/model_executor/layers/batch_invariant.py @yewentao256
/vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
@@ -48,6 +49,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/vllm/v1/attention/backends/mla @pavanimajety
/vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety
/vllm/v1/attention/backends/triton_attn.py @tdoublep
/vllm/v1/attention/backends/gdn_attn.py @ZJY0516
/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
/vllm/v1/sample @22quinn @houseroad @njhill
/vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni
@@ -142,6 +144,7 @@ mkdocs.yaml @hmellor
# Kernels
/vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
/vllm/v1/attention/ops/triton_unified_attention.py @tdoublep
/vllm/model_executor/layers/fla @ZJY0516
# ROCm related: specify owner with write access to notify AMD folks for careful code review
/vllm/**/*rocm* @tjtanaa
+30
View File
@@ -234,6 +234,36 @@ pull_request_rules:
add:
- rocm
- name: label-xpu
description: Automatically apply intel-gpu label
conditions:
- label != stale
- or:
- files~=^docker/Dockerfile.xpu
- files~=^\\.buildkite/intel_jobs/
- files=\.buildkite/ci_config_intel.yaml
- files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
- files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
- files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
- files=vllm/distributed/device_communicators/xpu_communicator.py
- files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py
- files=vllm/v1/attention/ops/xpu_mla_sparse.py
- files=vllm/v1/worker/xpu_worker.py
- files=vllm/v1/worker/xpu_model_runner.py
- files=vllm/_xpu_ops.py
- files~=^vllm/lora/ops/xpu_ops
- files=vllm/lora/punica_wrapper/punica_xpu.py
- files=vllm/platforms/xpu.py
- title~=(?i)Intel gpu
- title~=(?i)XPU
- title~=(?i)Intel
- title~=(?i)BMG
- title~=(?i)Arc
actions:
label:
add:
- intel-gpu
- name: label-cpu
description: Automatically apply cpu label
conditions:
+6 -6
View File
@@ -363,7 +363,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
# - sm80 doesn't support fp8 computation
# - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
# so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}")
# marlin arches for other files
cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
@@ -523,12 +523,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif()
# The cutlass_scaled_mm kernels for Geforce Blackwell SM120 (c3x, i.e. CUTLASS 3.x) require
# The cutlass_scaled_mm kernels for Blackwell SM12x (c3x, i.e. CUTLASS 3.x) require
# CUDA 12.8 or later
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0a" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
set(SRCS
@@ -616,12 +616,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif()
endif()
# The nvfp4_scaled_mm_sm120 kernels for Geforce Blackwell SM120 require
# The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require
# CUDA 12.8 or later
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(FP4_ARCHS "12.0a" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
set(SRCS
@@ -1050,7 +1050,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
# - sm80 doesn't support fp8 computation
# - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
# so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}")
# moe marlin arches for other files
cuda_archs_loose_intersection(MARLIN_MOE_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
if (MARLIN_MOE_OTHER_ARCHS)
+2 -8
View File
@@ -546,10 +546,7 @@ def main():
args.prefill_backends = yaml_config.get("prefill_backends", None)
# Check for special modes
if "mode" in yaml_config:
args.mode = yaml_config["mode"]
else:
args.mode = None
args.mode = yaml_config.get("mode", None)
# Batch specs and sizes
# Support both explicit batch_specs and generated batch_spec_ranges
@@ -572,10 +569,7 @@ def main():
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"]
else:
args.batch_sizes = None
args.batch_sizes = yaml_config.get("batch_sizes", None)
# Model config
if "model" in yaml_config:
+2 -3
View File
@@ -627,9 +627,8 @@ class BenchmarkWorker:
need_device_guard = True
with (
torch.accelerator.device_index(self.device_id)
if need_device_guard
else nullcontext()
# Ray restricts each worker to one GPU; use local index 0
torch.accelerator.device_index(0) if need_device_guard else nullcontext()
):
for idx, config in enumerate(tqdm(search_space)):
try:
+4 -4
View File
@@ -32,16 +32,16 @@ endif()
message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;10.0f" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(QUTLASS_ARCHS "10.0f;12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;10.0a;10.3a" "${CUDA_ARCHS}")
cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;12.1a;10.0a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
if(QUTLASS_ARCHS MATCHES "10\\.(0a|3a|0f)")
set(QUTLASS_TARGET_CC 100)
elseif(QUTLASS_ARCHS MATCHES "12\\.0a")
elseif(QUTLASS_ARCHS MATCHES "12\\.[01][af]?")
set(QUTLASS_TARGET_CC 120)
else()
message(FATAL_ERROR "[QUTLASS] internal error parsing CUDA_ARCHS='${QUTLASS_ARCHS}'.")
@@ -96,7 +96,7 @@ else()
"[QUTLASS] Skipping build: CUDA 12.8 or newer is required (found ${CMAKE_CUDA_COMPILER_VERSION}).")
else()
message(STATUS
"[QUTLASS] Skipping build: no supported arch (12.0a / 10.0a) found in "
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
"CUDA_ARCHS='${CUDA_ARCHS}'.")
endif()
endif()
+37 -2
View File
@@ -355,8 +355,11 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
list(REMOVE_DUPLICATES _PTX_ARCHS)
list(REMOVE_DUPLICATES _SRC_CUDA_ARCHS)
# If x.0a or x.0f is in SRC_CUDA_ARCHS and x.0 is in CUDA_ARCHS then we should
# remove x.0a or x.0f from SRC_CUDA_ARCHS and add x.0a or x.0f to _CUDA_ARCHS
# Handle architecture-specific suffixes (a/f) for SRC entries.
# First try exact base match (x.y), then cross-suffix match (x.ya / x.yf).
# For 'f' (family) suffix: if no exact/cross match, fall back to major-version
# match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
# family. The output uses TGT's value to preserve the user's compilation flags.
set(_CUDA_ARCHS)
foreach(_arch ${_SRC_CUDA_ARCHS})
if(_arch MATCHES "[af]$")
@@ -365,6 +368,38 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
list(APPEND _CUDA_ARCHS "${_arch}")
elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
list(APPEND _CUDA_ARCHS "${_base}a")
elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}f")
list(APPEND _CUDA_ARCHS "${_base}f")
elseif(_arch MATCHES "f$")
# Family suffix: match any TGT entry in the same major version family.
string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" _src_major "${_base}")
foreach(_tgt ${_TGT_CUDA_ARCHS})
string(REGEX REPLACE "[af]$" "" _tgt_base "${_tgt}")
string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" _tgt_major "${_tgt_base}")
if(_tgt_major STREQUAL _src_major)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_tgt}")
list(APPEND _CUDA_ARCHS "${_tgt}")
break()
endif()
endforeach()
endif()
endif()
endforeach()
# Symmetric handling: if TGT has x.ya/f and SRC has x.y (without suffix),
# preserve TGT's suffix in the output.
set(_tgt_copy ${_TGT_CUDA_ARCHS})
foreach(_arch ${_tgt_copy})
if(_arch MATCHES "[af]$")
string(REGEX REPLACE "[af]$" "" _base "${_arch}")
if ("${_base}" IN_LIST _SRC_CUDA_ARCHS)
list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_arch}")
list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_base}")
list(APPEND _CUDA_ARCHS "${_arch}")
endif()
endif()
endforeach()
+1 -1
View File
@@ -13,7 +13,7 @@
const int4 *__restrict__ b_bias_ptr, \
const float *__restrict__ a_scales_ptr, \
const int4 *__restrict__ scales_ptr, \
const uint16_t *__restrict__ global_scale_ptr, \
const float *__restrict__ global_scale_ptr, \
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
const int32_t *__restrict__ sorted_token_ids_ptr, \
const int32_t *__restrict__ expert_ids_ptr, \
+22 -13
View File
@@ -260,7 +260,7 @@ __global__ void Marlin(
// fp16 quantization scales. shape (k/groupsize, n)
const int4* __restrict__ scales_ptr,
// fp16 global scale (for nvfp4// only)
const uint16_t* __restrict__ global_scale_ptr,
const float* __restrict__ global_scale_ptr,
// 4bit packed zero-points of shape
// (k/groupsize, n/pack_factor)
const int4* __restrict__ zp_ptr,
@@ -308,7 +308,14 @@ __global__ void Marlin(
constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id();
static constexpr auto num_bits =
vllm::ScalarType::from_id(b_type_id).size_bits();
// Disable use_fp16_accum for NVFP4 and cases when group_size == -1 &&
// num_bits == 4
constexpr bool use_fp16_accum =
a_type_id == vllm::kFloat16.id() &&
(!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) &&
!(group_blocks == -1 && num_bits == 4));
#else
constexpr bool use_fp16_accum = false;
#endif
@@ -357,7 +364,7 @@ __global__ void Marlin(
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
has_zp && !is_zp_float && !(b_type == vllm::kU8);
c_scalar_t2 global_scale;
float global_scale_f32 = 1.0f;
constexpr bool has_act_order = group_blocks == 0;
@@ -507,11 +514,12 @@ __global__ void Marlin(
if (mul_topk_weights) {
idx = idx < prob_m_top_k ? idx : 0;
c_scalar_t2 topk_weight_val =
Cdtype::num2num2(Cdtype::float2num(topk_weights_ptr[idx]));
float topk_weight_tmp = topk_weights_ptr[idx];
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
topk_weight_val = __hmul2(topk_weight_val, global_scale);
topk_weight_tmp *= global_scale_f32;
}
c_scalar_t2 topk_weight_val =
Cdtype::num2num2(Cdtype::float2num(topk_weight_tmp));
sh_block_topk_weights[threadIdx.x] = topk_weight_val;
}
}
@@ -532,8 +540,7 @@ __global__ void Marlin(
expert_id = expert_ids_ptr[block_id];
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
uint16_t val = global_scale_ptr[expert_id];
global_scale = Cdtype::num2num2(*reinterpret_cast<c_scalar_t*>(&val));
global_scale_f32 = global_scale_ptr[expert_id];
}
B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4);
@@ -1784,6 +1791,13 @@ __global__ void Marlin(
// We first reorder in shared memory to guarantee the most efficient final
// global write patterns
auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
if (!mul_topk_weights) {
c0 *= global_scale_f32;
c1 *= global_scale_f32;
}
}
c_scalar_t2 res =
Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
@@ -1800,11 +1814,6 @@ __global__ void Marlin(
res = __hmul2(res, tmp_scale);
}
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
if (!mul_topk_weights) {
res = __hmul2(res, global_scale);
}
}
if (has_bias && last) {
c_scalar_t2 tmp_bias = b_bias[0];
if constexpr (m_block_size_8) {
+4 -4
View File
@@ -382,7 +382,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
const int4* bias_ptr = (const int4*)b_bias;
const float* a_s_ptr = (const float*)a_s;
const int4* b_s_ptr = (const int4*)b_s;
const uint16_t* g_s_ptr = (const uint16_t*)g_s;
const float* g_s_ptr = (const float*)g_s;
const int4* zp_ptr = (const int4*)zp;
const int* g_idx_ptr = (const int*)g_idx;
const int* perm_ptr = (const int*)perm;
@@ -759,7 +759,7 @@ torch::Tensor moe_wna16_marlin_gemm(
TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
"global_scale can only be used for nvfp4 format.");
} else {
global_scale = torch::empty({0}, options);
global_scale = torch::empty({0}, options_fp32);
TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
"the global_scale parameter must be passed for nvfp4 format.");
}
@@ -842,8 +842,8 @@ torch::Tensor moe_wna16_marlin_gemm(
TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
"scalar type of a_scales must be float");
TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
"scalar type of global_scale must be the same with c");
TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
"scalar type of global_scale must be float");
if (a_type.size_bits() == 16) {
TORCH_CHECK(
a.scalar_type() == c.scalar_type(),
-3
View File
@@ -189,10 +189,7 @@ __device__ __forceinline__ void cp_async_wait<0>() {
}
__device__ __forceinline__ float clip(float v, float mmin, float mmax) {
#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
return fminf(mmax, fmaxf(v, mmin));
#else
#endif
}
__device__ __forceinline__ __nv_bfloat16 clip(__nv_bfloat16 v,
+1 -1
View File
@@ -13,7 +13,7 @@
const int4 *__restrict__ b_bias_ptr, \
const float *__restrict__ a_scales_ptr, \
const int4 *__restrict__ scales_ptr, \
const uint16_t *__restrict__ global_scale_ptr, \
const float *__restrict__ global_scale_ptr, \
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \
bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \
+5 -5
View File
@@ -57,7 +57,7 @@ torch::Tensor marlin_gemm(
int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
bool is_zp_float) {
TORCH_CHECK_NOT_IMPLEMENTED(false,
"marlin_gemm(..) requires CUDA_ARCH >= 8.0");
"marlin_gemm(..) requires CUDA_ARCH >= 7.5");
return torch::empty({1, 1});
}
@@ -356,7 +356,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
const int4* bias_ptr = (const int4*)b_bias;
const float* a_s_ptr = (const float*)a_s;
const int4* b_s_ptr = (const int4*)b_s;
const uint16_t* g_s_ptr = (const uint16_t*)g_s;
const float* g_s_ptr = (const float*)g_s;
const int4* zp_ptr = (const int4*)zp;
const int* g_idx_ptr = (const int*)g_idx;
@@ -751,7 +751,7 @@ torch::Tensor marlin_gemm(
TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
"global_scale can only be used for nvfp4 format.");
} else {
global_scale = torch::empty({0}, options);
global_scale = torch::empty({0}, options_fp32);
TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
"the global_scale parameter must be passed for nvfp4 format.");
}
@@ -832,8 +832,8 @@ torch::Tensor marlin_gemm(
TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
"scalar type of a_scales must be float");
TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
"scalar type of global_scale must be the same with c");
TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
"scalar type of global_scale must be float");
if (a_type.size_bits() == 16) {
TORCH_CHECK(
a.scalar_type() == c.scalar_type(),
+15 -10
View File
@@ -251,8 +251,8 @@ __global__ void Marlin(
const float* __restrict__ a_scales_ptr,
// fp16 quantization scales. shape (k/groupsize, n)
const int4* __restrict__ scales_ptr,
// fp16 global scale (for nvfp4// only)
const uint16_t* __restrict__ global_scale_ptr,
// float global scale (for nvfp4// only)
const float* __restrict__ global_scale_ptr,
// 4bit packed zero-points of shape
// (k/groupsize, n/pack_factor)
const int4* __restrict__ zp_ptr,
@@ -292,7 +292,13 @@ __global__ void Marlin(
#endif
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id();
constexpr auto num_bits = vllm::ScalarType::from_id(b_type_id).size_bits();
// Disable use_fp16_accum for NVFP4 and cases when group_size == -1 &&
// num_bits == 4
constexpr bool use_fp16_accum =
a_type_id == vllm::kFloat16.id() &&
(!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) &&
!(group_blocks == -1 && num_bits == 4));
#else
constexpr bool use_fp16_accum = false;
#endif
@@ -342,11 +348,10 @@ __global__ void Marlin(
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
has_zp && !is_zp_float && !(b_type == vllm::kU8);
c_scalar_t2 global_scale;
float global_scale_f32 = 1.0f;
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
uint16_t val = global_scale_ptr[0];
global_scale = Cdtype::num2num2(*reinterpret_cast<c_scalar_t*>(&val));
global_scale_f32 = global_scale_ptr[0];
}
constexpr bool has_act_order = group_blocks == 0;
@@ -1644,6 +1649,10 @@ __global__ void Marlin(
// We first reorder in shared memory to guarantee the most efficient final
// global write patterns
auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
c0 *= global_scale_f32;
c1 *= global_scale_f32;
}
c_scalar_t2 res =
Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
@@ -1659,10 +1668,6 @@ __global__ void Marlin(
}
res = __hmul2(res, tmp_scale);
}
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
res = __hmul2(res, global_scale);
}
if (has_bias && last) {
c_scalar_t2 tmp_bias = b_bias[0];
if constexpr (m_block_size_8) {
+19
View File
@@ -596,6 +596,25 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
&& flashinfer show-config
# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
# which fails without internet. This step caches them at build time.
RUN python3 <<'PYEOF'
from flashinfer.jit import env as jit_env
from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
from flashinfer.artifacts import ArtifactPath, CheckSumHash
download_trtllm_headers(
'bmm',
jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
ArtifactPath.TRTLLM_GEN_BMM,
get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
)
print('FlashInfer TRTLLM BMM headers downloaded successfully')
PYEOF
# ============================================================
# OPENAI API SERVER DEPENDENCIES
# Pre-install these to avoid reinstalling on every vLLM wheel rebuild
+8 -2
View File
@@ -29,8 +29,11 @@ RUN if [ "$USE_SCCACHE" != "1" ]; then \
rm -f "$(which sccache)" || true; \
fi
# Install UV
RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR="/usr/local/bin" sh
# Install UV — download first, then run, so a curl failure is not masked by the pipe
RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \
&& env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \
&& rm -f /tmp/uv-install.sh \
&& uv --version
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
# Reference: https://github.com/astral-sh/uv/pull/1694
@@ -386,6 +389,9 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
# will not be imported by other tests
RUN mkdir src && mv vllm src/vllm
# This is a workaround to ensure pytest exits with the correct status code in CI tests.
RUN echo "import os\n\ndef pytest_sessionfinish(session, exitstatus):\n os._exit(int(exitstatus))" > /vllm-workspace/conftest.py
# -----------------------
# Final vLLM image
FROM base AS final
+13 -7
View File
@@ -1,7 +1,7 @@
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-complete
ARG TRITON_BRANCH="57c693b6"
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.1-complete
ARG TRITON_BRANCH="ba5c1517"
ARG TRITON_REPO="https://github.com/ROCm/triton.git"
ARG PYTORCH_BRANCH="89075173"
ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17
ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
ARG PYTORCH_VISION_BRANCH="v0.24.1"
ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
@@ -114,6 +114,8 @@ ARG TRITON_REPO
RUN git clone ${TRITON_REPO}
RUN cd triton \
&& git checkout ${TRITON_BRANCH} \
&& git config --global user.email "you@example.com" && git config --global user.name "Your Name" \
&& git cherry-pick 555d04f \
&& if [ ! -f setup.py ]; then cd python; fi \
&& python3 setup.py bdist_wheel --dist-dir=dist \
&& mkdir -p /app/install && cp dist/*.whl /app/install
@@ -142,10 +144,14 @@ ARG PYTORCH_VISION_REPO
ARG PYTORCH_AUDIO_REPO
ARG USE_SCCACHE
RUN apt-get update && apt-get install -y pkg-config liblzma-dev
RUN git clone ${PYTORCH_REPO} pytorch
RUN cd pytorch && git checkout ${PYTORCH_BRANCH} \
&& pip install -r requirements.txt && git submodule update --init --recursive \
&& python3 tools/amd_build/build_amd.py \
RUN cd pytorch && git checkout ${PYTORCH_BRANCH}
RUN cd pytorch \
&& pip install -r requirements.txt && git submodule update --init --recursive
RUN cd pytorch/third_party/kineto \
&& git remote add rocm https://github.com/ROCm/kineto && git fetch rocm && git checkout 2d73be3
RUN cd pytorch && python3 tools/amd_build/build_amd.py \
&& if [ "$USE_SCCACHE" = "1" ]; then \
export HIP_CLANG_PATH=/opt/sccache-wrappers \
&& export CMAKE_C_COMPILER_LAUNCHER=sccache \
@@ -239,7 +245,7 @@ RUN pip install pyyaml && cd aiter \
export HIP_CLANG_PATH=/opt/sccache-wrappers \
&& sccache --show-stats; \
fi \
&& GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
&& PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
&& if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
&& ls /app/aiter/dist/*.whl
RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
+2
View File
@@ -17,6 +17,8 @@ Before you begin, ensure that you have the following:
## Installing the chart
This guide uses the Helm chart at [examples/online_serving/chart-helm](../../../examples/online_serving/chart-helm).
To install the chart with the release name `test-vllm`:
```bash
+2 -2
View File
@@ -173,9 +173,9 @@ Priority is **1 = highest** (tried first).
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | 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_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | All | N/A |
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
| `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 |
+10 -10
View File
@@ -244,12 +244,12 @@ response = client.chat.completions.create(
Some models, such as [Qwen3](https://qwen.readthedocs.io/en/latest/getting_started/quickstart.html#thinking-budget), [DeepSeek](https://www.alibabacloud.com/help/en/model-studio/deep-thinking), and [Nemotron3](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16), support a thinking budget that limits the maximum number of tokens used for reasoning.
Token counting starts from `think_start_str`. Once the reasoning token count reaches the configured `thinking_token_budget`, vLLM forces the model to produce `think_end_str`, effectively terminating the reasoning block.
Token counting starts from `reasoning_start_str`. Once the reasoning token count reaches the configured `thinking_token_budget`, vLLM forces the model to produce `reasoning_end_str`, effectively terminating the reasoning block.
To use this feature:
- `--reasoning-parser` enables reasoning extraction.
- `--reasoning-config` defines the reasoning boundary tokens (e.g., `think_start_str`, `think_end_str`).
- `--reasoning-config` defines the reasoning boundary tokens (e.g., `reasoning_start_str`, `reasoning_end_str`).
- `thinking_token_budget` (a sampling parameter) sets the per-request reasoning token limit.
If `thinking_token_budget` is not specified, no explicit reasoning limit is applied beyond normal generation constraints such as `max_tokens`.
@@ -257,20 +257,20 @@ If `thinking_token_budget` is not specified, no explicit reasoning limit is appl
`--reasoning-config` accepts a JSON object corresponding to
[ReasoningConfig][vllm.config.ReasoningConfig] with the following fields:
| Field | Type | Description |
|-------------------|----------------|--------------------------------------------------|
| `think_start_str` | `str \| null` | String that marks the start of reasoning content |
| `think_end_str` | `str \| null` | String that marks the end of reasoning content |
| Field | Type | Description |
|-----------------------|----------------|--------------------------------------------------|
| `reasoning_start_str` | `str \| null` | String that marks the start of reasoning content |
| `reasoning_end_str` | `str \| null` | String that marks the end of reasoning content |
!!! note
`think_end_str` can include a transition phrase before the think end token. For example, setting `think_end_str` to `"I have to give the solution based on the thinking directly now.</think>"` instructs the model to emit that phrase when the budget is exhausted, making the reasoning termination more natural.
`reasoning_end_str` can include a transition phrase before the reasoning end token. For example, setting `reasoning_end_str` to `"I have to give the solution based on the reasoning directly now.</think>"` instructs the model to emit that phrase when the budget is exhausted, making the reasoning termination more natural.
### Online Serving
```bash
vllm serve Qwen/Qwen3-0.6B \
--reasoning-parser qwen3 \
--reasoning-config '{"think_start_str": "<think>", "think_end_str": "I have to give the solution based on the thinking directly now.</think>"}'
--reasoning-config '{"reasoning_start_str": "<think>", "reasoning_end_str": "I have to give the solution based on the reasoning directly now.</think>"}'
```
Then make a request with `thinking_token_budget` to limit the reasoning tokens:
@@ -298,8 +298,8 @@ from vllm.config import ReasoningConfig
llm = LLM(
model="Qwen/Qwen3-0.6B",
reasoning_config=ReasoningConfig(
think_start_str="<think>",
think_end_str="I have to give the solution based on the thinking directly now.</think>",
reasoning_start_str="<think>",
reasoning_end_str="I have to give the solution based on the thinking directly now.</think>",
),
)
@@ -172,8 +172,11 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.15.0/rocm700
--8<-- [end:build-wheel-from-source]
--8<-- [start:pre-built-images]
vLLM offers an official Docker image for deployment.
The image can be used to run OpenAI compatible server and is available on Docker Hub as [vllm/vllm-openai-rocm](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags).
vLLM offers official Docker images for deployment.
The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-rocm](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags).
- `vllm/vllm-openai-rocm:latest` — stable release
- `vllm/vllm-openai-rocm:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
```bash
docker run --rm \
@@ -186,30 +189,18 @@ docker run --rm \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8000:8000 \
--ipc=host \
vllm/vllm-openai-rocm:latest \
vllm/vllm-openai-rocm:<tag> \
--model Qwen/Qwen3-0.6B
```
#### Use AMD's Docker Images
#### Use AMD's Docker Images (Deprecated)
Prior to January 20th, 2026 when the official docker images are available on [upstream vLLM docker hub](https://hub.docker.com/v2/repositories/vllm/vllm-openai-rocm/tags/), the [AMD Infinity hub for vLLM](https://hub.docker.com/r/rocm/vllm/tags) offers a prebuilt, optimized
!!! warning "Deprecated"
AMD's Docker images (`rocm/vllm` and `rocm/vllm-dev`) are deprecated in favor of the official vLLM Docker images above (`vllm/vllm-openai-rocm`). Please migrate to the official images.
Prior to January 20th, 2026 when the official docker images became available on [upstream vLLM docker hub](https://hub.docker.com/v2/repositories/vllm/vllm-openai-rocm/tags/), the [AMD Infinity hub for vLLM](https://hub.docker.com/r/rocm/vllm/tags) offered a prebuilt, optimized
docker image designed for validating inference performance on the AMD Instinct MI300X™ accelerator.
AMD also offers nightly prebuilt docker image from [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev), which has vLLM and all its dependencies installed. The entrypoint of this docker image is `/bin/bash` (different from the vLLM's Official Docker Image).
```bash
docker pull rocm/vllm-dev:nightly # to get the latest image
docker run -it --rm \
--network=host \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v <path/to/your/models>:/app/models \
-e HF_HOME="/app/models" \
rocm/vllm-dev:nightly
```
AMD also offered nightly prebuilt docker image from [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev), which has vLLM and all its dependencies installed. The entrypoint of this docker image is `/bin/bash` (different from the vLLM's Official Docker Image).
!!! tip
Please check [LLM inference performance validation on AMD Instinct MI300X](https://rocm.docs.amd.com/en/latest/how-to/performance-validation/mi300x/vllm-benchmark.html)
+4 -1
View File
@@ -56,9 +56,12 @@ This guide will help you quickly get started with vLLM to perform:
!!! note
It currently supports Python 3.12, ROCm 7.0 and `glibc >= 2.35`.
!!! note
!!! note
Note that, previously, docker images were published using AMD's docker release pipeline and were located `rocm/vllm-dev`. This is being deprecated by using vLLM's docker release pipeline.
!!! tip
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
=== "Google TPU"
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
+1 -1
View File
@@ -153,7 +153,7 @@ class MarkdownFormatter(HelpFormatter):
heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
self._markdown_output.append(heading_md)
if action.choices or isinstance(action.metavar, (list, tuple)):
if action.choices or isinstance(action.metavar, list | tuple):
choices_iterable = action.choices or action.metavar
choices = f"`{'`, `'.join(str(c) for c in choices_iterable)}`"
self._markdown_output.append(f": Possible choices: {choices}\n\n")
+21 -2
View File
@@ -15,7 +15,7 @@ Many classification models support both (sequence) classification and token clas
!!! note
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (classify) is not
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (classify) is not
what you want, you need to manually specify it via `PoolerConfig(task="token_classify")` offline or
`--pooler-config.task token_classify` online.
@@ -29,6 +29,12 @@ Offline: [examples/pooling/token_classify/ner_offline.py](../../../examples/pool
Online: [examples/pooling/token_classify/ner_online.py](../../../examples/pooling/token_classify/ner_online.py)
### Forced Alignment
Forced alignment takes audio and reference text as input and produces word-level timestamps.
Offline: [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py)
### Sparse retrieval (lexical matching)
The BAAI/bge-m3 model leverages token classification for sparse retrieval. For more information, see [this page](specific_models.md#baaibge-m3).
@@ -43,12 +49,25 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m
| `Qwen3ForTokenClassification`<sup>C</sup> | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | |
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./README.md#model-conversion))
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./README.md#model-conversion))
\* Feature support is the same as that of the original model.
If your model is not in the above list, we will try to automatically convert the model using
[as_seq_cls_model][vllm.model_executor.models.adapters.as_seq_cls_model]. By default, the class probabilities are extracted from the softmaxed hidden state corresponding to the last token.
### Multimodal Models
!!! note
For more information about multimodal models inputs, see [this page](../supported_models.md#list-of-multimodal-language-models).
| Architecture | Models | Inputs | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) |
| --------------------------------------------- | ------------------- | ----------------- | ------------------------------------------ | ------------------------------ | ------------------------------------------ |
| `Qwen3ASRForcedAlignerForTokenClassification` | Qwen3-ForcedAligner | T + A<sup>+</sup> | `Qwen/Qwen3-ForcedAligner-0.6B` (see note) | | ✅︎ |
!!! note
Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`.
Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py).
### As Reward Models
Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md).
+12
View File
@@ -231,6 +231,18 @@ The most effective approach is to deploy vLLM behind a reverse proxy (such as ng
- Blocks all other endpoints, including the unauthenticated inference and operational control endpoints
- Implements additional authentication, rate limiting, and logging at the proxy layer
## Request Parameter Resource Limits
Certain API request parameters can have a large impact on resource consumption and may be abused to exhaust server resources. The `n` parameter in the `/v1/completions` and `/v1/chat/completions` endpoints controls how many independent output sequences are generated per request. A very large value causes the engine to allocate memory, CPU, and GPU time proportional to `n`, which can lead to out-of-memory conditions on the host and block the server from processing other requests.
To mitigate this, vLLM enforces a configurable upper bound on the `n` parameter via the `VLLM_MAX_N_SEQUENCES` environment variable (default: **16384**). Requests exceeding this limit are rejected before reaching the engine.
### Recommendations
- **Public-facing deployments:** Consider setting `VLLM_MAX_N_SEQUENCES` to a value appropriate for your workload (e.g., `64` or `128`) to limit the blast radius of a single request.
- **Reverse proxy layer:** In addition to vLLM's built-in limit, consider enforcing request body validation and rate limiting at your reverse proxy to further constrain abusive payloads.
- **Monitoring:** Monitor per-request resource consumption to detect anomalous patterns that may indicate abuse.
## Tool Server and MCP Security
vLLM supports connecting to external tool servers via the `--tool-server` argument. This enables models to call tools through the Responses API (`/v1/responses`). Tool server support works with all models — it is not limited to specific model architectures.
@@ -0,0 +1,90 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from Qwen3-ForcedAligner inference:
# https://github.com/QwenLM/Qwen3-ASR
"""
Offline forced alignment example using Qwen3-ForcedAligner-0.6B.
Forced alignment takes audio and reference text as input and produces
word-level timestamps. The model predicts a time bin at each <timestamp>
token position; multiplying by ``timestamp_segment_time`` gives milliseconds.
Usage::
python forced_alignment_offline.py \
--model Qwen/Qwen3-ForcedAligner-0.6B
"""
from argparse import Namespace
import numpy as np
from vllm import LLM, EngineArgs
from vllm.utils.argparse_utils import FlexibleArgumentParser
def parse_args():
parser = FlexibleArgumentParser()
parser = EngineArgs.add_cli_args(parser)
parser.set_defaults(
model="Qwen/Qwen3-ForcedAligner-0.6B",
runner="pooling",
enforce_eager=True,
hf_overrides={"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]},
)
return parser.parse_args()
def build_prompt(words: list[str]) -> str:
"""Build the forced alignment prompt from a word list.
Format: <|audio_start|><|audio_pad|><|audio_end|>
word1<timestamp><timestamp>word2<timestamp><timestamp>...
"""
body = "<timestamp><timestamp>".join(words) + "<timestamp><timestamp>"
return f"<|audio_start|><|audio_pad|><|audio_end|>{body}"
def main(args: Namespace):
llm = LLM(**vars(args))
config = llm.llm_engine.vllm_config.model_config.hf_config
timestamp_token_id = config.timestamp_token_id
timestamp_segment_time = config.timestamp_segment_time
# Example: align these words against a 5-second audio clip
words = ["Hello", "world"]
prompt = build_prompt(words)
# Use a 5-second silent audio as placeholder (replace with real audio)
sample_rate = 16000
audio = np.zeros(sample_rate * 5, dtype=np.float32)
outputs = llm.encode(
[{"prompt": prompt, "multi_modal_data": {"audio": audio}}],
pooling_task="token_classify",
)
for output in outputs:
logits = output.outputs.data # [num_tokens, classify_num]
predictions = logits.argmax(dim=-1)
token_ids = output.prompt_token_ids
# Extract timestamps at <timestamp> positions
ts_predictions = [
pred.item() * timestamp_segment_time
for tid, pred in zip(token_ids, predictions)
if tid == timestamp_token_id
]
# Pair up start/end times per word
for i, word in enumerate(words):
start_ms = ts_predictions[i * 2]
end_ms = ts_predictions[i * 2 + 1]
print(f"{word:15s} {start_ms / 1000:.3f}s - {end_ms / 1000:.3f}s")
if __name__ == "__main__":
args = parse_args()
main(args)
+154
View File
@@ -0,0 +1,154 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock
import pytest
import torch
from vllm.distributed.eplb.eplb_state import (
_commit_eplb_maps,
_commit_eplb_maps_for_layer,
)
def _make_model_state(
phy2log: torch.Tensor,
log2phy: torch.Tensor,
logcnt: torch.Tensor,
) -> MagicMock:
"""Build a minimal EplbModelState mock with only the three map tensors."""
state = MagicMock()
state.physical_to_logical_map = phy2log
state.logical_to_physical_map = log2phy
state.logical_replica_count = logcnt
return state
def test_commit_eplb_maps_shape_change():
"""
The normal path copies the physical_to_logical map in-place. When the number of
physical experts changes, the old map should be replaced entirely.
"""
num_layers, num_logical, num_physical = 2, 4, 6
max_replicas = 3
# Build current state tensors
model_state = _make_model_state(
phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
log2phy=torch.full(
(num_layers, num_logical, max_replicas), -1, dtype=torch.long
),
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
)
# The new map has two more physical experts. These new physical experts will
# automatically map to the first two logical experts
new_phy2log_larger = (
(torch.arange(num_physical + 2, dtype=torch.long) % num_logical)
.unsqueeze(0)
.expand(num_layers, -1)
)
_commit_eplb_maps(model_state, new_phy2log_larger)
# Check that the number of physical experts has been updated and that the values
# match
assert model_state.physical_to_logical_map.shape[1] == num_physical + 2
assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger)
def test_commit_eplb_maps_for_layer_logical_padding():
"""
Test that logical_to_physical_map is padded with -1 to fill the
pre-allocated slots when the new map has fewer replicas than the max.
"""
num_layers, num_logical, num_physical = 2, 4, 6
max_replicas = 3
model_state = _make_model_state(
phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
log2phy=torch.full(
(num_layers, num_logical, max_replicas), -1, dtype=torch.long
),
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
)
new_phy2log = (
(torch.arange(num_physical, dtype=torch.long) % num_logical)
.unsqueeze(0)
.expand(num_layers, -1)
.contiguous()
)
layer = 0
_commit_eplb_maps_for_layer(model_state, new_phy2log, layer)
assert torch.all(model_state.logical_to_physical_map[layer, :, 2] == -1)
def test_commit_eplb_maps_for_layer_shape_assert():
"""Test that a mismatched number of physical experts triggers an assertion error."""
num_layers, num_logical, num_physical = 2, 4, 6
model_state = _make_model_state(
phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
log2phy=torch.full((num_layers, num_logical, 2), -1, dtype=torch.long),
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
)
bad_phy2log = torch.zeros(num_layers, num_physical + 1, dtype=torch.long)
with pytest.raises(AssertionError):
_commit_eplb_maps_for_layer(model_state, bad_phy2log, layer=0)
def test_commit_eplb_maps():
"""Test that all values are copied correctly into model_state."""
num_layers, num_logical, num_physical, max_replicas = 2, 3, 4, 2
model_state = _make_model_state(
phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
log2phy=torch.full(
(num_layers, num_logical, max_replicas), -1, dtype=torch.long
),
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
)
new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long)
new_log2phy = torch.tensor(
[[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long
)
new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
_commit_eplb_maps(model_state, new_phy2log)
assert torch.equal(model_state.physical_to_logical_map, new_phy2log)
assert torch.equal(model_state.logical_to_physical_map, new_log2phy)
assert torch.equal(model_state.logical_replica_count, new_logcnt)
def test_commit_eplb_maps_for_layer():
"""Test that only the target layer is updated"""
num_layers, num_logical, max_replicas = 2, 3, 2
original_phy2log = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
model_state = _make_model_state(
phy2log=original_phy2log.clone(),
log2phy=torch.full(
(num_layers, num_logical, max_replicas), -1, dtype=torch.long
),
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
)
new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long)
new_log2phy = torch.tensor(
[[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long
)
new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
_commit_eplb_maps_for_layer(model_state, new_phy2log, layer=0)
# Layer 0 updated
assert torch.equal(model_state.physical_to_logical_map[0], new_phy2log[0])
assert torch.equal(model_state.logical_to_physical_map[0], new_log2phy[0])
assert torch.equal(model_state.logical_replica_count[0], new_logcnt[0])
# Layer 1 untouched
assert torch.equal(model_state.physical_to_logical_map[1], original_phy2log[1])
@@ -64,11 +64,12 @@ async def test_online_audio_in_video(
]
# multi-turn to test mm processor cache as well
for _ in range(2):
for turn in range(2):
chat_completion = await client.chat.completions.create(
model=MODEL_NAME,
messages=messages,
max_tokens=16,
max_tokens=8,
temperature=0.0,
extra_body={
"mm_processor_kwargs": {
"use_audio_in_video": True,
@@ -78,6 +79,12 @@ async def test_online_audio_in_video(
assert len(chat_completion.choices) == 1
choice = chat_completion.choices[0]
print(
f"[DEBUG][single-video] turn={turn} "
f"finish_reason={choice.finish_reason!r} "
f"content={choice.message.content!r} "
f"usage={chat_completion.usage}"
)
assert choice.finish_reason == "length"
@@ -111,11 +118,12 @@ async def test_online_audio_in_video_multi_videos(
]
# multi-turn to test mm processor cache as well
for _ in range(2):
for turn in range(2):
chat_completion = await client.chat.completions.create(
model=MODEL_NAME,
messages=messages,
max_tokens=16,
max_tokens=8,
temperature=0.0,
extra_body={
"mm_processor_kwargs": {
"use_audio_in_video": True,
@@ -125,6 +133,12 @@ async def test_online_audio_in_video_multi_videos(
assert len(chat_completion.choices) == 1
choice = chat_completion.choices[0]
print(
f"[DEBUG][multi-video] turn={turn} "
f"finish_reason={choice.finish_reason!r} "
f"content={choice.message.content!r} "
f"usage={chat_completion.usage}"
)
assert choice.finish_reason == "length"
@@ -1020,3 +1020,114 @@ def test_chat_completion_request_n_parameter_various_values():
assert sampling_params.n == n_value, (
f"Expected n={n_value}, got n={sampling_params.n}"
)
def test_chat_completion_request_n_parameter_exceeds_default_limit(
monkeypatch: pytest.MonkeyPatch,
):
"""Test that n values exceeding the default limit are rejected."""
import vllm.envs as envs
monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
if hasattr(envs.__getattr__, "cache_clear"):
envs.__getattr__.cache_clear()
max_n = envs.VLLM_MAX_N_SEQUENCES
request = ChatCompletionRequest(
model="test-model",
messages=[{"role": "user", "content": "Test"}],
n=max_n + 1,
max_tokens=10,
)
with pytest.raises(ValueError, match="n must be at most"):
request.to_sampling_params(
max_tokens=10,
default_sampling_params={},
)
def test_chat_completion_request_n_parameter_at_limit(
monkeypatch: pytest.MonkeyPatch,
):
"""Test that n at exactly the limit is accepted."""
import vllm.envs as envs
monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
if hasattr(envs.__getattr__, "cache_clear"):
envs.__getattr__.cache_clear()
max_n = envs.VLLM_MAX_N_SEQUENCES
request = ChatCompletionRequest(
model="test-model",
messages=[{"role": "user", "content": "Test"}],
n=max_n,
max_tokens=10,
)
sampling_params = request.to_sampling_params(
max_tokens=10,
default_sampling_params={},
)
assert sampling_params.n == max_n
def test_chat_completion_request_n_parameter_custom_limit(
monkeypatch: pytest.MonkeyPatch,
):
"""Test that VLLM_MAX_N_SEQUENCES env var overrides the default limit."""
import vllm.envs as envs
monkeypatch.setenv("VLLM_MAX_N_SEQUENCES", "128")
if hasattr(envs.__getattr__, "cache_clear"):
envs.__getattr__.cache_clear()
request = ChatCompletionRequest(
model="test-model",
messages=[{"role": "user", "content": "Test"}],
n=128,
max_tokens=10,
)
sampling_params = request.to_sampling_params(
max_tokens=10,
default_sampling_params={},
)
assert sampling_params.n == 128
request_over = ChatCompletionRequest(
model="test-model",
messages=[{"role": "user", "content": "Test"}],
n=129,
max_tokens=10,
)
with pytest.raises(ValueError, match="n must be at most 128"):
request_over.to_sampling_params(
max_tokens=10,
default_sampling_params={},
)
def test_chat_completion_request_n_parameter_massive_value(
monkeypatch: pytest.MonkeyPatch,
):
"""Test that astronomically large n values are rejected (CVE fix)."""
import vllm.envs as envs
monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
if hasattr(envs.__getattr__, "cache_clear"):
envs.__getattr__.cache_clear()
request = ChatCompletionRequest(
model="test-model",
messages=[{"role": "user", "content": "Test"}],
n=100_000_000,
max_tokens=1,
)
with pytest.raises(ValueError, match="n must be at most"):
request.to_sampling_params(
max_tokens=1,
default_sampling_params={},
)
@@ -55,6 +55,7 @@ class MockModelConfig:
skip_tokenizer_init = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
renderer_num_workers: int = 1
def get_diff_sampling_param(self):
return self.diff_sampling_param or {}
@@ -536,6 +536,7 @@ class MockModelConfig:
skip_tokenizer_init: bool = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
renderer_num_workers: int = 1
def get_diff_sampling_param(self):
return self.diff_sampling_param or {}
@@ -54,6 +54,7 @@ class MockModelConfig:
skip_tokenizer_init = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
renderer_num_workers: int = 1
def get_diff_sampling_param(self):
return self.diff_sampling_param or {}
@@ -54,6 +54,7 @@ class MockModelConfig:
skip_tokenizer_init: bool = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
renderer_num_workers: int = 1
def get_diff_sampling_param(self):
return self.diff_sampling_param or {}
@@ -14,13 +14,62 @@ import pytest_asyncio
import soundfile as sf
from tests.utils import RemoteOpenAIServer
from vllm.platforms import current_platform
MODEL_NAME = "openai/whisper-large-v3-turbo"
# Disable prefix caching on ROCm to reduce non-determinism in
# streaming-vs-non-streaming comparisons.
_ROCM_ARGS = ["--no-enable-prefix-caching"] if current_platform.is_rocm() else []
@pytest.fixture(scope="module")
def server():
with RemoteOpenAIServer(MODEL_NAME, []) as remote_server:
def _get_attention_backend_params() -> list[str | None]:
"""Return attention backends to parametrize the server fixture with.
On ROCm, we test multiple backends explicitly:
- None: default auto-selection (ROCM_ATTN for decoder self-attention,
falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for
cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER)
- TRITON_ATTN: always available on ROCm
- ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950
On non-ROCm platforms, we just run with the default backend.
"""
try:
from vllm.platforms import current_platform
if current_platform.is_rocm():
backends: list[str | None] = [None, "TRITON_ATTN"]
from vllm.platforms.rocm import _ON_MI3XX
if _ON_MI3XX:
backends.append("ROCM_AITER_UNIFIED_ATTN")
return backends
except Exception:
pass
return [None]
# Aiter backends need VLLM_ROCM_USE_AITER=1 (and MHA=1 for ROCM_AITER_FA)
# to be enabled in the server subprocess.
_AITER_ENV = {
"VLLM_ROCM_USE_AITER": "1",
"VLLM_ROCM_USE_AITER_MHA": "1",
}
_ATTN_BACKENDS = _get_attention_backend_params()
_ATTN_IDS = [b or "default" for b in _ATTN_BACKENDS]
@pytest.fixture(scope="module", params=_ATTN_BACKENDS, ids=_ATTN_IDS)
def server(request):
args = [*_ROCM_ARGS]
env_dict = None
if request.param is not None:
args += ["--attention-backend", request.param]
if "AITER" in request.param:
env_dict = _AITER_ENV
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
yield remote_server
@@ -57,16 +57,25 @@ def _openai_embed(
return [item["embedding"] for item in resp.json()["data"]]
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)))
def test_single_text_parity(server: RemoteOpenAIServer):
"""A single text should produce identical embeddings via both APIs."""
"""A single text should produce equivalent 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)
# Full-suite BF16 runs can introduce tiny numerical drift even when both
# endpoints are functionally equivalent, so compare semantic equivalence
# instead of exact elementwise equality.
cos = _cosine_sim(v2[0], v1[0])
assert cos > 0.9999, f"single-text parity failed, cosine={cos}"
def test_batch_parity(server: RemoteOpenAIServer):
"""A batch of texts should produce identical embeddings via both APIs,
"""A batch of texts should produce equivalent embeddings via both APIs,
in the same order."""
texts = [
"machine learning",
@@ -76,8 +85,18 @@ def test_batch_parity(server: RemoteOpenAIServer):
v2 = _cohere_embed(server, texts)
v1 = _openai_embed(server, texts)
assert len(v2) == len(v1) == 3
similarities = np.array(
[[_cosine_sim(v2_emb, v1_emb) for v1_emb in v1] for v2_emb in v2]
)
for i in range(3):
np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
assert int(np.argmax(similarities[i])) == i, (
f"batch parity order mismatch at index {i}: "
f"similarities={similarities[i].tolist()}"
)
assert similarities[i, i] > 0.9999, (
f"batch parity failed at index {i}, cosine={similarities[i, i]}"
)
def test_token_count_parity(server: RemoteOpenAIServer):
@@ -6,8 +6,11 @@ import pytest
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
from vllm.entrypoints.pooling.embed.protocol import (
CohereEmbedContent,
CohereEmbedInput,
CohereEmbedRequest,
)
from vllm.entrypoints.pooling.typing import PoolingServeContext
class TestResolveTruncation:
@@ -206,3 +209,116 @@ class TestValidateInputType:
handler = self._make_handler({"a": "", "b": ""})
with pytest.raises(ValueError, match="Supported values: a, b"):
handler._validate_input_type("z")
class TestPreProcessCohereOnline:
"""Unit tests for EmbedIOProcessor._pre_process_cohere_online."""
@staticmethod
def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
return PoolingServeContext(
request=CohereEmbedRequest(model="test", **request_kwargs),
model_name="test",
request_id="embd-test",
)
@staticmethod
def _make_handler():
handler = object.__new__(EmbedIOProcessor)
handler._validate_input_type = lambda _input_type: None
return handler
def test_text_only_without_task_prefix_uses_completion_path(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"])
calls: list[tuple[str, object]] = []
def preprocess_completion(request, prompt_input, prompt_embeds):
calls.append(("completion", prompt_input))
return ["completion"]
handler._get_task_instruction_prefix = lambda _input_type: None
handler._has_chat_template = lambda: False
handler._preprocess_completion_online = preprocess_completion
handler._batch_render_chat = lambda *_args, **_kwargs: (
pytest.fail("text-only request should not require chat rendering")
)
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["completion"]
assert calls == [("completion", ["hello"])]
def test_text_only_falls_back_to_prefixed_completion_without_template(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"], input_type="query")
calls: list[tuple[str, object]] = []
def preprocess_completion(request, prompt_input, prompt_embeds):
calls.append(("completion", prompt_input))
return ["fallback"]
handler._get_task_instruction_prefix = lambda _input_type: "query: "
handler._has_chat_template = lambda: False
handler._batch_render_chat = lambda *_args, **_kwargs: (
pytest.fail("chat rendering should be skipped without a template")
)
handler._preprocess_completion_online = preprocess_completion
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["fallback"]
assert calls == [("completion", ["query: hello"])]
def test_text_only_with_template_uses_chat_path(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"], input_type="query")
calls: list[tuple[str, object]] = []
def batch_render_chat(
request,
all_messages,
truncate_prompt_tokens,
truncation_side,
):
calls.append(
(
"chat",
{
"request": request,
"all_messages": all_messages,
"truncate_prompt_tokens": truncate_prompt_tokens,
"truncation_side": truncation_side,
},
)
)
return ["chat"]
handler._get_task_instruction_prefix = lambda _input_type: "query: "
handler._has_chat_template = lambda: True
handler._batch_render_chat = batch_render_chat
handler._preprocess_completion_online = lambda *_args, **_kwargs: (
pytest.fail("completion path should be skipped when a template exists")
)
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["chat"]
assert calls == [
(
"chat",
{
"request": ctx.request,
"all_messages": [
handler._mixed_input_to_messages(
CohereEmbedInput(
content=[CohereEmbedContent(type="text", text="hello")]
),
task_prefix="query: ",
)
],
"truncate_prompt_tokens": -1,
"truncation_side": None,
},
)
]
@@ -26,25 +26,6 @@ TEXTS_2 = [
]
def _assert_score_output_matches_hf(hf_model, server: RemoteOpenAIServer,
payload: dict, text_pairs: list[list[str]]):
score_response = requests.post(
server.url_for("score"),
json={"model": MODEL_NAME, **payload},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
assert len(vllm_outputs) == len(hf_outputs)
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
assert hf_output == pytest.approx(vllm_output, rel=0.01)
@pytest.fixture(scope="module")
def server():
args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
@@ -63,78 +44,207 @@ def hf_model():
@pytest.mark.asyncio
@pytest.mark.parametrize(
("payload", "text_pairs"),
[
pytest.param(
{
"queries": TEXTS_1[0],
"documents": TEXTS_2[0],
},
[[TEXTS_1[0], TEXTS_2[0]]],
id="queries-str-documents-str",
),
pytest.param(
{
"queries": TEXTS_1[0],
"documents": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[0], TEXTS_2[1]],
],
id="queries-str-documents-list",
),
pytest.param(
{
"queries": TEXTS_1,
"documents": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
],
id="queries-list-documents-list",
),
pytest.param(
{
"queries": TEXTS_1,
"items": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
],
id="queries-list-items-list",
),
pytest.param(
{
"text_1": TEXTS_1,
"text_2": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
],
id="text-1-vs-text-2",
),
pytest.param(
{
"data_1": TEXTS_1,
"data_2": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
],
id="data-1-vs-data-2",
),
],
)
async def test_score_api_request_formats(hf_model, server: RemoteOpenAIServer,
payload: dict,
text_pairs: list[list[str]]):
_assert_score_output_matches_hf(hf_model, server, payload, text_pairs)
async def test_score_api_queries_str_1_documents_str_1(
hf_model, server: RemoteOpenAIServer
):
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1[0],
"documents": TEXTS_2[0],
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict([[TEXTS_1[0], TEXTS_2[0]]]).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_queries_str_1_documents_str_n(
hf_model, server: RemoteOpenAIServer
):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[0], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1[0],
"documents": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_queries_str_n_documents_str_n(
hf_model, server: RemoteOpenAIServer
):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1,
"documents": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_queries_vs_documents(hf_model, server: RemoteOpenAIServer):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1,
"documents": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_queries_vs_items(hf_model, server: RemoteOpenAIServer):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1,
"items": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_text_1_vs_text_2(hf_model, server: RemoteOpenAIServer):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"text_1": TEXTS_1,
"text_2": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_data_1_vs_data_2(hf_model, server: RemoteOpenAIServer):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"data_1": TEXTS_1,
"data_2": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
@@ -28,25 +28,6 @@ TEXTS_2 = [
]
def _assert_score_output_matches_hf(hf_model, server: RemoteOpenAIServer,
payload: dict, text_pairs: list[list[str]]):
score_response = requests.post(
server.url_for("score"),
json={"model": MODEL_NAME, **payload},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
assert len(vllm_outputs) == len(hf_outputs)
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
assert hf_output == pytest.approx(vllm_output, rel=0.01)
@pytest.fixture(scope="module")
def server():
args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
@@ -80,78 +61,207 @@ async def test_basic(server: RemoteOpenAIServer):
@pytest.mark.asyncio
@pytest.mark.parametrize(
("payload", "text_pairs"),
[
pytest.param(
{
"queries": TEXTS_1[0],
"documents": TEXTS_2[0],
},
[[TEXTS_1[0], TEXTS_2[0]]],
id="queries-str-documents-str",
),
pytest.param(
{
"queries": TEXTS_1[0],
"documents": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[0], TEXTS_2[1]],
],
id="queries-str-documents-list",
),
pytest.param(
{
"queries": TEXTS_1,
"documents": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
],
id="queries-list-documents-list",
),
pytest.param(
{
"queries": TEXTS_1,
"items": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
],
id="queries-list-items-list",
),
pytest.param(
{
"text_1": TEXTS_1,
"text_2": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
],
id="text-1-vs-text-2",
),
pytest.param(
{
"data_1": TEXTS_1,
"data_2": TEXTS_2,
},
[
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
],
id="data-1-vs-data-2",
),
],
)
async def test_score_api_request_formats(hf_model, server: RemoteOpenAIServer,
payload: dict,
text_pairs: list[list[str]]):
_assert_score_output_matches_hf(hf_model, server, payload, text_pairs)
async def test_score_api_queries_str_1_documents_str_1(
hf_model, server: RemoteOpenAIServer
):
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1[0],
"documents": TEXTS_2[0],
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict([[TEXTS_1[0], TEXTS_2[0]]]).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_queries_str_1_documents_str_n(
hf_model, server: RemoteOpenAIServer
):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[0], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1[0],
"documents": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_queries_str_n_documents_str_n(
hf_model, server: RemoteOpenAIServer
):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1,
"documents": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_queries_vs_documents(hf_model, server: RemoteOpenAIServer):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1,
"documents": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_queries_vs_items(hf_model, server: RemoteOpenAIServer):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": TEXTS_1,
"items": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_text_1_vs_text_2(hf_model, server: RemoteOpenAIServer):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"text_1": TEXTS_1,
"text_2": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
async def test_score_api_data_1_vs_data_2(hf_model, server: RemoteOpenAIServer):
text_pairs = [
[TEXTS_1[0], TEXTS_2[0]],
[TEXTS_1[1], TEXTS_2[1]],
]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"data_1": TEXTS_1,
"data_2": TEXTS_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
hf_outputs = hf_model.predict(text_pairs).tolist()
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
@pytest.mark.asyncio
@@ -7,3 +7,4 @@ server_args: >-
--max-model-len 4096
--data-parallel-size 2
--enable-expert-parallel
--max-num-seqs 512
@@ -38,7 +38,10 @@ def test_move_metatensors():
def test_reload_lifecycle():
layer = torch.nn.Linear(2, 3)
info = LayerReloadingInfo(restore_metadata=capture_layer_to_meta(layer))
info = LayerReloadingInfo(
restore_metadata=capture_layer_to_meta(layer),
restore_device=torch.device("cpu"),
)
restore_layer_on_meta(layer, info)
for name, tensor in get_layer_tensors(layer).items():
@@ -48,7 +51,7 @@ def test_reload_lifecycle():
assert tensor.__class__ == meta_tensor.__class__
assert tensor.__dict__ == meta_tensor.__dict__
materialize_layer(layer)
materialize_layer(layer, info)
for name, tensor in get_layer_tensors(layer).items():
materialized_tensor = getattr(layer, name)
assert tensor.dtype == materialized_tensor.dtype
@@ -60,7 +63,10 @@ def test_reload_lifecycle():
def test_model_cleanup(dist_init, default_vllm_config):
layer = QKVParallelLinear(2, 3, 4)
assert layer.weight.weight_loader.__self__ is layer
info = LayerReloadingInfo(restore_metadata=capture_layer_to_meta(layer))
info = LayerReloadingInfo(
restore_metadata=capture_layer_to_meta(layer),
restore_device=torch.device("cpu"),
)
mock_info_dict: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = (
WeakKeyDictionary()
@@ -90,39 +96,46 @@ def test_get_numel_loaded():
assert ret == "value"
@pytest.mark.parametrize("tp_size", [2])
@pytest.mark.parametrize(
"tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
)
@pytest.mark.parametrize(
"base_model,mul_model,add_model",
[
(
pytest.param(
"Qwen/Qwen3-0.6B",
"inference-optimization/Qwen3-0.6B-debug-multiply",
"inference-optimization/Qwen3-0.6B-debug-add",
marks=[pytest.mark.slow_test],
),
(
pytest.param(
"inference-optimization/Qwen3-0.6B-FP8_BLOCK",
"inference-optimization/Qwen3-0.6B-debug-multiply-FP8_BLOCK",
"inference-optimization/Qwen3-0.6B-debug-add-FP8_BLOCK",
marks=[pytest.mark.slow_test],
),
(
pytest.param(
"inference-optimization/Qwen3-0.6B-W4A16-G128",
"inference-optimization/Qwen3-0.6B-debug-multiply-W4A16-G128",
"inference-optimization/Qwen3-0.6B-debug-add-W4A16-G128",
marks=[pytest.mark.slow_test],
),
(
pytest.param(
"inference-optimization/DeepSeek-V3-debug-empty",
"inference-optimization/DeepSeek-V3-debug-multiply",
"inference-optimization/DeepSeek-V3-debug-add",
marks=[pytest.mark.slow_test],
),
(
pytest.param(
"inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC",
"inference-optimization/DeepSeek-V3-debug-multiply-FP8_DYNAMIC",
"inference-optimization/DeepSeek-V3-debug-add-FP8_DYNAMIC",
),
(
pytest.param(
"inference-optimization/DeepSeek-V3-debug-empty-NVFP4A16",
"inference-optimization/DeepSeek-V3-debug-multiply-NVFP4A16",
"inference-optimization/DeepSeek-V3-debug-add-NVFP4A16",
marks=[pytest.mark.slow_test],
),
],
)
@@ -138,6 +151,75 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
tensor_parallel_size=tp_size,
enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),
enable_prefix_caching=False,
max_model_len=16,
max_num_seqs=1,
) as llm:
llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})
mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
assert mul_perp < add_perp
llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model})
mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
assert add_perp < mul_perp
@pytest.mark.parametrize(
"tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
)
@pytest.mark.parametrize(
"base_model,mul_model,add_model,quantization",
[
pytest.param(
"Qwen/Qwen3-0.6B",
"inference-optimization/Qwen3-0.6B-debug-multiply",
"inference-optimization/Qwen3-0.6B-debug-add",
"fp8",
),
pytest.param(
"inference-optimization/DeepSeek-V3-debug-empty",
"inference-optimization/DeepSeek-V3-debug-multiply",
"inference-optimization/DeepSeek-V3-debug-add",
"fp8",
marks=[pytest.mark.slow_test],
),
pytest.param(
"Qwen/Qwen3-0.6B",
"inference-optimization/Qwen3-0.6B-debug-multiply",
"inference-optimization/Qwen3-0.6B-debug-add",
"mxfp8",
marks=[pytest.mark.slow_test],
),
pytest.param(
"inference-optimization/DeepSeek-V3-debug-empty",
"inference-optimization/DeepSeek-V3-debug-multiply",
"inference-optimization/DeepSeek-V3-debug-add",
"mxfp8",
marks=[
pytest.mark.slow_test,
pytest.mark.xfail(reason="mxfp4 & mla is not supported yet"),
],
),
],
)
def test_online_quantize_reload(
base_model, mul_model, add_model, quantization, tp_size, vllm_runner
):
if cuda_device_count_stateless() < tp_size:
pytest.skip(reason="Not enough CUDA devices")
if quantization == "fp8" and not current_platform.supports_fp8():
pytest.skip(reason="Requires FP8 support")
with vllm_runner(
model_name=base_model,
quantization=quantization,
tensor_parallel_size=tp_size,
enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),
enable_prefix_caching=False,
max_model_len=16,
max_num_seqs=1,
) as llm:
llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})
mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
@@ -1,8 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import types
import pytest
import torch
import torch.nn as nn
@@ -11,6 +9,8 @@ from vllm.model_executor.models.bert import (
BertMLMHead,
SPLADESparsePooler,
)
from vllm.pooling_params import PoolingParams
from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
# ---------------------------------------------------------------------
# Functional test: SPLADE formula correctness (no HF download needed)
@@ -38,8 +38,12 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):
],
dtype=torch.long,
)
meta = types.SimpleNamespace(
prompt_lens=prompt_lens_tenser, prompt_token_ids=token_ids
meta = PoolingMetadata(
prompt_lens=prompt_lens_tenser,
prompt_token_ids=token_ids,
prompt_token_ids_cpu=token_ids,
pooling_params=[PoolingParams(task="embed")] * B,
pooling_states=[PoolingStates() for _ in range(B)],
)
# MLM head (prefer BertMLMHead, fallback to Linear if unavailable)
@@ -394,6 +394,22 @@ VLM_TEST_SETTINGS = {
vllm_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}},
patch_hf_runner=model_utils.gemma3_patch_hf_runner,
),
"gemma4": VLMTestInfo(
models=["google/gemma-4-E2B-it"],
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
prompt_formatter=lambda img_prompt: f"<bos><start_of_turn>user\n{img_prompt}<end_of_turn>\n<start_of_turn>model\n", # noqa: E501
single_image_prompts=IMAGE_ASSETS.prompts(
{
"stop_sign": "What's the content in the center of the image?",
"cherry_blossom": "What is the season?",
}
),
multi_image_prompt="Describe the two images in detail.",
max_model_len=4096,
max_num_seqs=2,
auto_cls=AutoModelForImageTextToText,
vllm_runner_kwargs={"limit_mm_per_prompt": {"image": 4}},
),
"granite_vision": VLMTestInfo(
models=["ibm-granite/granite-vision-3.3-2b"],
test_type=(VLMTestType.IMAGE),
@@ -15,6 +15,10 @@ from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
MODEL_NAME = "ModernVBERT/colmodernvbert-merged"
COLBERT_DIM = 128
DTYPE = "half"
# Fixme:
# Update colmodernvbert code to support the latest HF version
# and remove revision set.
REVISION = "4a0a9f3ac7a7992fec410bfa8e3d080ac9a5bcee"
# -----------------------------------------------------------------------
@@ -26,6 +30,7 @@ def test_colmodernvbert_text_token_embed(vllm_runner):
"""Text query produces per-token embeddings with shape (seq_len, 128)."""
with vllm_runner(
MODEL_NAME,
revision=REVISION,
runner="pooling",
dtype=DTYPE,
enforce_eager=True,
@@ -49,6 +54,7 @@ def test_colmodernvbert_text_relevance_ordering(vllm_runner):
with vllm_runner(
MODEL_NAME,
revision=REVISION,
runner="pooling",
dtype=DTYPE,
enforce_eager=True,
@@ -66,6 +72,7 @@ def test_colmodernvbert_text_late_interaction(vllm_runner):
with vllm_runner(
MODEL_NAME,
revision=REVISION,
runner="pooling",
dtype=DTYPE,
enforce_eager=True,
@@ -92,6 +99,7 @@ def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
"""Image input produces per-token embeddings including vision tokens."""
with vllm_runner(
MODEL_NAME,
revision=REVISION,
runner="pooling",
dtype=DTYPE,
enforce_eager=True,
@@ -0,0 +1,67 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import numpy as np
import pytest
import torch
MODEL = "Qwen/Qwen3-ForcedAligner-0.6B"
CLASSIFY_NUM = 5000
TIMESTAMP_TOKEN_ID = 151705
def build_prompt(words: list[str]) -> str:
body = "<timestamp><timestamp>".join(words) + "<timestamp><timestamp>"
return f"<|audio_start|><|audio_pad|><|audio_end|>{body}"
@pytest.mark.parametrize("model", [MODEL])
@pytest.mark.parametrize("dtype", ["bfloat16"])
@torch.inference_mode()
def test_qwen3_forced_aligner(
vllm_runner,
model: str,
dtype: str,
) -> None:
words = ["Hello", "world"]
prompt = build_prompt(words)
# 5-second silent audio at 16kHz
audio = np.zeros(16000 * 5, dtype=np.float32)
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
enforce_eager=True,
max_model_len=512,
hf_overrides={
"architectures": [
"Qwen3ASRForcedAlignerForTokenClassification",
],
},
) as vllm_model:
outputs = vllm_model.llm.encode(
[{"prompt": prompt, "multi_modal_data": {"audio": audio}}],
pooling_task="token_classify",
)
# Validate output structure
assert len(outputs) == 1
logits = outputs[0].outputs.data
assert logits.dim() == 2
assert logits.shape[1] == CLASSIFY_NUM
# Validate timestamp extraction
token_ids = outputs[0].prompt_token_ids
predictions = logits.argmax(dim=-1)
ts_indices = [i for i, t in enumerate(token_ids) if t == TIMESTAMP_TOKEN_ID]
# 2 words x 2 timestamps each (start + end) = 4
assert len(ts_indices) == 4
ts_preds = [predictions[i].item() for i in ts_indices]
assert all(p >= 0 for p in ts_preds)
# end >= start for each word
assert ts_preds[1] >= ts_preds[0] # Hello
assert ts_preds[3] >= ts_preds[2] # world
@@ -0,0 +1,44 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.multimodal import MULTIMODAL_REGISTRY
from ....conftest import ImageTestAssets
from ...utils import build_model_context
# TODO: to be updated to "google/gemma-4-e2b-it" once the models are available
GEMMA4_MODEL_ID = "google/gemma-4-E2B-it"
@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
def test_limit_mm_per_prompt(
image_assets: ImageTestAssets,
model_id: str,
):
"""Test that limit_mm_per_prompt accurately restricts multiple images."""
# We only allow 1 image
ctx = build_model_context(
model_id,
mm_processor_kwargs={},
limit_mm_per_prompt={"image": 1},
)
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
# Provide 2 images in the prompt
prompt = "<image><image>"
# image_assets usually has multiple images
images = [asset.pil_image for asset in image_assets][:2]
if len(images) < 2:
images = [images[0], images[0]]
mm_data = {"image": images}
# Expect ValueError when exceeding limit
with pytest.raises(ValueError, match="At most 1 image"):
processor(
prompt,
mm_items=processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs={},
)
+16
View File
@@ -277,6 +277,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"google/gemma-2-9b", extras={"tiny": "google/gemma-2-2b-it"}
),
"Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),
"Gemma4ForCausalLM": _HfExamplesInfo(
"google/gemma-4-E2B-it",
min_transformers_version="5.0.0",
),
"Gemma3nForCausalLM": _HfExamplesInfo("google/gemma-3n-E2B-it"),
"GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"),
"Glm4ForCausalLM": _HfExamplesInfo("zai-org/GLM-4-9B-0414"),
@@ -636,6 +640,7 @@ _LATE_INTERACTION_EXAMPLE_MODELS = {
# [Multimodal]
"ColModernVBertForRetrieval": _HfExamplesInfo(
"ModernVBERT/colmodernvbert-merged",
revision="4a0a9f3ac7a7992fec410bfa8e3d080ac9a5bcee",
),
"ColPaliForRetrieval": _HfExamplesInfo("vidore/colpali-v1.3-hf"),
"ColQwen3": _HfExamplesInfo(
@@ -791,6 +796,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"Ernie4_5_VLMoeForConditionalGeneration": _HfExamplesInfo(
"baidu/ERNIE-4.5-VL-28B-A3B-PT",
trust_remote_code=True,
revision="refs/pr/17",
),
"FireRedASR2ForConditionalGeneration": _HfExamplesInfo(
"allendou/FireRedASR2-LLM-vllm",
@@ -803,6 +809,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
),
"FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
"Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),
"Gemma4ForConditionalGeneration": _HfExamplesInfo(
"google/gemma-4-E2B-it",
min_transformers_version="5.5.0",
),
"Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it"),
"GlmAsrForConditionalGeneration": _HfExamplesInfo(
"zai-org/GLM-ASR-Nano-2512",
@@ -1093,6 +1103,12 @@ _MULTIMODAL_EXAMPLE_MODELS = {
min_transformers_version="4.57",
hf_overrides={"architectures": ["Qwen3ASRRealtimeGeneration"]},
),
"Qwen3ASRForcedAlignerForTokenClassification": _HfExamplesInfo(
"Qwen/Qwen3-ForcedAligner-0.6B",
max_model_len=4096,
min_transformers_version="4.57",
hf_overrides={"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]},
),
"RForConditionalGeneration": _HfExamplesInfo("YannQi/R-4B", trust_remote_code=True),
"SkyworkR1VChatModel": _HfExamplesInfo(
"Skywork/Skywork-R1V-38B", trust_remote_code=True
+61 -9
View File
@@ -239,6 +239,17 @@ def test_video_media_io_backend_env_var_fallback(monkeypatch: pytest.MonkeyPatch
assert metadata_missing["video_backend"] == "test_video_backend_override_2"
def _make_jpeg_b64_frames(n: int, width: int = 8, height: int = 8) -> list[str]:
"""Return *n* tiny base64-encoded JPEG frames."""
frames: list[str] = []
for i in range(n):
img = Image.new("RGB", (width, height), color=(i % 256, 0, 0))
buf = io.BytesIO()
img.save(buf, format="JPEG")
frames.append(pybase64.b64encode(buf.getvalue()).decode("ascii"))
return frames
def test_load_base64_jpeg_returns_metadata():
"""Regression test: load_base64 with video/jpeg must return metadata.
@@ -248,16 +259,8 @@ def test_load_base64_jpeg_returns_metadata():
"""
num_test_frames = 3
frame_width, frame_height = 8, 8
# Build a few tiny JPEG frames and base64-encode them
b64_frames = []
for i in range(num_test_frames):
img = Image.new("RGB", (frame_width, frame_height), color=(i * 80, 0, 0))
buf = io.BytesIO()
img.save(buf, format="JPEG")
b64_frames.append(pybase64.b64encode(buf.getvalue()).decode("ascii"))
b64_frames = _make_jpeg_b64_frames(num_test_frames)
data = ",".join(b64_frames)
imageio = ImageMediaIO()
@@ -287,3 +290,52 @@ def test_load_base64_jpeg_returns_metadata():
# Default fps=1 → duration == num_frames
assert metadata["fps"] == 1.0
assert metadata["duration"] == float(num_test_frames)
def test_load_base64_jpeg_enforces_num_frames_limit():
"""Frames beyond num_frames must be truncated in the video/jpeg path.
Without the limit an attacker can send thousands of base64 JPEG frames
in a single request and exhaust server memory (OOM).
"""
num_frames_limit = 4
sent_frames = 20
b64_frames = _make_jpeg_b64_frames(sent_frames)
data = ",".join(b64_frames)
imageio = ImageMediaIO()
videoio = VideoMediaIO(imageio, num_frames=num_frames_limit)
frames, metadata = videoio.load_base64("video/jpeg", data)
assert frames.shape[0] == num_frames_limit
assert metadata["total_num_frames"] == num_frames_limit
assert metadata["frames_indices"] == list(range(num_frames_limit))
def test_load_base64_jpeg_no_limit_when_num_frames_negative():
"""When num_frames is -1, all frames should be loaded without truncation."""
sent_frames = 10
b64_frames = _make_jpeg_b64_frames(sent_frames)
data = ",".join(b64_frames)
imageio = ImageMediaIO()
videoio = VideoMediaIO(imageio, num_frames=-1)
frames, metadata = videoio.load_base64("video/jpeg", data)
assert frames.shape[0] == sent_frames
assert metadata["total_num_frames"] == sent_frames
assert metadata["frames_indices"] == list(range(sent_frames))
def test_load_base64_jpeg_raises_on_zero_num_frames():
"""num_frames=0 is invalid and should raise ValueError."""
b64_frames = _make_jpeg_b64_frames(3)
data = ",".join(b64_frames)
imageio = ImageMediaIO()
videoio = VideoMediaIO(imageio, num_frames=0)
with pytest.raises(ValueError, match="num_frames must be greater than 0 or -1"):
videoio.load_base64("video/jpeg", data)
+1
View File
@@ -11,6 +11,7 @@ MODELS = [
"TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ",
"TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx
"Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx
"RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp
]
DTYPE = ["bfloat16"]
+23
View File
@@ -466,3 +466,26 @@ def test_fp8_reloading(
weight_loader(param, torch.zeros(shape)) # cannot use empty
method.process_weights_after_loading(layer)
@pytest.mark.skipif(
not is_quant_method_supported("fp8"),
reason="FP8 is not supported on this GPU type.",
)
def test_kv_cache_dtype_skip_layers(vllm_runner, monkeypatch):
"""Test that kv_cache_dtype_skip_layers skips quantization for specified layers."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(
"facebook/opt-125m",
kv_cache_dtype="fp8",
kv_cache_dtype_skip_layers=["0", "2"],
enforce_eager=True,
) as llm:
def check_layers(model):
for i, layer in enumerate(model.model.decoder.layers):
expected = "auto" if str(i) in ["0", "2"] else "fp8"
assert layer.self_attn.attn.kv_cache_dtype == expected
llm.apply_model(check_layers)
@@ -0,0 +1,196 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from tests.reasoning.utils import run_reasoning_extraction
from vllm.reasoning import ReasoningParser, ReasoningParserManager
# Using mistral tokenizer as a generic mock since the actual model is not on HF
from vllm.tokenizers.registry import get_tokenizer
parser_name = "gemma4"
@pytest.fixture(scope="module")
def generic_tokenizer():
return get_tokenizer("google/gemma-4-E2B-it")
INVALID_SIMPLE_NONSTREAMING = {
"output": "This is a reasoning section<channel|>This is the rest",
"reasoning": "This is a reasoning section",
"content": "This is the rest",
"is_reasoning_end": True,
}
INVALID_SIMPLE_STREAMING = {
"output": "This is a reasoning section<channel|>This is the rest",
"reasoning": None,
"content": "This is a reasoning sectionThis is the rest",
"is_reasoning_end": True,
}
INVALID_COMPLETE_NONSTREAMING = {
"output": "This is a reasoning section<channel|>",
"reasoning": "This is a reasoning section",
"content": None,
"is_reasoning_end": True,
}
INVALID_COMPLETE_STREAMING = {
"output": "This is a reasoning section<channel|>",
"reasoning": None,
"content": "This is a reasoning section",
"is_reasoning_end": True,
}
NO_CONTENT = {
"output": "<|channel>This is reasoning",
"reasoning": "This is reasoning",
"content": None,
"is_reasoning_end": False,
}
NO_REASONING = {
"output": "This is content",
"reasoning": None,
"content": "This is content",
"is_reasoning_end": False,
}
REASONING_WITH_CHANNEL = {
"output": "<|channel>This is a reasoning section<channel|>This is the rest",
"reasoning": "This is a reasoning section",
"content": "This is the rest",
"is_reasoning_end": True,
}
COMPLETE_REASONING_WITH_CHANNEL = {
"output": "<|channel>This is a reasoning section<channel|>",
"reasoning": "This is a reasoning section",
"content": None,
"is_reasoning_end": True,
}
MULTIPLE_LINES_WITH_CHANNEL = {
"output": "<|channel>This\nThat<channel|>This is the rest\nThat",
"reasoning": "This\nThat",
"content": "This is the rest\nThat",
"is_reasoning_end": True,
}
CHANNEL_NO_END = {
"output": "<|channel>This is a reasoning section",
"reasoning": "This is a reasoning section",
"content": None,
"is_reasoning_end": False,
}
EMPTY = {
"output": "",
"reasoning": None,
"content": "",
"is_reasoning_end": False,
}
NEW_LINE_NONSTREAMING = {
"output": (
"Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest"
),
"reasoning": "This is a reasoning section",
"content": "\nThis is the rest",
"is_reasoning_end": True,
}
NEW_LINE_STREAMING = {
"output": (
"Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest"
),
"reasoning": "This is a reasoning section",
"content": "Before\n\nThis is the rest",
"is_reasoning_end": True,
}
TEST_CASES = [
pytest.param(False, INVALID_SIMPLE_NONSTREAMING, id="invalid_simple"),
pytest.param(True, INVALID_SIMPLE_STREAMING, id="invalid_simple_streaming"),
pytest.param(False, INVALID_COMPLETE_NONSTREAMING, id="invalid_complete"),
pytest.param(True, INVALID_COMPLETE_STREAMING, id="invalid_complete_streaming"),
pytest.param(False, NO_CONTENT, id="no_content"),
pytest.param(False, NO_REASONING, id="no_reasoning"),
pytest.param(False, REASONING_WITH_CHANNEL, id="reasoning"),
pytest.param(True, REASONING_WITH_CHANNEL, id="reasoning_streaming"),
pytest.param(False, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning"),
pytest.param(
True, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning_streaming"
),
pytest.param(False, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines"),
pytest.param(True, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines_streaming"),
pytest.param(False, CHANNEL_NO_END, id="no_end"),
pytest.param(True, CHANNEL_NO_END, id="no_end_streaming"),
pytest.param(False, EMPTY, id="empty"),
pytest.param(False, NEW_LINE_NONSTREAMING, id="new_line"),
pytest.param(True, NEW_LINE_STREAMING, id="new_line_streaming"),
]
@pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
def test_gemma4_reasoning(
streaming: bool,
param_dict: dict,
generic_tokenizer,
):
output = param_dict["output"]
# Resolve token IDs dynamically from the real tokenizer
vocab = generic_tokenizer.get_vocab()
start_token_id = vocab["<|channel>"]
end_token_id = vocab["<channel|>"]
index_start = output.find("<|channel>")
len_start = len("<|channel>")
index_end = output.find("<channel|>")
len_end = len("<channel|>")
output_tokens = []
def _encode(text: str) -> list[int]:
if not text:
return []
# Handle both raw transformers and vLLM wrappers
enc = getattr(generic_tokenizer, "tokenizer", generic_tokenizer)
try:
return enc.encode(text, add_special_tokens=False)
except TypeError:
return enc.encode(text)
if index_start != -1:
output_before = output[:index_start]
output_tokens += _encode(output_before)
output_tokens += [start_token_id]
if index_end != -1:
output_middle = output[index_start + len_start : index_end]
output_after = output[index_end + len_end :]
output_tokens += _encode(output_middle)
output_tokens += [end_token_id]
output_tokens += _encode(output_after)
else:
output_middle = output[index_start + len_start :]
output_tokens += _encode(output_middle)
elif index_end != -1:
output_before = output[:index_end]
output_after = output[index_end + len_end :]
output_tokens += _encode(output_before)
output_tokens += [end_token_id]
output_tokens += _encode(output_after)
else:
output_tokens += _encode(output)
parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
generic_tokenizer
)
# We use the generic run_reasoning_extraction from utils
# Use decode per token to get standard spaces instead of
# SentencePiece space characters
output_token_strings = [generic_tokenizer.decode([t]) for t in output_tokens]
reasoning, content = run_reasoning_extraction(
parser, output_token_strings, streaming=streaming
)
assert reasoning == param_dict["reasoning"]
assert content == param_dict["content"]
# Test is_reasoning_end
is_reasoning_end = parser.is_reasoning_end(output_tokens)
assert is_reasoning_end == param_dict["is_reasoning_end"]
+1
View File
@@ -38,6 +38,7 @@ class MockModelConfig:
skip_tokenizer_init: bool = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
renderer_num_workers: int = 1
@dataclass
+1
View File
@@ -37,6 +37,7 @@ class MockModelConfig:
skip_tokenizer_init: bool = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
renderer_num_workers: int = 1
@dataclass
+22
View File
@@ -1131,6 +1131,28 @@ def test_needs_dp_coordination(
assert vllm_config.needs_dp_coordinator == expected_needs_coordinator
def test_renderer_num_workers_with_mm_cache():
"""Disallow renderer_num_workers > 1 when mm processor cache is enabled,
since neither cache type is thread-safe."""
mm_model = "Qwen/Qwen2-VL-2B-Instruct"
# Should raise: multi-worker + cache enabled (default cache_gb=4)
with pytest.raises(ValueError, match="renderer-num-workers"):
ModelConfig(mm_model, renderer_num_workers=4)
# Should raise: multi-worker + explicit cache size
with pytest.raises(ValueError, match="renderer-num-workers"):
ModelConfig(mm_model, renderer_num_workers=2, mm_processor_cache_gb=1.0)
# Should pass: multi-worker + cache disabled
config = ModelConfig(mm_model, renderer_num_workers=4, mm_processor_cache_gb=0)
assert config.renderer_num_workers == 4
# Should pass: single worker + cache enabled (default)
config = ModelConfig(mm_model, renderer_num_workers=1)
assert config.renderer_num_workers == 1
def test_eagle_draft_model_config():
"""Test that EagleDraft model config is correctly set."""
target_model_config = ModelConfig(
+52
View File
@@ -454,3 +454,55 @@ class TestVllmConfigureLogging:
with pytest.raises(ValueError, match="invalid literal for int"):
_ = envs.VLLM_CONFIGURE_LOGGING
class TestVllmMaxNSequences:
def test_default_value(self):
"""Test that VLLM_MAX_N_SEQUENCES defaults to 64."""
with patch.dict(os.environ, {}, clear=False):
os.environ.pop("VLLM_MAX_N_SEQUENCES", None)
if hasattr(envs.__getattr__, "cache_clear"):
envs.__getattr__.cache_clear()
assert envs.VLLM_MAX_N_SEQUENCES == 16384
def test_custom_value(self, monkeypatch: pytest.MonkeyPatch):
"""Test that VLLM_MAX_N_SEQUENCES can be overridden."""
monkeypatch.setenv("VLLM_MAX_N_SEQUENCES", "128")
if hasattr(envs.__getattr__, "cache_clear"):
envs.__getattr__.cache_clear()
assert envs.VLLM_MAX_N_SEQUENCES == 128
def test_sampling_params_respects_limit(
self,
monkeypatch: pytest.MonkeyPatch,
):
"""Test that SamplingParams rejects n above the limit."""
from vllm.sampling_params import SamplingParams
monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
if hasattr(envs.__getattr__, "cache_clear"):
envs.__getattr__.cache_clear()
max_n = envs.VLLM_MAX_N_SEQUENCES
SamplingParams(n=max_n)
with pytest.raises(ValueError, match="n must be at most"):
SamplingParams(n=max_n + 1)
def test_sampling_params_respects_custom_limit(
self,
monkeypatch: pytest.MonkeyPatch,
):
"""Test that SamplingParams uses the overridden env var limit."""
from vllm.sampling_params import SamplingParams
monkeypatch.setenv("VLLM_MAX_N_SEQUENCES", "128")
if hasattr(envs.__getattr__, "cache_clear"):
envs.__getattr__.cache_clear()
SamplingParams(n=128)
with pytest.raises(ValueError, match="n must be at most 128"):
SamplingParams(n=129)
@@ -11,6 +11,7 @@ from unittest.mock import MagicMock
import pytest
from vllm.tokenizers import get_tokenizer
from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser
# ---------------------------------------------------------------------------
@@ -474,3 +475,183 @@ class TestExtractToolCallsStreaming:
deltas = self._stream(parser, partial_text)
# Should have no tool call deltas yet
assert all(not d.tool_calls for d in deltas)
@pytest.fixture(scope="module")
def deepseekv32_tokenizer():
return get_tokenizer(tokenizer_name="deepseek-ai/DeepSeek-V3.2")
@pytest.fixture
def parser(deepseekv32_tokenizer):
return DeepSeekV32ToolParser(deepseekv32_tokenizer)
def test_convert_param_value_single_types(parser):
"""Test _convert_param_value with single type parameters."""
# Test string type
assert parser._convert_param_value("hello", "string") == "hello"
assert parser._convert_param_value("123", "string") == "123"
# Test integer type - valid integers
assert parser._convert_param_value("123", "integer") == 123
assert parser._convert_param_value("456", "int") == 456
# Invalid integer should return original string (due to exception catch)
assert parser._convert_param_value("abc", "integer") == "abc"
# Test float/number type
assert parser._convert_param_value("123.45", "float") == 123.45
assert (
parser._convert_param_value("123.0", "number") == 123
) # Should be int when whole number
assert parser._convert_param_value("123.5", "number") == 123.5
# Invalid float should return original string
assert parser._convert_param_value("abc", "float") == "abc"
# Test boolean type - valid boolean values
assert parser._convert_param_value("true", "boolean") is True
assert parser._convert_param_value("false", "bool") is False
assert parser._convert_param_value("1", "boolean") is True
assert parser._convert_param_value("0", "boolean") is False
# Invalid boolean should return original string
assert parser._convert_param_value("yes", "boolean") == "yes"
assert parser._convert_param_value("no", "bool") == "no"
# Test null value
assert parser._convert_param_value("null", "string") is None
assert parser._convert_param_value("null", "integer") is None
# Test object/array type (JSON)
assert parser._convert_param_value('{"key": "value"}', "object") == {"key": "value"}
assert parser._convert_param_value("[1, 2, 3]", "array") == [1, 2, 3]
# Invalid JSON should return original string
assert parser._convert_param_value("{invalid}", "object") == "{invalid}"
# Test fallback for unknown type (tries json.loads, then returns original)
assert parser._convert_param_value('{"key": "value"}', "unknown") == {
"key": "value"
}
assert parser._convert_param_value("plain text", "unknown") == "plain text"
def test_convert_param_value_multi_typed_values(parser):
"""Test _convert_param_value with multi-typed values (list of types)."""
# Test with list of types where first type succeeds
assert parser._convert_param_value("123", ["integer", "string"]) == 123
assert parser._convert_param_value("true", ["boolean", "string"]) is True
assert parser._convert_param_value('{"x": 1}', ["object", "string"]) == {"x": 1}
# Test with list of types where first type fails but second succeeds
# "abc" is not a valid integer, so should try string next
assert parser._convert_param_value("abc", ["integer", "string"]) == "abc"
# Test with list of types where all fail - should return original value
# "invalid json" is not valid JSON, last type is "object" which will fail JSON parse
result = parser._convert_param_value("invalid json", ["integer", "object"])
assert result == "invalid json" # Returns original value after all types fail
# Test with three types
assert parser._convert_param_value("123.5", ["integer", "float", "string"]) == 123.5
assert parser._convert_param_value("true", ["integer", "boolean", "string"]) is True
# Test with null in multi-type list
assert parser._convert_param_value("null", ["integer", "string"]) is None
assert parser._convert_param_value("null", ["boolean", "object"]) is None
# Test nested type conversion - boolean fails, integer succeeds
value = parser._convert_param_value("123", ["boolean", "integer", "string"])
assert value == 123 # Should be integer, not boolean
# Test that order matters
assert (
parser._convert_param_value("123", ["string", "integer"]) == "123"
) # String first
assert (
parser._convert_param_value("123", ["integer", "string"]) == 123
) # Integer first
# Test with all types failing - returns original value
assert (
parser._convert_param_value("not_a_number", ["integer", "float", "boolean"])
== "not_a_number"
)
def test_convert_param_value_stricter_type_checking(parser):
"""Test stricter type checking in the updated implementation."""
# Boolean now has stricter validation
assert parser._convert_param_value("true", "boolean") is True
assert parser._convert_param_value("false", "boolean") is False
assert parser._convert_param_value("1", "boolean") is True
assert parser._convert_param_value("0", "boolean") is False
# These should return original string (not valid boolean values)
assert parser._convert_param_value("yes", "boolean") == "yes"
assert parser._convert_param_value("no", "boolean") == "no"
assert parser._convert_param_value("TRUE", "boolean") is True
assert parser._convert_param_value("FALSE", "boolean") is False
# Integer and float now raise exceptions for invalid values
assert parser._convert_param_value("123abc", "integer") == "123abc"
assert parser._convert_param_value("123.45.67", "float") == "123.45.67"
# JSON parsing is stricter - invalid JSON returns original
assert parser._convert_param_value("{invalid: json}", "object") == "{invalid: json}"
assert parser._convert_param_value("[1, 2,", "array") == "[1, 2,"
# Test multi-type with stricter checking
# "yes" is not valid boolean, but string would accept it
assert parser._convert_param_value("yes", ["boolean", "string"]) == "yes"
# "123abc" is not valid integer or float, but string accepts it
assert (
parser._convert_param_value("123abc", ["integer", "float", "string"])
== "123abc"
)
def test_convert_param_value_edge_cases(parser):
"""Test edge cases for _convert_param_value."""
# Empty string
assert parser._convert_param_value("", "string") == ""
assert (
parser._convert_param_value("", "integer") == ""
) # Invalid int returns original
# Whitespace - trimmed by conversion functions
assert parser._convert_param_value(" 123 ", "integer") == 123
assert parser._convert_param_value(" true ", "boolean") is True
# Numeric strings with special characters
assert parser._convert_param_value("123.45.67", "float") == "123.45.67"
assert parser._convert_param_value("123abc", "integer") == "123abc"
# JSON with whitespace - should parse correctly
assert parser._convert_param_value(' { "key" : "value" } ', "object") == {
"key": "value"
}
# Invalid JSON returns original
assert parser._convert_param_value("{invalid}", "object") == "{invalid}"
assert parser._convert_param_value("[1, 2,", "array") == "[1, 2,"
def test_convert_param_value_checked_helper(parser):
"""Test the _convert_param_value_checked helper function indirectly."""
# This tests the behavior through the main function
# Valid conversions should work
assert parser._convert_param_value("123", "integer") == 123
assert parser._convert_param_value("123.45", "float") == 123.45
assert parser._convert_param_value("true", "boolean") is True
assert parser._convert_param_value('{"x": 1}', "object") == {"x": 1}
# Invalid conversions should return original value (exception caught)
assert parser._convert_param_value("abc", "integer") == "abc"
assert parser._convert_param_value("abc", "float") == "abc"
assert parser._convert_param_value("yes", "boolean") == "yes"
assert parser._convert_param_value("{invalid}", "object") == "{invalid}"
# Test that null handling works in checked function
assert parser._convert_param_value("null", "integer") is None
assert parser._convert_param_value("null", "boolean") is None
assert parser._convert_param_value("null", "object") is None
@@ -0,0 +1,504 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from typing import Any
from unittest.mock import MagicMock
import pytest
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.tool_parsers.gemma4_tool_parser import (
TOOL_CALL_END,
TOOL_CALL_START,
Gemma4ToolParser,
_parse_gemma4_args,
_parse_gemma4_array,
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def mock_tokenizer():
tokenizer = MagicMock()
tokenizer.encode.return_value = [1, 2, 3]
# Include the tool call start token in the vocab for the parser
tokenizer.get_vocab.return_value = {TOOL_CALL_START: 48, TOOL_CALL_END: 49}
return tokenizer
@pytest.fixture
def parser(mock_tokenizer):
return Gemma4ToolParser(mock_tokenizer)
@pytest.fixture
def mock_request():
request = MagicMock(spec=ChatCompletionRequest)
request.tools = []
request.tool_choice = "auto"
return request
# ---------------------------------------------------------------------------
# Unit tests for _parse_gemma4_args (shared parser logic)
# ---------------------------------------------------------------------------
class TestParseGemma4Args:
def test_empty_string(self):
assert _parse_gemma4_args("") == {}
def test_whitespace_only(self):
assert _parse_gemma4_args(" ") == {}
def test_single_string_value(self):
result = _parse_gemma4_args('location:<|"|>Paris<|"|>')
assert result == {"location": "Paris"}
def test_string_value_with_comma(self):
result = _parse_gemma4_args('location:<|"|>Paris, France<|"|>')
assert result == {"location": "Paris, France"}
def test_multiple_string_values(self):
result = _parse_gemma4_args(
'location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|>'
)
assert result == {"location": "San Francisco", "unit": "celsius"}
def test_integer_value(self):
result = _parse_gemma4_args("count:42")
assert result == {"count": 42}
def test_float_value(self):
result = _parse_gemma4_args("score:3.14")
assert result == {"score": 3.14}
def test_boolean_true(self):
result = _parse_gemma4_args("flag:true")
assert result == {"flag": True}
def test_boolean_false(self):
result = _parse_gemma4_args("flag:false")
assert result == {"flag": False}
def test_mixed_types(self):
result = _parse_gemma4_args(
'name:<|"|>test<|"|>,count:42,active:true,score:3.14'
)
assert result == {
"name": "test",
"count": 42,
"active": True,
"score": 3.14,
}
def test_nested_object(self):
result = _parse_gemma4_args('nested:{inner:<|"|>value<|"|>}')
assert result == {"nested": {"inner": "value"}}
def test_array_of_strings(self):
result = _parse_gemma4_args('items:[<|"|>a<|"|>,<|"|>b<|"|>]')
assert result == {"items": ["a", "b"]}
def test_unterminated_string(self):
"""Unterminated strings should take everything after the delimiter."""
result = _parse_gemma4_args('key:<|"|>unterminated')
assert result == {"key": "unterminated"}
def test_empty_value(self):
"""Key with no value after colon."""
result = _parse_gemma4_args("key:")
assert result == {"key": ""}
class TestParseGemma4Array:
def test_string_array(self):
result = _parse_gemma4_array('<|"|>a<|"|>,<|"|>b<|"|>')
assert result == ["a", "b"]
def test_empty_array(self):
result = _parse_gemma4_array("")
assert result == []
def test_bare_values(self):
result = _parse_gemma4_array("42,true,3.14")
assert result == [42, True, 3.14]
# ---------------------------------------------------------------------------
# Non-streaming extraction tests
# ---------------------------------------------------------------------------
class TestExtractToolCalls:
def test_no_tool_calls(self, parser, mock_request):
model_output = "Hello, how can I help you today?"
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is False
assert result.tool_calls == []
assert result.content == model_output
def test_single_tool_call(self, parser, mock_request):
model_output = (
'<|tool_call>call:get_weather{location:<|"|>London<|"|>}<tool_call|>'
)
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 1
assert result.tool_calls[0].function.name == "get_weather"
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"location": "London"}
def test_multiple_arguments(self, parser, mock_request):
model_output = (
"<|tool_call>call:get_weather{"
'location:<|"|>San Francisco<|"|>,'
'unit:<|"|>celsius<|"|>}'
"<tool_call|>"
)
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 1
assert result.tool_calls[0].function.name == "get_weather"
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"location": "San Francisco", "unit": "celsius"}
def test_text_before_tool_call(self, parser, mock_request):
model_output = (
"Let me check the weather for you. "
'<|tool_call>call:get_weather{location:<|"|>Paris<|"|>}'
"<tool_call|>"
)
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert result.content == "Let me check the weather for you."
assert len(result.tool_calls) == 1
assert result.tool_calls[0].function.name == "get_weather"
def test_multiple_tool_calls(self, parser, mock_request):
model_output = (
'<|tool_call>call:get_weather{location:<|"|>London<|"|>}'
"<tool_call|>"
'<|tool_call>call:get_time{location:<|"|>London<|"|>}'
"<tool_call|>"
)
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 2
assert result.tool_calls[0].function.name == "get_weather"
assert result.tool_calls[1].function.name == "get_time"
def test_nested_arguments(self, parser, mock_request):
model_output = (
"<|tool_call>call:complex_function{"
'nested:{inner:<|"|>value<|"|>},'
'list:[<|"|>a<|"|>,<|"|>b<|"|>]}'
"<tool_call|>"
)
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 1
assert result.tool_calls[0].function.name == "complex_function"
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"nested": {"inner": "value"}, "list": ["a", "b"]}
def test_tool_call_with_number_and_boolean(self, parser, mock_request):
model_output = (
"<|tool_call>call:set_status{"
"is_active:true,"
"count:42,"
"score:3.14}"
"<tool_call|>"
)
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 1
assert result.tool_calls[0].function.name == "set_status"
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"is_active": True, "count": 42, "score": 3.14}
def test_incomplete_tool_call(self, parser, mock_request):
model_output = '<|tool_call>call:get_weather{location:<|"|>London'
result = parser.extract_tool_calls(model_output, mock_request)
# Incomplete — no <tool_call|> end marker, regex won't match
assert result.tools_called is False
assert result.content == model_output
def test_hyphenated_function_name(self, parser, mock_request):
"""Ensure function names with hyphens are parsed correctly."""
model_output = (
'<|tool_call>call:get-weather{location:<|"|>London<|"|>}<tool_call|>'
)
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert result.tool_calls[0].function.name == "get-weather"
def test_dotted_function_name(self, parser, mock_request):
"""Ensure function names with dots are parsed correctly."""
model_output = (
'<|tool_call>call:weather.get{location:<|"|>London<|"|>}<tool_call|>'
)
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert result.tool_calls[0].function.name == "weather.get"
def test_no_arguments(self, parser, mock_request):
"""Tool calls with empty arguments."""
model_output = "<|tool_call>call:get_status{}<tool_call|>"
result = parser.extract_tool_calls(model_output, mock_request)
assert result.tools_called is True
assert result.tool_calls[0].function.name == "get_status"
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {}
# ---------------------------------------------------------------------------
# Streaming extraction tests
# ---------------------------------------------------------------------------
class TestStreamingExtraction:
"""Tests for the streaming tool call extraction.
These simulate the token-by-token streaming that vLLM performs,
feeding incremental text to extract_tool_calls_streaming() and
verifying that the accumulated argument deltas form valid JSON.
"""
def _simulate_streaming(
self, parser: Gemma4ToolParser, mock_request: Any, chunks: list[str]
) -> list[tuple[Any, str]]:
"""Feed chunks through the streaming parser and collect results.
Returns a list of (delta_message, accumulated_text) tuples.
"""
results: list[tuple[Any, str]] = []
previous_text: str = ""
previous_token_ids: list[int] = []
for chunk in chunks:
current_text = previous_text + chunk
# Use token ID 48 for tool_call start, 49 for end, 0 otherwise
delta_token_ids: list[int] = []
if TOOL_CALL_START in chunk:
delta_token_ids.append(48)
elif TOOL_CALL_END in chunk:
delta_token_ids.append(49)
else:
delta_token_ids.append(0)
current_token_ids = previous_token_ids + delta_token_ids
delta = parser.extract_tool_calls_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=chunk,
previous_token_ids=tuple(previous_token_ids),
current_token_ids=tuple(current_token_ids),
delta_token_ids=tuple(delta_token_ids),
request=mock_request,
)
results.append((delta, current_text))
previous_text = current_text
previous_token_ids = list(current_token_ids)
return results
def _collect_arguments(self, results):
"""Collect all argument deltas from streaming results into one string."""
args_text = ""
for delta, _ in results:
if delta and delta.tool_calls:
for tc in delta.tool_calls:
func = tc.function if isinstance(tc.function, dict) else tc.function
if isinstance(func, dict):
arg = func.get("arguments", "")
else:
arg = getattr(func, "arguments", "") or ""
if arg:
args_text += arg
return args_text
def _collect_function_name(self, results):
"""Extract the function name from streaming results."""
for delta, _ in results:
if delta and delta.tool_calls:
for tc in delta.tool_calls:
func = tc.function if isinstance(tc.function, dict) else tc.function
if isinstance(func, dict):
name = func.get("name")
else:
name = getattr(func, "name", None)
if name:
return name
return None
def test_basic_streaming_single_tool(self, parser, mock_request):
"""Simulate the exact streaming scenario from the bug report.
Model generates:
<|tool_call>call:get_weather{location:<|"|>Paris, France<|"|>}<tool_call|>
Expected: arguments should be valid JSON {"location": "Paris, France"}
"""
chunks = [
"<|tool_call>",
"call:get_weather{",
'location:<|"|>Paris',
", France",
'<|"|>}',
"<tool_call|>",
]
results = self._simulate_streaming(parser, mock_request, chunks)
# Verify function name
name = self._collect_function_name(results)
assert name == "get_weather", f"Expected 'get_weather', got '{name}'"
# Verify arguments form valid JSON
args_text = self._collect_arguments(results)
assert args_text, "No arguments were streamed"
parsed_args = json.loads(args_text)
assert parsed_args == {"location": "Paris, France"}
def test_streaming_multi_arg(self, parser, mock_request):
"""Streaming with multiple arguments."""
chunks = [
"<|tool_call>",
"call:get_weather{",
'location:<|"|>Tokyo<|"|>,',
'unit:<|"|>celsius<|"|>}',
"<tool_call|>",
]
results = self._simulate_streaming(parser, mock_request, chunks)
name = self._collect_function_name(results)
assert name == "get_weather"
args_text = self._collect_arguments(results)
assert args_text
parsed_args = json.loads(args_text)
assert parsed_args == {"location": "Tokyo", "unit": "celsius"}
def test_streaming_no_extra_brace(self, parser, mock_request):
"""Verify the closing } is NOT leaked into arguments (Bug #2)."""
chunks = [
"<|tool_call>",
"call:get_weather{",
'location:<|"|>London<|"|>}',
"<tool_call|>",
]
results = self._simulate_streaming(parser, mock_request, chunks)
args_text = self._collect_arguments(results)
assert args_text
# The args text must be valid JSON (no extra })
parsed = json.loads(args_text)
assert parsed == {"location": "London"}
# Specifically assert no double-brace
assert args_text.count("}") <= 1, (
f"Arguments contain extra closing brace: {args_text!r}"
)
def test_streaming_no_unquoted_keys(self, parser, mock_request):
"""Verify keys are properly quoted in JSON (Bug #1)."""
chunks = [
"<|tool_call>",
"call:get_weather{",
'location:<|"|>Paris<|"|>}',
"<tool_call|>",
]
results = self._simulate_streaming(parser, mock_request, chunks)
args_text = self._collect_arguments(results)
# Must start with { and contain quoted key
assert args_text.lstrip().startswith("{"), (
f"Arguments don't start with '{{': {args_text!r}"
)
assert '"location"' in args_text, (
f"Key 'location' not properly quoted: {args_text!r}"
)
def test_streaming_name_no_call_prefix(self, parser, mock_request):
"""Verify function name has no 'call:' prefix."""
chunks = [
"<|tool_call>",
"call:get_weather{",
'location:<|"|>Paris<|"|>}',
"<tool_call|>",
]
results = self._simulate_streaming(parser, mock_request, chunks)
name = self._collect_function_name(results)
assert name == "get_weather"
assert not name.startswith("call:"), f"Name has 'call:' prefix: {name!r}"
def test_streaming_text_before_tool_call(self, parser, mock_request):
"""Text before tool call should be emitted as content."""
chunks = [
"Let me check ",
"the weather. ",
"<|tool_call>",
"call:get_weather{",
'location:<|"|>London<|"|>}',
"<tool_call|>",
]
results = self._simulate_streaming(parser, mock_request, chunks)
# First chunks should be content
content_parts = []
for delta, _ in results:
if delta and delta.content:
content_parts.append(delta.content)
assert "".join(content_parts).strip().startswith("Let me check")
def test_streaming_numeric_args(self, parser, mock_request):
"""Streaming with numeric and boolean argument values."""
chunks = [
"<|tool_call>",
"call:set_config{",
"count:42,",
"active:true}",
"<tool_call|>",
]
results = self._simulate_streaming(parser, mock_request, chunks)
args_text = self._collect_arguments(results)
if args_text:
parsed_args = json.loads(args_text)
assert parsed_args["count"] == 42
assert parsed_args["active"] is True
def test_streaming_empty_args(self, parser, mock_request):
"""Tool call with no arguments."""
chunks = [
"<|tool_call>",
"call:get_status{}",
"<tool_call|>",
]
results = self._simulate_streaming(parser, mock_request, chunks)
name = self._collect_function_name(results)
assert name == "get_status"
@@ -152,6 +152,175 @@ def test_hermes_parser_streaming(
}
def _simulate_streaming(
tokenizer: TokenizerLike,
parser: ToolParser,
request: ChatCompletionRequest,
text: str,
stream_interval: int = 1,
) -> list:
"""Simulate streaming with a given stream_interval.
Tokens are batched into chunks of `stream_interval` tokens,
mimicking how the output processor delivers them.
Returns a list of non-None DeltaMessages.
"""
tokens = tokenizer.encode(text)
previous_text = ""
delta_messages = []
for i in range(0, len(tokens), stream_interval):
chunk_ids = tokens[i : i + stream_interval]
delta_text = tokenizer.decode(chunk_ids)
current_text = previous_text + delta_text
delta = parser.extract_tool_calls_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=delta_text,
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=chunk_ids,
request=request,
)
previous_text = current_text
if delta is not None:
delta_messages.append(delta)
return delta_messages
@pytest.mark.parametrize("stream_interval", [2, 3, 5, 8])
def test_hermes_streaming_tool_call_with_stream_interval(
qwen_tokenizer: TokenizerLike,
any_chat_request: ChatCompletionRequest,
stream_interval: int,
) -> None:
"""Tool call streaming must produce correct name + args at any interval."""
text = (
'<tool_call>{"name": "get_current_temperature", '
'"arguments": {"location": "San Francisco", "unit": "celsius"}}'
"</tool_call>"
)
parser = Hermes2ProToolParser(qwen_tokenizer)
deltas = _simulate_streaming(
qwen_tokenizer, parser, any_chat_request, text, stream_interval
)
# Flatten all DeltaToolCalls across all deltas.
tool_deltas = [tc for d in deltas if d.tool_calls for tc in d.tool_calls]
assert tool_deltas, "Expected at least one tool call delta"
assert tool_deltas[0].function.name == "get_current_temperature"
# Concatenated arguments must be valid JSON matching the original.
args_str = "".join(tc.function.arguments or "" for tc in tool_deltas)
assert json.loads(args_str) == {
"location": "San Francisco",
"unit": "celsius",
}
@pytest.mark.parametrize("stream_interval", [2, 3, 5, 8])
def test_hermes_streaming_content_then_tool_call_with_stream_interval(
qwen_tokenizer: TokenizerLike,
any_chat_request: ChatCompletionRequest,
stream_interval: int,
) -> None:
"""Content before a tool call must be fully streamed, then tool call."""
text = (
"Sure, let me check the weather."
'<tool_call>{"name": "get_weather", '
'"arguments": {"city": "NYC"}}</tool_call>'
)
parser = Hermes2ProToolParser(qwen_tokenizer)
deltas = _simulate_streaming(
qwen_tokenizer, parser, any_chat_request, text, stream_interval
)
content_deltas = [d for d in deltas if d.content]
tool_deltas = [d for d in deltas if d.tool_calls]
# Content must reconstruct the prefix.
content_str = "".join(d.content for d in content_deltas)
assert content_str == "Sure, let me check the weather."
# Tool call must be correct.
tool_calls = [tc for d in tool_deltas for tc in d.tool_calls]
assert tool_calls[0].function.name == "get_weather"
args_str = "".join(tc.function.arguments or "" for tc in tool_calls)
assert json.loads(args_str) == {"city": "NYC"}
@pytest.mark.parametrize("stream_interval", [1, 2, 4])
def test_hermes_streaming_multiple_tool_calls_with_stream_interval(
qwen_tokenizer: TokenizerLike,
any_chat_request: ChatCompletionRequest,
stream_interval: int,
) -> None:
"""Multiple sequential tool calls must each be streamed correctly."""
text = (
'<tool_call>{"name": "search", "arguments": {"q": "cats"}}</tool_call>'
'<tool_call>{"name": "search", "arguments": {"q": "dogs"}}</tool_call>'
)
parser = Hermes2ProToolParser(qwen_tokenizer)
deltas = _simulate_streaming(
qwen_tokenizer, parser, any_chat_request, text, stream_interval
)
# Flatten all DeltaToolCalls across all deltas.
all_tool_calls = [tc for d in deltas if d.tool_calls for tc in d.tool_calls]
# Separate by tool index.
tool0 = [tc for tc in all_tool_calls if tc.index == 0]
tool1 = [tc for tc in all_tool_calls if tc.index == 1]
assert tool0[0].function.name == "search"
args0 = "".join(tc.function.arguments or "" for tc in tool0)
assert json.loads(args0) == {"q": "cats"}
assert tool1[0].function.name == "search"
args1 = "".join(tc.function.arguments or "" for tc in tool1)
assert json.loads(args1) == {"q": "dogs"}
@pytest.mark.parametrize("stream_interval", [2, 5])
def test_hermes_streaming_boolean_args_with_stream_interval(
qwen_tokenizer: TokenizerLike,
any_chat_request: ChatCompletionRequest,
stream_interval: int,
) -> None:
"""Regression test for bug #19056 with stream_interval > 1."""
text = (
"<tool_call>\n"
'{"name": "final_answer", "arguments": {"trigger": true}}\n'
"</tool_call>"
)
parser = Hermes2ProToolParser(qwen_tokenizer)
deltas = _simulate_streaming(
qwen_tokenizer, parser, any_chat_request, text, stream_interval
)
tool_calls = [tc for d in deltas if d.tool_calls for tc in d.tool_calls]
assert tool_calls[0].function.name == "final_answer"
args_str = "".join(tc.function.arguments or "" for tc in tool_calls)
assert json.loads(args_str) == {"trigger": True}
@pytest.mark.parametrize("stream_interval", [2, 3, 5])
def test_hermes_streaming_just_forward_text_with_stream_interval(
qwen_tokenizer: TokenizerLike,
any_chat_request: ChatCompletionRequest,
stream_interval: int,
) -> None:
"""Plain text with no tool calls must be fully forwarded."""
text = "This is plain text with no tool calling involved."
parser = Hermes2ProToolParser(qwen_tokenizer)
deltas = _simulate_streaming(
qwen_tokenizer, parser, any_chat_request, text, stream_interval
)
for d in deltas:
assert not d.tool_calls
assert "".join(d.content for d in deltas) == text
def test_hermes_parser_non_streaming_no_tool_call(
hermes_parser: ToolParser,
any_chat_request: ChatCompletionRequest,
@@ -218,3 +387,28 @@ def test_hermes_parser_non_streaming_tool_call_invalid_json(
assert tool_call is not None
assert not tool_call.tools_called
def test_hermes_streaming_content_and_tool_call_in_single_chunk(
qwen_tokenizer: TokenizerLike,
any_chat_request: ChatCompletionRequest,
) -> None:
"""Content + complete tool call in one chunk must both be emitted."""
text = 'Hi!<tool_call>{"name": "f", "arguments": {"x": 1}}</tool_call>'
# Use a stream_interval large enough to guarantee a single chunk.
parser = Hermes2ProToolParser(qwen_tokenizer)
deltas = _simulate_streaming(
qwen_tokenizer,
parser,
any_chat_request,
text,
stream_interval=9999,
)
content_parts = [d.content for d in deltas if d.content]
tool_parts = [tc for d in deltas if d.tool_calls for tc in d.tool_calls]
assert "".join(content_parts) == "Hi!"
assert tool_parts[0].function.name == "f"
args_str = "".join(tc.function.arguments or "" for tc in tool_parts)
assert json.loads(args_str) == {"x": 1}
@@ -42,6 +42,7 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import (
FlashMLASparseBackend,
triton_convert_req_index_to_global_index,
)
from vllm.v1.attention.backends.mla.indexer import split_indexer_prefill_chunks
from vllm.v1.attention.backends.utils import split_prefill_chunks
from vllm.v1.attention.ops import flashmla
@@ -716,6 +717,81 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected):
assert out == expected
@pytest.mark.parametrize(
"seq_lens,query_lens,workspace_size,max_logits_bytes,expected",
[
# Logits constraint triggers split (M*N exceeds budget)
# req0: M=10, N=100 -> 1000 elems (4000 bytes) - fits in 5000
# req1: adding M=10, N=100 -> new_M=20, new_N=200 -> 4000 elems > 1250
(
torch.tensor([100, 100, 100]),
torch.tensor([10, 10, 10]),
1000, # workspace allows all
5000, # 1250 float32 elems -> forces split
[
(slice(0, 1), slice(0, 10)),
(slice(1, 2), slice(0, 10)),
(slice(2, 3), slice(0, 10)),
],
),
# Both constraints satisfied - all fit in one chunk
(
torch.tensor([10, 10, 10]),
torch.tensor([5, 5, 5]),
100,
10000, # 2500 elems, M*N = 15*30 = 450 < 2500
[(slice(0, 3), slice(0, 15))],
),
# Workspace constraint triggers first
(
torch.tensor([50, 50, 50]),
torch.tensor([1, 1, 1]),
50, # workspace only fits one at a time
1000000, # logits budget is huge
[
(slice(0, 1), slice(0, 1)),
(slice(1, 2), slice(0, 1)),
(slice(2, 3), slice(0, 1)),
],
),
# Greedy filling: first two fit, third doesn't
# req0: M=5, N=10 -> 50 elems
# req0+1: M=10, N=20 -> 200 elems <= 250
# req0+1+2: M=15, N=30 -> 450 elems > 250
(
torch.tensor([10, 10, 10]),
torch.tensor([5, 5, 5]),
100,
1000, # 250 elems
[(slice(0, 2), slice(0, 10)), (slice(2, 3), slice(0, 5))],
),
],
)
def test_split_indexer_prefill_chunks(
seq_lens, query_lens, workspace_size, max_logits_bytes, expected
):
out = split_indexer_prefill_chunks(
seq_lens,
query_lens,
workspace_size,
max_logits_bytes,
)
assert out == expected
def test_split_indexer_prefill_chunks_single_request_overflow():
"""Test that single request exceeding budget is sub-chunked on query dim."""
seq_lens = torch.tensor([1000, 50])
query_lens = torch.tensor([100, 5])
out = split_indexer_prefill_chunks(seq_lens, query_lens, 2000, 1000)
# max_logits_elems = 250, N=1000 -> max_q = 1 -> 100 query sub-chunks
expected = [(slice(0, 1), slice(i, i + 1)) for i in range(100)]
# req1: M=5, N=50 -> 250 elems fits budget
expected.append((slice(1, 2), slice(0, 5)))
assert out == expected
def test_triton_convert_returns_valid_counts():
"""Test that return_valid_counts correctly counts non-negative indices."""
device = torch.device("cuda")
+8 -2
View File
@@ -36,14 +36,20 @@ MESSAGES = [
]
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
@pytest.mark.parametrize(
"model_name",
[
pytest.param("Qwen/Qwen3.5-4B", marks=[large_gpu_mark(min_gb=40)]),
pytest.param(
"nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8",
marks=[large_gpu_mark(min_gb=80)] + multi_gpu_marks(num_gpus=4),
marks=[large_gpu_mark(min_gb=80)]
+ multi_gpu_marks(num_gpus=4)
+ [
pytest.mark.skipif(
not current_platform.is_cuda(),
reason="modelopt quantization is supported only on CUDA",
)
],
),
],
)
+19 -11
View File
@@ -259,8 +259,25 @@ async def test_abort_during_final_step(async_scheduling: bool):
# Wait for generation to complete
await gen_task
# Give the scheduler a moment to finish cleanup
await asyncio.sleep(0.1)
# Poll for the KV connector to record the finish status
timeout = 5.0
start = time.time()
captured_statuses = []
while time.time() - start < timeout:
with open(status_file) as f4:
status_lines = f4.read().strip().split("\n")
captured_statuses = [
line
for line in status_lines
if line and line.startswith("FINISHED_")
]
if captured_statuses:
break
await asyncio.sleep(0.05)
else:
raise TimeoutError(
"Timeout waiting for KV connector to record finish status."
)
# Verify we got output
assert len(outputs) > 0, "Should have received at least one output"
@@ -275,15 +292,6 @@ async def test_abort_during_final_step(async_scheduling: bool):
f"'{final_output.outputs[0].finish_reason}'. "
)
with open(status_file) as f4:
status_lines = f4.read().strip().split("\n")
# Filter for actual finish statuses (not INIT or empty lines)
captured_statuses = [
line
for line in status_lines
if line and line.startswith("FINISHED_")
]
assert len(captured_statuses) >= 1, (
f"Expected at least 1 captured finish status, got "
f"{len(captured_statuses)}. File content: {status_lines}"
@@ -20,7 +20,7 @@ def server():
"--reasoning-parser",
"qwen3",
"--reasoning-config",
'{"think_start_str": "<think>", "think_end_str": "</think>"}',
'{"reasoning_start_str": "<think>", "reasoning_end_str": "</think>"}',
"--max-model-len",
"2048",
"--enforce-eager",
@@ -0,0 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from tests.v1.kv_connector.unit.offloading_connector.utils import (
request_runner,
)
__all__ = ["request_runner"]
@@ -0,0 +1,151 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
OffloadingConnectorStats,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
OffloadingConnector,
)
def test_build_kv_connector_stats_with_none():
"""Test that build_kv_connector_stats returns empty stats when given None."""
stats = OffloadingConnector.build_kv_connector_stats(data=None)
assert stats is not None
assert isinstance(stats, OffloadingConnectorStats)
assert len(stats.data) == 0
assert stats.is_empty()
def test_build_kv_connector_stats_with_empty_dict():
"""Test that build_kv_connector_stats returns empty stats with empty dict."""
stats = OffloadingConnector.build_kv_connector_stats(data={})
assert stats is not None
assert isinstance(stats, OffloadingConnectorStats)
assert len(stats.data) == 0
assert stats.is_empty()
def test_build_kv_connector_stats_reconstructs_offload_stats():
"""Test that OffloadingConnector stats are properly reconstructed with
correct data."""
serialized_data = {
"CPU_to_GPU": [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
],
"GPU_to_CPU": [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
],
}
stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
offload_connector_stats = stats
assert isinstance(offload_connector_stats, OffloadingConnectorStats)
assert offload_connector_stats.data["CPU_to_GPU"] == [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
]
assert offload_connector_stats.data["GPU_to_CPU"] == [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
]
def test_aggregate_same_connector():
"""Test aggregating stats from the same connector type."""
stats1 = OffloadingConnectorStats(
data={
"CPU_to_GPU": [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
],
"GPU_to_CPU": [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
],
}
)
stats2 = OffloadingConnectorStats(
data={
"CPU_to_GPU": [
{"op_size": 3, "op_time": 0.2},
{"op_size": 7, "op_time": 0.9},
],
"GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
}
)
result = stats1.aggregate(stats2)
assert result is stats1 # Should return self
offload_connector_stats = result
assert offload_connector_stats.data["CPU_to_GPU"] == [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
{"op_size": 3, "op_time": 0.2},
{"op_size": 7, "op_time": 0.9},
]
assert offload_connector_stats.data["GPU_to_CPU"] == [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
{"op_size": 16, "op_time": 2},
]
def test_reduce():
"""Test that reduce() correctly reduces all nested connector stats."""
stats = OffloadingConnectorStats(
data={
"CPU_to_GPU": [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
{"op_size": 3, "op_time": 0.2},
{"op_size": 7, "op_time": 0.9},
],
"GPU_to_CPU": [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
{"op_size": 16, "op_time": 2},
],
}
)
reduced = stats.reduce()
assert isinstance(reduced, dict)
# Check that the stats were reduced (should have aggregated values)
assert "CPU_to_GPU_total_bytes" in reduced
assert "CPU_to_GPU_total_time" in reduced
assert "GPU_to_CPU_total_bytes" in reduced
assert "GPU_to_CPU_total_time" in reduced
assert reduced["CPU_to_GPU_total_bytes"] == 34
assert reduced["CPU_to_GPU_total_time"] == 2.6
assert reduced["GPU_to_CPU_total_time"] == 2.3
assert reduced["GPU_to_CPU_total_bytes"] == 19
def test_reset():
"""Test that reset() resets all nested connector stats."""
offload_connector_stats = OffloadingConnectorStats(
data={
"CPU_to_GPU": [
{"op_size": 3, "op_time": 0.2},
{"op_size": 7, "op_time": 0.9},
],
"GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
}
)
assert not offload_connector_stats.is_empty()
offload_connector_stats.reset()
# After reset, stats should be empty
assert offload_connector_stats.is_empty()
assert len(offload_connector_stats.data) == 0
@@ -0,0 +1,341 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
import pytest
from tests.v1.kv_connector.unit.offloading_connector.utils import (
generate_store_output,
)
from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
from vllm.distributed.kv_events import BlockRemoved, BlockStored
from vllm.v1.core.kv_cache_utils import BlockHash
from vllm.v1.kv_offload.abstract import OffloadingEvent
from vllm.v1.request import RequestStatus
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_offloading_connector(request_runner, async_scheduling: bool):
offloaded_block_size = 12
gpu_block_size = 4
num_gpu_blocks = 100
block_size_factor = offloaded_block_size // gpu_block_size
runner = request_runner(
offloaded_block_size=offloaded_block_size,
gpu_block_size=gpu_block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
# 3 blocks, store just the middle block (skip first and last)
# blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
runner.new_request(token_ids=[0] * offloaded_block_size * 3)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
)
runner.run(decoded_tokens=[0])
# add block missing 1 token -> no offload
runner.run(
decoded_tokens=[0] * (offloaded_block_size - 1),
expected_stored_gpu_block_indexes=(3, 4, 5),
)
runner.manager.prepare_store.assert_not_called()
# +1 token -> single block, fail prepare_store
runner.manager.prepare_store.side_effect = lambda block_hashes: None
runner.run(decoded_tokens=[0])
runner.manager.prepare_store.assert_called()
# 1 more block (+ token for async scheduling)
# now set block_hashes_to_store = []
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
# 1 more block (+ token for kicking off offloading)
# now check touch was called with all 6 blocks
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[0] * (offloaded_block_size + 1),
expected_stored_gpu_block_indexes=(15, 16, 17),
)
runner.manager.touch.assert_called()
block_hashes1 = list(runner.manager.touch.call_args.args[0])
assert len(block_hashes1) == 6
# terminate request
runner.run(decoded_tokens=[EOS_TOKEN_ID])
# create a new request differing only on the last token
runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
runner.run(decoded_tokens=[0])
runner.manager.touch.assert_called()
block_hashes2 = list(runner.manager.touch.call_args.args[0])
assert len(block_hashes2) == 6
# verify hashes are the same, except for the last block
assert block_hashes1[:5] == block_hashes2[:5]
assert block_hashes1[5] != block_hashes2[5]
# terminate request
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
)
# full_block_tokens - num_computed_tokens < offloaded_block_size
runner.new_request(
token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.run(decoded_tokens=[EOS_TOKEN_ID])
runner.manager.lookup.assert_not_called()
# single block lookup with no hits
runner.new_request(token_ids=[1] * offloaded_block_size)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.run(decoded_tokens=[EOS_TOKEN_ID])
runner.manager.lookup.assert_called()
assert len(list(runner.manager.lookup.call_args.args[0])) == 1
# single block lookup with a hit
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
)
# single block lookup with a hit in a middle block
runner.new_request(
token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
)
# test take_events
def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
return [BlockHash(str(i).encode()) for i in int_hashes]
def take_events() -> Iterable[OffloadingEvent]:
yield OffloadingEvent(
block_hashes=to_hashes([1, 2, 3]), block_size=16, medium="A", removed=False
)
yield OffloadingEvent(
block_hashes=to_hashes([4, 5, 6]), block_size=32, medium="B", removed=True
)
runner.manager.take_events.side_effect = take_events
events = list(runner.scheduler_connector.take_events())
assert len(events) == 2
event = events[0]
assert isinstance(event, BlockStored)
assert event.block_hashes == to_hashes([1, 2, 3])
assert event.block_size == 16
assert event.medium == "A"
assert event.token_ids == []
assert event.parent_block_hash is None
assert event.lora_id is None
assert event.lora_name is None
event = events[1]
assert isinstance(event, BlockRemoved)
assert event.block_hashes == to_hashes([4, 5, 6])
assert event.medium == "B"
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_request_preemption(request_runner, async_scheduling: bool):
offloaded_block_size = 12
gpu_block_size = 4
num_gpu_blocks = 100
runner = request_runner(
offloaded_block_size=offloaded_block_size,
gpu_block_size=gpu_block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
num_free_blocks_empty = free_block_queue.num_free_blocks
# 2 blocks, store all, without flushing
# blocks = [0, 1, 2], [3, 4, 5]
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[0],
complete_transfers=False,
)
# decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
complete_transfers=False,
)
# simulate KV cache running out of space
free_block_queue.num_free_blocks = 0
# request should be preempted now
runner.run(
decoded_tokens=[],
complete_transfers=False,
expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
)
# restore KV cache space and reset GPU prefix cache
free_block_queue.num_free_blocks = num_free_blocks_empty
runner.scheduler.reset_prefix_cache()
# request should now return from preemption
# re-load [0, ..., 8] from the CPU and store [9, 10, 11]
runner.manager.lookup.return_value = 3
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[0] * gpu_block_size,
expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=(9, 10, 11),
)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
offloaded_block_size = 12
gpu_block_size = 4
num_gpu_blocks = 100
runner = request_runner(
offloaded_block_size=offloaded_block_size,
gpu_block_size=gpu_block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
# store 1 blocks
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=(0, 1, 2),
)
# start a request to load the first block, but don't complete
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[],
complete_transfers=False,
)
# request triggered a load
transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
assert transfer_jobs
# start a new request to load the same first block
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[],
complete_transfers=False,
)
# request did not trigger a load
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
# complete transfers
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_loaded_gpu_block_indexes=(0, 1, 2),
)
# second request will use the GPU prefix cache
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_abort_loading_requests(request_runner, async_scheduling: bool):
offloaded_block_size = 12
gpu_block_size = 4
num_gpu_blocks = 100
runner = request_runner(
offloaded_block_size=offloaded_block_size,
gpu_block_size=gpu_block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
# store 1 blocks
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=(0, 1, 2),
)
# start a request to load the first block, but don't complete
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[],
complete_transfers=False,
)
# request triggered a load
transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
assert transfer_jobs
# abort request
req_id = str(runner.req_id)
runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
# verify request is not deleted
assert req_id in runner.scheduler.requests
# complete loading request
runner.run(
decoded_tokens=[],
expected_loaded_gpu_block_indexes=(0, 1, 2),
)
# assert request is deleted
assert req_id not in runner.scheduler.requests
@@ -0,0 +1,504 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections import defaultdict
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm.platforms import current_platform
from vllm.utils.torch_utils import get_dtype_size
from vllm.v1.attention.backend import AttentionBackend
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.attention.backends.utils import set_kv_cache_layout
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
MambaSpec,
MLAAttentionSpec,
UniformTypeKVCacheSpecs,
)
from vllm.v1.kv_offload.spec import (
CanonicalKVCacheRef,
CanonicalKVCaches,
OffloadingSpec,
)
NUM_BLOCKS = 10
BLOCK_SIZE = 16
NUM_KV_HEADS = 4
HEAD_SIZE = 64
DTYPE = torch.float16
# Attention backends to test
ATTN_BACKENDS: list[str] = []
if current_platform.is_cuda():
ATTN_BACKENDS = [
"FLASH_ATTN",
"FLEX_ATTENTION",
"FLASHINFER",
"TRITON_ATTN",
]
elif current_platform.is_rocm():
ATTN_BACKENDS = ["TRITON_ATTN"]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _allocate_and_reshape_kv_caches(
kv_cache_config: KVCacheConfig,
attn_groups: list[list],
device: torch.device,
):
"""
Use the real GPUModelRunner allocation and reshape methods to produce
kv_caches, just like the model runner does during initialization.
"""
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
# Some backends (e.g. FlashAttention) query the KV cache layout during
# reshape, which ultimately calls get_current_vllm_config(). Setting
# the layout override avoids needing a full VllmConfig context.
set_kv_cache_layout("NHD")
try:
runner = object.__new__(GPUModelRunner)
runner.device = device
runner.runner_only_attn_layers = set()
runner.attn_groups = attn_groups
runner.kv_cache_config = kv_cache_config
runner.cache_config = MagicMock(cache_dtype="auto")
runner.shared_kv_cache_layers = {}
runner.model_config = MagicMock()
runner.model_config.hf_config.model_type = ""
runner.compilation_config = MagicMock(
static_forward_context=defaultdict(MagicMock)
)
runner.kv_caches = []
kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
return runner.initialize_kv_cache_tensors(kv_cache_config, kernel_block_sizes)
finally:
set_kv_cache_layout(None)
def _make_mock_layer(backend_cls: type[AttentionBackend]):
"""
Create a mock AttentionLayerBase whose get_attn_backend returns backend_cls.
"""
layer = MagicMock()
layer.get_attn_backend.return_value = backend_cls
return layer
def _make_worker(kv_cache_config: KVCacheConfig):
"""
Create an OffloadingConnectorWorker with mocked dependencies.
"""
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
OffloadingConnectorWorker,
)
spec = MagicMock(spec=OffloadingSpec)
spec.kv_cache_config = kv_cache_config
spec.vllm_config = MagicMock()
spec.get_handlers.return_value = iter([])
worker = OffloadingConnectorWorker(spec=spec)
worker.worker = MagicMock()
return worker, spec
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("backend", ATTN_BACKENDS)
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.offloading"
".worker.get_layers_from_vllm_config"
)
def test_register_kv_caches(mock_get_layers, backend):
"""Test register_kv_caches with multiple groups covering all layer types.
Creates one FullAttention group, one MLA group, one Mamba group, and
one Mamba-padded group. Each group has GROUP_SIZE layers.
KVCacheTensors are shared across all groups mirroring the real allocation
in kv_cache_utils.py: tensor i is shared by layer i from every group.
The padded-mamba group has a different page size so its layers get their
own dedicated tensors.
Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
kv_caches, which automatically applies
_update_hybrid_attention_mamba_layout for hybrid models.
Verifies that the canonicalized CanonicalKVCaches has the correct
block tensors, tensor_idx references, and page sizes across all groups.
"""
from vllm.v1.attention.backends.mla.indexer import (
DeepseekV32IndexerBackend,
)
from vllm.v1.worker.utils import AttentionGroup
MLA_HEAD_SIZE = NUM_KV_HEADS * HEAD_SIZE * 2
# padded mamba (missing HEAD_SIZE)
CONV_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS, HEAD_SIZE)
UNALIGNED_SSM_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS - 1, HEAD_SIZE)
PAGE_SIZE_BYTES = 2 * BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * get_dtype_size(DTYPE)
unaligned_mamba_page_size = PAGE_SIZE_BYTES - HEAD_SIZE * get_dtype_size(DTYPE)
# unpadded mamba (fills page exactly)
ALIGNED_SSM_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS, HEAD_SIZE)
backend_cls = AttentionBackendEnum[backend].get_class()
attn_spec = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=NUM_KV_HEADS,
head_size=HEAD_SIZE,
dtype=DTYPE,
)
mla_spec = MLAAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=1,
head_size=MLA_HEAD_SIZE,
dtype=DTYPE,
)
unaligned_mamba_spec = MambaSpec(
block_size=BLOCK_SIZE,
shapes=(CONV_STATE_SHAPE, UNALIGNED_SSM_STATE_SHAPE),
dtypes=(DTYPE, DTYPE),
page_size_padded=PAGE_SIZE_BYTES,
)
aligned_mamba_spec = MambaSpec(
block_size=BLOCK_SIZE,
shapes=(CONV_STATE_SHAPE, ALIGNED_SSM_STATE_SHAPE),
dtypes=(DTYPE, DTYPE),
page_size_padded=PAGE_SIZE_BYTES,
)
assert attn_spec.page_size_bytes == PAGE_SIZE_BYTES
assert mla_spec.page_size_bytes == PAGE_SIZE_BYTES
assert unaligned_mamba_spec.page_size_bytes == PAGE_SIZE_BYTES
assert aligned_mamba_spec.page_size_bytes == PAGE_SIZE_BYTES
GROUP_SIZE = 3
# -- Build per-group layer info ----------------------------------------
layer_idx = 0
attn_layer_names = []
for _ in range(GROUP_SIZE):
attn_layer_names.append(f"model.layers.{layer_idx}.self_attn")
layer_idx += 1
mla_layer_names = []
for _ in range(GROUP_SIZE):
mla_layer_names.append(f"model.layers.{layer_idx}.self_attn")
layer_idx += 1
unaligned_mamba_layer_names = []
for _ in range(GROUP_SIZE):
unaligned_mamba_layer_names.append(f"model.layers.{layer_idx}.mamba_unpadded")
layer_idx += 1
aligned_mamba_layer_names = []
for _ in range(GROUP_SIZE - 1):
aligned_mamba_layer_names.append(f"model.layers.{layer_idx}.mamba_padded")
layer_idx += 1
layer_groups = [
attn_layer_names,
mla_layer_names,
unaligned_mamba_layer_names,
aligned_mamba_layer_names,
]
kv_cache_tensors: list[KVCacheTensor] = []
for i in range(GROUP_SIZE):
shared_by: list[str] = []
for group_layer_names in layer_groups:
if len(group_layer_names) > i:
shared_by.append(group_layer_names[i])
kv_cache_tensors.append(
KVCacheTensor(
size=PAGE_SIZE_BYTES * NUM_BLOCKS,
shared_by=shared_by,
)
)
kv_cache_groups = [
KVCacheGroupSpec(layer_names=attn_layer_names, kv_cache_spec=attn_spec),
KVCacheGroupSpec(layer_names=mla_layer_names, kv_cache_spec=mla_spec),
KVCacheGroupSpec(
layer_names=unaligned_mamba_layer_names, kv_cache_spec=unaligned_mamba_spec
),
KVCacheGroupSpec(
layer_names=aligned_mamba_layer_names, kv_cache_spec=aligned_mamba_spec
),
]
attn_groups = [
[
AttentionGroup(
backend=backend_cls,
layer_names=attn_layer_names,
kv_cache_spec=attn_spec,
kv_cache_group_id=0,
),
AttentionGroup(
backend=DeepseekV32IndexerBackend,
layer_names=mla_layer_names,
kv_cache_spec=mla_spec,
kv_cache_group_id=1,
),
AttentionGroup(
backend=DeepseekV32IndexerBackend, # unused for mamba
layer_names=unaligned_mamba_layer_names,
kv_cache_spec=unaligned_mamba_spec,
kv_cache_group_id=2,
),
AttentionGroup(
backend=DeepseekV32IndexerBackend, # unused for mamba
layer_names=aligned_mamba_layer_names,
kv_cache_spec=aligned_mamba_spec,
kv_cache_group_id=3,
),
]
]
kv_cache_config = KVCacheConfig(
num_blocks=NUM_BLOCKS,
kv_cache_tensors=kv_cache_tensors,
kv_cache_groups=kv_cache_groups,
)
kv_caches = _allocate_and_reshape_kv_caches(
kv_cache_config,
attn_groups,
device=torch.device("cuda:0"),
)
mock_layers: dict[str, MagicMock] = {}
for layer_name in attn_layer_names:
mock_layers[layer_name] = _make_mock_layer(backend_cls)
for layer_name in mla_layer_names:
mock_layers[layer_name] = _make_mock_layer(DeepseekV32IndexerBackend)
mock_get_layers.return_value = mock_layers
worker, spec = _make_worker(kv_cache_config)
worker.register_kv_caches(kv_caches)
canonical = spec.get_handlers.call_args[0][0]
assert isinstance(canonical, CanonicalKVCaches)
# -- Expected block tensors ----------------------------------------------
# All tensors have the same padded page size (PAGE_SIZE_BYTES).
# Tensor 0: shared by attn[0], mla[0], mamba_unaligned[0], mamba_aligned[0]
# Tensor 1: shared by attn[1], mla[1], mamba_unaligned[1], mamba_aligned[1]
# Tensor 2: shared by attn[2], mla[2], mamba_unaligned[2]
# (mamba_aligned has only GROUP_SIZE-1 = 2 layers)
expected_tensors = [
(NUM_BLOCKS, PAGE_SIZE_BYTES),
(NUM_BLOCKS, PAGE_SIZE_BYTES),
(NUM_BLOCKS, PAGE_SIZE_BYTES),
]
# -- Expected group data refs (order matches kv_cache_groups) -------------
ref = CanonicalKVCacheRef
expected_group_refs = [
# attn group: layers attn[0..2] → tensors 0,1,2 with full page size
[
ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
ref(tensor_idx=2, page_size_bytes=PAGE_SIZE_BYTES),
],
# mla group: layers mla[0..2] → tensors 0,1,2 with full page size
[
ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
ref(tensor_idx=2, page_size_bytes=PAGE_SIZE_BYTES),
],
# unaligned mamba group: layers [0..2] → tensors 0,1,2 with unaligned page
[
ref(tensor_idx=0, page_size_bytes=unaligned_mamba_page_size),
ref(tensor_idx=1, page_size_bytes=unaligned_mamba_page_size),
ref(tensor_idx=2, page_size_bytes=unaligned_mamba_page_size),
],
# aligned mamba group: layers [0..1] → tensors 0,1 with full page size
[
ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
],
]
# Verify block tensors
assert len(canonical.tensors) == len(expected_tensors)
for block_tensor, (exp_num_blocks, exp_page_size) in zip(
canonical.tensors, expected_tensors
):
tensor = block_tensor.tensor
assert tensor.dtype == torch.int8
assert tensor.shape == (exp_num_blocks, exp_page_size)
assert block_tensor.page_size_bytes == exp_page_size
# Verify group data refs
assert len(canonical.group_data_refs) == len(expected_group_refs)
for actual_refs, exp_refs in zip(canonical.group_data_refs, expected_group_refs):
assert len(actual_refs) == len(exp_refs)
for actual, expected in zip(actual_refs, exp_refs):
assert actual.tensor_idx == expected.tensor_idx
assert actual.page_size_bytes == expected.page_size_bytes
@pytest.mark.parametrize("backend", ATTN_BACKENDS)
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.offloading"
".worker.get_layers_from_vllm_config"
)
def test_register_kv_caches_uniform_type(mock_get_layers, backend):
"""Test register_kv_caches with UniformTypeKVCacheSpecs.
Two attention layers use the same backend but different num_kv_heads,
giving them different per-layer page sizes. Each has its own
KVCacheTensor and are wrapped in a UniformTypeKVCacheSpecs group.
Verifies that each layer gets the correct tensor_idx and
page_size_bytes in its block data ref.
"""
from vllm.v1.worker.utils import AttentionGroup
backend_cls = AttentionBackendEnum[backend].get_class()
layer_a = "model.layers.0.self_attn"
layer_b = "model.layers.1.self_attn"
spec_a = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=NUM_KV_HEADS,
head_size=HEAD_SIZE,
dtype=DTYPE,
)
spec_b = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=NUM_KV_HEADS * 2,
head_size=HEAD_SIZE,
dtype=DTYPE,
)
assert spec_a.page_size_bytes != spec_b.page_size_bytes
uniform_spec = UniformTypeKVCacheSpecs(
block_size=BLOCK_SIZE,
kv_cache_specs={layer_a: spec_a, layer_b: spec_b},
)
kv_cache_config = KVCacheConfig(
num_blocks=NUM_BLOCKS,
kv_cache_tensors=[
KVCacheTensor(
size=spec_a.page_size_bytes * NUM_BLOCKS,
shared_by=[layer_a],
),
KVCacheTensor(
size=spec_b.page_size_bytes * NUM_BLOCKS,
shared_by=[layer_b],
),
],
kv_cache_groups=[
KVCacheGroupSpec(
layer_names=[layer_a, layer_b],
kv_cache_spec=uniform_spec,
)
],
)
attn_groups = [
[
AttentionGroup(
backend=backend_cls,
layer_names=[layer_a],
kv_cache_spec=spec_a,
kv_cache_group_id=0,
),
AttentionGroup(
backend=backend_cls,
layer_names=[layer_b],
kv_cache_spec=spec_b,
kv_cache_group_id=0,
),
]
]
kv_caches = _allocate_and_reshape_kv_caches(
kv_cache_config,
attn_groups,
device=torch.device("cuda:0"),
)
mock_get_layers.return_value = {
layer_a: _make_mock_layer(backend_cls),
layer_b: _make_mock_layer(backend_cls),
}
worker, spec = _make_worker(kv_cache_config)
worker.register_kv_caches(kv_caches)
canonical = spec.get_handlers.call_args[0][0]
assert isinstance(canonical, CanonicalKVCaches)
unbinds = backend_cls.get_name() in ("FLASH_ATTN", "FLEX_ATTENTION")
tensors_per_layer = 2 if unbinds else 1
for block_tensor in canonical.tensors:
assert block_tensor.tensor.dtype == torch.int8
# Single group with refs from both layers
assert len(canonical.group_data_refs) == 1
group_refs = canonical.group_data_refs[0]
assert len(group_refs) == 2 * tensors_per_layer
if unbinds:
half_a = spec_a.page_size_bytes // 2
half_b = spec_b.page_size_bytes // 2
assert len(canonical.tensors) == 4
assert canonical.tensors[0].page_size_bytes == half_a
assert canonical.tensors[1].page_size_bytes == half_a
assert canonical.tensors[2].page_size_bytes == half_b
assert canonical.tensors[3].page_size_bytes == half_b
assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, half_a)
assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, half_a)
assert canonical.tensors[2].tensor.shape == (NUM_BLOCKS, half_b)
assert canonical.tensors[3].tensor.shape == (NUM_BLOCKS, half_b)
assert group_refs[0] == CanonicalKVCacheRef(
tensor_idx=0, page_size_bytes=half_a
)
assert group_refs[1] == CanonicalKVCacheRef(
tensor_idx=1, page_size_bytes=half_a
)
assert group_refs[2] == CanonicalKVCacheRef(
tensor_idx=2, page_size_bytes=half_b
)
assert group_refs[3] == CanonicalKVCacheRef(
tensor_idx=3, page_size_bytes=half_b
)
else:
assert len(canonical.tensors) == 2
assert canonical.tensors[0].page_size_bytes == spec_a.page_size_bytes
assert canonical.tensors[1].page_size_bytes == spec_b.page_size_bytes
assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
assert group_refs[0] == CanonicalKVCacheRef(
tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
)
assert group_refs[1] == CanonicalKVCacheRef(
tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
)
@@ -9,16 +9,17 @@ from unittest.mock import MagicMock
import pytest
import torch
from tests.v1.kv_connector.unit.utils import (
EOS_TOKEN_ID,
create_model_runner_output,
create_vllm_config,
)
from vllm import SamplingParams
from vllm.config import KVTransferConfig, VllmConfig
from vllm.distributed.kv_events import BlockRemoved, BlockStored
from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
OffloadingConnectorMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
OffloadingConnectorStats,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
OffloadingConnector,
)
@@ -39,7 +40,6 @@ from vllm.v1.kv_cache_interface import (
)
from vllm.v1.kv_offload.abstract import (
LoadStoreSpec,
OffloadingEvent,
OffloadingManager,
PrepareStoreOutput,
)
@@ -51,15 +51,9 @@ from vllm.v1.kv_offload.worker.worker import (
TransferSpec,
)
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
from vllm.v1.request import Request, RequestStatus
from vllm.v1.request import Request
from vllm.v1.structured_output import StructuredOutputManager
from .utils import (
EOS_TOKEN_ID,
create_model_runner_output,
create_vllm_config,
)
class MockLoadStoreSpec(LoadStoreSpec):
def __init__(self, block_hashes: Iterable[BlockHash]):
@@ -125,7 +119,7 @@ class MockOffloadingSpec(OffloadingSpec):
return self.manager
def get_handlers(
self, _, __
self, _
) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler
yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler
@@ -179,7 +173,7 @@ class RequestRunner:
kv_role="kv_both",
kv_connector_extra_config={
"spec_name": "MockOffloadingSpec",
"spec_module_path": "tests.v1.kv_connector.unit.test_offloading_connector", # noqa: E501
"spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501
"block_size": offloaded_block_size,
},
)
@@ -217,10 +211,12 @@ class RequestRunner:
)
# register worker kv_caches to enable OffloadingWorker creations
self.worker_connector.register_cross_layers_kv_cache(
kv_cache=torch.empty(0),
attn_backend=FlashAttentionBackend,
)
# set_current_vllm_config is needed for get_kv_cache_layout() to work
with set_current_vllm_config(vllm_config):
self.worker_connector.register_cross_layers_kv_cache(
kv_cache=torch.empty(0),
attn_backend=FlashAttentionBackend,
)
# extract connector of scheduler
scheduler_connector = self.scheduler.connector
@@ -521,471 +517,3 @@ def generate_store_output(block_hashes: Iterable[BlockHash]):
store_spec=MockLoadStoreSpec(block_hashes),
block_hashes_evicted=[],
)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_offloading_connector(request_runner, async_scheduling: bool):
offloaded_block_size = 12
gpu_block_size = 4
num_gpu_blocks = 100
block_size_factor = offloaded_block_size // gpu_block_size
runner = request_runner(
offloaded_block_size=offloaded_block_size,
gpu_block_size=gpu_block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
# 3 blocks, store just the middle block (skip first and last)
# blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
runner.new_request(token_ids=[0] * offloaded_block_size * 3)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
)
runner.run(decoded_tokens=[0])
# add block missing 1 token -> no offload
runner.run(
decoded_tokens=[0] * (offloaded_block_size - 1),
expected_stored_gpu_block_indexes=(3, 4, 5),
)
runner.manager.prepare_store.assert_not_called()
# +1 token -> single block, fail prepare_store
runner.manager.prepare_store.side_effect = lambda block_hashes: None
runner.run(decoded_tokens=[0])
runner.manager.prepare_store.assert_called()
# 1 more block (+ token for async scheduling)
# now set block_hashes_to_store = []
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
# 1 more block (+ token for kicking off offloading)
# now check touch was called with all 6 blocks
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[0] * (offloaded_block_size + 1),
expected_stored_gpu_block_indexes=(15, 16, 17),
)
runner.manager.touch.assert_called()
block_hashes1 = list(runner.manager.touch.call_args.args[0])
assert len(block_hashes1) == 6
# terminate request
runner.run(decoded_tokens=[EOS_TOKEN_ID])
# create a new request differing only on the last token
runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
runner.run(decoded_tokens=[0])
runner.manager.touch.assert_called()
block_hashes2 = list(runner.manager.touch.call_args.args[0])
assert len(block_hashes2) == 6
# verify hashes are the same, except for the last block
assert block_hashes1[:5] == block_hashes2[:5]
assert block_hashes1[5] != block_hashes2[5]
# terminate request
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
)
# full_block_tokens - num_computed_tokens < offloaded_block_size
runner.new_request(
token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.run(decoded_tokens=[EOS_TOKEN_ID])
runner.manager.lookup.assert_not_called()
# single block lookup with no hits
runner.new_request(token_ids=[1] * offloaded_block_size)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.run(decoded_tokens=[EOS_TOKEN_ID])
runner.manager.lookup.assert_called()
assert len(list(runner.manager.lookup.call_args.args[0])) == 1
# single block lookup with a hit
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
)
# single block lookup with a hit in a middle block
runner.new_request(
token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
)
# test take_events
def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
return [BlockHash(str(i).encode()) for i in int_hashes]
def take_events() -> Iterable[OffloadingEvent]:
yield OffloadingEvent(
block_hashes=to_hashes([1, 2, 3]), block_size=16, medium="A", removed=False
)
yield OffloadingEvent(
block_hashes=to_hashes([4, 5, 6]), block_size=32, medium="B", removed=True
)
runner.manager.take_events.side_effect = take_events
events = list(runner.scheduler_connector.take_events())
assert len(events) == 2
event = events[0]
assert isinstance(event, BlockStored)
assert event.block_hashes == to_hashes([1, 2, 3])
assert event.block_size == 16
assert event.medium == "A"
assert event.token_ids == []
assert event.parent_block_hash is None
assert event.lora_id is None
assert event.lora_name is None
event = events[1]
assert isinstance(event, BlockRemoved)
assert event.block_hashes == to_hashes([4, 5, 6])
assert event.medium == "B"
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_request_preemption(request_runner, async_scheduling: bool):
offloaded_block_size = 12
gpu_block_size = 4
num_gpu_blocks = 100
runner = request_runner(
offloaded_block_size=offloaded_block_size,
gpu_block_size=gpu_block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
num_free_blocks_empty = free_block_queue.num_free_blocks
# 2 blocks, store all, without flushing
# blocks = [0, 1, 2], [3, 4, 5]
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[0],
complete_transfers=False,
)
# decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
complete_transfers=False,
)
# simulate KV cache running out of space
free_block_queue.num_free_blocks = 0
# request should be preempted now
runner.run(
decoded_tokens=[],
complete_transfers=False,
expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
)
# restore KV cache space and reset GPU prefix cache
free_block_queue.num_free_blocks = num_free_blocks_empty
runner.scheduler.reset_prefix_cache()
# request should now return from preemption
# re-load [0, ..., 8] from the CPU and store [9, 10, 11]
runner.manager.lookup.return_value = 3
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[0] * gpu_block_size,
expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=(9, 10, 11),
)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
offloaded_block_size = 12
gpu_block_size = 4
num_gpu_blocks = 100
runner = request_runner(
offloaded_block_size=offloaded_block_size,
gpu_block_size=gpu_block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
# store 1 blocks
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=(0, 1, 2),
)
# start a request to load the first block, but don't complete
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[],
complete_transfers=False,
)
# request triggered a load
transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
assert transfer_jobs
# start a new request to load the same first block
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[],
complete_transfers=False,
)
# request did not trigger a load
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
# complete transfers
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output([])
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_loaded_gpu_block_indexes=(0, 1, 2),
)
# second request will use the GPU prefix cache
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_abort_loading_requests(request_runner, async_scheduling: bool):
offloaded_block_size = 12
gpu_block_size = 4
num_gpu_blocks = 100
runner = request_runner(
offloaded_block_size=offloaded_block_size,
gpu_block_size=gpu_block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
)
# store 1 blocks
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = (
lambda block_hashes: generate_store_output(block_hashes)
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=(0, 1, 2),
)
# start a request to load the first block, but don't complete
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[],
complete_transfers=False,
)
# request triggered a load
transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
assert transfer_jobs
# abort request
req_id = str(runner.req_id)
runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
# verify request is not deleted
assert req_id in runner.scheduler.requests
# complete loading request
runner.run(
decoded_tokens=[],
expected_loaded_gpu_block_indexes=(0, 1, 2),
)
# assert request is deleted
assert req_id not in runner.scheduler.requests
class TestOffloadingConnectorStats:
"""Tests for OffloadingConnector stats reconstruction and operations."""
def test_build_kv_connector_stats_with_none(self):
"""Test that build_kv_connector_stats returns empty stats when given None."""
stats = OffloadingConnector.build_kv_connector_stats(data=None)
assert stats is not None
assert isinstance(stats, OffloadingConnectorStats)
assert len(stats.data) == 0
assert stats.is_empty()
def test_build_kv_connector_stats_with_empty_dict(self):
"""Test that build_kv_connector_stats returns empty stats with empty dict."""
stats = OffloadingConnector.build_kv_connector_stats(data={})
assert stats is not None
assert isinstance(stats, OffloadingConnectorStats)
assert len(stats.data) == 0
assert stats.is_empty()
def test_build_kv_connector_stats_reconstructs_offload_stats(self):
"""Test that OffloadingConnector stats are properly reconstructed with
correct data."""
serialized_data = {
"CPU_to_GPU": [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
],
"GPU_to_CPU": [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
],
}
stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
offload_connector_stats = stats
assert isinstance(offload_connector_stats, OffloadingConnectorStats)
assert offload_connector_stats.data["CPU_to_GPU"] == [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
]
assert offload_connector_stats.data["GPU_to_CPU"] == [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
]
def test_aggregate_same_connector(self):
"""Test aggregating stats from the same connector type."""
stats1 = OffloadingConnectorStats(
data={
"CPU_to_GPU": [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
],
"GPU_to_CPU": [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
],
}
)
stats2 = OffloadingConnectorStats(
data={
"CPU_to_GPU": [
{"op_size": 3, "op_time": 0.2},
{"op_size": 7, "op_time": 0.9},
],
"GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
}
)
result = stats1.aggregate(stats2)
assert result is stats1 # Should return self
offload_connector_stats = result
assert offload_connector_stats.data["CPU_to_GPU"] == [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
{"op_size": 3, "op_time": 0.2},
{"op_size": 7, "op_time": 0.9},
]
assert offload_connector_stats.data["GPU_to_CPU"] == [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
{"op_size": 16, "op_time": 2},
]
def test_reduce(self):
"""Test that reduce() correctly reduces all nested connector stats."""
stats = OffloadingConnectorStats(
data={
"CPU_to_GPU": [
{"op_size": 16, "op_time": 1.0},
{"op_size": 8, "op_time": 0.5},
{"op_size": 3, "op_time": 0.2},
{"op_size": 7, "op_time": 0.9},
],
"GPU_to_CPU": [
{"op_size": 1, "op_time": 0.1},
{"op_size": 2, "op_time": 0.2},
{"op_size": 16, "op_time": 2},
],
}
)
reduced = stats.reduce()
assert isinstance(reduced, dict)
# Check that the stats were reduced (should have aggregated values)
assert "CPU_to_GPU_total_bytes" in reduced
assert "CPU_to_GPU_total_time" in reduced
assert "GPU_to_CPU_total_bytes" in reduced
assert "GPU_to_CPU_total_time" in reduced
assert reduced["CPU_to_GPU_total_bytes"] == 34
assert reduced["CPU_to_GPU_total_time"] == 2.6
assert reduced["GPU_to_CPU_total_time"] == 2.3
assert reduced["GPU_to_CPU_total_bytes"] == 19
def test_reset(self):
"""Test that reset() resets all nested connector stats."""
offload_connector_stats = OffloadingConnectorStats(
data={
"CPU_to_GPU": [
{"op_size": 3, "op_time": 0.2},
{"op_size": 7, "op_time": 0.9},
],
"GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
}
)
assert not offload_connector_stats.is_empty()
offload_connector_stats.reset()
# After reset, stats should be empty
assert offload_connector_stats.is_empty()
assert len(offload_connector_stats.data) == 0
@@ -0,0 +1,756 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import contextlib
import time
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import torch
import zmq.asyncio
from vllm.config import set_current_vllm_config
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
KVConnectorRole,
MooncakeConnector,
MooncakeConnectorMetadata,
MooncakeXferMetadata,
MooncakeXferResponse,
MooncakeXferResponseStatus,
PullReqMeta,
SendBlockMeta,
)
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
MooncakeBootstrapServer,
)
from vllm.utils.network_utils import get_open_port
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
from vllm.v1.request import RequestStatus
from .utils import create_request, create_scheduler, create_vllm_config
class FakeMooncakeWrapper:
"""Mock Mooncake TransferEngine for unit testing environments."""
def __init__(self, *args, **kwargs):
pass
def initialize(self, local_hostname, metadata_server, protocol, device_name) -> int:
return 0
def get_rpc_port(self) -> int:
return 12345
def batch_transfer_sync_write(
self, target_hostname, buffers, peer_buffer_addresses, lengths
) -> int:
return 0
def batch_register_memory(self, buffer_addresses, capacities) -> int:
return 0
def test_basic_interface():
"""Unit test for basic MooncakeConnector interface functionality."""
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector", kv_role="kv_consumer"
)
scheduler = create_scheduler(vllm_config)
# 2 Full Blocks and 1 Half Block.
BLOCK_SIZE = vllm_config.cache_config.block_size
NUM_EXTERNAL_FULL_BLOCKS = 2
NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
request = create_request(
request_id=1,
block_size=BLOCK_SIZE,
num_tokens=NUM_TOKENS,
do_remote_prefill=True,
)
request_id = request.request_id
request.kv_transfer_params.update(
{
"transfer_id": request_id,
"remote_bootstrap_addr": 54321,
}
)
scheduler.add_request(request)
# Remote Prefill, triggers NixlConnectorMetadata.
scheduler_output = scheduler.schedule()
kv_connector_metadata = scheduler_output.kv_connector_metadata
assert kv_connector_metadata is not None
assert isinstance(kv_connector_metadata, MooncakeConnectorMetadata)
assert len(kv_connector_metadata.reqs_to_recv) == 1
assert request_id in kv_connector_metadata.reqs_to_recv["my-engine-id"]
req_meta = kv_connector_metadata.reqs_to_recv["my-engine-id"][request_id]
for block_id, block in zip(
req_meta.local_block_ids,
scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
request_id
],
):
assert block_id == block.block_id
def test_prompt_less_than_block_size():
"""Test that we can handle case where prompt is < block."""
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector", kv_role="kv_consumer"
)
scheduler = create_scheduler(vllm_config)
# Half of a block.
BLOCK_SIZE = vllm_config.cache_config.block_size
NUM_TOKENS = int(BLOCK_SIZE * 0.5)
# Request will have 1 partial remote block.
request = create_request(
request_id=1,
block_size=BLOCK_SIZE,
num_tokens=NUM_TOKENS,
do_remote_prefill=True,
num_remote_blocks=1,
)
request.kv_transfer_params.update(
{
"transfer_id": request.request_id,
"remote_bootstrap_addr": 54321,
}
)
scheduler.add_request(request)
scheduler_output = scheduler.schedule()
# This request will read async.
kv_connector_metadata = scheduler_output.kv_connector_metadata
assert kv_connector_metadata is not None
assert isinstance(kv_connector_metadata, MooncakeConnectorMetadata)
assert len(kv_connector_metadata.reqs_to_recv["my-engine-id"]) == 1
assert len(scheduler_output.scheduled_new_reqs) == 0
@pytest.fixture
def bootstrap_server():
"""Fixture to launch and cleanup a Mooncake Bootstrap HTTP Server."""
port = get_open_port()
server = MooncakeBootstrapServer("127.0.0.1", port)
server.start()
yield server
server.shutdown()
@pytest.mark.asyncio
async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
"""
Tests the bootstrap server's api for worker registration and querying.
Validates DP/TP/PP rank indexing and error handling for duplicate registrations.
"""
import httpx
base_url = f"http://127.0.0.1:{bootstrap_server.port}"
# Query when empty
async with httpx.AsyncClient() as client:
response = await client.get(f"{base_url}/query")
assert response.status_code == 200
assert response.json() == {}
# Register a worker
payload1 = {
"engine_id": "eng-1",
"dp_rank": 0,
"tp_rank": 0,
"pp_rank": 0,
"addr": "tcp://1.1.1.1:1111",
}
async with httpx.AsyncClient() as client:
response = await client.post(f"{base_url}/register", json=payload1)
assert response.status_code == 200
assert response.json() == {"status": "ok"}
# Query after registration
async with httpx.AsyncClient() as client:
response = await client.get(f"{base_url}/query")
assert response.status_code == 200
data = response.json()
assert "0" in data
assert data["0"]["engine_id"] == "eng-1"
assert data["0"]["worker_addr"]["0"]["0"] == "tcp://1.1.1.1:1111"
# Test failure: re-registering the same worker
async with httpx.AsyncClient() as client:
response = await client.post(f"{base_url}/register", json=payload1)
assert response.status_code == 400
assert "is already registered" in response.text
# Test failure: engine_id mismatch for same dp_rank
payload3_fail = {
"engine_id": "eng-2",
"dp_rank": 0,
"tp_rank": 1,
"pp_rank": 0,
"addr": "tcp://3.3.3.3:3333",
}
async with httpx.AsyncClient() as client:
response = await client.post(f"{base_url}/register", json=payload3_fail)
assert response.status_code == 400
assert "Engine ID mismatch" in response.text
def test_scheduler_request_finished():
"""
Tests the scheduler-side logic when a request finishes.
Differentiates between 'Finished' (requires transfer)
and 'Aborted' (immediate free).
"""
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector", kv_role="kv_producer"
)
scheduler = create_scheduler(vllm_config)
scheduler_connector = scheduler.get_kv_connector().connector_scheduler
request = create_request(request_id=1, do_remote_decode=True)
request.kv_transfer_params["transfer_id"] = request.request_id
# Case: Capped length (Successful prefill, need to send to decoder)
request.status = RequestStatus.FINISHED_LENGTH_CAPPED
delay_free, _ = scheduler_connector.request_finished(request, block_ids=[10, 11])
assert delay_free is True
assert "id-1" in scheduler_connector._reqs_need_send
assert scheduler_connector._reqs_need_send["id-1"][1] == [10, 11]
# Case: Aborted (No need to transfer, free blocks immediately)
scheduler_connector._reqs_need_send.clear()
request.status = RequestStatus.FINISHED_ABORTED
delay_free, _ = scheduler_connector.request_finished(request, block_ids=[12])
assert delay_free is False
assert len(scheduler_connector._reqs_need_send) == 0
assert "id-1" in scheduler_connector._reqs_not_processed
@contextlib.contextmanager
def patch_worker_dependencies():
"""Helper to mock all distributed and network dependencies for Worker tests."""
with (
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
FakeMooncakeWrapper,
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_ip",
return_value="127.0.0.1",
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_tensor_model_parallel_rank",
return_value=0,
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_tensor_model_parallel_world_size",
return_value=1,
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_pp_group"
) as mock_pp,
patch("vllm.distributed.parallel_state.is_local_first_rank", return_value=True),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.should_launch_bootstrap_server",
return_value=False,
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.make_zmq_socket"
) as mock_make_zmq,
patch("httpx.AsyncClient") as mock_async_client,
):
# Mock PP group
mock_pp_group = MagicMock()
mock_pp_group.rank_in_group = 0
mock_pp.return_value = mock_pp_group
# Mock ZMQ socket
mock_socket_object = AsyncMock()
mock_socket_object.setsockopt = MagicMock()
mock_socket_ctx = MagicMock()
mock_socket_ctx.__enter__.return_value = mock_socket_object
mock_make_zmq.return_value = mock_socket_ctx
# Mock httpx client
mock_http_client_instance = AsyncMock()
mock_async_client.return_value = mock_http_client_instance
yield {
"mock_make_zmq": mock_make_zmq,
"mock_socket_object": mock_socket_object,
"mock_async_client": mock_async_client,
"mock_http_client": mock_http_client_instance,
}
@pytest.mark.asyncio
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
FakeMooncakeWrapper,
)
async def test_kv_producer(monkeypatch):
"""
Simulates a Producer Worker (Prefiller) receiving a transfer request
from a Consumer (Decoder).
Verifies memory offset calculation: ptr = base_addr + block_id * block_len.
"""
monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector", kv_role="kv_producer"
)
with set_current_vllm_config(vllm_config), patch_worker_dependencies():
prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
prefill_worker = prefill_connector.connector_worker
prefill_worker.kv_caches_base_addr = [0x1000]
block_len = 4096
prefill_worker.block_len_per_layer = [block_len]
# Override loop to use current test loop
origin_sender_loop = prefill_worker.sender_loop
prefill_worker.sender_loop = asyncio.get_event_loop()
# A request is finished on Producer and ready to be sent.
transfer_id = "xfer-req-1"
send_meta = SendBlockMeta(
p_req_id="p-req-1",
transfer_id=transfer_id,
local_block_ids=[10, 11],
ready=asyncio.Event(),
)
prefill_worker.reqs_need_send[transfer_id] = send_meta
send_meta.ready.set()
# Remote consumer request metadata
xfer_meta = MooncakeXferMetadata(
remote_hostname="consumer-host",
remote_port=54321,
remote_tp_size=1,
remote_tp_rank=0,
req_blocks={"d-req-1": (transfer_id, [20, 21])},
kv_caches_base_addr=[0x2000],
block_lens=[block_len],
)
mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
mock_socket.send_multipart = AsyncMock()
identity = b"consumer-id"
with patch.object(
prefill_worker, "_send_blocks", return_value=0
) as mock_send_blocks:
# Normal case: 2 blocks to 2 blocks
# Worker processes the consumer's request
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
# Verify transfer parameters are correct
src_ptr = 0x1000 + 10 * block_len
dst_ptr = 0x2000 + 20 * block_len
length = 2 * block_len
mock_send_blocks.assert_called_once_with(
"consumer-host:54321", [src_ptr], [dst_ptr], [length]
)
mock_socket.send_multipart.assert_called_once()
# Verify the response sent back to the consumer
sent_call = mock_socket.send_multipart.call_args[0][0]
sent_identity, sent_payload = sent_call
assert sent_identity == identity
response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
assert response.status == MooncakeXferResponseStatus.FINISH
assert response.ok_reqs == ["d-req-1"]
# Verify internal state cleanup
assert transfer_id not in prefill_worker.reqs_need_send
assert "p-req-1" in prefill_worker.finished_sending_reqs
# More cases:
# Consumer only needs 1 block (less than P)
mock_send_blocks.reset_mock()
mock_socket.send_multipart.reset_mock()
prefill_worker.reqs_need_send[transfer_id] = send_meta
send_meta.sent = 0
send_meta.ready.set()
xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20])
# Worker processes the consumer's request
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
# Verify transfer parameters are correct: 11 to 20
src_ptr = 0x1000 + 11 * block_len
dst_ptr = 0x2000 + 20 * block_len
length = 1 * block_len
mock_send_blocks.assert_called_once_with(
"consumer-host:54321", [src_ptr], [dst_ptr], [length]
)
mock_socket.send_multipart.assert_called_once()
# Consumer needs 3 blocks (more than P, error case)
mock_send_blocks.reset_mock()
mock_socket.send_multipart.reset_mock()
prefill_worker.reqs_need_send[transfer_id] = send_meta
send_meta.sent = 0
send_meta.ready.set()
xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21, 22])
# Worker processes the consumer's request
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
# This should not be called because error.
mock_send_blocks.assert_not_called()
mock_socket.send_multipart.assert_called_once()
_, sent_payload = mock_socket.send_multipart.call_args[0][0]
response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
assert response.err_msg == "P num blocks less than D"
assert response.err_reqs == ["d-req-1"]
# Timeout
mock_send_blocks.reset_mock()
mock_socket.send_multipart.reset_mock()
prefill_worker.reqs_need_send[transfer_id] = send_meta
send_meta.sent = 0
send_meta.ready.clear()
xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21])
# Worker processes the consumer's request
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
# This should not be called because timeout.
mock_send_blocks.assert_not_called()
mock_socket.send_multipart.assert_called_once()
_, sent_payload = mock_socket.send_multipart.call_args[0][0]
response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
assert response.err_msg == "Timeout waiting for P side ready."
assert response.err_reqs == ["d-req-1"]
# Transfer error
with patch.object(
prefill_worker, "_send_blocks", return_value=123
) as mock_send_blocks:
mock_socket.send_multipart.reset_mock()
prefill_worker.reqs_need_send[transfer_id] = send_meta
send_meta.sent = 0
send_meta.ready.set()
xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21])
# Worker processes the consumer's request
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
mock_send_blocks.assert_called_once()
mock_socket.send_multipart.assert_called_once()
_, sent_payload = mock_socket.send_multipart.call_args[0][0]
response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
assert response.err_msg == "Mooncake transfer engine returned 123"
assert response.err_reqs == ["d-req-1"]
# Clean up
prefill_worker.sender_loop = origin_sender_loop
prefill_worker.shutdown()
@pytest.mark.asyncio
async def test_kv_consumuer(monkeypatch):
"""
Simulates a Consumer Worker (Decoder) initiating a pull from a Producer.
Verifies that MooncakeXferMetadata is correctly serialized and sent via ZMQ.
"""
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector", kv_role="kv_consumer"
)
with set_current_vllm_config(vllm_config), patch_worker_dependencies() as mocks:
decode_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
decode_worker = decode_connector.connector_worker
decode_worker.kv_caches_base_addr = [0x1000]
decode_worker.rpc_port = 54321
# A request to pull data arrives.
pull_metas = {
"d-req-1": PullReqMeta(
d_req_id="d-req-1",
transfer_id="xfer-req-1",
local_block_ids=[100, 101],
remote_engine_id="p-engine",
remote_bootstrap_addr="http://bootstrap:33333",
pull_tasks_count=1,
)
}
decode_worker._remote_agents = {"p-engine": {0: {0: "tcp://producer:1234"}}}
decode_worker._tp_size["p-engine"] = 1
# Mock the response from the producer.
mock_response = MooncakeXferResponse(
status=MooncakeXferResponseStatus.FINISH, ok_reqs=["d-req-1"]
)
encoded_response = decode_worker._encoder.encode(mock_response)
mocks["mock_socket_object"].recv.return_value = encoded_response
# Trigger the receive logic.
decode_worker.receive_kv("p-engine", pull_metas)
await asyncio.sleep(1) # Allow async task to run
# Verify the metadata sent to the producer.
mocks["mock_make_zmq"].assert_called_with(
decode_worker.async_zmq_ctx,
"tcp://producer:1234",
zmq.DEALER,
bind=False,
linger=0,
)
sent_payload = mocks["mock_socket_object"].send.call_args[0][0]
sent_meta = decode_worker._xfer_meta_decoder.decode(sent_payload)
assert sent_meta.remote_hostname == "127.0.0.1"
assert sent_meta.remote_port == 54321
assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [100, 101])
# Verify internal state is updated correctly.
assert "d-req-1" in decode_worker.finished_recving_reqs
# Clean up
decode_worker.shutdown()
@pytest.mark.asyncio
async def test_worker_get_finished_timeout(monkeypatch):
"""Tests the cleanup mechanism for requests."""
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector", kv_role="kv_producer"
)
with set_current_vllm_config(vllm_config), patch_worker_dependencies():
prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
prefill_worker = prefill_connector.connector_worker
# Add an expired request (expire_time is in the past).
prefill_worker.reqs_need_send["tx-expired"] = SendBlockMeta(
p_req_id="p-req-expired",
transfer_id="tx-expired",
local_block_ids=[1, 2],
ready=MagicMock(),
expire_time=time.perf_counter() - 100,
)
# Add a non-expired request.
prefill_worker.reqs_need_send["tx-active"] = SendBlockMeta(
p_req_id="p-req-active",
transfer_id="tx-active",
local_block_ids=[3, 4],
ready=MagicMock(),
expire_time=time.perf_counter() + 100,
)
finished_reqs = await prefill_worker.fetch_finished_sending_reqs()
assert "p-req-expired" in finished_reqs
assert "p-req-active" not in finished_reqs
assert "tx-expired" not in prefill_worker.reqs_need_send
assert "tx-active" in prefill_worker.reqs_need_send
def test_register_kv_caches():
"""Tests the memory registration logic with the underlying Mooncake engine."""
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector", kv_role="kv_consumer"
)
with (
set_current_vllm_config(vllm_config),
patch_worker_dependencies(),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event"
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread"
) as mock_thread,
):
connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
worker = connector.connector_worker
mock_thread.return_value.is_alive.return_value = False
kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
)
tensor1 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
tensor2 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
kv_caches = {"layer0": tensor1, "layer1": tensor2}
with patch.object(
worker.engine, "batch_register_memory", return_value=0
) as mock_batch_register:
connector.register_kv_caches(kv_caches)
mock_batch_register.assert_called_once()
registered_ptrs, registered_lens = mock_batch_register.call_args[0]
expected_ptrs = {
tensor.data_ptr()
for kv_pair in kv_caches.values()
for tensor in kv_pair
}
assert set(registered_ptrs) == expected_ptrs
assert set(registered_lens) == {tensor1[0].nbytes}
# Verify block_len_per_layer is set correctly.
assert len(worker.block_len_per_layer) == len(registered_ptrs)
for bl in worker.block_len_per_layer:
assert bl == tensor1[0].nbytes // tensor1.shape[1]
@pytest.mark.asyncio
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
"mooncake_connector.TransferEngine",
FakeMooncakeWrapper,
)
@pytest.mark.parametrize("d_tp_size", [1, 4], ids=["p_tp2_d_tp1", "p_tp2_d_tp4"])
async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
"""
Tests heterogeneous TP support in the producer transfer path.
Verifies correct pointer and offset calculation when producer TP=2
sends to consumer with TP=1 (P>D) or TP=4 (P<D).
Parametrized cases:
- P TP=2 > D TP=1: one D rank receives; dst_offset based on P rank
- P TP=2 < D TP=4: two D ranks receive; src_offset based on D rank
"""
P_TP_SIZE = 2
P_TP_RANK = 0
LOCAL_BLOCK_LEN = 4096
local_block_len = LOCAL_BLOCK_LEN
remote_block_len = LOCAL_BLOCK_LEN * P_TP_SIZE // d_tp_size
monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector", kv_role="kv_producer"
)
with set_current_vllm_config(vllm_config), patch_worker_dependencies():
prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
prefill_worker = prefill_connector.connector_worker
# Override TP rank/size to simulate P TP=2
prefill_worker.tp_rank = P_TP_RANK
prefill_worker.tp_size = P_TP_SIZE
# Update shared dict so kv_topo sees correct TP size
prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE
prefill_worker.kv_topo.tp_rank = P_TP_RANK
prefill_worker.kv_caches_base_addr = [0x1000]
prefill_worker.block_len_per_layer = [local_block_len]
origin_sender_loop = prefill_worker.sender_loop
prefill_worker.sender_loop = asyncio.get_event_loop()
transfer_id = "xfer-hetero-1"
local_block_ids = [10, 11]
send_meta = SendBlockMeta(
p_req_id="p-req-h1",
transfer_id=transfer_id,
local_block_ids=local_block_ids,
ready=asyncio.Event(),
)
prefill_worker.reqs_need_send[transfer_id] = send_meta
send_meta.ready.set()
# Compute target D ranks using the production code path
target_d_ranks = prefill_worker.kv_topo.get_target_remote_ranks(d_tp_size)
mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
mock_socket.send_multipart = AsyncMock()
identity = b"consumer-hetero"
# Assign different remote block IDs per D rank
d_rank_remote_blocks = {
rank: [20 + i * 10, 21 + i * 10] for i, rank in enumerate(target_d_ranks)
}
with patch.object(
prefill_worker, "_send_blocks", return_value=0
) as mock_send_blocks:
for d_rank in target_d_ranks:
remote_block_ids = d_rank_remote_blocks[d_rank]
xfer_meta = MooncakeXferMetadata(
remote_hostname="consumer-host",
remote_port=54321,
remote_tp_size=d_tp_size,
remote_tp_rank=d_rank,
req_blocks={
f"d-req-h1-r{d_rank}": (
transfer_id,
remote_block_ids,
)
},
kv_caches_base_addr=[0x2000],
block_lens=[remote_block_len],
)
mock_send_blocks.reset_mock()
mock_socket.send_multipart.reset_mock()
await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
# Verify _send_blocks was called
mock_send_blocks.assert_called_once()
call_args = mock_send_blocks.call_args[0]
src_ptrs = call_args[1]
dst_ptrs = call_args[2]
lengths = call_args[3]
# Heterogeneous TP: blocks cannot be coalesced because
# local and remote block_lens differ
assert len(src_ptrs) == len(local_block_ids)
assert len(dst_ptrs) == len(local_block_ids)
assert len(lengths) == len(local_block_ids)
# Compute expected offsets based on TP ratio
if d_tp_size <= P_TP_SIZE:
tp_ratio = P_TP_SIZE // d_tp_size
expected_src_off = 0
expected_dst_off = (P_TP_RANK % tp_ratio) * local_block_len
expected_xfer_len = local_block_len
else:
ratio_abs = d_tp_size // P_TP_SIZE
expected_src_off = (d_rank % ratio_abs) * remote_block_len
expected_dst_off = 0
expected_xfer_len = remote_block_len
for idx, (lblk, rblk) in enumerate(
zip(local_block_ids, remote_block_ids)
):
assert src_ptrs[idx] == (
0x1000 + lblk * local_block_len + expected_src_off
)
assert dst_ptrs[idx] == (
0x2000 + rblk * remote_block_len + expected_dst_off
)
assert lengths[idx] == expected_xfer_len
# Verify successful response sent back to consumer
mock_socket.send_multipart.assert_called_once()
_, sent_payload = mock_socket.send_multipart.call_args[0][0]
response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
assert response.status == MooncakeXferResponseStatus.FINISH
assert response.ok_reqs == [f"d-req-h1-r{d_rank}"]
# After serving all D ranks, the request should be complete
assert transfer_id not in prefill_worker.reqs_need_send
assert "p-req-h1" in prefill_worker.finished_sending_reqs
prefill_worker.sender_loop = origin_sender_loop
prefill_worker.shutdown()
@@ -91,6 +91,9 @@ def clear_kv_transfer():
yield
if has_kv_transfer_group():
ensure_kv_transfer_shutdown()
# Reset any KV cache layout override set during tests so it doesn't
# leak into tests in other modules.
set_kv_cache_layout(None)
def get_default_xfer_telemetry(
+4 -2
View File
@@ -100,6 +100,8 @@ def create_vllm_config(
hf_overrides: dict[str, Any] | None = None,
attention_backend: str | None = None,
kv_load_failure_policy: Literal["recompute", "fail"] = "fail",
kv_connector: str = "NixlConnector",
kv_role: str = "kv_both",
) -> VllmConfig:
"""Initialize VllmConfig For Testing."""
model_config = ModelConfig(
@@ -124,8 +126,8 @@ def create_vllm_config(
enable_prefix_caching=True,
)
kv_transfer_config = KVTransferConfig(
kv_connector="NixlConnector",
kv_role="kv_both",
kv_connector=kv_connector,
kv_role=kv_role,
enable_permute_local_kv=enable_permute_local_kv,
kv_connector_extra_config=kv_connector_extra_config or {},
kv_load_failure_policy=kv_load_failure_policy,
+87 -126
View File
@@ -6,32 +6,20 @@ import time
import pytest
import torch
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
from vllm.v1.kv_offload.spec import (
CanonicalKVCacheRef,
CanonicalKVCaches,
CanonicalKVCacheTensor,
)
from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
BACKENDS_TO_TEST = [FlashAttentionBackend]
if not current_platform.is_rocm():
from vllm.v1.attention.backends.flashinfer import FlashInferBackend
BACKENDS_TO_TEST.append(FlashInferBackend)
from vllm.v1.attention.backends.mla.flashattn_mla import FlashAttnMLABackend
BACKENDS_TO_TEST.append(FlashAttnMLABackend)
NUM_GPU_BLOCKS = [64]
NUM_CPU_BLOCKS = [256]
KERNEL_BLOCK_SIZES = [16]
LOGICAL_BLOCK_SIZES = [16, 32]
LOGICAL_BLOCKS_PER_CPU_BLOCK = [1, 3]
HEAD_SIZES = [64]
NUM_HEADS = [8]
NUM_LAYERS = [4]
DTYPES = [torch.bfloat16]
GPU_PAGE_SIZES = [512, 1024]
BLOCK_SIZE_FACTORS = [1, 3]
NUM_TENSORS = [4]
SEEDS = [0]
CUDA_DEVICES = ["cuda:0"]
NUM_MAPPINGS = [3]
@@ -39,15 +27,11 @@ NUM_MAPPINGS = [3]
@pytest.mark.parametrize("gpu_to_cpu", [True, False])
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("kernel_block_size", KERNEL_BLOCK_SIZES)
@pytest.mark.parametrize("logical_block_size", LOGICAL_BLOCK_SIZES)
@pytest.mark.parametrize("logical_blocks_per_cpu_block", LOGICAL_BLOCKS_PER_CPU_BLOCK)
@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS)
@pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
@pytest.mark.parametrize("num_layers", NUM_LAYERS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("num_tensors", NUM_TENSORS)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
@@ -55,113 +39,89 @@ def test_transfer(
default_vllm_config,
gpu_to_cpu: bool,
num_mappings: int,
head_size: int,
num_heads: int,
kernel_block_size: int,
logical_block_size: int,
logical_blocks_per_cpu_block: int,
gpu_page_size_bytes: int,
block_size_factor: int,
num_gpu_blocks: int,
num_cpu_blocks: int,
num_layers: int,
dtype: torch.dtype,
num_tensors: int,
seed: int,
device: str,
) -> None:
set_random_seed(seed)
# create per-layer GPU KV caches based on available attn_backends
attn_backends_list = BACKENDS_TO_TEST
assert logical_block_size % kernel_block_size == 0
kernel_blocks_per_gpu_block = logical_block_size // kernel_block_size
num_gpu_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
gpu_caches = {}
attn_backends = {}
for i in range(num_layers):
layer_name = f"layer {i}"
attn_backend = attn_backends_list[i % len(attn_backends_list)]
attn_backends[layer_name] = attn_backend
gpu_cache_shape = attn_backend.get_kv_cache_shape(
num_gpu_kernel_blocks, kernel_block_size, num_heads, head_size
# build CanonicalKVCacheTensor list: one per tensor
kv_cache_tensors: list[CanonicalKVCacheTensor] = []
for i in range(num_tensors):
gpu_tensor = torch.randint(
-128,
127,
(num_gpu_blocks, gpu_page_size_bytes),
dtype=torch.int8,
device=device,
)
kv_cache_tensors.append(
CanonicalKVCacheTensor(
tensor=gpu_tensor,
page_size_bytes=gpu_page_size_bytes,
)
)
gpu_caches[layer_name] = torch.rand(gpu_cache_shape, dtype=dtype, device=device)
# create handler
cpu_block_size = logical_blocks_per_cpu_block * logical_block_size
kernel_blocks_per_cpu_block = cpu_block_size // kernel_block_size
# one group containing all tensors, one data ref per tensor
kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]] = [
[
CanonicalKVCacheRef(
tensor_idx=i,
page_size_bytes=gpu_page_size_bytes,
)
for i in range(num_tensors)
]
]
kv_caches = CanonicalKVCaches(
tensors=kv_cache_tensors,
group_data_refs=kv_cache_groups_data_refs,
)
handlers = CpuGpuOffloadingHandlers(
attn_backends=attn_backends,
gpu_block_size=logical_block_size,
cpu_block_size=cpu_block_size,
kv_caches=kv_caches,
block_size_factor=block_size_factor,
num_cpu_blocks=num_cpu_blocks,
gpu_caches=gpu_caches,
)
# select block mappings
gpu_blocks = random.sample(
range(num_gpu_blocks), num_mappings * logical_blocks_per_cpu_block
)
gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * block_size_factor)
cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings)
# convert gpu blocks to kernel block size
gpu_blocks_in_kernel_block_size = []
for gpu_block in gpu_blocks:
base_block_id = gpu_block * kernel_blocks_per_gpu_block
for i in range(kernel_blocks_per_gpu_block):
gpu_blocks_in_kernel_block_size.append(i + base_block_id)
# expand cpu blocks to gpu-page granularity for uniform comparison:
# each cpu block maps to block_size_factor consecutive sub-blocks
cpu_blocks_expanded = [
cpu_block * block_size_factor + j
for cpu_block in cpu_blocks
for j in range(block_size_factor)
]
# convert cpu blocks to gpu block size
cpu_blocks_in_kernel_block_size = []
for cpu_block in cpu_blocks:
base_block_id = cpu_block * kernel_blocks_per_cpu_block
for i in range(kernel_blocks_per_cpu_block):
cpu_blocks_in_kernel_block_size.append(i + base_block_id)
# maybe skip some GPU block to test reading from the middle of a CPU block
# maybe skip some GPU blocks to test reading from the middle of a CPU block
if not gpu_to_cpu:
gpu_blocks_to_skip = logical_blocks_per_cpu_block - 1
gpu_blocks = gpu_blocks[gpu_blocks_to_skip:]
kernel_blocks_to_skip = gpu_blocks_to_skip * kernel_blocks_per_gpu_block
gpu_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size[
kernel_blocks_to_skip:
]
cpu_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size[
kernel_blocks_to_skip:
]
blocks_to_skip = block_size_factor - 1
gpu_blocks = gpu_blocks[blocks_to_skip:]
cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:]
# set transfer direction
if gpu_to_cpu:
handler = handlers.gpu_to_cpu_handler
src_blocks = gpu_blocks
dst_blocks = cpu_blocks
src_spec = GPULoadStoreSpec(src_blocks, group_sizes=(len(src_blocks),))
dst_spec = CPULoadStoreSpec(dst_blocks)
src_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
dst_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
dst_size_in_kernel_blocks = num_cpu_blocks * kernel_blocks_per_cpu_block
src_spec = GPULoadStoreSpec(gpu_blocks, group_sizes=(len(gpu_blocks),))
dst_spec = CPULoadStoreSpec(cpu_blocks)
dst_to_src = dict(zip(cpu_blocks_expanded, gpu_blocks))
num_dst_sub_blocks = num_cpu_blocks * block_size_factor
else:
handler = handlers.cpu_to_gpu_handler
src_blocks = cpu_blocks
dst_blocks = gpu_blocks
src_spec = CPULoadStoreSpec(src_blocks)
dst_spec = GPULoadStoreSpec(dst_blocks, group_sizes=(len(dst_blocks),))
src_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
dst_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
dst_size_in_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
# build dst -> src mapping
dst_to_src = {}
for src_block, dst_block in zip(
src_blocks_in_kernel_block_size, dst_blocks_in_kernel_block_size
):
dst_to_src[dst_block] = src_block
src_spec = CPULoadStoreSpec(cpu_blocks)
dst_spec = GPULoadStoreSpec(gpu_blocks, group_sizes=(len(gpu_blocks),))
dst_to_src = dict(zip(gpu_blocks, cpu_blocks_expanded))
num_dst_sub_blocks = num_gpu_blocks
# clone src and dst tensors before transfer
orig_src_caches = [x.clone() for x in handler.src_tensors]
orig_dst_caches = [x.clone() for x in handler.dst_tensors]
orig_src_tensors = [x.clone() for x in handler.src_tensors]
orig_dst_tensors = [x.clone() for x in handler.dst_tensors]
# call transfer function
start_time = time.time()
@@ -180,11 +140,8 @@ def test_transfer(
if gpu_to_cpu
else ("CPU", "GPU")
)
assert (
finished[0].transfer_size
== handler.total_block_size_in_bytes
* handler.dst_block_size_factor
* len(dst_blocks)
assert finished[0].transfer_size == (
len(gpu_blocks) * handler.group_block_size_in_bytes[0]
)
assert finished[0].transfer_time > 0
assert finished[0].transfer_time < (time.time() - start_time)
@@ -192,19 +149,23 @@ def test_transfer(
time.sleep(0.1)
# verify src tensors did not change
for orig_tensor, tensor in zip(orig_src_caches, handler.src_tensors):
for orig_tensor, tensor in zip(orig_src_tensors, handler.src_tensors):
assert torch.equal(orig_tensor, tensor)
# verify dst tensors
for dst_block in range(dst_size_in_kernel_blocks):
src_block_candidate = dst_to_src.get(dst_block)
for src_cache, dst_cache, orig_dst_cache in zip(
handler.src_tensors,
handler.dst_tensors,
orig_dst_caches,
):
if src_block_candidate is not None:
expected_value = src_cache[src_block_candidate]
# verify dst tensors at gpu-page granularity.
for src_tensor, dst_tensor, orig_dst_tensor in zip(
handler.src_tensors,
handler.dst_tensors,
orig_dst_tensors,
):
# view both GPU and CPU tensors as (n, gpu_page_size_bytes) for comparison.
src_view = src_tensor.view(-1, gpu_page_size_bytes)
dst_view = dst_tensor.view(-1, gpu_page_size_bytes)
orig_dst_view = orig_dst_tensor.view(-1, gpu_page_size_bytes)
for dst_sub_block in range(num_dst_sub_blocks):
src_sub_block = dst_to_src.get(dst_sub_block)
if src_sub_block is not None:
expected = src_view[src_sub_block]
else:
expected_value = orig_dst_cache[dst_block]
torch.testing.assert_close(dst_cache[dst_block].cpu(), expected_value.cpu())
expected = orig_dst_view[dst_sub_block]
torch.testing.assert_close(dst_view[dst_sub_block].cpu(), expected.cpu())
@@ -103,8 +103,8 @@ class LogitsProcsRequestParams:
class MockReasoningConfig:
"""Mock reasoning config for testing ThinkingTokenBudgetLogitsProcessor."""
think_start_token_ids = [THINK_START_TOKEN_ID]
think_end_token_ids = [THINK_END_TOKEN_ID]
reasoning_start_token_ids = [THINK_START_TOKEN_ID]
reasoning_end_token_ids = [THINK_END_TOKEN_ID]
def _generate_fake_sampling_metadata(
@@ -491,7 +491,7 @@ def _thinking_budget_validate(
# Find if thinking has started in output tokens
thinking_started = False
start_tokens = tb_processor.think_start_token_ids
start_tokens = tb_processor.reasoning_start_token_ids
if len(start_tokens) > 0:
for i in range(len(output_tokens) - len(start_tokens) + 1):
@@ -518,7 +518,7 @@ def _thinking_budget_validate(
)
# Validate that only end tokens are allowed
end_tokens = tb_processor.think_end_token_ids
end_tokens = tb_processor.reasoning_end_token_ids
if len(end_tokens) > 0:
expected_end_token_id = end_tokens[
min(state["end_count"], len(end_tokens) - 1)
@@ -0,0 +1,193 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Integration tests for SimpleCPUOffloadConnector with real models."""
import time
import pytest
from vllm import LLM, SamplingParams, TokensPrompt
from vllm.config import KVTransferConfig
from vllm.platforms import current_platform
if not current_platform.is_cuda():
pytest.skip("Requires CUDA", allow_module_level=True)
# Small models for default CI / local runs (accuracy only).
SMALL_MODELS = [
"meta-llama/Llama-3.2-1B-Instruct",
"google/gemma-3-1b-it",
]
# Large models for optional perf runs only (slow to load and execute).
PERF_MODELS = [
"meta-llama/Llama-3.1-8B",
"openai/gpt-oss-20b",
]
def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM:
kv_transfer_config = KVTransferConfig(
kv_connector="SimpleCPUOffloadConnector",
kv_role="kv_both",
kv_connector_extra_config={
"cpu_bytes_to_use": cpu_bytes_to_use,
"lazy_offload": lazy,
},
)
return LLM(
model=model,
kv_cache_memory_bytes=40 << 30, # 40 GiB
disable_hybrid_kv_cache_manager=False,
enable_prefix_caching=True,
kv_transfer_config=kv_transfer_config,
)
def _flush_gpu_cache(llm: LLM, sampling_params: SamplingParams, seed: int = 0):
"""Generate enough filler requests to allocate the entire GPU KV cache.
This pushes all prior blocks through the free queue so that the lazy
cursor offloads them to CPU before they are evicted.
"""
cache_config = llm.llm_engine.vllm_config.cache_config
num_gpu_blocks = cache_config.num_gpu_blocks
block_size = cache_config.block_size
# Use 1.2x GPU capacity to give the lazy cursor enough scheduling steps
# to walk past all target blocks near the tail of the free queue.
total_tokens_needed = int(num_gpu_blocks * block_size * 1.5)
# Use token-id prompts so each filler is unique (no prefix sharing).
# Split into multiple requests to stay under max_model_len.
max_tokens_per_req = 4096
num_fillers = (total_tokens_needed + max_tokens_per_req - 1) // max_tokens_per_req
batch_size = 10
for i in range(0, num_fillers, batch_size):
batch_end = min(i + batch_size, num_fillers)
filler_prompts = []
for j in range(i, batch_end):
ids = [seed * num_fillers + j + 1] * max_tokens_per_req
filler_prompts.append(TokensPrompt(prompt_token_ids=ids))
llm.generate(filler_prompts, sampling_params, use_tqdm=False)
def _accuracy_test(llm: LLM, lazy: bool = False):
"""Verify that CPU-loaded KV produces correct output."""
sampling_params = SamplingParams(max_tokens=1, temperature=0)
prompt = "hi " * 2000 + "Let's count to ten. One, two, three, "
# Cold run — populate GPU cache and trigger CPU offload
cold_output = llm.generate(prompt, sampling_params, use_tqdm=False)[0]
# CPU hit runs
test_count = 10
success_count = 0
expected = cold_output.outputs[0].text
for i in range(test_count):
if lazy:
_flush_gpu_cache(llm, sampling_params, seed=i)
time.sleep(2) # let engine core drain pending transfers
# Reset GPU prefix cache so next run must load from CPU
if not llm.reset_prefix_cache():
print(f"GPU prefix cache reset failed for iteration {i}")
output = llm.generate(prompt, sampling_params, use_tqdm=False)[0]
if output.outputs[0].text == expected:
success_count += 1
assert success_count >= 0.5 * test_count, (
f"Accuracy too low: {success_count}/{test_count} matched '{expected}'"
)
def _latency_test(llm: LLM, lazy: bool = False):
"""Verify CPU cache hit is faster than cold compute."""
sampling_params = SamplingParams(max_tokens=1, seed=42)
prompt_token_ids = [0] * 10001
num_times_cpu_better = 0
num_tests = 10
for i in range(num_tests):
prompt_token_ids[0] = i
prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
# Cold
time.sleep(2) # let engine core drain pending transfers
if not llm.reset_prefix_cache():
print(f"GPU prefix cache reset failed for iteration {i}")
start = time.time()
llm.generate(prompts, sampling_params, use_tqdm=False)
cold_time = time.time() - start
if lazy:
_flush_gpu_cache(llm, sampling_params, seed=i)
else:
# Eager mode: GPU hit ensures store completion is processed.
llm.generate(prompts, sampling_params, use_tqdm=False)
time.sleep(2) # let engine core drain pending transfers
if not llm.reset_prefix_cache():
print(f"GPU prefix cache reset failed for iteration {i}")
# CPU hit
start = time.time()
llm.generate(prompts, sampling_params, use_tqdm=False)
cpu_time = time.time() - start
if cpu_time < cold_time:
num_times_cpu_better += 1
assert num_times_cpu_better >= 0.8 * num_tests, (
f"CPU hit only faster {num_times_cpu_better}/{num_tests} times"
)
@pytest.mark.optional
@pytest.mark.slow_test
@pytest.mark.parametrize("model", SMALL_MODELS)
def test_simple_cpu_offload_accuracy(model: str):
"""Store to CPU, reset GPU, load from CPU; verify output matches baseline."""
llm = _make_llm(model, False, 1 << 30) # 1GB
try:
_accuracy_test(llm, lazy=False)
finally:
del llm
@pytest.mark.optional
@pytest.mark.slow_test
@pytest.mark.parametrize("model", PERF_MODELS)
def test_simple_cpu_offload_perf_latency(model: str):
"""CPU KV hit should beat cold prefill on long context (large models only)."""
llm = _make_llm(model, False, 10 << 30) # 10GB
try:
_latency_test(llm, lazy=False)
finally:
del llm
@pytest.mark.optional
@pytest.mark.slow_test
@pytest.mark.parametrize("model", SMALL_MODELS)
def test_simple_cpu_offload_accuracy_lazy(model: str):
"""Lazy mode: flush GPU cache to trigger CPU offload, then verify hit."""
# CPU must be larger than GPU KV cache to avoid evicting offloaded blocks.
llm = _make_llm(model, True, 80 << 30) # 80GB
try:
_accuracy_test(llm, lazy=True)
finally:
del llm
@pytest.mark.optional
@pytest.mark.slow_test
@pytest.mark.parametrize("model", PERF_MODELS)
def test_simple_cpu_offload_perf_latency_lazy(model: str):
"""Lazy mode: CPU KV hit should beat cold prefill (large models only)."""
# CPU must be larger than GPU KV cache to avoid evicting offloaded blocks.
llm = _make_llm(model, True, 80 << 30) # 80GB
try:
_latency_test(llm, lazy=True)
finally:
del llm
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,147 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression tests for the backup token fix in prepare_next_token_ids_padded.
Fixes #38098: with async scheduling, seq_lens_cpu is inflated by unaccepted
draft token placeholders, causing get_token_id() to return -1.
"""
from __future__ import annotations
import numpy as np
import pytest
import torch
class _FakeRequest:
def __init__(self, prompt_tokens: list[int], output_tokens: list[int]):
self.num_prompt_tokens = len(prompt_tokens)
self._prompt = prompt_tokens
self._output = output_tokens
@property
def num_tokens(self) -> int:
return self.num_prompt_tokens + len(self._output)
def get_token_id(self, idx: int) -> int:
if idx < self.num_prompt_tokens:
return self._prompt[idx]
out_idx = idx - self.num_prompt_tokens
if out_idx < len(self._output):
return self._output[out_idx]
return -1 # out of range
class _FakeInputBatch:
def __init__(
self,
req_ids: list[str],
num_tokens_no_spec: list[int],
vocab_size: int = 32000,
):
self.req_ids = req_ids
self.num_reqs = len(req_ids)
self.vocab_size = vocab_size
self.num_tokens_no_spec = np.array(num_tokens_no_spec, dtype=np.int64)
def _make_requests(
req_ids: list[str],
prompt_lens: list[int],
output_lens: list[int],
) -> dict[str, _FakeRequest]:
requests = {}
for rid, plen, olen in zip(req_ids, prompt_lens, output_lens):
requests[rid] = _FakeRequest(list(range(plen)), list(range(1000, 1000 + olen)))
return requests
def _backup_buggy(
seq_lens_cpu: torch.Tensor,
requests: dict[str, _FakeRequest],
batch: _FakeInputBatch,
) -> list[int]:
"""Old logic: uses seq_lens_cpu directly (may be inflated)."""
n = batch.num_reqs
return [
requests[batch.req_ids[i]].get_token_id(int(seq_lens_cpu[i])) for i in range(n)
]
def _backup_fixed(
requests: dict[str, _FakeRequest],
batch: _FakeInputBatch,
) -> list[int]:
"""New logic: uses num_tokens_no_spec - 1 (last committed token)."""
n = batch.num_reqs
idx = (batch.num_tokens_no_spec[:n] - 1).tolist()
return [requests[batch.req_ids[i]].get_token_id(int(idx[i])) for i in range(n)]
class TestBackupTokenAsyncSpec:
def test_no_inflation_fixed_returns_last_token(self):
req_ids = ["r0", "r1"]
requests = _make_requests(req_ids, [3, 3], [2, 2])
batch = _FakeInputBatch(req_ids, [5, 5])
# idx = 5-1 = 4 → output[1] = 1001
assert _backup_fixed(requests, batch) == [1001, 1001]
def test_inflation_buggy_returns_placeholder(self):
req_ids = ["r0", "r1"]
requests = _make_requests(req_ids, [3, 3], [2, 2])
batch = _FakeInputBatch(req_ids, [5, 5])
# inflated by 3 spec tokens → idx 8 is out of range
seq_lens = torch.tensor([8, 8], dtype=torch.int64)
assert _backup_buggy(seq_lens, requests, batch) == [-1, -1]
def test_inflation_fixed_returns_correct_token(self):
req_ids = ["r0", "r1"]
requests = _make_requests(req_ids, [3, 3], [2, 2])
batch = _FakeInputBatch(req_ids, [5, 5])
assert _backup_fixed(requests, batch) == [1001, 1001]
def test_mixed_inflation_per_request(self):
req_ids = ["r0", "r1", "r2"]
requests = {
"r0": _FakeRequest([0, 1], [1000, 1001, 1002]),
"r1": _FakeRequest([0, 1, 2, 3], [2000]),
"r2": _FakeRequest([0], [3000, 3001, 3002, 3003]),
}
batch = _FakeInputBatch(req_ids, [5, 5, 5])
seq_lens = torch.tensor([7, 9, 5], dtype=torch.int64)
assert _backup_buggy(seq_lens, requests, batch) == [-1, -1, -1]
assert _backup_fixed(requests, batch) == [1002, 2000, 3003]
def test_prefill_only_request(self):
"""No output tokens yet — backup should be the last prompt token."""
req_ids = ["r0"]
requests = {"r0": _FakeRequest([10, 20, 30], [])}
batch = _FakeInputBatch(req_ids, [3])
# idx = 3-1 = 2 → prompt[2] = 30
assert _backup_fixed(requests, batch) == [30]
@pytest.mark.parametrize("num_spec_tokens", [1, 2, 3, 4, 5])
def test_various_spec_token_counts(self, num_spec_tokens: int):
req_ids = ["r0"]
requests = {"r0": _FakeRequest([0, 1, 2], list(range(1000, 1005)))}
batch = _FakeInputBatch(req_ids, [8])
# idx = 8-1 = 7 → output[4] = 1004
assert _backup_fixed(requests, batch) == [1004]
def test_buggy_code_was_always_off_by_one(self):
"""The original code used seq_len as index, which is always one past
the end of output_token_ids even without async inflation."""
req_ids = ["r0"]
requests = {"r0": _FakeRequest([0, 1, 2], [1000, 1001])}
batch = _FakeInputBatch(req_ids, [5])
# no inflation: seq_len == num_tokens == 5 → idx 5 is out of range
seq_lens = torch.tensor([5], dtype=torch.int64)
assert _backup_buggy(seq_lens, requests, batch) == [-1]
assert _backup_fixed(requests, batch) == [1001]
# with inflation: still -1, fixed still correct
seq_lens_inf = torch.tensor([8], dtype=torch.int64)
assert _backup_buggy(seq_lens_inf, requests, batch) == [-1]
assert _backup_fixed(requests, batch) == [1001]
+4 -12
View File
@@ -3,6 +3,7 @@
from unittest import mock
import numpy as np
import pytest
import torch
@@ -111,16 +112,14 @@ def test_prepare_next_token_ids():
num_requests = 4
num_speculative_tokens = 4
batch_spec = BatchSpec(
seq_lens=[num_speculative_tokens + 1] * num_requests,
query_lens=[num_speculative_tokens + 1] * num_requests,
)
req_ids = [f"req_{i + 1}" for i in range(num_requests)]
mock_input_batch = mock.MagicMock(spec=InputBatch)
mock_input_batch.req_ids = req_ids
mock_input_batch.num_reqs = num_requests
mock_input_batch.vocab_size = 100
mock_input_batch.num_tokens_no_spec = np.array(
[num_speculative_tokens + 1] * num_requests
)
mock_num_scheduled_tokens = {req_id: 0 for req_id in req_ids}
mock_requests = {}
@@ -165,19 +164,12 @@ def test_prepare_next_token_ids():
assert torch.equal(next_token_ids_from_cpu, expected_next_token_ids_tensor)
common_attn_metadata = create_common_attn_metadata(
batch_spec,
block_size=BLOCK_SIZE,
device=device,
)
expected_valid_sampled_tokens_count = torch.tensor(
[2, 5, 0, 0], dtype=torch.int32, device=device
)
next_token_ids_from_padded, valid_sampled_tokens_count = (
proposer.prepare_next_token_ids_padded(
common_attn_metadata.seq_lens_cpu,
sampled_token_ids_tensor,
mock_requests,
mock_input_batch,
@@ -3,6 +3,7 @@
from unittest import mock
import numpy as np
import pytest
import torch
@@ -132,16 +133,12 @@ def test_prepare_next_token_ids_padded():
device = torch.device(current_platform.device_type)
num_requests = 4
batch_spec = BatchSpec(
seq_lens=[5] * num_requests,
query_lens=[5] * num_requests,
)
req_ids = [f"req_{i + 1}" for i in range(num_requests)]
mock_input_batch = mock.MagicMock(spec=InputBatch)
mock_input_batch.req_ids = req_ids
mock_input_batch.num_reqs = num_requests
mock_input_batch.vocab_size = 100
mock_input_batch.num_tokens_no_spec = np.array([5] * num_requests)
mock_requests = {}
for req_id in req_ids:
@@ -174,12 +171,6 @@ def test_prepare_next_token_ids_padded():
proposer = _create_proposer(num_speculative_tokens=1)
common_attn_metadata = create_common_attn_metadata(
batch_spec,
block_size=16,
device=device,
)
# valid_sampled_tokens_count tracks if token is valid (not -1 and in vocab range)
# It doesn't depend on whether the request is discarded
expected_valid_sampled_tokens_count = torch.tensor(
@@ -187,7 +178,6 @@ def test_prepare_next_token_ids_padded():
)
next_token_ids, valid_sampled_tokens_count = proposer.prepare_next_token_ids_padded(
common_attn_metadata.seq_lens_cpu,
sampled_token_ids,
mock_requests,
mock_input_batch,
+47 -2
View File
@@ -380,7 +380,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
_compare_objs(input_batch, ref_input_batch)
def _construct_pooling_request(req_id_suffix: int):
def _construct_pooling_request(req_id_suffix: int, pooling_params=None):
from vllm.pooling_params import PoolingParams
prompt_token_ids = [
@@ -391,7 +391,7 @@ def _construct_pooling_request(req_id_suffix: int):
req_id=f"pool_req_{req_id_suffix}",
prompt_token_ids=prompt_token_ids,
sampling_params=None,
pooling_params=PoolingParams(task="classify"),
pooling_params=pooling_params or PoolingParams(task="classify"),
mm_features=[],
block_ids=([],),
generator=None,
@@ -440,3 +440,48 @@ def test_pooling_prompt_lens_not_aliased(device: str):
"mutations to num_prompt_tokens_cpu_tensor corrupted prompt_lens. "
f"Expected {prompt_lens_snapshot}, got {metadata.prompt_lens}"
)
@pytest.mark.parametrize(
("pooling_params", "expect_device_prompt_token_ids", "expect_cpu_prompt_token_ids"),
[
({"task": "classify"}, False, False),
({"task": "classify", "requires_token_ids": True}, True, True),
],
)
def test_pooling_metadata_token_id_buffers(
pooling_params: dict[str, object],
expect_device_prompt_token_ids: bool,
expect_cpu_prompt_token_ids: bool,
):
from vllm.pooling_params import PoolingParams
input_batch = InputBatch(
max_num_reqs=1,
max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
device=torch.device("cpu"),
pin_memory=False,
vocab_size=VOCAB_SIZE,
block_sizes=[16],
kernel_block_sizes=[16],
is_pooling_model=True,
)
req = _construct_pooling_request(0, PoolingParams(**pooling_params))
input_batch.add_request(req)
input_batch.refresh_metadata()
metadata = input_batch.get_pooling_metadata()
if expect_device_prompt_token_ids:
assert input_batch.sampling_metadata.prompt_token_ids is not None
assert metadata.prompt_token_ids is not None
assert metadata.get_prompt_token_ids()[0].tolist() == req.prompt_token_ids
else:
assert input_batch.sampling_metadata.prompt_token_ids is None
assert metadata.prompt_token_ids is None
if expect_cpu_prompt_token_ids:
assert metadata.prompt_token_ids_cpu is not None
assert metadata.get_prompt_token_ids_cpu()[0].tolist() == req.prompt_token_ids
else:
assert metadata.prompt_token_ids_cpu is None
@@ -446,7 +446,7 @@ def parse_attention_types(node: ast.ClassDef) -> str:
if not types:
return "Decoder"
return "All" if len(types) >= 3 else ", ".join(sorted(types))
return "All" if types >= set(type_map.values()) else ", ".join(sorted(types))
def parse_impl_bool_attr(
@@ -17,14 +17,14 @@
#
# Environment variables:
# S3_BUCKET - Bucket name (default: vllm-wheels)
# VARIANT - ROCm variant (default: rocm700)
# VARIANT - ROCm variant (default: rocm721)
# DRY_RUN - Set to 1 for preview mode (same as --dry-run)
set -euo pipefail
# ======== Configuration ========
BUCKET="${S3_BUCKET:-vllm-wheels}"
VARIANT="${VARIANT:-rocm700}"
VARIANT="${VARIANT:-rocm721}"
DRY_RUN="${DRY_RUN:-0}"
FORCE_VERSION=""
+15 -13
View File
@@ -66,27 +66,29 @@ class CacheConfig:
enable_prefix_caching: bool = True
"""Whether to enable prefix caching."""
prefix_caching_hash_algo: PrefixCachingHashAlgo = "sha256"
"""Set the hash algorithm for prefix caching:\n
- "sha256" uses Pickle for object serialization before hashing. This is the
current default, as SHA256 is the most secure choice to avoid potential
hash collisions.\n
"""Set the hash algorithm for prefix caching:
- "sha256" uses Pickle for object serialization before hashing. This is the current
default, as SHA256 is the most secure choice to avoid potential hash collisions.
- "sha256_cbor" provides a reproducible, cross-language compatible hash. It
serializes objects using canonical CBOR and hashes them with SHA-256.\n
serializes objects using canonical CBOR and hashes them with SHA-256.
- "xxhash" uses Pickle serialization with xxHash (128-bit) for faster,
non-cryptographic hashing. Requires the optional ``xxhash`` package.
IMPORTANT: Use of a hashing algorithm that is not considered
cryptographically secure theoretically increases the risk of hash collisions,
which can cause undefined behavior or even leak private information in
multi-tenant environments. Even if collisions are still very unlikely, it is
important to consider your security risk tolerance against the performance
benefits before turning this on.\n
non-cryptographic hashing. Requires the optional ``xxhash`` package.
IMPORTANT: Use of a hashing algorithm that is not considered cryptographically
secure theoretically increases the risk of hash collisions, which can cause
undefined behavior or even leak private information in multi-tenant environments.
Even if collisions are still very unlikely, it is important to consider your
security risk tolerance against the performance benefits before turning this on.
- "xxhash_cbor" combines canonical CBOR serialization with xxHash for
reproducible hashing. Requires the optional ``xxhash`` package."""
reproducible hashing. Requires the optional ``xxhash`` package."""
calculate_kv_scales: bool = False
"""Deprecated: This option is deprecated and will be removed in v0.19.
It enables dynamic calculation of `k_scale` and `v_scale` when
kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model
checkpoint if available. Otherwise, the scales will default to 1.0."""
kv_cache_dtype_skip_layers: list[str] = field(default_factory=list)
"""Layer patterns to skip KV cache quantization. Accepts layer indices
(e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window')."""
cpu_kvcache_space_bytes: int | None = None
"""(CPU backend only) CPU key-value cache space."""
mamba_page_size_padded: int | None = None
+8 -8
View File
@@ -32,14 +32,14 @@ class KernelConfig:
moe_backend: MoEBackend = "auto"
"""Backend for MoE expert computation kernels. Available options:
- "auto": Automatically select the best backend based on model and hardware\n
- "triton": Use Triton-based fused MoE kernels\n
- "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)\n
- "cutlass": Use vLLM CUTLASS kernels\n
- "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels\n
- "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels\n
- "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)\n
- "marlin": Use Marlin kernels (weight-only quantization)\n
- "auto": Automatically select the best backend based on model and hardware
- "triton": Use Triton-based fused MoE kernels
- "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)
- "cutlass": Use vLLM CUTLASS kernels
- "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels
- "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
- "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
- "marlin": Use Marlin kernels (weight-only quantization)
- "aiter": Use AMD AITer kernels (ROCm only)"""
@field_validator("moe_backend", mode="before")
+1 -1
View File
@@ -51,7 +51,7 @@ class LoadConfig:
- "gguf" will load weights from GGUF format files (details specified in
https://github.com/ggml-org/ggml/blob/master/docs/gguf.md).
- "mistral" will load weights from consolidated safetensors files used by
Mistral models.\n
Mistral models.
- Other custom values can be supported via plugins.
"""
download_dir: str | None = None
+58 -35
View File
@@ -125,26 +125,28 @@ class ModelConfig:
"""Name or path of the Hugging Face tokenizer to use. If unspecified, model
name or path will be used."""
tokenizer_mode: TokenizerMode | str = "auto"
"""Tokenizer mode:\n
"""Tokenizer mode:
- "auto" will use the tokenizer from `mistral_common` for Mistral models
if available, otherwise it will use the "hf" tokenizer.\n
- "hf" will use the fast tokenizer if available.\n
- "slow" will always use the slow tokenizer.\n
- "mistral" will always use the tokenizer from `mistral_common`.\n
- "deepseek_v32" will always use the tokenizer from `deepseek_v32`.\n
- "qwen_vl" will always use the tokenizer from `qwen_vl`.\n
if available, otherwise it will use the "hf" tokenizer.
- "hf" will use the fast tokenizer if available.
- "slow" will always use the slow tokenizer.
- "mistral" will always use the tokenizer from `mistral_common`.
- "deepseek_v32" will always use the tokenizer from `deepseek_v32`.
- "qwen_vl" will always use the tokenizer from `qwen_vl`.
- Other custom values can be supported via plugins."""
trust_remote_code: bool = False
"""Trust remote code (e.g., from HuggingFace) when downloading the model
and tokenizer."""
dtype: ModelDType | torch.dtype = "auto"
"""Data type for model weights and activations:\n
"""Data type for model weights and activations:
- "auto" will use FP16 precision for FP32 and FP16 models, and BF16
precision for BF16 models.\n
- "half" for FP16. Recommended for AWQ quantization.\n
- "float16" is the same as "half".\n
- "bfloat16" for a balance between precision and range.\n
- "float" is shorthand for FP32 precision.\n
precision for BF16 models.
- "half" for FP16. Recommended for AWQ quantization.
- "float16" is the same as "half".
- "bfloat16" for a balance between precision and range.
- "float" is shorthand for FP32 precision.
- "float32" for FP32 precision."""
seed: int = 0
"""Random seed for reproducibility.
@@ -182,13 +184,14 @@ class ModelConfig:
automatically derived from the model config.
When passing via `--max-model-len`, supports k/m/g/K/M/G in human-readable
format. Examples:\n
- 1k -> 1000\n
- 1K -> 1024\n
- 25.6k -> 25,600\n
format. Examples:
- 1k -> 1000
- 1K -> 1024
- 25.6k -> 25,600
- -1 or 'auto' -> Automatically choose the maximum model length that fits in
GPU memory. This will use the model's maximum context length if it fits,
otherwise it will find the largest length that can be accommodated."""
GPU memory. This will use the model's maximum context length if it fits,
otherwise it will find the largest length that can be accommodated."""
spec_target_max_model_len: int | None = None
"""Specify the maximum length for spec decoding draft models."""
quantization: QuantizationMethods | str | None = None
@@ -248,10 +251,11 @@ class ModelConfig:
prometheus metrics, if multiple names provided, metrics tag will take the
first one."""
config_format: str | ConfigFormat = "auto"
"""The format of the model config to load:\n
"""The format of the model config to load:
- "auto" will try to load the config in hf format if available after trying
to load in mistral format.\n
- "hf" will load the config in hf format.\n
to load in mistral format.
- "hf" will load the config in hf format.
- "mistral" will load the config in mistral format."""
hf_token: bool | str | None = None
"""The token to use as HTTP bearer authorization for remote files . If
@@ -276,12 +280,12 @@ class ModelConfig:
"""Enable sleep mode for the engine (only cuda and
hip platforms are supported)."""
model_impl: str | ModelImpl = "auto"
"""Which implementation of the model to use:\n
- "auto" will try to use the vLLM implementation, if it exists, and fall
back to the Transformers implementation if no vLLM implementation is
available.\n
- "vllm" will use the vLLM model implementation.\n
- "transformers" will use the Transformers model implementation.\n
"""Which implementation of the model to use:
- "auto" will try to use the vLLM implementation, if it exists, and fall back to the
Transformers implementation if no vLLM implementation is available.
- "vllm" will use the vLLM model implementation.
- "transformers" will use the Transformers model implementation.
- "terratorch" will use the TerraTorch model implementation.
"""
override_attention_dtype: str | None = None
@@ -291,6 +295,10 @@ class ModelConfig:
definitions"""
io_processor_plugin: str | None = None
"""IOProcessor plugin name to load at model startup"""
renderer_num_workers: int = 1
"""Number of worker threads in the renderer thread pool. This pool
handles async tokenization, chat template rendering, and multimodal
preprocessing."""
# Pooler config
pooler_config: PoolerConfig | None = None
@@ -639,6 +647,19 @@ class ModelConfig:
self.multimodal_config = MultiModalConfig(**mm_config_kwargs) # type: ignore[arg-type]
if (
self.renderer_num_workers > 1
and self.multimodal_config.mm_processor_cache_gb > 0
):
raise ValueError(
"Cannot use --renderer-num-workers > 1 with the "
"multimodal processor cache enabled. The cache is "
"not thread-safe and does not support concurrent "
"renderer workers. Please set "
"--renderer-num-workers 1 (the default), or "
"disable the cache with --mm-processor-cache-gb 0."
)
# Multimodal GGUF models must use original repo for mm processing
if is_gguf(self.tokenizer) and self.is_multimodal_model:
raise ValueError(
@@ -1508,10 +1529,11 @@ class ModelConfig:
@property
def score_type(self) -> ScoreType:
"""
Scoring API handles score/rerank for:\n
- "classify" task (score_type: cross-encoder models)\n
- "embed" task (score_type: bi-encoder models)\n
- "token_embed" task (score_type: late interaction models)\n
Scoring API handles score/rerank for:
- "classify" task (score_type: cross-encoder models)
- "embed" task (score_type: bi-encoder models)
- "token_embed" task (score_type: late interaction models)
"""
# fixme: self._model_info.score_type is the score type before
# as_seq_cls_model, which is "bi-encoder", rather than the
@@ -1589,9 +1611,10 @@ class ModelConfig:
such as the lm_head in a generation model,
or the score or classifier in a classification model.
`head_dtype` currently only supports pooling models.\n
- The pooling model defaults to using fp32 head,
you can use --hf-overrides '{"head_dtype": "model"}' to disable it.
`head_dtype` currently only supports pooling models.
- The pooling model defaults to using fp32 head, you can use
--hf-overrides '{"head_dtype": "model"}' to disable it.
"""
head_dtype = _get_head_dtype(
+7 -7
View File
@@ -146,14 +146,14 @@ class MultiModalConfig:
parallelism (TP).
- `"weights"`: Within the same vLLM engine, split the weights of
each layer across TP ranks. (default TP behavior)\n
each layer across TP ranks. (default TP behavior)
- `"data"`: Within the same vLLM engine, split the batched input data
across TP ranks to process the data in parallel, while hosting
the full weights on each TP rank.
This batch-level DP is not to be confused with API request-level
DP (which is controlled by `--data-parallel-size`).
This is only supported on a per-model basis and falls back to
`"weights"` if the encoder does not support DP."""
across TP ranks to process the data in parallel, while hosting
the full weights on each TP rank.
This batch-level DP is not to be confused with API request-level
DP (which is controlled by `--data-parallel-size`).
This is only supported on a per-model basis and falls back to
`"weights"` if the encoder does not support DP."""
mm_encoder_attn_backend: AttentionBackendEnum | None = None
"""Optional override for the multi-modal encoder attention backend when
using vision transformers. Accepts any value from
+8 -7
View File
@@ -148,10 +148,11 @@ class ParallelConfig:
eplb_config: EPLBConfig = Field(default_factory=EPLBConfig)
"""Expert parallelism configuration."""
expert_placement_strategy: ExpertPlacementStrategy = "linear"
"""The expert placement strategy for MoE layers:\n
"""The expert placement strategy for MoE layers:
- "linear": Experts are placed in a contiguous manner. For example, with 4
experts and 2 ranks, rank 0 will have experts [0, 1] and rank 1 will have
experts [2, 3].\n
experts [2, 3].
- "round_robin": Experts are placed in a round-robin manner. For example,
with 4 experts and 2 ranks, rank 0 will have experts [0, 2] and rank 1
will have experts [1, 3]. This strategy can help improve load balancing
@@ -159,11 +160,11 @@ class ParallelConfig:
all2all_backend: All2AllBackend = "allgather_reducescatter"
"""All2All backend for MoE expert parallel communication. Available options:
- "allgather_reducescatter": All2all based on allgather and reducescatter\n
- "deepep_high_throughput": Use deepep high-throughput kernels\n
- "deepep_low_latency": Use deepep low-latency kernels\n
- "mori": Use mori kernels\n
- "nixl_ep": Use nixl-ep kernels\n
- "allgather_reducescatter": All2all based on allgather and reducescatter
- "deepep_high_throughput": Use deepep high-throughput kernels
- "deepep_low_latency": Use deepep low-latency kernels
- "mori": Use mori kernels
- "nixl_ep": Use nixl-ep kernels
- "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl
- "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels"""
+1 -1
View File
@@ -37,7 +37,7 @@ class ProfilerConfig:
profiler: ProfilerKind | None = None
"""Which profiler to use. Defaults to None. Options are:
- 'torch': Use PyTorch profiler.\n
- 'torch': Use PyTorch profiler.
- 'cuda': Use CUDA profiler."""
torch_profiler_dir: str = ""

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