Compare commits

...
Author SHA1 Message Date
yewentao256 7d98b72bd3 fix ci pre-commit
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-07-23 18:15:27 +00:00
75ccdf3145 [Core] Update PyTorch to 2.13.0, torchvision to 0.28.0, triton to 3.7.1 (#48155)
Signed-off-by: Andrey Talman <atalman@users.noreply.github.com>
Co-authored-by: Andrey Talman <atalman@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-23 11:09:36 -07:00
Andrey TalmanandGitHub c6fe94b4d5 [CI] Bump PyTorch Compilation Unit Tests timeout to 150 min (#49606) 2026-07-23 11:08:45 -07:00
46f01a50ac [CI][Bugfix] Fix test isolation in block_int8/ptpc_fp8 MoE kernel tests (#49609)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-23 12:26:02 -05:00
Andreas KaratzasandGitHub f00efc5265 [CI] Isolate cudagraph tests in child processes (#49510)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
2026-07-23 17:16:19 +00:00
Wentao YeandGitHub b0cb1da1bd [DSv4 Perf] Skip topk and router when not needed, 3.4% E2E TTFT improvement for Decode case (#49486)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-07-23 13:08:08 -04:00
Andreas KaratzasandGitHub 0e36e3bbd1 [CI] Use explicit devices in quantization tests (#49512)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
2026-07-23 10:54:26 -06:00
494845e79f Revert "[MRV2] Always build attn metadata at capture time" (#49364) (#49451)
Co-authored-by: vllm-agent CI bot <ci-bot@vllm-agent.local>
2026-07-23 09:51:33 -07:00
yue.yuandGitHub 0416dab275 [Bugfix][Structured Output][Spec Decode] Advance grammar across reasoning boundary (#44993)
Signed-off-by: Allen.Yu <yuyue0225sc@163.com>
2026-07-23 09:14:15 -07:00
c8db00b16c Fix GPTQ quantized Qwen3.5 MTP weight loading with spec decode (#48816)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Co-authored-by: noobHappylife <64898326+noobHappylife@users.noreply.github.com>
2026-07-23 06:59:48 -07:00
Guan-Ming ChiuandGitHub 80c7683923 [Perf] Defer MM embeds loading off the event loop (#49477)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
2026-07-23 13:58:50 +00:00
638d6e9757 [Bugfix][CI/Build] Fix Plamo2 HF runner crash on transformers v5 (_tied_weights_keys list→dict) (#44239)
Signed-off-by: Nikhil Kulkarni <nikhilkulkarni1755@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-23 12:25:14 +00:00
Junpu YuandGitHub 1ad84fea86 [Bugfix][Spec Decode] Select earliest-completing stop string in check_stop_strings (#49391)
Signed-off-by: Junpu Yu <davidyu@nvidia.com>
2026-07-23 18:14:06 +08:00
12213c6795 [Bugfix] handle grammar compilation failures to avoid engine crash (#47312)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-07-23 18:13:49 +08:00
10c75477b0 [Bugfix][Core] shm_broadcast: bound idle reader waits and release read slots (#45224)
Signed-off-by: Chaemin Lim <chaemin.lim@mangoboost.io>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Edwin Lim <edwin.lim@mangoboost.io>
Co-authored-by: Jaeyoun Kim <jaeyoun.kim@mangoboost.io>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-07-23 18:13:29 +08:00
ac36a7a1e7 [MRV2][Spec Decode] Avoid rejection sampler OOM by chunking (#48630)
Signed-off-by: mgoin <mgoin64@gmail.com>
Signed-off-by: Michael Goin <mgoin64@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-07-23 18:13:13 +08:00
Nick HillGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
521aa80f71 [Core] Simplify KVBlockZeroer index tensor handling (#48399)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-23 18:12:46 +08:00
a76df87db8 [MooncakeStore] Re-derive full external hits on stored boundaries (#49481)
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-23 18:12:28 +08:00
a4904ba903 [Perf][KVConnector][Mooncake] Vectorize prepare_value on the KV load path (#48531)
Signed-off-by: girasoley <girasoley@inferact.ai>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: girasoley <girasoley@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 18:12:00 +08:00
Rehan KhanandGitHub f83de6d44c [CPU][Docs] Update docs and dockerfile for s390x (#49523)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
2026-07-23 07:17:24 +00:00
Umut PolatandGitHub 239fc73553 [Misc] Use VLLMValidationError in chat_utils content-part validation (#49217)
Signed-off-by: Umut Polat <52835619+umut-polat@users.noreply.github.com>
2026-07-23 05:48:37 +00:00
Mike GandGitHub 76bf55240c [Bugfix] Fix DeepSeek-V4 DSpark draft shared-expert padding for TP > 8 (#49415)
Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com>
2026-07-23 05:06:21 +00:00
9a698f3255 [Performance][Model] Avoid transient Inkling result allocations (performance, and OOM prevention on smaller memory configurations) (#49487)
Signed-off-by: Michael Gschwind <mgschwind@nvidia.com>
Co-authored-by: Michael Gschwind <mgschwind@nvidia.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 05:03:34 +00:00
4080263bb2 [Bugfix][Model] Remove SciPy dependency from Inkling scale planning (#49485)
Signed-off-by: Michael Gschwind <mgschwind@nvidia.com>
Co-authored-by: Michael Gschwind <mgschwind@nvidia.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 04:34:57 +00:00
jcotant-inferactandGitHub fc5fda105f [Docs] Re-add Reo.dev analytics beacon (#49474) 2026-07-23 03:03:32 +00:00
Matej SirovatkaandGitHub b07ec92faa [Bugfix] Make shared NVFP4 MoE scales writable (#49489)
Signed-off-by: S1ro1 <matej.sirovatka@gmail.com>
2026-07-22 19:17:53 -07:00
Colin ZGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
27ffbfde8d Fused Shared Expert Support for AMD Quark DeepSeek-V4 Model Checkpoints (#48044)
Signed-off-by: Colin Zeng <Colin.Zeng@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-23 00:34:17 +00:00
Nick HillandGitHub 229e01e9e1 [BugFix] Handle per-group prefix-hit divergence for hybrid models with KV connector (#48425) 2026-07-22 17:19:11 -07:00
191146dba5 Add quantization label automation (#49492)
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 19:59:15 -04:00
Summer YangandGitHub f3a920a076 [Core][DSV4] Compact MXFP4 indexer KV cache and packed group overlays (#48993) 2026-07-22 16:58:44 -07:00
Nils MattesonGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
149daf0d72 [Bugfix] Exclude location-derived path vars from torch.compile cache factors (#47573)
Signed-off-by: Nils Matteson <nilsmatteson@icloud.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-22 16:56:09 -07:00
Michael GoinandGitHub 917fdb5bf7 [Bugfix] Fix DeepGEMM warmup when using FlashInferFp8DeepGEMMDynamicBlockScaledKernel (#49467)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-07-22 16:28:49 -07:00
stefankoncarevicandGitHub 4b594b4aa1 [Bugfix][CI] Fix topk_softplus_sqrt no-op on non-XPU platforms (#49452)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
2026-07-22 15:36:17 -07:00
7d10a4cfce [Bugfix] Retry config read to survive concurrent HF cache refresh (#49001)
Signed-off-by: pei.zhang <pei.zhang@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-22 15:35:38 -07:00
Nick HillandGitHub 910cc8543a [Bugfix] Restore gather_and_maybe_dequant_cache OOB guard (#49427)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 13:47:37 -07:00
Nick HillandGitHub 431934522b [CI] Fix stale/fragile untethered kernels-root tests (#49423)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 14:43:07 -06:00
61a09532f2 Bump Flashinfer version to 0.6.15 (#48914)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Signed-off-by: Wei Zhao <weizha@oci-aga-slurm-1-vscode-02.cm.cluster>
Co-authored-by: Wei Zhao <weizha@oci-aga-slurm-1-vscode-02.cm.cluster>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-07-22 13:32:05 -07:00
Ben BrowningGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
3de4b2bf3c [Bugfix][Parser] Fix special tokens (EOS/BOS) leaking into reasoning content (#48748)
Signed-off-by: Ben Browning <56071+bbrowning@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-22 16:10:55 -04:00
b44311b6ef [CI] stabilize GDN prefill CuTeDSL test (#49388)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: Codex <noreply@openai.com>
2026-07-22 09:03:46 -07:00
Nick HillandGitHub b0d7875180 [CI] Increase timeout of pytorch-compilation-unit-tests (#49450)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 15:39:09 +00:00
Divakar VermaandGitHub 53c2f20dd9 [ROCm][CI] skip moe weight padding for eplb (#49350)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
2026-07-22 10:18:01 -05:00
Wentao YeandGitHub 37e370fe93 [DSv4 Perf] Skip empty c128 kernel launch, around 2x kernel performance improvement. (#48957)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-07-22 10:55:20 -04:00
Guan-Ming ChiuandGitHub 2dc5a72e7e [Bugfix][Renderer] Rebuild vision chunk UUIDs in async render path (#49400)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
2026-07-22 14:11:24 +00:00
Andrey TalmanandGitHub c79ff5f918 [Build] Bump vllm-flash-attn to C++20-compatible commit for torch-nightly (#49326)
Signed-off-by: Andrey Talman <atalman@fb.com>
2026-07-22 13:51:59 +00:00
Teresa ChenandGitHub 1a659a0c37 Upgrade tpu-inference to v0.25.0 (#49431) 2026-07-22 11:52:56 +00:00
SageandGitHub 0f6cf7f628 [Rust Frontend] Extract request preparation from the inference path (#49045)
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
2026-07-22 11:31:36 +00:00
c79ad3ae21 [Rust Frontend][gRPC] Add abort control RPC (#49255)
Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Connor Carpenter <connorc@nvidia.com>
2026-07-22 11:31:01 +00:00
wang.yuqiandGitHub 61c9ef986a [Frontend] Parallelize preprocessing within the same request for pooling models online serving. (#49153)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-07-22 10:56:23 +00:00
LiangqiusongandGitHub d6dbdb9b0d [XPU] WA of topk_softplus_sqrt arg mismatch on XPU (#49408)
Signed-off-by: xiaolong <xiaolong.guo@intel.com>
2026-07-22 16:16:13 +08:00
liuzhenweiandGitHub 06da482fb4 [XPU] WA of topk_softmax arg mismatch on XPU (#49395)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
2026-07-22 01:05:01 -07:00
2f75e7f712 [CI] Increase timeouts for jobs exceeding current limits (#49374)
Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 00:16:16 -07:00
Ziming HuangandGitHub 7c21548ce3 [PD][Bugfix] Fix NIXL hybrid MLA+mamba heterogeneous TP (#49297)
Signed-off-by: ZeldaHuang <zelda.huanghuang@gmail.com>
2026-07-22 07:11:39 +00:00
Guan-Ming ChiuandGitHub 9df2f91232 [Renderer] Offload derender CPU work to renderer thread pool (#49396)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
2026-07-22 06:55:03 +00:00
387189c429 [ROCm] Remove redundant AITER fused_qk_rmsnorm probe (avoids config-time HIP init) (#47992)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-21 22:58:14 -05:00
Kunshang JiandGitHub 75576c63be Add auto label for xpu relate issue (#49398)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-22 03:33:20 +00:00
Andreas KaratzasandGitHub 16aca639b7 [ROCm] Upgrade NIXL and UCX (#49251)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
2026-07-21 22:30:02 -05:00
Woosuk KwonandGitHub 6049424b7e [MRV2] Always build attn metadata at capture time (#49364)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-07-21 19:27:16 -07:00
060b5f61dc [Bugfix][Attention] Ignore empty MLA context chunks during merge (#49294)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 02:23:48 +00:00
ec59c1579f [MoE Refactor] Migrate MoeWNA16Method quantization method over to using the new MK oracle scheme. (#44120)
Signed-off-by: Bill Nell <bnell@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: mgoin <mgoin64@gmail.com>
2026-07-21 19:20:30 -07:00
Isotr0pyandGitHub 1750e443f2 [Misc] Move PyNvVideoCodec stuff out of gpu worker (#49322)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
2026-07-22 10:11:56 +08:00
ba18929079 [Bugfix][SpecDecode] Scope MTP completeness checks outside bucketed updates (#49178)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
2026-07-21 18:56:39 -07:00
0500ca6a58 [CI][Bugfix] Fix ROCm FP8 KV cache dtype in attention backend test (#49380)
Signed-off-by: pei.zhang <pei.zhang@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-22 01:34:49 +00:00
a1c15bcb0f [CI][Bugfix] Fix and wire streaming-input tests (#49356)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-22 00:35:23 +00:00
Nick HillandGitHub 4809de7317 [Misc] Fix terminal output logo coloring (#49344)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 00:21:29 +00:00
stefankoncarevicandGitHub 05781e21dd [ROCm][CI] Fix order-dependent failure in test_flash_attn_accepts_handled_fp8_variants (MI355) (#49329)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
2026-07-21 18:54:33 -05:00
gnovackandGitHub 85f638a2b8 skip cudagraph/DP padding in topk (#48979)
Signed-off-by: gnovack <novackgm@gmail.com>
2026-07-21 15:20:30 -07:00
Nick HillandGitHub 08e5067561 [CI] Bump timeout of entrypoints-integration-api-server-openai-part-2 (#49359)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-21 14:54:03 -07:00
a7d00ec051 [Bugfix] DFlash fc sized wrong when num_target_layers != num_hidden_layers (#48524)
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
2026-07-21 14:42:52 -07:00
Michael GoinandGitHub b8fb56d970 [CI] Add gemma-4-E4B-it-assistant to CI gsm8k for GemmaMTP (#49243)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-07-21 16:51:04 -04:00
96a739289e [Bugfix] fix cutalss version upgrade bug, need update MSG new commit (#49016)
Signed-off-by: lengrongfu <lenronfu@gmail.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
2026-07-21 13:42:14 -07:00
60d443f738 [CI/Build][The Rock][BugFix] Use fork method in test_multiproc_executor_multi_node for py 3.14 compat and fix test_multiproc_executor_shutdown_cleanup (#48655)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-21 15:40:29 -05:00
1dca300653 [CI] Fix and wire encoder/manager cudagraph unit tests (#49339)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 20:41:16 +01:00
Flora FengandGitHub fca252d59e [CI][Bugfix] Reduce max_model_len in OOT embedding test to fix KV-cache OOM on small GPUs (#49351)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-07-21 15:00:06 -04:00
33178f9006 Fix Qwen3-VL M-RoPE on the Transformers modeling backend (grids + compile) (#49292)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-21 18:28:39 +00:00
b2b8f679d0 [Bugfix][Spec Decode] Restrict embedding-width share guard to EAGLE drafts (#47953)
Signed-off-by: Evan Takahashi <evan10takahashi@gmail.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
2026-07-21 10:53:37 -07:00
de6ec294ef [Bugfix] Fix DSA crash under breakable piecewise cudagraphs (#49302)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 18:28:13 +01:00
stefankoncarevicandGitHub 61e10f0116 [ROCm][CI] Fix AITER MLA fp8 decode metadata regression test (#48845)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
2026-07-21 12:19:45 -05:00
6e96891ba0 [ROCm] Bump AITER to v0.1.16.post5 (#48683)
Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com>
Signed-off-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-21 10:39:11 -05:00
47f1b47a73 Ci/add laguna xs gsm8k (#49241)
Signed-off-by: Joe Rowell <joerowell4@gmail.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Joe Rowell <joerowell4@gmail.com>
2026-07-21 11:30:42 -04:00
5aab491bc9 [CI] Wire tests/models/inkling into a B200 job (#49325)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 15:43:20 +01:00
5812e1a66b [Test] Add DeepSeek MTP parallel-load tests (#41653)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
2026-07-21 10:38:20 -04:00
8950394e0a [Bugfix] Prefix-cache metrics double-counted when a KV connector defers requests (#48860)
Signed-off-by: Seiji Eicher <58963096+eicherseiji@users.noreply.github.com>
Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Lixi Zhang <muxixibbb@gmail.com>
Co-authored-by: harjothkhara <harjoth.khara@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-07-21 15:02:03 +01:00
Roberto L. CastroandGitHub 7bb49be4d1 [Bugfix] Handle MLA fallback during FA4 JIT warmup (#49306) 2026-07-21 13:58:05 +00:00
c67650f04b [XPU][DeepSeekV4]Add DeepSeek-V4 fuse_index_q SYCL kernel path (#45991)
Signed-off-by: Avery Miao <avery.miao@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-21 21:22:34 +08:00
f890e1dbe2 [BugFix] Set graph_pool_id before FULL CUDA graph capture in ModelRunner V2 (#48843)
Signed-off-by: Markov Ilya <markovilya19@gmail.com>
Co-authored-by: Markov Ilya <markovilya19@gmail.com>
2026-07-21 08:59:26 -04:00
234 changed files with 6494 additions and 2485 deletions
+28 -28
View File
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_METADATA_VERSION="1"
IMAGE_EXISTED_BEFORE_BUILD=0
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
return 0
fi
for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
if [[ -n "${!var:-}" ]]; then
echo "Using provided ${var}: ${!var}"
continue
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
compute_dependency_cache_keys() {
local bake_dir=""
local dockerfile_rocm=""
local rixl_branch=""
local nixl_branch=""
local ucx_branch=""
local rocshmem_branch=""
local deepep_branch=""
local rixl_material=""
local nixl_material=""
local rocshmem_material=""
local deepep_material=""
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
RIXL_CACHE_KEY=$(
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
NIXL_CACHE_KEY=$(
compose_dependency_cache_key \
"${rixl_branch}-ucx-${ucx_branch}" \
"${rixl_material}"
"${nixl_branch}-ucx-${ucx_branch}" \
"${nixl_material}"
)
export RIXL_CACHE_KEY
echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
export NIXL_CACHE_KEY
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
fi
if [[ -n "${rocshmem_branch}" ]]; then
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
case "${target}" in
rixl-rocm-ci)
if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
elif [[ -n "${RIXL_BRANCH:-}" ]]; then
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
nixl-rocm-ci)
if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
elif [[ -n "${NIXL_BRANCH:-}" ]]; then
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
fi
;;
rocshmem-rocm-ci)
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
resolve_ci_base_dependency_targets() {
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
local rixl_ref=""
local nixl_ref=""
local rocshmem_ref=""
local deepep_ref=""
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
case "${mode}" in
always)
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
add_dependency_cache_target "${target}"
fi
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
;;
esac
if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
if dependency_cache_ref_exists "${rixl_ref}"; then
echo "RIXL dependency cache exists: ${rixl_ref}"
if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
if dependency_cache_ref_exists "${nixl_ref}"; then
echo "NIXL dependency cache exists: ${nixl_ref}"
else
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
add_dependency_cache_target "rixl-rocm-ci"
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
add_dependency_cache_target "nixl-rocm-ci"
fi
fi
@@ -1,10 +1,11 @@
#!/bin/bash
set -euox pipefail
export VLLM_CPU_KVCACHE_SPACE=1
export VLLM_CPU_KVCACHE_SPACE=1
export VLLM_CPU_CI_ENV=1
# Reduce sub-processes for acceleration
export TORCH_COMPILE_DISABLE=1
# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of
# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op
# when callers specify fullgraph=True.
export VLLM_ENABLE_V1_MULTIPROCESSING=0
SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
@@ -49,15 +50,15 @@ wait_for_pid_and_check_log() {
}
# Test Sky Lake (AVX512F)
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 &
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 &
PID_TEST_0=$!
# Test Cascade Lake (AVX512F + VNNI)
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 &
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 &
PID_TEST_1=$!
# Test Cooper Lake (AVX512F + VNNI + BF16)
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
PID_TEST_2=$!
wait_for_pid_and_check_log $PID_TEST_0 test_0.log
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Basic Correctness
key: basic-correctness
timeout_in_minutes: 45
timeout_in_minutes: 68
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Benchmarks CLI Test
key: benchmarks-cli-test
timeout_in_minutes: 30
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
+3
View File
@@ -26,7 +26,10 @@ steps:
- vllm/v1/cudagraph_dispatcher.py
- vllm/config/compilation.py
- vllm/compilation
- vllm/v1/worker/encoder_cudagraph.py
- vllm/v1/worker/encoder_cudagraph_defs.py
commands:
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
- pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
+1 -1
View File
@@ -51,7 +51,7 @@ steps:
- label: e2e Scheduling (1 GPU)
key: e2e-scheduling-1-gpu
timeout_in_minutes: 35
timeout_in_minutes: 53
device: h200_18gb
source_file_dependencies:
- vllm/v1/
+4 -4
View File
@@ -39,7 +39,7 @@ steps:
- label: Entrypoints Integration (API Server)
key: entrypoints-integration-api-server
device: h200_35gb
timeout_in_minutes: 50
timeout_in_minutes: 75
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -59,7 +59,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 1)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 45
timeout_in_minutes: 68
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -78,7 +78,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 2)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 45
timeout_in_minutes: 83
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -156,7 +156,7 @@ steps:
- label: Entrypoints Integration (Pooling)
device: h200_35gb
key: entrypoints-integration-pooling
timeout_in_minutes: 50
timeout_in_minutes: 75
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
+4 -1
View File
@@ -31,7 +31,7 @@ steps:
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 45
timeout_in_minutes: 83
device: h200_18gb
source_file_dependencies:
- vllm/config/
@@ -90,6 +90,7 @@ steps:
- tests/v1/kv_offload
- tests/v1/simple_kv_offload
- tests/v1/worker
- tests/v1/streaming_input
- tests/v1/kv_connector/unit
- tests/v1/ec_connector/unit
- tests/v1/metrics
@@ -103,6 +104,7 @@ steps:
- pytest -v -s v1/kv_offload
- pytest -v -s v1/simple_kv_offload
- pytest -v -s v1/worker
- pytest -v -s v1/streaming_input
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
@@ -143,6 +145,7 @@ steps:
- pytest -v -s -m 'cpu_test' v1/core
- pytest -v -s v1/structured_output
- pytest -v -s v1/test_serial_utils.py
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'cpu_test' v1/metrics
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: Model Executor
device: h200_35gb
key: model-executor
timeout_in_minutes: 45
timeout_in_minutes: 60
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
+13
View File
@@ -46,6 +46,19 @@ steps:
depends_on:
- image-build-amd
- label: Inkling Unit Tests (B200)
key: inkling-unit-tests-b200
timeout_in_minutes: 40
device: b200-k8s
source_file_dependencies:
- vllm/models/inkling/
- vllm/cute_utils/
- cmake/external_projects/tml_fa4.cmake
- tests/models/inkling/
commands:
# FA4 kernel tests require SM100; the suite skips them elsewhere.
- pytest -v -s models/inkling
- label: Basic Models Test (Other CPU) # 5min
key: basic-models-test-other-cpu
depends_on:
+1 -1
View File
@@ -137,7 +137,7 @@ steps:
- label: Language Models Test (MTEB)
key: language-models-test-mteb
timeout_in_minutes: 45
timeout_in_minutes: 68
device: h200_18gb
optional: true
source_file_dependencies:
+4 -4
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 45
timeout_in_minutes: 68
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -20,7 +20,7 @@ steps:
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 50
timeout_in_minutes: 75
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -54,7 +54,7 @@ steps:
- label: "Multi-Modal Models (Standard) 4: other + whisper"
device: h200_35gb
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 50
timeout_in_minutes: 75
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -85,7 +85,7 @@ steps:
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 65
timeout_in_minutes: 98
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: PyTorch Compilation Unit Tests
device: h200_35gb
key: pytorch-compilation-unit-tests
timeout_in_minutes: 90
timeout_in_minutes: 150
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
+16
View File
@@ -170,3 +170,19 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
key: spec-decode-deepseek-mtp-parallel-load-b200
timeout_in_minutes: 30
device: b200-k8s
optional: true
num_devices: 2
source_file_dependencies:
- vllm/v1/spec_decode/llm_base_proposer.py
- vllm/v1/spec_decode/eagle.py
- vllm/v1/worker/gpu/spec_decode/eagle/
- vllm/model_executor/models/deepseek_mtp.py
- vllm/model_executor/models/deepseek_v2.py
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
commands:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
+12
View File
@@ -181,6 +181,18 @@ pull_request_rules:
add:
- performance
- name: label-quantization
description: Automatically apply quantization label
conditions:
- label != stale
- or:
- files~=^vllm/model_executor/layers/quantization/
- title~=(?i)quant
actions:
label:
add:
- quantization
- name: label-qwen
description: Automatically apply qwen label
conditions:
+42 -1
View File
@@ -130,6 +130,47 @@ jobs:
},
],
},
quantization: {
keywords: [
{
term: "quantization",
searchIn: "both"
},
{
term: "quantized",
searchIn: "both"
},
],
},
"intel-gpu": {
// Keyword search - matches whole words only (with word boundaries)
keywords: [
{
term: "B50",
searchIn: "both"
},
{
term: "B60",
searchIn: "both"
},
{
term: "B70",
searchIn: "both"
},
{
term: "intel gpu",
searchIn: "both"
},
{
term: "Arc GPU",
searchIn: "both"
},
{
term: "BMG",
searchIn: "both"
},
],
},
// Add more label configurations here as needed
// example: {
// keywords: [...],
@@ -491,4 +532,4 @@ jobs:
issue_number: context.issue.number,
body: message,
});
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
+2 -2
View File
@@ -68,8 +68,8 @@ endif()
# requirements.txt files and should be kept consistent. The ROCm torch
# versions are derived from docker/Dockerfile.rocm
#
set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported-
# version check would always warn. Only treat it as a nightly build when the
# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which
+1 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+3
View File
@@ -1025,6 +1025,9 @@ __global__ void gather_and_maybe_dequant_cache(
batch_offset += offset;
int32_t block_table_id = batch_offset / block_size;
int32_t slot_id = batch_offset % block_size;
// seq_starts may push the block index past the end of the batch's block
// table row.
if (block_table_id >= block_table_stride) continue;
int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
int32_t block_id = block_table[block_table_offset];
int64_t cache_offset =
+6 -3
View File
@@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias);
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> is_padding);
void topk_sigmoid(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor);
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding);
void topk_softplus_sqrt(
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
@@ -25,7 +27,8 @@ void topk_softplus_sqrt(
double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid);
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding);
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
std::optional<torch::stable::Tensor> topk_ids,
@@ -174,7 +174,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int end_expert,
const bool renormalize,
const float* bias,
const double routed_scaling_factor)
const double routed_scaling_factor,
const bool* is_padding)
{
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int expert = result_kvp.key;
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool is_pad_row = is_padding != nullptr && is_padding[block_row];
const int idx = k * block_row + k_idx;
// Return the unbiased scores for output weights
output[idx] = inputs_after_softmax[thread_read_offset + expert];
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : num_experts);
assert(is_pad_row || indices[idx] >= 0);
source_rows[idx] = k_idx * num_rows + block_row;
if (renormalize) {
selected_sum += inputs_after_softmax[thread_read_offset + expert];
@@ -277,7 +280,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias, const double routed_scaling_factor)
const float* bias, const double routed_scaling_factor, const bool* is_padding)
{
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Add a guard to ignore experts not included by this node
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
// The lead thread from each sub-group will write out the final results to global memory. (This will be a
// single) thread per row of the input/output matrices.
const int idx = k * thread_row + k_idx;
output[idx] = max_val;
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -605,7 +610,7 @@ struct TopkConstants
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias, const double routed_scaling_factor, cudaStream_t stream)
const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)
{
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
@@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding);
}
#ifndef USE_ROCM
@@ -627,7 +632,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream);
bias, routed_scaling_factor, stream, is_padding);
#else
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -635,13 +640,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream); \
bias, routed_scaling_factor, stream, is_padding); \
} else if (WARP_SIZE == 32) { \
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream); \
bias, routed_scaling_factor, stream, is_padding); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -661,7 +666,8 @@ void topkGatingKernelLauncher(
const bool renormalize,
const float* bias,
const double routed_scaling_factor,
cudaStream_t stream) {
cudaStream_t stream,
const bool* is_padding) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
#ifndef USE_ROCM
@@ -736,7 +742,7 @@ void topkGatingKernelLauncher(
}
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding);
}
}
}
@@ -755,7 +761,8 @@ void dispatch_topk_launch(
int num_tokens, int num_experts, int topk, bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
cudaStream_t stream)
cudaStream_t stream,
std::optional<torch::stable::Tensor> is_padding)
{
const float* bias_ptr = nullptr;
if (bias.has_value()) {
@@ -769,6 +776,18 @@ void dispatch_topk_launch(
bias_ptr = bias_tensor.const_data_ptr<float>();
}
const bool* is_padding_ptr = nullptr;
if (is_padding.has_value()) {
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,
"is_padding tensor must be bool");
STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
"is_padding size mismatch, expected: ", num_tokens);
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
}
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -777,7 +796,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream);
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
} else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -786,7 +805,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream);
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
} else {
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
@@ -796,7 +815,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream);
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
}
}
@@ -806,7 +825,8 @@ void topk_softmax(
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias)
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> is_padding)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -825,15 +845,15 @@ void topk_softmax(
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream);
bias, 1.0, stream, is_padding);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream);
bias, 1.0, stream, is_padding);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream);
bias, 1.0, stream, is_padding);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -846,7 +866,8 @@ void topk_sigmoid(
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor)
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -865,15 +886,15 @@ void topk_sigmoid(
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream);
bias, routed_scaling_factor, stream, is_padding);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream);
bias, routed_scaling_factor, stream, is_padding);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream);
bias, routed_scaling_factor, stream, is_padding);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -80,22 +80,27 @@ __launch_bounds__(128) __global__
OutIndType* indices, int num_rows,
int num_experts, float routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid) {
const HashIndType* tid2eid,
const bool* is_padding) {
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int lane = threadIdx.x % 32;
if (warp >= num_rows) return;
const int64_t token_id = load_index_as_int64(input_ids, warp);
const bool is_pad_row = is_padding != nullptr && is_padding[warp];
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
int expert = 0;
float weight = 0.f;
if (lane < 6) {
if (lane < 6 && !is_pad_row) {
// only load and calculate for 6 experts
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
const float x = input[warp * num_experts + expert];
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
if (isnan(weight)) {
weight = 0.f;
}
}
float weight_sum = weight;
#pragma unroll
@@ -111,7 +116,8 @@ __launch_bounds__(128) __global__
const int offset = warp * 6 + lane;
output[offset] =
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
indices[offset] = static_cast<OutIndType>(expert);
indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
: static_cast<OutIndType>(-1);
}
}
@@ -120,7 +126,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
int num_rows, int num_experts,
double routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream) {
const HashIndType* tid2eid, cudaStream_t stream,
const bool* is_padding) {
if (num_rows == 0) return;
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
cudaLaunchConfig_t config = {};
@@ -134,7 +141,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
config.numAttrs = 1;
const float scale = static_cast<float>(routed_scaling_factor);
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
num_experts, scale, input_ids, tid2eid);
num_experts, scale, input_ids, tid2eid, is_padding);
}
#endif
@@ -166,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
const int num_rows, IndType* indices, int* source_rows, const int k,
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const HashIndType* input_ids, const HashIndType* tid2eid) {
const HashIndType* input_ids, const HashIndType* tid2eid,
const bool* is_padding) {
static_assert(std::is_same_v<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -231,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
return;
}
const bool row_is_active = finished ? !finished[thread_row] : true;
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
// We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read.
@@ -249,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
cudaGridDependencySynchronize();
#endif
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
// to float
if constexpr (std::is_same_v<InputType, float>) {
if (is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = 0.f;
}
} else if constexpr (std::is_same_v<InputType, float>) {
using VecType = AlignedArray<float, ELTS_PER_LDG>;
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
const VecType* vec_thread_read_ptr =
@@ -315,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
if constexpr (USE_HASH) {
const int64_t token_id = load_index_as_int64(input_ids, thread_row);
const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
if (!is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
row_chunk[ii] = sqrtf(val);
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
// Dummy/padding tokens can result in NaN values, so
// clamp them to 0.0. Note: this clamp could likely be removed if
// 'is_padding' is made mandatory
if (isnan(val)) {
val = 0.f;
}
row_chunk[ii] = val;
}
}
float selected_sum = 0.f;
#pragma unroll
@@ -335,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
if (expert == expert_idx) {
indices[idx] = static_cast<IndType>(expert);
indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
: static_cast<IndType>(-1);
selected_sum += row_chunk[ii];
break;
}
@@ -379,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
#endif
return;
} else {
if (!is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
// Dummy/padding tokens can result in NaN values, so
// clamp them to 0.0. Note: this clamp could likely be removed if
// 'is_padding' is made mandatory
if (isnan(val)) {
val = 0.f;
}
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
}
row_chunk[ii] = val;
}
// Original TopK path: find top-k experts by score
@@ -450,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Add a guard to ignore experts not included by this node
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool should_process_row =
row_is_active && node_uses_expert && !is_pad_row;
// The lead thread from each sub-group will write out the final results
// to global memory. (This will be a single) thread per row of the
// input/output matrices.
const int idx = k * thread_row + k_idx;
if (correction_bias != nullptr) {
if (correction_bias != nullptr && should_process_row) {
max_val -= correction_bias[expert];
}
output[idx] = max_val;
indices[idx] =
should_process_row ? (expert - start_expert) : NUM_EXPERTS;
!is_pad_row ? expert - start_expert : static_cast<IndType>(-1);
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -544,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const bool use_hash, const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream) {
const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
static constexpr int BYTES_PER_LDG =
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants =
@@ -573,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper(
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
indices, source_row, k, start_expert, end_expert,
renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid);
input_ids, tid2eid, is_padding);
#else
kernel<<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert,
end_expert, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid);
input_ids, tid2eid, is_padding);
#endif
})
}
@@ -592,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
stream);
stream, is_padding);
#else
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -601,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream); \
tid2eid, stream, is_padding); \
} else if (WARP_SIZE == 32) { \
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
MAX_BYTES>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream); \
tid2eid, stream, is_padding); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -622,14 +654,14 @@ void topkGatingSoftplusSqrtKernelLauncher(
const int topk, const bool renormalize, double routed_scaling_factor,
const float* correction_bias, const bool use_hash,
const HashIndType* input_ids, const HashIndType* tid2eid,
cudaStream_t stream) {
cudaStream_t stream, const bool* is_padding) {
#ifndef USE_ROCM
if constexpr (std::is_same_v<InputType, float>) {
if (use_hash && topk == 6 && renormalize &&
(num_experts == 256 || num_experts == 384)) {
launchDsv4HashTopk<IndType, HashIndType>(
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
routed_scaling_factor, input_ids, tid2eid, stream);
routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
return;
}
}
@@ -728,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch(
int num_experts, int topk, bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
const std::optional<torch::stable::Tensor>& is_padding) {
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
bias_ptr = correction_bias.value().const_data_ptr<float>();
@@ -737,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch(
auto launch = [&](auto* topk_indices_ptr) {
using OutIndType =
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
const bool* is_padding_ptr = nullptr;
if (is_padding.has_value()) {
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==
torch::headeronly::ScalarType::Bool,
"is_padding tensor must be bool");
STD_TORCH_CHECK(is_padding_tensor.dim() == 1,
"is_padding tensor must be 1D");
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
"is_padding size mismatch, expected: ", num_tokens);
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),
"is_padding tensor must be contiguous");
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
}
if (tid2eid.has_value()) {
STD_TORCH_CHECK(input_ids.has_value(),
"input_ids is required for hash MoE");
@@ -751,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
tid2eid.value().const_data_ptr<int64_t>(), stream);
tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
} else {
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
torch::headeronly::ScalarType::Int);
@@ -761,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
tid2eid.value().const_data_ptr<int>(), stream);
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
}
} else {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -769,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
static_cast<const OutIndType*>(nullptr), stream);
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
}
};
@@ -793,7 +842,8 @@ void topk_softplus_sqrt(
bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid) {
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding) {
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
const int topk = topk_weights.size(-1);
@@ -806,21 +856,22 @@ void topk_softplus_sqrt(
dispatch_topk_softplus_sqrt_launch<float>(
gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
token_expert_indices, num_tokens, num_experts, topk, renormalize,
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
is_padding);
} else if (gating_output.scalar_type() ==
torch::headeronly::ScalarType::Half) {
dispatch_topk_softplus_sqrt_launch<__half>(
reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream);
input_ids, tid2eid, stream, is_padding);
} else if (gating_output.scalar_type() ==
torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream);
input_ids, tid2eid, stream, is_padding);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
gating_output.scalar_type());
+3 -3
View File
@@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
"bias) -> ()");
"bias, Tensor? is_padding) -> ()");
// Apply topk sigmoid to the gating outputs.
m.def(
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, "
"Tensor? bias, float routed_scaling_factor) -> ()");
"Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()");
m.def(
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, float "
"routed_scaling_factor, Tensor? "
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
"bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
// Calculate the result of moe by summing up the partial results
// from all selected experts. topk_ids/expert_map are optional and, when
+2 -2
View File
@@ -22,7 +22,7 @@
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
# =============================================================================
ARG CUDA_VERSION=13.0.2
ARG CUDA_VERSION=13.0.3
ARG PYTHON_VERSION=3.12
ARG UBUNTU_VERSION=22.04
@@ -793,7 +793,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
# https://docs.flashinfer.ai/installation.html
# From versions.json: .flashinfer.version
ARG FLASHINFER_VERSION=0.6.14
ARG FLASHINFER_VERSION=0.6.15.post1
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+53 -31
View File
@@ -339,18 +339,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
# RIXL/UCX build stages
FROM base AS build_rixl
ARG RIXL_BRANCH="39be1de8"
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
ARG UCX_BRANCH="bfb51733"
# NIXL/UCX build stages
FROM base AS build_nixl
ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ENV ROCM_PATH=/opt/rocm
ENV UCX_HOME=/usr/local/ucx
ENV RIXL_HOME=/usr/local/rixl
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
ENV NIXL_HOME=/usr/local/nixl
# RIXL build system dependences and RDMA support
# NIXL build system dependencies and RDMA support
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
libgrpc-dev \
libgrpc++-dev \
@@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
&& rm -rf /var/lib/apt/lists/*
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system meson auditwheel patchelf tomlkit
uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
RUN --mount=type=cache,target=/root/.cache/ccache \
cd /usr/local/src && \
@@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
RUN --mount=type=cache,target=/root/.cache/ccache \
git clone ${RIXL_REPO} /opt/rixl && \
cd /opt/rixl && \
git checkout ${RIXL_BRANCH} && \
git clone ${NIXL_REPO} /opt/nixl && \
cd /opt/nixl && \
git checkout ${NIXL_BRANCH} && \
CC="ccache gcc" CXX="ccache g++" \
meson setup build --prefix=${RIXL_HOME} \
meson setup build --prefix=${NIXL_HOME} \
-Ducx_path=${UCX_HOME} \
-Drocm_path=${ROCM_PATH} && \
-Dwheel_variant=rocm \
-Dbuild_tests=false \
-Dbuild_examples=false && \
cd build && \
ninja -j$(nproc) && \
ninja install
ninja install && \
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
> /etc/ld.so.conf.d/nixl.conf && \
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
>> /etc/ld.so.conf.d/nixl.conf && \
ldconfig
# Generate RIXL wheel
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
# so configure the ROCm wheel variant directly through Meson.
# Exclude libcore and libpull from auditwheel: transitive dependencies
# that are not shipped in the wheel and vary across base images.
RUN cd /opt/rixl && \
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
contrib/build-wheel.sh && \
mkdir -p /app/install && \
_ucx_install_dir=${UCX_HOME} \
./contrib/build-wheel.sh \
--output-dir /app/install \
--rocm-dir ${ROCM_PATH} \
RUN cd /opt/nixl && \
./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
CC="ccache gcc" CXX="ccache g++" \
uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
--python ${PYTHON_VERSION} \
-Csetup-args=-Ducx_path=${UCX_HOME} \
-Csetup-args=-Dwheel_variant=rocm \
-Csetup-args=-Dbuild_tests=false \
-Csetup-args=-Dbuild_examples=false && \
mkdir -p /tmp/nixl_wheels/repaired /app/install && \
auditwheel repair \
--exclude 'libamdhip64*' \
--exclude 'libcore*' \
--exclude 'libpull*' \
/tmp/nixl_wheels/nixl_rocm*.whl \
--plat manylinux_2_34_$(uname -m) \
--wheel-dir /tmp/nixl_wheels/repaired && \
./contrib/wheel_add_ucx_plugins.py \
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
/tmp/nixl_wheels/repaired/*.whl && \
cp /tmp/nixl_wheels/repaired/*.whl /app/install
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
# invalidate the slow ROCShmem build.
@@ -660,10 +680,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
fi
# Install RIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
# Install NIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
# Copy ROCShmem runtime libraries.
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
@@ -724,6 +744,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
# ROCm profiler limits workaround.
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -796,9 +817,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
&& pip uninstall -y vllm \
&& uv pip install --system *.whl
# Install RIXL wheel
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
uv pip install --system /rixl_install/*.whl
# Install NIXL ROCm wheel
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
uv pip install --system /nixl_install/*.whl
ARG COMMON_WORKDIR
ARG BASE_IMAGE
@@ -813,6 +834,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV TOKENIZERS_PARALLELISM=false
+1 -1
View File
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
ARG FA_BRANCH="0e60e394"
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
ARG AITER_BRANCH="v0.1.16.post3"
ARG AITER_BRANCH="v0.1.16.post5"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
ARG MORI_BRANCH="v1.1.0"
ARG MORI_REPO="https://github.com/ROCm/mori.git"
+25 -1
View File
@@ -86,6 +86,29 @@ RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p /tmp/hf-xet/dist && \
cp dist/*.whl /tmp/hf-xet/dist/
# Build LLVM 20 from source for llvmlite (system repos ship LLVM 21 which
# llvmlite v0.47 does not support; only SystemZ target is needed).
FROM base AS llvm20-build
ARG LLVM_VERSION=20.1.8
WORKDIR /tmp
RUN microdnf install -y ninja-build gcc gcc-c++ python3 xz && \
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}/llvm-project-${LLVM_VERSION}.src.tar.xz && \
tar -xf llvm-project-${LLVM_VERSION}.src.tar.xz && \
cmake -G Ninja -S llvm-project-${LLVM_VERSION}.src/llvm -B build \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
-DLLVM_ENABLE_RTTI=ON \
-DLLVM_BUILD_TOOLS=OFF \
-DLLVM_BUILD_UTILS=ON \
-DLLVM_BUILD_EXAMPLES=OFF \
-DLLVM_BUILD_TESTS=OFF \
-DLLVM_INCLUDE_TESTS=OFF \
-DLLVM_INCLUDE_EXAMPLES=OFF \
-DLLVM_INCLUDE_BENCHMARKS=OFF && \
ninja -C build install && \
rm -rf build llvm-project-${LLVM_VERSION}.src*
# Build numba
FROM python-install AS numba-builder
@@ -96,11 +119,13 @@ WORKDIR /tmp
# Clone all required dependencies
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,from=llvm20-build,source=/opt/llvm20,target=/opt/llvm20 \
microdnf install ninja-build gcc gcc-c++ -y && \
git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
cd llvmlite && \
uv pip install 'cmake<4' 'setuptools<70' numpy && \
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
python setup.py bdist_wheel && \
cd ../numba && \
if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
@@ -158,7 +183,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
uv pip install -v \
$ARROW_WHL_FILE \
$VISION_WHL_FILE \
$HF_XET_WHL_FILE \
$LLVM_WHL_FILE \
+13 -13
View File
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
}
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
# stage and is irrelevant when building --target ci_base itself.
variable "CI_BASE_IMAGE" {
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
# functions produce no entries when the variable is empty.
variable "RIXL_BRANCH" {
variable "NIXL_BRANCH" {
default = ""
}
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
default = ""
}
variable "RIXL_CACHE_KEY" {
variable "NIXL_CACHE_KEY" {
default = ""
}
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
])
}
# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
# These stages are pinned to specific upstream commit hashes, so cache keys use
# those hashes rather than the Buildkite commit. This means the cache persists
# across all vLLM commits as long as the upstream dependency pins don't change.
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
function "get_cache_from_rocm_deps" {
params = []
result = compact([
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
])
}
function "get_cache_to_rocm_rixl" {
function "get_cache_to_rocm_nixl" {
params = []
result = compact([
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
])
}
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
# (e.g., requirements change), these stages are cache hits if their upstream
# pins haven't changed -- saving ~35min of compilation.
target "rixl-rocm-ci" {
target "nixl-rocm-ci" {
inherits = ["_common-rocm", "_ci-rocm"]
target = "build_rixl"
target = "build_nixl"
cache-from = get_cache_from_rocm_deps()
cache-to = get_cache_to_rocm_rixl()
cache-to = get_cache_to_rocm_nixl()
output = ["type=cacheonly"]
}
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
output = ["type=cacheonly"]
}
# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
# files drifts from the remote image label. Per-PR builds then pull the result
# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
]),
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
# Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
# are cache hits even when ci_base itself needs rebuilding.
get_cache_from_rocm_deps(),
)
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
# Group for ci_base builds -- exports dependency stage caches alongside the
# ci_base image so future rebuilds can reuse them independently.
group "ci-base-rocm-ci-with-deps" {
targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
}
+2 -2
View File
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
# overrides) and writes a small HCL override before invoking CI targets.
variable "RIXL_BRANCH" {
variable "NIXL_BRANCH" {
default = ""
}
@@ -106,7 +106,7 @@ target "test-rocm" {
output = ["type=docker"]
}
# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
# torchcodec, requirements, etc.). Used by the weekly scheduled build and
# the auto-rebuild trigger when requirements change in a PR.
target "ci-base-rocm" {
+4 -4
View File
@@ -2,7 +2,7 @@
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
"variable": {
"CUDA_VERSION": {
"default": "13.0.2"
"default": "13.0.3"
},
"PYTHON_VERSION": {
"default": "3.12"
@@ -11,10 +11,10 @@
"default": "22.04"
},
"BUILD_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
"default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
},
"FINAL_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
"default": "nvidia/cuda:13.0.3-base-ubuntu22.04"
},
"BUILD_OS": {
"default": "ubuntu"
@@ -68,7 +68,7 @@
"default": "true"
},
"FLASHINFER_VERSION": {
"default": "0.6.14"
"default": "0.6.15.post1"
},
"GDRCOPY_CUDA_VERSION": {
"default": "12.8"
+1 -5
View File
@@ -13,11 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files
For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
+1 -1
View File
@@ -315,7 +315,7 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
- vLLM CPU supports quantizations:
- AWQ (x86 only)
- GPTQ (x86 only)
- compressed-tensor INT8 W8A8 (x86, s390x)
- compressed-tensor INT8 W8A8 (x86 only)
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
@@ -11,7 +11,7 @@ Currently, the CPU implementation for s390x architecture supports FP32, BF16 and
- OS: `Linux`
- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`
--8<-- [end:requirements]
--8<-- [start:set-up-using-python]
@@ -28,13 +28,24 @@ Install the following packages from the package manager before building the vLLM
```bash
dnf install -y \
which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
which procps findutils tar vim git patch xz ninja-build \
gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \
clang llvm-devel llvm-static clang-devel
```
Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
Build and install `numactl` from source:
```bash
curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz
tar -xvzf v2.0.19.tar.gz
cd numactl-2.0.19
./autogen.sh && ./configure && make && make install
cd ..
```
Install rust>=1.80 which is needed for `outlines-core`, `uvloop`, and `hf-xet` python packages installation.
```bash
curl https://sh.rustup.rs -sSf | sh -s -- -y && \
@@ -44,26 +55,79 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
Execute the following commands to build and install vLLM from source.
!!! tip
Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
Pre-built wheels are not available for s390x for the following packages. Build them from source before building vLLM: `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`.
See `docker/Dockerfile.s390x` for exact versions and build commands used in each multi-stage build.
!!! note "LLVM 20 required for llvmlite"
`llvmlite v0.47` requires LLVM 20, but UBI 9.6 repos ship LLVM 21 which is
not compatible. You must build LLVM 20 from source before building `llvmlite`:
```bash
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-20.1.8/llvm-project-20.1.8.src.tar.xz
tar -xf llvm-project-20.1.8.src.tar.xz
cmake -G Ninja -S llvm-project-20.1.8.src/llvm -B llvm-build \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
-DLLVM_ENABLE_RTTI=ON \
-DLLVM_BUILD_TOOLS=OFF \
-DLLVM_BUILD_UTILS=ON \
-DLLVM_BUILD_EXAMPLES=OFF \
-DLLVM_BUILD_TESTS=OFF \
-DLLVM_INCLUDE_TESTS=OFF \
-DLLVM_INCLUDE_EXAMPLES=OFF \
-DLLVM_INCLUDE_BENCHMARKS=OFF
ninja -C llvm-build install
```
Then build `llvmlite` pointing to LLVM 20:
```bash
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
python setup.py bdist_wheel
```
```bash
uv pip install -v \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt \
--torch-backend cpu \
--index-strategy unsafe-best-match && \
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
uv pip install dist/*.whl
uv pip install -v \
/path/to/torchvision.whl \
/path/to/llvmlite.whl \
/path/to/numba.whl \
/path/to/opencv_python_headless.whl \
/path/to/hf_xet.whl \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt \
--torch-backend cpu \
--index-strategy unsafe-best-match && \
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
uv pip install dist/*.whl
```
??? console "pip"
```bash
pip install -v \
--extra-index-url https://download.pytorch.org/whl/cpu \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt \
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
pip install dist/*.whl
pip install -v \
--extra-index-url https://download.pytorch.org/whl/cpu \
/path/to/torchvision.whl \
/path/to/llvmlite.whl \
/path/to/numba.whl \
/path/to/opencv_python_headless.whl \
/path/to/hf_xet.whl \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt && \
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
pip install dist/*.whl
```
!!! warning "Protobuf workaround for s390x"
The C++ protobuf extension crashes on s390x. After installation, set the
following environment variable and remove the C++ extensions:
```bash
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
# Remove C++ protobuf extensions that crash on s390x
SITE_PKGS=$(python -c "import site; print(site.getsitepackages()[0])")
rm -rf "$SITE_PKGS/google/_upb/"*.so \
"$SITE_PKGS/google/protobuf/pyext/"*.so 2>/dev/null || true
```
--8<-- [end:build-wheel-from-source]
@@ -80,19 +144,20 @@ docker build -f docker/Dockerfile.s390x \
# Launch OpenAI server
docker run --rm \
--privileged true \
--security-opt seccomp=unconfined \
--cap-add SYS_NICE \
--shm-size 4g \
-p 8000:8000 \
-e VLLM_CPU_KVCACHE_SPACE=<KV cache space> \
-e VLLM_CPU_OMP_THREADS_BIND=<CPU cores for inference> \
vllm-cpu-env \
--model meta-llama/Llama-3.2-1B-Instruct \
--dtype float \
--dtype bfloat16 \
other vLLM OpenAI server arguments
```
!!! tip
An alternative of `--privileged true` is `--cap-add SYS_NICE --security-opt seccomp=unconfined`.
Alternatively, `--privileged=true` also works but is broader and not generally recommended.
--8<-- [end:build-image-from-source]
--8<-- [start:extra-information]
+3
View File
@@ -0,0 +1,3 @@
// Reo.Dev documentation tracking
// https://docs.reo.dev/integrations/input-sources/developer-insights/documentation
!function(){var e,t,n;e="d5c4337961ef0ac",t=function(){Reo.init({clientID:"d5c4337961ef0ac", enableThirdPartyTracking: true})},(n=document.createElement("script")).src="https://static.reo.dev/"+e+"/reo.js",n.defer=!0,n.onload=t,document.head.appendChild(n)}();
+1
View File
@@ -160,3 +160,4 @@ extra_javascript:
- https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
- mkdocs/javascript/edit_and_feedback.js
- mkdocs/javascript/slack_and_forum.js
- mkdocs/javascript/reo.js
+1 -1
View File
@@ -7,7 +7,7 @@ requires = [
"setuptools>=77.0.3,<81.0.0",
"setuptools-scm>=8.0",
"setuptools-rust>=1.9.0",
"torch == 2.11.0",
"torch == 2.13.0",
"wheel",
"jinja2",
]
+2 -2
View File
@@ -4,8 +4,8 @@ packaging>=24.2
setuptools==77.0.3 # this version can reuse CMake build dir
setuptools-scm>=8
setuptools-rust>=1.9.0
torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
wheel
jinja2>=3.1.6
regex
+1 -1
View File
@@ -5,7 +5,7 @@ packaging>=24.2
setuptools>=77.0.3,<81.0.0
setuptools-scm>=8
setuptools-rust>=1.9.0
torch==2.11.0
torch==2.13.0
wheel
jinja2>=3.1.6
regex
+2 -2
View File
@@ -6,8 +6,8 @@ setuptools==77.0.3 # this version can reuse CMake build dir
numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding
# Dependencies for CPUs
torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
# required for the image processor of minicpm-o-2_6, this must be updated alongside torch
torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"
+5 -5
View File
@@ -4,18 +4,18 @@
numba == 0.65.0 # Required for N-gram speculative decoding
# Dependencies for NVIDIA GPUs
torch==2.11.0
torch==2.13.0
torchaudio==2.11.0
# These must be updated alongside torch
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
torchcodec >= 0.14
PyNvVideoCodec==2.0.4
# FlashInfer should be updated together with the Dockerfile
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
# install_requires so the published wheel does not carry an unresolvable pin
--extra-index-url https://flashinfer.ai/whl/
flashinfer-python==0.6.14
flashinfer-cubin==0.6.14
flashinfer-python==0.6.15.post1
flashinfer-cubin==0.6.15.post1
apache-tvm-ffi==0.1.10
tilelang==0.1.9
nvidia-cudnn-frontend>=1.19.1
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
nvidia-cutlass-dsl[cu13]==4.6.0
quack-kernels>=0.4.0 # Required for tml-fa4
quack-kernels>=0.6.1 # Required for CUTLASS DSL 4.6 by MSA
# Tokenspeed_MLA for faster mla with spec decode
tokenspeed-mla==0.1.8; platform_system == "Linux"
+2 -2
View File
@@ -1107,7 +1107,7 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/cuda.in
# transformers
torch==2.11.0+cpu
torch==2.13.0+cpu
# via
# -r requirements/test/cuda.in
# accelerate
@@ -1134,7 +1134,7 @@ torchaudio==2.11.0+cpu
# vocos
torchcodec==0.14.0+cpu
# via -r requirements/test/cuda.in
torchvision==0.26.0+cpu
torchvision==0.28.0+cpu
# via
# -r requirements/test/cuda.in
# open-clip-torch
+2 -2
View File
@@ -28,9 +28,9 @@ soundfile # required for audio tests
jiwer # required for audio tests
tblib # for pickling test exceptions
timm >=1.0.17 # required for internvl and gemma3n-mm test
torch==2.11.0
torch==2.13.0
torchaudio==2.11.0
torchvision==0.26.0
torchvision==0.28.0
transformers_stream_generator # required for qwen-vl test
matplotlib # required for qwen-vl test
mistral_common[image,audio] >= 1.11.5 # required for voxtral test
+11 -9
View File
@@ -159,7 +159,7 @@ cuda-bindings==13.0.3
# via torch
cuda-pathfinder==1.3.3
# via cuda-bindings
cuda-toolkit==13.0.2
cuda-toolkit==13.0.3.0
# via torch
cupy-cuda12x==13.6.0
# via ray
@@ -599,7 +599,7 @@ numpy==2.2.6
# tritonclient
# vocos
# xgrammar
nvidia-cublas==13.1.0.3
nvidia-cublas==13.1.1.3
# via
# cuda-toolkit
# nvidia-cudnn-cu13
@@ -607,10 +607,12 @@ nvidia-cublas==13.1.0.3
nvidia-cuda-cupti==13.0.85
# via cuda-toolkit
nvidia-cuda-nvrtc==13.0.88
# via cuda-toolkit
# via
# cuda-toolkit
# nvidia-cublas
nvidia-cuda-runtime==13.0.96
# via cuda-toolkit
nvidia-cudnn-cu13==9.19.0.56
nvidia-cudnn-cu13==9.20.0.48
# via torch
nvidia-cufft==12.0.0.61
# via cuda-toolkit
@@ -624,9 +626,9 @@ nvidia-cusparse==12.6.3.3
# via
# cuda-toolkit
# nvidia-cusolver
nvidia-cusparselt-cu13==0.8.0
nvidia-cusparselt-cu13==0.8.1
# via torch
nvidia-nccl-cu13==2.28.9
nvidia-nccl-cu13==2.29.7
# via torch
nvidia-nvjitlink==13.0.88
# via
@@ -1202,7 +1204,7 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/cuda.in
# transformers
torch==2.11.0+cu130
torch==2.13.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
@@ -1233,7 +1235,7 @@ torchcodec==0.14.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
torchvision==0.26.0+cu130
torchvision==0.28.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
@@ -1270,7 +1272,7 @@ transformers==5.13.1
# xgrammar
transformers-stream-generator==0.0.5
# via -r requirements/test/cuda.in
triton==3.6.0
triton==3.7.1
# via
# torch
# xgrammar
+1 -1
View File
@@ -12,4 +12,4 @@ ray[data]
setuptools==78.1.0
setuptools-rust>=1.9.0
nixl==0.3.0
tpu-inference==0.24.0
tpu-inference==0.25.0
+13
View File
@@ -14,6 +14,10 @@ service Generate {
rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
}
service Control {
rpc Abort (AbortRequest) returns (AbortResponse) {}
}
// ======================================================================================
// Generate Request
// ======================================================================================
@@ -201,3 +205,12 @@ message TokenIds {
repeated uint32 ids = 1;
}
// ======================================================================================
// Control
// ======================================================================================
message AbortRequest {
repeated string request_ids = 1;
}
message AbortResponse {}
+106 -61
View File
@@ -54,6 +54,7 @@ mod stream;
use vllm_engine_core_client::EngineCoreClient;
use vllm_engine_core_client::protocol::dtype::ModelDtype;
use vllm_engine_core_client::protocol::multimodal::MmFeatures;
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
use vllm_llm::Llm;
use vllm_text::{Prompt, TextLlm, TextRequest};
@@ -88,6 +89,92 @@ pub fn validate_parser_overrides(
Ok(())
}
/// Chat request preparation shared by inference and render-only frontends.
pub struct ChatRequestProcessor {
backend: DynChatBackend,
/// Effective model dtype reported by the engine.
/// Absent for text-only frontends without an engine handshake.
model_dtype: Option<ModelDtype>,
}
impl ChatRequestProcessor {
/// Create a processor with multimodal support using the effective model
/// dtype reported by the engine.
fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self {
Self {
backend,
model_dtype: Some(model_dtype),
}
}
/// Create a render-only processor that rejects multimodal requests.
pub fn render_only(backend: DynChatBackend) -> Self {
Self {
backend,
model_dtype: None,
}
}
async fn finalize_rendered_prompt(
&self,
request: &ChatRequest,
rendered: RenderedPrompt,
) -> Result<(Prompt, Option<MmFeatures>)> {
match self.model_dtype {
Some(model_dtype) => {
multimodal::finalize_rendered_prompt(
request,
rendered,
self.backend.multimodal_model_info(),
model_dtype,
)
.await
}
None if !request.has_multimodal() => Ok((rendered.prompt, None)),
None => Err(Error::UnsupportedMultimodalRenderer),
}
}
/// Prepare one chat request without submitting it to an engine.
pub async fn prepare(
&self,
mut request: ChatRequest,
options: NewChatOutputProcessorOptions<'_>,
) -> Result<(TextRequest, DynChatOutputProcessor)> {
request.validate()?;
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
let arrival_time = vllm_llm::current_unix_timestamp_secs();
let output_processor = self.backend.new_chat_output_processor(&mut request, options)?;
let rendered = self.backend.chat_renderer().render(&request)?;
let reasoning_parser_kwargs =
request
.sampling_params
.structured_outputs
.is_some()
.then(|| ReasoningParserKwargs {
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
});
let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
let text_request = TextRequest {
request_id: request.request_id,
prompt,
mm_features,
sampling_params: request.sampling_params,
decode_options: request.decode_options,
intermediate: request.intermediate,
priority: request.priority,
cache_salt: request.cache_salt,
add_special_tokens: request.add_special_tokens,
data_parallel_rank: request.data_parallel_rank,
reasoning_parser_kwargs,
lora_request: request.lora_request,
arrival_time: Some(arrival_time),
};
Ok((text_request, output_processor))
}
}
/// Structured chat facade above [`TextLlm`].
///
/// This layer stays above raw text semantics: it takes care of chat-template
@@ -95,9 +182,7 @@ pub fn validate_parser_overrides(
/// request semantics such as tool calls.
pub struct ChatLlm {
text: TextLlm,
backend: DynChatBackend,
/// Effective model dtype reported by the engine.
model_dtype: ModelDtype,
processor: ChatRequestProcessor,
/// Tool-call parser selection.
tool_call_parser: ParserSelection,
/// Reasoning parser selection.
@@ -112,8 +197,7 @@ impl ChatLlm {
Self {
text,
backend,
model_dtype,
processor: ChatRequestProcessor::new(backend, model_dtype),
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
}
@@ -140,7 +224,7 @@ impl ChatLlm {
/// Override the effective model dtype used for multimodal tensor encoding.
pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
self.model_dtype = model_dtype;
self.processor.model_dtype = Some(model_dtype);
self
}
@@ -172,57 +256,23 @@ impl ChatLlm {
}
/// Render, tokenize, and submit one chat request.
pub async fn chat(&self, mut request: ChatRequest) -> Result<ChatEventStream> {
request.validate()?;
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
let arrival_time = vllm_llm::current_unix_timestamp_secs();
let output_processor = self.backend.new_chat_output_processor(
&mut request,
NewChatOutputProcessorOptions {
tool_call_parser: &self.tool_call_parser,
reasoning_parser: &self.reasoning_parser,
},
)?;
let rendered = self.backend.chat_renderer().render(&request)?;
let reasoning_parser_kwargs =
request
.sampling_params
.structured_outputs
.is_some()
.then(|| ReasoningParserKwargs {
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
});
let (prompt, mm_features) = multimodal::finalize_rendered_prompt(
&request,
rendered,
self.backend.multimodal_model_info(),
self.model_dtype,
)
.await?;
let text_request = TextRequest {
request_id: request.request_id.clone(),
prompt,
mm_features,
sampling_params: request.sampling_params,
decode_options: request.decode_options,
intermediate: request.intermediate,
priority: request.priority,
cache_salt: request.cache_salt,
add_special_tokens: request.add_special_tokens,
data_parallel_rank: request.data_parallel_rank,
reasoning_parser_kwargs,
lora_request: request.lora_request,
arrival_time: Some(arrival_time),
};
pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
let (text_request, output_processor) = self
.processor
.prepare(
request,
NewChatOutputProcessorOptions {
tool_call_parser: &self.tool_call_parser,
reasoning_parser: &self.reasoning_parser,
},
)
.await?;
let request_id = text_request.request_id.clone();
let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
let structured_stream = output_processor.process(decoded_stream)?;
Ok(ChatEventStream::new(request.request_id, structured_stream))
Ok(ChatEventStream::new(request_id, structured_stream))
}
/// Render through the chat template and tokenize, without submitting to the engine.
@@ -233,14 +283,9 @@ impl ChatLlm {
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
request.validate()?;
let rendered = self.backend.chat_renderer().render(&request)?;
let (prompt, _mm_features) = multimodal::finalize_rendered_prompt(
&request,
rendered,
self.backend.multimodal_model_info(),
self.model_dtype,
)
.await?;
let rendered = self.processor.backend.chat_renderer().render(&request)?;
let (prompt, _mm_features) =
self.processor.finalize_rendered_prompt(&request, rendered).await?;
let tokenizer = self.text.tokenizer();
let token_ids = match prompt {
+6 -11
View File
@@ -8,7 +8,7 @@ use tonic_health::ServingStatus;
use tonic_health::server::HealthReporter;
use tracing::{info, warn};
use super::GenerateGrpcService;
use super::{ControlGrpcService, GenerateGrpcService};
pub(crate) async fn monitor_health(
mut health_reporter: HealthReporter,
@@ -16,21 +16,18 @@ pub(crate) async fn monitor_health(
shutdown: CancellationToken,
) {
let generate_service = GenerateGrpcService::NAME;
let control_service = ControlGrpcService::NAME;
let status = ServingStatus::NotServing;
let health_event_first = tokio::select! {
result = engine_health.wait_for(|healthy| !*healthy) => {
match result {
Ok(_) => warn!(
generate_service,
overall_service = true,
status = ?status,
reason = "engine_unhealthy",
"marking gRPC health services as not serving"
),
Err(error) => warn!(
%error,
generate_service,
overall_service = true,
status = ?status,
reason = "health_channel_closed",
"engine health channel closed; marking gRPC health services as not serving"
@@ -40,8 +37,6 @@ pub(crate) async fn monitor_health(
}
_ = shutdown.cancelled() => {
info!(
generate_service,
overall_service = true,
status = ?status,
reason = "server_shutdown",
"server shutting down; marking gRPC health services as not serving"
@@ -51,20 +46,20 @@ pub(crate) async fn monitor_health(
};
health_reporter.set_not_serving::<GenerateGrpcService>().await;
// Generate is currently the only engine-backed gRPC service, so overall
// server health intentionally mirrors it.
health_reporter.set_not_serving::<ControlGrpcService>().await;
// Both gRPC services use the same engine client, so overall server health
// mirrors their shared engine health.
health_reporter.set_service_status("", status).await;
if health_event_first {
shutdown.cancelled().await;
info!(
generate_service,
overall_service = true,
reason = "server_shutdown",
"server shutting down; closing gRPC health watches"
);
}
health_reporter.clear_service_status(generate_service).await;
health_reporter.clear_service_status(control_service).await;
health_reporter.clear_service_status("").await;
}
+32
View File
@@ -26,8 +26,10 @@ pub mod pb {
}
pub(crate) use health::monitor_health;
pub use pb::control_server::ControlServer;
pub use pb::generate_server::GenerateServer;
pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
#[cfg(test)]
@@ -44,6 +46,36 @@ impl GenerateServiceImpl {
}
}
/// gRPC control service backed by the shared application state.
pub struct ControlServiceImpl {
state: Arc<AppState>,
}
impl ControlServiceImpl {
pub fn new(state: Arc<AppState>) -> Self {
Self { state }
}
}
#[tonic::async_trait]
impl pb::control_server::Control for ControlServiceImpl {
async fn abort(
&self,
request: Request<pb::AbortRequest>,
) -> Result<Response<pb::AbortResponse>, Status> {
let request_ids = request.into_inner().request_ids;
if request_ids.is_empty() {
return Ok(Response::new(pb::AbortResponse {}));
}
self.state
.chat
.abort(&request_ids)
.await
.map_err(|error| Status::internal(error.to_report_string()))?;
Ok(Response::new(pb::AbortResponse {}))
}
}
#[tonic::async_trait]
impl pb::generate_server::Generate for GenerateServiceImpl {
type GenerateStreamStream =
+122 -18
View File
@@ -38,8 +38,9 @@ use vllm_tokenizer::test_utils::TestTokenizer;
use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
use super::pb::control_client::ControlClient;
use super::pb::generate_client::GenerateClient;
use super::{GenerateServer, GenerateServiceImpl, pb};
use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb};
use crate::listener::{Listener, MaybeTlsListener};
use crate::state::AppState;
use crate::tls;
@@ -153,10 +154,6 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
dealer.recv().await.expect("recv engine message").into_vec()
}
fn test_llm(client: EngineCoreClient) -> Llm {
Llm::new(client).with_request_id_randomization(false)
}
#[derive(Clone, Debug)]
struct FakeTextBackend;
@@ -206,6 +203,7 @@ async fn setup_grpc_service(
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> (
GenerateServer<GenerateServiceImpl>,
ControlServer<ControlServiceImpl>,
tokio::sync::watch::Receiver<bool>,
MockEngineTask,
) {
@@ -243,12 +241,13 @@ async fn setup_grpc_service(
let engine_health = client.subscribe_health();
let chat = ChatLlm::from_shared_backend(
test_llm(client),
Llm::new(client),
Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
);
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
(
GenerateServer::new(GenerateServiceImpl::new(state)),
GenerateServer::new(GenerateServiceImpl::new(state.clone())),
ControlServer::new(ControlServiceImpl::new(state)),
engine_health,
engine_task,
)
@@ -264,9 +263,11 @@ async fn grpc_test_server(
tokio::task::JoinHandle<()>,
MockEngineTask,
) {
let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (generate_service, control_service, engine_health, engine_task) =
setup_grpc_service(engine_id, output_specs).await;
let (channel, server_task) = start_grpc_test_server(
svc,
generate_service,
control_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
@@ -276,11 +277,13 @@ async fn grpc_test_server(
async fn start_grpc_test_server(
generate_service: GenerateServer<GenerateServiceImpl>,
control_service: ControlServer<ControlServiceImpl>,
engine_health: tokio::sync::watch::Receiver<bool>,
shutdown: tokio_util::sync::CancellationToken,
) -> (Channel, tokio::task::JoinHandle<()>) {
let (health_reporter, health_service) = health_reporter();
health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr");
@@ -289,6 +292,7 @@ async fn start_grpc_test_server(
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
let server = TonicServer::builder()
.add_service(health_service)
.add_service(control_service)
.add_service(generate_service)
.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
let health_monitor =
@@ -319,7 +323,8 @@ async fn grpc_tls_test_server(
certs: &TestCerts,
cert_reqs: i32,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (generate_service, control_service, _engine_health, engine_task) =
setup_grpc_service(engine_id, output_specs).await;
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
.expect("build grpc tls config");
@@ -329,7 +334,8 @@ async fn grpc_tls_test_server(
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
TonicServer::builder()
.add_service(svc)
.add_service(control_service)
.add_service(generate_service)
.serve_with_incoming(incoming)
.await
.expect("grpc tls server");
@@ -409,7 +415,7 @@ async fn grpc_server_with_keepalive(
engine_id: impl Into<EngineId>,
keepalive: Option<Duration>,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, _engine_health, engine_task) =
let (generate_service, control_service, _engine_health, engine_task) =
setup_grpc_service(engine_id, default_stream_output_specs()).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
@@ -425,7 +431,8 @@ async fn grpc_server_with_keepalive(
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
builder
.add_service(svc)
.add_service(control_service)
.add_service(generate_service)
.serve_with_incoming(incoming)
.await
.expect("grpc server");
@@ -1073,14 +1080,106 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn control_abort_resolves_external_id_and_empty_is_noop() {
let (generate_service, control_service, engine_health, engine_task) =
setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await;
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
let mut generate_client = GenerateClient::new(channel.clone());
let mut control_client = ControlClient::new(channel);
let request_id = "test-abort-active";
let mut stream = generate_client
.generate_stream(pb::GenerateRequest {
request_id: request_id.to_string(),
model: "test-model".to_string(),
prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
stopping: Some(pb::StoppingCriteria {
max_new_tokens: 10,
..Default::default()
}),
..Default::default()
})
.await
.expect("start generation")
.into_inner();
loop {
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for active generation output")
.expect("read active generation output")
.expect("generation ended before producing output");
if let Some(output) = response.outputs {
assert!(
output.finish_info.is_none(),
"generation finished before abort behavior was exercised"
);
break;
}
}
control_client
.abort(pb::AbortRequest::default())
.await
.expect("empty abort should be a no-op");
assert!(
tokio::time::timeout(Duration::from_millis(100), stream.message())
.await
.is_err(),
"empty abort unexpectedly ended the active generation"
);
control_client
.abort(pb::AbortRequest {
request_ids: vec![
request_id.to_string(),
request_id.to_string(),
"unknown".to_string(),
],
})
.await
.expect("abort active generation");
let finish_reason = loop {
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for aborted generation")
.expect("read aborted generation")
.expect("generation ended without an aborted response");
if let Some(finish_info) = response.outputs.and_then(|output| output.finish_info) {
break finish_info.finish_reason;
}
};
assert_eq!(finish_reason, pb::finish_info::FinishReason::Aborted as i32);
control_client
.abort(pb::AbortRequest {
request_ids: vec![request_id.to_string()],
})
.await
.expect("repeated abort should be idempotent");
engine_task.await.expect("mock engine task");
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
let (generate_service, _connected_engine_health, _engine_task) =
let (generate_service, control_service, _connected_engine_health, _engine_task) =
setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
@@ -1088,7 +1187,7 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
let mut health_client = HealthClient::new(channel);
let mut health_streams = Vec::new();
for service in ["vllm.Generate", ""] {
for service in ["vllm.Generate", "vllm.Control", ""] {
let service_label = if service.is_empty() {
"overall"
} else {
@@ -1143,14 +1242,19 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_watch_closes_on_graceful_shutdown() {
let (generate_service, engine_health, _engine_task) = setup_grpc_service(
let (generate_service, control_service, engine_health, _engine_task) = setup_grpc_service(
b"engine-grpc-health-shutdown",
default_stream_output_specs(),
)
.await;
let shutdown = tokio_util::sync::CancellationToken::new();
let (channel, server_task) =
start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
shutdown.clone(),
)
.await;
let mut health_client = HealthClient::new(channel);
let mut stream = health_client
.watch(HealthCheckRequest {
+4
View File
@@ -207,6 +207,9 @@ where
let (health_reporter, health_service) = health_reporter();
let engine_health = state.engine_core_client().subscribe_health();
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
health_reporter.set_serving::<grpc::ControlGrpcService>().await;
let control_service =
grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()));
let generate_service =
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
let svc = TonicServer::builder()
@@ -214,6 +217,7 @@ where
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
.layer(middleware::request_runtime_layer(state.clone()))
.add_service(health_service)
.add_service(control_service)
.add_service(generate_service);
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
+108 -72
View File
@@ -38,32 +38,22 @@ trait_set! {
pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
}
/// Raw text facade above [`Llm`].
///
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
/// in, decoded text deltas and terminal metadata flow out.
pub struct TextLlm {
/// Generate-only client owned by this text facade.
llm: Llm,
/// Text request preparation shared by inference and render-only frontends.
pub struct TextRequestProcessor {
/// Tokenizer/model metadata backend responsible for prompt encode/decode
/// and sampling hints.
backend: DynTextBackend,
/// Runtime context window size reported by the engine startup handshake.
/// Render-only frontends supply the downstream engine's effective value.
max_model_len: u32,
/// Maximum number of top log probabilities accepted by this text facade.
max_logprobs: i32,
}
impl TextLlm {
/// Create a new text-generation facade from a shared LLM client plus a text
/// backend.
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
// The engine-reported value reflects the post-profiling, auto-fitted
// KV cache limit used at runtime.
let max_model_len = llm.engine_core_client().max_model_len();
impl TextRequestProcessor {
/// Create a processor with the effective model context length.
pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self {
Self {
llm,
backend,
max_model_len,
max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
@@ -78,61 +68,18 @@ impl TextLlm {
self
}
/// Return the backend model ID.
pub fn model_id(&self) -> &str {
self.backend.model_id()
}
/// Expose the underlying engine-core client for low-level utility/admin
/// calls.
pub fn engine_core_client(&self) -> &EngineCoreClient {
self.llm.engine_core_client()
}
/// Return the tokenizer used by this text backend.
/// Return the tokenizer used by this processor.
pub fn tokenizer(&self) -> DynTokenizer {
self.backend.tokenizer()
}
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
pub fn tokenizer_vocab_size(&self) -> usize {
self.backend.tokenizer_vocab_size()
/// Return the effective model context length.
pub fn max_model_len(&self) -> u32 {
self.max_model_len
}
/// Model vocabulary size from the model config, used to bound generated
/// token IDs and logits-domain sampling controls.
pub fn model_vocab_size(&self) -> usize {
self.backend.model_vocab_size()
}
/// Tokenize if needed, lower to a generate request, and return the raw
/// token stream.
pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
let (_, raw_stream) = self.generate_inner(request).await?;
Ok(raw_stream)
}
/// Tokenize if needed, lower to a generate request, and stream
/// incrementally decoded text.
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
let (text_request, raw_stream) = self.generate_inner(request).await?;
let tokenizer = self.backend.tokenizer();
let decoded_stream = output::decoded_text_event_stream(
text_request.request_id,
tokenizer,
raw_stream,
text_request.decode_options,
text_request.intermediate,
);
Ok(decoded_stream)
}
async fn generate_inner(
&self,
mut request: TextRequest,
) -> Result<(TextRequest, GenerateOutputStream)> {
/// Tokenize and lower one request without submitting it to an engine.
pub fn prepare(&self, mut request: TextRequest) -> Result<PreparedTextRequest> {
request.validate()?;
if request.arrival_time.is_none() {
@@ -146,7 +93,6 @@ impl TextLlm {
// and infra workloads bypass chat rendering and tokenizer overhead entirely.
Prompt::TokenIds(token_ids) => token_ids,
};
let sampling_hints = self.backend.sampling_hints()?;
let sampling_limits = SamplingLimits {
max_model_len: self.max_model_len,
@@ -155,16 +101,106 @@ impl TextLlm {
tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
};
let PreparedTextRequest {
text_request,
generate_request,
} = lower_text_request(
lower_text_request(
request,
prompt_token_ids,
sampling_hints,
sampling_limits,
&*tokenizer,
)?;
tokenizer.as_ref(),
)
}
}
/// Raw text facade above [`Llm`].
///
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
/// in, decoded text deltas and terminal metadata flow out.
pub struct TextLlm {
/// Generate-only client owned by this text facade.
llm: Llm,
/// Shared engine-free request preparation.
processor: TextRequestProcessor,
}
impl TextLlm {
/// Create a new text-generation facade from a shared LLM client plus a text
/// backend.
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
// The engine-reported value reflects the post-profiling, auto-fitted
// KV cache limit used at runtime.
let max_model_len = llm.engine_core_client().max_model_len();
Self {
llm,
processor: TextRequestProcessor::new(backend, max_model_len),
}
}
/// Override the maximum accepted logprobs count.
pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
self.processor = self.processor.with_max_logprobs(max_logprobs);
self
}
/// Return the backend model ID.
pub fn model_id(&self) -> &str {
self.processor.backend.model_id()
}
/// Expose the underlying engine-core client for low-level utility/admin
/// calls.
pub fn engine_core_client(&self) -> &EngineCoreClient {
self.llm.engine_core_client()
}
/// Return the tokenizer used by this text backend.
pub fn tokenizer(&self) -> DynTokenizer {
self.processor.tokenizer()
}
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
pub fn tokenizer_vocab_size(&self) -> usize {
self.processor.backend.tokenizer_vocab_size()
}
/// Model vocabulary size from the model config, used to bound generated
/// token IDs and logits-domain sampling controls.
pub fn model_vocab_size(&self) -> usize {
self.processor.backend.model_vocab_size()
}
/// Tokenize if needed, lower to a generate request, and return the raw
/// token stream.
pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
let (_, raw_stream) = self.generate_inner(request).await?;
Ok(raw_stream)
}
/// Tokenize if needed, lower to a generate request, and stream
/// incrementally decoded text.
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
let (text_request, raw_stream) = self.generate_inner(request).await?;
let tokenizer = self.processor.tokenizer();
let decoded_stream = output::decoded_text_event_stream(
text_request.request_id,
tokenizer,
raw_stream,
text_request.decode_options,
text_request.intermediate,
);
Ok(decoded_stream)
}
async fn generate_inner(
&self,
request: TextRequest,
) -> Result<(TextRequest, GenerateOutputStream)> {
let PreparedTextRequest {
text_request,
generate_request,
} = self.processor.prepare(request)?;
let raw_stream = self.llm.generate(generate_request).await?;
Ok((text_request, raw_stream))
+62
View File
@@ -214,3 +214,65 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs():
base_hash = CacheConfig().compute_hash()
assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash
assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash
def test_envs_compile_factors_relocation_invariant(tmp_path):
"""Relocating HOME or the XDG roots must not change the compile-cache
env hash.
Location-derived env vars (VLLM_XLA_CACHE_PATH from XDG_CACHE_HOME,
VLLM_CONFIG_ROOT from XDG_CONFIG_HOME/HOME) carry no information about
compiled artifacts, only about where directories live. When they leak
into compile_factors(), a cache produced under one HOME/XDG layout
silently misses under another - which defeats copying or pre-baking a
compile cache into a container image.
"""
import os
import subprocess
import sys
code = """
import sys
import logging
logging.disable(logging.CRITICAL)
from vllm import envs
from vllm.config.utils import hash_factors
print(hash_factors(envs.compile_factors()))
"""
def hash_with(extra_env):
env = {**dict(os.environ), "VLLM_LOGGING_LEVEL": "ERROR"}
# Drop explicit overrides so the derived defaults are what is
# exercised, then apply the relocation under test.
for key in ("VLLM_XLA_CACHE_PATH", "VLLM_CONFIG_ROOT", "VLLM_CACHE_ROOT"):
env.pop(key, None)
env.update(extra_env)
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
check=True,
env=env,
)
return result.stdout.strip()
xdg_cache = tmp_path / "relocated-xdg-cache"
xdg_config = tmp_path / "relocated-xdg-config"
new_home = tmp_path / "relocated-home"
for d in (xdg_cache, xdg_config, new_home):
d.mkdir()
base = hash_with({})
relocated_xdg = hash_with(
{"XDG_CACHE_HOME": str(xdg_cache), "XDG_CONFIG_HOME": str(xdg_config)}
)
relocated_home = hash_with({"HOME": str(new_home)})
assert relocated_xdg == base, (
"XDG_CACHE_HOME/XDG_CONFIG_HOME relocation changed the compile-cache "
"env hash - a location-only derived var is leaking into the key"
)
assert relocated_home == base, (
"HOME relocation changed the compile-cache env hash - a "
"location-only derived var is leaking into the key"
)
+9
View File
@@ -43,6 +43,15 @@ def test_language_model_only_affects_model_hash():
assert base_hash != lm_only_hash
@pytest.mark.parametrize("backend_arg", ["video_backend", "backend"])
def test_use_gpu_video_backend_from_media_io_kwargs(backend_arg: str):
config = MultiModalConfig(
media_io_kwargs={"video": {backend_arg: "pynvvideocodec"}}
)
assert config.use_gpu_video_backend()
def test_mm_encoder_fp8_scale_path_requires_fp8():
with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
+30
View File
@@ -349,6 +349,20 @@ _T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
_R = TypeVar("_R")
def _fix_v4_tied_weights_keys(model_cls: type) -> None:
"""Convert a v4 list-format _tied_weights_keys to the transformers v5 dict form."""
tied = getattr(model_cls, "_tied_weights_keys", None)
if not isinstance(tied, list) or not tied:
return
result = {
k: "model.embed_tokens.weight"
for k in tied
if "lm_head" in k and k.endswith(".weight")
}
if result:
setattr(model_cls, "_tied_weights_keys", result)
class HfRunner:
def get_default_device(self):
from vllm.platforms import current_platform
@@ -474,6 +488,22 @@ class HfRunner:
trust_remote_code=trust_remote_code,
)
else:
if trust_remote_code and hasattr(self.config, "auto_map"):
cls_ref = self.config.auto_map.get(auto_cls.__name__)
if cls_ref is not None:
from vllm.transformers_utils.dynamic_module import (
try_get_class_from_dynamic_module,
)
model_cls = try_get_class_from_dynamic_module(
cls_ref,
model_name,
trust_remote_code=trust_remote_code,
warn_on_fail=False,
)
if model_cls is not None:
_fix_v4_tied_weights_keys(model_cls)
model = cast(
nn.Module,
auto_cls.from_pretrained(
@@ -0,0 +1,76 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for check_stop_strings.
These are pure-function tests (no model / GPU). They pin down which stop
string is selected when several stop strings match within the text that was
appended in a single step -- which happens under speculative decoding, where
multiple tokens (and therefore multiple stop strings) can be appended at once.
"""
import pytest
from vllm.v1.engine.detokenizer import check_stop_strings
@pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]])
def test_earliest_completing_stop_wins_regardless_of_list_order(stop):
# " The user is a": " is a" (5 chars) was appended in one step. Both "is"
# (index 10) and " a" (index 13) land in the same window. "is" completes
# earlier in the text, so it must win over list order.
text = " The user is a"
new_char_count = len(" is a")
assert check_stop_strings(text, new_char_count, stop, include_in_output=False) == (
"is",
10,
)
@pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]])
def test_earliest_completing_stop_include_in_output(stop):
text = " The user is a"
new_char_count = len(" is a")
# Truncate to the end of "is" (index 12) -> " The user is".
assert check_stop_strings(text, new_char_count, stop, include_in_output=True) == (
"is",
12,
)
def test_completion_position_not_start_position():
# "b" starts later than "abc" but completes earlier, so it must win.
text = "abc"
assert check_stop_strings(
text, len(text), ["abc", "b"], include_in_output=False
) == ("b", 1)
@pytest.mark.parametrize(
"stop,expected",
[
(["ab", "b"], ("ab", 0)),
(["b", "ab"], ("b", 1)),
],
)
def test_ties_broken_by_list_order(stop, expected):
# "ab" and "b" both complete at index 2; list order decides the winner.
text = "ab"
assert (
check_stop_strings(text, len(text), stop, include_in_output=False) == expected
)
def test_single_stop_in_window_unchanged():
# The common case (one stop in the window) is unaffected by the change.
text = "hello world."
assert check_stop_strings(text, 1, ["."], include_in_output=False) == (".", 11)
# Stop completes at the very end -> no truncation needed (-1).
assert check_stop_strings(text, 1, ["."], include_in_output=True) == (".", -1)
def test_no_match_and_empty_inputs_return_none():
assert check_stop_strings("hello", 5, ["zzz"], include_in_output=False) is None
assert check_stop_strings("hello", 0, ["h"], include_in_output=False) is None
assert check_stop_strings("hello", 5, [], include_in_output=False) is None
+8 -5
View File
@@ -254,8 +254,8 @@ def test_multiproc_executor_shutdown_cleanup():
for worker in executor.workers:
assert not worker.proc.is_alive(), "Worker processes should be terminated"
# Verify shutdown event is set
assert executor.shutdown_event.is_set(), "Shutdown event should be set"
# Verify shutdown flag is set
assert executor.shutting_down, "Shutdown flag should be set"
# Multiple shutdowns should be safe (idempotent)
executor.shutdown()
@@ -292,7 +292,6 @@ def test_multiproc_executor_pipeline_parallel():
"Max concurrent batches should follow the configured PP/async "
"scheduling policy"
)
finally:
# Clean up
executor.shutdown()
@@ -338,6 +337,10 @@ def test_multiproc_executor_multi_node():
- Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
Total world_size = 4, nnodes = 2
"""
# Python 3.14+ changed default multiprocessing start method to 'forkserver'
# which cannot pickle nested functions. Use 'fork' for this test.
mp_ctx = multiprocessing.get_context("fork")
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
port = s.getsockname()[1]
@@ -405,12 +408,12 @@ def test_multiproc_executor_multi_node():
executor.shutdown()
# Create a queue to collect results from both processes
result_queue: multiprocessing.Queue[dict[str, int | bool]] = multiprocessing.Queue()
result_queue: multiprocessing.Queue[dict[str, int | bool]] = mp_ctx.Queue()
# Start both node processes
processes = []
for node_rank in range(2):
p = multiprocessing.Process(
p = mp_ctx.Process(
target=run_node,
args=(node_rank, result_queue, port),
name=f"Node{node_rank}",
+130
View File
@@ -348,6 +348,136 @@ def test_message_queue_busy_to_idle():
distributed_run(worker_fn_test_busy_to_idle, 4)
@pytest.mark.parametrize("should_warn", [False, True])
def test_reader_timeout_caps_indefinite_waits(should_warn):
with (
mock.patch(
"vllm.distributed.device_communicators.shm_broadcast."
"SHM_READER_RECHECK_INTERVAL_MS",
new=7,
),
mock.patch(
"vllm.distributed.device_communicators.shm_broadcast."
"VLLM_RINGBUFFER_WARNING_INTERVAL",
new=60,
),
):
timeout = MessageQueue.ReadTimeoutWithWarnings(
timeout=None, should_warn=should_warn
)
assert timeout.timeout_ms() == 7
def test_reader_rechecks_shm_after_idle_wait_timeout_without_notify():
writer = MessageQueue(
n_reader=1,
n_local_reader=1,
max_chunk_bytes=1024 * 1024,
max_chunks=1,
)
reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
payload = 123
poll_started = threading.Event()
allow_timeout = threading.Event()
result = {}
def acquire_read_in_thread():
try:
with reader.acquire_read(indefinite=True) as buf:
result["value"] = buf[0]
except Exception as exc:
result["exc"] = exc
def poll_timeout(*, timeout: int | None = None):
poll_started.set()
assert allow_timeout.wait(timeout=5)
return []
try:
writer.wait_until_ready()
reader.wait_until_ready()
reader._spin_condition.last_read = 0
reader._spin_condition.busy_loop_s = 0
with (
mock.patch(
"vllm.distributed.device_communicators.shm_broadcast."
"SHM_READER_RECHECK_INTERVAL_MS",
new=50,
),
mock.patch(
"vllm.distributed.device_communicators.shm_broadcast."
"VLLM_RINGBUFFER_WARNING_INTERVAL",
new=60,
),
mock.patch.object(
reader._spin_condition.poller,
"poll",
side_effect=poll_timeout,
) as poll,
):
read_thread = threading.Thread(target=acquire_read_in_thread, daemon=True)
read_thread.start()
assert poll_started.wait(timeout=5)
with writer.acquire_write(timeout=0.1) as buf:
buf[0] = payload
allow_timeout.set()
read_thread.join(timeout=5)
assert not read_thread.is_alive()
poll.assert_called_once_with(timeout=50)
if "exc" in result:
raise result["exc"]
assert result["value"] == payload
with writer.buffer.get_metadata(0) as metadata_buffer:
assert metadata_buffer[0] == 1
assert metadata_buffer[1] == 1
finally:
writer.shutdown()
reader.shutdown()
for socket in (
writer.local_socket,
writer._spin_condition.local_notify_socket,
reader.local_socket,
reader._spin_condition.local_notify_socket,
reader._spin_condition.read_cancel_socket,
reader._spin_condition.write_cancel_socket,
):
socket.close(linger=0)
def test_acquire_read_releases_slot_when_reader_raises():
writer = MessageQueue(
n_reader=1,
n_local_reader=1,
max_chunk_bytes=1024 * 1024,
max_chunks=1,
)
reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
try:
writer.wait_until_ready()
reader.wait_until_ready()
writer.enqueue({"payload": "first"})
with (
pytest.raises(RuntimeError, match="reader failed"),
reader.acquire_read(timeout=0.1),
):
raise RuntimeError("reader failed")
with writer.buffer.get_metadata(0) as metadata_buffer:
assert metadata_buffer[0] == 1
assert metadata_buffer[1] == 1
with writer.acquire_write(timeout=0.1) as buf:
buf[0] = 0
finally:
writer.shutdown()
reader.shutdown()
def test_warning_logs(caplog_vllm):
"""
Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals
@@ -142,8 +142,19 @@ def qwen2audio_aligned_content_and_embeds_b64() -> tuple[str, str]:
@pytest.mark.asyncio
@pytest.mark.parametrize(
"audio_first",
[True, False],
ids=["audio_embeds-then-text", "text-then-audio_embeds"],
[
pytest.param(True, id="audio_embeds-then-text"),
pytest.param(
False,
id="text-then-audio_embeds",
marks=pytest.mark.xfail(
reason="torch 2.12 regression: prompt_embeds output diverges "
"from raw-text when text precedes audio; "
"https://github.com/pytorch/pytorch/issues/184431",
strict=True,
),
),
],
)
async def test_text_content_and_prompt_embeds_match_with_audio_embeds(
qwen2audio_client: openai.AsyncOpenAI,
@@ -9,8 +9,6 @@ from pydantic import TypeAdapter, ValidationError
from vllm import PoolingParams
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
from vllm.entrypoints.pooling.embed.protocol import (
CohereEmbedContent,
CohereEmbedInput,
CohereEmbedRequest,
EmbeddingBatchChatInputRequest,
EmbeddingBatchChatRequest,
@@ -19,7 +17,10 @@ from vllm.entrypoints.pooling.embed.protocol import (
EmbeddingCompletionRequest,
EmbeddingRequest,
)
from vllm.entrypoints.pooling.typing import PoolingServeContext
from vllm.entrypoints.pooling.typing import (
PoolingEngineInput,
PoolingServeContext,
)
from vllm.outputs import PoolingOutput, PoolingRequestOutput
@@ -410,6 +411,7 @@ class TestChunkedEmbeddingProcessing:
def _make_handler(cls):
handler = object.__new__(EmbedIOProcessor)
handler.model_config = cls._FakeModelConfig()
handler.enable_chunked_processing = True
return handler
@staticmethod
@@ -421,15 +423,29 @@ class TestChunkedEmbeddingProcessing:
}
)
assert isinstance(request, EmbeddingCompletionRequest)
pooling_params = PoolingParams()
return PoolingServeContext(
request=request,
pooling_params=PoolingParams(),
pooling_params=pooling_params,
model_name="test",
request_id="embd-client-prompt-999-chunk-888",
engine_inputs=[
{"prompt_token_ids": [0, 1, 2, 3, 4]},
{"prompt_token_ids": [10, 11]},
PoolingEngineInput(
prompts={"prompt_token_ids": [0, 1, 2, 3, 4]},
params=pooling_params,
lora_requests=None,
priorities=0,
),
PoolingEngineInput(
prompts={"prompt_token_ids": [10, 11]},
params=pooling_params,
lora_requests=None,
priorities=0,
),
],
lora_request=None,
priorities=0,
prompt_extras=None,
)
@staticmethod
@@ -450,7 +466,7 @@ class TestChunkedEmbeddingProcessing:
handler = self._make_handler()
ctx = self._make_context()
handler._pre_process_chunked(ctx)
handler.maybe_pre_process_chunked(ctx)
assert ctx.prompt_request_ids == [
"embd-client-prompt-999-chunk-888-prompt-0-chunk-0",
@@ -488,227 +504,3 @@ class TestChunkedEmbeddingProcessing:
ctx.final_res_batch[1].outputs.data,
torch.tensor([9.0, 9.0]),
)
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),
pooling_params=PoolingParams(),
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_cmpl_online(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_cmpl_online = preprocess_cmpl_online
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_cmpl(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_cmpl_online = preprocess_cmpl
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_cmpl_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,
},
)
]
class TestPreProcessOpenAIEmbeddingChatOnline:
"""Unit tests for OpenAI embedding chat preprocessing."""
class _FakeModelConfig:
max_model_len = 128
encoder_config: dict[str, object] = {}
pooler_config = None
multimodal_config = None
is_encoder_decoder = False
class _FakeRenderer:
tokenizer = object()
def __init__(self):
self.calls = []
def render_chat(
self,
all_messages,
chat_params,
tok_params,
prompt_extras=None,
):
self.calls.append(
{
"all_messages": all_messages,
"chat_params": chat_params,
"tok_params": tok_params,
"prompt_extras": prompt_extras,
}
)
return all_messages, [
{"prompt_token_ids": [index]} for index, _ in enumerate(all_messages)
]
@classmethod
def _make_handler(cls, renderer):
handler = object.__new__(EmbedIOProcessor)
handler.renderer = renderer
handler.model_config = cls._FakeModelConfig()
handler.chat_template = "template"
handler.chat_template_content_format = "auto"
handler.trust_request_chat_template = False
handler.enable_chunked_processing = False
return handler
@staticmethod
def _make_context(
request: (
EmbeddingChatRequest
| EmbeddingBatchChatRequest
| EmbeddingChatInputRequest
| EmbeddingBatchChatInputRequest
),
) -> PoolingServeContext[
EmbeddingChatRequest
| EmbeddingBatchChatRequest
| EmbeddingChatInputRequest
| EmbeddingBatchChatInputRequest
]:
return PoolingServeContext(
request=request,
pooling_params=PoolingParams(),
model_name="test",
request_id="embd-test",
)
def test_chat_template_kwargs_forwarded_for_batched_input_messages(self):
request = TypeAdapter(EmbeddingRequest).validate_python(
{
"model": "test",
"input": [
[{"role": "user", "content": "hello"}],
[{"role": "user", "content": "goodbye"}],
],
"add_generation_prompt": True,
"chat_template_kwargs": {"instruction": "Represent the query: "},
"mm_processor_kwargs": {"max_pixels": 1},
"cache_salt": "salt",
}
)
assert isinstance(request, EmbeddingBatchChatInputRequest)
renderer = self._FakeRenderer()
handler = self._make_handler(renderer)
ctx = self._make_context(request)
handler.pre_process_online(ctx)
assert ctx.engine_inputs == [
{"prompt_token_ids": [0]},
{"prompt_token_ids": [1]},
]
assert len(renderer.calls) == 1
call = renderer.calls[0]
assert call["all_messages"] == request.messages
assert call["prompt_extras"] == {
"mm_processor_kwargs": {"max_pixels": 1},
"cache_salt": "salt",
}
chat_template_kwargs = call["chat_params"].chat_template_kwargs
assert chat_template_kwargs["instruction"] == "Represent the query: "
assert chat_template_kwargs["add_generation_prompt"] is True
assert chat_template_kwargs["continue_final_message"] is False
assert "tools" not in chat_template_kwargs
assert chat_template_kwargs["tokenize"] is False
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
from types import SimpleNamespace
import pytest
import torch
@@ -10,10 +9,7 @@ import torch
from tests.models.utils import softmax
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
from vllm.entrypoints.pooling.scoring.typing import ScoringData
from vllm.platforms import current_platform
from vllm.renderers import TokenizeParams
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
PROMPT = "The chef prepared a delicious meal."
@@ -145,45 +141,6 @@ def test_max_tokens_per_doc(llm: LLM):
assert with_limit_tokens < no_limit_tokens
def test_token_type_ids_follow_post_tokenization():
processor = object.__new__(CrossEncoderIOProcessor)
processor.tokenizer = SimpleNamespace(truncation_side="right", pad_token_id=-1)
processor.renderer = SimpleNamespace(process_for_engine=lambda prompt, _: prompt)
processor.model_config = None
processor.get_score_prompt = lambda **_: (
"",
{
"prompt_token_ids": list(range(32)),
"token_type_ids": [0] * 16 + [1] * 16,
},
)
engine_inputs, pooling_params = processor._pre_process(
ScoringData(data_1=["query"], data_2=["document"]),
TokenizeParams(
max_total_tokens=None,
truncate_prompt_tokens=16,
truncation_side="left",
),
PoolingParams(task="classify", extra_kwargs={"cache_salt": "salt"}),
)
assert engine_inputs[0]["prompt_token_ids"] == list(range(16, 32))
assert pooling_params[0].extra_kwargs == {
"cache_salt": "salt",
"compressed_token_type_ids": 0,
}
engine_inputs, pooling_params = processor._pre_process(
ScoringData(data_1=["query"], data_2=["document"]),
TokenizeParams(max_total_tokens=None, pad_prompt_tokens=40),
PoolingParams(task="classify"),
)
assert engine_inputs[0]["prompt_token_ids"] == list(range(32)) + [-1] * 8
assert pooling_params[0].extra_kwargs == {"compressed_token_type_ids": 16}
def test_pooling_params(llm: LLM):
def get_outputs(use_activation):
outputs = llm.score(
@@ -161,7 +161,9 @@ def _make_pooling_serving(lora_name: str) -> _ConcretePoolingServing:
return serving
def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
def _make_pooling_ctx(
model_name: str, serving: PoolingBaseServing
) -> PoolingServeContext:
mock_request = MagicMock()
mock_request.model = model_name
return PoolingServeContext(
@@ -169,6 +171,9 @@ def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
model_name=MODEL_NAME,
request_id="test-id",
pooling_params=PoolingParams(),
lora_request=serving._maybe_get_adapters(mock_request),
priorities=0,
prompt_extras=None,
)
@@ -176,9 +181,7 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
"""LoRA adapter name must populate ctx.lora_request without raising."""
lora_name = "bot-embed-lora"
serving = _make_pooling_serving(lora_name)
ctx = _make_pooling_ctx(lora_name)
ctx.lora_request = serving._maybe_get_adapters(ctx.request)
ctx = _make_pooling_ctx(lora_name, serving)
assert ctx.lora_request is not None
assert ctx.lora_request.lora_name == lora_name
@@ -187,7 +190,6 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
def test_pooling_maybe_get_adapters_unknown_model_raises():
"""An unrecognised model name must still raise VLLMNotFoundError."""
serving = _make_pooling_serving("some-lora")
ctx = _make_pooling_ctx("unknown-model")
with pytest.raises(VLLMNotFoundError):
serving._maybe_get_adapters(ctx.request)
_make_pooling_ctx("unknown-model", serving)
@@ -6,6 +6,7 @@ max_concurrency: 100
server_args: >-
--enforce-eager
--max-model-len 4096
--max-num-batched-tokens 32768
--safetensors-load-strategy prefetch
--moe-backend flashinfer_cutlass
--prefill-context-parallel-size 4
@@ -6,6 +6,7 @@ max_concurrency: 100
server_args: >-
--enforce-eager
--max-model-len 4096
--max-num-batched-tokens 32768
--safetensors-load-strategy prefetch
--moe-backend flashinfer_cutlass
--tensor-parallel-size 2
@@ -0,0 +1,9 @@
model_name: "poolside/Laguna-XS.2-NVFP4"
accuracy_threshold: 0.86
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
server_args: >-
--enforce-eager
--max-model-len 4096
--trust-remote-code
@@ -2,4 +2,7 @@ model_name: "google/gemma-4-E4B-it-qat-mobile-ct"
accuracy_threshold: 0.50
num_questions: 1319
num_fewshot: 5
server_args: "--enforce-eager --max-model-len 4096"
server_args: >-
--enforce-eager
--max-model-len 4096
--speculative-config '{"method":"mtp","model":"google/gemma-4-E4B-it-assistant","num_speculative_tokens":4}'
@@ -3,3 +3,4 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
Qwen1.5-MoE-W4A16-CT.yaml
DeepSeek-V2-Lite-Instruct-FP8.yaml
Qwen3-30B-A3B-NVFP4.yaml
Laguna-XS.2-NVFP4.yaml
@@ -546,8 +546,11 @@ def test_flash_attn_accepts_handled_fp8_variants(
):
"""FlashAttentionBackend must accept the two fp8 dtypes it can actually
handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""
import vllm.v1.attention.backends.flash_attn as fa_mod
import vllm.v1.attention.backends.fa_utils as fa_utils_mod
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
monkeypatch.setattr(fa_mod.current_platform, "is_xpu", lambda: True)
# The fp8 decision is made in fa_utils, using its own current_platform
# binding, so patch is_xpu there (not on flash_attn's) to stay robust to
# import order across earlier tests that patch vllm.platforms.current_platform.
monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True)
assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
@@ -11,6 +11,9 @@ from vllm._custom_ops import (
scaled_fp8_quant,
)
from vllm.platforms import current_platform
from vllm.v1.attention.ops.triton_merge_attn_states import (
mask_empty_context,
)
from vllm.v1.attention.ops.triton_merge_attn_states import (
merge_attn_states as merge_attn_states_triton,
)
@@ -73,6 +76,59 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]
all_case_info: list[tuple] = []
def test_mask_empty_context() -> None:
query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
query_start_loc = torch.cat(
(torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
).cuda()
context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
context_start_loc = torch.cat(
(torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
).cuda()
num_heads, num_tokens, head_dim = 4, 165, 16
lse = torch.randn(num_heads, num_tokens, device="cuda")
output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
# Empty-context rows carry undefined (possibly non-finite) attention output.
output[33:164] = float("nan")
expected_lse = lse.clone()
expected_lse[:, 33:164] = float("-inf")
expected_output = output.clone()
expected_output[33:164] = 0.0
mask_empty_context(lse, output, query_start_loc, context_start_loc)
torch.testing.assert_close(lse, expected_lse)
torch.testing.assert_close(output, expected_output)
@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
"""When a token is empty on both sides (both LSE -inf), the 0/0 softmax
scales must not surface as NaN in the merged output."""
num_tokens, num_heads, head_size = 6, 8, 128
prefix_output = torch.zeros(
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
)
prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")
suffix_output = torch.zeros(
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
)
suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
# Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
# their outputs and set both LSEs to -inf).
empty = slice(2, 4)
prefix_lse[:, empty] = float("-inf")
suffix_lse[:, empty] = float("-inf")
output = torch.empty_like(prefix_output)
merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
assert not output.isnan().any()
def generate_markdown_table():
global all_case_info
table_header = (
@@ -40,10 +40,11 @@ BATCH_SIZE = 128
CONTEXT_LEN = 8192
PAGE_SIZE = 1
# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8
# KV-cache -> fp8_e4m3 kv.
EXPECTED_Q_DTYPE = torch.bfloat16
EXPECTED_KV_DTYPE = torch.float8_e4m3fn
# On the fp8 KV-cache path the builder forwards the platform fp8 dtype (aiter
# dtypes.fp8) for both q and kv. Mirror it via current_platform.fp8_dtype()
# instead of hardcoding a literal (see #47276).
EXPECTED_Q_DTYPE = current_platform.fp8_dtype()
EXPECTED_KV_DTYPE = current_platform.fp8_dtype()
# The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
# is excluded: it holds raw device pointers, never equal across allocations.
+15
View File
@@ -0,0 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
@pytest.fixture(autouse=True)
def reset_default_torch_device():
"""Several kernel tests call torch.set_default_device without restoring
it, which poisons subsequent tests in the same pytest run (e.g. CPU
tensors silently created on CUDA). Restore the factory default after
every test.
"""
yield
torch.set_default_device(None)
+14 -19
View File
@@ -33,12 +33,10 @@ from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402
@pytest.mark.parametrize("num_seqs", [1, 5, 257])
@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
seq_lens = torch.randint(
1,
130,
(num_seqs,),
dtype=torch.int32,
)
rng_cpu = torch.Generator("cpu").manual_seed(1234)
rng = torch.Generator("cuda").manual_seed(2345)
seq_lens = torch.randint(1, 130, (num_seqs,), dtype=torch.int32, generator=rng_cpu)
cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32)
cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0)
total_tokens = int(cu_seqlens[-1].item())
@@ -56,8 +54,9 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_k_dim,
device="cuda",
dtype=dtype,
generator=rng,
)
k = torch.randn_like(q)
k = torch.randn_like(q, generator=rng)
v = torch.randn(
1,
total_tokens,
@@ -65,29 +64,24 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_v_dim,
device="cuda",
dtype=dtype,
generator=rng,
)
q = F.normalize(q.float(), p=2, dim=-1).to(dtype)
k = F.normalize(k.float(), p=2, dim=-1).to(dtype)
a = torch.randn(
1,
total_tokens,
num_v_heads,
device="cuda",
dtype=dtype,
1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
)
b = torch.randn(
1,
total_tokens,
num_v_heads,
device="cuda",
dtype=dtype,
1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
)
# Match upstream FLA GatedDeltaNet synthetic initialization:
# https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16)
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(
0, 16, generator=rng
)
A_log = torch.log(A)
dt = torch.exp(
torch.rand(num_v_heads, device="cuda", dtype=torch.float32)
torch.rand(num_v_heads, device="cuda", dtype=torch.float32, generator=rng)
* (math.log(0.1) - math.log(0.001))
+ math.log(0.001)
)
@@ -105,6 +99,7 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_k_dim,
device="cuda",
dtype=state_dtype,
generator=rng,
)
* 0.05
)
+2 -2
View File
@@ -82,9 +82,9 @@ def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
).sum(dim=1)
@pytest.fixture(autouse=True, scope="module")
@pytest.fixture(autouse=True)
def setup_cuda():
"""Sets the default CUDA device for all tests in this module."""
"""Sets the default CUDA device before each test in this module."""
torch.set_default_device("cuda")
@@ -6,6 +6,7 @@ import pytest
import torch
import torch.nn.functional as F
import vllm._custom_ops as ops
from vllm.model_executor.layers.fused_moe.config import (
RoutingMethodType,
get_routing_method_type,
@@ -231,3 +232,119 @@ def test_dsv4_fast_topk(
atol=2e-5,
rtol=2e-5,
)
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="This test is skipped on non-CUDA platform.",
)
@pytest.mark.parametrize("use_hash", [False, True])
@pytest.mark.parametrize("use_bias", [False, True])
@pytest.mark.parametrize("use_padding_mask", [False, True])
@pytest.mark.parametrize("pad_with_nan", [False, True])
@pytest.mark.parametrize("num_experts", [128, 256, 384])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
def test_fused_topk_softplus_sqrt_padding(
use_hash: bool,
use_bias: bool,
use_padding_mask: bool,
pad_with_nan: bool,
num_experts: int,
dtype: torch.dtype,
):
"""Verify explicit padding and NaN-padded rows do not affect real rows."""
torch.manual_seed(0)
num_tokens = 8
topk = 6
indices_dtype = torch.int32
gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
padding_rows = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
padding_rows[1::2] = True
if pad_with_nan:
gating_output[padding_rows] = float("nan")
is_padding = padding_rows if use_padding_mask else None
# A negative correction bias makes explicit pad rows look selectable unless
# the kernel uses the is_padding guard.
e_score_correction_bias = None
if use_bias:
e_score_correction_bias = (
-torch.rand((num_experts,), dtype=torch.float32, device="cuda") - 1.0
)
input_ids = None
hash_indices_table = None
if use_hash:
vocab_size = 64
hash_indices_table = torch.stack(
[torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
).to(device="cuda", dtype=indices_dtype)
input_ids = torch.randint(
0, vocab_size, (num_tokens,), dtype=indices_dtype, device="cuda"
)
topk_weights = torch.empty(num_tokens, topk, dtype=torch.float32, device="cuda")
topk_ids = torch.empty(num_tokens, topk, dtype=indices_dtype, device="cuda")
token_expert_indices = torch.empty(
num_tokens, topk, dtype=torch.int32, device="cuda"
)
ops.topk_hash_softplus_sqrt(
topk_weights,
topk_ids,
token_expert_indices,
gating_output,
renormalize=True,
routed_scaling_factor=1.0,
e_score_correction_bias=e_score_correction_bias,
input_tokens=input_ids,
hash_indices_table=hash_indices_table,
is_padding=is_padding,
)
if use_padding_mask:
pad_ids = topk_ids[padding_rows]
pad_weights = topk_weights[padding_rows]
assert torch.equal(pad_ids, torch.full_like(pad_ids, -1)), (
f"Explicit pad rows should contain only -1 ids, got {pad_ids.tolist()}"
)
assert (pad_weights == 0).all(), (
"Explicit pad rows should have all-zero weights, "
f"got {pad_weights.tolist()}"
)
if pad_with_nan:
nan_pad_weights = topk_weights[padding_rows]
assert torch.isfinite(nan_pad_weights).all(), (
f"NaN-padded rows have non-finite weights, got {nan_pad_weights.tolist()}"
)
assert (nan_pad_weights == 0).all(), (
"NaN-padded rows should have all-zero weights, "
f"got {nan_pad_weights.tolist()}"
)
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
gating_output=gating_output,
topk=topk,
renormalize=True,
routed_scaling_factor=1.0,
e_score_correction_bias=e_score_correction_bias,
input_ids=input_ids,
hash_indices_table=hash_indices_table,
)
rows_to_compare = torch.ones(num_tokens, dtype=torch.bool, device="cuda")
if use_padding_mask or pad_with_nan:
rows_to_compare = ~padding_rows
sorted_ref_ids, idx_ref = topk_ids_ref[rows_to_compare].sort(dim=-1)
sorted_ids, idx_ops = topk_ids[rows_to_compare].sort(dim=-1)
torch.testing.assert_close(
sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0
)
sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref)
sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops)
torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
@@ -102,9 +102,9 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
).sum(dim=1)
@pytest.fixture(autouse=True, scope="module")
@pytest.fixture(autouse=True)
def setup_cuda():
"""Sets the default CUDA device for all tests in this module."""
"""Sets the default CUDA device before each test in this module."""
torch.set_default_device("cuda")
+10 -10
View File
@@ -28,12 +28,6 @@ BLOCK_SIZE = [[128, 128]]
SEEDS = [0]
@pytest.fixture(autouse=True, scope="module")
def setup_cuda():
"""Sets the default CUDA device for all tests in this module."""
torch.set_default_device("cuda")
@pytest.mark.parametrize(
"M,N,K,block_size,out_dtype,seed",
itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS),
@@ -41,22 +35,28 @@ def setup_cuda():
@torch.inference_mode()
def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed):
torch.manual_seed(seed)
device = current_platform.device_type
factor_for_scale = 1e-2
int8_info = torch.iinfo(torch.int8)
int8_max, int8_min = int8_info.max, int8_info.min
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * int8_max
A_fp32 = torch.rand(M, K, dtype=torch.float32, device=device)
A_fp32 = (A_fp32 - 0.5) * 2 * int8_max
A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * int8_max
B_fp32 = torch.rand(N, K, dtype=torch.float32, device=device)
B_fp32 = (B_fp32 - 0.5) * 2 * int8_max
B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
block_n, block_k = block_size[0], block_size[1]
n_tiles = (N + block_n - 1) // block_n
k_tiles = (K + block_k - 1) // block_k
As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale
Bs = (
torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device)
* factor_for_scale
)
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
+12 -10
View File
@@ -82,12 +82,6 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids
).sum(dim=1)
@pytest.fixture(autouse=True, scope="module")
def setup_cuda():
"""Sets the default CUDA device for all tests in this module."""
torch.set_default_device("cuda")
DTYPES = [torch.half, torch.bfloat16]
M = [1, 33]
N = [128, 1024]
@@ -104,6 +98,7 @@ SEEDS = [0]
@torch.inference_mode()
def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
torch.manual_seed(seed)
device = current_platform.device_type
# Initialize int8 quantization parameters
factor_for_scale = 1e-2
int8_max = 127
@@ -111,19 +106,26 @@ def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
# Input tensor
# M * K
a = torch.randn((M, K), dtype=dtype) / 10
a = torch.randn((M, K), dtype=dtype, device=device) / 10
# Generate int8 weights
w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
w1_fp32 = (
torch.rand(
(E, 2 * N, K),
dtype=torch.float32,
device=device,
)
- 0.5
) * 2
w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32, device=device) - 0.5) * 2
w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
# Generate scale for each column (per-column quantization)
w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
score = torch.randn((M, E), dtype=dtype)
score = torch.randn((M, E), dtype=dtype, device=device)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weights, topk_ids = torch.topk(score, topk)
+5 -3
View File
@@ -21,9 +21,9 @@ def test_gather_cache_oob():
seq_starts causes the block_table offset to read out of bounds.
"""
batch_size = 1
block_size = 64
entry_size = 128
# The kernel only supports the MLA entry sizes.
entry_size = 576
block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
@@ -34,6 +34,7 @@ def test_gather_cache_oob():
seq_len = 65
cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
token_to_seq = torch.zeros(seq_len, dtype=torch.int32, device="cuda")
# src_cache: [num_blocks, block_size, entry_size]
num_blocks = 5
@@ -51,7 +52,8 @@ def test_gather_cache_oob():
dst,
block_table,
cu_seq_lens,
batch_size,
token_to_seq,
seq_len,
"auto", # kv_cache_dtype
scale,
seq_starts,
+21
View File
@@ -13,6 +13,7 @@ These tests cover:
"""
import math
from types import SimpleNamespace
import pytest
import torch
@@ -27,6 +28,7 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
_launch_two_stage_sparse_attn_compressor,
)
from vllm.models.deepseek_v4.compressor import _get_c128_boundary
from vllm.platforms import current_platform
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
@@ -58,6 +60,25 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
return x_fp8, scales
@pytest.mark.parametrize(
("starts", "query_start_loc", "expected"),
[
([0], [0, 127], False),
([0], [0, 128], True),
([127], [0, 1], True),
([128], [0, 127], False),
([1, 255], [0, 1, 2], True),
(None, [0, 1], None),
],
)
def test_get_c128_boundary(starts, query_start_loc, expected):
metadata = SimpleNamespace(
_num_computed_tokens_cpu=None if starts is None else torch.tensor(starts),
query_start_loc_cpu=torch.tensor(query_start_loc),
)
assert _get_c128_boundary(metadata) is expected
# ── Test A: DeepseekV4 Attention path ──────────────────────────────────────────────
+1 -1
View File
@@ -264,7 +264,7 @@ def test_block_mask_direct_vs_slow_path():
device = torch.device("cuda")
vllm_config = create_vllm_config(
model_name="meta-llama/Meta-Llama-3-8B", block_size=16, max_model_len=1024
model_name="Qwen/Qwen2.5-1.5B-Instruct", block_size=16, max_model_len=1024
)
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
+13 -4
View File
@@ -726,14 +726,19 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
This catches stride/layout bugs that only manifest when the einsum
kernel actually consumes the quantized activations.
"""
from deep_gemm.utils.math import ceil_div
from vllm.utils.deep_gemm import (
fp8_einsum,
is_deep_gemm_supported,
per_block_cast_to_fp8,
transform_sf_into_required_layout,
)
if not is_deep_gemm_supported():
pytest.skip("DeepGEMM not supported on this platform")
def ceil_div(a: int, b: int) -> int:
return (a + b - 1) // b
heads_per_group = num_heads // n_groups
d = heads_per_group * HEAD_DIM
o_lora_rank = 1024
@@ -809,8 +814,12 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
# -- Checks --
# Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
# come from fp32 ordering and UE8M0 boundary shifts only.
# Use relative diff (same metric as test_fp8_einsum.py).
from deep_gemm.testing import calc_diff
# Use relative diff (same metric as deep_gemm.testing.calc_diff).
def calc_diff(x, y):
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
z_diff = calc_diff(z_fused, z_ref)
assert z_diff < 0.01, (
@@ -84,6 +84,22 @@ def norm_rope_ref(x, weight, positions, cos_sin_cache, eps):
return roped
def assert_fp8_cache_close(kv_cache, expected_kv_cache):
"""Compare two e4m3 caches allowing 1 ulp.
On CUDA the fused kernel quantizes K from its fp32 intermediate, while the
reshape_and_cache_flash reference quantizes the bf16-materialized value, so
rounding-boundary values may differ by one e4m3 code.
"""
byte_diff = (kv_cache.int() - expected_kv_cache.int()).abs()
got = kv_cache.view(torch.float8_e4m3fn).float()
exp = expected_kv_cache.view(torch.float8_e4m3fn).float()
ok = (byte_diff <= 1) | ((got == 0) & (exp == 0))
assert bool(ok.all()), (
f"fp8 cache differs by more than 1 ulp in {int((~ok).sum())} elements"
)
# ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
@@ -265,7 +281,7 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
scale,
scale,
)
torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
assert_fp8_cache_close(kv_cache, expected_kv_cache)
else:
for t in range(num_tokens):
s = slot_mapping[t].item()
@@ -383,7 +399,7 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
scale,
scale,
)
torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
assert_fp8_cache_close(kv_cache, expected_kv_cache)
else:
k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)
@@ -41,11 +41,12 @@ def test_fused_recurrent_packed_decode_matches_reference(
A_log = torch.randn((HV,), device=device, dtype=dtype)
dt_bias = torch.randn((HV,), device=device, dtype=dtype)
# Continuous batching indices (include PAD_SLOT_ID=-1 cases).
ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32)
# Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is
# reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1.
ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32)
ssm_state_indices[-3:] = -1
state0 = torch.randn((B, HV, V, K), device=device, dtype=dtype)
state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype)
state_ref = state0.clone()
state_packed = state0.clone()
@@ -94,5 +95,8 @@ def test_fused_recurrent_packed_decode_matches_reference(
atol = 2e-2 if dtype != torch.float32 else 1e-4
rtol = 1e-2 if dtype != torch.float32 else 1e-4
torch.testing.assert_close(out_packed, out_ref, rtol=rtol, atol=atol)
# Output rows for PAD_SLOT_ID entries are never written (uninitialized in
# both paths), so compare only the valid rows.
valid = ssm_state_indices > 0
torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol)
torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)
@@ -58,10 +58,12 @@ def test_fused_sigmoid_gating_delta_rule_update_non_spec(
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
# Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
# state indices start at 1.
ssm_state = torch.rand(
total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
)
state_indices = torch.randperm(total_entries, dtype=torch.int32)[:num_tokens]
state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[:num_tokens]
cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32)
beta = b.sigmoid()
@@ -144,13 +146,14 @@ def test_fused_sigmoid_gating_delta_rule_update_spec(
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
# Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
# state indices start at 1.
ssm_state = torch.rand(
total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
)
state_indices = torch.randperm(
total_entries,
dtype=torch.int32,
)[:num_tokens].view(num_reqs, num_speculative_tokens + 1)
state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[
:num_tokens
].view(num_reqs, num_speculative_tokens + 1)
num_accepted_tokens = torch.randint(
1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32
)
@@ -0,0 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.model_executor.model_loader.mtp_validation import (
disable_mtp_completeness_check,
is_mtp_completeness_check_enabled,
)
def test_disable_mtp_completeness_check_is_scoped():
assert is_mtp_completeness_check_enabled()
with pytest.raises(RuntimeError), disable_mtp_completeness_check():
assert not is_mtp_completeness_check_enabled()
raise RuntimeError
assert is_mtp_completeness_check_enabled()
@@ -6,6 +6,7 @@ import pytest
from vllm.config.compilation import CompilationConfig, CUDAGraphMode
from vllm.models.inkling.common.mm_preprocess import InklingMultiModalDataParser
from vllm.models.inkling.common.towers import plan_out_scales
from vllm.models.inkling.configs import (
InklingAudioConfig,
InklingModelConfig,
@@ -17,6 +18,22 @@ from vllm.models.inkling.nvidia.sconv_swa_attn import (
from vllm.v1.attention.backend import AttentionCGSupport
def test_vision_scale_plan_matches_released_config():
assert plan_out_scales(2, 40, 4) == [
(1, 1, 1, 3),
(1, 5, 5, 128),
(1, 10, 10, 320),
(1, 40, 40, 4800),
(2, 40, 40, 9600),
]
def test_vision_scale_plan_breaks_assignment_ties_in_order():
reductions = [np.prod(scale[:-1]) for scale in plan_out_scales(2, 52, 4)]
assert reductions == sorted(set(reductions))
@pytest.mark.parametrize(
("config_cls", "kwargs", "missing"),
[
@@ -908,7 +908,15 @@ VLM_TEST_SETTINGS = {
multi_image_prompt="Picture 1: <vlm_image>\nPicture 2: <vlm_image>\nDescribe these two images with one paragraph respectively.", # noqa: E501
max_model_len=4096,
max_num_seqs=2,
num_logprobs=10,
# torch 2.13 accumulates CPU numerical drift in the qwen2_vl multi-image
# path: HF and vLLM agree for a long prefix (~69 tokens) then a token
# flips outside vLLM's top-N only near the end of the generation. The
# window is already at the max_logprobs=20 cap, so widening it further is
# not possible. Treat this as acceptable drift and cap max_tokens on CPU
# so the compared prefix stays before the divergence, keeping the
# multi-image path under test. See pytorch/pytorch#187735.
max_tokens=64 if current_platform.is_cpu() else 128,
num_logprobs=20 if current_platform.is_cpu() else 10,
auto_cls=AutoModelForImageTextToText,
vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output,
image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
-7
View File
@@ -487,13 +487,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"Plamo2ForCausalLM": _HfExamplesInfo(
"pfnet/plamo-2-1b",
trust_remote_code=True,
max_transformers_version="4.57",
transformers_version_reason={
"hf": (
"Custom model code uses `_tied_weight_keys: list[str]` but "
"Transformers v5 now expects `_tied_weight_keys: dict[str, str]`"
)
},
),
"Plamo3ForCausalLM": _HfExamplesInfo(
"pfnet/plamo-3-nict-2b-base",
+97
View File
@@ -6,15 +6,43 @@ import time
import pytest
import vllm.config.multimodal as multimodal_config_module
from vllm.config.multimodal import MultiModalConfig
from vllm.multimodal.gpu_ipc_memory import (
MultiModalGPUMemoryPool,
get_mm_gpu_ipc_pool,
maybe_init_mm_gpu_ipc_pool,
reserve_mm_ipc_gpu_memory,
set_mm_gpu_ipc_pool,
)
from vllm.multimodal.video import (
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
PYNVVIDEOCODEC_MAX_RETAINED_DECODERS,
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
from vllm.utils.mem_constants import GiB_bytes
def _mm_config(
*,
mm_ipc_gpu_memory_gb: float = 0,
video_backend: str | None = None,
) -> MultiModalConfig:
video_kwargs = {} if video_backend is None else {"video_backend": video_backend}
return MultiModalConfig(
mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb,
media_io_kwargs={"video": video_kwargs} if video_kwargs else {},
)
def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int:
return api_process_count * (
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS
+ PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES
)
def test_acquire_release_accounting():
pool = MultiModalGPUMemoryPool(total_bytes=100)
assert pool.available_bytes == 100
@@ -143,3 +171,72 @@ def test_global_pool_splits_budget_across_api_processes():
def test_global_pool_rejects_invalid_api_process_count():
with pytest.raises(ValueError):
maybe_init_mm_gpu_ipc_pool(2, api_process_count=0)
@pytest.mark.parametrize("video_backend", [None, "opencv"])
def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
monkeypatch: pytest.MonkeyPatch,
video_backend: str | None,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
mm_config = _mm_config(
mm_ipc_gpu_memory_gb=0.25,
video_backend=video_backend,
)
assert reserve_mm_ipc_gpu_memory(GiB_bytes, mm_config) == int(0.75 * GiB_bytes)
def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
mm_config = _mm_config(
mm_ipc_gpu_memory_gb=0.25,
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
)
available_bytes = 4 * GiB_bytes
assert reserve_mm_ipc_gpu_memory(available_bytes, mm_config) == (
available_bytes - int(0.25 * GiB_bytes) - _pynvvideocodec_decoder_budget()
)
def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
available_bytes = 4 * GiB_bytes
assert reserve_mm_ipc_gpu_memory(available_bytes, _mm_config()) == (
available_bytes - _pynvvideocodec_decoder_budget()
)
def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
available_bytes = 8 * GiB_bytes
assert reserve_mm_ipc_gpu_memory(
available_bytes,
_mm_config(),
api_process_count=3,
) == available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3)
+42
View File
@@ -920,3 +920,45 @@ class TestDelegatingParserLargeDelta:
assert output.tool_calls[0]["name"] == "get_weather"
args = json.loads(output.tool_calls[0]["arguments"])
assert args == {"location": "Berlin"}
@pytest.mark.parametrize(
"chunk_size",
[1, 2, 3, 5, None],
ids=lambda c: f"chunk={c}",
)
def test_eos_not_leaked_when_reasoning_never_ends(self, chunk_size):
"""EOS must not leak into reasoning_content when the model never
emits </think> (generation ends while still in REASONING state)."""
eos_text = "<end▁of▁sentence>"
eos_id = 128801
vocab = {
**_DSV4_FULL_VOCAB,
eos_text: eos_id,
}
reasoning_text = "Good morning! How can I help you today?"
tokens: list[tuple[int, str]] = []
tid = 100
for word in reasoning_text.split(" "):
prefix = " " if tokens else ""
tokens.append((tid, prefix + word))
tid += 1
tokens.append((eos_id, eos_text))
tokenizer = MockTokenizer(vocab=vocab, tokens=tokens)
parser = _DeepSeekV4Delegating(
tokenizer,
chat_template_kwargs={"thinking": True},
)
deltas = replay_streaming(
parser,
tokens,
chunk_size=chunk_size,
finished_on_last=True,
)
output = collect_output(deltas)
assert reasoning_text in output.reasoning
assert eos_text not in output.reasoning
assert output.content == ""
assert output.tool_calls == []
+29 -7
View File
@@ -1624,19 +1624,41 @@ class TestDropSpecialTokens:
assert delta is not None
assert "<bos>" in delta.reasoning
def test_drops_suppressed_with_skip_tool_parsing(self):
"""When skip_tool_parsing is active, drop tokens are preserved
as content so a later tool-call pass can see them."""
def test_drops_applied_with_skip_tool_parsing(self):
"""Drop tokens are always dropped, even with skip_tool_parsing.
DROP_TERMINALs have no transitions by construction, so no parser
pass can use them."""
for initial_state in (ParserState.REASONING, ParserState.CONTENT):
engine = _make_engine(
vocab=_DROP_VOCAB,
special_tokens=list(_DROP_VOCAB.keys()),
)
engine._engine.skip_tool_parsing = True
engine._engine.reset(initial_state=initial_state)
events = engine._engine.feed("hello<bos>world", [72, 204, 73])
delta = engine._events_to_delta(events)
assert delta is not None
output = (delta.reasoning or "") + (delta.content or "")
assert "<bos>" not in output, f"<bos> leaked in state {initial_state}"
def test_transitions_unaffected_by_drop_in_reasoning_with_skip_tool_parsing(self):
"""With skip_tool_parsing in REASONING state, drop tokens are
removed but configured terminals still fire their transitions."""
engine = _make_engine(
vocab=_DROP_VOCAB,
special_tokens=list(_DROP_VOCAB.keys()),
)
engine._engine.skip_tool_parsing = True
engine._engine.reset()
events = engine._engine.feed("hello<bos>world", [72, 204, 73])
delta = engine._events_to_delta(events)
assert delta is not None
assert "<bos>" in delta.reasoning
events = engine._engine.feed("thought<bos></think>answer", [72, 204, 201, 73])
types = [e.type for e in events]
assert EventType.REASONING_CHUNK in types
assert EventType.REASONING_END in types
assert EventType.TEXT_CHUNK in types
reasoning_text = "".join(
e.value for e in events if e.type == EventType.REASONING_CHUNK
)
assert "<bos>" not in reasoning_text
def test_drops_in_tool_args_state(self):
"""Drop tokens in TOOL_ARGS state are silently discarded."""
@@ -50,7 +50,7 @@ def test_oot_registration_embedding(
m.setenv("VLLM_PLUGINS", "register_dummy_model")
prompts = ["Hello, my name is", "The text does not matter"]
llm = LLM(
model=dummy_gemma2_embedding_path, load_format="dummy", max_model_len=2048
model=dummy_gemma2_embedding_path, load_format="dummy", max_model_len=512
)
outputs = llm.embed(prompts)
+79
View File
@@ -5,13 +5,17 @@
Run `pytest tests/quantization/test_auto_gptq.py -v -s`.
"""
from types import SimpleNamespace
import pytest
import torch
from tests.quantization.utils import is_quant_method_supported
from vllm.model_executor.layers.fused_moe import RoutedExperts
from vllm.model_executor.layers.quantization.auto_gptq import (
AutoGPTQConfig,
AutoGPTQLinearMethod,
AutoGPTQMoEMethod,
)
PROMPT = "On the surface of Mars, we found"
@@ -54,3 +58,78 @@ def test_auto_gptq_quantization_method(vllm_runner, model_id: str, monkeypatch):
def test_auto_gptq_config_get_name():
"""Test that AutoGPTQConfig.get_name() returns 'auto_gptq'."""
assert AutoGPTQConfig.get_name() == "auto_gptq"
def test_auto_gptq_moe_creates_zero_initialized_expert_biases():
method = object.__new__(AutoGPTQMoEMethod)
method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
method.input_dtype = None
method.experts_cls = None
layer = torch.nn.Module()
method.create_weights(
layer=layer,
num_experts=2,
hidden_size=8,
intermediate_size_per_partition=4,
params_dtype=torch.float16,
intermediate_size_full=4,
weight_loader=lambda *args, **kwargs: None,
)
assert layer.w13_bias.shape == (2, 8)
assert layer.w2_bias.shape == (2, 8)
assert torch.count_nonzero(layer.w13_bias) == 0
assert torch.count_nonzero(layer.w2_bias) == 0
def test_routed_experts_loads_per_expert_biases():
class Loader:
quant_config = None
quant_method = object()
moe_config = SimpleNamespace(
is_act_and_mul=True,
tp_rank=0,
moe_parallel_config=SimpleNamespace(tp_size=1),
)
_get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim)
_narrow_expert_data_for_padding = staticmethod(
RoutedExperts._narrow_expert_data_for_padding
)
_load_w13 = RoutedExperts._load_w13
_loaded_expert_biases = set()
@staticmethod
def _map_global_expert_id_to_local_expert_id(expert_id):
return expert_id
loader = Loader()
w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False)
w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False)
for shard_id, loaded in (
("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])),
("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])),
):
assert RoutedExperts.weight_loader(
loader,
w13_bias,
loaded,
weight_name="model.layers.0.mlp.experts.w13_bias",
shard_id=shard_id,
expert_id=0,
return_success=True,
)
assert RoutedExperts.weight_loader(
loader,
w2_bias,
torch.tensor([9.0, 10.0, 11.0, 12.0]),
weight_name="model.layers.0.mlp.experts.w2_bias",
shard_id="w2",
expert_id=0,
return_success=True,
)
assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8))
assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4))
assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"}
+181 -32
View File
@@ -5,45 +5,194 @@ from types import SimpleNamespace
import pytest
import torch
from compressed_tensors.quantization import (
ActivationOrdering,
QuantizationArgs,
QuantizationStrategy,
QuantizationType,
)
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method
from vllm.platforms import current_platform
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
WNA16MoEBackend,
_backend_incompatibility_reason,
_convert_moe_wna16_humming_tensors,
convert_to_wna16_moe_kernel_format,
map_wna16_backend,
)
from vllm.model_executor.layers.quantization import moe_wna16
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
from vllm.model_executor.layers.quantization.moe_wna16 import (
MoeWNA16Config,
MoeWNA16Method,
)
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
def test_moe_wna16_apply_passes_layer_activation(monkeypatch):
captured_kwargs = {}
def test_map_wna16_backend_supports_triton():
assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON
def fake_fused_experts(*args, **kwargs):
captured_kwargs.update(kwargs)
return torch.empty(1, 2)
monkeypatch.setattr(
"vllm.model_executor.layers.fused_moe.fused_experts",
fake_fused_experts,
@pytest.mark.parametrize(
("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"),
[
(
WNA16MoEBackend.TRITON,
AutoAWQConfig(4, 128, True, False),
True,
False,
"AutoAWQ weight layout",
),
(
WNA16MoEBackend.TRITON,
AutoGPTQConfig(4, 128, True, True, False, {}, {}),
False,
False,
"activation ordering",
),
(
WNA16MoEBackend.TRITON,
QuantizationArgs(
num_bits=4,
type=QuantizationType.INT,
strategy=QuantizationStrategy.GROUP,
symmetric=True,
dynamic=False,
group_size=128,
actorder=ActivationOrdering.GROUP,
),
False,
False,
"activation ordering",
),
(
WNA16MoEBackend.TRITON,
AutoGPTQConfig(4, 128, False, True, False, {}, {}),
False,
True,
"bias",
),
(
WNA16MoEBackend.MARLIN,
MoeWNA16Config(
linear_quant_method="gptq",
weight_bits=4,
group_size=128,
has_zp=False,
lm_head_quantized=False,
modules_to_not_convert=None,
full_config={},
),
False,
False,
"MoeWNA16 checkpoint layout",
),
],
)
def test_wna16_oracle_rejects_incompatible_quant_structures(
backend, quant_config, may_have_zp, may_have_bias, expected
):
reason = _backend_incompatibility_reason(
backend=backend,
quant_config=quant_config,
may_have_zp=may_have_zp,
may_have_bias=may_have_bias,
)
assert reason is not None
assert expected in reason
def test_compressed_tensors_weights_are_transposed_for_triton():
quant_config = QuantizationArgs(
num_bits=4,
type=QuantizationType.INT,
strategy=QuantizationStrategy.GROUP,
symmetric=True,
dynamic=False,
group_size=32,
)
w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8)
w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6)
w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8)
w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6)
converted = convert_to_wna16_moe_kernel_format(
backend=WNA16MoEBackend.TRITON,
layer=torch.nn.Module(),
quant_config=quant_config,
input_dtype=None,
w13=w13,
w2=w2,
w13_scale=w13_scale,
w2_scale=w2_scale,
)
assert converted is not None
assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8))
assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8))
assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous())
assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous())
def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):
method = object.__new__(MoeWNA16Method)
method.experts_cls = object
method.wna16_backend = WNA16MoEBackend.HUMMING
method.moe = object()
quant_config = object()
method.get_fused_moe_quant_config = lambda layer: quant_config
layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None))
captured = {}
kernel = object()
def fake_make_wna16_moe_kernel(**kwargs):
captured.update(kwargs)
return kernel
monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel)
method._setup_kernel(layer)
assert method.moe_kernel is kernel
assert captured["backend"] == WNA16MoEBackend.HUMMING
assert captured["layer"] is layer
def test_moe_wna16_humming_adapter_repacks_uint8_tensors():
qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8)
scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4)
qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2)
converted = _convert_moe_wna16_humming_tensors(
{"qweight": qweight, "scales": scales, "qzeros": qzeros},
has_zero_point=True,
)
assert torch.equal(converted["weight"], qweight.view(torch.int32))
assert converted["weight"].shape == (1, 4, 2)
assert torch.equal(converted["weight_scale"], scales)
expected_qzeros = (
qzeros.transpose(-1, -2)
.contiguous()
.view(torch.int32)
.transpose(-1, -2)
.contiguous()
)
assert torch.equal(converted["zero_point"], expected_qzeros)
assert converted["zero_point"].shape == (1, 2, 2)
def test_moe_wna16_uses_humming_quant_config(monkeypatch):
from vllm.model_executor.layers.quantization.utils import humming_utils
method = object.__new__(MoeWNA16Method)
method.moe = SimpleNamespace(disable_inplace=False)
method.moe_quant_config = object()
layer = SimpleNamespace(
w13_qweight=torch.empty(1, 2),
w2_qweight=torch.empty(1, 2),
activation=MoEActivation.GELU_TANH,
apply_router_weight_on_input=False,
global_num_experts=1,
expert_map=None,
method.wna16_backend = WNA16MoEBackend.HUMMING
layer = object()
quant_config = object()
monkeypatch.setattr(
humming_utils,
"get_humming_moe_quant_config",
lambda actual_layer: quant_config if actual_layer is layer else None,
)
output = method.apply(
layer,
x=torch.empty(1, 2),
topk_weights=torch.empty(1, 1),
topk_ids=torch.empty(1, 1, dtype=torch.int32),
shared_experts=None,
shared_experts_input=None,
)
assert output.shape == (1, 2)
assert captured_kwargs["activation"] is MoEActivation.GELU_TANH
assert method.get_fused_moe_quant_config(layer) is quant_config
@@ -1,13 +1,55 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import torch
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
from vllm.model_executor.layers.quantization.utils import flashinfer_fp4_moe
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
prepare_nvfp4_moe_layer_for_fi_or_cutlass,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
align_trtllm_fp4_moe_hidden_dim_for_fi,
)
def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):
monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x)
num_experts = 3
layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False))
w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
w2_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
weight_scale = torch.ones(num_experts)
outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass(
backend=NvFp4MoeBackend.FLASHINFER_CUTLASS,
layer=layer,
w13=w13,
w13_scale=w13_scale,
w13_scale_2=weight_scale,
a13_scale=torch.tensor([1.0, 2.0, 3.0]),
w2=w2,
w2_scale=w2_scale,
w2_scale_2=weight_scale,
a2_scale=torch.tensor([4.0, 5.0, 6.0]),
is_act_and_mul=False,
)
a13_scale, a2_scale = outputs[3], outputs[7]
torch.testing.assert_close(a13_scale, torch.full((num_experts,), 3.0))
torch.testing.assert_close(a2_scale, torch.full((num_experts,), 6.0))
distinct_values = torch.arange(num_experts, dtype=torch.float32)
a13_scale.copy_(distinct_values)
a2_scale.copy_(distinct_values)
torch.testing.assert_close(a13_scale, distinct_values)
torch.testing.assert_close(a2_scale, distinct_values)
def test_align_trtllm_fp4_moe_hidden_dim_noop():
w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)
@@ -46,9 +46,12 @@ BACKENDS_TO_TEST = [
DEVICE_TYPE = current_platform.device_type
# Use the platform's preferred FP8 type so the stored cache matches what the
# backends reinterpret at runtime. On ROCm gfx94x this is e4m3fnuz, not e4m3fn;
# storing e4m3fn bytes there would be re-read as fnuz and produce NaNs.
FP8_KV_CACHE_DTYPES = {
"fp8": torch.float8_e4m3fn,
"fp8_e4m3": torch.float8_e4m3fn,
"fp8": current_platform.fp8_dtype(),
"fp8_e4m3": current_platform.fp8_dtype(),
}
# Remove flashinfer from the list if it's not available
+3 -1
View File
@@ -9,6 +9,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import RequestStatus
from vllm.v1.structured_output import StructuredOutputGrammar
from vllm.v1.utils import ConstantList
from .utils import create_requests, create_scheduler
@@ -262,7 +263,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler = object.__new__(AsyncScheduler)
request = create_requests(num_requests=1, num_tokens=1)[0]
request.structured_output_request = Mock()
request.structured_output_request.grammar = Mock()
request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
request.structured_output_request.grammar.accept_tokens.return_value = False
request.status = RequestStatus.RUNNING
request.num_computed_tokens = request.num_tokens
@@ -284,6 +285,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler.kv_event_publisher = Mock()
scheduler.finished_req_ids = set()
scheduler.finished_req_ids_dict = None
scheduler.grammar_compile_error_reqs = set()
scheduler.vllm_config = Mock()
scheduler.vllm_config.model_config.enable_return_routed_experts = False
scheduler.enable_return_routed_experts = False
+168 -20
View File
@@ -9,6 +9,7 @@ import torch
from vllm.v1.core.kv_cache_utils import (
_get_kv_cache_config_packed,
_get_kv_cache_groups_uniform_groups,
get_kv_cache_config_from_groups,
)
from vllm.v1.kv_cache_interface import (
@@ -16,6 +17,7 @@ from vllm.v1.kv_cache_interface import (
KVCacheGroupSpec,
KVCacheTensor,
MLAAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
UniformTypeKVCacheSpecs,
)
@@ -109,7 +111,132 @@ def _page_sizes_by_layer(
return page_sizes
def _packing_by_layer(
tensors: list[KVCacheTensor],
) -> dict[str, tuple[int, int]]:
return {
layer_name: (tensor.offset, tensor.block_stride)
for tensor in tensors
for layer_name in tensor.shared_by
}
def _make_views(
groups: list[KVCacheGroupSpec],
num_blocks: int,
tensors: list[KVCacheTensor],
) -> dict[str, torch.Tensor]:
page_sizes = _page_sizes_by_layer(groups)
packing = _packing_by_layer(tensors)
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
return {
layer_name: torch.as_strided(
backing,
size=(num_blocks, page_size),
stride=(packing[layer_name][1], 1),
storage_offset=packing[layer_name][0],
)
for layer_name, page_size in page_sizes.items()
}
def _make_page_group(prefix: str, page_sizes: list[int]) -> KVCacheGroupSpec:
specs = {
f"{prefix}.{i}": MagicMock(page_size_bytes=page_size)
for i, page_size in enumerate(page_sizes)
}
return KVCacheGroupSpec(
layer_names=list(specs),
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=specs),
)
class TestInterleavedPacking:
def test_compact_cache_overlays_fp32_state_group(self):
full_specs = {}
state_specs = {}
for i in range(2):
full_specs[f"mla.{i}"] = MLAAttentionSpec(
block_size=256,
num_kv_heads=1,
head_size=512,
dtype=torch.uint8,
page_size_padded=32768,
indexes_kv_by_block_stride=True,
compress_ratio=4,
)
full_specs[f"indexer.{i}"] = MLAAttentionSpec(
block_size=256,
num_kv_heads=1,
head_size=68,
dtype=torch.uint8,
page_size_padded=4608,
compress_ratio=4,
)
state_specs[f"mla_state.{i}"] = SlidingWindowMLASpec(
block_size=4,
num_kv_heads=1,
head_size=2048,
dtype=torch.float32,
sliding_window=8,
indexes_kv_by_block_stride=True,
)
state_specs[f"indexer_state.{i}"] = SlidingWindowMLASpec(
block_size=4,
num_kv_heads=1,
head_size=512,
dtype=torch.float32,
sliding_window=8,
indexes_kv_by_block_stride=True,
)
grouped_specs = [
UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=full_specs),
UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=state_specs),
]
groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
assert len(groups) == 2
assert {full_specs[f"indexer.{i}"].page_size_bytes for i in range(2)} == {4608}
assert {full_specs[f"indexer.{i}"].real_page_size_bytes for i in range(2)} == {
4352
}
assert {
state_specs[f"indexer_state.{i}"].page_size_bytes for i in range(2)
} == {8192}
full_group_bytes = 2 * (32768 + 4608)
state_group_bytes = 2 * (32768 + 8192)
bytes_per_block = max(full_group_bytes, state_group_bytes)
num_blocks, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, bytes_per_block * 32
)
assert num_blocks == 32
assert {tensor.block_stride for tensor in tensors} == {bytes_per_block}
packing = _packing_by_layer(tensors)
assert packing["mla.0"][0] == packing["mla_state.0"][0] == 0
assert packing["indexer.0"][0] == 32768
assert packing["indexer_state.0"][0] == 32768
def test_deepseek_v4_pro_stride(self):
groups = [
_make_page_group("full", [32768, 4608] * 30 + [1024] * 31),
_make_page_group("c4_state", [32768, 8192] * 30),
_make_page_group("c128_state", [32768] * 31),
_make_page_group("swa.0", [32768] * 31),
_make_page_group("swa.1", [32768] * 30),
]
expected_stride = 1_228_800
num_blocks, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, expected_stride * 32
)
assert num_blocks == 32
assert {tensor.block_stride for tensor in tensors} == {expected_stride}
assert {tensor.size for tensor in tensors} == {expected_stride * 32}
def test_all_tensors_have_block_stride(self):
_, tensors = _run()
for t in tensors:
@@ -122,9 +249,30 @@ class TestInterleavedPacking:
assert sizes.pop() > 0
def test_offsets_within_one_block(self):
_, tensors = _run()
for t in tensors:
assert t.offset < t.block_stride
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
_, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
page_sizes = _page_sizes_by_layer(groups)
packing = _packing_by_layer(tensors)
for layer_name, page_size in page_sizes.items():
offset, block_stride = packing[layer_name]
assert offset + page_size <= block_stride
def test_layouts_are_disjoint_within_each_group(self):
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
_, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
page_sizes = _page_sizes_by_layer(groups)
packing = _packing_by_layer(tensors)
for group in groups:
ranges = sorted(
(packing[name][0], packing[name][0] + page_sizes[name])
for name in group.layer_names
)
assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
def test_all_layers_accounted_for(self):
n_c4, n_c128, n_swa = 5, 4, 7
@@ -135,29 +283,29 @@ class TestInterleavedPacking:
expected = n_c4 * 2 + n_c128 + n_swa
assert len(all_names) == expected
def test_strided_views_are_independent(self):
def test_group_owned_blocks_do_not_alias(self):
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
page_sizes = _page_sizes_by_layer(groups)
num_blocks, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
views = []
for t in tensors:
page_size = page_sizes[t.shared_by[0]]
v = torch.as_strided(
backing,
size=(num_blocks, page_size),
stride=(t.block_stride, 1),
storage_offset=t.offset,
)
views.append(v)
views = _make_views(groups, num_blocks, tensors)
for i, v in enumerate(views):
v.fill_(i + 1)
expected = {}
value = 1
for block_id, group in enumerate(groups):
for layer_name in group.layer_names:
views[layer_name][block_id].fill_(value)
expected[layer_name] = (block_id, value)
value += 1
for i, v in enumerate(views):
assert (v == i + 1).all(), f"View {i} was corrupted"
for layer_name, (block_id, value) in expected.items():
assert (views[layer_name][block_id] == value).all()
# Once the first group releases its block, another group may reuse it.
for layer_name in groups[1].layer_names:
views[layer_name][0].fill_(255)
for layer_name in groups[1].layer_names:
assert (views[layer_name][0] == 255).all()
def test_hma_attention_groups_keep_default_backing(self):
full = _make_full_spec()
+388 -2
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import dataclasses
from concurrent.futures import Future
from unittest.mock import Mock
import pytest
@@ -25,18 +26,21 @@ from vllm.multimodal.inputs import (
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
from vllm.utils.hashing import sha256
from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
from vllm.v1.engine import FinishReason
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
MambaSpec,
)
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
from vllm.v1.structured_output import StructuredOutputManager
from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager
from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
@@ -1068,6 +1072,141 @@ def test_preempt_during_execution():
assert requests[1].output_token_ids[0] == 42
def test_prefix_cache_query_not_inflated_by_connector_defer():
"""The GPU prefix-cache query is recorded at admission, so a request the
connector defers several times is counted once, not once per retry."""
num_defers_before_matching = 3
scheduler = create_scheduler(
enable_prefix_caching=True,
use_kv_connector=mock_kv(
matched_tokens=0,
is_async=False,
num_defers_before_matching=num_defers_before_matching,
),
)
request = create_requests(num_requests=1, num_tokens=32, block_size=16)[0]
scheduler.add_request(request)
# Each deferred step re-runs the lookup but records nothing.
for _ in range(num_defers_before_matching):
assert not scheduler.schedule().scheduled_new_reqs
output = scheduler.schedule()
assert any(r.req_id == request.request_id for r in output.scheduled_new_reqs)
stats = scheduler.kv_cache_manager.prefix_cache_stats
assert stats is not None
assert stats.requests == 1
assert stats.queries == request.num_tokens
def test_preemption_re_records_prefix_cache_query():
"""A preempted request re-enters the lookup on resume, so its recomputation
is counted again into the preempted stats."""
scheduler = create_scheduler(enable_prefix_caching=True)
request = create_requests(num_requests=1)[0]
scheduler.add_request(request)
scheduler.schedule()
stats = scheduler.kv_cache_manager.prefix_cache_stats
assert stats is not None
assert (stats.requests, stats.preempted_requests) == (1, 0)
scheduler.running.remove(request)
scheduler._preempt_request(request, 0.0)
assert request.status == RequestStatus.PREEMPTED
scheduler.schedule()
assert request.status == RequestStatus.RUNNING
assert stats.preempted_requests == 1
def test_prefix_cache_stats_not_recorded_when_caching_disabled():
"""With prefix caching off there is no local lookup, so admitting a request
records no phantom miss."""
scheduler = create_scheduler(enable_prefix_caching=False)
for request in create_requests(num_requests=2):
scheduler.add_request(request)
scheduler.schedule()
stats = scheduler.kv_cache_manager.prefix_cache_stats
assert stats is not None
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request():
"""A real cache hit rejected once by allocate_slots and admitted on the next
step is counted exactly once, hits included."""
block_size = 16
scheduler = create_scheduler(
enable_prefix_caching=True,
enable_chunked_prefill=False,
block_size=block_size,
)
# Seed the cache so the next request with the same prompt hits it.
seed = create_requests(
num_requests=1,
num_tokens=block_size * 2,
max_tokens=2,
same_prompt=True,
block_size=block_size,
req_ids=["seed"],
)[0]
scheduler.add_request(seed)
_step_until_done(
scheduler,
scheduler.schedule(),
ModelRunnerOutput(
req_ids=["seed"],
req_id_to_index={"seed": 0},
sampled_token_ids=[[1000]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
),
)
# The seeding step swapped in a fresh accumulator, so re-read it; the
# retried request must be the only thing recorded from here on.
stats = scheduler.kv_cache_manager.prefix_cache_stats
assert stats is not None
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
retried = create_requests(
num_requests=1,
num_tokens=block_size * 3,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["retried"],
)[0]
scheduler.add_request(retried)
# Reject the first allocation attempt, then delegate to the real one.
orig_allocate_slots = scheduler.kv_cache_manager.allocate_slots
allocate_results: list = []
def spy_allocate_slots(*args, **kwargs):
result = None if not allocate_results else orig_allocate_slots(*args, **kwargs)
allocate_results.append(result)
return result
scheduler.kv_cache_manager.allocate_slots = spy_allocate_slots
assert not scheduler.schedule().scheduled_new_reqs
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
assert "retried" in scheduler.schedule().num_scheduled_tokens
assert allocate_results[0] is None and allocate_results[1] is not None
assert (stats.requests, stats.queries, stats.hits) == (
1,
retried.num_tokens,
block_size * 2,
)
def test_scheduler_reset_prefix_cache():
scheduler = create_scheduler(enable_prefix_caching=True)
requests = create_requests(num_requests=10)
@@ -3006,6 +3145,58 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
assert len(scheduler.skipped_waiting) == 1
@pytest.mark.parametrize("async_grammar", [True, False])
def test_grammar_compile_error_finishes_only_request(async_grammar: bool):
scheduler = create_scheduler()
manager = scheduler.structured_output_manager
manager.backend = Mock()
manager.backend.compile_grammar.side_effect = RuntimeError(
"forced FSM compilation error"
)
manager._use_async_grammar_compilation = async_grammar
sampling_params = SamplingParams(
max_tokens=16,
structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),
)
sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
request = Request(
request_id="grammar-error",
prompt_token_ids=[0, 1],
sampling_params=sampling_params,
pooling_params=None,
)
manager.grammar_init(request)
assert request.structured_output_request is not None
grammar_future = request.structured_output_request._grammar
assert isinstance(grammar_future, Future)
assert isinstance(grammar_future.exception(timeout=5), RuntimeError)
scheduler.add_request(request)
scheduler_output = scheduler.schedule()
assert not scheduler_output.num_scheduled_tokens
engine_core_outputs = scheduler.update_from_output(
scheduler_output,
ModelRunnerOutput(req_ids=[], req_id_to_index={}),
)
assert request.status == RequestStatus.FINISHED_ERROR
assert request.request_id not in scheduler.requests
output = engine_core_outputs[0].outputs[0]
assert output.request_id == request.request_id
assert output.finish_reason == FinishReason.ERROR
assert output.stop_reason is None
healthy_request = create_requests(num_requests=1, req_ids=["healthy-request"])[0]
scheduler.add_request(healthy_request)
next_output = scheduler.schedule()
assert [req.req_id for req in next_output.scheduled_new_reqs] == [
healthy_request.request_id
]
def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler = object.__new__(Scheduler)
sampling_params = SamplingParams(ignore_eos=True, max_tokens=4)
@@ -3019,7 +3210,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
pooling_params=None,
)
request.structured_output_request = Mock()
request.structured_output_request.grammar = Mock()
request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
request.structured_output_request.grammar.accept_tokens.return_value = False
request.status = RequestStatus.RUNNING
request.num_computed_tokens = request.num_tokens
@@ -3040,6 +3231,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler.kv_event_publisher = Mock()
scheduler.finished_req_ids = set()
scheduler.finished_req_ids_dict = None
scheduler.grammar_compile_error_reqs = set()
scheduler.vllm_config = Mock()
scheduler.vllm_config.model_config.enable_return_routed_experts = False
scheduler.enable_return_routed_experts = False
@@ -5329,3 +5521,197 @@ def test_async_load_reservation_prevents_wedge_e2e():
assert b.status == RequestStatus.WAITING
assert b.num_preemptions == 0
assert b.request_id not in req_to_blocks
def _create_hybrid_mamba_connector_scheduler(
matched_tokens: int,
block_size: int = 16,
num_blocks: int = 100,
) -> Scheduler:
"""FA + Mamba ("all" cache mode) scheduler with a MockKVConnector."""
model_config = ModelConfig(
model="facebook/opt-125m",
trust_remote_code=True,
dtype="float16",
seed=42,
skip_tokenizer_init=True,
)
vllm_config = VllmConfig(
scheduler_config=SchedulerConfig(
max_num_seqs=4,
max_num_batched_tokens=8192,
max_model_len=8192,
enable_chunked_prefill=True,
is_encoder_decoder=False,
watermark=0.0,
),
model_config=model_config,
cache_config=CacheConfig(
block_size=block_size,
enable_prefix_caching=True,
mamba_cache_mode="all",
),
kv_transfer_config=KVTransferConfig(
kv_connector="MockKVConnector",
kv_role="kv_both",
kv_connector_extra_config={
"matched_tokens": matched_tokens,
"is_async": False,
},
),
)
vllm_config.cache_config.num_gpu_blocks = num_blocks
kv_cache_config = KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["fa"],
FullAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
),
KVCacheGroupSpec(
["mamba"],
MambaSpec(
block_size=block_size,
shapes=((1, 1),),
dtypes=(torch.float32,),
mamba_cache_mode="all",
),
),
],
)
register_all_kvcache_specs(vllm_config)
return Scheduler(
vllm_config=vllm_config,
kv_cache_config=kv_cache_config,
structured_output_manager=StructuredOutputManager(vllm_config),
block_size=block_size,
hash_block_size=block_size,
log_stats=True,
)
@pytest.mark.parametrize(
"matched_tokens,expected_num_computed",
[
# No external hit: resume on the deepest locally-consistent boundary
# (block 0's state survives for both groups).
(0, 16),
# One external block on top of the reconciled local boundary.
(16, 32),
],
)
def test_hybrid_per_group_hit_divergence_with_connector(
matched_tokens: int, expected_num_computed: int
):
"""Per-group prefix hits can diverge for hybrid models with a connector
(#46453): under block pressure the FA prefix tail is evicted while a
deeper Mamba state block survives. The scheduler must not report the
deeper hit as locally computed (evicted FA blocks are not resident ->
engine crash / dirty KV); it falls back to the reconciled boundary that
every group is consistent at.
"""
block_size = 16
scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens)
manager = scheduler.kv_cache_manager
assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
# Seed a 4-block prefix so both groups cache all four boundaries
# (mamba cache mode "all" caches every block's state densely).
[fill] = create_requests(
num_requests=1,
num_tokens=4 * block_size,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["fill"],
)
computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
blocks = manager.allocate_slots(
fill, fill.num_tokens, num_computed, computed_blocks
)
fa_ids = [b.block_id for b in blocks.blocks[0]]
mamba_ids = [b.block_id for b in blocks.blocks[1]]
manager.free(fill)
# Evict the FA tail and the middle mamba states; block 0 (both groups)
# and the deep mamba state at block 3 survive.
manager.block_pool.evict_blocks({fa_ids[2], fa_ids[3], mamba_ids[1], mamba_ids[2]})
# A replay of the prefix plus one extra block now sees diverged
# per-group hits: FA stops at the evicted tail, while the mamba lookup
# finds the deeper surviving state.
[replay] = create_requests(
num_requests=1,
num_tokens=5 * block_size,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["replay"],
)
_, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
replay.block_hashes, replay.num_tokens - 1
)
assert per_group_hits == (2 * block_size, 4 * block_size) # diverged
scheduler.add_request(replay)
output = scheduler.schedule()
num_scheduled = output.num_scheduled_tokens[replay.request_id]
assert replay.num_tokens - num_scheduled == expected_num_computed
def test_hybrid_per_group_hit_divergence_fa_deeper_no_external():
"""The opposite divergence: the FA prefix survives deeper than the Mamba
state and the connector supplies nothing (ext == 0). Reporting the deep FA
hit as locally computed would resume with no valid Mamba state at that
boundary (silent bad output). The scheduler must fall back to the
convergent boundary that every group agrees on (block 0's surviving state).
"""
block_size = 16
scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens=0)
manager = scheduler.kv_cache_manager
assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
# Seed a 4-block prefix in both groups.
[fill] = create_requests(
num_requests=1,
num_tokens=4 * block_size,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["fill"],
)
computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
blocks = manager.allocate_slots(
fill, fill.num_tokens, num_computed, computed_blocks
)
mamba_ids = [b.block_id for b in blocks.blocks[1]]
manager.free(fill)
# Keep all FA blocks; evict every mamba state but block 0. FA reaches 4
# blocks, the mamba hit only reaches 1 -> diverged (FA > Mamba).
manager.block_pool.evict_blocks({mamba_ids[1], mamba_ids[2], mamba_ids[3]})
[replay] = create_requests(
num_requests=1,
num_tokens=5 * block_size,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["replay"],
)
_, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
replay.block_hashes, replay.num_tokens - 1
)
assert per_group_hits == (4 * block_size, 1 * block_size) # FA deeper
scheduler.add_request(replay)
output = scheduler.schedule()
num_scheduled = output.num_scheduled_tokens[replay.request_id]
# Must resume at the convergent boundary (block 0), not the deep FA hit.
assert replay.num_tokens - num_scheduled == block_size
+9 -2
View File
@@ -38,8 +38,12 @@ from vllm.v1.structured_output import StructuredOutputManager
EOS_TOKEN_ID = 50256
def mock_kv(matched_tokens: int, is_async: bool):
return MockKVConfig(matched_tokens=matched_tokens, is_async=is_async)
def mock_kv(matched_tokens: int, is_async: bool, num_defers_before_matching: int = 0):
return MockKVConfig(
matched_tokens=matched_tokens,
is_async=is_async,
num_defers_before_matching=num_defers_before_matching,
)
def create_scheduler(
@@ -115,6 +119,9 @@ def create_scheduler(
kv_connector_extra_config={
"matched_tokens": use_kv_connector.matched_tokens,
"is_async": use_kv_connector.is_async,
"num_defers_before_matching": (
use_kv_connector.num_defers_before_matching
),
},
)
elif isinstance(use_kv_connector, str):

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