Compare commits

..
Author SHA1 Message Date
khluu b7a9ca4ebf 2.24.0
Signed-off-by: khluu <khluu000@gmail.com>
2026-03-06 10:55:24 -08:00
khluu e9b140aea6 upgrade openai
Signed-off-by: khluu <khluu000@gmail.com>
2026-03-06 02:20:55 -08:00
lailooandGitHub 097eb544e9 [Bugfix] Improve engine ready timeout error message (#35616)
Signed-off-by: damaozi <1811866786@qq.com>
2026-03-04 05:54:32 +00:00
7cdba98edf [BugFix] Support tool_choice=none in the Anthropic API (#35835)
Signed-off-by: ZhongsJie <zhongsjie@gmail.com>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
2026-03-04 05:24:46 +00:00
Charlie FuandGitHub 3c85cd9d74 [Rocm][CI] Fix ROCm LM Eval Large Models (8 Card) (#35913)
Signed-off-by: charlifu <charlifu@amd.com>
2026-03-04 04:50:13 +00:00
Andreas KaratzasandGitHub edba15045a [Bugfix] Guard mm_token_type_ids kwarg in get_mrope_input_positions (#35711)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-04 04:12:51 +00:00
Cyrus LeungandGitHub e379396167 [Refactor] Clean up processor kwargs extraction (#35872)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
2026-03-03 19:53:53 -08:00
Isotr0pyandGitHub 6e9f21e8a2 [Chore] Remove debug code in model implementation (#35883)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-03 19:50:58 -08:00
c1d963403c [model] support FireRedASR2 (#35727)
Signed-off-by: zixiao <shunli.dsl@alibaba-inc.com>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: zixiao <shunli.dsl@alibaba-inc.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-03 19:41:30 -08:00
Shanshan ShenandGitHub 77e6dcbbfa [PluggableLayer][MM] Add PluggableLayer for RelPosAttention (#33753)
Signed-off-by: shen-shanshan <467638484@qq.com>
2026-03-03 19:41:27 -08:00
William ZhangandGitHub 70c73df69e [Bugfix] Fix EVS implementation for Qwen3 VL (#33607)
Signed-off-by: 2ez4bz <133824995+2ez4bz@users.noreply.github.com>
2026-03-04 02:18:11 +00:00
9a9d442464 Enable bnb for multiple indices weight (#35838)
Signed-off-by: xjx <493337577@qq.com>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-04 01:46:47 +00:00
Andreas KaratzasandGitHub f7da9cdffc [ROCm][CI] Support async weight transfer example with platform-aware determinism (#35710)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-04 09:44:14 +08:00
JaewonandGitHub f22ff2958c [Bugfix] Fix coord_socket assertion in DPEngineCoreProc for offline DP mode (#35916)
Signed-off-by: Jaewon Lee <jaewon@meta.com>
2026-03-04 00:10:11 +00:00
Nick HillandGitHub d15c3b90fc [Core] Move save_tensorized_model logic to Worker (#35825)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-03-03 15:31:59 -08:00
97286a20ed [Model Runner V2] support dp & ep for spec decoding (#35294)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai>
2026-03-03 15:19:45 -08:00
Amr MahdiandGitHub 12b38c0f45 [CI/Build] Allow mounting AWS credentials for sccache S3 auth (#35912)
Signed-off-by: Amr Mahdi <amrmahdi@meta.com>
2026-03-03 14:30:47 -08:00
Woosuk KwonandGitHub 467886a0c4 [Model Runner V2] Fix inputs_embeds=None bug for MM models (#35917)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-03-03 13:47:45 -08:00
a9b8b13e5c [Bugfix] Fix misnamed parameter in compressed_tensors_moe.py (#35813)
Signed-off-by: Bill Nell <bnell@redhat.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
2026-03-03 16:29:57 -05:00
Micah WilliamsonandGitHub e7213003cb [ROCm][CI] Fix TP size issue for test_gpt_oss (#35887)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-03-03 20:57:34 +00:00
Rohan PotdarandGitHub 3a8eef5869 [ROCm][Bugfix]: Disable AITER Triton ROPE by default (#35601)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
2026-03-03 13:43:56 -06:00
97995f6376 [MoE Refactor] Create MK for TRTLLM Kernels (#32564)
Signed-off-by: Robert Shaw <robshaw@redhat.com>
Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
Co-authored-by: Robert Shaw <robshaw@redhat.com>
Co-authored-by: Robert Shaw <rshaw@neuralmagic.com>
2026-03-03 10:39:50 -08:00
Robert ShawandGitHub 881a6b011b [CI] Temporarily Disable Llama4 MoE Refactor Test (#35870)
Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
2026-03-03 10:36:15 -08:00
Matthew BonanniandGitHub 8e1fd5baf0 [CI] Bump num_speculative_tokens to 3 in nightly DeepSeek tests (#35882)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
2026-03-03 09:26:44 -08:00
ae88468bcc fix: Ensure invalid audio files return 400 error (#34715)
Signed-off-by: Jason Ozuzu <jasonozuzu@cohere.com>
Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
2026-03-03 08:47:39 -08:00
e05cb3b93e TRTLLM gen-full attn Test Coverage (#34986)
Signed-off-by: Anshika Ojha <anshikao@nvidia.com>
Co-authored-by: Anshika Ojha <anshikao@gb-nvl-059-compute09.nvidia.com>
2026-03-03 11:35:34 -05:00
28ef9ba399 [BugFix] Add support for MTP num_speculative_tokens > 1 with sparse MLA (#34552)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-03-03 07:21:57 -08:00
fb7fdc49c4 [ROCm] [CI] Add new fusion test cases that are relevant to vLLM IR Ops (#34307)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com>
2026-03-03 06:24:21 -08:00
ea463978bb [Frontend][1/n] Improve pooling entrypoints | classify. (#35604)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
Signed-off-by: wang.yuqi <noooop@126.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
2026-03-03 06:05:36 -08:00
Li, JiangandGitHub 440f0e7dc6 [Bugfix] Avoid src/dst as None in irecv/isend_tensor_dict (#35754)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
2026-03-03 05:56:08 -08:00
wang.yuqiGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
fd4a90f337 [CI] And PPL test for Qwen3.5. (#35853)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
Signed-off-by: wang.yuqi <noooop@126.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-03-03 13:15:51 +00:00
Thomas ParnellandGitHub ad9d09e2b8 [Perf] [Hybrid] Copy num_accepted_tokens in non-blocking way when not using prefix caching (#35442)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
2026-03-03 04:15:43 -08:00
4beebfd146 [CI/Build][Intel] Add new performance benchmarks for Intel Gaudi 3 (#31025)
Signed-off-by: Szymon Reginis <sreginis@habana.ai>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-03-03 19:48:24 +08:00
halleriteandGitHub b8401cde0e add regression test (#35834)
Signed-off-by: hallerite <git@hallerite.com>
2026-03-03 07:32:15 +00:00
TJianandGitHub 5dfc5abe94 [ROCm] [Release] Change the package from aiter to amd-aiter (#35198)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-03-02 23:13:39 -08:00
lin-shhandGitHub 8fa68a8ce4 Fix TYPE_CHECKING stub defaults in envs.py to match actual runtime defaults (#35645) 2026-03-02 21:59:43 -08:00
lin-shhandGitHub 35a6f0bfe2 [Misc] Fix typos in comments: explict→explicit, paramaters→parameters (#35648) 2026-03-02 21:59:14 -08:00
Taneem IbrahimandGitHub 3a6cbf16e2 [MISC] Removed unused function find_all_indices() from tool_parsers/utils.py (#35683)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
2026-03-03 13:58:42 +08:00
f44d1ddc8c [BugFix] Fix cmake based incremental install (wrong vllm install dir) (#35773)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
2026-03-02 21:58:16 -08:00
Cyrus LeungandGitHub 48a54c1e0d [CI/Build] Trigger processor tests on registry update (#35824)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
2026-03-03 13:55:57 +08:00
Micah WilliamsonandGitHub 8b9e8b7454 [ROCm][CI] Fix Assertion Logic For test_gpt_oss (#35806)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-03-03 05:08:04 +00:00
c21d0039ec [Refactor] Fix maxsim cuda platform and add cli to control it (#35427)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-03-03 12:48:31 +08:00
Isotr0pyandGitHub 7d8bbe6f42 [CI/Build] Automatically patch video metadata for multimodal processor test (#35822)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-03 04:27:45 +00:00
aykoppolandGitHub 25e02647c2 [Core] Add optional flags to check for repetitive token patterns in engine output (#35451)
Signed-off-by: aykoppol <aykoppol+git@gmail.com>
2026-03-03 12:23:25 +08:00
Woosuk KwonandGitHub a0a5178ab4 [Model Runner V2] Use ModelState.prepare_attn() for cuda graph capture [5/N] (#35774)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-03-02 20:06:27 -08:00
Isotr0pyandGitHub 8ea8ba275e [V0 deprecation] Remove Swin model (#35821)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-02 20:03:41 -08:00
Woosuk KwonandGitHub 4f85bae9d6 [Docs][Model Runner V2] Add Design Docs (#35819)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-03-02 19:58:14 -08:00
Andy LoandGitHub 0a7165fd71 [ModelRunnerV2] Rename sampler functions and variables for clarity (#35459)
Signed-off-by: Andy Lo <andy@mistral.ai>
2026-03-02 19:48:56 -08:00
6521ccf286 [CI] Temporarily Disable Nightly Failures (#35770)
Signed-off-by: Robert Shaw <robshaw@redhat.com>
Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Co-authored-by: Robert Shaw <robshaw@redhat.com>
2026-03-03 01:49:13 +00:00
8ebd872f50 [Tool Parser] Fix Qwen3Coder streaming parameter loss with speculative decode (#35615)
Signed-off-by: Martin Vit <martin@voipmonitor.org>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-03 09:40:37 +08:00
zhrrrandGitHub 168ee03e1c [Model Runner V2][Perf] align dummy_run tokens to uniform decode for dp cudagraph (#35376)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
2026-03-02 17:10:47 -08:00
9dd656f0ea [XPU][NIXL] Add GPUDirect RDMA support for XPU (#35270)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-03-03 08:42:49 +08:00
Jakub ZakrzewskiandGitHub c8b678e53e [Model] Add support for nvidia/llama-nemotron-rerank-vl-1b-v2 (#35735)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com>
2026-03-03 08:32:14 +08:00
Andreas KaratzasandGitHub 18c29c746b [ROCm][CI] Fix backslash-continuation in pytest marker re-quoting and treat exit code 5 as success (#35798)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-02 16:07:51 -08:00
Hanjie QiuandGitHub 96fc09503a [All Reduce] Change default backend of Flashinfer All Reduce to trtllm (#35793)
Signed-off-by: hjjq <hanjieq@nvidia.com>
2026-03-02 18:57:38 -05:00
Roger WangandGitHub 1b82b433fc [Bugfix] Fix MM processor test for Qwen3.5 (#35797)
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-03-02 23:05:08 +00:00
9319044ee9 [MoE][Perf] Wrap DSV3 QKVAProj GEMM in custom op for torch.compile (#35751)
Signed-off-by: Robert Shaw <robshaw@redhat.com>
Co-authored-by: Robert Shaw <robshaw@redhat.com>
2026-03-02 23:03:49 +00:00
Boyuan FengandGitHub c42dc402c1 clean unused cudagraph_batch_sizes (#35552)
Signed-off-by: Boyuan Feng <boyuan@meta.com>
2026-03-02 22:00:16 +00:00
Ye (Charlotte) QiandGitHub fa6a6be519 [Bugfix] Fix missing sequence_lengths in qwen3_omni_moe_thinker (#35741)
Signed-off-by: Ye (Charlotte) Qi <yeq@meta.com>
2026-03-02 21:11:56 +00:00
216 changed files with 21501 additions and 4530 deletions
@@ -1,2 +1 @@
Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml
Qwen3-235B-A22B-Instruct-2507-FP8.yaml
@@ -51,5 +51,56 @@
"max-model-len": 256,
"async-scheduling": ""
}
},
{
"test_name": "latency_deepseek_r1",
"environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"parameters": {
"model": "deepseek-ai/DeepSeek-R1",
"tensor_parallel_size": 8,
"load_format": "dummy",
"max-model-len": 2048,
"dtype": "bfloat16"
}
},
{
"test_name": "latency_llama4_maverick_17b128e_instruct_fp8",
"environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"parameters": {
"model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
"tensor_parallel_size": 8,
"max-model-len": 512,
"max-num-seqs": 128,
"async-scheduling": "",
"gpu-memory-utilization": 0.95,
"enable_expert_parallel": ""
}
},
{
"test_name": "latency_qwen3_8b",
"environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"parameters": {
"model": "Qwen/Qwen3-8B",
"tensor_parallel_size": 1,
"max-model-len": 2048,
"max-num-seqs": 128,
"dtype": "bfloat16",
"async-scheduling": ""
}
}
]
@@ -78,5 +78,84 @@
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
"num_prompts": 200
}
},
{
"test_name": "serving_deepseek_r1",
"qps_list": [1, 4, 16, "inf"],
"server_environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"server_parameters": {
"model": "deepseek-ai/DeepSeek-R1",
"tensor_parallel_size": 8,
"swap_space": 16,
"disable_log_stats": "",
"load_format": "dummy",
"max-model-len": 2048,
"max-num-seqs": 200,
"async-scheduling": "",
"dtype": "bfloat16"
},
"client_parameters": {
"model": "deepseek-ai/DeepSeek-R1",
"backend": "vllm",
"dataset_name": "sharegpt",
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
"num_prompts": 200
}
},
{
"test_name": "serving_llama4_maverick_17b128e_instruct_fp8",
"qps_list": [1, 4, 16, "inf"],
"server_environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"server_parameters": {
"model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
"tensor_parallel_size": 8,
"disable_log_stats": "",
"max-model-len": 2048,
"max-num-seqs": 128,
"async-scheduling": "",
"enable_expert_parallel": "",
"max-num-batched-tokens": 4096
},
"client_parameters": {
"model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
"backend": "vllm",
"dataset_name": "sharegpt",
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
"num_prompts": 200
}
},
{
"test_name": "serving_qwen3_8b",
"qps_list": [1, 4, 10, "inf"],
"server_environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"server_parameters": {
"model": "Qwen/Qwen-3-8B",
"tensor_parallel_size": 1,
"dtype": "bfloat16",
"disable_log_stats": "",
"async-scheduling": ""
},
"client_parameters": {
"model": "Qwen/Qwen-3-8B",
"backend": "vllm",
"dataset_name": "sharegpt",
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
"num_prompts": 200
}
}
]
@@ -57,5 +57,67 @@
"max-num-seqs": 512,
"async-scheduling": ""
}
},
{
"test_name": "throughput_deepseek_r1",
"environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"parameters": {
"model": "deepseek-ai/DeepSeek-R1",
"tensor_parallel_size": 8,
"load_format": "dummy",
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
"dataset_name": "sharegpt",
"num_prompts": 1000,
"backend": "vllm",
"max-model-len": 2048,
"max-num-seqs": 384,
"async-scheduling": ""
}
},
{
"test_name": "throughput_llama4_maverick_17b128e_instruct_fp8",
"environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"parameters": {
"model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
"tensor_parallel_size": 8,
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
"dataset_name": "sharegpt",
"num_prompts": 1000,
"backend": "vllm",
"max-model-len": 2048,
"max-num-seqs": 512,
"async-scheduling": "",
"enable_expert_parallel": ""
}
},
{
"test_name": "throughput_qwen3_8b",
"environment_variables": {
"PT_HPU_LAZY_MODE": 1,
"PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
"VLLM_CONTIGUOUS_PA": 1,
"VLLM_DEFRAG": 1
},
"parameters": {
"model": "Qwen/Qwen-3-8B",
"tensor_parallel_size": 1,
"load_format": "dummy",
"dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
"dataset_name": "sharegpt",
"num_prompts": 1000,
"max-num-seqs": 512,
"backend": "vllm",
"async-scheduling": ""
}
}
]
+2 -2
View File
@@ -68,7 +68,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl .
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl .
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl .
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl .
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amd_aiter-*.whl .
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl .
\`\`\`
@@ -80,7 +80,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-
- **torchvision**: TorchVision for ROCm PyTorch
- **torchaudio**: Torchaudio for ROCm PyTorch
- **amdsmi**: AMD SMI Python bindings
- **aiter**: Aiter for ROCm
- **amd_aiter**: Aiter for ROCm
- **flash-attn**: Flash Attention for ROCm
### :warning: Notes
+23 -2
View File
@@ -99,6 +99,15 @@ is_multi_node() {
return 1
}
handle_pytest_exit() {
local exit_code=$1
if [ "$exit_code" -eq 5 ]; then
echo "Pytest exit code 5 (no tests collected) - treating as success."
exit 0
fi
exit "$exit_code"
}
###############################################################################
# Pytest marker/keyword re-quoting
#
@@ -135,8 +144,9 @@ re_quote_pytest_markers() {
local collecting=false
local marker_buf=""
# Flatten newlines for consistent tokenization
local flat="${input//$'\n'/ }"
# Strip backslash-newline continuations, then flatten remaining newlines
local flat="${input//$'\\\n'/ }"
flat="${flat//$'\n'/ }"
# Disable globbing to prevent *.py etc. from expanding during read -ra
local restore_glob
@@ -164,6 +174,9 @@ re_quote_pytest_markers() {
local is_boundary=false
case "$word" in
# Line-continuation artifact
"\\")
is_boundary=true ;;
# Command separators
"&&"|"||"|";"|"|")
is_boundary=true ;;
@@ -204,6 +217,9 @@ re_quote_pytest_markers() {
if [[ "$word" == "-m" || "$word" == "-k" ]]; then
output+="${word} "
collecting=true
# Drop stray backslash tokens silently
elif [[ "$word" == "\\" ]]; then
:
else
output+="${word} "
fi
@@ -453,7 +469,9 @@ if is_multi_node "$commands"; then
done
/bin/bash -c "${composite_command}"
exit_code=$?
cleanup_network
handle_pytest_exit "$exit_code"
else
echo "Multi-node job detected but failed to parse bracket command syntax."
echo "Expected format: prefix ; [node0_cmd1, node0_cmd2] && [node1_cmd1, node1_cmd2]"
@@ -480,4 +498,7 @@ else
--name "${container_name}" \
"${image_name}" \
/bin/bash -c "${commands}"
exit_code=$?
handle_pytest_exit "$exit_code"
fi
@@ -1,26 +1,43 @@
#!/bin/bash
set -euox pipefail
export VLLM_CPU_CI_ENV=0
echo "--- PP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename tp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid &
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/tp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
echo "--- DP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename dp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid &
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/dp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
+118 -43
View File
@@ -610,6 +610,8 @@ steps:
--ignore=lora/test_qwen3moe_tp.py
parallelism: 4
##### .buildkite/test_areas/pytorch.yaml #####
# corresponds to .buildkite/test_areas/pytorch.yaml
- label: PyTorch Compilation Unit Tests # 15min
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction]
@@ -627,6 +629,20 @@ steps:
# they do not suffer from https://github.com/vllm-project/vllm/issues/28965
- "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
# corresponds to .buildkite/test_areas/pytorch.yaml
- label: PyTorch Compilation Passes Unit Tests
timeout_in_minutes: 20
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
source_file_dependencies:
- vllm/
- tests/compile/passes
commands:
# TODO: clean up this comment if not needed. It is used to
# keep track of the tests changes during vLLM IR Ops refactoring.
# Use `find` to launch multiple instances of pytest.
- "find compile/passes -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
- label: PyTorch Fullgraph Smoke Test # 15min
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction]
@@ -1044,6 +1060,7 @@ steps:
source_file_dependencies:
- vllm/
- tests/models/multimodal
- tests/models/registry.py
no_gpu: true
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
@@ -1057,6 +1074,7 @@ steps:
source_file_dependencies:
- vllm/
- tests/models/multimodal
- tests/models/registry.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/processing
@@ -1209,41 +1227,6 @@ steps:
- pytest -v -s tests/kernels/moe/test_flashinfer.py
- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
- label: Blackwell Fusion and Compile Tests # 30 min
timeout_in_minutes: 40
working_dir: "/vllm-workspace/"
gpu: b200
source_file_dependencies:
- csrc/quantization/fp4/
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
- vllm/v1/attention/backends/flashinfer.py
- vllm/v1/worker/
- vllm/v1/cudagraph_dispatcher.py
- vllm/compilation/
# can affect pattern matching
- vllm/model_executor/layers/layernorm.py
- vllm/model_executor/layers/activation.py
- vllm/model_executor/layers/quantization/input_quant_fp8.py
- tests/compile/passes/test_fusion_attn.py
- tests/compile/passes/test_silu_mul_quant_fusion.py
- tests/compile/passes/distributed/test_fusion_all_reduce.py
- tests/compile/fullgraph/test_full_graph.py
commands:
- nvidia-smi
- pytest -v -s tests/compile/passes/test_fusion_attn.py
- pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py
# this runner has 2 GPUs available even though num_gpus=2 is not set
- pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
# # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
# # Wrap with quotes to escape yaml
# - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
# Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
# in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
# test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
- pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
- label: Blackwell GPT-OSS Eval
timeout_in_minutes: 60
working_dir: "/vllm-workspace/"
@@ -1356,6 +1339,7 @@ steps:
- tests/v1/entrypoints/openai/test_multi_api_servers.py
- tests/v1/shutdown
- tests/v1/worker/test_worker_memory_snapshot.py
- examples/offline_inference/new_weight_syncing/
commands:
# Work around HIP bug tracked here: https://github.com/ROCm/hip/issues/3876
# TODO: Remove when the bug is fixed in a future ROCm release
@@ -1369,7 +1353,6 @@ steps:
- pytest -v -s ./compile/test_wrapper.py
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
- pytest -v -s compile/correctness_e2e/test_sequence_parallel.py
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
@@ -1599,16 +1582,16 @@ steps:
commands:
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
- pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
- pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
# TODO: this test is not supported on ROCm, there are aiter kernels for this.
# - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
#- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
# - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
# Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
# in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- pytest -v -s tests/distributed/test_context_parallel.py
- HIP_VISIBLE_DEVICES=0,1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
- pytest -v -s tests/v1/distributed/test_dbo.py
# this test is not supported on ROCm
# - pytest -v -s tests/v1/distributed/test_dbo.py
##### B200 test #####
- label: Distributed Tests (B200) # optional
@@ -1719,6 +1702,93 @@ steps:
commands:
- bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040
##### .buildkite/test_areas/compile.yaml #####
# Slowly setting up the tests so that it is also easier for the
# CI team to review and upstream to the pipelinev2.
# The following tests are important for vLLM IR Ops refactoring,
# which affects fusion passes on ROCm. So we have to
# enable them as as soon as possible.
## TODO: Enable the test in this group
# # corresponds to .buildkite/test_areas/compile.yaml
# - label: Fusion and Compile Unit Tests (2xMI325 GPUs)
# timeout_in_minutes: 20
# working_dir: "/vllm-workspace/"
# mirror_hardwares: [amdexperimental, amdproduction, tj]
# agent_pool: mi325_1 # changed to 1 GPU until the fusion all reduce is enabled then only revert back to 2 GPUs
# source_file_dependencies:
# - csrc/quantization/fp4/
# - vllm/model_executor/layers/quantization/
# - vllm/model_executor/layers/layernorm.py
# - vllm/model_executor/layers/activation.py
# - vllm/model_executor/layers/attention/attention.py
# - vllm/v1/attention/backends/flashinfer.py
# - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes
# - tests/compile/test_fusion_attn.py
# - tests/compile/test_silu_mul_quant_fusion.py
# - tests/compile/distributed/test_fusion_all_reduce.py
# - tests/compile/fullgraph/test_full_graph.py
# commands:
# - rocm-smi
# # we run all backend tests on ROCm
# # These two tests are covered in "PyTorch Compilation Passes Unit Tests"
# # - "pytest -v -s tests/compile/passes/test_fusion_attn.py"
# # - "pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py"
# # TODO: this test is not supported on ROCm, there are aiter kernels for this.
# # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
# # TODO: find out more details
# # - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
# corresponds to .buildkite/test_areas/compile.yaml
- label: Fusion E2E Quick (MI325)
timeout_in_minutes: 15
working_dir: "/vllm-workspace/"
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
num_devices: 1
source_file_dependencies:
- csrc/quantization/
- vllm/model_executor/
- vllm/v1/attention/
- vllm/compilation/
- tests/compile/fusions_e2e/
commands:
- rocm-smi
# Run all models and attn backends but only Inductor partition and native custom ops
- "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and not +rms_norm and not +quant_fp8'"
# Different from CUDA, Qwen requires +rms_norm and +quant_fp8 as rms+quant fusion is only supported on AITER
- "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and +rms_norm and +quant_fp8 and qwen3'"
# corresponds to .buildkite/test_areas/compile.yaml
- label: Fusion E2E Config Sweep (MI325)
timeout_in_minutes: 30
working_dir: "/vllm-workspace/"
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
num_devices: 1
source_file_dependencies:
- csrc/quantization/
- vllm/compilation/
# can affect pattern matching
- vllm/model_executor/layers/layernorm.py
- vllm/model_executor/layers/activation.py
- vllm/model_executor/layers/attention/attention.py
- vllm/model_executor/layers/quantization/input_quant_fp8.py
- tests/compile/fusions_e2e/
commands:
- rocm-smi
# Run just llama3 (fp8) for all config combinations
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
## There are no ops on ROCm for these tests.
## The test still passes but the logs are not useful.
## fused ops just call torch.ops.symm_mem which
## exists in ROCm even though they don't work
# - label: AsyncTP Correctness Tests (2xMI325 GPUs)
# - label: Fusion E2E TP2 Quick (MI325)
# - label: Fusion E2E TP2 AsyncTP Config Sweep (MI325)
# - label: Fusion E2E TP2 (MI325)
# - label: Sequence Parallel Correctness Tests (2xMI325 GPUs)
#####################################################################################################################################
@@ -1901,8 +1971,10 @@ steps:
- label: Distributed Tests (4 GPUs) # 35min
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental]
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_4
optional: true
# grade: Blocking
working_dir: "/vllm-workspace/tests"
num_gpus: 4
source_file_dependencies:
@@ -1956,7 +2028,8 @@ steps:
- popd
# NEW rlhf examples
- pushd ../examples/offline_inference/new_weight_syncing
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
- popd
@@ -2920,8 +2993,10 @@ steps:
- label: Distributed Tests (2 GPUs) # 68min
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental]
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_2
optional: true
# grade: Blocking
working_dir: "/vllm-workspace/tests"
num_gpus: 2
source_file_dependencies:
+1 -1
View File
@@ -146,7 +146,7 @@ steps:
num_devices: 2
commands:
- pytest -v -s tests/distributed/test_context_parallel.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py
# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py --- failing, need to re-enable
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
+2 -1
View File
@@ -44,7 +44,8 @@ steps:
- vllm/envs.py
- vllm/config
commands:
- pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 2
- label: Kernels Mamba Test
+11 -11
View File
@@ -11,17 +11,17 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
- label: LM Eval Large Models (4 GPUs)(A100)
device: a100
optional: true
num_devices: 4
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
# - label: LM Eval Large Models (4 GPUs)(A100)
# device: a100
# optional: true
# num_devices: 4
# working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
# source_file_dependencies:
# - csrc/
# - vllm/model_executor/layers/quantization
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
- label: LM Eval Large Models (4 GPUs)(H100)
device: h100
+2 -9
View File
@@ -20,6 +20,7 @@ steps:
source_file_dependencies:
- vllm/
- tests/models/multimodal
- tests/models/registry.py
device: cpu
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
@@ -30,6 +31,7 @@ steps:
source_file_dependencies:
- vllm/
- tests/models/multimodal
- tests/models/registry.py
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
@@ -70,12 +72,3 @@ steps:
commands:
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
# This test is used only in PR development phase to test individual models and should never run on main
- label: Custom Models
optional: true
commands:
- echo 'Testing custom models...'
# PR authors can temporarily add commands below to test individual models
# e.g. pytest -v -s models/encoder_decoder/vision_language/test_mllama.py
# *To avoid merge conflicts, remember to REMOVE (not just comment out) them before merging the PR*
+10 -10
View File
@@ -13,13 +13,13 @@ steps:
commands:
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
- label: Weight Loading Multiple GPU - Large Models # optional
working_dir: "/vllm-workspace/tests"
num_devices: 2
device: a100
optional: true
source_file_dependencies:
- vllm/
- tests/weight_loading
commands:
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
# - label: Weight Loading Multiple GPU - Large Models # optional
# working_dir: "/vllm-workspace/tests"
# num_devices: 2
# device: a100
# optional: true
# source_file_dependencies:
# - vllm/
# - tests/weight_loading
# commands:
# - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
+17 -11
View File
@@ -12,12 +12,12 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.kernels.moe.utils import make_dummy_moe_config
from vllm import _custom_ops as ops
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.platforms import current_platform
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.v1.worker.workspace import init_workspace_manager
@@ -137,15 +137,21 @@ def bench_run(
per_out_ch_quant=per_out_ch,
)
fn = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
moe_config = make_dummy_moe_config(
num_experts=num_experts,
hidden_dim=k,
intermediate_size_per_partition=n,
in_dtype=a.dtype,
)
fn = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
CutlassExpertsFp8(
moe_config=make_dummy_moe_config(
num_experts=num_experts,
hidden_dim=k,
intermediate_size_per_partition=n,
in_dtype=a.dtype,
),
moe_config=moe_config,
quant_config=quant_config,
),
)
@@ -15,6 +15,9 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.kernels.moe.utils import make_dummy_moe_config
from vllm import _custom_ops as ops
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
fp8_w8a8_moe_quant_config,
nvfp4_moe_quant_config,
@@ -23,9 +26,6 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
CutlassExpertsFp4,
)
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.scalar_type import scalar_types
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.v1.worker.workspace import init_workspace_manager
@@ -196,10 +196,21 @@ def bench_run(
g2_alphas=w2_gs,
)
kernel = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
moe_config = make_dummy_moe_config(
num_experts=num_experts,
hidden_dim=k,
intermediate_size_per_partition=n,
in_dtype=a.dtype,
)
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
CutlassExpertsFp4(
make_dummy_moe_config(),
moe_config=moe_config,
quant_config=quant_config,
),
)
@@ -240,11 +251,17 @@ def bench_run(
g1_alphas=w1_gs,
g2_alphas=w2_gs,
)
moe_config = make_dummy_moe_config()
kernel = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
CutlassExpertsFp4(
make_dummy_moe_config(),
moe_config=moe_config,
quant_config=quant_config,
),
)
@@ -9,15 +9,15 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.kernels.moe.utils import make_dummy_moe_config
from vllm import _custom_ops as ops
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.fused_moe import (
fused_experts,
fused_topk,
)
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.v1.worker.workspace import init_workspace_manager
@@ -131,16 +131,22 @@ def bench_run(
w2_scale=w2_scale,
per_act_token_quant=per_act_token,
)
moe_config = make_dummy_moe_config(
num_experts=w2.shape[0],
hidden_dim=w2.shape[1],
intermediate_size_per_partition=w2.shape[2],
in_dtype=a.dtype,
)
fn = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
fn = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
CutlassExpertsFp8(
moe_config=make_dummy_moe_config(
num_experts=w2.shape[0],
hidden_dim=w2.shape[1],
intermediate_size_per_partition=w2.shape[2],
in_dtype=a.dtype,
),
moe_config=moe_config,
quant_config=quant_config,
),
)
@@ -163,16 +169,22 @@ def bench_run(
w2_scale=w2_scale,
per_act_token_quant=per_act_token,
)
moe_config = make_dummy_moe_config(
num_experts=w2.shape[0],
hidden_dim=w2.shape[1],
intermediate_size_per_partition=w2.shape[2],
in_dtype=a.dtype,
)
fn = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
fn = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
CutlassExpertsFp8(
moe_config=make_dummy_moe_config(
num_experts=w2.shape[0],
hidden_dim=w2.shape[1],
intermediate_size_per_partition=w2.shape[2],
in_dtype=a.dtype,
),
moe_config=moe_config,
quant_config=quant_config,
),
)
+37 -17
View File
@@ -17,6 +17,9 @@ from ray.experimental.tqdm_ray import tqdm
from vllm.model_executor.layers.fused_moe import fused_topk
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
@@ -242,24 +245,33 @@ def benchmark_config(
deep_gemm_experts = None
if use_deep_gemm:
deep_gemm_experts = mk.FusedMoEModularKernel(
prepare_finalize=MoEPrepareAndFinalizeNoEP(),
moe_config = (
FusedMoEConfig(
num_experts=num_experts,
experts_per_token=topk,
hidden_dim=hidden_size,
intermediate_size_per_partition=shard_intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
activation=MoEActivation.SILU,
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
in_dtype=init_dtype,
routing_method=RoutingMethodType.TopK,
device="cuda",
),
)
deep_gemm_experts = mk.FusedMoEKernel(
prepare_finalize=maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
fused_experts=TritonOrDeepGemmExperts(
moe_config=FusedMoEConfig(
num_experts=num_experts,
experts_per_token=topk,
hidden_dim=hidden_size,
intermediate_size_per_partition=shard_intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
activation=MoEActivation.SILU,
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
in_dtype=init_dtype,
routing_method=RoutingMethodType.TopK,
device="cuda",
),
moe_config=moe_config,
quant_config=quant_config,
),
inplace=not disable_inplace(),
)
with override_config(config):
@@ -269,8 +281,16 @@ def benchmark_config(
inplace = not disable_inplace()
if use_deep_gemm:
return deep_gemm_experts(
x, w1, w2, topk_weights, topk_ids, inplace=inplace
return deep_gemm_experts.apply(
x,
w1,
w2,
topk_weights,
topk_ids,
activation=MoEActivation.SILU,
global_num_experts=num_experts,
apply_router_weight_on_input=False,
expert_map=False,
)
return fused_experts(
x,
+8 -12
View File
@@ -46,24 +46,20 @@ else()
)
endif()
# Install rules for FA components need the install prefix nested under vllm/
# These run at install time, before the FA library's own install rules
foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" COMPONENT ${_FA_COMPONENT})
install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" COMPONENT ${_FA_COMPONENT})
install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" COMPONENT ${_FA_COMPONENT})
endforeach()
# Make sure vllm-flash-attn install rules are nested under vllm/
# ALL_COMPONENTS ensures the save/modify/restore runs exactly once regardless
# of how many components are being installed, avoiding double-append of /vllm/.
install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" ALL_COMPONENTS)
install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_COMPONENTS)
# Fetch the vllm-flash-attn library
FetchContent_MakeAvailable(vllm-flash-attn)
message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}")
# Restore the install prefix after FA's install rules
foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" COMPONENT ${_FA_COMPONENT})
install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" COMPONENT ${_FA_COMPONENT})
endforeach()
install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
# Install shared Python files for both FA2 and FA3 components
foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
+2
View File
@@ -262,7 +262,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
# Build the vLLM wheel
# if USE_SCCACHE is set, use sccache to speed up compilation
# AWS credentials mounted at ~/.aws/credentials for sccache S3 auth (optional)
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
if [ "$USE_SCCACHE" = "1" ]; then \
echo "Installing sccache..." \
&& case "${TARGETPLATFORM}" in \
+51 -3
View File
@@ -115,9 +115,57 @@ RUN --mount=type=cache,target=/root/.cache/uv \
# install development dependencies (for testing)
RUN uv pip install -e tests/vllm_test_utils
# install nixl from source code
ENV NIXL_VERSION=0.7.0
RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
# install NIXL and UCX from source code
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
ARG NIXL_VERSION=0.7.0
RUN apt-get update && apt-get install -y \
pciutils \
net-tools \
iproute2 \
hwloc \
numactl \
wget \
curl \
git \
build-essential \
autoconf \
automake \
libtool \
pkg-config \
rdma-core \
libibverbs-dev \
ibverbs-utils \
libibverbs1 \
librdmacm-dev \
librdmacm1 \
libibumad-dev \
libibumad3 \
libibmad-dev \
libibmad5 \
infiniband-diags \
perftest \
ibutils \
libmlx5-1 \
libmlx4-1 \
ibverbs-providers \
librdmacm1t64
ENV PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig:${PKG_CONFIG_PATH}
ENV LD_LIBRARY_PATH=/tmp/ucx_install/lib:${LD_LIBRARY_PATH}
RUN --mount=type=cache,target=/root/.cache/uv \
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
bash autogen.sh && \
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
cd /tmp/nixl_source && \
uv pip install --upgrade meson pybind11 patchelf && \
uv pip install -r requirements.txt && \
uv pip install . && \
rm -rf /tmp/ucx_source /tmp/nixl_source
# FIX triton
RUN --mount=type=cache,target=/root/.cache/uv \
Binary file not shown.

After

Width:  |  Height:  |  Size: 130 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 128 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 254 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 73 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 65 KiB

+2
View File
@@ -54,6 +54,8 @@ For example:
--8<-- "vllm/model_executor/layers/attention/mm_encoder_attention.py:mm_encoder_attn"
--8<-- "vllm/model_executor/layers/mla.py:multi_head_latent_attention"
--8<-- "vllm/model_executor/models/deepencoder.py:rel_pos_attention"
```
**2. Activation:**
+1 -1
View File
@@ -81,7 +81,7 @@ The current implementation has all `dbo_yield` and `dbo_maybe_run_recv_hook` cal
The `make_ubatch_context` function initializes two `UBatchContexts`, one for each UBatch thread. It takes two CUDA streams, the preexisting `ForwardContexts` and a CPU thread barrier. This function should be used exclusively to instantiate `UBatchContexts`. It will handle all of the event initialization.
The `dbo_register_recv_hook` method registers a callback that can be returned by the `FusedMoEPrepareAndFinalize` class in the other UBatch threads `UBatchContext`. The callback will be run when the other thread calls `dbo_maybe_run_recv_hook`. This is typically used to wait on an all-to-all kernel.
The `dbo_register_recv_hook` method registers a callback that can be returned by the `FusedMoEPrepareAndFinalizeModular` class in the other UBatch threads `UBatchContext`. The callback will be run when the other thread calls `dbo_maybe_run_recv_hook`. This is typically used to wait on an all-to-all kernel.
The `dbo_maybe_run_recv_hook` method runs a callback thats set by the `dbo_register_recv_hook` function if that callback exists.
+52 -52
View File
@@ -37,31 +37,31 @@ The rest of the document will focus on the Contiguous / Non-Batched case. Extrap
FusedMoEModularKernel splits the FusedMoE operation into 3 parts,
1. TopKWeightAndReduce
2. FusedMoEPrepareAndFinalize
3. FusedMoEPermuteExpertsUnpermute
2. FusedMoEPrepareAndFinalizeModular
3. FusedMoEExpertsModular
### TopKWeightAndReduce
The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEPermuteExpertsUnpermute` is responsible for the Unpermute and `FusedMoEPrepareAndFinalize` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEPermuteExpertsUnpermute`. But some implementations choose to do it `FusedMoEPrepareAndFinalize`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class.
The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEExpertsModular` is responsible for the Unpermute and `FusedMoEPrepareAndFinalizeModular` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEExpertsModular`. But some implementations choose to do it `FusedMoEPrepareAndFinalizeModular`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class.
Please find the implementations of TopKWeightAndReduce [here](../../vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py).
`FusedMoEPrepareAndFinalize::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method.
The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEPermuteExpertsUnpermute` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens.
`FusedMoEPrepareAndFinalizeModular::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method.
The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEExpertsModular` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens.
* `FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEPermuteExpertsUnpermute` implementation does the weight application and reduction itself.
* `FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEPermuteExpertsUnpermute` implementation needs the `FusedMoEPrepareAndFinalize::finalize()` to do the weight application and reduction.
* `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEExpertsModular` implementation does the weight application and reduction itself.
* `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEExpertsModular` implementation needs the `FusedMoEPrepareAndFinalizeModular::finalize()` to do the weight application and reduction.
### FusedMoEPrepareAndFinalize
### FusedMoEPrepareAndFinalizeModular
The `FusedMoEPrepareAndFinalize` abstract class exposes `prepare`, `prepare_no_receive` and `finalize` functions.
The `prepare` function is responsible for input activation Quantization and All2All Dispatch. If implemented, The `prepare_no_receive` is like `prepare` except it does not wait to receive results from other workers. Instead it returns a "receiver" callback that must be invoked to wait for the final results of worker. It is not required that this method is supported by all `FusedMoEPrepareAndFinalize` classes, but if it is available, it can be used to interleave work with the initial all to all communication, e.g. interleaving shared experts with fused experts. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section)
The `FusedMoEPrepareAndFinalizeModular` abstract class exposes `prepare`, `prepare_no_receive` and `finalize` functions.
The `prepare` function is responsible for input activation Quantization and All2All Dispatch. If implemented, The `prepare_no_receive` is like `prepare` except it does not wait to receive results from other workers. Instead it returns a "receiver" callback that must be invoked to wait for the final results of worker. It is not required that this method is supported by all `FusedMoEPrepareAndFinalizeModular` classes, but if it is available, it can be used to interleave work with the initial all to all communication, e.g. interleaving shared experts with fused experts. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section)
![FusedMoEPrepareAndFinalize Blocks](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png)
![FusedMoEPrepareAndFinalizeModular Blocks](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png)
### FusedMoEPermuteExpertsUnpermute
### FusedMoEExpertsModular
The `FusedMoEPermuteExpertsUnpermute` class is where the crux of the MoE operations happen. The `FusedMoEPermuteExpertsUnpermute` abstract class exposes a few important functions,
The `FusedMoEExpertsModular` class is where the crux of the MoE operations happen. The `FusedMoEExpertsModular` abstract class exposes a few important functions,
* apply()
* workspace_shapes()
@@ -81,25 +81,25 @@ The `apply` method is where the implementations perform
#### workspace_shapes()
The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEPermuteExpertsUnpermute::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation.
The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation.
#### finalize_weight_and_reduce_impl()
It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEPermuteExpertsUnpermute::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section.
`FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalize::finalize()` to use.
It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEExpertsModular::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section.
`FusedMoEExpertsModular::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalizeModular::finalize()` to use.
![FusedMoEPermuteExpertsUnpermute Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png)
![FusedMoEExpertsModular Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png)
### FusedMoEModularKernel
`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` objects.
`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` objects.
`FusedMoEModularKernel` pseudocode/sketch,
```py
class FusedMoEModularKernel:
def __init__(self,
prepare_finalize: FusedMoEPrepareAndFinalize,
fused_experts: FusedMoEPermuteExpertsUnpermute):
prepare_finalize: FusedMoEPrepareAndFinalizeModular,
fused_experts: FusedMoEExpertsModular):
self.prepare_finalize = prepare_finalize
self.fused_experts = fused_experts
@@ -128,53 +128,53 @@ class FusedMoEModularKernel:
## How-To
### How To Add a FusedMoEPrepareAndFinalize Type
### How To Add a FusedMoEPrepareAndFinalizeModular Type
Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & Combine implementation / kernel. For example,
Typically a FusedMoEPrepareAndFinalizeModular type is backed by an All2All Dispatch & Combine implementation / kernel. For example,
* DeepEPHTPrepareAndFinalize type is backed by DeepEP High-Throughput All2All kernels, and
* DeepEPLLPrepareAndFinalize type is backed by DeepEP Low-Latency All2All kernels.
#### Step 1: Add an All2All manager
The purpose of the All2All Manager is to set up the All2All kernel implementations. The `FusedMoEPrepareAndFinalize` implementations typically fetch a kernel-implementation "handle" from the All2All Manager to invoke the Dispatch and Combine functions. Please look at the All2All Manager implementations [here](../../vllm/distributed/device_communicators/all2all.py).
The purpose of the All2All Manager is to set up the All2All kernel implementations. The `FusedMoEPrepareAndFinalizeModular` implementations typically fetch a kernel-implementation "handle" from the All2All Manager to invoke the Dispatch and Combine functions. Please look at the All2All Manager implementations [here](../../vllm/distributed/device_communicators/all2all.py).
#### Step 2: Add a FusedMoEPrepareAndFinalize Type
#### Step 2: Add a FusedMoEPrepareAndFinalizeModular Type
This section describes the significance of the various functions exposed by the `FusedMoEPrepareAndFinalize` abstract class.
This section describes the significance of the various functions exposed by the `FusedMoEPrepareAndFinalizeModular` abstract class.
`FusedMoEPrepareAndFinalize::prepare()`: The prepare method implements the Quantization and All2All Dispatch. Typically the Dispatch function from the relevant All2All Manager is invoked.
`FusedMoEPrepareAndFinalizeModular::prepare()`: The prepare method implements the Quantization and All2All Dispatch. Typically the Dispatch function from the relevant All2All Manager is invoked.
`FusedMoEPrepareAndFinalize::has_prepare_no_receive()`: Indicates whether or not this subclass implements `prepare_no_receive`. Defaults to False.
`FusedMoEPrepareAndFinalizeModular::has_prepare_no_receive()`: Indicates whether or not this subclass implements `prepare_no_receive`. Defaults to False.
`FusedMoEPrepareAndFinalize::prepare_no_receive()`: The prepare_no_receive method implements the Quantization and All2All Dispatch. It does not wait for the result of the dispatch operation but instead returns a thunk that can be invoked to wait for the final results. Typically the Dispatch function from the relevant All2All Manager is invoked.
`FusedMoEPrepareAndFinalizeModular::prepare_no_receive()`: The prepare_no_receive method implements the Quantization and All2All Dispatch. It does not wait for the result of the dispatch operation but instead returns a thunk that can be invoked to wait for the final results. Typically the Dispatch function from the relevant All2All Manager is invoked.
`FusedMoEPrepareAndFinalize::finalize()`: Maybe perform TopK Weight Application and Reduction and All2All Combine. Typically the Combine function from the relevant All2AllManager is invoked.
`FusedMoEPrepareAndFinalizeModular::finalize()`: Maybe perform TopK Weight Application and Reduction and All2All Combine. Typically the Combine function from the relevant All2AllManager is invoked.
`FusedMoEPrepareAndFinalize::activation_format()`: Return `FusedMoEActivationFormat.BatchedExperts` if the output of the prepare method (i.e. the All2All dispatch) is Batched. Return `FusedMoEActivationFormat.Standard` otherwise.
`FusedMoEPrepareAndFinalizeModular::activation_format()`: Return `FusedMoEActivationFormat.BatchedExperts` if the output of the prepare method (i.e. the All2All dispatch) is Batched. Return `FusedMoEActivationFormat.Standard` otherwise.
`FusedMoEPrepareAndFinalize::topk_indices_dtype()`: Data type of the TopK ids. Some All2All kernels have strict requirements pertaining to the data type of the TopK ids. This requirement is passed on to the `FusedMoe::select_experts` function so it could be respected. If there are no strict requirements return None.
`FusedMoEPrepareAndFinalizeModular::topk_indices_dtype()`: Data type of the TopK ids. Some All2All kernels have strict requirements pertaining to the data type of the TopK ids. This requirement is passed on to the `FusedMoe::select_experts` function so it could be respected. If there are no strict requirements return None.
`FusedMoEPrepareAndFinalize::max_num_tokens_per_rank()`: This is the maximum number of tokens that would be submitted to the All2All Dispatch at once.
`FusedMoEPrepareAndFinalizeModular::max_num_tokens_per_rank()`: This is the maximum number of tokens that would be submitted to the All2All Dispatch at once.
`FusedMoEPrepareAndFinalize::num_dispatchers()`: Total number of dispatching units. This value determines the size of the Dispatch output. The Dispatch output is of shape (num_local_experts, max_num_tokens, K). Here max_num_tokens = num_dispatchers() * max_num_tokens_per_rank().
`FusedMoEPrepareAndFinalizeModular::num_dispatchers()`: Total number of dispatching units. This value determines the size of the Dispatch output. The Dispatch output is of shape (num_local_experts, max_num_tokens, K). Here max_num_tokens = num_dispatchers() * max_num_tokens_per_rank().
We suggest picking an already existing `FusedMoEPrepareAndFinalize` implementation that matches your All2All implementation closely and using it as a reference.
We suggest picking an already existing `FusedMoEPrepareAndFinalizeModular` implementation that matches your All2All implementation closely and using it as a reference.
### How To Add a FusedMoEPermuteExpertsUnpermute Type
### How To Add a FusedMoEExpertsModular Type
FusedMoEPermuteExpertsUnpermute performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows,
FusedMoEExpertsModular performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows,
`FusedMoEPermuteExpertsUnpermute::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format.
`FusedMoEExpertsModular::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format.
`FusedMoEPermuteExpertsUnpermute::supports_chunking()`: Return True if the implementation supports chunking. Typically
`FusedMoEExpertsModular::supports_chunking()`: Return True if the implementation supports chunking. Typically
implementations that input `FusedMoEActivationFormat.Standard` support chunking and `FusedMoEActivationFormat.BatchedExperts` do not.
`FusedMoEPermuteExpertsUnpermute::supports_expert_map()`: Return True if the implementation supports expert map.
`FusedMoEExpertsModular::supports_expert_map()`: Return True if the implementation supports expert map.
`FusedMoEPermuteExpertsUnpermute::workspace_shapes()` /
`FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` /
`FusedMoEPermuteExpertsUnpermute::apply`: Refer to `FusedMoEPermuteExpertsUnpermute` section above.
`FusedMoEExpertsModular::workspace_shapes()` /
`FusedMoEExpertsModular::finalize_weight_and_reduce_impl` /
`FusedMoEExpertsModular::apply`: Refer to `FusedMoEExpertsModular` section above.
### FusedMoEModularKernel Initialization
@@ -186,14 +186,14 @@ implementations that input `FusedMoEActivationFormat.Standard` support chunking
#### maybe_make_prepare_finalize
The `maybe_make_prepare_finalize` method is responsible for constructing an instance of `FusedMoEPrepareAndFinalize` when appropriate based on the current all2all backend, e.g. when EP + DP is enabled. The base class method currently constructs all the `FusedMoEPrepareAndFinalize` objects for the EP+DP case. Derived classes can override this method to construct prepare/finalize objects for different scenarios, e.g. `ModelOptNvFp4FusedMoE` can construct a `FlashInferCutlassMoEPrepareAndFinalize` for the EP+TP case.
The `maybe_make_prepare_finalize` method is responsible for constructing an instance of `FusedMoEPrepareAndFinalizeModular` when appropriate based on the current all2all backend, e.g. when EP + DP is enabled. The base class method currently constructs all the `FusedMoEPrepareAndFinalizeModular` objects for the EP+DP case. Derived classes can override this method to construct prepare/finalize objects for different scenarios, e.g. `ModelOptNvFp4FusedMoE` can construct a `FlashInferCutlassMoEPrepareAndFinalize` for the EP+TP case.
Please refer to the implementations in,
* `ModelOptNvFp4FusedMoE`
#### select_gemm_impl
The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEPermuteExpertsUnpermute` object.
The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEExpertsModular` object.
Please refer to the implementations in,
* `UnquantizedFusedMoEMethod`
@@ -205,7 +205,7 @@ derived classes.
#### init_prepare_finalize
Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalize` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEPermuteExpertsUnpermute` object and builds the `FusedMoEModularKernel` object
Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalizeModular` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEModularKernel` object
Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188).
**Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used.
@@ -214,9 +214,9 @@ Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vl
We have `FusedMoEModularKernel` unit tests at [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py).
The unit test iterates through all combinations of `FusedMoEPrepareAndFinalize` and `FusedMoEPremuteExpertsUnpermute` types and if they are
The unit test iterates through all combinations of `FusedMoEPrepareAndFinalizeModular` and `FusedMoEPremuteExpertsUnpermute` types and if they are
compatible, runs some correctness tests.
If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEPermuteExpertsUnpermute` implementations,
If you are adding some `FusedMoEPrepareAndFinalizeModular` / `FusedMoEExpertsModular` implementations,
1. Add the implementation type to `MK_ALL_PREPARE_FINALIZE_TYPES` and `MK_FUSED_EXPERT_TYPES` in [mk_objects.py](../../tests/kernels/moe/modular_kernel_tools/mk_objects.py) respectively.
2. Update `Config::is_batched_prepare_finalize()`, `Config::is_batched_fused_experts()`, `Config::is_standard_fused_experts()`,
@@ -225,24 +225,24 @@ If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEPermuteExpertsUnp
Doing this will add the new implementation to the test suite.
### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` Compatibility
### How To Check `FusedMoEPrepareAndFinalizeModular` & `FusedMoEExpertsModular` Compatibility
The unit test file [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py) can also be executed as a standalone script.
Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts`
As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` compatibility. When invoked
As a side effect, this script can be used to test `FusedMoEPrepareAndFinalizeModular` & `FusedMoEExpertsModular` compatibility. When invoked
with incompatible types, the script will error.
### How To Profile
Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py)
The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible
`FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` types.
`FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` types.
Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts`
## FusedMoEPrepareAndFinalize Implementations
## FusedMoEPrepareAndFinalizeModular Implementations
See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-modular-all2all-backends) for a list of all the available modular prepare and finalize subclasses.
## FusedMoEPermuteExpertsUnpermute
## FusedMoEExpertsModular
See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-experts-kernels) for a list of all the available modular experts.
+198
View File
@@ -0,0 +1,198 @@
# Model Runner V2 Design Document
## Introduction
Since vLLM V1 was first implemented, we discovered several fundamental design mistakes and accumulated significant technical debt. Many features were bolted on that were not considered in the original design. We also gained valuable insights into sampling techniques (for example, Gumbel-max sampling), tools (for example, Triton), and CUDA features (for example, UVA). With this knowledge, we implemented Model Runner V2 (MRV2) from first principles to be cleaner, more efficient, and more modular.
In hindsight, many of V1's design choices were suboptimal. While MRV2 is not yet feature-complete, not rigorously tested, and still has open design decisions, we believe it is a substantial improvement over V1.
This document describes the design of MRV2.
## 1. Persistent Batch
One significant source of friction in V1 is its persistent batch implementation.
### Background
V1 introduced persistent batches to minimize CPU overhead during input preparation. When requests are scheduled for a step, the model runner must construct contiguous input tensors (for example, block tables and per-request temperature values) to feed into the model. Building these tensors from scratch each step is often very slow in Python, especially for large tensors like block tables.
The persistent batch optimization exploits the fact that request batches in consecutive steps are mostly identical. Only a few requests (if any) join or finish per step. By maintaining persistent state tensors and applying incremental diffs instead of reconstructing inputs from scratch, CPU overhead can be reduced significantly.
### Problems with V1's Approach
While efficient, V1's persistent batch design introduced unnecessary complexity due to coupling persistent state with input tensors. V1 uses persistent state tensors directly as model and sampler inputs, which imposes strict layout and ordering requirements. When requests join or finish, this often requires complex tensor-wide reordering rather than simple row insertion/removal.
V1 also had to maintain `CachedRequestState`, a redundant backup copy of request state, because rows in persistent tensors can be overwritten while requests are still active.
The result is complex bookkeeping that becomes more difficult under async scheduling.
![Persistent Batch in V1](../assets/design/model_runner_v2/persistent_batch_v1.png)
### MRV2's Solution
MRV2 decouples persistent state tensors from per-step input tensors. Given request ordering for the step (usually determined by the attention backend), MRV2 gathers input tensors from persistent state.
1. Pre-allocate a fixed-size tensor with `max_num_reqs` rows (1024 by default on most platforms).
2. Assign each request a permanent row for its active lifetime (until finish or preemption).
3. Treat preemption as completion. On resume, re-add request data as fresh state.
This removes the need for `CachedRequestState` and simplifies bookkeeping. Large state tensors are mostly stored on GPU memory, so gather runs in parallel on the GPU with low overhead.
![Persistent Batch in MRV2](../assets/design/model_runner_v2/persistent_batch_mrv2.png)
## 2. Async-First
vLLM now relies heavily on asynchronous scheduling. The scheduler and worker prepare inputs for step `N+1` while the GPU executes step `N`, overlapping CPU and GPU work to maximize utilization.
V1 was not originally designed with async scheduling in mind, and support required retrofitted behavior and hacks. MRV2 instead assumes the core model execution loop is a CUDA stream with no CPU synchronization points. CPU entrypoints queue work onto the stream.
![Async execution timeline](../assets/design/model_runner_v2/async_sched.png)
## 3. Removing Async Barrier
A key requirement for async execution is that CPU operations remain non-blocking. Both explicit sync (for example, `torch.cuda.synchronize`) and implicit sync (for example, unpinned `.to("cuda")`) must be avoided.
However, async execution can introduce race conditions when CPU and GPU concurrently touch the same memory.
Example (unsafe):
```python
class ModelRunner:
def __init__(self, ...):
# Pinned buffer
self.states = torch.zeros(
max_num_reqs, dtype=torch.int32, device="cpu", pin_memory=True
)
def execute_step(self, ...):
self.states[req_idx] = new_req.data
states = self.states.to("cuda", non_blocking=True)
```
The CPU may modify `self.states` while GPU is still reading from it via async copy.
V1 addresses this with an async barrier around critical sections. That avoids races but has drawbacks:
1. Easy to miss protected buffers (bug-prone).
2. Inflexible organization (all CPU work must stay inside barrier).
3. Potentially less overlap due to synchronization.
![Race condition with shared CPU buffer](../assets/design/model_runner_v2/async_race_condition.png)
### MRV2's Solution: Eliminate the Race
MRV2 separates persistent CPU state from the copied tensor:
```python
class ModelRunner:
def __init__(self, ...):
# Not pinned
self.states = torch.zeros(
max_num_reqs, dtype=torch.int32, device="cpu", pin_memory=False
)
def execute_step(self, ...):
self.states[req_idx] = new_req.data
tmp_states = self.states.pin_memory()
states = tmp_states.to("cuda", non_blocking=True)
```
Now CPU writes to `self.states` while GPU reads from `tmp_states`, eliminating the race without explicit synchronization.
![No race with temporary pinned copy](../assets/design/model_runner_v2/async_no_race_condition.png)
## 4. StagedWriteTensor
For large tensors like block tables, MRV2 avoids full CPU-to-GPU copies each step by using `StagedWriteTensor`:
1. Keep the base tensor on GPU.
2. Stage diffs on CPU.
3. Pack diffs into contiguous buffers.
4. Copy packed diffs to GPU.
5. Launch one kernel to apply diffs.
Example usage:
```python
# Initialize state on GPU
state = StagedWriteTensor(size=(1024, 1000), dtype=torch.int32, device="cuda")
# Write [3, 1, 2] into row 2, starting at index 3
state.stage_write(row=2, start=3, value=[3, 1, 2])
# Write [-1, -2, -5] into row 0, starting at index 1
state.stage_write(row=0, start=1, value=[-1, -2, -5])
# Apply staged changes
state.apply_write()
```
This supports ragged updates with no CPU-GPU synchronization and minimal kernel launches. It is especially useful for block tables and mixed CPU/GPU-written states such as `num_computed_tokens`.
## 5. GPU-Native Input Metadata Preparation and Output Processing
MRV2 uses Triton kernels to prepare inputs such as `input_ids`, `positions`, `query_start_loc`, and `seq_lens`.
Benefits:
1. Better async behavior: GPU can derive values (for example with speculative decoding) that CPU may not know yet.
2. Lower CPU overhead: input prep is very cheap on GPU and avoids Python bottlenecks.
### Universal Virtual Addressing (UVA)
MRV2 uses UVA in some paths to let GPU kernels access large CPU-resident tensors directly (for example `prefill_token_ids`) without duplicating those tensors into GPU memory.
## 6. Triton-Native Sampler
MRV2 reimplements sampling mostly in Triton for better numeric/memory control and optimization.
### Gumbel Sampling Kernel
MRV2 introduces a Triton Gumbel sampling kernel that avoids explicit softmax materialization and uses stateless in-kernel RNG from seed input.
### Efficient Top-K Logprobs
V1 materializes full-vocabulary logprobs before top-k. MRV2 identifies top-k tokens from logits first, then computes logprobs only for selected tokens. This reduces peak GPU memory usage.
### Memory-Efficient Prompt Logprobs
MRV2 supports finer-grained chunking, including chunking inside a single prompt, to avoid memory spikes on long prompts.
### Better Compatibility with Speculative Decoding
Instead of expanding per-request sampling states to match per-logit shapes, MRV2 uses indirection (`idx_mapping`) inside kernels to map each logits vector to the right request state. This simplifies support for complex sampling parameters and logits processors.
## 7. Modularity
MRV2 emphasizes modularity. Compared to V1's large, entangled `gpu_model_runner.py`, MRV2 splits feature logic across dedicated files (for example, `mrope_utils.py`, `penalties.py`, and many others).
It also consolidates model inputs into an `InputBatch` class and reduces direct model-runner attribute coupling.
## 8. No Abuse of `dummy_run`
In V1, `dummy_run` handled too many responsibilities:
- Initial memory profiling and `torch.compile`
- CUDA graph capture
- Warmups
- Empty DP forward passes for EP+DP
MRV2 simplifies this:
1. `execute_model` supports dummy runs without affecting state.
2. `dummy_run` delegates to `execute_model` for profiling, warmup, and empty DP forward passes.
3. CUDA graph capture uses a separate dedicated path.
This reduces complexity and removes bugs caused by divergence between `execute_model` and `dummy_run` behavior.
## 9. Explicit CUDA Graph Management
V1's CUDA graph handling is implicit and hard to reason about. MRV2 uses a `CUDAGraphManager` that explicitly captures and launches full CUDA graphs through standard PyTorch APIs.
This makes graph lifecycle and execution mode decisions more understandable and easier to extend. Example: MRV2 can capture multiple draft-model forward passes into one CUDA graph.
## Development Philosophy
MRV2 changes should meet a higher code quality bar. As feature gaps with V1 are filled, features should be reconsidered from first principles in the MRV2 design context instead of quickly porting V1 behavior.
A key requirement is preserving modularity and clean abstraction boundaries, even if that requires more upfront design iteration.
+7 -9
View File
@@ -4,17 +4,17 @@ The purpose of this document is to provide an overview of the various MoE kernel
## Fused MoE Modular All2All backends
There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalize` subclasses provide an interface for each all2all backend.
There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalizeModular` subclasses provide an interface for each all2all backend.
The following table describes the relevant features of each backend, i.e. activation format, supported quantization schemes and async support.
The output activation format (standard or batched) corresponds to the output of the prepare step of the `FusedMoEPrepareAndFinalize` subclass, and the finalize step requires the same format. All the backend `prepare` methods expect activations in the standard format and all the `finalize` methods return activations in standard format. More details on the formats can be found in the [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) document.
The output activation format (standard or batched) corresponds to the output of the prepare step of the `FusedMoEPrepareAndFinalizeModular` subclass, and the finalize step requires the same format. All the backend `prepare` methods expect activations in the standard format and all the `finalize` methods return activations in standard format. More details on the formats can be found in the [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) document.
The quantization types and formats enumerate which quantization schemes are supported by each `FusedMoEPrepareAndFinalize` class. The quantization can happen before or after the dispatch based on the format the all2all backend supports, e.g. deepep_high_throughput supports only block-quantized fp8 format. Any other format will result in dispatching in higher precision and quantizing afterwards. The output of the prepare step for each backend is the quantized type. The finalize step generally requires the same input type as the original activations, e.g. if the original input is bfloat16 and the quantization scheme is fp8 with per-tensor scales, `prepare` will return fp8/per-tensor scale activations and `finalize` will take bfloat16 activations. See the diagrams in [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) for more details on the types and formats of activations at each step of the MoE process. If no quantization type is specified, the kernel operates on float16 and/or bfloat16.
The quantization types and formats enumerate which quantization schemes are supported by each `FusedMoEPrepareAndFinalizeModular` class. The quantization can happen before or after the dispatch based on the format the all2all backend supports, e.g. deepep_high_throughput supports only block-quantized fp8 format. Any other format will result in dispatching in higher precision and quantizing afterwards. The output of the prepare step for each backend is the quantized type. The finalize step generally requires the same input type as the original activations, e.g. if the original input is bfloat16 and the quantization scheme is fp8 with per-tensor scales, `prepare` will return fp8/per-tensor scale activations and `finalize` will take bfloat16 activations. See the diagrams in [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) for more details on the types and formats of activations at each step of the MoE process. If no quantization type is specified, the kernel operates on float16 and/or bfloat16.
Async backends support the use of DBO (Dual Batch Overlap) and shared expert overlap (where shared experts are computed during the combine step).
Certain models require the topk weights to be applied to the input activations rather than the output activations when topk==1, e.g. Llama. For modular kernels, this feature is supported by the `FusedMoEPrepareAndFinalize` subclass. For non-modular kernels, it is up to the experts function to deal with this flag.
Certain models require the topk weights to be applied to the input activations rather than the output activations when topk==1, e.g. Llama. For modular kernels, this feature is supported by the `FusedMoEPrepareAndFinalizeModular` subclass. For non-modular kernels, it is up to the experts function to deal with this flag.
Unless otherwise specified, backends are controlled via the `--all2all-backend` command-line argument (or the `all2all_backend` parameter in `ParallelConfig`). All backends except `flashinfer` only work with EP+DP or EP+TP. `Flashinfer` can work with EP or DP without EP.
@@ -36,8 +36,6 @@ th {
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |
| MoEPrepareAndFinalizeNoEP<sup>5</sup> | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoEP`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoEP] |
| BatchedPrepareAndFinalize<sup>5</sup> | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] |
!!! info "Table key"
1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -75,9 +73,9 @@ Each experts kernel supports one or more activation functions, e.g. silu or gelu
As with the backends, some experts support applying topk weights on the input activations. The entries in the column in this table only apply to the non-modular experts.
Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEPermuteExpertsUnpermute`.
Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEExpertsModular`.
To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats.
To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats.
| Kernel | Input act. format | Quant. types | Quant. format | Activation function | Apply Weight On Input | Modular | Source |
|--------|-------------------|--------------|---------------|---------------------|-----------------------|---------|--------|
@@ -106,7 +104,7 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels
The following table shows "families" of modular kernels that are intended to work together. There are some combinations which may work but have not yet been tested, e.g. flashinfer with other fp8 experts. Note that the "naive" backend will work with any non-modular experts.
| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses |
| backend | `FusedMoEPrepareAndFinalizeModular` subclasses | `FusedMoEExpertsModular` subclasses |
|---------|-----------------------------------------|----------------------------------------------|
| deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8`, </br>`MarlinExperts` |
| deepep_low_latency | `DeepEPLLPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`CutlassBatchedExpertsFp8`,</br>`BatchedMarlinExperts` |
+67 -1
View File
@@ -498,7 +498,9 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
- Multi-vector retrieval: [examples/pooling/token_embed/colqwen3_token_embed_online.py](../../examples/pooling/token_embed/colqwen3_token_embed_online.py)
- Reranking (text + multi-modal): [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
### Llama Nemotron Multimodal Embedding Models
### Llama Nemotron Multimodal
#### Embedding Model
Llama Nemotron VL Embedding models combine the bidirectional Llama embedding backbone
(from `nvidia/llama-nemotron-embed-1b-v2`) with SigLIP as the vision encoder to produce
@@ -559,6 +561,70 @@ curl -s http://localhost:8000/v1/embeddings -H "Content-Type: application/json"
}'
```
#### Reranker Model
Llama Nemotron VL reranker models combine the same bidirectional Llama + SigLIP
backbone with a sequence-classification head for cross-encoder scoring and reranking.
| Architecture | Backbone | Example HF Models |
|---|---|---|
| `LlamaNemotronVLForSequenceClassification` | Bidirectional Llama + SigLIP | `nvidia/llama-nemotron-rerank-vl-1b-v2` |
Start the server:
```shell
vllm serve nvidia/llama-nemotron-rerank-vl-1b-v2 \
--runner pooling \
--trust-remote-code \
--chat-template examples/pooling/score/template/nemotron-vl-rerank.jinja
```
!!! note
The chat template bundled with this checkpoint's tokenizer is not suitable
for the Score/Rerank APIs. Use the provided override template when serving:
`examples/pooling/score/template/nemotron-vl-rerank.jinja`.
Score a text query against an image document:
```shell
curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
"model": "nvidia/llama-nemotron-rerank-vl-1b-v2",
"data_1": "Find diagrams about autonomous robots",
"data_2": [
{
"content": [
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
{"type": "text", "text": "Robotics workflow diagram."}
]
}
]
}'
```
Rerank image documents by a text query:
```shell
curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
"model": "nvidia/llama-nemotron-rerank-vl-1b-v2",
"query": "Find diagrams about autonomous robots",
"documents": [
{
"content": [
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64_1>"}},
{"type": "text", "text": "Robotics workflow diagram."}
]
},
{
"content": [
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64_2>"}},
{"type": "text", "text": "General skyline photo."}
]
}
],
"top_n": 2
}'
```
### BAAI/bge-m3
The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
+2
View File
@@ -793,6 +793,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|--------------|--------|-------------------|----------------------|---------------------------|
| `FireRedASR2ForConditionalGeneration` | FireRedASR2 | `allendou/FireRedASR2-LLM-vllm`, etc. | | |
| `FunASRForConditionalGeneration` | FunASR | `allendou/Fun-ASR-Nano-2512-vllm`, etc. | | |
| `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
| `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
@@ -842,6 +843,7 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
| Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|--------------|--------|--------|-------------------|----------------------|---------------------------|
| `JinaVLForSequenceClassification` | JinaVL-based | T + I<sup>E+</sup> | `jinaai/jina-reranker-m0`, etc. | ✅︎ | ✅︎ |
| `LlamaNemotronVLForSequenceClassification` | Llama Nemotron Reranker + SigLIP | T + I<sup>E+</sup> | `nvidia/llama-nemotron-rerank-vl-1b-v2` | | |
| `Qwen3VLForSequenceClassification` | Qwen3-VL-Reranker | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-Reranker-2B`(see note), etc. | ✅︎ | ✅︎ |
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./pooling_models.md#model-conversion))
@@ -47,12 +47,14 @@ from vllm.distributed.weight_transfer.nccl_engine import (
NCCLWeightTransferInitInfo,
NCCLWeightTransferUpdateInfo,
)
from vllm.platforms import current_platform
from vllm.utils.network_utils import get_ip, get_open_port
from vllm.v1.executor import Executor
MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base"
MODEL_NAME_V2 = "Qwen/Qwen3-1.7B"
PAUSE_TOKEN_THRESHOLD = 10
ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN"
class MyLLM(vllm.AsyncLLMEngine):
@@ -116,10 +118,16 @@ class TrainModel:
from vllm.model_executor.layers.batch_invariant import (
init_batch_invariance,
)
from vllm.platforms import current_platform
from vllm.v1.attention.backends.registry import AttentionBackendEnum
# need to init all env vars for batch invariance which affect nccl ops
init_batch_invariance(AttentionBackendEnum.FLASH_ATTN)
attn_backend = (
AttentionBackendEnum.TRITON_ATTN
if current_platform.is_rocm()
else AttentionBackendEnum.FLASH_ATTN
)
init_batch_invariance(attn_backend)
self.model = AutoModelForCausalLM.from_pretrained(
model_name, dtype=torch.bfloat16
@@ -175,23 +183,48 @@ class TrainModel:
return new_token_ids
ray.init(
runtime_env={
"env_vars": {
# enable batch invariance for deterministic outputs
"VLLM_BATCH_INVARIANT": "1",
# prevent ray from setting CUDA_VISIBLE_DEVICES
"RAY_EXPERIMENTAL_NOSET_CUDA_ENV_VAR": "1",
}
}
)
# Build platform-specific env vars for Ray
ray_env_vars = {
# Prevent Ray from setting CUDA_VISIBLE_DEVICES
"RAY_EXPERIMENTAL_NOSET_CUDA_ENV_VAR": "1",
}
if current_platform.is_rocm():
# For ROCm, BATCH_INVARIANT vllm is not supported
ray_env_vars["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
else:
# Enable batch invariance for deterministic outputs on NVIDIA
ray_env_vars["VLLM_BATCH_INVARIANT"] = "1"
ray.init(runtime_env={"env_vars": ray_env_vars})
# Launch the training model actor. Ray's resource scheduler will allocate
# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
train_model = TrainModel.remote(MODEL_NAME_V2)
# Launch the vLLM inference engine. The `enforce_eager` flag reduces
# start-up latency.
rocm_determinism_kwargs = {}
if current_platform.is_rocm():
# ROCm: To minimize non-determinism, we set fixed seed, no prefix caching, and
# sequential request processing (max_num_seqs=1).
rocm_determinism_kwargs = {
"seed": 0,
"enable_prefix_caching": False,
"max_num_seqs": 1,
}
# Build platform-specific LLM kwargs
llm_kwargs = dict(
model=MODEL_NAME_V1,
enforce_eager=True,
max_model_len=8192,
distributed_executor_backend="ray",
attention_backend=ATTN_BACKEND,
gpu_memory_utilization=0.75,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)
llm_kwargs.update(rocm_determinism_kwargs)
# Launch the vLLM inference engine.
# With data_parallel_backend="ray", vLLM's CoreEngineActorManager creates
# its own placement groups internally for each DP rank, so we must NOT
# create an outer placement group (it would reserve GPUs and hide them
@@ -199,15 +232,7 @@ train_model = TrainModel.remote(MODEL_NAME_V2)
llm = ray.remote(
num_cpus=0,
num_gpus=0,
)(MyLLM).remote(
model=MODEL_NAME_V1,
enforce_eager=True,
max_model_len=8192,
distributed_executor_backend="ray",
attention_backend="FLASH_ATTN",
gpu_memory_utilization=0.75,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)
)(MyLLM).remote(**llm_kwargs)
PROMPTS = [
"The president of the United States is",
@@ -304,25 +329,42 @@ for i, (output, pause_idx) in enumerate(results):
print(f" New weights ({n_after} tokens): {after_text!r}")
# ── Phase 2: validate with a fresh V2 vLLM instance ────────────────
# This validation relies on batch-invariant (deterministic) generation to
# compare outputs from the weight-synced engine against a fresh V2 instance.
# On NVIDIA, batch invariance is fully supported, so we require 100% exact
# token match. On ROCm, batch invariance is not yet fully implemented
# (see https://github.com/vllm-project/vllm/issues/27433 and
# https://github.com/vllm-project/vllm/issues/33123), so residual
# non-determinism (e.g. GEMM accumulation order, missing kernel overrides)
# can cause single-token divergences that don't indicate a weight-sync
# failure. We relax the pass rate to 90% on ROCm to accommodate this; a
# real regression (broken weight transfer) would cause ~0% pass rate, not 90%+.
MIN_PASS_RATE = 1.0 if not current_platform.is_rocm() else 0.9
print(f"\n{'=' * 50}")
print("VALIDATION: comparing weight-synced vLLM with fresh V2 instance")
if current_platform.is_rocm():
print(f" (ROCm mode: requiring >= {MIN_PASS_RATE:.0%} exact match rate)")
print(f"{'=' * 50}")
ray.get(llm.shutdown.remote())
ray.kill(llm)
ray.kill(train_model)
llm_v2 = ray.remote(
num_cpus=0,
num_gpus=0,
)(MyLLM).remote(
llm_v2_kwargs = dict(
model=MODEL_NAME_V2,
enforce_eager=True,
max_model_len=8192,
gpu_memory_utilization=0.75,
distributed_executor_backend="ray",
attention_backend="FLASH_ATTN",
attention_backend=ATTN_BACKEND,
)
llm_v2_kwargs.update(rocm_determinism_kwargs)
llm_v2 = ray.remote(
num_cpus=0,
num_gpus=0,
)(MyLLM).remote(**llm_v2_kwargs)
val_futures = [
llm_v2.do_generate.remote(
@@ -335,16 +377,17 @@ val_futures = [
]
val_results = ray.get(val_futures)
all_pass = True
num_pass = 0
num_total = len(results)
for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)):
expected = list(output.outputs[0].token_ids)[pause_idx:]
actual = list(val_output.outputs[0].token_ids)
match = actual == expected
if match:
num_pass += 1
print(f" [PASS] {PROMPTS[i]!r}")
else:
all_pass = False
print(f" [FAIL] {PROMPTS[i]!r}")
print(f" weight-synced vLLM: {tokenizer.decode(expected)!r}")
print(f" V2 vLLM: {tokenizer.decode(actual)!r}")
@@ -359,5 +402,14 @@ for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_resu
ray.get(llm_v2.shutdown.remote())
ray.kill(llm_v2)
assert all_pass, "Some prompts failed validation, see above for details"
pass_rate = num_pass / num_total
print(f"\n Result: {num_pass}/{num_total} prompts passed ({pass_rate:.0%})")
print(f" Required: >= {MIN_PASS_RATE:.0%}")
assert pass_rate >= MIN_PASS_RATE, (
f"Validation pass rate {pass_rate:.0%} ({num_pass}/{num_total}) "
f"is below the required {MIN_PASS_RATE:.0%} threshold. "
f"See failures above for details."
)
print("=" * 50)
@@ -0,0 +1,15 @@
{%- set query_msg = (messages | selectattr('role', 'equalto', 'query') | list | first) -%}
{%- set doc_msg = (messages | selectattr('role', 'equalto', 'document') | list | first) -%}
{%- set q = query_msg['content'] -%}
{%- set d = doc_msg['content'] -%}
{# If the doc contains <image> anywhere, hoist a single <image> to the front #}
{%- set has_image = ("<image>" in d) -%}
{%- set d_clean = d | replace("<image>", "") -%}
{%- set q_clean = q | replace("<image>", "") -%}
{%- if has_image -%}<image>{{ " " }}{%- endif -%}
question:{{ q_clean }}{{ " " }}
{{ " " }}
{{ " " }}passage:{{ d_clean }}
+12505
View File
File diff suppressed because it is too large Load Diff
+2 -1
View File
@@ -12,7 +12,7 @@ tokenizers >= 0.21.1 # Required for fast incremental detokenization.
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
aiohttp >= 3.13.3
openai >= 1.99.1 # For Responses API with reasoning content
openai == 2.24.0 # For Responses API with reasoning content
pydantic >= 2.12.0
prometheus_client >= 0.18.0
pillow # Required for image processing
@@ -57,3 +57,4 @@ opentelemetry-sdk >= 1.27.0
opentelemetry-api >= 1.27.0
opentelemetry-exporter-otlp >= 1.27.0
opentelemetry-semantic-conventions-ai >= 0.4.1
kaldi-native-fbank >= 1.18.7
+4
View File
@@ -13,6 +13,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
class Matches(NamedTuple):
# simple pointwise
aiter_rms_quant_fusion: int = 0
rms_quant_fusion: int = 0
act_quant_fusion: int = 0
norm_rope_fusion: int = 0
@@ -82,6 +83,9 @@ INDUCTOR_GRAPH_PARTITION = [
]
FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
"aiter_rms_quant_fusion": re.compile(
r"RocmAiterRMSNormQuantFusionPass Replaced (\d+) patterns"
),
"rms_quant_fusion": re.compile(r"rms_quant_fusion.py:\d+] Replaced (\d+) patterns"),
"act_quant_fusion": re.compile(r"act_quant_fusion.py:\d+] Replaced (\d+) patterns"),
"norm_rope_fusion": re.compile(
+5
View File
@@ -63,9 +63,14 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
compilation_config: dict,
matches_check: list[str],
use_deepgemm: bool = False,
use_aiter: bool = False,
tp_size: int = 1,
):
monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
from vllm._aiter_ops import rocm_aiter_ops
rocm_aiter_ops.refresh_env_variables()
# Disable, compile cache to make sure custom passes run.
# Otherwise, we can't verify fusion happened through the logs.
+20 -2
View File
@@ -2,6 +2,8 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm._aiter_ops import is_aiter_found_and_supported
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -24,6 +26,24 @@ TRITON_ATTN = pytest.param(
AttentionBackendCase(backend=AttentionBackendEnum.TRITON_ATTN), id="TRITON_ATTN"
)
ROCM_ATTN = pytest.param(
AttentionBackendCase(backend=AttentionBackendEnum.ROCM_ATTN),
id="ROCM_ATTN",
marks=pytest.mark.skipif(
not current_platform.is_rocm(),
reason="ROCm attention only for AMD",
),
)
ROCM_AITER_UNIFIED_ATTN = pytest.param(
AttentionBackendCase(backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN),
id="ROCM_AITER_UNIFIED_ATTN",
marks=pytest.mark.skipif(
not is_aiter_found_and_supported(),
reason="ROCM_AITER_UNIFIED_ATTN only for AMD when AITER is installed",
),
)
# Models
llama3_8b = ModelFusionInfo(
model_name="meta-llama/Llama-3.1-8B-Instruct",
@@ -49,7 +69,6 @@ llama3_8b_fp8 = ModelFusionInfo(
llama3_8b_fp4 = ModelFusionInfo(
model_name="nvidia/Llama-3.1-8B-Instruct-FP4",
matches=lambda n_layers: Matches(
rms_quant_fusion=0,
act_quant_fusion=n_layers,
attn_quant_fusion=n_layers,
ar_rms_fusion=n_layers * 2 + 1,
@@ -79,7 +98,6 @@ llama4_scout_fp4 = ModelFusionInfo(
model_name="nvidia/Llama-4-Scout-17B-16E-Instruct-NVFP4",
hf_overrides=lambda n_layers: {"text_config": {"num_hidden_layers": n_layers}},
matches=lambda n_layers: Matches(
rms_quant_fusion=0,
attn_quant_fusion=n_layers,
ar_rms_fusion=n_layers * 2,
sequence_parallel=n_layers * 2,
+39 -3
View File
@@ -5,6 +5,7 @@ from collections.abc import Callable
import pytest
from vllm.config import PassConfig
from vllm.platforms import current_platform
from vllm.utils.flashinfer import is_flashinfer_fp8_blockscale_gemm_supported
from .common import (
@@ -16,6 +17,8 @@ from .common import (
)
from .models import (
FLASHINFER_ATTN,
ROCM_AITER_UNIFIED_ATTN,
ROCM_ATTN,
TRITON_ATTN,
llama3_8b_fp4,
llama3_8b_fp8,
@@ -29,12 +32,33 @@ from .models import (
"model_name, matches_fn, model_kwargs, hf_overrides, use_deepgemm",
[
(*llama3_8b_fp8, False),
(*llama4_scout_fp8, False),
(*qwen3_a3b_fp8, False),
(*qwen3_a3b_fp8, True),
pytest.param(
*llama4_scout_fp8,
False,
marks=pytest.mark.skipif(
not current_platform.is_cuda(),
reason="Llama4 Scout FP8 only supported on CUDA",
),
),
pytest.param(
*qwen3_a3b_fp8,
True,
marks=pytest.mark.skipif(
not current_platform.is_cuda(), reason="DeepGemm only supported on CUDA"
),
),
],
)
@pytest.mark.parametrize(
"attn_backend",
[
TRITON_ATTN,
FLASHINFER_ATTN,
ROCM_ATTN,
ROCM_AITER_UNIFIED_ATTN,
],
)
@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
@pytest.mark.parametrize("n_layers", [6])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
@@ -81,6 +105,8 @@ def test_tp1_fp8_fusions(
),
)
use_aiter = current_platform.is_rocm() and ("qwen" in model_name.lower())
matches_check = [
"rms_quant_fusion",
"act_quant_fusion",
@@ -88,6 +114,15 @@ def test_tp1_fp8_fusions(
"attn_quant_fusion",
]
if use_aiter:
matches_check[0] = "aiter_rms_quant_fusion"
matches = matches._replace(aiter_rms_quant_fusion=matches.rms_quant_fusion)
# TODO: enable the `norm_rope_fusion` test,
# On ROCm norm_rope_fusion is only supported without
# enabling AITER.
matches_check.remove("norm_rope_fusion")
run_e2e_fusion_test(
model_name,
matches,
@@ -96,6 +131,7 @@ def test_tp1_fp8_fusions(
compilation_config,
matches_check,
use_deepgemm=use_deepgemm,
use_aiter=use_aiter,
)
@@ -5,6 +5,7 @@ from collections.abc import Callable
import pytest
from vllm.config import PassConfig
from vllm.platforms import current_platform
from ...utils import multi_gpu_test
from .common import (
@@ -26,6 +27,8 @@ from .models import (
qwen3_a3b_fp8,
)
pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
@@ -5,6 +5,7 @@ from collections.abc import Callable
import pytest
from vllm.config import PassConfig
from vllm.platforms import current_platform
from ...utils import multi_gpu_test
from .common import (
@@ -23,6 +24,8 @@ from .models import (
qwen3_a3b,
)
pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
@@ -36,6 +36,8 @@ from vllm.platforms import current_platform
from vllm.utils.system_utils import update_environment_variables
from vllm.utils.torch_utils import set_random_seed
pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
FP8_DTYPE = current_platform.fp8_dtype()
prompts = [
"Hello, my name is",
@@ -182,8 +182,24 @@ TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS
"model_class, enable_quant_fp8_custom_op, force_kernel",
list(itertools.product([TestSiluMulFp8QuantModel], [True, False], TEST_KERNELS))
+ [
(TestSiluMulNvfp4QuantModel, False, None),
(TestSiluMulGroupFp8QuantModel, False, None),
pytest.param(
TestSiluMulNvfp4QuantModel,
False,
None,
marks=pytest.mark.skipif(
not current_platform.is_cuda(), reason="CUDA only"
),
),
# GroupFP8Quant fusion only works with AITER on ROCm.
# and the enable_quant_fp8_custom_op must be True.
pytest.param(
TestSiluMulGroupFp8QuantModel,
True,
None,
marks=pytest.mark.skipif(
not current_platform.is_rocm(), reason="ROCm only"
),
),
],
)
@pytest.mark.skipif(
@@ -201,6 +217,7 @@ def test_fusion_silu_and_mul_quant(
enable_silu_mul_custom_op: bool,
enable_quant_fp8_custom_op: bool,
force_kernel: FP8ScaledMMLinearKernel | None,
monkeypatch: pytest.MonkeyPatch,
):
if model_class is TestSiluMulNvfp4QuantModel and not is_nvfp4_supported():
pytest.skip("NVFP4 is not supported on this GPU.")
@@ -227,13 +244,16 @@ def test_fusion_silu_and_mul_quant(
),
)
with set_current_vllm_config(config):
with set_current_vllm_config(config), monkeypatch.context() as m:
fusion_passes = [ActivationQuantFusionPass(config)]
if IS_AITER_FOUND:
if IS_AITER_FOUND and model_class is TestSiluMulGroupFp8QuantModel:
from vllm._aiter_ops import rocm_aiter_ops
from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
RocmAiterSiluMulFp8GroupQuantFusionPass,
)
m.setenv("VLLM_ROCM_USE_AITER", "1")
rocm_aiter_ops.refresh_env_variables()
fusion_passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
passes = [NoOpEliminationPass(config), *fusion_passes, PostCleanupPass(config)]
@@ -0,0 +1,61 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Regression test: ``/tokenize`` must expand image placeholders for VLM models.
Fixed by PR #34560 ("Move InputPreprocessor into Renderer (2/2)").
Before that change, ``/tokenize`` returned ~26 tokens for a message with an
image instead of the expected 1451. Confirmed broken on 0.15.1 and 0.16.0.
"""
import json
import pytest
import requests
from ...utils import RemoteOpenAIServer
MODEL_NAME = "Qwen/Qwen2.5-VL-3B-Instruct"
@pytest.fixture(scope="module")
def server():
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"4096",
"--max-num-seqs",
"5",
"--enforce-eager",
"--limit-mm-per-prompt",
json.dumps({"image": 1}),
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
def test_tokenize_chat_expands_image_placeholders(
server: RemoteOpenAIServer,
local_asset_server,
):
image_url = local_asset_server.url_for("stop_sign.jpg")
messages = [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": image_url}},
{"type": "text", "text": "Describe this image."},
],
}
]
response = requests.post(
server.url_for("tokenize"),
json={"model": MODEL_NAME, "messages": messages},
)
response.raise_for_status()
# stop_sign.jpg (1300x876) produces 1451 tokens after expansion.
# Without expansion the count would be ~26 (text + one placeholder).
assert response.json()["count"] == 1451
@@ -108,6 +108,23 @@ async def test_long_audio_request(mary_had_lamb, whisper_client):
assert out_usage["seconds"] == 161, out_usage["seconds"]
@pytest.mark.asyncio
async def test_invalid_audio_file(whisper_client):
"""Corrupted audio should surface as HTTP 400."""
invalid_audio = io.BytesIO(b"not a valid audio file")
invalid_audio.name = "invalid.wav"
with pytest.raises(openai.BadRequestError) as exc_info:
await whisper_client.audio.transcriptions.create(
model=MODEL_NAME,
file=invalid_audio,
language="en",
)
assert exc_info.value.status_code == 400
assert "Invalid or unsupported audio file" in exc_info.value.message
@pytest.mark.asyncio
async def test_completion_endpoints(whisper_client):
# text to text model
@@ -8,4 +8,4 @@ server_args: >-
--max-model-len 4096
--data-parallel-size 8
--enable-expert-parallel
--speculative-config '{"method":"mtp","num_speculative_tokens":1}'
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
@@ -8,4 +8,4 @@ server_args: >-
--max-model-len 4096
--tensor-parallel-size 8
--enable-expert-parallel
--speculative-config '{"method":"mtp","num_speculative_tokens":1}'
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
@@ -8,4 +8,4 @@ server_args: >-
--max-model-len 4096
--data-parallel-size 8
--enable-expert-parallel
--speculative-config '{"method":"mtp","num_speculative_tokens":1}'
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
@@ -8,4 +8,4 @@ server_args: >-
--max-model-len 4096
--tensor-parallel-size 8
--enable-expert-parallel
--speculative-config '{"method":"mtp","num_speculative_tokens":1}'
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'
@@ -8,8 +8,5 @@ Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml
Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml
Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml
Qwen3-30B-A3B-Fp8-CT-Channel-vllm-cutlass.yaml
Llama-4-Scout-Fp8-ModelOpt-fi-cutlass.yaml
Llama-4-Scout-Fp8-ModelOpt-marlin.yaml
Llama-4-Scout-Fp8-ModelOpt-triton.yaml
Qwen3-30B-A3B-BF16-fi-cutlass.yaml
Qwen3-30B-A3B-BF16-triton.yaml
Qwen3-30B-A3B-BF16-triton.yaml
@@ -0,0 +1,196 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import patch
import pytest
import torch
from vllm.utils.flashinfer import (
can_use_trtllm_attention,
supports_trtllm_attention,
use_trtllm_attention,
)
MODEL_CONFIGS = {
"Llama-3-70B": dict(num_qo_heads=64, num_kv_heads=8),
"Llama-3-8B": dict(num_qo_heads=32, num_kv_heads=8),
"Qwen2.5-0.5B": dict(num_qo_heads=14, num_kv_heads=2),
"Mistral-7B": dict(num_qo_heads=32, num_kv_heads=8),
"Gemma-2-9B": dict(num_qo_heads=8, num_kv_heads=4),
"Falcon-40B": dict(num_qo_heads=128, num_kv_heads=8),
}
def get_config(model: str) -> dict:
"""Return the attention config for a model."""
return MODEL_CONFIGS[model]
DEFAULT_KWARGS = dict(
**get_config("Llama-3-70B"),
num_tokens=128,
max_seq_len=4096,
dcp_world_size=1,
kv_cache_dtype="auto",
q_dtype=torch.bfloat16,
is_prefill=False,
force_use_trtllm=None,
has_sinks=False,
has_spec=False,
)
def _call(**overrides) -> bool:
kwargs = {**DEFAULT_KWARGS, **overrides}
return use_trtllm_attention(**kwargs)
@pytest.fixture(autouse=True)
def _clear_supports_cache():
"""Clear functools.cache to ensure each test runs independently."""
supports_trtllm_attention.cache_clear()
# supports_trtllm_attention
@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=True)
def test_supports_batch_invariant_disables(_mock):
assert supports_trtllm_attention() is False
@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
@patch(
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
return_value=True,
)
@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
def test_supports_sm100_with_artifactory(_art, _cap, _bi):
assert supports_trtllm_attention() is True
@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
@patch(
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
return_value=False,
)
def test_supports_non_sm100_platform(_cap, _bi):
assert supports_trtllm_attention() is False
@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
@patch(
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
return_value=True,
)
@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False)
def test_supports_sm100_without_artifactory(_art, _cap, _bi):
assert supports_trtllm_attention() is False
# can_use_trtllm_attention
@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=False)
def test_can_use_force_disabled(_mock):
cfg = get_config("Llama-3-70B")
assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is False
@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_can_use_compatible_heads(_sup, _force):
cfg = get_config("Llama-3-70B")
assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is True
@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_can_use_incompatible_heads(_sup, _force):
assert can_use_trtllm_attention(40, 6) is False
@pytest.mark.parametrize("model", list(MODEL_CONFIGS.keys()))
@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
def test_can_use_platform_unsupported(_sup, _force, model):
cfg = get_config(model)
assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is False
# use_trtllm_attention
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_force_off(_mock):
assert _call(force_use_trtllm=False) is False
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_dcp_fallback(_mock):
assert _call(dcp_world_size=2) is False
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
def test_use_platform_unsupported(_mock):
assert _call() is False
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
def test_use_platform_unsupported_force_on_still_false(_mock):
assert _call(force_use_trtllm=True) is False
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_incompatible_heads(_mock):
assert _call(num_qo_heads=40, num_kv_heads=6) is False
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_incompatible_heads_force_on_still_false(_mock):
assert _call(num_qo_heads=40, num_kv_heads=6, force_use_trtllm=True) is False
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_spec_decode_enables(_mock):
assert _call(has_spec=True, is_prefill=False) is True
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
@patch(
"vllm.utils.flashinfer.current_platform.fp8_dtype",
return_value=torch.float8_e4m3fn,
)
def test_use_fp8_query_forces_trtllm(_fp8, _sup):
assert _call(q_dtype=torch.float8_e4m3fn) is True
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_sinks_force_trtllm(_mock):
assert _call(has_sinks=True) is True
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_auto_prefill_kv_auto(_mock):
assert _call(is_prefill=True, kv_cache_dtype="auto") is True
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_auto_prefill_kv_fp8(_mock):
assert _call(is_prefill=True, kv_cache_dtype="fp8") is False
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_auto_decode_small_batch(_mock):
assert _call(is_prefill=False, num_tokens=128, kv_cache_dtype="auto") is True
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_auto_decode_large_batch(_mock):
assert _call(is_prefill=False, num_tokens=512, kv_cache_dtype="auto") is False
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
def test_use_force_on(_mock):
assert _call(force_use_trtllm=True) is True
@@ -17,13 +17,13 @@ from .mk_objects import (
def make_config_arg_parser(description: str):
def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalize:
def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalizeModular:
for pf in MK_ALL_PREPARE_FINALIZE_TYPES:
if pf.__name__ == s:
return pf
raise ValueError(f"Cannot find a PrepareFinalize type that matches {s}")
def to_experts_class_type(s: str) -> mk.FusedMoEPermuteExpertsUnpermute:
def to_experts_class_type(s: str) -> mk.FusedMoEExpertsModular:
for fe in MK_FUSED_EXPERT_TYPES:
if fe.__name__ == s:
return fe
@@ -66,7 +66,7 @@ class Config:
quant_config: TestMoEQuantConfig | None
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize
fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute
fused_experts_type: mk.FusedMoEExperts
fused_moe_chunk_size: int | None
world_size: int
@@ -566,7 +566,7 @@ def make_modular_kernel(
config: Config,
vllm_config: VllmConfig,
quant_config: FusedMoEQuantConfig,
) -> mk.FusedMoEModularKernel:
) -> mk.FusedMoEKernel:
def next_power_of_2(x):
import math
@@ -613,7 +613,7 @@ def make_modular_kernel(
config.N,
)
modular_kernel = mk.FusedMoEModularKernel(
modular_kernel = mk.FusedMoEKernel(
prepare_finalize=prepare_finalize,
fused_experts=fused_experts,
inplace=False,
@@ -667,6 +667,7 @@ def run_modular_kernel(
"w2": rank_weights.w2,
"topk_weights": rank_tensors.topk_weights,
"topk_ids": topk_ids,
"activation": MoEActivation.SILU,
"expert_map": rank_tensors.expert_map,
"global_num_experts": config.E,
"apply_router_weight_on_input": config.topk == 1
@@ -684,6 +685,6 @@ def run_modular_kernel(
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp,
):
out = mk.forward(**mk_kwargs)
out = mk.apply(**mk_kwargs)
return out
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
NaiveBatchedExperts,
)
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
MoEPrepareAndFinalizeNoDPEPModular,
)
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
TritonOrDeepGemmExperts,
@@ -71,12 +71,14 @@ class ExpertInfo:
needs_aiter: bool = False
PREPARE_FINALIZE_INFO: dict[mk.FusedMoEPrepareAndFinalize, PrepareFinalizeInfo] = {}
EXPERT_INFO: dict[mk.FusedMoEPermuteExpertsUnpermute, ExpertInfo] = {}
MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []
MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []
MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []
MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEPermuteExpertsUnpermute] = []
PREPARE_FINALIZE_INFO: dict[
mk.FusedMoEPrepareAndFinalizeModular, PrepareFinalizeInfo
] = {}
EXPERT_INFO: dict[mk.FusedMoEExpertsModular, ExpertInfo] = {}
MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []
MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []
MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []
MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEExpertsModular] = []
standard_format = mk.FusedMoEActivationFormat.Standard
batched_format = mk.FusedMoEActivationFormat.BatchedExperts
@@ -162,7 +164,7 @@ def expert_info(kind) -> ExpertInfo:
register_prepare_and_finalize(
MoEPrepareAndFinalizeNoEP,
MoEPrepareAndFinalizeNoDPEPModular,
standard_format,
common_float_types,
blocked_quantization_support=True,
@@ -239,14 +241,14 @@ if has_mori():
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501
FlashInferCutlassMoEPrepareAndFinalize,
FlashInferA2APrepareAndFinalize,
)
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
FlashInferExperts,
)
register_prepare_and_finalize(
FlashInferCutlassMoEPrepareAndFinalize,
FlashInferA2APrepareAndFinalize,
standard_format,
nvfp4_types + fp8_types,
blocked_quantization_support=True,
@@ -430,12 +432,12 @@ def make_cutlass_strides(
def make_fused_experts(
fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,
fused_experts_type: mk.FusedMoEExpertsModular,
moe: FusedMoEConfig,
quant_config: FusedMoEQuantConfig,
num_dispatchers: int,
N: int,
) -> mk.FusedMoEPermuteExpertsUnpermute:
) -> mk.FusedMoEExpertsModular:
if (
fused_experts_type.activation_format()
== mk.FusedMoEActivationFormat.BatchedExperts
@@ -72,7 +72,7 @@ def profile_modular_kernel(
"apply_router_weight_on_input": config.topk == 1,
}
do_profile(mk.forward, mk_kwargs, pgi, config)
do_profile(mk.apply, mk_kwargs, pgi, config)
def rank_worker(
+12 -5
View File
@@ -4,6 +4,7 @@
import pytest
import torch
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
BatchedDeepGemmExperts,
)
@@ -12,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
BatchedPrepareAndFinalize,
BatchedTritonExperts,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported
from .test_deepgemm import make_block_quant_fp8_weights
@@ -74,19 +75,22 @@ def test_batched_deepgemm_vs_triton(
quant_config=quant_config,
moe_config=make_dummy_moe_config(),
)
mk_triton = FusedMoEModularKernel(
mk_triton = FusedMoEKernel(
prep_finalize,
triton_experts,
inplace=False,
)
out_triton = mk_triton(
out_triton = mk_triton.apply(
hidden_states=a,
w1=w1,
w2=w2,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=MoEActivation.SILU,
global_num_experts=E,
expert_map=None,
apply_router_weight_on_input=False,
)
# deepgemm
@@ -96,19 +100,22 @@ def test_batched_deepgemm_vs_triton(
quant_config=quant_config,
moe_config=make_dummy_moe_config(),
)
mk_deepgemm = FusedMoEModularKernel(
mk_deepgemm = FusedMoEKernel(
prep_finalize,
deepgemm_experts,
inplace=False,
)
out_deepgemm = mk_deepgemm(
out_deepgemm = mk_deepgemm.apply(
hidden_states=a,
w1=w1,
w2=w2,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=MoEActivation.SILU,
global_num_experts=E,
expert_map=None,
apply_router_weight_on_input=False,
)
diff = calc_diff(out_deepgemm, out_triton)
+29 -8
View File
@@ -21,15 +21,16 @@ from vllm.model_executor.layers.fused_moe import (
fused_experts,
fused_topk,
)
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
fp8_w8a8_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
_valid_deep_gemm_shape,
)
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
TritonOrDeepGemmExperts,
)
@@ -193,7 +194,17 @@ def test_w8a8_block_fp8_fused_moe(
a, w1, w2, topk_weights, topk_ids, quant_config=quant_config
)
m_out = m_fused_moe(a, w1, w2, topk_weights, topk_ids)
m_out = m_fused_moe.apply(
a,
w1,
w2,
topk_weights,
topk_ids,
activation=MoEActivation.SILU,
apply_router_weight_on_input=False,
expert_map=None,
global_num_experts=w1.shape[0],
)
# 0.039 only needed for M >= 8192
tol = 0.035 if M < 8192 else 0.039
@@ -252,23 +263,33 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)
w2_scale=w2_s,
block_shape=block_size,
)
moe_config = make_dummy_moe_config()
deep_gemm_experts = mk.FusedMoEModularKernel(
prepare_finalize=MoEPrepareAndFinalizeNoEP(),
deep_gemm_experts = mk.FusedMoEKernel(
prepare_finalize=maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
fused_experts=TritonOrDeepGemmExperts(
moe_config=make_dummy_moe_config(),
moe_config=moe_config,
quant_config=quant_config,
),
inplace=False,
)
def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids):
return deep_gemm_experts(
return deep_gemm_experts.apply(
hidden_states=a,
w1=w1,
w2=w2,
topk_weights=topk_weights,
topk_ids=topk_ids,
global_num_experts=E,
activation=MoEActivation.SILU,
apply_router_weight_on_input=False,
expert_map=False,
)
# Set the context to avoid lots of warning spam.
+37 -21
View File
@@ -13,6 +13,9 @@ from vllm import _custom_ops as ops
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
FUSED_MOE_UNQUANTIZED_CONFIG,
FusedMoEQuantConfig,
@@ -22,9 +25,6 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
CutlassExpertsFp8,
run_cutlass_moe_fp8,
)
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
@@ -197,20 +197,26 @@ def run_with_expert_maps(
for kwargs, new_quant_config in slice_experts():
w2 = kwargs["w2"]
a = kwargs["hidden_states"]
kernel = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
moe_config = make_dummy_moe_config(
num_experts=w2.shape[0],
hidden_dim=w2.shape[1],
intermediate_size_per_partition=w2.shape[2],
in_dtype=a.dtype,
)
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=new_quant_config,
allow_new_interface=True,
use_monolithic=False,
),
CutlassExpertsFp8(
moe_config=make_dummy_moe_config(
num_experts=w2.shape[0],
hidden_dim=w2.shape[1],
intermediate_size_per_partition=w2.shape[2],
in_dtype=a.dtype,
),
moe_config=moe_config,
quant_config=new_quant_config,
),
inplace=False,
)
out_tensor = out_tensor + kernel(**kwargs)
out_tensor = out_tensor + kernel.apply(**kwargs)
return out_tensor
@@ -252,25 +258,35 @@ def run_8_bit(
"w2": moe_tensors.w2_q, # type: ignore[union-attr]
"topk_weights": topk_weights,
"topk_ids": topk_ids,
"global_num_experts": moe_tensors.w1_q.shape[0], # type: ignore[union-attr]
"activation": MoEActivation.SILU,
"expert_map": None,
"apply_router_weight_on_input": False,
}
num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined]
with_ep = num_local_experts is not None or num_local_experts == num_experts
if not with_ep:
kernel = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
moe_config = make_dummy_moe_config(
num_experts=moe_tensors.w2_q.shape[0], # type: ignore[union-attr]
hidden_dim=moe_tensors.w2_q.shape[1], # type: ignore[union-attr]
intermediate_size_per_partition=moe_tensors.w2_q.shape[2], # type: ignore[union-attr]
in_dtype=moe_tensors.a.dtype,
)
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
CutlassExpertsFp8(
moe_config=make_dummy_moe_config(
num_experts=moe_tensors.w2_q.shape[0], # type: ignore[union-attr]
hidden_dim=moe_tensors.w2_q.shape[1], # type: ignore[union-attr]
intermediate_size_per_partition=moe_tensors.w2_q.shape[2], # type: ignore[union-attr]
in_dtype=moe_tensors.a.dtype,
),
moe_config=moe_config,
quant_config=quant_config,
),
inplace=False,
)
return kernel(**kwargs)
return kernel.apply(**kwargs)
assert num_local_experts is not None
return run_with_expert_maps(
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
fp8_w8a8_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
from vllm.utils.deep_gemm import (
get_mk_alignment_for_contiguous_layout,
is_deep_gemm_e8m0_used,
@@ -170,7 +170,7 @@ def make_ll_modular_kernel(
q_dtype: torch.dtype | None,
test_config: TestConfig,
quant_config: FusedMoEQuantConfig,
) -> FusedMoEModularKernel:
) -> FusedMoEKernel:
assert test_config.low_latency
assert test_config.use_fp8_dispatch is not None
@@ -195,7 +195,7 @@ def make_ll_modular_kernel(
quant_config=quant_config,
moe_config=make_dummy_moe_config(),
)
return FusedMoEModularKernel(
return FusedMoEKernel(
prepare_finalize=a2a,
fused_experts=fused_experts,
inplace=False,
@@ -210,7 +210,7 @@ def make_ht_modular_kernel(
q_dtype: torch.dtype | None,
test_config: TestConfig,
quant_config: FusedMoEQuantConfig,
) -> FusedMoEModularKernel:
) -> FusedMoEKernel:
assert not test_config.low_latency
assert test_config.use_fp8_dispatch is None
@@ -228,7 +228,7 @@ def make_ht_modular_kernel(
moe_config=make_dummy_moe_config(),
quant_config=quant_config,
)
return FusedMoEModularKernel(
return FusedMoEKernel(
prepare_finalize=a2a,
fused_experts=fused_experts,
inplace=False,
@@ -242,11 +242,11 @@ def make_modular_kernel(
num_local_experts: int,
test_tensors: TestTensors,
quant_config: FusedMoEQuantConfig,
) -> FusedMoEModularKernel:
) -> FusedMoEKernel:
q_dtype = torch.float8_e4m3fn
test_config = test_tensors.config
mk: FusedMoEModularKernel
mk: FusedMoEKernel
# Make modular kernel
if test_config.low_latency:
max_tokens_per_rank = max(64, next_power_of_2(test_tensors.rank_tokens.size(0)))
@@ -307,7 +307,7 @@ def deepep_deepgemm_moe_impl(
)
# Make modular kernel
mk: FusedMoEModularKernel = make_modular_kernel(
mk: FusedMoEKernel = make_modular_kernel(
pg=pg,
pgi=pgi,
dp_size=dp_size,
@@ -319,7 +319,7 @@ def deepep_deepgemm_moe_impl(
with with_dp_metadata(
M=test_tensors.rank_tokens.size(0), world_size=pgi.world_size
):
out = mk.forward(
out = mk.apply(
hidden_states=test_tensors.rank_tokens,
w1=w1,
w2=w2,
+5 -5
View File
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
)
from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
per_token_group_quant_fp8,
)
@@ -135,7 +135,7 @@ def make_modular_kernel(
q_dtype: torch.dtype | None,
use_fp8_dispatch: bool,
quant_config: FusedMoEQuantConfig,
) -> FusedMoEModularKernel:
) -> FusedMoEKernel:
ht_args: DeepEPHTArgs | None = None
ll_args: DeepEPLLArgs | None = None
@@ -180,7 +180,7 @@ def make_modular_kernel(
quant_config=quant_config,
)
mk = FusedMoEModularKernel(
mk = FusedMoEKernel(
prepare_finalize=a2a,
fused_experts=fused_experts,
inplace=False,
@@ -242,7 +242,7 @@ def deep_ep_moe_impl(
)
# Make modular kernel
mk: FusedMoEModularKernel = make_modular_kernel(
mk: FusedMoEKernel = make_modular_kernel(
pg,
pgi,
low_latency_mode,
@@ -255,7 +255,7 @@ def deep_ep_moe_impl(
quant_config,
)
out = mk.forward(
out = mk.apply(
hidden_states=rank_tokens_chunk,
w1=w1,
w2=w2,
+20 -7
View File
@@ -14,13 +14,16 @@ import torch
# vLLM fused-expert reference (Triton fallback + DeepGEMM option)
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.kernels.moe.utils import make_dummy_moe_config
from vllm.model_executor.layers.fused_moe.activation import (
MoEActivation,
)
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
fp8_w8a8_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
TritonOrDeepGemmExperts,
)
@@ -108,11 +111,17 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
a1_scale=a1_scale,
block_shape=block_size,
)
moe_config = make_dummy_moe_config()
deep_gemm_experts = mk.FusedMoEModularKernel(
prepare_finalize=MoEPrepareAndFinalizeNoEP(),
deep_gemm_experts = mk.FusedMoEKernel(
prepare_finalize=maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
fused_experts=TritonOrDeepGemmExperts(
moe_config=make_dummy_moe_config(),
moe_config=moe_config,
quant_config=quant_config,
),
inplace=False,
@@ -130,12 +139,16 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
)
# DeepGemm
out_deepgemm = deep_gemm_experts(
out_deepgemm = deep_gemm_experts.apply(
hidden_states=tokens_bf16,
w1=w1,
w2=w2,
topk_weights=topk_weights,
topk_ids=topk_ids,
global_num_experts=num_experts,
activation=MoEActivation.SILU,
apply_router_weight_on_input=False,
expert_map=None,
)
diff = calc_diff(out_deepgemm, out_triton)
assert diff < 0.001, f"Diff exceeded 1%: {diff}"
+49 -21
View File
@@ -8,6 +8,9 @@ import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
@@ -15,16 +18,14 @@ from vllm.model_executor.layers.fused_moe.config import (
RoutingMethodType,
fp8_w8a8_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
TrtLlmFp8Experts,
)
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
FlashInferExperts,
)
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
apply_fi_trtllm_fp8_per_tensor_moe,
register_scales_for_trtllm_fp8_per_tensor_moe,
rotate_weights_for_fi_trtllm_fp8_per_tensor_moe,
swap_w13_to_w31,
)
@@ -115,6 +116,7 @@ class TestData:
e: int,
is_trtllm: bool,
activation: MoEActivation = MoEActivation.SILU,
topk: int = 1,
) -> "TestData":
is_gated = activation.is_gated
@@ -152,13 +154,6 @@ class TestData:
rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(
layer.w13_weight, layer.w2_weight, is_gated
)
register_scales_for_trtllm_fp8_per_tensor_moe(
layer,
layer.w13_weight_scale,
layer.w13_input_scale,
layer.w2_weight_scale,
layer.w2_input_scale,
)
layer.custom_routing_function = Llama4MoE.custom_routing_function
layer.routing_method_type = RoutingMethodType.Llama4
layer.renormalize = False
@@ -166,6 +161,21 @@ class TestData:
layer.ep_rank = 0
layer.local_num_experts = e
layer.moe = FusedMoEConfig(
num_experts=e,
experts_per_token=topk,
hidden_dim=k,
intermediate_size_per_partition=n,
num_local_experts=e,
num_logical_experts=e,
moe_parallel_config=layer.moe_parallel_config,
in_dtype=hidden_states.dtype,
is_act_and_mul=is_gated,
routing_method=layer.routing_method_type,
activation=activation,
device=w13_quantized.device,
)
return TestData(
hidden_states=hidden_states,
w13_quantized=w13_quantized,
@@ -230,16 +240,29 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
quant_config=quant_config,
)
flashinfer_output = apply_fi_trtllm_fp8_per_tensor_moe(
layer=td.layer,
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=td.layer.moe,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=True,
),
TrtLlmFp8Experts(
moe_config=td.layer.moe,
quant_config=quant_config,
),
)
flashinfer_output = kernel.apply_monolithic(
hidden_states=td.hidden_states,
w1=td.layer.w13_weight,
w2=td.layer.w2_weight,
router_logits=score,
routing_bias=None,
activation=activation,
global_num_experts=e,
top_k=topk,
num_expert_group=None,
topk_group=None,
expert_map=None,
apply_router_weight_on_input=True,
routed_scaling_factor=1.0,
)
check_accuracy(
@@ -329,8 +352,13 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
routing_method=RoutingMethodType.TopK,
)
kernel = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
FlashInferExperts(
moe_config=moe_config,
quant_config=quant_config,
@@ -338,7 +366,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
inplace=False,
)
flashinfer_cutlass_output = kernel(
flashinfer_cutlass_output = kernel.apply(
td.hidden_states,
td.layer.w13_weight,
td.layer.w2_weight,
+15 -7
View File
@@ -14,6 +14,9 @@ from vllm import _custom_ops as ops
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe import fused_topk
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
@@ -23,10 +26,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
FlashInferExperts,
is_valid_flashinfer_cutlass_fused_moe,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
from vllm.utils.torch_utils import set_random_seed
@@ -107,19 +107,27 @@ def test_flashinfer_fp4_moe_no_graph(
routing_method=RoutingMethodType.TopK,
)
flashinfer_experts = FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
flashinfer_experts = FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
FlashInferExperts(moe_config=moe_config, quant_config=quant_config),
inplace=False,
)
flashinfer_output = flashinfer_experts(
flashinfer_output = flashinfer_experts.apply(
hidden_states=a,
w1=w1_q,
w2=w2_q,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=activation,
global_num_experts=e,
expert_map=None,
apply_router_weight_on_input=False,
)
# Reference check:
@@ -221,16 +221,16 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group
)
marlin_output = fused_marlin_moe(
a,
w1_marlin,
w2_marlin,
None,
None,
w1_scales_marlin,
w2_scales_marlin,
None, # gating_output not needed when topk_weights/ids provided
topk_weights,
topk_ids,
hidden_states=a,
w1=w1_marlin,
w2=w2_marlin,
bias1=None,
bias2=None,
w1_scale=w1_scales_marlin,
w2_scale=w2_scales_marlin,
topk_weights=topk_weights,
topk_ids=topk_ids,
quant_type_id=scalar_types.uint4b8.id,
global_num_experts=e,
expert_map=None,
global_scale1=None,
@@ -244,7 +244,6 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group
w1_zeros=None,
w2_zeros=None,
input_dtype=dtype,
quant_type_id=scalar_types.uint4b8.id,
is_k_full=True,
)
@@ -168,7 +168,6 @@ FUSED_MOE_CHUNK_SIZEs = [None, 16]
def is_nyi_config(config: Config) -> bool:
# We know these configs to be legitimate. but still fail.
info = expert_info(config.fused_experts_type)
if info.needs_matching_quant:
# The triton kernels expect both per-act-token-quant and
# per-out-ch-quant or neither.
@@ -259,7 +258,7 @@ def test_modular_kernel_combinations_multigpu(
dtype: torch.dtype,
quant_config: TestMoEQuantConfig | None,
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,
fused_experts_type: mk.FusedMoEExperts,
chunk_size: int | None,
world_size: int,
pytestconfig,
@@ -301,7 +300,7 @@ def test_modular_kernel_combinations_singlegpu(
dtype: torch.dtype,
quant_config: TestMoEQuantConfig | None,
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,
fused_experts_type: mk.FusedMoEExperts,
chunk_size: int | None,
world_size: int,
pytestconfig,
@@ -7,6 +7,7 @@ Test modular OAI Triton MoE
import pytest
import torch
from tests.utils import wait_for_gpu_memory_to_clear
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.utils.import_utils import has_triton_kernels
@@ -24,15 +25,15 @@ from triton_kernels.tensor_details import layout
from triton_kernels.testing import assert_close
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
OAITritonExperts,
UnfusedOAITritonExperts,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
@@ -174,19 +175,25 @@ def oai_triton_moe_impl(
w1_scale=w1_scale,
w2_scale=w2_scale,
)
moe_config = make_dummy_moe_config()
if unfused:
fused_experts = UnfusedOAITritonExperts(make_dummy_moe_config(), quant_config)
fused_experts = UnfusedOAITritonExperts(moe_config, quant_config)
else:
fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config)
fused_experts = OAITritonExperts(moe_config, quant_config)
mk = FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
mk = FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
fused_experts,
inplace=False,
)
return mk.forward(
return mk.apply(
hidden_states=x,
w1=w1,
w2=w2,
@@ -217,6 +224,7 @@ def test_oai_triton_moe(
unfused: bool,
workspace_init,
):
wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1)
set_random_seed(0)
(
w1,
+6 -2
View File
@@ -346,14 +346,16 @@ def test_fused_moe(
expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
return m_fused_moe_fn(
return m_fused_moe_fn.apply(
a,
w1,
w2,
topk_weights,
topk_ids,
activation=MoEActivation.SILU,
global_num_experts=global_num_experts,
expert_map=expert_map,
apply_router_weight_on_input=False,
)
fused_moe_fn = functools.partial(fused_moe, renormalize=False)
@@ -500,14 +502,16 @@ def test_naive_block_assignment_moe(
expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
return m_fused_moe_fn(
return m_fused_moe_fn.apply(
a,
w1,
w2,
topk_weights,
topk_ids,
activation=MoEActivation.SILU,
global_num_experts=global_num_experts,
expert_map=expert_map,
apply_router_weight_on_input=False,
)
fused_moe_fn = functools.partial(fused_moe, renormalize=False)
+24 -8
View File
@@ -15,12 +15,15 @@ from vllm import _custom_ops as ops
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe import fused_topk
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
CutlassExpertsFp4,
)
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
make_moe_prepare_and_finalize_no_dp_ep,
)
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
@@ -89,22 +92,32 @@ def test_cutlass_fp4_moe_no_graph(
w1_scale=w1_blockscale,
w2_scale=w2_blockscale,
)
moe_config = make_dummy_moe_config()
kernel = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
CutlassExpertsFp4(
moe_config=make_dummy_moe_config(),
moe_config=moe_config,
quant_config=quant_config,
),
inplace=False,
)
cutlass_output = kernel(
cutlass_output = kernel.apply(
hidden_states=a,
w1=w1_q,
w2=w2_q,
topk_weights=topk_weights,
topk_ids=topk_ids,
global_num_experts=e,
activation=mk.MoEActivation.SILU,
apply_router_weight_on_input=False,
expert_map=None,
)
# Reference check:
@@ -207,8 +220,8 @@ def test_cutlass_fp4_moe_swiglustep(
w2_scale=w2_blockscale,
)
kernel = mk.FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
kernel = mk.FusedMoEKernel(
make_moe_prepare_and_finalize_no_dp_ep(use_monolithic=False),
CutlassExpertsFp4(
moe_config=make_dummy_moe_config(),
quant_config=quant_config,
@@ -216,13 +229,16 @@ def test_cutlass_fp4_moe_swiglustep(
inplace=False,
)
cutlass_output = kernel(
cutlass_output = kernel.apply(
hidden_states=a,
w1=w1_q,
w2=w2_q,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=MoEActivation.SWIGLUSTEP,
global_num_experts=e,
expert_map=None,
apply_router_weight_on_input=False,
)
# Reference: dequantize everything and run torch_moe with swiglustep
+41 -13
View File
@@ -8,6 +8,9 @@ from tests.kernels.quant_utils import per_block_cast_to_int8
from tests.kernels.quantization.nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
@@ -23,10 +26,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
TritonExperts,
fused_experts,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoEP,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
from vllm.utils.deep_gemm import per_block_cast_to_fp8
@@ -125,7 +125,9 @@ def batched_moe(
a2_scale=a2_scale,
)
fused_experts = FusedMoEModularKernel(
moe_config = make_dummy_moe_config()
fused_experts = FusedMoEKernel(
BatchedPrepareAndFinalize(
max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0
),
@@ -133,12 +135,22 @@ def batched_moe(
max_num_tokens=max_num_tokens,
num_dispatchers=1,
quant_config=quant_config,
moe_config=make_dummy_moe_config(),
moe_config=moe_config,
),
inplace=False,
)
return fused_experts(a, w1, w2, topk_weight, topk_ids)
return fused_experts.apply(
a,
w1,
w2,
topk_weight,
topk_ids,
global_num_experts=w1.shape[0],
activation=moe_config.activation,
apply_router_weight_on_input=False,
expert_map=None,
)
def naive_batched_moe(
@@ -166,8 +178,9 @@ def naive_batched_moe(
a1_scale=a1_scale,
a2_scale=a2_scale,
)
moe_config = make_dummy_moe_config()
fused_experts = FusedMoEModularKernel(
fused_experts = FusedMoEKernel(
BatchedPrepareAndFinalize(
max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0
),
@@ -175,12 +188,22 @@ def naive_batched_moe(
max_num_tokens=max_num_tokens,
num_dispatchers=1,
quant_config=quant_config,
moe_config=make_dummy_moe_config(),
moe_config=moe_config,
),
inplace=False,
)
return fused_experts(a, w1, w2, topk_weight, topk_ids)
return fused_experts.apply(
a,
w1,
w2,
topk_weight,
topk_ids,
global_num_experts=w1.shape[0],
activation=moe_config.activation,
apply_router_weight_on_input=False,
expert_map=None,
)
def chunk_scales(
@@ -581,9 +604,14 @@ def modular_triton_fused_moe(
moe_config: FusedMoEConfig,
quant_config: FusedMoEQuantConfig,
shared_experts: torch.nn.Module | None = None,
) -> FusedMoEModularKernel:
return FusedMoEModularKernel(
MoEPrepareAndFinalizeNoEP(),
) -> FusedMoEKernel:
return FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=False,
),
TritonExperts(moe_config, quant_config),
shared_experts,
inplace=False,
+237
View File
@@ -0,0 +1,237 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import dataclasses
import random
from dataclasses import dataclass
import pytest
import torch
from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
from vllm.multimodal.inputs import (
MultiModalFeatureSpec,
MultiModalFieldElem,
MultiModalKwargsItem,
PlaceholderRange,
)
@pytest.fixture(autouse=True, scope="module")
def _force_cpu_default_device():
# _get_mrope_input_positions returns CPU tensors (via torch.from_numpy).
# Ensure the default device is CPU so the rest of the test tensors match.
original = torch.get_default_device()
torch.set_default_device("cpu")
yield
torch.set_default_device(original)
IMAGE_TOKEN_ID = 999
VIDEO_TOKEN_ID = 888
VISION_START_TOKEN_ID = 777
VISION_END_TOKEN_ID = 778
@dataclass
class DummyVisionConfig:
spatial_merge_size: int = 1
@dataclass
class DummyConfig:
image_token_id: int = IMAGE_TOKEN_ID
video_token_id: int = VIDEO_TOKEN_ID
vision_start_token_id: int = VISION_START_TOKEN_ID
vision_end_token_id: int = VISION_END_TOKEN_ID
vision_config: DummyVisionConfig = dataclasses.field(
default_factory=DummyVisionConfig
)
def make_video_embedding(
t, h, w, interleave_text_tokens: tuple[int, int], video_pruning_rate: float = 0.0
):
"""
Helper function to make a video embedding for a given video size and pruning rate.
Args:
t: Number of frames.
h: Number of rows.
w: Number of columns.
interleave_text_tokens: Tuple of minimum and maximum number of text tokens to
interleave with the video.
video_pruning_rate: Pruning rate for the video.
Returns:
Tuple of (unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask)
"""
unpruned_tokens_sequence = []
population = list(range(1, 100))
for _ in range(t):
num_prefix_tokens = random.randint(
interleave_text_tokens[0], interleave_text_tokens[1]
)
prefix_tokens = random.choices(population, k=num_prefix_tokens)
vision_tokens = (
[VISION_START_TOKEN_ID] + [VIDEO_TOKEN_ID] * h * w + [VISION_END_TOKEN_ID]
)
unpruned_tokens_sequence.extend(prefix_tokens)
unpruned_tokens_sequence.extend(vision_tokens)
unpruned_tokens_sequence = torch.tensor(unpruned_tokens_sequence, dtype=torch.long)
video_token_mask = unpruned_tokens_sequence == VIDEO_TOKEN_ID
pruning_mask = torch.bernoulli(video_token_mask.float() * video_pruning_rate).bool() # type: ignore[attr-defined]
# Sanity check that we don't prune what should not be pruned.
assert not pruning_mask[~video_token_mask].any()
retention_mask = ~pruning_mask
pruned_tokens_sequence = unpruned_tokens_sequence[retention_mask]
return unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask
@pytest.mark.parametrize("spatial_merge_size", [1, 2])
@pytest.mark.parametrize("grid_thw", [[3, 8, 7], [128, 10, 12]])
@pytest.mark.parametrize("num_prefix_tokens", [1, 11])
@pytest.mark.parametrize("num_suffix_tokens", [0, 7])
@pytest.mark.parametrize("video_pruning_rate", [0, 0.25, 0.75])
@pytest.mark.parametrize("interleave_text_tokens", [(0, 0), (1, 4)])
def test_match_qwen3vl_mrope_evs_on(
spatial_merge_size: int,
num_prefix_tokens: int,
grid_thw: tuple[int, int, int],
num_suffix_tokens: int,
video_pruning_rate: float,
interleave_text_tokens: tuple[int, int],
):
hf_config = DummyConfig()
hf_config.vision_config.spatial_merge_size = spatial_merge_size
t, h, w = grid_thw
population = list(range(1, 100))
prefix_tokens = random.choices(population, k=num_prefix_tokens)
suffix_tokens = random.choices(population, k=num_suffix_tokens)
video_tokens, video_tokens_pruned, retention_mask = make_video_embedding(
t,
h // spatial_merge_size,
w // spatial_merge_size,
interleave_text_tokens=interleave_text_tokens,
video_pruning_rate=video_pruning_rate,
)
assert len(video_tokens) == len(retention_mask)
input_tokens = prefix_tokens + video_tokens.tolist() + suffix_tokens
input_tokens_pruned = prefix_tokens + video_tokens_pruned.tolist() + suffix_tokens
whole_sequence_retention_mask = torch.cat(
[
torch.ones(len(prefix_tokens), dtype=torch.bool),
retention_mask,
torch.ones(len(suffix_tokens), dtype=torch.bool),
],
dim=0,
)
# Build the GT mrope for unpruned input.
mm_feature = MultiModalFeatureSpec(
data=MultiModalKwargsItem(
{
"video_grid_thw": MultiModalFieldElem(
data=torch.tensor(grid_thw),
field=None, # HACK.
),
}
),
modality="video",
identifier="DUMMY",
mm_position=PlaceholderRange(offset=0, length=len(input_tokens)),
)
expected_mrope, _ = Qwen3VLForConditionalGeneration._get_mrope_input_positions(
input_tokens=input_tokens,
mm_features=[mm_feature],
config=hf_config,
)
# Compute mrope for a video-only media (unpruned).
mm_feature = MultiModalFeatureSpec(
data=MultiModalKwargsItem(
{
"video_grid_thw": MultiModalFieldElem(
data=torch.tensor(grid_thw),
field=None, # HACK.
),
}
),
modality="video",
identifier="DUMMY",
mm_position=PlaceholderRange(offset=0, length=video_tokens.numel()),
)
video_mrope, _ = Qwen3VLForConditionalGeneration._get_mrope_input_positions(
input_tokens=video_tokens.tolist(),
mm_features=[mm_feature],
config=hf_config,
)
video_mrope = video_mrope.permute(1, 0) # [N, 3]
hidden_size = 16
is_video_embed = torch.isin(
video_tokens_pruned, torch.tensor([VIDEO_TOKEN_ID], dtype=torch.long)
)
expanded_positions = torch.full(
(len(video_tokens_pruned), 5),
fill_value=-100,
device=video_mrope.device,
dtype=torch.long,
)
expanded_positions[is_video_embed, :3] = video_mrope[retention_mask][is_video_embed]
expanded_positions[~is_video_embed, :3] = video_mrope[retention_mask][
~is_video_embed
]
is_vision_start = video_tokens_pruned == VISION_START_TOKEN_ID
expanded_positions[..., 3] = is_vision_start
expanded_positions[..., 4] = is_video_embed
# Check that all positions were filled, since we initialized them as negative.
assert (expanded_positions >= 0).all()
video_embeddings = torch.empty(
(len(video_tokens_pruned), hidden_size), device=video_mrope.device
)
video_embeddings = torch.cat(
[
video_embeddings,
expanded_positions.float(),
],
dim=1,
)
multimodal_embeddings = [video_embeddings]
expected_mrope_masked = expected_mrope[:, whole_sequence_retention_mask]
# Initialize computed_mrope with sequential positions for all prefix tokens
computed_mrope = torch.empty((3, len(input_tokens_pruned)), dtype=torch.long)
computed_mrope[:, 0 : len(prefix_tokens)] = expected_mrope[
:, 0 : len(prefix_tokens)
]
# Paranoia check that computed_mrope is wrong.
assert not torch.equal(computed_mrope, expected_mrope_masked)
_, actual_mrope, _ = Qwen3VLForConditionalGeneration._recompute_mrope_positions(
input_ids=input_tokens_pruned,
multimodal_embeddings=multimodal_embeddings,
mrope_positions=computed_mrope,
num_computed_tokens=len(prefix_tokens),
vision_start_token_id=hf_config.vision_start_token_id,
image_token_id=hf_config.image_token_id,
video_token_id=hf_config.video_token_id,
)
assert torch.equal(actual_mrope, expected_mrope_masked)
@@ -7,9 +7,9 @@ from tests.models.utils import GenerateModelInfo
from .ppl_utils import wikitext_ppl_test
MODELS = [
GenerateModelInfo("google/gemma-2b"),
GenerateModelInfo("google/gemma-2-2b"),
GenerateModelInfo("google/gemma-3-4b-it"),
GenerateModelInfo("google/gemma-2b", hf_ppl=21.48524284362793),
GenerateModelInfo("google/gemma-2-2b", hf_ppl=102.59290313720703),
GenerateModelInfo("google/gemma-3-4b-it", hf_ppl=27.79648208618164),
]
@@ -6,7 +6,7 @@ from tests.models.utils import GenerateModelInfo
from .ppl_utils import wikitext_ppl_test
MODELS = [GenerateModelInfo("openai-community/gpt2-large")]
MODELS = [GenerateModelInfo("openai-community/gpt2-large", hf_ppl=19.457056045532227)]
@pytest.mark.parametrize("model_info", MODELS)
@@ -8,14 +8,20 @@ from tests.models.utils import GenerateModelInfo
from .ppl_utils import wikitext_ppl_test
MODELS = [
GenerateModelInfo("Qwen/Qwen3-0.6B"),
GenerateModelInfo("Qwen/Qwen3-0.6B-FP8"),
# transformers:
# Loading a GPTQ quantized model requires optimum, gptqmodel
# GenerateModelInfo("Qwen/Qwen3-0.6B-GPTQ-Int8"),
# for Qwen3
GenerateModelInfo("Qwen/Qwen3-0.6B", hf_ppl=23.864173889160156),
GenerateModelInfo("Qwen/Qwen3-0.6B-FP8", hf_ppl=24.313045501708984),
# for Qwen3.5
GenerateModelInfo("Qwen/Qwen3.5-0.8B", hf_ppl=19.38858413696289),
]
@pytest.mark.parametrize("model_info", MODELS)
def test_ppl(hf_runner, vllm_runner, model_info: GenerateModelInfo):
wikitext_ppl_test(hf_runner, vllm_runner, model_info)
vllm_extra_kwargs = {}
if model_info.name == "Qwen/Qwen3.5-0.8B":
vllm_extra_kwargs["language_model_only"] = True
wikitext_ppl_test(
hf_runner, vllm_runner, model_info, vllm_extra_kwargs=vllm_extra_kwargs
)
@@ -0,0 +1,355 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Tests for the LlamaNemotronVL model family:
- nvidia/llama-nemotron-embed-vl-1b-v2 (LlamaNemotronVLForCausalLM / embed)
- nvidia/llama-nemotron-rerank-vl-1b-v2
(LlamaNemotronVLForSequenceClassification / rerank)
Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone.
"""
import base64
from io import BytesIO
from pathlib import Path
import pytest
import torch
from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor
from vllm.entrypoints.chat_utils import (
ChatCompletionContentPartImageParam,
ChatCompletionContentPartTextParam,
)
from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
from ...utils import check_embeddings_close
# Prefixes used by the model API
QUERY_PREFIX = "query: "
PASSAGE_PREFIX = "passage: "
# Text prompts for text-only embedding
HF_TEXT_PROMPTS = [
# T -> X (text embedding queries)
f"{QUERY_PREFIX}The label of the object is stop sign",
f"{QUERY_PREFIX}cherry blossom",
]
# Image prompts using the model's expected format
HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
{
# I -> X (image embedding as passage/document)
"stop_sign": f"{PASSAGE_PREFIX}<image>",
"cherry_blossom": f"{PASSAGE_PREFIX}<image>",
}
)
MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"]
def _run_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
input_texts: list[str],
input_images: PromptImageInput,
model: str,
*,
dtype: str,
) -> None:
"""Run embedding comparison test between HF and vLLM.
NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
"""
# Run vLLM inference first
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=2048,
enforce_eager=True,
trust_remote_code=True,
) as vllm_model:
vllm_outputs = vllm_model.embed(input_texts, images=input_images)
# Run HF inference using the model's encode_queries/encode_documents API
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
hf_outputs = []
for text, image in zip(input_texts, input_images):
with torch.inference_mode():
if text.startswith(QUERY_PREFIX):
# Strip prefix and use encode_queries for query texts
query_text = text[len(QUERY_PREFIX) :]
embedding = hf_model.model.encode_queries([query_text])
elif text.startswith(PASSAGE_PREFIX):
# Strip prefix and use encode_documents for passages/images
passage_text = text[len(PASSAGE_PREFIX) :]
if image is not None:
# Image document - pass image to encode_documents
embedding = hf_model.model.encode_documents(
images=[image],
texts=[passage_text],
)
else:
# Text-only document
embedding = hf_model.model.encode_documents(
texts=[passage_text]
)
else:
raise ValueError(
f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'"
)
hf_outputs.append(embedding[0].tolist())
check_embeddings_close(
embeddings_0_lst=hf_outputs,
embeddings_1_lst=vllm_outputs,
name_0="hf",
name_1="vllm",
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_models_text(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
"""Test text-only embedding."""
input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images, # type: ignore
model,
dtype=dtype,
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_models_image(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
"""Test image embedding."""
input_texts_images = [
(text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images,
model,
dtype=dtype,
)
# ---------------------------------------------------------------------------
# Reranker tests — nvidia/llama-nemotron-rerank-vl-1b-v2
# ---------------------------------------------------------------------------
RERANKER_MODELS = ["nvidia/llama-nemotron-rerank-vl-1b-v2"]
# The tokenizer's built-in chat template is not suitable for the Score/Rerank
# APIs (it's inherited from the base LLM). We must use the provided override.
_RERANKER_SCORE_TEMPLATE = (
Path(__file__).parents[4]
/ "examples/pooling/score/template/nemotron-vl-rerank.jinja"
).read_text()
RERANKER_TEXT_QUERY = "How is AI improving the intelligence and capabilities of robots?"
RERANKER_TEXT_DOCS = [
"AI enables robots to perceive, plan, and act autonomously.",
(
"A biological foundation model designed to analyze DNA, RNA, "
"and protein sequences."
),
]
RERANKER_IMAGE_QUERY = "photo of a red stop sign on a street"
def _pil_to_data_uri(image) -> str:
buf = BytesIO()
image.save(buf, format="PNG")
b64 = base64.b64encode(buf.getvalue()).decode()
return f"data:image/png;base64,{b64}"
def _run_hf_reranker(
hf_runner: type[HfRunner],
model: str,
dtype: str,
query: str,
docs: list,
) -> list[float]:
"""Run HF reranker inference; docs is a list of (doc_text, doc_image|None)."""
with hf_runner(
model,
dtype=dtype,
trust_remote_code=True,
auto_cls=AutoModelForSequenceClassification,
) as hf_model:
processor = AutoProcessor.from_pretrained(
model,
trust_remote_code=True,
max_input_tiles=6,
use_thumbnail=True,
rerank_max_length=2048,
)
examples = [
{
"question": query,
"doc_text": doc_text if doc_text is not None else "",
"doc_image": doc_image if doc_image is not None else "",
}
for doc_text, doc_image in docs
]
batch_dict = processor.process_queries_documents_crossencoder(examples)
batch_dict = {
k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v
for k, v in batch_dict.items()
}
with torch.inference_mode():
logits = hf_model.model(**batch_dict, return_dict=True).logits
# vLLM applies sigmoid activation to the raw logits before returning
# scores; apply the same here so both sides are comparable.
scores = torch.sigmoid(logits.squeeze(-1).float())
return scores.detach().cpu().tolist()
def _run_vllm_reranker(
vllm_runner: type[VllmRunner],
model: str,
dtype: str,
query: str,
docs: list,
) -> list[float]:
"""Run vLLM reranker inference; docs is a list of (doc_text, doc_image|None)."""
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=2048,
enforce_eager=True,
trust_remote_code=True,
) as vllm_model:
has_images = any(img is not None for _, img in docs)
if not has_images:
# Text-only path: use the simple string score API.
queries = [query] * len(docs)
doc_texts = [doc_text for doc_text, _ in docs]
outputs = vllm_model.score(
queries,
doc_texts,
chat_template=_RERANKER_SCORE_TEMPLATE,
)
else:
# Multimodal path: build ScoreMultiModalParam for each pair.
query_params = [
ScoreMultiModalParam(
content=[
ChatCompletionContentPartTextParam(
type="text",
text=query,
)
]
)
] * len(docs)
doc_params = []
for doc_text, doc_image in docs:
content: list = []
if doc_image is not None:
content.append(
ChatCompletionContentPartImageParam(
type="image_url",
image_url={"url": _pil_to_data_uri(doc_image)},
)
)
if doc_text:
content.append(
ChatCompletionContentPartTextParam(
type="text",
text=doc_text,
)
)
doc_params.append(ScoreMultiModalParam(content=content))
raw_outputs = vllm_model.llm.score(
query_params,
doc_params,
chat_template=_RERANKER_SCORE_TEMPLATE,
)
outputs = [o.outputs.score for o in raw_outputs]
return outputs
def _run_reranker_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
model: str,
dtype: str,
query: str,
docs: list,
) -> None:
"""Compare HF and vLLM reranker scores.
NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
"""
vllm_scores = _run_vllm_reranker(vllm_runner, model, dtype, query, docs)
hf_scores = _run_hf_reranker(hf_runner, model, dtype, query, docs)
assert len(hf_scores) == len(vllm_scores), (
f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}"
)
for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)):
assert hf_score == pytest.approx(vllm_score, rel=0.02), (
f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}"
)
@pytest.mark.parametrize("model", RERANKER_MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_reranker_text(
hf_runner,
vllm_runner,
model: str,
dtype: str,
) -> None:
"""Test reranking with text-only query and text documents."""
docs = [(text, None) for text in RERANKER_TEXT_DOCS]
_run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_TEXT_QUERY, docs)
@pytest.mark.parametrize("model", RERANKER_MODELS)
@pytest.mark.parametrize("dtype", ["half"])
def test_reranker_image_doc(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
"""Test reranking with text query against image documents."""
docs = [(None, asset.pil_image) for asset in image_assets]
_run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_IMAGE_QUERY, docs)
@@ -1,148 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Tests for LlamaNemotronVL embedding model (nvidia/llama-nemotron-embed-vl-1b-v2).
This model uses SigLIP vision encoder with bidirectional LLaMA for embeddings.
"""
import pytest
import torch
from transformers import AutoModel
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
from ...utils import check_embeddings_close
# Prefixes used by the model API
QUERY_PREFIX = "query: "
PASSAGE_PREFIX = "passage: "
# Text prompts for text-only embedding
HF_TEXT_PROMPTS = [
# T -> X (text embedding queries)
f"{QUERY_PREFIX}The label of the object is stop sign",
f"{QUERY_PREFIX}cherry blossom",
]
# Image prompts using the model's expected format
HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
{
# I -> X (image embedding as passage/document)
"stop_sign": f"{PASSAGE_PREFIX}<image>",
"cherry_blossom": f"{PASSAGE_PREFIX}<image>",
}
)
MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"]
def _run_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
input_texts: list[str],
input_images: PromptImageInput,
model: str,
*,
dtype: str,
) -> None:
"""Run embedding comparison test between HF and vLLM.
NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
"""
# Run vLLM inference first
with vllm_runner(
model,
runner="pooling",
dtype=dtype,
max_model_len=2048,
enforce_eager=True,
trust_remote_code=True,
) as vllm_model:
vllm_outputs = vllm_model.embed(input_texts, images=input_images)
# Run HF inference using the model's encode_queries/encode_documents API
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
hf_outputs = []
for text, image in zip(input_texts, input_images):
with torch.inference_mode():
if text.startswith(QUERY_PREFIX):
# Strip prefix and use encode_queries for query texts
query_text = text[len(QUERY_PREFIX) :]
embedding = hf_model.model.encode_queries([query_text])
elif text.startswith(PASSAGE_PREFIX):
# Strip prefix and use encode_documents for passages/images
passage_text = text[len(PASSAGE_PREFIX) :]
if image is not None:
# Image document - pass image to encode_documents
embedding = hf_model.model.encode_documents(
images=[image],
texts=[passage_text],
)
else:
# Text-only document
embedding = hf_model.model.encode_documents(
texts=[passage_text]
)
else:
raise ValueError(
f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'"
)
hf_outputs.append(embedding[0].tolist())
check_embeddings_close(
embeddings_0_lst=hf_outputs,
embeddings_1_lst=vllm_outputs,
name_0="hf",
name_1="vllm",
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["bfloat16"])
def test_models_text(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
"""Test text-only embedding."""
input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images, # type: ignore
model,
dtype=dtype,
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("dtype", ["bfloat16"])
def test_models_image(
hf_runner,
vllm_runner,
image_assets,
model: str,
dtype: str,
) -> None:
"""Test image embedding."""
input_texts_images = [
(text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
]
input_texts = [text for text, _ in input_texts_images]
input_images = [image for _, image in input_texts_images]
_run_test(
hf_runner,
vllm_runner,
input_texts,
input_images,
model,
dtype=dtype,
)
@@ -33,32 +33,9 @@ from ...registry import (
)
def glm4_1v_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict:
def add_video_metadata(mm_data: MultiModalDataDict) -> MultiModalDataDict:
"""
Patch the multimodal data for GLM4.1V model.
"""
# Ensure video metadata is included
if "video" in mm_data:
# GLM4.1V doesn't support multiple videos
video = mm_data["video"]
num_frames = len(video)
mm_data["video"] = (
video,
{
"total_num_frames": num_frames,
"fps": num_frames,
"duration": 1,
"frames_indices": [i for i in range(num_frames)],
"video_backend": "opencv",
"do_sample_frames": True,
},
)
return mm_data
def qwen3_vl_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict:
"""
Patch the multimodal data for Qwen3-VL model.
Add metadata to video mm_data
"""
def create_metadata(frames: np.ndarray):
@@ -119,16 +96,7 @@ _IGNORE_MM_KEYS = {
}
MM_DATA_PATCHES = {
# Ernie4.5-VL, GLM4.1V and Qwen3-VL requires video metadata
"ernie4_5_moe_vl": qwen3_vl_patch_mm_data,
"glm4v": glm4_1v_patch_mm_data,
"glm4v_moe": glm4_1v_patch_mm_data,
"glm_ocr": glm4_1v_patch_mm_data,
"glmasr": glmasr_patch_mm_data,
"interns1_pro": qwen3_vl_patch_mm_data,
"molmo2": qwen3_vl_patch_mm_data,
"qwen3_vl": qwen3_vl_patch_mm_data,
"qwen3_vl_moe": qwen3_vl_patch_mm_data,
}
@@ -174,6 +142,9 @@ def get_text_token_prompts(
tokenizer: TokenizerLike = processor.info.get_tokenizer()
model_config = processor.info.ctx.model_config
if processor.info.data_parser.video_needs_metadata:
mm_data = add_video_metadata(mm_data)
model_type = model_config.hf_config.model_type
if model_type in MM_DATA_PATCHES:
mm_data = MM_DATA_PATCHES[model_type](mm_data)
+10 -5
View File
@@ -12,8 +12,8 @@ Config:
Run: pytest tests/models/quantization/test_gpt_oss.py
"""
import importlib
import importlib.metadata
import importlib.util
from dataclasses import dataclass
import huggingface_hub
@@ -21,6 +21,8 @@ import lm_eval
import pytest
from packaging import version
from vllm.utils.torch_utils import cuda_device_count_stateless
MODEL_ACCURACIES = {
# Full quantization: attention linears and MoE linears
"amd/gpt-oss-20b-WFP8-AFP8-KVFP8": 0.89,
@@ -83,6 +85,9 @@ class EvaluationConfig:
def test_gpt_oss_attention_quantization(
model_name: str, tp_size: int, expected_accuracy: float
):
if tp_size > cuda_device_count_stateless():
pytest.skip("Not enough GPUs to run this test case")
model_args = EvaluationConfig(model_name).get_model_args(tp_size)
extra_run_kwargs = {
@@ -104,7 +109,7 @@ def test_gpt_oss_attention_quantization(
)
rtol = 0.02
assert (
measured_accuracy - rtol < expected_accuracy
and measured_accuracy + rtol > expected_accuracy
), f"Expected: {expected_accuracy} | Measured: {measured_accuracy}"
assert measured_accuracy >= expected_accuracy - rtol, (
f"Accuracy {measured_accuracy:.4f} is below threshold "
f"{expected_accuracy - rtol:.4f} (expected >= {expected_accuracy} - {rtol})"
)
+6
View File
@@ -653,6 +653,9 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
"LlamaBidirectionalForSequenceClassification": _HfExamplesInfo(
"nvidia/llama-nemotron-rerank-1b-v2", trust_remote_code=True
),
"LlamaNemotronVLForSequenceClassification": _HfExamplesInfo(
"nvidia/llama-nemotron-rerank-vl-1b-v2", trust_remote_code=True
),
"ModernBertForSequenceClassification": _HfExamplesInfo(
"Alibaba-NLP/gte-reranker-modernbert-base"
),
@@ -740,6 +743,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"baidu/ERNIE-4.5-VL-28B-A3B-PT",
trust_remote_code=True,
),
"FireRedASR2ForConditionalGeneration": _HfExamplesInfo(
"allendou/FireRedASR2-LLM-vllm",
),
"FunASRForConditionalGeneration": _HfExamplesInfo(
"allendou/Fun-ASR-Nano-2512-vllm",
),
+80
View File
@@ -127,6 +127,14 @@ def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch):
)
def test_deepseek_fp8_block_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch):
can_initialize(
"deepseek-ai/DeepSeek-V3.1",
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=triton"],
)
@pytest.mark.skip(
reason=(
"Known issue: lack of kernel support. "
@@ -149,6 +157,14 @@ def test_deepseek_fp8_block_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatc
)
def test_deepseek_nvfp4_moe_flashinfer_vllm(monkeypatch: pytest.MonkeyPatch):
can_initialize(
"nvidia/DeepSeek-R1-0528-FP4-v2",
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=cutlass"],
)
def test_deepseek_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
can_initialize(
"nvidia/DeepSeek-R1-0528-FP4-v2",
@@ -200,3 +216,67 @@ def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=flashinfer_trtllm"],
)
## NemoTron ##
def test_nemotron_fp8_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch):
can_initialize(
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=flashinfer_cutlass"],
)
@pytest.mark.skip(
reason=(
"FP8 MoE backend FLASHINFER_TRTLLM does not support the "
"deployment configuration since kernel does not support "
"no act_and_mul MLP layer."
)
)
def test_nemotron_fp8_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch):
can_initialize(
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=flashinfer_trtllm"],
)
@pytest.mark.skip(
reason=(
"FP8 MoE backend TRITON does not support the "
"deployment configuration since kernel does not support "
"no act_and_mul MLP layer."
)
)
def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch):
can_initialize(
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=triton"],
)
def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch):
can_initialize(
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=flashinfer_cutlass"],
)
@pytest.mark.skip(
reason=(
"FP4 MoE backend FLASHINFER_TRTLLM does not support the "
"deployment configuration since kernel does not support "
"hidden_dim % 512 != 0."
)
)
def test_nemotron_fp4_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch):
can_initialize(
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=flashinfer_trtllm"],
)
+6 -5
View File
@@ -7,7 +7,8 @@ from transformers.processing_utils import ProcessingKwargs
from typing_extensions import Unpack
from vllm.transformers_utils.processor import (
get_processor_kwargs_from_processor,
get_processor_kwargs_keys,
get_processor_kwargs_type,
)
@@ -35,7 +36,7 @@ def _assert_has_all_expected(keys: set[str]) -> None:
assert k in keys
# Path 1: __call__ method has kwargs: Unpack[*ProcessingKwargs]
# Path 1: __call__ method has kwargs: Unpack[*ProcessorKwargs]
class _ProcWithUnpack:
def __call__(self, *args, **kwargs: Unpack[_FakeProcessorKwargs]): # type: ignore
return None
@@ -43,11 +44,11 @@ class _ProcWithUnpack:
def test_get_processor_kwargs_from_processor_unpack_path_returns_full_union():
proc = _ProcWithUnpack()
keys = get_processor_kwargs_from_processor(proc)
keys = get_processor_kwargs_keys(get_processor_kwargs_type(proc))
_assert_has_all_expected(keys)
# ---- Path 2: No Unpack, fallback to scanning *ProcessingKwargs in module ----
# ---- Path 2: No Unpack, fallback to scanning *ProcessorKwargs in module ----
class _ProcWithoutUnpack:
@@ -62,5 +63,5 @@ def test_get_processor_kwargs_from_processor_module_scan_returns_full_union():
assert hasattr(mod, "_FakeProcessorKwargs")
proc = _ProcWithoutUnpack()
keys = get_processor_kwargs_from_processor(proc)
keys = get_processor_kwargs_keys(get_processor_kwargs_type(proc))
_assert_has_all_expected(keys)
@@ -0,0 +1,360 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Integration tests for TRTLLM gen-full attention through FlashInfer."""
import unittest.mock
from functools import partial
import pytest
import torch
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
from tests.v1.attention.utils import (
BatchSpec,
create_common_attn_metadata,
create_vllm_config,
)
from vllm.config import set_current_vllm_config
from vllm.platforms import current_platform
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.attention.backends.utils import (
PerLayerParameters,
get_kv_cache_layout,
set_kv_cache_layout,
)
from vllm.v1.kv_cache_interface import FullAttentionSpec
if not current_platform.is_device_capability_family(100):
pytest.skip(
"TRTLLM integration tests require NVIDIA Blackwell (SM100).",
allow_module_level=True,
)
from vllm.v1.attention.backends.flashinfer import ( # noqa: E402
FlashInferImpl,
FlashInferMetadataBuilder,
TRTLLMDecode,
TRTLLMPrefill,
)
class MockAttentionLayer:
"""Minimal mock of an attention layer for testing."""
def __init__(self, device: torch.device):
self._q_scale = torch.tensor(1.0, device=device)
self._k_scale = torch.tensor(1.0, device=device)
self._v_scale = torch.tensor(1.0, device=device)
self._q_scale_float = 1.0
self._k_scale_float = 1.0
self._v_scale_float = 1.0
self._o_scale_float = None
MODEL = "Qwen/Qwen2.5-0.5B"
BLOCK_SIZE = 16
NUM_GPU_BLOCKS = 8192
BATCH_SPECS = {
"decode_only": BatchSpec(
seq_lens=[128, 256, 512],
query_lens=[1, 1, 1],
),
"prefill_only": BatchSpec(
seq_lens=[64, 128, 256],
query_lens=[16, 32, 16],
),
"mixed": BatchSpec(
seq_lens=[128, 256, 512, 128],
query_lens=[1, 1, 8, 16],
),
}
def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
head_size = vllm_config.model_config.get_head_size()
return {
name: PerLayerParameters(
window_left=-1,
logits_soft_cap=0.0,
sm_scale=1.0 / (head_size**0.5),
)
for name in layer_names
}
def _create_hnd_kv_cache(
k_contexts,
v_contexts,
block_size,
num_kv_heads,
head_size,
dtype,
device,
num_blocks,
common_attn_metadata,
):
"""Create and populate a KV cache with HND-compatible strides.
The returned tensor has logical shape
(num_blocks, 2, block_size, num_kv_heads, head_size) but is physically
laid out as (num_blocks, 2, num_kv_heads, block_size, head_size) so that
``kv_cache.permute(0, 1, 3, 2, 4)`` yields a contiguous HND view.
"""
seq_lens = common_attn_metadata.seq_lens.cpu()
query_lens = (
common_attn_metadata.query_start_loc_cpu[1:]
- common_attn_metadata.query_start_loc_cpu[:-1]
)
block_table = common_attn_metadata.block_table_tensor
slot_mapping = common_attn_metadata.slot_mapping
batch_size = len(k_contexts)
# Build cache in (2, num_blocks, block_size, num_kv_heads, head_size)
# then convert to HND format (same approach as test_attention_backends.py).
kv_cache_raw = torch.zeros(
2,
num_blocks,
block_size,
num_kv_heads,
head_size,
dtype=dtype,
device=device,
)
kv_cache_flat = kv_cache_raw.view(2, -1, num_kv_heads, head_size)
start_block_idx = 1
for i in range(batch_size):
k_ctx, v_ctx = k_contexts[i], v_contexts[i]
start = start_block_idx * block_size
end = start + k_ctx.shape[0]
kv_cache_flat[0, start:end] = k_ctx
kv_cache_flat[1, start:end] = v_ctx
start_block_idx += cdiv(int(seq_lens[i]), block_size)
blocks_end = start_block_idx
# Randomly permute blocks (starting from block 1; block 0 is null).
perm = torch.randperm(blocks_end - 1) + 1
inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
inv_perm[1:] = torch.argsort(perm) + 1
kv_cache_raw[:, 1:blocks_end] = kv_cache_raw[:, perm]
# Build block table.
start_block_idx = 1
for i in range(batch_size):
n_blocks = cdiv(int(seq_lens[i]), block_size)
block_table[i, :n_blocks] = inv_perm[
start_block_idx : start_block_idx + n_blocks
]
start_block_idx += n_blocks
# Build slot mapping that is consistent with the block table.
for i in range(batch_size):
ctx_len = int(seq_lens[i]) - int(query_lens[i])
token_offsets = torch.arange(int(query_lens[i])) + ctx_len
block_indices = token_offsets // block_size
intra_block_offsets = token_offsets % block_size
start = common_attn_metadata.query_start_loc_cpu[i]
end = common_attn_metadata.query_start_loc_cpu[i + 1]
slot_mapping[start:end] = block_table[
i, block_indices
] * block_size + intra_block_offsets.to(device)
# Transpose to FlashInfer logical shape then make HND-strided.
kv_cache = kv_cache_raw.transpose(0, 1)
kv_cache = kv_cache.transpose(2, 3).contiguous().transpose(2, 3)
return kv_cache
def _run_trtllm_integration(batch_spec):
"""Run TRTLLM attention through the full FlashInfer pipeline
and compare against an SDPA reference."""
set_random_seed(42)
device = torch.device("cuda:0")
vllm_config = create_vllm_config(
model_name=MODEL,
max_model_len=max(batch_spec.seq_lens),
block_size=BLOCK_SIZE,
num_gpu_blocks=NUM_GPU_BLOCKS,
)
vllm_config.attention_config.use_trtllm_attention = True
num_q_heads = vllm_config.model_config.get_num_attention_heads(
vllm_config.parallel_config
)
num_kv_heads = vllm_config.model_config.get_num_kv_heads(
vllm_config.parallel_config
)
head_size = vllm_config.model_config.get_head_size()
dtype = vllm_config.model_config.dtype
scale = 1.0 / (head_size**0.5)
# 1. Generate data and compute SDPA reference
all_q, all_k, all_v = [], [], []
all_sdpa_out = []
k_contexts, v_contexts = [], []
for i in range(batch_spec.batch_size):
s_len = batch_spec.seq_lens[i]
q_len = batch_spec.query_lens[i]
ctx_len = s_len - q_len
q = torch.randn(q_len, num_q_heads, head_size, dtype=dtype, device=device)
k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
# SDPA reference (N=1, H, L, D)
q_sdpa = q.unsqueeze(0).transpose(1, 2)
k_sdpa = k_full.unsqueeze(0).transpose(1, 2)
v_sdpa = v_full.unsqueeze(0).transpose(1, 2)
if num_q_heads != num_kv_heads:
repeats = num_q_heads // num_kv_heads
k_sdpa = k_sdpa.repeat_interleave(repeats, dim=1)
v_sdpa = v_sdpa.repeat_interleave(repeats, dim=1)
def causal_mask_mod(b, h, q_idx, kv_idx, *, context_len):
return (q_idx + context_len) >= kv_idx
mask_fn = partial(causal_mask_mod, context_len=ctx_len)
block_mask = create_block_mask(
mask_fn, B=None, H=None, Q_LEN=q_len, KV_LEN=s_len, device=device
)
sdpa_out = flex_attention(
q_sdpa,
k_sdpa,
v_sdpa,
block_mask=block_mask,
scale=scale,
enable_gqa=True,
)
all_sdpa_out.append(sdpa_out.transpose(1, 2).squeeze(0))
all_q.append(q)
all_k.append(k_full[ctx_len:])
all_v.append(v_full[ctx_len:])
k_contexts.append(k_full[:ctx_len])
v_contexts.append(v_full[:ctx_len])
query_vllm = torch.cat(all_q, dim=0)
key_vllm = torch.cat(all_k, dim=0)
value_vllm = torch.cat(all_v, dim=0)
sdpa_output = torch.cat(all_sdpa_out, dim=0)
common_attn_metadata = create_common_attn_metadata(batch_spec, BLOCK_SIZE, device)
# 2. Create HND KV cache
kv_cache = _create_hnd_kv_cache(
k_contexts,
v_contexts,
BLOCK_SIZE,
num_kv_heads,
head_size,
dtype,
device,
NUM_GPU_BLOCKS,
common_attn_metadata,
)
# 3. Run through FlashInfer with TRTLLM enabled
set_kv_cache_layout("HND")
get_kv_cache_layout.cache_clear()
try:
kv_cache_spec = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=num_kv_heads,
head_size=head_size,
dtype=dtype,
)
layer_names = ["test_layer_0"]
with (
set_current_vllm_config(vllm_config),
unittest.mock.patch(
"vllm.utils.flashinfer.supports_trtllm_attention",
return_value=True,
),
unittest.mock.patch(
"vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
_mock_get_per_layer_parameters,
),
):
builder = FlashInferMetadataBuilder(
kv_cache_spec, layer_names, vllm_config, device
)
attn_metadata = builder.build(
common_prefix_len=0,
common_attn_metadata=common_attn_metadata,
)
# Verify the correct TRTLLM metadata types were produced.
has_prefills = any(ql > 1 for ql in batch_spec.query_lens)
has_decodes = any(ql == 1 for ql in batch_spec.query_lens)
if has_prefills:
assert isinstance(attn_metadata.prefill, TRTLLMPrefill), (
f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}"
)
if has_decodes:
assert isinstance(attn_metadata.decode, TRTLLMDecode), (
f"Expected TRTLLMDecode, got {type(attn_metadata.decode)}"
)
impl = FlashInferImpl(
num_heads=num_q_heads,
head_size=head_size,
scale=scale,
num_kv_heads=num_kv_heads,
alibi_slopes=None,
sliding_window=None,
kv_cache_dtype="auto",
)
mock_layer = MockAttentionLayer(device)
output = torch.empty_like(query_vllm)
impl.do_kv_cache_update(
mock_layer,
key_vllm,
value_vllm,
kv_cache,
attn_metadata.slot_mapping,
)
output = impl.forward(
mock_layer,
query_vllm,
key_vllm,
value_vllm,
kv_cache,
attn_metadata,
output=output,
)
# 4. Compare against SDPA reference
torch.testing.assert_close(
output,
sdpa_output,
atol=1e-2,
rtol=1e-2,
)
finally:
set_kv_cache_layout(None)
get_kv_cache_layout.cache_clear()
@pytest.mark.parametrize(
"batch_spec_name",
list(BATCH_SPECS.keys()),
)
@torch.inference_mode()
def test_trtllm_gen_full_attention_integration(batch_spec_name: str):
"""Test TRTLLM gen-full attention through the full FlashInfer
MetadataBuilder.build() -> FlashInferImpl.forward() pipeline,
with real TRTLLM kernels on Blackwell."""
_run_trtllm_integration(BATCH_SPECS[batch_spec_name])
+290
View File
@@ -0,0 +1,290 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.sampling_params import RepetitionDetectionParams, SamplingParams
from vllm.v1.core.sched.utils import check_sequence_repetition, check_stop
from vllm.v1.request import Request, RequestStatus
pytestmark = pytest.mark.cpu_test
# ============================================================================
# UNIT TESTS - check_sequence_repetition function
# ============================================================================
class TestCheckSequenceRepetition:
"""Unit tests for the check_sequence_repetition function"""
def test_simple_repetition_detected(self):
"""Test detection of simple repetitive patterns"""
token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 3]
params = RepetitionDetectionParams(
max_pattern_size=3,
min_pattern_size=2,
min_count=3,
)
assert check_sequence_repetition(token_ids, params)
def test_repetition_below_min_count(self):
"""Test that pattern below min_count is not detected"""
token_ids = [1, 2, 3, 1, 2, 3]
params = RepetitionDetectionParams(
max_pattern_size=3,
min_pattern_size=2,
min_count=3,
)
assert not check_sequence_repetition(token_ids, params)
def test_two_token_pattern(self):
"""Test detection of 2-token patterns"""
token_ids = [1, 2, 1, 2, 1, 2, 1, 2]
params = RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=2,
min_count=4,
)
assert check_sequence_repetition(token_ids, params)
def test_no_repetition_varied_sequence(self):
"""Test that non-repetitive sequences are not flagged"""
token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9]
params = RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=2,
min_count=2,
)
assert not check_sequence_repetition(token_ids, params)
def test_partial_repetition_not_detected(self):
"""Test that incomplete repetitions are not detected"""
token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 4]
params = RepetitionDetectionParams(
max_pattern_size=3,
min_pattern_size=2,
min_count=3,
)
assert not check_sequence_repetition(token_ids, params)
def test_empty_token_list(self):
"""Test with empty token list"""
params = RepetitionDetectionParams(
max_pattern_size=3,
min_pattern_size=2,
min_count=2,
)
assert not check_sequence_repetition([], params)
def test_detection_disabled_max_size_zero(self):
"""Test that zero max_pattern_size disables detection"""
token_ids = [1, 2, 1, 2, 1, 2]
params = RepetitionDetectionParams()
assert not check_sequence_repetition(token_ids, params)
def test_invalid_min_count(self):
"""Test that min_count < 2 returns False"""
token_ids = [1, 2, 1, 2]
params = RepetitionDetectionParams()
assert not check_sequence_repetition(token_ids, params)
def test_repetition_at_end_of_sequence(self):
"""Test detection when repetition occurs at the end"""
token_ids = [1, 2, 3, 4, 5, 6, 5, 6, 5, 6]
params = RepetitionDetectionParams(
max_pattern_size=3,
min_pattern_size=2,
min_count=3,
)
assert check_sequence_repetition(token_ids, params)
def test_large_pattern_many_repetitions(self):
"""Test large pattern repeated many times"""
token_ids = [1, 2, 3, 4, 5, 6, 7, 8] * 5
params = RepetitionDetectionParams(
max_pattern_size=10,
min_pattern_size=2,
min_count=3,
)
assert check_sequence_repetition(token_ids, params)
# ============================================================================
# INTEGRATION TESTS - check_stop with repetition detection
# ============================================================================
class TestRepetitionDetectionIntegration:
"""Integration tests for repetition detection in check_stop"""
def test_basic_repetition_stops_generation(self):
"""Test that repetition is detected and stops generation"""
params = SamplingParams(
max_tokens=100,
repetition_detection=RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=2,
min_count=3,
),
)
request = Request(
request_id="test",
prompt_token_ids=[1, 2, 3],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
assert check_stop(request, max_model_len=1024)
assert request.status == RequestStatus.FINISHED_REPETITION
assert request.stop_reason == "repetition_detected"
def test_detection_disabled_no_stop(self):
"""Test that disabled detection doesn't stop generation"""
params = SamplingParams(
max_tokens=100,
)
request = Request(
request_id="test",
prompt_token_ids=[1, 2, 3],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
assert not check_stop(request, max_model_len=1024)
def test_repetition_respects_min_tokens(self):
"""Test that repetition detection respects min_tokens"""
params = SamplingParams(
min_tokens=10,
max_tokens=100,
repetition_detection=RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=2,
min_count=3,
),
)
request = Request(
request_id="test",
prompt_token_ids=[1, 2, 3],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
assert not check_stop(request, max_model_len=1024)
def test_no_repetition_continues_generation(self):
"""Test that non-repetitive tokens don't stop generation"""
params = SamplingParams(
max_tokens=100,
repetition_detection=RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=2,
min_count=3,
),
)
request = Request(
request_id="test",
prompt_token_ids=[1, 2, 3],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([10, 20, 30, 40, 50, 60])
assert not check_stop(request, max_model_len=1024)
def test_pattern_at_size_boundary(self):
"""Test detection at exact pattern size boundary"""
params = SamplingParams(
max_tokens=100,
repetition_detection=RepetitionDetectionParams(
max_pattern_size=3,
min_pattern_size=3,
min_count=2,
),
)
request = Request(
request_id="test",
prompt_token_ids=[1, 2],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([10, 20, 30, 10, 20, 30])
assert check_stop(request, max_model_len=1024)
assert request.status == RequestStatus.FINISHED_REPETITION
def test_multiple_pattern_sizes_checked(self):
"""Test that function checks pattern sizes in range"""
params = SamplingParams(
max_tokens=100,
repetition_detection=RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=2,
min_count=3,
),
)
request = Request(
request_id="test",
prompt_token_ids=[1],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([7, 8, 9, 10, 7, 8, 9, 10, 7, 8, 9, 10])
assert check_stop(request, max_model_len=1024)
assert request.status == RequestStatus.FINISHED_REPETITION
def test_eos_takes_precedence_over_repetition(self):
"""Test that EOS token stops before repetition check"""
params = SamplingParams(
max_tokens=100,
stop_token_ids=[999],
repetition_detection=RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=2,
min_count=3,
),
)
request = Request(
request_id="test",
prompt_token_ids=[1, 2, 3],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([10, 20, 10, 20, 999])
assert check_stop(request, max_model_len=1024)
assert request.status == RequestStatus.FINISHED_STOPPED
def test_min_pattern_size_filters_small_patterns(self):
"""Test that min_pattern_size filters out smaller patterns"""
params = SamplingParams(
max_tokens=100,
repetition_detection=RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=3,
min_count=3,
),
)
request = Request(
request_id="test",
prompt_token_ids=[1],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
assert not check_stop(request, max_model_len=1024)
def test_high_repetition_threshold(self):
"""Test that high min_count requires many repetitions"""
params = SamplingParams(
max_tokens=100,
repetition_detection=RepetitionDetectionParams(
max_pattern_size=5,
min_pattern_size=2,
min_count=5,
),
)
request = Request(
request_id="test",
prompt_token_ids=[1],
sampling_params=params,
pooling_params=None,
)
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
assert not check_stop(request, max_model_len=1024)
+6 -39
View File
@@ -40,26 +40,6 @@ from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
pytestmark = pytest.mark.cpu_test
def _get_remote_waiting_queue(scheduler: Scheduler):
return getattr(scheduler, "waiting_for_remote_kvs", None)
def _num_waiting_requests(scheduler: Scheduler) -> int:
remote_waiting = _get_remote_waiting_queue(scheduler)
return len(scheduler.waiting) + (len(remote_waiting) if remote_waiting else 0)
def _get_remote_waiting_requests(scheduler: Scheduler) -> list[Request]:
remote_waiting = _get_remote_waiting_queue(scheduler)
if remote_waiting is not None:
return list(remote_waiting)
return [
req
for req in scheduler.waiting
if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
]
def test_add_requests():
scheduler = create_scheduler()
requests = create_requests(num_requests=10)
@@ -1140,8 +1120,7 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
# Requests should first transition to WAITING_FOR_REMOTE_KVS
output = scheduler.schedule()
assert _num_waiting_requests(scheduler) == len(req_ids)
assert len(_get_remote_waiting_requests(scheduler)) == len(req_ids)
assert len(scheduler.waiting) == len(req_ids)
assert len(scheduler.running) == 0
assert len(output.scheduled_new_reqs) == 0
for req in scheduler.requests.values():
@@ -1160,8 +1139,7 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
# Simulate KV transfer completion using KVConnectorOutput.finished_recving
output = scheduler.schedule()
assert _num_waiting_requests(scheduler) == len(req_ids)
assert len(_get_remote_waiting_requests(scheduler)) == len(req_ids)
assert len(scheduler.waiting) == len(req_ids)
assert len(scheduler.running) == 0
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
@@ -1568,7 +1546,7 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
# All can be scheduled - 1st token.
output = scheduler.schedule()
if is_async:
assert _num_waiting_requests(scheduler) == 2
assert len(scheduler.waiting) == 2
assert scheduler.running == []
_step_until_kv_transfer_finished(scheduler, req_ids)
output = scheduler.schedule()
@@ -1626,9 +1604,7 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
# This will have a local and remote cache hit.
output = scheduler.schedule()
if is_async:
waiting_req_ids = [
req.request_id for req in _get_remote_waiting_requests(scheduler)
]
waiting_req_ids = [req.request_id for req in scheduler.waiting]
assert len(waiting_req_ids) == 1
_step_until_kv_transfer_finished(scheduler, waiting_req_ids)
output = scheduler.schedule()
@@ -3638,10 +3614,6 @@ def test_prepend_skipped_requests_order():
# simulate first 2 waiting requests are waiting for remote KVs
for req in expected_waiting_reqs[:2]:
req.status = RequestStatus.WAITING_FOR_REMOTE_KVS
remote_waiting = _get_remote_waiting_queue(scheduler)
if remote_waiting is not None:
scheduler.waiting.remove_request(req)
remote_waiting.add_request(req)
# schedule step
# expect the first 2 waiting to be skipped, the third running,
@@ -3651,13 +3623,8 @@ def test_prepend_skipped_requests_order():
# pop the third request which is expected to be running
expected_waiting_reqs.pop(2)
# verify waiting order is preserved for schedulable requests.
remote_waiting = _get_remote_waiting_queue(scheduler)
if remote_waiting is not None:
assert list(scheduler.waiting) == expected_waiting_reqs[2:]
assert list(remote_waiting) == expected_waiting_reqs[:2]
else:
assert list(scheduler.waiting) == expected_waiting_reqs
# verify waiting order is preserved
assert list(scheduler.waiting) == expected_waiting_reqs
def test_abort_request_waiting_for_remote_kvs():
@@ -17,26 +17,6 @@ from .utils import (
)
def _get_remote_waiting_queue(scheduler: Scheduler):
return getattr(scheduler, "waiting_for_remote_kvs", None)
def _num_waiting_requests(scheduler: Scheduler) -> int:
remote_waiting = _get_remote_waiting_queue(scheduler)
return len(scheduler.waiting) + (len(remote_waiting) if remote_waiting else 0)
def _get_remote_waiting_requests(scheduler: Scheduler) -> list[Request]:
remote_waiting = _get_remote_waiting_queue(scheduler)
if remote_waiting is not None:
return list(remote_waiting)
return [
req
for req in scheduler.waiting
if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
]
def _make_get_num_new_matched_tokens(
req_num_new_matched_tokens: dict[str, int],
async_load,
@@ -96,8 +76,8 @@ def test_async_load_failure(
scheduler_output = scheduler.schedule()
assert _num_waiting_requests(scheduler) == 3
for request in _get_remote_waiting_requests(scheduler):
assert len(scheduler.waiting) == 3
for request in scheduler.waiting:
assert request.num_computed_tokens == 0
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
@@ -116,8 +96,8 @@ def test_async_load_failure(
min_invalid_block_idx = min(invalid_block_idxs)
assert _num_waiting_requests(scheduler) == 3
for request in _get_remote_waiting_requests(scheduler):
assert len(scheduler.waiting) == 3
for request in scheduler.waiting:
if request.request_id == request2.request_id:
assert request.num_computed_tokens == (
min_invalid_block_idx * scheduler.block_size
@@ -323,10 +303,8 @@ def test_async_progressive_load_failure(
scheduler_output = scheduler.schedule()
assert _num_waiting_requests(scheduler) == 1
remote_waiting_reqs = _get_remote_waiting_requests(scheduler)
assert len(remote_waiting_reqs) == 1
assert remote_waiting_reqs[0].request_id == request.request_id
assert len(scheduler.waiting) == 1
assert scheduler.waiting.peek_request().request_id == request.request_id
assert request.num_computed_tokens == 0
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
@@ -347,10 +325,8 @@ def test_async_progressive_load_failure(
min_invalid_block_idx = min(min_invalid_block_idx, invalid_block_idx)
assert _num_waiting_requests(scheduler) == 1
remote_waiting_reqs = _get_remote_waiting_requests(scheduler)
assert len(remote_waiting_reqs) == 1
assert remote_waiting_reqs[0].request_id == request.request_id
assert len(scheduler.waiting) == 1
assert scheduler.waiting.peek_request().request_id == request.request_id
assert request.num_computed_tokens == (
min_invalid_block_idx * scheduler.block_size
)
+24 -29
View File
@@ -476,12 +476,12 @@ def test_set_inputs_first_pass_draft_model():
proposer.max_num_tokens, dtype=torch.bool, device=device
)
# Mock the attn_metadata_builder to avoid needing the full model setup
# Mock draft_attn_groups to avoid needing the full model setup
mock_kv_cache_spec = mock.MagicMock()
mock_kv_cache_spec.block_size = block_size
mock_builder = mock.MagicMock()
mock_builder.kv_cache_spec = mock_kv_cache_spec
proposer.attn_metadata_builder = mock_builder
mock_attn_group = mock.MagicMock()
mock_attn_group.kv_cache_spec = mock_kv_cache_spec
proposer.draft_attn_groups = [mock_attn_group]
# Request 0: query_len=3 (but 1 rejected), Request 1: query_len=2
batch_spec = BatchSpec(
@@ -616,12 +616,12 @@ def test_set_inputs_first_pass_parallel_drafting():
proposer.max_num_tokens, dtype=torch.bool, device=device
)
# Mock the attn_metadata_builder
# Mock draft_attn_groups
mock_kv_cache_spec = mock.MagicMock()
mock_kv_cache_spec.block_size = block_size
mock_builder = mock.MagicMock()
mock_builder.kv_cache_spec = mock_kv_cache_spec
proposer.attn_metadata_builder = mock_builder
mock_attn_group = mock.MagicMock()
mock_attn_group.kv_cache_spec = mock_kv_cache_spec
proposer.draft_attn_groups = [mock_attn_group]
# Request 0: query_len=4 (1 rejected), Request 1: query_len=4 (all valid)
batch_spec = BatchSpec(
@@ -916,7 +916,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
proposer.model = model_mock
# Assign draft attn_layer_names since load_model is not invoked
proposer.attn_layer_names = ["layer.0"]
proposer._draft_attn_layer_names = {"layer.0"}
# Create input tensors
batch_spec = BatchSpec(
@@ -961,20 +961,18 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
attn_metadata_builder = attn_metadata_builder_cls(
kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
layer_names=proposer.attn_layer_names,
layer_names=proposer._draft_attn_layer_names,
vllm_config=proposer.vllm_config,
device=device,
)
# Mock runner for attention metadata building
# Mock runner and draft_attn_groups for attention metadata building
proposer.runner = mock.MagicMock()
proposer.runner.attn_groups.append([mock.MagicMock()])
proposer.runner.attn_groups[0][
0
].get_metadata_builder.return_value = attn_metadata_builder
proposer._get_attention_metadata_builder = mock.MagicMock(
return_value=attn_metadata_builder
)
mock_attn_group = mock.MagicMock()
mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder
mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)
mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec
proposer.draft_attn_groups = [mock_attn_group]
result = proposer.propose(
target_token_ids=target_token_ids,
@@ -1089,7 +1087,7 @@ def test_propose_tree(spec_token_tree):
proposer.model = model_mock
# Assign draft attn_layer_names since load_model is not invoked
proposer.attn_layer_names = ["layer.0"]
proposer._draft_attn_layer_names = {"layer.0"}
# Get the tree attention metadata builder.
attn_metadata_builder_cls, _ = try_get_attention_backend(
@@ -1097,21 +1095,18 @@ def test_propose_tree(spec_token_tree):
)
attn_metadata_builder = attn_metadata_builder_cls(
kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
layer_names=proposer.attn_layer_names,
layer_names=proposer._draft_attn_layer_names,
vllm_config=proposer.vllm_config,
device=device,
)
# Mock runner for attention metadata building.
# Mock runner and draft_attn_groups for attention metadata building.
proposer.runner = mock.MagicMock()
proposer.runner.attn_groups.append([mock.MagicMock()])
proposer.runner.attn_groups[0][0].metadata_builders = [attn_metadata_builder]
proposer.runner.attn_groups[0][
0
].get_metadata_builder.return_value = attn_metadata_builder
proposer._get_attention_metadata_builder = mock.MagicMock(
return_value=attn_metadata_builder
)
mock_attn_group = mock.MagicMock()
mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder
mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)
mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec
proposer.draft_attn_groups = [mock_attn_group]
# Setup inputs for the proposer.
target_token_ids = torch.randint(0, vocab_size, (total_tokens,), device=device)
+7 -3
View File
@@ -162,7 +162,7 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):
model_mock.compute_logits.side_effect = logits_returns
proposer.model = model_mock
proposer.attn_layer_names = ["layer.0"]
proposer._draft_attn_layer_names = {"layer.0"}
# Prepare inputs
batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=seq_lens)
@@ -190,13 +190,17 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):
attn_metadata_builder = attn_metadata_builder_cls(
kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
layer_names=proposer.attn_layer_names,
layer_names=list(proposer._draft_attn_layer_names),
vllm_config=proposer.vllm_config,
device=device,
)
proposer.runner = mock.MagicMock()
proposer.attn_metadata_builder = attn_metadata_builder
mock_attn_group = mock.MagicMock()
mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder
mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)
mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec
proposer.draft_attn_groups = [mock_attn_group]
# Run propose
result = proposer.propose(
+1 -1
View File
@@ -64,7 +64,7 @@ def get_custom_wheel_versions(install_dir: str) -> dict[str, str]:
("torchaudio-", "torchaudio"), # Match torchaudio-
("amdsmi-", "amdsmi"), # Match amdsmi-
("flash_attn-", "flash-attn"), # Match flash_attn-
("aiter-", "aiter"), # Match aiter-
("amd_aiter-", "amd-aiter"), # Match amd_aiter-
]
for wheel_file in install_path.glob("*.whl"):
@@ -5,7 +5,6 @@ import torch
import torch._inductor.pattern_matcher as pm
from torch import fx
from torch._inductor.pattern_matcher import PatternMatcherPass
from torch._ops import OpOverload
import vllm.model_executor.layers.quantization.utils.fp8_utils # noqa: F401
from vllm._aiter_ops import rocm_aiter_ops
@@ -15,6 +14,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
QuantKey,
ScaleDesc,
kFp8Dynamic128Sym,
)
from vllm.platforms import current_platform
@@ -312,7 +312,9 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
@VllmInductorPass.time_and_log
def __call__(self, graph: fx.Graph) -> None:
self.matched_count = self.patterns.apply(graph)
logger.debug("Replaced %s patterns", self.matched_count)
logger.debug(
"%s Replaced %s patterns", self.__class__.__name__, self.matched_count
)
def uuid(self) -> str:
fusion_patterns = [
@@ -332,9 +334,11 @@ class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
FUSED_SILU_MUL_QUANT_OP = rocm_aiter_ops.get_act_mul_fused_fp8_group_quant_op()
def __init__(self, quant_op: OpOverload) -> None:
def __init__(self) -> None:
self.silu_and_mul_matcher = MatcherSiluAndMul()
self.quant_op = quant_op
self.quant_matcher = MatcherQuantFP8(
quant_key=kFp8Dynamic128Sym, match_rocm_aiter=True
)
def get_inputs(self) -> list[torch.Tensor]:
return [
@@ -346,7 +350,7 @@ class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
input: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
at1 = self.silu_and_mul_matcher(input)
at2 = self.quant_op(at1, 128)
at2 = self.quant_matcher(at1)
return at2[0], at2[1]
def replacement(
@@ -370,11 +374,6 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980
"""
AITER_GROUP_FP8_QUANT_OP = rocm_aiter_ops.get_group_quant_op()
TRITON_GROUP_FP8_QUANT_OP = torch.ops.vllm.triton_per_token_group_quant_fp8.default
QUANT_OPS = [AITER_GROUP_FP8_QUANT_OP, TRITON_GROUP_FP8_QUANT_OP]
@enable_fake_mode
def __init__(self, config: VllmConfig) -> None:
super().__init__(config)
@@ -383,8 +382,7 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
pass_name="rocm_aiter_silu_mul_fp8_group_quant_fusion_pass"
)
for quant_op in self.QUANT_OPS:
AiterSiluMulFp8GroupQuantPattern(quant_op).register(self.patterns)
AiterSiluMulFp8GroupQuantPattern().register(self.patterns)
self.dump_patterns(config, self.patterns)
@@ -135,7 +135,10 @@ _NIXL_SUPPORTED_DEVICE = {
"cpu",
),
"tpu": ("cpu",),
"xpu": ("cpu",),
"xpu": (
"cpu",
"xpu",
),
"cpu": ("cpu",),
}
# support for oot platform by providing mapping in current_platform
@@ -945,7 +948,7 @@ class NixlConnectorWorker:
# type based on kv_buffer_device
nixl_memory_type = current_platform.get_nixl_memory_type()
if nixl_memory_type is None:
if self.kv_buffer_device == "cuda":
if self.kv_buffer_device in ["cuda", "xpu"]:
nixl_memory_type = "VRAM"
elif self.kv_buffer_device == "cpu":
nixl_memory_type = "DRAM"
+9 -8
View File
@@ -851,6 +851,10 @@ class GroupCoordinator:
if self.world_size <= 1:
return []
if dst is None:
dst = (self.rank_in_group + 1) % self.world_size
assert dst < self.world_size, f"Invalid dst rank ({dst})"
if self.use_cpu_custom_send_recv:
if self.device_communicator is None:
raise ValueError("No device communicator found")
@@ -868,10 +872,6 @@ class GroupCoordinator:
group = self.device_group
metadata_group = self.cpu_group
if dst is None:
dst = (self.rank_in_group + 1) % self.world_size
assert dst < self.world_size, f"Invalid dst rank ({dst})"
metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
self.send_object(metadata_list, dst=dst)
@@ -948,6 +948,11 @@ class GroupCoordinator:
]:
if not torch.distributed.is_initialized() or self.world_size == 1:
return None, [], []
if src is None:
src = (self.rank_in_group - 1) % self.world_size
assert src < self.world_size, f"Invalid src rank ({src})"
if self.use_cpu_custom_send_recv:
if self.device_communicator is None:
raise ValueError("No device communicator found")
@@ -965,10 +970,6 @@ class GroupCoordinator:
group = self.device_group
metadata_group = self.cpu_group
if src is None:
src = (self.rank_in_group - 1) % self.world_size
assert src < self.world_size, f"Invalid src rank ({src})"
recv_metadata_list = self.recv_object(src=src)
tensor_dict: dict[str, Any] = {}
handles: list[Handle] = []
+1 -1
View File
@@ -77,7 +77,7 @@ class AnthropicTool(BaseModel):
class AnthropicToolChoice(BaseModel):
"""Tool Choice definition"""
type: Literal["auto", "any", "tool"]
type: Literal["auto", "any", "tool", "none"]
name: str | None = None
@model_validator(mode="after")
+2
View File
@@ -349,6 +349,8 @@ class AnthropicServingMessages(OpenAIServingChat):
req.tool_choice = "auto"
elif tool_choice_type == "any":
req.tool_choice = "required"
elif tool_choice_type == "none":
req.tool_choice = "none"
elif tool_choice_type == "tool":
req.tool_choice = ChatCompletionNamedToolChoiceParam.model_validate(
{
+8
View File
@@ -7,6 +7,7 @@ import warnings
from abc import ABC, abstractmethod
from collections import Counter, defaultdict
from collections.abc import Awaitable, Callable, Iterable
from dataclasses import dataclass
from functools import cached_property, lru_cache, partial
from itertools import accumulate
from pathlib import Path
@@ -1024,6 +1025,13 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):
self._add_placeholder("video", placeholder)
@dataclass
class ChatTemplateConfig:
chat_template: str | None = None
chat_template_content_format: ChatTemplateContentFormatOption = "auto"
trust_request_chat_template: bool = False
def validate_chat_template(chat_template: Path | str | None):
"""Raises if the provided chat template appears invalid."""
if chat_template is None:
+6
View File
@@ -220,6 +220,12 @@ def run_multi_api_server(args: argparse.Namespace):
num_api_servers: int = args.api_server_count
assert num_api_servers > 0
if num_api_servers > 1 and getattr(args, "use_gpu_for_pooling_score", False):
# TODO(wentao): remove this once well tested
raise ValueError(
"--use-gpu-for-pooling-score cannot be used with api_server_count > 1 now"
)
if num_api_servers > 1:
setup_multiprocess_prometheus()
+68 -26
View File
@@ -3,6 +3,7 @@
import itertools
from collections.abc import Callable, Iterable, Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Any
import cloudpickle
@@ -40,8 +41,11 @@ from vllm.distributed.weight_transfer.base import (
from vllm.engine.arg_utils import EngineArgs
from vllm.entrypoints.chat_utils import (
ChatCompletionMessageParam,
ChatTemplateConfig,
ChatTemplateContentFormatOption,
load_chat_template,
)
from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors
from vllm.entrypoints.pooling.score.utils import (
ScoreData,
ScoreMultiModalParam,
@@ -145,6 +149,7 @@ class LLM:
a tag name, or a commit id.
tokenizer_revision: The specific tokenizer version to use. It can be a
branch name, a tag name, or a commit id.
chat_template: The chat template to apply.
seed: The seed to initialize the random number generator for sampling.
gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to
reserve for the model weights, activations, and KV cache. Higher
@@ -232,6 +237,7 @@ class LLM:
quantization: QuantizationMethods | None = None,
revision: str | None = None,
tokenizer_revision: str | None = None,
chat_template: Path | str | None = None,
seed: int = 0,
gpu_memory_utilization: float = 0.9,
swap_space: float = 4,
@@ -384,9 +390,16 @@ class LLM:
self.model_config = self.llm_engine.model_config
self.renderer = self.llm_engine.renderer
self.chat_template = load_chat_template(chat_template)
self.io_processor = self.llm_engine.io_processor
self.input_processor = self.llm_engine.input_processor
self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template)
self.init_pooling_io_processors = init_pooling_io_processors(
supported_tasks=supported_tasks,
model_config=self.model_config,
renderer=self.renderer,
chat_template_config=self.chat_template_config,
)
# Cache for __repr__ to avoid repeated collective_rpc calls
self._cached_repr: str | None = None
@@ -1086,7 +1099,7 @@ class LLM:
"pooling model."
)
if use_io_processor := (isinstance(prompts, dict) and "data" in prompts):
if isinstance(prompts, dict) and "data" in prompts:
if self.io_processor is None:
raise ValueError(
"No IOProcessor plugin installed. Please refer "
@@ -1120,6 +1133,31 @@ class LLM:
for p in params_seq:
if p.task is None:
p.task = "plugin"
outputs = self._run_completion(
prompts=prompts_seq,
params=params_seq,
output_type=PoolingRequestOutput,
use_tqdm=use_tqdm,
lora_request=lora_request,
tokenization_kwargs=tokenization_kwargs,
)
# get the post-processed model outputs
assert self.io_processor is not None
processed_outputs = self.io_processor.post_process(outputs)
return [
PoolingRequestOutput[Any](
request_id="",
outputs=processed_outputs,
num_cached_tokens=getattr(
processed_outputs, "num_cached_tokens", 0
),
prompt_token_ids=[],
finished=True,
)
]
else:
if pooling_params is None:
# Use default pooling params.
@@ -1137,32 +1175,36 @@ class LLM:
)
raise ValueError(msg)
outputs = self._run_completion(
prompts=prompts_seq,
params=params_seq,
output_type=PoolingRequestOutput,
use_tqdm=use_tqdm,
lora_request=lora_request,
tokenization_kwargs=tokenization_kwargs,
)
if use_io_processor:
# get the post-processed model outputs
assert self.io_processor is not None
processed_outputs = self.io_processor.post_process(outputs)
return [
PoolingRequestOutput[Any](
request_id="",
outputs=processed_outputs,
num_cached_tokens=getattr(
processed_outputs, "num_cached_tokens", 0
),
prompt_token_ids=[],
finished=True,
if pooling_task in self.init_pooling_io_processors:
io_processor = self.init_pooling_io_processors[pooling_task]
processor_inputs = io_processor.pre_process_offline(
prompts_seq, tokenization_kwargs
)
]
seq_lora_requests = self._lora_request_to_seq(
lora_request, len(prompts_seq)
)
seq_priority = self._priority_to_seq(None, len(prompts))
self._render_and_add_requests(
prompts=processor_inputs,
params=params_seq,
lora_requests=seq_lora_requests,
priorities=seq_priority,
)
outputs = self._run_engine(
use_tqdm=use_tqdm, output_type=PoolingRequestOutput
)
outputs = io_processor.post_process(outputs)
else:
outputs = self._run_completion(
prompts=prompts_seq,
params=params_seq,
output_type=PoolingRequestOutput,
use_tqdm=use_tqdm,
lora_request=lora_request,
tokenization_kwargs=tokenization_kwargs,
)
return outputs
def embed(
@@ -38,6 +38,7 @@ from vllm.logprobs import Logprob
from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
from vllm.sampling_params import (
BeamSearchParams,
RepetitionDetectionParams,
RequestOutputKind,
SamplingParams,
StructuredOutputsParams,
@@ -336,6 +337,16 @@ class ChatCompletionRequest(OpenAIBaseModel):
),
)
repetition_detection: RepetitionDetectionParams | None = Field(
default=None,
description="Parameters for detecting repetitive N-gram patterns "
"in output tokens. If such repetition is detected, generation will "
"be ended early. LLMs can sometimes generate repetitive, unhelpful "
"token patterns, stopping only when they hit the maximum output length "
"(e.g. 'abcdabcdabcd...' or '\emoji \emoji \emoji ...'). This feature "
"can detect such behavior and terminate early, saving time and tokens.",
)
# --8<-- [end:chat-completion-extra-params]
def build_chat_params(
@@ -499,6 +510,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
allowed_token_ids=self.allowed_token_ids,
extra_args=extra_args or None,
skip_clone=True, # Created fresh per request, safe to skip clone
repetition_detection=self.repetition_detection,
)
@model_validator(mode="before")
@@ -1249,13 +1249,23 @@ class OpenAIServingChat(OpenAIServing):
)
# get the expected call based on partial JSON
# parsing which "autocompletes" the JSON
expected_call = json.dumps(
tool_parser.prev_tool_call_arr[index].get(
"arguments", {}
),
ensure_ascii=False,
# parsing which "autocompletes" the JSON.
# Tool parsers (e.g. Qwen3Coder) store
# arguments as a JSON string in
# prev_tool_call_arr. Calling json.dumps()
# on an already-serialized string would
# double-serialize it (e.g. '{"k":1}' becomes
# '"{\\"k\\":1}"'), which then causes the
# replace() below to fail and append the
# entire double-serialized string as a
# spurious final delta.
args = tool_parser.prev_tool_call_arr[index].get(
"arguments", {}
)
if isinstance(args, str):
expected_call = args
else:
expected_call = json.dumps(args, ensure_ascii=False)
# get what we've streamed so far for arguments
# for the current tool
+4
View File
@@ -278,6 +278,10 @@ class FrontendArgs(BaseFrontendArgs):
Enable offline FastAPI documentation for air-gapped environments.
Uses vendored static assets bundled with vLLM.
"""
use_gpu_for_pooling_score: bool = False
"""If set, run pooling score MaxSim on GPU in the API server process.
Can significantly improve late-interaction scoring performance.
https://github.com/vllm-project/vllm/pull/35330"""
@classmethod
def _customize_cli_kwargs(
@@ -26,6 +26,7 @@ from vllm.logprobs import Logprob
from vllm.renderers import TokenizeParams
from vllm.sampling_params import (
BeamSearchParams,
RepetitionDetectionParams,
RequestOutputKind,
SamplingParams,
StructuredOutputsParams,
@@ -166,6 +167,16 @@ class CompletionRequest(OpenAIBaseModel):
),
)
repetition_detection: RepetitionDetectionParams | None = Field(
default=None,
description="Parameters for detecting repetitive N-gram patterns "
"in output tokens. If such repetition is detected, generation will "
"be ended early. LLMs can sometimes generate repetitive, unhelpful "
"token patterns, stopping only when they hit the maximum output length "
"(e.g. 'abcdabcdabcd...' or '\emoji \emoji \emoji ...'). This feature "
"can detect such behavior and terminate early, saving time and tokens.",
)
# --8<-- [end:completion-extra-params]
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
@@ -310,6 +321,7 @@ class CompletionRequest(OpenAIBaseModel):
allowed_token_ids=self.allowed_token_ids,
extra_args=extra_args or None,
skip_clone=True, # Created fresh per request, safe to skip clone
repetition_detection=self.repetition_detection,
)
@model_validator(mode="before")

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