Compare commits

..
Author SHA1 Message Date
Isotr0pyandkhluu 5a532c7e0b Revert "[Misc] Move pyav and soundfile to common requirements" (#40276)
Co-authored-by: Roger Wang <hey@rogerw.io>
(cherry picked from commit 67eb6083e3)
2026-04-22 02:00:47 -07:00
Wentao Yeandkhluu 2f3555bf53 [Deprecation] Deprecate cprofile and cprofile_context (#39100)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
(cherry picked from commit 301024aa9c)
2026-04-22 02:00:12 -07:00
Jhao-Ting Chenandkhluu 4699f1bf8b fix: clamp NaN/Inf in topk_softmax to prevent duplicate expert IDs (#39391)
Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
(cherry picked from commit 28c222157b)
2026-04-22 01:59:52 -07:00
Nicolò Lucchesiandkhluu 681a6371cc [Bugfix][CI] Fix tests/distributed/test_torchrun_example_moe.py (#40349)
Signed-off-by: NickLucche <nlucches@redhat.com>
(cherry picked from commit 304d5ba1a0)
2026-04-22 01:59:37 -07:00
TJianandkhluu e86d349053 [ROCm] [Wheel] [Bugfix] [Critical] Remove any packages installed from github from rocm.txt e.g fastsafetensors as it is incompatible with uv pip (#40461)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
(cherry picked from commit 583e6f2226)
2026-04-22 01:58:51 -07:00
Julien DenizeandGitHub 6097afb9bd [BUGFIX] Fix Pixtral consolidated format vision weight loading (#39916)
Signed-off-by: Julien Denize <julien.denize@mistral.ai>
Signed-off-by: juliendenize <julien.denize@mistral.ai>
2026-04-19 22:25:03 -07:00
4f4713f96e [XPU] [torch.compile] Skipping CUDA graph memory estimation to avoid startup errors. (#39977)
Signed-off-by: chaojun-zhang <chaojun.zhang@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-04-20 13:04:39 +08:00
Tao HeandGitHub 8936118134 [Qwen][Bugfix] Fixes sigmoid activation in torch impl of RMSNormGated. (#40245)
Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com>
2026-04-20 04:28:19 +00:00
Yuan TangandGitHub 67ed01c353 fix: Do not make function calls when request has no tools for /v1/responses (#40314)
Signed-off-by: Yuan Tang <terrytangyuan@gmail.com>
2026-04-20 04:17:30 +00:00
6e10cb54f6 [Bugfix][Responses API] Fix streaming tool calls on /v1/responses (#39892)
Signed-off-by: Hoang Nguyen <118159510+hnt2601@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-04-20 11:24:52 +08:00
Shinichi HemmiGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
fcb31c1ac3 [Bugfix] Properly initialize PerTensorScaleParameter for fused-on-disk checkpoints (#39765)
Signed-off-by: Hemmi Shinichi <shemmi@preferred.jp>
Signed-off-by: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-04-20 02:53:04 +00:00
LxxandGitHub d886c26d4d [Doc] Fix typos in token_embed pooling documentation (#40266)
Signed-off-by: YifanLi3 <lyfqlx3@gmail.com>
2026-04-19 19:27:32 -07:00
898beca5a8 [BugFix][XPU] fix lora ops bgmv_expand size not match (#39989)
Signed-off-by: Ma, Liangliang <liangliang.ma@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-04-20 08:24:50 +08:00
Kevin H. LuuandGitHub 629d45eacb [ci] Make ecr authenticate non blocking (#40305)
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2026-04-19 15:37:53 -07:00
Andrew BarnesandGitHub f150107efd [ROCm] Fix cu_seqlens_q off-by-one in AITER FA speculative decode path (#39120)
Signed-off-by: Bortlesboat <bortstheboat@gmail.com>
2026-04-19 18:34:33 +00:00
daniserebandGitHub d1135a5087 Fix MoE backend selection for LoRA (unquantized MoE) (#40273)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>
2026-04-19 17:18:40 +00:00
982beae809 Optimize nemotron VL image/video preprocessing (#40283)
Signed-off-by: milesial <milesial@users.noreply.github.com>
Co-authored-by: milesial <milesial@users.noreply.github.com>
2026-04-19 15:06:20 +00:00
TJianandGitHub 45232a454e [FEAT] [Perf] [Gemma4] Fused Gemma4 Routing Function Triton (#39083)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-04-19 09:57:39 +00:00
Flora FengandGitHub 03ce1c6ed9 [Bugfix] Kimi-K2 tool parser streaming - fix token leakage, argument truncation, and content dropping (#38579)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-04-19 01:30:27 -07:00
omerpaz95andGitHub 4353c9cb4a [KV Offload] Pass request context (#39185)
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
2026-04-19 08:54:59 +03:00
4b7f5ea1a0 [KV Connector] Allow metrics of multiple connectors of same types in multi connector. (#40010)
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
2026-04-19 07:49:10 +03:00
38907e4391 [Frontend] Preserve structured output special tokens in offline LLM.chat (#39352)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>
2026-04-18 19:46:07 -04:00
d0359f3e04 [Bugfix] Guard mxfp4_experts_quant bindings on ENABLE_NVFP4_SM100 (#40191)
Signed-off-by: ultranationalism <www913363043@gmail.com>
Signed-off-by: mgoin <mike.goin12@gmail.com>
Co-authored-by: mgoin <mike.goin12@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-18 13:58:46 -07:00
Dan AlistarhandGitHub ed0622e3a8 [Attention] TurboQuant: remove redundant random signs, add prior art attribution (#40194)
Signed-off-by: Dan Alistarh <d.alistarh@gmail.com>
2026-04-18 14:31:59 -04:00
Yusuf MohammadandGitHub b5f6c5f834 Added general ND x ND matmul and unit test for it (#39909)
Signed-off-by: Yusuf <yusufmohammad@live.com>
2026-04-18 10:05:21 -04:00
Jee Jee LiandGitHub bfde49e287 [DOC] Add fuse_minimax_qk_norm (#39782)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
2026-04-18 00:41:37 -07:00
153ba7f0f3 [Refactor] Drop direct dependency on librosa (#39079)
Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-04-18 06:55:38 +00:00
Chinmay-Kulkarni-AMDandGitHub 87518c3027 [ZenCPU] AMD Zen CPU Backend with supported dtypes via zentorch weekly (#39967)
Signed-off-by: Chinmay Kulkarni <Chinmay.Kulkarni@amd.com>
2026-04-18 06:22:37 +00:00
Rishapveer SinghandGitHub aeee7ef939 [Bugfix] Fix k_proj's bias for GLM-ASR (#40160)
Signed-off-by: Rishapveer Singh <singhrishapveer@gmail.com>
2026-04-17 22:34:33 -07:00
z1yingandGitHub cda19ecf4d [Doc] Fix outdated source reference comment in anthropic/serving.py (#40189)
Signed-off-by: z1ying <tzzying@outlook.com>
2026-04-17 22:31:13 -07:00
80b18230e0 [Frontend] Add multimodal support to /inference/v1/generate endpoint (#38405)
Signed-off-by: Nithin Chalapathi <nithin.ch10@gmail.com>
Signed-off-by: Nithin Chalapathi <nithinc@berkeley.edu>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-04-17 20:31:56 -07:00
z1yingandGitHub d0697cc7b6 [Doc] Add Realtime Transcription section to supported_models.md (#39845)
Signed-off-by: Ziying Tao <tzzying@outlook.com>
2026-04-18 03:26:14 +00:00
milesialGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
b0755523dc [Core] Reduce mm scheduler, get_num_embed overhead (#40143)
Signed-off-by: milesial <milesial@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-04-18 11:25:49 +08:00
Chaojun ZhangGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
993859ceb0 [XPU] fix all_reduce all-zero accuracy issue under torch.compile (#39844)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-04-18 02:33:07 +00:00
Michael GoinandGitHub 48a65ccb02 [CI] Speed up test_fused_marlin_moe (#40178)
Signed-off-by: mgoin <mgoin64@gmail.com>
Signed-off-by: Michael Goin <mgoin64@gmail.com>
2026-04-17 19:26:21 -07:00
55842a8d69 [XPU]fake impl for xpu fp8_gemm (#39984)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-04-18 08:53:56 +08:00
Michael GoinandGitHub 1f45e83756 Remove outdated tests test_mixtral_moe and test_duplicated_ignored_sequence_group (#40175)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-04-17 16:49:43 -07:00
Michael GoinandGitHub a8bffaa133 [Kernel] Add MXFP4 W4A4 CUTLASS MoE kernel for SM100 (#37463)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-04-17 16:42:32 -07:00
5cdddddd4a [Kernel] [Helion] Force disable HOP path due to performance regression (#40171)
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>
Co-authored-by: Claude Sonnet 4 <noreply@anthropic.com>
2026-04-17 17:36:49 -04:00
aditi-amdandGitHub 6ef1efd51f [ROCm] Fix TurboQuant on ROCm: backend routing, flash-attn compat, int64 overflow (#39953)
Signed-off-by: aditi <aditi.rana@amd.com>
2026-04-17 13:08:37 -07:00
Ryan RockandGitHub 58da4ee047 [AMD][CI] Update DeepEP branch (#38396)
Signed-off-by: Ryan Rock <ryan.rock@amd.com>
2026-04-17 14:30:20 -05:00
Andreas KaratzasandGitHub 1ae11e2bfc [ROCm][CI] Build fastsafetensors from source so it links against libamdhip64 (#39978)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-04-17 14:30:08 -05:00
251c18d1f8 skip fp8e4b15 on xpu (#39957)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-04-17 16:55:08 +00:00
512765d52d [Misc][UX] Map mimo reasoning and tooling parsers (#40089)
Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
2026-04-17 16:49:21 +00:00
allgatherGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
640cc9dd7d feat: Add LoRA support for Gemma4ForConditionalGeneration (#39291)
Signed-off-by: allgather <all2allops@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-04-17 09:39:19 -07:00
ceade1952c [BugFix] Support custom tool parsers when tool_choice is required and named function (#39870)
Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: sfeng33 <4florafeng@gmail.com>
Co-authored-by: sfeng33 <4florafeng@gmail.com>
2026-04-17 16:38:10 +00:00
747256bb5d [Bugfix][Core] Fix stuck chunked pipeline parallelism with async scheduling (#38726)
Signed-off-by: Jing Wang <jingwang96@qq.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-04-17 16:02:50 +00:00
Michael GoinandGitHub 1174723eba Fix TURBOQUANT backend selection in cuda.py (#40060)
Signed-off-by: Michael Goin <mgoin64@gmail.com>
2026-04-17 07:31:41 -07:00
sychen52andGitHub 6b2b7bd0eb Add nvfp4 support to reshape_and_cache_flash (#37332)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
2026-04-17 07:28:00 -07:00
Ben BrowningandGitHub 70770268c3 Add @bbrowning to CODEOWNERS (#40141)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-04-17 09:51:48 -04:00
ChaunceyandGitHub 7a51b3e415 [Bugfix] Fix empty delta detection in Qwen3XMLToolParser streaming (#40090)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-04-17 13:34:55 +00:00
Li, JiangandGitHub d02421a7db [CPU] Refactor CPU affinity and memory management (#39781)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
2026-04-17 21:01:08 +08:00
Lukas GeigerandGitHub b1dc87a098 [Models][Gemma4] Prevent GPU/CPU sync in embed_input_ids (#39234)
Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>
2026-04-17 12:37:21 +00:00
Or OzeriandGitHub 79a5b63253 [kv_offload]: Fix num CPU blocks for UniformTypeKVCacheSpecs (#39617)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
2026-04-17 15:13:55 +03:00
MaralandGitHub c0c98b8b9a [Bugfix] Add Marlin kernel in block scaled mm kernel selection. (#40105)
Signed-off-by: maral <maralbahari.98@gmail.com>
2026-04-17 10:20:32 +00:00
wang.yuqiandGitHub 8d2cff8140 [Examples] Resettle Observability examples. (#40123)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-04-17 03:13:31 -07:00
Cyrus LeungandGitHub 4f436782af [Misc] Improve new PR bot trigger condition (#40114)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
2026-04-17 16:56:22 +08:00
z1yingandGitHub bf45e6d0a5 [Doc] Add Gemma 4 to supported models list (#39607)
Signed-off-by: z1ying <tzzying@outlook.com>
Signed-off-by: Ziying Tao <tzzying@outlook.com>
2026-04-17 13:42:52 +08:00
978a4462bb [CI Failure] Fix Plugin Tests (2 GPUs) Failure (#40083)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
Co-authored-by: Michele Gazzetti <michele.gazzetti1@ibm.com>
2026-04-17 04:17:39 +00:00
Michael GoinandGitHub 1948d0c467 [UX] Defer some imports on CLI paths to save ~2s (#40056)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-04-16 19:48:37 -07:00
Shinichi HemmiandGitHub 4c47710bf7 [CI/Build] Apply ruff formatter to pass pre-commit (#40078)
Signed-off-by: Hemmi Shinichi <shemmi@preferred.jp>
2026-04-17 08:54:32 +08:00
Giancarlo DelfinandGitHub bf9a5ddb24 [MLA] Optimize mla indexer prepare uniform decode for MTP > 1 (#39458)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
2026-04-16 16:27:51 -07:00
bnellnmandGitHub 79e799ebbd [Bugfix] Temporarily disable B200 fp4 MoE layer tests (#40057)
Signed-off-by: Bill Nell <bnell@redhat.com>
2026-04-16 19:26:55 -04:00
Netanel HaberandGitHub c4e601c73c Bugfix: Parakeet: .conv.pointwise/depthwise_conv1/2.bias weigths can exist even if convolution_bias=False (#40007)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
2026-04-16 23:22:05 +00:00
BadrBasowidandGitHub 29057d3bee [Compilation] Add Unit Tests for VllmFusionPatternMatcherPass (#39692)
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
2026-04-16 22:57:16 +00:00
Matthew BonanniandGitHub 219bb5b8c0 [Misc] Update committers.md (#40058)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
2026-04-16 13:48:41 -07:00
Asaf GardinandGitHub ad2b1277f9 [Quantization] Consolidate experts_int8 with fp8 online quantization (#38463)
Signed-off-by: Josephasafg <ajgard7@gmail.com>
2026-04-16 13:12:20 -07:00
roikoren755andGitHub b897f00c9c Gate SSU dispatch setup (#40039)
Signed-off-by: Roi Koren <roik@nvidia.com>
2026-04-16 13:06:01 -07:00
adf9bb3c57 [CI] Add weight transfer tests to CI (#39821)
Signed-off-by: SumanthRH <sumanthrh99@gmail.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-04-16 15:51:45 -04:00
Flora FengandGitHub b16fda62b7 [Misc] Add @sfeng33 to CODEOWNERS (#40048)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-04-16 12:25:29 -07:00
Yufeng HeandGitHub de111f3246 [Bugfix] Fix bench_serve UTF-8 decode crash on split multi-byte chars (#38732) 2026-04-16 12:01:25 -07:00
Jared WenandGitHub afabb5f45a [bugfix] Normalize tool message content from array to string format (#39899)
Signed-off-by: JaredforReal <w13431838023@gmail.com>
2026-04-16 11:54:39 -07:00
Roger WangandGitHub 3abb7560c0 [Bugfix] Fix audioflamingo test (#40052)
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-04-16 11:53:58 -07:00
187 changed files with 6807 additions and 3549 deletions
+2 -2
View File
@@ -92,8 +92,8 @@ check_and_skip_if_image_exists() {
}
ecr_login() {
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true
}
prepare_cache_tags() {
+1 -1
View File
@@ -11,7 +11,7 @@ REPO=$2
BUILDKITE_COMMIT=$3
# authenticate with AWS ECR
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
# skip build if image already exists
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
@@ -11,7 +11,7 @@ REPO=$2
BUILDKITE_COMMIT=$3
# authenticate with AWS ECR
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
# skip build if image already exists
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then
@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
exit 1
fi
#echo "--- DP+TP"
#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
#server_pid=$!
#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
#vllm bench serve \
# --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; 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
echo "--- DP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--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; 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
+1 -1
View File
@@ -2613,6 +2613,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- export TORCH_NCCL_BLOCKING_WAIT=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
@@ -3601,7 +3602,6 @@ steps:
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
- pytest -v -s tests/distributed/test_context_parallel.py
- 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
View File
@@ -196,6 +196,8 @@ steps:
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
- 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
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
- pytest -v -s tests/distributed/test_packed_tensor.py
- label: Distributed Tests (2 GPUs)(B200)
device: b200
+1
View File
@@ -141,6 +141,7 @@ steps:
- pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
- pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
- pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
- pytest -v -s tests/kernels/moe/test_flashinfer.py
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
+7 -3
View File
@@ -44,8 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/vllm/pooling_params.py @noooop @DarkLight1337
/vllm/tokenizers @DarkLight1337 @njhill
/vllm/renderers @DarkLight1337 @njhill
/vllm/reasoning @aarnphm @chaunceyjiang
/vllm/tool_parsers @aarnphm @chaunceyjiang
/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning
# vLLM V1
/vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -91,7 +92,10 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/tests/v1/kv_connector/nixl_integration @NickLucche
/tests/v1/kv_connector @ApostaC @orozery
/tests/v1/kv_offload @ApostaC @orozery
/tests/v1/determinism @yewentao256
/tests/v1/determinism @yewentao256
/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning
# Transformers modeling backend
/vllm/model_executor/models/transformers @hmellor
+1
View File
@@ -45,6 +45,7 @@ jobs:
- name: Smoke test vllm serve
run: |
# Start server in background
VLLM_CPU_KVCACHE_SPACE=1 \
vllm serve Qwen/Qwen3-0.6B \
--max-model-len=2K \
--load-format=dummy \
+5 -5
View File
@@ -62,14 +62,14 @@ jobs:
const prAuthor = context.payload.pull_request.user.login;
const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`,
per_page: 1,
});
const authorPRCount = searchResults.total_count;
console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
const mergedPRCount = searchResults.total_count;
console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`);
if (authorPRCount === 1) {
if (mergedPRCount === 0) {
console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
await github.rest.issues.createComment({
owner,
@@ -98,5 +98,5 @@ jobs:
].join('\n'),
});
} else {
console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`);
}
+17 -1
View File
@@ -923,6 +923,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
# nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
# from cache_kernels.cu, so it belongs in _C rather than _C_stable.
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -944,11 +952,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
+15 -14
View File
@@ -30,6 +30,21 @@ else()
list(APPEND CXX_COMPILE_FLAGS
"-fopenmp"
"-DVLLM_CPU_EXTENSION")
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
# and create a local shim dir with it
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
find_library(OPEN_MP
NAMES gomp
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
NO_DEFAULT_PATH
REQUIRED
)
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
if (OPEN_MP)
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
endif()
endif()
if (NOT MACOSX_FOUND)
@@ -175,20 +190,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
if(NOT NPROC)
set(NPROC 4)
endif()
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
# and create a local shim dir with it
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
find_library(OPEN_MP
NAMES gomp
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
NO_DEFAULT_PATH
REQUIRED
)
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
if (OPEN_MP)
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
endif()
# Fetch and populate ACL
if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
+22 -2
View File
@@ -724,6 +724,28 @@ void reshape_and_cache_flash(
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (kv_cache_dtype == "nvfp4") {
#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
// NVFP4 dispatch is compiled separately for SM100+.
extern void reshape_and_cache_nvfp4_dispatch(
torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
torch::Tensor & value_cache, torch::Tensor & slot_mapping,
torch::Tensor & k_scale, torch::Tensor & v_scale);
reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
slot_mapping, k_scale, v_scale);
return;
#else
TORCH_CHECK(false,
"NVFP4 KV cache requires SM100+ (Blackwell). "
"Please rebuild vllm with a Blackwell-compatible CUDA target.");
#endif
}
// Original FP8/auto path.
int block_size = key_cache.size(1);
int64_t key_stride = key.stride(0);
@@ -741,8 +763,6 @@ void reshape_and_cache_flash(
dim3 grid(num_tokens);
dim3 block(std::min(num_heads * head_size, 512));
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
CALL_RESHAPE_AND_CACHE_FLASH);
+4
View File
@@ -141,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping,
const int64_t block_size);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
void eagle_prepare_inputs_padded_kernel_impl(
const torch::Tensor& cu_num_draft_tokens,
@@ -431,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"block_size) -> ()",
&compute_slot_mapping_kernel_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels
ops.def(
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
+73 -6
View File
@@ -13,13 +13,80 @@
#include "cpu/utils.hpp"
#ifdef VLLM_NUMA_DISABLED
std::string init_cpu_threads_env(const std::string& cpu_ids) {
return std::string(
"Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
"no effect to setup thread affinity.");
}
void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
#else
void init_cpu_memory_env(std::vector<int64_t> node_ids) {
// Memory node binding
if (numa_available() != -1) {
// Concatenate all node_ids into a single comma-separated string
if (!node_ids.empty()) {
std::string node_ids_str;
for (const int node_id : node_ids) {
if (!node_ids_str.empty()) {
node_ids_str += ",";
}
node_ids_str += std::to_string(node_id);
}
#endif
bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
bitmask* src_mask = numa_get_mems_allowed();
int pid = getpid();
if (mask && src_mask) {
// move all existing pages to the specified numa node.
*(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
int page_num = numa_migrate_pages(pid, src_mask, mask);
if (page_num == -1) {
TORCH_WARN("numa_migrate_pages failed. errno: " +
std::to_string(errno));
}
// Restrict memory allocation to the selected NUMA node(s).
// Enhances memory locality for the threads bound to those NUMA CPUs.
if (node_ids.size() > 1) {
errno = 0;
numa_set_interleave_mask(mask);
if (errno != 0) {
TORCH_WARN("numa_set_interleave_mask failed. errno: " +
std::to_string(errno));
} else {
TORCH_WARN(
"NUMA binding: Using INTERLEAVE policy for memory "
"allocation across multiple NUMA nodes (nodes: " +
node_ids_str +
"). Memory allocations will be "
"interleaved across the specified NUMA nodes.");
}
} else {
errno = 0;
numa_set_membind(mask);
if (errno != 0) {
TORCH_WARN("numa_set_membind failed. errno: " +
std::to_string(errno));
} else {
TORCH_WARN(
"NUMA binding: Using MEMBIND policy for memory "
"allocation on the NUMA nodes (" +
node_ids_str +
"). Memory allocations will be "
"strictly bound to these NUMA nodes.");
}
}
numa_set_strict(1);
numa_free_nodemask(mask);
numa_free_nodemask(src_mask);
} else {
TORCH_WARN(
"numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
std::to_string(errno));
}
}
}
}
#endif // VLLM_NUMA_DISABLED
namespace cpu_utils {
ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
+9
View File
@@ -134,4 +134,13 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
torch::stable::Tensor& input,
torch::stable::Tensor& input_global_scale);
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
const torch::stable::Tensor& a,
const torch::stable::Tensor& b,
const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets);
#endif
@@ -0,0 +1,468 @@
/*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
* Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
*/
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include <cutlass/arch/arch.h>
#include "cutlass_extensions/common.hpp"
#include "cute/tensor.hpp"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/collective/default_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/group_array_problem_shape.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/util/packed_stride.hpp"
#include <cassert>
using namespace cute;
// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
template <typename ElementAB, typename ElementC, typename ElementSF,
typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
__global__ void __mxfp4_get_group_gemm_starts(
ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
ElementAB* a_base_as_int, ElementAB* b_base_as_int,
ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
const int64_t a_stride_val, const int64_t b_stride_val,
const int64_t c_stride_val, const int K, const int N) {
int64_t expert_id = threadIdx.x;
if (expert_id >= gridDim.x * blockDim.x) {
return;
}
int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
int64_t group_size = 32;
int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
"unexpected problem sizes");
int64_t half_k = static_cast<int64_t>(k / 2);
int64_t group_k = static_cast<int64_t>(k / group_size);
// Shape of A as uint8/byte = [M, K // 2]
a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
// Shape of B as uint8/byte = [E, N, K // 2]
b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
// Shape of C = [M, N]
out_offsets[expert_id] = out_base_as_int + expert_offset * n;
// Shape of a_scale = [sum(sf_sizes), K // group_size]
a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
0 &&
"TMA requires 128-byte alignment");
// Shape of B scale = [E, N, K // group_size]
b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
0 &&
"TMA requires 128-byte alignment");
// Initialize strides
a_strides[expert_id] = a_stride_val;
b_strides[expert_id] = b_stride_val;
c_strides[expert_id] = c_stride_val;
LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
*layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
*layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
}
#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE, \
TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
LayoutSFB, ScaleConfig) \
else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
__mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
LayoutSFB, ScaleConfig> \
<<<1, num_experts, 0, stream>>>( \
static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \
static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()), \
static_cast<C_TYPE**>(out_starts.data_ptr()), \
static_cast<SF_TYPE**>(a_scales_starts.data_ptr()), \
static_cast<SF_TYPE**>(b_scales_starts.data_ptr()), \
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \
static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()), \
static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()), \
static_cast<C_TYPE*>(out_tensors.data_ptr()), \
static_cast<SF_TYPE*>(a_scales.data_ptr()), \
static_cast<SF_TYPE*>(b_scales.data_ptr()), \
static_cast<int32_t*>(expert_offsets.data_ptr()), \
static_cast<int32_t*>(sf_offsets.data_ptr()), \
static_cast<int32_t*>(problem_sizes.data_ptr()), \
static_cast<int64_t*>(a_strides.data_ptr()), \
static_cast<int64_t*>(b_strides.data_ptr()), \
static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val, \
b_stride_val, c_stride_val, K, N); \
}
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
void mxfp4_run_get_group_gemm_starts(
const torch::stable::Tensor& a_starts,
const torch::stable::Tensor& b_starts,
const torch::stable::Tensor& out_starts,
const torch::stable::Tensor& a_scales_starts,
const torch::stable::Tensor& b_scales_starts,
const torch::stable::Tensor& layout_sfa,
const torch::stable::Tensor& layout_sfb,
const torch::stable::Tensor& a_strides,
const torch::stable::Tensor& b_strides,
const torch::stable::Tensor& c_strides, int64_t a_stride_val,
int64_t b_stride_val, int64_t c_stride_val,
torch::stable::Tensor const& a_tensors,
torch::stable::Tensor const& b_tensors,
torch::stable::Tensor const& out_tensors,
torch::stable::Tensor const& a_scales,
torch::stable::Tensor const& b_scales,
torch::stable::Tensor const& expert_offsets,
torch::stable::Tensor const& sf_offsets,
torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
int num_experts = (int)expert_offsets.size(0);
auto stream = get_current_cuda_stream(a_tensors.get_device_index());
STD_TORCH_CHECK(out_tensors.size(1) == N,
"Output tensor shape doesn't match expected shape");
STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
"b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
" dimension must match");
if (false) {
}
// MXFP4 uses E8M0 (float_ue8m0_t) scale factors
__CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
torch::headeronly::ScalarType::BFloat16,
cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
ScaleConfig)
__CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
torch::headeronly::ScalarType::Half, half,
LayoutSFA, LayoutSFB, ScaleConfig)
else {
STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
}
}
template <typename OutType>
void run_mxfp4_blockwise_scaled_group_mm_sm100(
torch::stable::Tensor& output, const torch::stable::Tensor& a,
const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
using ProblemShape =
cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
using ElementType = cutlass::float_e2m1_t;
using ElementSFType = cutlass::float_ue8m0_t;
using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementC = OutType;
using ElementD = ElementC;
using ElementAccumulator = float;
// Layout definitions
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::RowMajor;
using LayoutD = LayoutC;
static constexpr int AlignmentA = 32;
static constexpr int AlignmentB = 32;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
// Architecture definitions
using ArchTag = cutlass::arch::Sm100;
using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
using StageCountType = cutlass::gemm::collective::StageCountAuto;
using ClusterShape = Shape<_1, _1, _1>;
struct MMA1SMConfig {
using MmaTileShape = Shape<_128, _128, _128>;
using KernelSchedule =
cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
};
using CollectiveEpilogue =
typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
ClusterShape, Shape<_128, _64>, ElementAccumulator,
ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
LayoutC*, AlignmentD,
typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
using CollectiveMainloop =
typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
ElementB, LayoutB*, AlignmentB, ElementAccumulator,
typename MMA1SMConfig::MmaTileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
sizeof(typename CollectiveEpilogue::SharedStorage))>,
typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
using GemmKernel =
cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
CollectiveEpilogue>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
using StrideA = typename Gemm::GemmKernel::InternalStrideA;
using StrideB = typename Gemm::GemmKernel::InternalStrideB;
using StrideC = typename Gemm::GemmKernel::InternalStrideC;
using StrideD = typename Gemm::GemmKernel::InternalStrideD;
using LayoutSFA =
typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
using LayoutSFB =
typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
using ScaleConfig =
typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
int num_experts = static_cast<int>(expert_offsets.size(0));
torch::stable::Tensor a_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor b_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor out_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor a_scales_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor b_scales_ptrs =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor layout_sfa = torch::stable::empty(
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
a.device());
torch::stable::Tensor layout_sfb = torch::stable::empty(
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
a.device());
torch::stable::Tensor a_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor b_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
torch::stable::Tensor c_strides1 =
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
std::nullopt, a.device());
mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
// Create an instance of the GEMM
Gemm gemm_op;
UnderlyingProblemShape* problem_sizes_as_shapes =
static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
// Set the Scheduler info
cutlass::KernelHardwareInfo hw_info;
using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
PersistentTileSchedulerSm100GroupParams<
typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
scheduler.raster_order = RasterOrderOptions::AlongM;
hw_info.device_id = a.get_device_index();
static std::unordered_map<int, int> cached_sm_counts;
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
cached_sm_counts[hw_info.device_id] =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
hw_info.device_id);
}
hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
// Mainloop Arguments
typename GemmKernel::MainloopArguments mainloop_args{
static_cast<const ElementType**>(a_ptrs.data_ptr()),
static_cast<StrideA*>(a_strides1.data_ptr()),
static_cast<const ElementType**>(b_ptrs.data_ptr()),
static_cast<StrideB*>(b_strides1.data_ptr()),
static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
// Epilogue Arguments
typename GemmKernel::EpilogueArguments epilogue_args{
{}, // epilogue.thread
nullptr,
static_cast<StrideC*>(c_strides1.data_ptr()),
static_cast<ElementD**>(out_ptrs.data_ptr()),
static_cast<StrideC*>(c_strides1.data_ptr())};
auto& fusion_args = epilogue_args.thread;
// Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
fusion_args.alpha_ptr = nullptr;
fusion_args.beta_ptr = nullptr;
fusion_args.alpha = 1.0f;
fusion_args.alpha_ptr_array = nullptr;
fusion_args.dAlpha = {_0{}, _0{}, 0};
fusion_args.beta = 0.0f;
fusion_args.beta_ptr_array = nullptr;
fusion_args.dBeta = {_0{}, _0{}, 0};
// Gemm Arguments
typename GemmKernel::Arguments args{
cutlass::gemm::GemmUniversalMode::kGrouped,
{num_experts, problem_sizes_as_shapes, nullptr},
mainloop_args,
epilogue_args,
hw_info,
scheduler};
size_t workspace_size = Gemm::get_workspace_size(args);
auto workspace =
torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
std::nullopt, a.device());
const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
auto can_implement_status = gemm_op.can_implement(args);
STD_TORCH_CHECK(
can_implement_status == cutlass::Status::kSuccess,
"Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
// Run the GEMM
auto status = gemm_op.initialize(args, workspace.data_ptr());
STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
"Failed to initialize MXFP4 GEMM: status=", (int)status,
" workspace_size=", workspace_size,
" num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
status = gemm_op.run(args, workspace.data_ptr(), stream);
STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
"Failed to run MXFP4 GEMM");
}
template <typename OutType>
void run_mxfp4_blockwise_scaled_group_mm(
torch::stable::Tensor& output, const torch::stable::Tensor& a,
const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
int32_t version_num = get_sm_version_num();
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
if (version_num >= 100 && version_num < 120) {
run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
output, a, b, a_blockscale, b_blockscales, problem_sizes,
expert_offsets, sf_offsets, M, N, K);
return;
}
#endif
STD_TORCH_CHECK_NOT_IMPLEMENTED(
false,
"No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
version_num, ". Required capability: 100");
}
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
// E8M0 scale factors stored as uint8
constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
#endif
#define CHECK_TYPE(x, st, m) \
STD_TORCH_CHECK(x.scalar_type() == st, \
": Inconsistency of torch::stable::Tensor type:", m)
#define CHECK_TH_CUDA(x, m) \
STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
#define CHECK_CONTIGUOUS(x, m) \
STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
#define CHECK_INPUT(x, st, m) \
CHECK_TH_CUDA(x, m); \
CHECK_CONTIGUOUS(x, m); \
CHECK_TYPE(x, st, m)
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
const torch::stable::Tensor& a,
const torch::stable::Tensor& b,
const torch::stable::Tensor& a_blockscale,
const torch::stable::Tensor& b_blockscales,
const torch::stable::Tensor& problem_sizes,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets) {
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
// Input validation
CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
// MXFP4 uses E8M0 scale factors (stored as uint8)
CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
STD_TORCH_CHECK(
a_blockscale.dim() == 2,
"expected a_blockscale to be of shape [num_experts, rounded_m,"
" k // group_size], observed rank: ",
a_blockscale.dim())
STD_TORCH_CHECK(b_blockscales.dim() == 3,
"expected b_blockscale to be of shape: "
" [num_experts, n, k // group_size], observed rank: ",
b_blockscales.dim())
STD_TORCH_CHECK(problem_sizes.dim() == 2,
"problem_sizes must be a 2D tensor");
STD_TORCH_CHECK(problem_sizes.size(1) == 3,
"problem_sizes must have the shape (num_experts, 3)");
STD_TORCH_CHECK(
problem_sizes.size(0) == expert_offsets.size(0),
"Number of experts in problem_sizes must match expert_offsets");
STD_TORCH_CHECK(
problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
"problem_sizes must be int32.");
int M = static_cast<int>(a.size(0));
int N = static_cast<int>(b.size(1));
int E = static_cast<int>(b.size(0));
int K = static_cast<int>(2 * b.size(2));
if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
output, a, b, a_blockscale, b_blockscales, problem_sizes,
expert_offsets, sf_offsets, M, N, K);
} else {
run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
output, a, b, a_blockscale, b_blockscales, problem_sizes,
expert_offsets, sf_offsets, M, N, K);
}
#else
STD_TORCH_CHECK_NOT_IMPLEMENTED(
false,
"No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
"SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
#endif
}
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
}
@@ -0,0 +1,432 @@
/*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* MXFP4 activation quantization kernel for MoE experts.
* Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
* over 32-element groups.
*
* Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
* - Block size 32 (2 threads per SF in PACK16 mode)
* - E8M0 (power-of-two) scale factors
* - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
*/
// MXFP4 requires PACK16 mode (16 elements per thread) so that
// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
// Must be defined before any header that (transitively) includes
// nvfp4_utils.cuh.
#define NVFP4_ENABLE_ELTS16 1
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include <cuda_fp8.h>
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include "cuda_vec_utils.cuh"
#include "cuda_utils.h"
#include "nvfp4_utils.cuh"
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
#include "launch_bounds_utils.h"
namespace vllm {
// MXFP4 block size constants
static constexpr int MXFP4_SF_VEC_SIZE = 32;
// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
static constexpr int MXFP4_NUM_THREADS_PER_SF =
MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
// MXFP4 quantization kernel for experts.
// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
// SiLU(gate)*up before quantization.
template <class Type, bool FUSE_SILU_MUL = false,
bool SMALL_NUM_EXPERTS = false>
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
fp4_packed_t* out, uint32_t* SFout,
uint32_t* input_offset_by_experts,
uint32_t* output_scale_offset_by_experts,
int n_experts, bool low_latency) {
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
"Vec size is not matched.");
// MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
int32_t const numKTiles = (numCols + 127) / 128;
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
globalIdx += gridDim.x * blockDim.x) {
int rowIdx = globalIdx / colsPerRow;
int colIdx = globalIdx % colsPerRow;
int rowIdx_in_expert = 0;
int expert_idx = 0;
if constexpr (SMALL_NUM_EXPERTS) {
for (int i = 0; i < n_experts; i++) {
uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
if (rowIdx >= current_offset && rowIdx < next_offset) {
rowIdx_in_expert = rowIdx - current_offset;
expert_idx = i;
break;
}
}
} else {
uint32_t local_offsets[17];
for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
*reinterpret_cast<int4*>(local_offsets) =
__ldca(reinterpret_cast<const int4*>(
&input_offset_by_experts[chunk_start]));
*reinterpret_cast<int4*>(local_offsets + 4) =
__ldca(reinterpret_cast<const int4*>(
&input_offset_by_experts[chunk_start + 4]));
*reinterpret_cast<int4*>(local_offsets + 8) =
__ldca(reinterpret_cast<const int4*>(
&input_offset_by_experts[chunk_start + 8]));
*reinterpret_cast<int4*>(local_offsets + 12) =
__ldca(reinterpret_cast<const int4*>(
&input_offset_by_experts[chunk_start + 12]));
local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
#pragma unroll
for (int i = 0; i < 16; i++) {
if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
rowIdx_in_expert = rowIdx - local_offsets[i];
expert_idx = chunk_start + i;
break;
}
}
}
}
// Load input and optionally apply fused SiLU+Mul
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
PackedVec quant_input;
if constexpr (FUSE_SILU_MUL) {
PackedVec in_vec_up =
reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
quant_input = compute_silu_mul(in_vec, in_vec_up);
} else {
quant_input = in_vec;
}
// In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
int64_t outOffset = rowIdx * colsPerRow + colIdx;
auto& out_pos = out[outOffset];
uint32_t* SFout_in_expert =
SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
// Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
auto sf_out =
cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
// Block E8M0 scales only; no extra tensor-level scale in this path
constexpr float SFScaleVal = 1.0f;
// UE8M0_SF=true for MXFP4 E8M0 scale factors
out_pos =
cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
quant_input, SFScaleVal, sf_out);
}
}
// Large M_topk variant using shared memory for expert offsets
template <class Type, bool FUSE_SILU_MUL = false,
bool SMALL_NUM_EXPERTS = false>
__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
fp4_packed_t* out, uint32_t* SFout,
uint32_t* input_offset_by_experts,
uint32_t* output_scale_offset_by_experts,
int n_experts) {
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
"Vec size is not matched.");
// MXFP4: numKTiles = ceil(numCols / 128)
int32_t const numKTiles = (numCols + 127) / 128;
extern __shared__ uint32_t shared_input_offsets[];
if constexpr (SMALL_NUM_EXPERTS) {
for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
shared_input_offsets[i] = input_offset_by_experts[i];
}
} else {
for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
*reinterpret_cast<int4*>(&shared_input_offsets[i]) =
*reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
}
if (threadIdx.x == 0) {
shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
}
}
__syncthreads();
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
globalIdx += gridDim.x * blockDim.x) {
int rowIdx = globalIdx / colsPerRow;
int colIdx = globalIdx % colsPerRow;
int rowIdx_in_expert = 0;
int expert_idx = 0;
// Binary search through experts using shared memory
int left = 0, right = n_experts - 1;
while (left <= right) {
int mid = (left + right) / 2;
uint32_t mid_offset = shared_input_offsets[mid];
uint32_t next_offset = shared_input_offsets[mid + 1];
if (rowIdx >= mid_offset && rowIdx < next_offset) {
rowIdx_in_expert = rowIdx - mid_offset;
expert_idx = mid;
break;
} else if (rowIdx < mid_offset) {
right = mid - 1;
} else {
left = mid + 1;
}
}
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
PackedVec quant_input;
if constexpr (FUSE_SILU_MUL) {
PackedVec in_vec_up =
reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
quant_input = compute_silu_mul(in_vec, in_vec_up);
} else {
quant_input = in_vec;
}
int64_t outOffset = rowIdx * colsPerRow + colIdx;
auto& out_pos = out[outOffset];
// MXFP4 has no global scale - only block-level E8M0 scale factors
constexpr float SFScaleVal = 1.0f;
uint32_t* SFout_in_expert =
SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
auto sf_out =
cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
out_pos =
cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
quant_input, SFScaleVal, sf_out);
}
}
template <typename T, bool FUSE_SILU_MUL = false>
void mxfp4_quant_impl(void* output, void* output_scale, void* input,
void* input_offset_by_experts,
void* output_scale_offset_by_experts, int m_topk, int k,
int n_experts, cudaStream_t stream) {
int multiProcessorCount =
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
int const workSizePerRow = k / ELTS_PER_THREAD;
int const totalWorkSize = m_topk * workSizePerRow;
dim3 block(std::min(workSizePerRow, 512));
int const numBlocksPerSM =
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
multiProcessorCount * numBlocksPerSM));
while (grid.x <= multiProcessorCount && block.x > 64) {
grid.x *= 2;
block.x = (block.x + 1) / 2;
}
int const blockRepeat =
(totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
if (blockRepeat > 1) {
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
if (n_experts >= 4) {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
<<<grid, block, shared_mem_size, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
reinterpret_cast<fp4_packed_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts);
} else {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
<<<grid, block, shared_mem_size, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
reinterpret_cast<fp4_packed_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts);
}
} else {
if (n_experts >= 16) {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
<<<grid, block, 0, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
reinterpret_cast<fp4_packed_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts, /* bool low_latency */ true);
} else {
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
reinterpret_cast<fp4_packed_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
n_experts, /* bool low_latency */ true);
}
}
}
} // namespace vllm
/*Quantization entry for mxfp4 experts quantization*/
#define CHECK_TH_CUDA(x, m) \
STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x, m) \
STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
#define CHECK_INPUT(x, m) \
CHECK_TH_CUDA(x, m); \
CHECK_CONTIGUOUS(x, m);
constexpr auto HALF = torch::headeronly::ScalarType::Half;
constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
constexpr auto INT = torch::headeronly::ScalarType::Int;
constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
static constexpr int MXFP4_BLOCK_SIZE = 32;
static void validate_mxfp4_experts_quant_inputs(
torch::stable::Tensor const& output,
torch::stable::Tensor const& output_scale,
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
int64_t n_experts, int64_t m_topk, int64_t k) {
CHECK_INPUT(output, "output");
CHECK_INPUT(output_scale, "output_scale");
CHECK_INPUT(input, "input");
CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
STD_TORCH_CHECK(output.dim() == 2);
STD_TORCH_CHECK(output_scale.dim() == 2);
STD_TORCH_CHECK(input.dim() == 2);
STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
// output is uint8 (two mxfp4 values packed into one uint8)
// output_scale is int32 (four E8M0 values packed into one int32)
STD_TORCH_CHECK(output.scalar_type() == UINT8);
STD_TORCH_CHECK(output_scale.scalar_type() == INT);
STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
STD_TORCH_CHECK(output.size(0) == m_topk);
STD_TORCH_CHECK(output.size(1) == k / 2);
int scales_k = k / MXFP4_BLOCK_SIZE;
// K-dimension scale columns padded to a multiple of 4 for swizzle layout
int padded_k = (scales_k + (4 - 1)) / 4 * 4;
// 4 = 4 E8M0 values packed into one int32
STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
}
void mxfp4_experts_quant(
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
int64_t n_experts) {
auto m_topk = input.size(0);
auto k = input.size(1);
validate_mxfp4_experts_quant_inputs(
output, output_scale, input, input_offset_by_experts,
output_scale_offset_by_experts, n_experts, m_topk, k);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
input_offset_by_experts.data_ptr(),
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
stream);
});
}
void silu_and_mul_mxfp4_experts_quant(
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
torch::stable::Tensor const& input,
torch::stable::Tensor const& input_offset_by_experts,
torch::stable::Tensor const& output_scale_offset_by_experts,
int64_t n_experts) {
auto m_topk = input.size(0);
auto k_times_2 = input.size(1);
STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
auto k = k_times_2 / 2;
validate_mxfp4_experts_quant_inputs(
output, output_scale, input, input_offset_by_experts,
output_scale_offset_by_experts, n_experts, m_topk, k);
const torch::stable::accelerator::DeviceGuard device_guard(
input.get_device_index());
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
input_offset_by_experts.data_ptr(),
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
stream);
});
}
// Registered here (not torch_bindings.cpp) because VLLM_GPU_FLAGS is applied
// only under COMPILE_LANGUAGE:CUDA, so ENABLE_NVFP4_SM100 is invisible to
// .cpp files and cannot gate the registration from there.
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
m.impl("silu_and_mul_mxfp4_experts_quant",
TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
}
+21 -3
View File
@@ -116,6 +116,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
" Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
// cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
ops.def(
"cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
" Tensor a_blockscale, Tensor b_blockscales,"
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
// Compute NVFP4 block quantized tensor.
ops.def(
"scaled_fp4_quant(Tensor input,"
@@ -149,6 +155,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
"Tensor output_scale_offset_by_experts) -> ()");
// Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
ops.def(
"mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
"Tensor input, Tensor input_offset_by_experts,"
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
// Fused SiLU+Mul+MXFP4 experts quantization.
ops.def(
"silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
"output_scale,"
"Tensor input, Tensor input_offset_by_experts,"
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
// Fused SiLU+Mul+NVFP4 quantization.
ops.def(
"silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
@@ -233,9 +252,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("silu_and_mul_scaled_fp4_experts_quant",
TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
// W4A8 ops: impl registrations are in the source files
// (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
// mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
// W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
#endif
}
+19 -2
View File
@@ -126,7 +126,9 @@ __launch_bounds__(TPB) __global__
{
const int idx = thread_row_offset + ii;
const float val = toFloat(input[idx]);
const float softmax_val = expf(val - float_max) * normalizing_factor;
float softmax_val = expf(val - float_max) * normalizing_factor;
// Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f;
output[idx] = softmax_val;
}
}
@@ -147,7 +149,9 @@ __launch_bounds__(TPB) __global__
{
const int idx = thread_row_offset + ii;
const float val = toFloat(input[idx]);
const float sigmoid_val = 1.0f / (1.0f + __expf(-val));
float sigmoid_val = 1.0f / (1.0f + __expf(-val));
// Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f;
output[idx] = sigmoid_val;
}
}
@@ -442,6 +446,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
}
}
// Fix: clamp NaN/Inf values to 0 to prevent duplicate expert IDs.
// NaN gating (from degenerate hidden states in CUDA graph padding) causes
// softmax to produce all-NaN, which makes the argmax loop always pick
// expert 0 for every top-k slot, producing duplicate expert IDs that
// crash FlashInfer's three-step MoE sort.
// With 0s, the argmax uses index tie-breaking to pick [0,1,2,...,k-1].
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) {
row_chunk[ii] = 0.f;
}
}
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
// If bias is not null, use biased value for selection
+275
View File
@@ -0,0 +1,275 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
// NVFP4 KV cache store kernel.
// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them
// into the paged KV cache.
//
// Per page layout: [K_data | K_scale | V_data | V_scale]
// Both data and scale regions are contiguous per head, enabling direct
// TMA descriptor use.
//
// Reuses device functions from nvfp4_utils.cuh:
// - cvt_warp_fp16_to_fp4() for bf16 → fp4 quantization + block scale
// - pack_fp4() for packing float pairs to fp4
// - reciprocal_approximate_ftz() for fast reciprocal
#define NVFP4_ENABLE_ELTS16 1
#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include "dispatch_utils.h"
namespace vllm {
// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
// The swizzle pattern for HND layout is:
// [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
// where T = block_size (page_size), S = scale_dim = head_size // 16.
//
// For a linear (t, s) position, the swizzled position is:
// swizzled_t = (t / 4) * 4 + (s / (S / 4))
// swizzled_s = (s % (S / 4)) * 4 + (t % 4)
__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
int scale_dim) {
int s_group = scale_dim / 4;
int swizzled_t = (t / 4) * 4 + (s / s_group);
int swizzled_s = (s % s_group) * 4 + (t % 4);
return swizzled_t * scale_dim + swizzled_s;
}
// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
//
// Takes separate data and scale cache pointers for K and V.
// Within each KV side, data and scale are separate contiguous regions.
//
// Threading: one CUDA block per token, threads process heads and
// groups of 16 elements within each head.
template <typename scalar_t>
__global__ void reshape_and_cache_nvfp4_kernel(
const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size]
const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size]
uint8_t* __restrict__ key_data_cache, // data region for K
uint8_t* __restrict__ value_data_cache, // data region for V
uint8_t* __restrict__ key_scale_cache, // scale region for K
uint8_t* __restrict__ value_scale_cache, // scale region for V
const int64_t* __restrict__ slot_mapping, // [num_actual_tokens]
const float* __restrict__ k_scale_ptr, // pointer to checkpoint k_scale
const float* __restrict__ v_scale_ptr, // pointer to checkpoint v_scale
const int64_t key_stride, // key.stride(0) in elements
const int64_t value_stride, // value.stride(0) in elements
const int num_heads, const int head_size, const int block_size,
const int64_t data_block_stride, // data cache stride for dim 0
const int64_t data_head_stride, // data cache stride for heads
const int64_t data_block_offset_stride, // data cache stride for tokens
const int64_t scale_block_stride, // scale cache stride for dim 0
const int64_t scale_head_stride, // scale cache stride for heads
const int64_t scale_block_offset_stride // scale cache stride for tokens
) {
using CudaType = typename CUDATypeConverter<scalar_t>::Type;
using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD; // 16 or 8
static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;
const int64_t token_idx = blockIdx.x;
const int64_t slot_idx = slot_mapping[token_idx];
if (slot_idx < 0) return;
const int64_t block_idx = slot_idx / block_size;
const int block_offset = static_cast<int>(slot_idx % block_size);
const int scale_dim = head_size / 16;
const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
const int total_groups = num_heads * groups_per_head;
const int tid = threadIdx.x;
const int num_thread_groups = blockDim.x / THREADS_PER_SF;
const int tg_id = tid / THREADS_PER_SF;
const int tg_lane = tid % THREADS_PER_SF;
// Process both K (kv=0) and V (kv=1)
#pragma unroll
for (int kv = 0; kv < 2; kv++) {
const scalar_t* __restrict__ src = (kv == 0) ? key : value;
const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
uint8_t* __restrict__ data_cache =
(kv == 0) ? key_data_cache : value_data_cache;
uint8_t* __restrict__ sc_cache =
(kv == 0) ? key_scale_cache : value_scale_cache;
// Source pointer for this token (use actual stride, not assumed contiguous)
const CudaType* __restrict__ token_src =
reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;
// Destination bases in data and scale caches for this token's block
uint8_t* __restrict__ data_block =
data_cache + block_idx * data_block_stride;
uint8_t* __restrict__ scale_block =
sc_cache + block_idx * scale_block_stride;
for (int g = tg_id; g < total_groups; g += num_thread_groups) {
const int head = g / groups_per_head;
const int group_in_head = g % groups_per_head;
// Load 16 (or 8) bf16 elements from source
PVec in_vec;
const CudaType* __restrict__ src_ptr =
token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
tg_lane * ELTS;
#pragma unroll
for (int i = 0; i < ELTS / 2; i++) {
in_vec.elts[i] = reinterpret_cast<
const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
}
// Quantize: produces packed fp4 and writes scale factor.
uint8_t sf_val;
uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
in_vec, global_scale, sf_out_ptr);
// Write packed FP4 data to data cache
uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
block_offset * data_block_offset_stride;
#if CVT_FP4_PACK16
{
// 16 elements → 8 bytes (u32x2)
int data_byte_offset = group_in_head * 8;
reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
(uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
}
#else
{
// 8 elements → 4 bytes (uint32_t)
int data_byte_offset =
group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
}
#endif
// Write block scale to scale cache.
// K (kv==0): linear layout (no swizzle).
// V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel.
if (sf_out_ptr != nullptr) {
int scale_idx = group_in_head;
uint8_t* __restrict__ scale_dst;
if (kv == 0) {
scale_dst = scale_block + head * scale_head_stride +
block_offset * scale_block_offset_stride + scale_idx;
} else {
int swizzled_offset =
swizzle_scale_offset(block_offset, scale_idx, scale_dim);
int swizzled_t = swizzled_offset / scale_dim;
int swizzled_s = swizzled_offset % scale_dim;
scale_dst = scale_block + head * scale_head_stride +
swizzled_t * scale_block_offset_stride + swizzled_s;
}
*scale_dst = sf_val;
}
}
}
}
} // namespace vllm
// Non-template entry point callable from cache_kernels.cu.
// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
// Each KV side contains both data and scale:
// page = [K_data | K_scale | V_data | V_scale]
void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
torch::Tensor& k_scale,
torch::Tensor& v_scale) {
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
int data_dim = head_size / 2;
int scale_dim = head_size / 16;
int full_dim = data_dim + scale_dim;
// key_cache is kv_cache[:, 0] with shape
// [num_blocks, block_size, num_heads, full_dim] in logical order.
// Strides encode the physical layout (HND or NHD).
TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
TORCH_CHECK(key_cache.size(3) == full_dim,
"key_cache last dim must be data_dim + scale_dim, got ",
key_cache.size(3), " expected ", full_dim);
int block_size = key_cache.size(1);
TORCH_CHECK(head_size % 16 == 0,
"head_size must be divisible by 16 for NVFP4 KV cache");
TORCH_CHECK(block_size % 4 == 0,
"block_size must be divisible by 4 for NVFP4 KV cache swizzle");
// Detect physical layout from strides (based on full_dim).
// HND: head stride > block_offset stride.
bool is_hnd = key_cache.stride(2) > key_cache.stride(1);
int64_t data_block_stride = key_cache.stride(0); // page_bytes
int64_t data_head_stride, data_block_offset_stride;
if (is_hnd) {
data_head_stride = (int64_t)block_size * data_dim;
data_block_offset_stride = data_dim;
} else {
data_head_stride = data_dim;
data_block_offset_stride = (int64_t)num_heads * data_dim;
}
// Page layout: [K_data | K_scale | V_data | V_scale]
// Scale follows data within each KV side.
int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
// Scale strides: same page stride, inner strides from layout.
int64_t scale_block_stride = data_block_stride;
int64_t scale_head_stride, scale_block_offset_stride;
if (is_hnd) {
scale_head_stride = (int64_t)block_size * scale_dim;
scale_block_offset_stride = scale_dim;
} else {
scale_head_stride = scale_dim;
scale_block_offset_stride = (int64_t)num_heads * scale_dim;
}
const float* k_scale_ptr = k_scale.data_ptr<float>();
const float* v_scale_ptr = v_scale.data_ptr<float>();
int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
int total_groups = num_heads * groups_per_head;
constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
num_threads = ((num_threads + 31) / 32) * 32;
dim3 grid(num_tokens);
dim3 block(num_threads);
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_REDUCED_FLOATING_TYPES(
key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
<<<grid, block, 0, stream>>>(
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
key_scale_ptr, value_scale_ptr,
slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
key.stride(0), value.stride(0), num_heads, head_size,
block_size, data_block_stride, data_head_stride,
data_block_offset_stride, scale_block_stride, scale_head_stride,
scale_block_offset_stride);
});
}
+2 -1
View File
@@ -173,7 +173,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements/dev.txt && \
uv pip install -r requirements/lint.txt && \
uv pip install -r requirements/test/cpu.txt && \
pre-commit install --hook-type pre-commit --hook-type commit-msg
ENTRYPOINT ["bash"]
+7 -8
View File
@@ -192,9 +192,10 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
FROM base AS build_deep
ARG ROCSHMEM_BRANCH="ba0bf0f3"
ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
ARG DEEPEP_BRANCH="e84464ec"
ARG DEEPEP_BRANCH="5d90af8b"
ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
ARG DEEPEP_NIC="cx7"
ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
ENV ROCSHMEM_DIR=/opt/rocshmem
RUN git clone ${ROCSHMEM_REPO} \
@@ -202,13 +203,11 @@ RUN git clone ${ROCSHMEM_REPO} \
&& git checkout ${ROCSHMEM_BRANCH} \
&& mkdir -p projects/rocshmem/build \
&& cd projects/rocshmem/build \
&& cmake .. \
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-DROCM_PATH=/opt/rocm \
-DCMAKE_POSITION_INDEPENDENT_CODE=ON \
-DUSE_EXTERNAL_MPI=OFF \
&& make -j \
&& make install
&& bash ../scripts/build_configs/all_backends \
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-DROCM_PATH=/opt/rocm \
-DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-DUSE_EXTERNAL_MPI=OFF
# Build DeepEP wheel.
# DeepEP looks for rocshmem at ROCSHMEM_DIR.
+3 -3
View File
@@ -193,7 +193,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics:
The API server takes care of basic audio I/O and optional chunking before building prompts:
- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `librosa`.
- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`.
- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`.
- Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words.
@@ -206,8 +206,8 @@ Relevant server logic:
async def _preprocess_speech_to_text(...):
language = self.model_cls.validate_language(request.language)
...
y, sr = librosa.load(bytes_, sr=self.asr_config.sample_rate)
duration = librosa.get_duration(y=y, sr=sr)
y, sr = load_audio(bytes_, sr=self.asr_config.sample_rate)
duration = get_audio_duration(y=y, sr=sr)
do_split_audio = (self.asr_config.allow_audio_chunking
and duration > self.asr_config.max_audio_clip_s)
chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
+2 -2
View File
@@ -206,8 +206,8 @@ Both the `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_cont
used to profile a section of code.
!!! note
The legacy import paths `vllm.utils.cprofile` and `vllm.utils.cprofile_context` are deprecated.
Please use `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_context` instead.
The `vllm.utils.profiling` helpers are deprecated and will be removed in
`v0.21`. Please use Python's `cProfile` module directly instead.
### Example usage - decorator
+2
View File
@@ -106,6 +106,7 @@ Priority is **1 = highest** (tried first).
| 2 | `FLASH_ATTN` |
| 3 | `TRITON_ATTN` |
| 4 | `FLEX_ATTENTION` |
| 5 | `TURBOQUANT` |
**Ampere/Hopper (SM 8.x-9.x):**
@@ -115,6 +116,7 @@ Priority is **1 = highest** (tried first).
| 2 | `FLASHINFER` |
| 3 | `TRITON_ATTN` |
| 4 | `FLEX_ATTENTION` |
| 5 | `TURBOQUANT` |
### MLA Attention (DeepSeek-style)
+34
View File
@@ -21,6 +21,7 @@ or just on the low or high end.
| Fusion | `PassConfig` flag | Fused operations | Default at | E2E Speedup | Fullgraph | `num_tokens` |
| ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
| [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low |
| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm) | `fuse_minimax_qk_norm` | Q/K variance all-reduce → Q/K RMSNorm | Off by default | 2-3% | No | Low |
| [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always |
| [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always |
| [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low |
@@ -40,6 +41,7 @@ The table below lists the quantization schemes supported by each fusion on each
| Fusion | SM100 (Blackwell) | SM90 (Hopper) | SM89 (Ada) | SM80 (Ampere) | ROCm |
| ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
| `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* |
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
@@ -54,6 +56,9 @@ The table below lists the quantization schemes supported by each fusion on each
fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
for per-backend details.
\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
`enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
@@ -184,6 +189,35 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
- Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)
### MiniMax QK Norm (`fuse_minimax_qk_norm`)
!!! info
This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
optimization level.
**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
per-token Q/K variances before applying RMS normalization to Q and K.
This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
Example:
```bash
vllm serve MiniMaxAI/MiniMax-M2.5 \
--tensor-parallel-size 4 \
--compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
```
**Code locations.**
- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
### Sequence Parallelism (`enable_sp`)
**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,
+2 -2
View File
@@ -42,7 +42,7 @@ These are documented under [Inferencing and Serving -> Production Metrics](../us
### Grafana Dashboard
vLLM also provides [a reference example](../../examples/online_serving/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
vLLM also provides [a reference example](../../examples/observability/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
The subset of metrics exposed in the Grafana dashboard gives us an indication of which metrics are especially important:
@@ -657,7 +657,7 @@ vLLM has support for OpenTelemetry tracing:
- Added by <https://github.com/vllm-project/vllm/pull/4687> and reinstated by <https://github.com/vllm-project/vllm/pull/20372>
- Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces`
- [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/)
- [User-facing docs](../../examples/online_serving/opentelemetry/README.md)
- [User-facing docs](../../examples/observability/opentelemetry/README.md)
- [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f)
- [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview)
+3 -3
View File
@@ -300,12 +300,12 @@ Full example: [examples/offline_inference/audio_language.py](../../examples/offl
Speech-to-text models like Whisper have a maximum audio length they can process (typically 30 seconds). For longer audio files, vLLM provides a utility to intelligently split audio into chunks at quiet points to minimize cutting through speech.
```python
import librosa
from vllm import LLM, SamplingParams
from vllm.multimodal.audio import split_audio
from vllm.multimodal.media.audio import load_audio
# Load long audio file
audio, sr = librosa.load("long_audio.wav", sr=16000)
audio, sr = load_audio("long_audio.wav", sr=16000)
# Split into chunks at low-energy (quiet) regions
chunks = split_audio(
@@ -832,7 +832,7 @@ Then, you can use the OpenAI client as follows:
base_url=openai_api_base,
)
# Any format supported by librosa is supported
# Any format supported by soundfile/PyAV is supported
audio_url = AudioAsset("winning_call").url
audio_base64 = encode_base64_content_from_url(audio_url)
+9 -6
View File
@@ -14,6 +14,7 @@ Sorted alphabetically by GitHub handle:
- [@aarnphm](https://github.com/aarnphm): Structured output
- [@alexm-redhat](https://github.com/alexm-redhat): Performance
- [@ApostaC](https://github.com/ApostaC): Connectors, offloading
- [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
- [@benchislett](https://github.com/benchislett): Engine core and spec decode
- [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
- [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
@@ -31,6 +32,7 @@ Sorted alphabetically by GitHub handle:
- [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance
- [@luccafong](https://github.com/luccafong): Llama models, speculative decoding, distributed
- [@markmc](https://github.com/markmc): Observability
- [@MatthewBonanni](https://github.com/MatthewBonanni): Kernels and performance
- [@mgoin](https://github.com/mgoin): Quantization and performance
- [@NickLucche](https://github.com/NickLucche): KV connector
- [@njhill](https://github.com/njhill): Distributed, API server, engine core
@@ -41,6 +43,7 @@ Sorted alphabetically by GitHub handle:
- [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg
- [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support
- [@russellb](https://github.com/russellb): Structured output, engine core, security
- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser
- [@sighingnow](https://github.com/sighingnow): Qwen models, new model support
- [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community
- [@tdoublep](https://github.com/tdoublep): State space models
@@ -86,7 +89,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints
- @robertgshaw2-redhat, @njhill, @russellb
- ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni
- KV Connector: Connector interface and implementation for KV cache offload and transfer
- @robertgshaw2-redhat, @njhill, @KuntaiDu, @NickLucche, @ApostaC
- Distributed, Parallelism, Process Management: Process launchers managing each worker, and assign them to the right DP/TP/PP/EP ranks
@@ -105,7 +108,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- Custom Layers: Utility layers in vLLM such as rotary embedding and rms norms
- @ProExpertProg
- Attention: Attention interface for paged attention
- @WoosukKwon, @LucasWilkinson, @heheda12345
- @WoosukKwon, @LucasWilkinson, @heheda12345, @MatthewBonanni
- FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB
- @tlrmchlsmth
- Quantization: Various quantization config, weight loading, and kernel.
@@ -119,7 +122,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- State space models: The state space models implementation in vLLM
- @tdoublep, @tlrmchlsmth
- Reasoning and tool calling parsers
- @chaunceyjiang, @aarnphm
- @chaunceyjiang, @aarnphm, @sfeng33, @bbrowning
### Entrypoints
@@ -133,7 +136,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
### Features
- Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP.
- @WoosukKwon, @benchislett, @luccafong
- @WoosukKwon, @benchislett, @luccafong, @MatthewBonanni
- Structured Output: The structured output implementation
- @russellb, @aarnphm
- RL: The RL related features such as collective rpc, sleep mode, etc.
@@ -153,8 +156,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
### External Kernels Integration
- FlashAttention: @LucasWilkinson
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon
- FlashAttention: @LucasWilkinson, @MatthewBonanni
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni
- Blackwell Kernels: @mgoin, @yewentao256
- DeepEP/DeepGEMM: @mgoin, @yewentao256
+2 -2
View File
@@ -9,14 +9,14 @@
- Online APIs:
- Pooling API (`/pooling`)
The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs a embedding for each token.
The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs an embedding for each token.
Many embedding models support both (sequence) embedding and token embedding. For further details on (sequence) embedding, please refer to [this page](embed.md).
!!! note
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (embed) is not
what you want, you need to manually specify it via via `PoolerConfig(task="token_embed")` offline or
what you want, you need to manually specify it via `PoolerConfig(task="token_embed")` offline or
`--pooler-config.task token_embed` online.
## Typical Use Cases
+26
View File
@@ -400,6 +400,7 @@ th {
| `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ |
| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ |
| `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ |
| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
| `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ |
| `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ |
@@ -554,6 +555,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |
| `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>E+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |
| `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I<sup>+</sup> + V + A<sup>*</sup> | `google/gemma-4-E2B-it`, etc. | | ✅︎ |
| `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ |
| `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
@@ -633,6 +635,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
<sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
<sup>+</sup> Multiple items can be inputted per text prompt for this modality.
<sup>*</sup> Only specific variants of the model support this modality (see notes below).</br>
!!! note
`Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
@@ -643,6 +646,11 @@ Some models are supported only via the [Transformers modeling backend](#transfor
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
- There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
!!! note
For `Gemma4ForConditionalGeneration`:
- audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants.
- The model does not ingest videos directly. However, vLLMs Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model.
!!! note
For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
@@ -674,6 +682,24 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
!!! note
`VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed.
#### Realtime Transcription
Speech models that support streaming transcription via the
[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api)
WebSocket endpoint.
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | |
| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | |
!!! note
`VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`.
`Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`.
You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'`
when serving.
## Pooling Models
See [this page](pooling_models/README.md) for more information on how to use pooling models.
+1 -1
View File
@@ -60,7 +60,7 @@ We currently support the following OpenAI APIs:
- [Translation API](#translations-api) (`/v1/audio/translations`)
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
- [Realtime API](#realtime-api) (`/v1/realtime`)
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription).
In addition, we have the following custom APIs:
@@ -74,8 +74,8 @@ percli apply -f perses/performance_statistics.yaml
For detailed deployment instructions and platform-specific options, see:
- **[Grafana Documentation](./grafana)** - JSON dashboards, operator usage, manual import
- **[Perses Documentation](./perses)** - YAML specs, CLI usage, operator wrapping
- **[Grafana Documentation](grafana)** - JSON dashboards, operator usage, manual import
- **[Perses Documentation](perses)** - YAML specs, CLI usage, operator wrapping
## Contributing
@@ -0,0 +1,117 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Disaggregated multimodal serving: render → generate round-trip.
Demonstrates the two-phase disaggregated flow:
1. /v1/chat/completions/render preprocesses a multimodal chat request
into token IDs and serialized tensor features.
2. /inference/v1/generate runs inference on the preprocessed tokens.
The render response is passed *directly* to generate with only
``sampling_params`` added, showing that the two endpoints compose with
zero client-side transformation.
Launch the server first:
vllm serve Qwen/Qwen3-VL-2B-Instruct \
--dtype bfloat16 --max-model-len 4096 --enforce-eager
Then run this script:
python example_mm_serve.py
"""
import io
import pybase64 as base64
import requests
from PIL import Image
from transformers import AutoTokenizer
BASE_URL = "http://localhost:8000"
MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
def make_data_url(image: Image.Image) -> str:
"""Encode a PIL image as a base64 data URL."""
buf = io.BytesIO()
image.save(buf, format="PNG")
b64 = base64.b64encode(buf.getvalue()).decode()
return f"data:image/png;base64,{b64}"
def main():
# -- Step 1: Create a test image (solid red) -------------------------
image = Image.new("RGB", (224, 224), color=(255, 0, 0))
data_url = make_data_url(image)
print("Created 224x224 red test image")
# -- Step 2: Render (preprocess) -------------------------------------
render_payload = {
"model": MODEL_NAME,
"messages": [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": data_url}},
{
"type": "text",
"text": "What color is this image? Answer in one word.",
},
],
}
],
}
print("\n--- Render ---")
render_resp = requests.post(
f"{BASE_URL}/v1/chat/completions/render", json=render_payload
)
render_resp.raise_for_status()
render_data = render_resp.json()
print(f"Response keys: {list(render_data.keys())}")
print(f"Number of token_ids: {len(render_data['token_ids'])}")
features = render_data.get("features")
if features and features.get("kwargs_data"):
print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}")
for modality, items in features["kwargs_data"].items():
print(
f" {modality}: {len(items)} item(s), "
f"first item type: {type(items[0])} length: {len(items[0])}"
if items
else "First item: (empty)"
)
else:
print("WARNING: no kwargs_data in render response")
# -- Step 3: Generate (inference) ------------------------------------
# Pass the render output directly — only add sampling_params.
generate_payload = render_data
generate_payload["sampling_params"] = {
"max_tokens": 20,
"temperature": 0.0,
}
print("\n--- Generate ---")
gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload)
gen_resp.raise_for_status()
gen_data = gen_resp.json()
# -- Step 4: Decode & print ------------------------------------------
output_ids = gen_data["choices"][0]["token_ids"]
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
text = tokenizer.decode(output_ids, skip_special_tokens=True)
print(f"Output token count: {len(output_ids)}")
print(f"Generated text: {text!r}")
if "red" in text.lower():
print("\nModel correctly identified the red image.")
else:
print(f"\nWARNING: Expected 'red' in output, got: {text!r}")
if __name__ == "__main__":
main()
@@ -267,7 +267,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
{
"type": "input_audio",
"input_audio": {
# Any format supported by librosa is supported
# Any format supported by soundfile/PyAV is supported
"data": audio_base64,
"format": "wav",
},
@@ -292,7 +292,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
{
"type": "audio_url",
"audio_url": {
# Any format supported by librosa is supported
# Any format supported by soundfile/PyAV is supported
"url": audio_url
},
},
@@ -316,7 +316,7 @@ def run_audio(model: str, max_completion_tokens: int) -> None:
{
"type": "audio_url",
"audio_url": {
# Any format supported by librosa is supported
# Any format supported by soundfile/PyAV is supported
"url": f"data:audio/ogg;base64,{audio_base64}"
},
},
@@ -12,7 +12,6 @@ model, for example:
Requirements:
- vllm with audio support
- websockets
- librosa
- numpy
The script:
@@ -26,12 +25,12 @@ import argparse
import asyncio
import json
import librosa
import numpy as np
import pybase64 as base64
import websockets
from vllm.assets.audio import AudioAsset
from vllm.multimodal.media.audio import load_audio
def audio_to_pcm16_base64(audio_path: str) -> str:
@@ -39,7 +38,7 @@ def audio_to_pcm16_base64(audio_path: str) -> str:
Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
"""
# Load audio and resample to 16kHz mono
audio, _ = librosa.load(audio_path, sr=16000, mono=True)
audio, _ = load_audio(audio_path, sr=16000, mono=True)
# Convert to PCM16
pcm16 = (audio * 32767).astype(np.int16)
# Encode as base64
+1
View File
@@ -170,6 +170,7 @@ eles = "eles"
datas = "datas"
ser = "ser"
ure = "ure"
VALU = "VALU"
# Walsh-Hadamard Transform
wht = "wht"
WHT = "WHT"
-2
View File
@@ -32,9 +32,7 @@ pyzmq >= 25.0.0
msgspec
gguf >= 0.17.0
mistral_common[image] >= 1.11.0
av # required for audio in video IO
opencv-python-headless >= 4.13.0 # required for video IO
soundfile # required for audio IO
pyyaml
six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
+1 -3
View File
@@ -20,6 +20,4 @@ conch-triton-kernels==1.2.1
timm>=1.0.17
# amd-quark: required for Quark quantization on ROCm
# To be consistent with test_quark.py
amd-quark>=0.8.99
# Required for faster safetensors model loading
fastsafetensors >= 0.2.2
amd-quark>=0.8.99
+1 -1
View File
@@ -55,7 +55,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
numba==0.61.2 # Required for N-gram speculative decoding
numpy
runai-model-streamer[s3,gcs,azure]==0.15.7
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
instanttensor>=0.1.5
pydantic>=2.12 # 2.11 leads to error on python 3.13
decord==0.6.0
+3 -8
View File
@@ -76,9 +76,7 @@ attrs==26.1.0
audioread==3.0.1
# via librosa
av==16.1.0
# via
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
# via -r requirements/test/rocm.in
azure-core==1.39.0
# via
# azure-identity
@@ -277,10 +275,8 @@ fastar==0.10.0
# via fastapi-cloud-cli
fastparquet==2026.3.0
# via genai-perf
fastsafetensors==0.2.2
# via
# -c requirements/rocm.txt
# -r requirements/test/rocm.in
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
# via -r requirements/test/rocm.in
filelock==3.25.2
# via
# -c requirements/common.txt
@@ -1333,7 +1329,6 @@ sortedcontainers==2.4.0
# via hypothesis
soundfile==0.13.1
# via
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
# genai-perf
# librosa
+5 -1
View File
@@ -1085,14 +1085,18 @@ setup(
install_requires=get_requirements(),
extras_require={
# AMD Zen CPU optimizations via zentorch
"zen": ["zentorch"],
"zen": [
"zentorch-weekly==5.2.1.dev20260408"
], # Zentorch has weekly releases. This pulls the known-good version.
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.10.1"],
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
"instanttensor": ["instanttensor >= 0.1.5"],
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
"audio": [
"av",
"scipy",
"soundfile",
"mistral_common[audio]",
], # Required for audio processing
"video": [], # Kept for backwards compatibility
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import vllm.config
from tests.compile.backend import TestBackend
from vllm.compilation.passes.vllm_inductor_pass import (
VllmFusionPatternMatcherPass,
VllmPatternMatcherPass,
VllmPatternReplacement,
)
from vllm.config import CompilationConfig, CompilationMode, VllmConfig
from vllm.platforms import current_platform
class ReluToAbsPattern(VllmPatternReplacement):
"""Replaces relu(x) with abs(x) — a minimal test fixture."""
@property
def pattern(self):
def _pattern(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.relu.default(x)
return _pattern
@property
def replacement(self):
def _replacement(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.abs.default(x)
return _replacement
def get_inputs(self) -> list[torch.Tensor]:
return [self.empty_fp32(4)]
class ExpToSqrtPattern(VllmPatternReplacement):
"""A second distinct pattern type — used to test uuid differentiation."""
@property
def pattern(self):
def _pattern(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.exp.default(x)
return _pattern
@property
def replacement(self):
def _replacement(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.sqrt.default(x)
return _replacement
def get_inputs(self) -> list[torch.Tensor]:
return [self.empty_fp32(4)]
class ReluFusionPass(VllmFusionPatternMatcherPass):
def __init__(self, config: VllmConfig) -> None:
super().__init__(config, "test_relu_fusion")
self.register(ReluToAbsPattern())
class TwoPatternFusionPass(VllmFusionPatternMatcherPass):
def __init__(self, config: VllmConfig) -> None:
super().__init__(config, "test_two_pattern_fusion")
self.register(ReluToAbsPattern())
self.register(ExpToSqrtPattern())
@pytest.fixture
def vllm_config():
return VllmConfig(
compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE),
)
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
def test_register_tracks_patterns(vllm_config):
"""register() appends each VllmPatternReplacement to _pattern_replacements."""
with vllm.config.set_current_vllm_config(vllm_config):
single = ReluFusionPass(vllm_config)
two = TwoPatternFusionPass(vllm_config)
assert len(single._pattern_replacements) == 1
assert len(two._pattern_replacements) == 2
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
def test_uuid_stable(vllm_config):
"""Two instances of the same pass class produce identical uuids."""
with vllm.config.set_current_vllm_config(vllm_config):
p1 = ReluFusionPass(vllm_config)
p2 = ReluFusionPass(vllm_config)
p3 = TwoPatternFusionPass(vllm_config)
assert p1.uuid() == p2.uuid()
assert p1.uuid() != p3.uuid()
assert p2.uuid() != p3.uuid()
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
@pytest.mark.parametrize("N", [1, 2, 4])
def test_matched_count_and_match_table(vllm_config, N):
"""matched_count and match_table reflect the number of matched patterns."""
class Model(torch.nn.Module):
def forward(self, *inputs):
# N independent relus
return sum(torch.relu(x) for x in inputs)
with vllm.config.set_current_vllm_config(vllm_config):
torch.set_default_device("cuda")
torch.set_default_dtype(torch.float32)
fusion_pass = ReluFusionPass(vllm_config)
backend = TestBackend(fusion_pass)
model = torch.compile(Model(), backend=backend)
inputs = [torch.rand(8) for _ in range(N)]
model(*inputs)
assert fusion_pass.matched_count == N
assert VllmPatternMatcherPass.match_table["test_relu_fusion"] >= N
@@ -38,6 +38,8 @@ llm = LLM(
distributed_executor_backend="external_launcher",
gpu_memory_utilization=random.uniform(0.7, 0.9),
seed=0,
max_model_len=1024,
max_num_seqs=16,
)
outputs = llm.generate(prompts, sampling_params)
@@ -41,6 +41,7 @@ def create_mock_parallel_config(
config.rank = rank
config.world_size = world_size
config.data_parallel_rank = dp_rank
config.data_parallel_index = dp_rank
return config
@@ -283,6 +284,7 @@ def inference_receive_tensor(
parallel_config.rank = 0
parallel_config.world_size = 1
parallel_config.data_parallel_rank = 0
parallel_config.data_parallel_index = 0
engine = NCCLWeightTransferEngine(config, parallel_config)
@@ -666,6 +668,7 @@ def inference_receive_ipc_tensor(
parallel_config.rank = 0
parallel_config.world_size = 1
parallel_config.data_parallel_rank = 0
parallel_config.data_parallel_index = 0
engine = IPCWeightTransferEngine(config, parallel_config)
@@ -13,7 +13,6 @@ import io
import time
from statistics import mean, median
import librosa
import pytest
import soundfile
import torch
@@ -21,6 +20,7 @@ from datasets import load_dataset
from evaluate import load
from transformers.models.whisper.english_normalizer import EnglishTextNormalizer
from vllm.multimodal.audio import get_audio_duration
from vllm.tokenizers import get_tokenizer
from ....models.registry import HF_EXAMPLE_MODELS
@@ -84,7 +84,7 @@ async def process_dataset(model, client, data, concurrent_request):
trust_remote_code=model_info.trust_remote_code,
)
# Warmup call as the first `librosa.load` server-side is quite slow.
# Warmup call as the first `load_audio` server-side is quite slow.
audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"]
_ = await bound_transcribe(sem, client, tokenizer, (audio, sr), "")
@@ -118,7 +118,7 @@ def print_performance_metrics(results, total_time):
def add_duration(sample):
y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
sample["duration_ms"] = librosa.get_duration(y=y, sr=sr) * 1000
sample["duration_ms"] = get_audio_duration(y=y, sr=sr) * 1000
return sample
@@ -5,7 +5,6 @@ import asyncio
import json
import warnings
import librosa
import numpy as np
import pybase64 as base64
import pytest
@@ -14,6 +13,7 @@ import websockets
from tests.entrypoints.openai.conftest import add_attention_backend
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
from vllm.assets.audio import AudioAsset
from vllm.multimodal.media.audio import load_audio
# Increase engine iteration timeout for ROCm where first-use JIT compilation
# can exceed the default 60s, causing a silent deadlock in feed_tokens.
@@ -56,7 +56,7 @@ async def send_event(ws, event: dict) -> None:
def mary_had_lamb_audio_chunks() -> list[str]:
"""Audio split into ~1 second chunks for streaming."""
path = AudioAsset("mary_had_lamb").get_local_path()
audio, _ = librosa.load(str(path), sr=16000, mono=True)
audio, _ = load_audio(str(path), sr=16000, mono=True)
# Split into ~0.1 second chunks (1600 samples at 16kHz)
chunk_size = 1600
@@ -6,7 +6,6 @@ import asyncio
import io
import json
import librosa
import numpy as np
import openai
import pytest
@@ -14,6 +13,7 @@ import pytest_asyncio
import soundfile as sf
from tests.utils import RemoteOpenAIServer
from vllm.multimodal.media.audio import load_audio
from vllm.platforms import current_platform
MODEL_NAME = "openai/whisper-large-v3-turbo"
@@ -134,7 +134,7 @@ async def test_bad_requests(mary_had_lamb, whisper_client):
@pytest.mark.asyncio
async def test_long_audio_request(mary_had_lamb, whisper_client):
mary_had_lamb.seek(0)
audio, sr = librosa.load(mary_had_lamb)
audio, sr = load_audio(mary_had_lamb)
# Add small silence after each audio for repeatability in the split process
audio = np.pad(audio, (0, 1600))
repeated_audio = np.tile(audio, 10)
@@ -7,7 +7,6 @@ import io
import json
import httpx
import librosa
import numpy as np
import openai
import pytest
@@ -17,6 +16,7 @@ import soundfile as sf
from tests.entrypoints.openai.conftest import add_attention_backend
from tests.utils import RemoteOpenAIServer
from vllm.logger import init_logger
from vllm.multimodal.media.audio import load_audio
logger = init_logger(__name__)
@@ -264,7 +264,7 @@ async def test_long_audio_request(foscolo, client_and_model):
if model_name == "google/gemma-3n-E2B-it":
pytest.skip("Gemma3n does not support long audio requests")
foscolo.seek(0)
audio, sr = librosa.load(foscolo)
audio, sr = load_audio(foscolo)
repeated_audio = np.tile(audio, 2)
# Repeated audio to buffer
buffer = io.BytesIO()
+111
View File
@@ -0,0 +1,111 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Roundtrip tests for multimodal serde used by the disagg generate endpoint."""
import torch
from vllm.entrypoints.serve.disagg.mm_serde import (
decode_mm_kwargs_item,
encode_mm_kwargs_item,
)
from vllm.entrypoints.serve.disagg.protocol import (
MultiModalFeatures,
PlaceholderRangeInfo,
)
from vllm.multimodal.inputs import (
MultiModalBatchedField,
MultiModalFieldElem,
MultiModalFlatField,
MultiModalKwargsItem,
MultiModalSharedField,
)
def test_mm_kwargs_item_roundtrip():
"""Full roundtrip test with all three field types and multiple dtypes."""
e1 = MultiModalFieldElem(
data=torch.zeros(1000, dtype=torch.bfloat16),
field=MultiModalBatchedField(),
)
e2 = MultiModalFieldElem(
data=torch.ones(100, dtype=torch.int32),
field=MultiModalSharedField(batch_size=4),
)
e3 = MultiModalFieldElem(
data=torch.randn(20, dtype=torch.float32),
field=MultiModalFlatField(slices=[slice(0, 10), slice(10, 20)], dim=0),
)
item = MultiModalKwargsItem({"pixel_values": e1, "grid_thw": e2, "embeds": e3})
encoded = encode_mm_kwargs_item(item)
# Encoded result is a base64 string
assert isinstance(encoded, str)
decoded = decode_mm_kwargs_item(encoded)
assert set(decoded.keys()) == {"pixel_values", "grid_thw", "embeds"}
assert torch.equal(item["pixel_values"].data, decoded["pixel_values"].data)
assert torch.equal(item["grid_thw"].data, decoded["grid_thw"].data)
assert torch.equal(item["embeds"].data, decoded["embeds"].data)
assert isinstance(decoded["pixel_values"].field, MultiModalBatchedField)
assert isinstance(decoded["grid_thw"].field, MultiModalSharedField)
assert isinstance(decoded["embeds"].field, MultiModalFlatField)
def test_mm_kwargs_item_none_data():
"""Roundtrip with None data field."""
elem = MultiModalFieldElem(
data=None,
field=MultiModalSharedField(batch_size=2),
)
item = MultiModalKwargsItem({"empty": elem})
encoded = encode_mm_kwargs_item(item)
decoded = decode_mm_kwargs_item(encoded)
assert decoded["empty"].data is None
assert isinstance(decoded["empty"].field, MultiModalSharedField)
def test_mm_kwargs_item_nested_tensors():
"""Roundtrip with nested tensor data."""
nested = [torch.randn(3, 4), torch.randn(5, 4)]
elem = MultiModalFieldElem(
data=nested,
field=MultiModalBatchedField(),
)
item = MultiModalKwargsItem({"nested": elem})
encoded = encode_mm_kwargs_item(item)
decoded = decode_mm_kwargs_item(encoded)
decoded_data = decoded["nested"].data
assert len(decoded_data) == 2
assert torch.equal(nested[0], decoded_data[0])
assert torch.equal(nested[1], decoded_data[1])
def test_mm_features_with_kwargs_data():
"""Test that MultiModalFeatures can carry serialized tensor data."""
elem = MultiModalFieldElem(
data=torch.randn(5, 3, dtype=torch.float32),
field=MultiModalBatchedField(),
)
item = MultiModalKwargsItem({"pixel_values": elem})
encoded = encode_mm_kwargs_item(item)
features = MultiModalFeatures(
mm_hashes={"image": ["abc123"]},
mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=10)]},
kwargs_data={"image": [encoded]},
)
# JSON roundtrip
json_str = features.model_dump_json()
features2 = MultiModalFeatures.model_validate_json(json_str)
assert features2.mm_hashes == {"image": ["abc123"]}
assert features2.kwargs_data is not None
assert len(features2.kwargs_data["image"]) == 1
decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0])
assert torch.equal(elem.data, decoded["pixel_values"].data)
@@ -0,0 +1,158 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for multimodal features through the /inference/v1/generate endpoint.
Mirrors test_serving_tokens.py but exercises the multimodal piping
using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render ->
/generate -> /detokenize path. Intentionally avoids running the HF
processor in the pytest parent process to keep os.fork() in sibling
tests (e.g. test_weight_transfer_llm.py) deadlock-free.
"""
import os
import httpx
import pytest
import pytest_asyncio
from PIL import Image
from tests.utils import RemoteOpenAIServer
from vllm.multimodal.utils import encode_image_url
MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
GEN_ENDPOINT = "/inference/v1/generate"
RENDER_ENDPOINT = "/v1/chat/completions/render"
DETOKENIZE_ENDPOINT = "/detokenize"
@pytest.fixture(scope="module")
def test_image():
return Image.new("RGB", (224, 224), color=(255, 0, 0))
@pytest.fixture(scope="module")
def server():
args = [
"--dtype",
"bfloat16",
"--max-model-len",
"4096",
"--enforce-eager",
"--no-enable-prefix-caching",
]
envs = os.environ.copy()
envs["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=envs) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client(server: RemoteOpenAIServer):
transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
async with httpx.AsyncClient(
transport=transport,
base_url=server.url_root,
timeout=600,
headers=headers,
) as c:
yield c
@pytest.mark.asyncio
async def test_render_to_generate_roundtrip(client, test_image):
"""End-to-end: render a multimodal chat -> feed into generate -> decode.
All preprocessing and detokenization happens in the server subprocess;
the pytest parent never imports transformers or touches torch tensors.
"""
data_url = encode_image_url(test_image, format="PNG")
render_payload = {
"model": MODEL_NAME,
"messages": [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": data_url}},
{
"type": "text",
"text": "What color is this image? Answer in one word.",
},
],
}
],
}
render_resp = await client.post(RENDER_ENDPOINT, json=render_payload)
render_resp.raise_for_status()
render_data = render_resp.json()
# Validate render output structure: keys exist and values are non-empty
# and well-typed.
assert "token_ids" in render_data
assert isinstance(render_data["token_ids"], list)
assert len(render_data["token_ids"]) > 0
assert all(isinstance(t, int) for t in render_data["token_ids"])
assert "features" in render_data
features = render_data["features"]
assert features is not None
assert isinstance(features, dict)
assert "mm_hashes" in features
assert "image" in features["mm_hashes"]
image_hashes = features["mm_hashes"]["image"]
assert isinstance(image_hashes, list)
assert len(image_hashes) > 0
assert all(isinstance(h, str) and h for h in image_hashes)
assert "mm_placeholders" in features
assert "image" in features["mm_placeholders"]
image_placeholders = features["mm_placeholders"]["image"]
assert isinstance(image_placeholders, list)
assert len(image_placeholders) > 0
for p in image_placeholders:
assert isinstance(p.get("offset"), int)
assert isinstance(p.get("length"), int)
assert p["length"] > 0
assert "kwargs_data" in features
assert "image" in features["kwargs_data"]
assert len(features["kwargs_data"]["image"]) > 0
# Build generate request from render output
generate_payload = render_data
generate_payload["sampling_params"] = {
"max_tokens": 10,
"temperature": 0.0,
}
gen_resp = await client.post(GEN_ENDPOINT, json=generate_payload)
gen_resp.raise_for_status()
gen_data = gen_resp.json()
assert "choices" in gen_data
assert isinstance(gen_data["choices"], list)
assert len(gen_data["choices"]) >= 1
choice = gen_data["choices"][0]
assert "token_ids" in choice
assert isinstance(choice["token_ids"], list)
assert len(choice["token_ids"]) > 0
assert all(isinstance(t, int) for t in choice["token_ids"])
detok_resp = await client.post(
DETOKENIZE_ENDPOINT,
json={"model": MODEL_NAME, "tokens": choice["token_ids"]},
)
detok_resp.raise_for_status()
detok_data = detok_resp.json()
assert "prompt" in detok_data
text = detok_data["prompt"]
assert isinstance(text, str)
assert len(text) > 0
assert "red" in text.lower(), (
f"Expected model to identify the red image, got: {text!r}"
)
+99 -21
View File
@@ -10,7 +10,7 @@ from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
from vllm.utils.torch_utils import nvfp4_kv_cache_split_views, set_random_seed
COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
DTYPES = [torch.bfloat16, torch.float]
@@ -172,7 +172,7 @@ def test_reshape_and_cache(
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE + ["nvfp4"])
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@@ -202,6 +202,25 @@ def test_reshape_and_cache_flash(
if kv_scale_type == "attn_head" and implementation != "cuda":
pytest.skip("Only CUDA implementation supports attn_head scaling.")
if kv_cache_dtype == "nvfp4":
if not current_platform.has_device_capability(100):
pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).")
if implementation != "cuda":
pytest.skip("NVFP4 only supports CUDA implementation.")
if kv_scale_type != "tensor":
pytest.skip("NVFP4 only supports per-tensor scaling.")
if head_size % 16 != 0:
pytest.skip("NVFP4 requires head_size divisible by 16.")
if (head_size // 16) % 4 != 0:
pytest.skip(
"NVFP4 requires (head_size // 16) divisible by 4 "
"for 4x4 block scale swizzle."
)
if block_size % 4 != 0:
pytest.skip("NVFP4 requires block_size divisible by 4.")
if dtype not in (torch.float16, torch.bfloat16):
pytest.skip("NVFP4 quantization only supports fp16/bf16 input.")
# fp8 conversion requires continugous memory buffer. Reduce the number of
# blocks and tokens to consume less memory.
num_tokens = num_tokens // 2
@@ -229,7 +248,23 @@ def test_reshape_and_cache_flash(
del key_caches
del value_caches
if kv_scale_type == "tensor":
# For nvfp4, the factory returns kv[:, 0] and kv[:, 1] like all dtypes.
# Split views are still needed for dequant verification.
key_scale_cache = None
value_scale_cache = None
nvfp4_key_data = None
nvfp4_value_data = None
if kv_cache_dtype == "nvfp4":
(nvfp4_key_data,), (key_scale_cache,) = nvfp4_kv_cache_split_views(key_cache)
(nvfp4_value_data,), (value_scale_cache,) = nvfp4_kv_cache_split_views(
value_cache
)
if kv_cache_dtype == "nvfp4":
# Global scale = amax / 448 (per-tensor)
k_scale = (key.abs().amax() / 448.0).to(torch.float32)
v_scale = (value.abs().amax() / 448.0).to(torch.float32)
elif kv_scale_type == "tensor":
k_scale = (key.amax() / 64.0).to(torch.float32)
v_scale = (value.amax() / 64.0).to(torch.float32)
else: # "attn_head"
@@ -240,8 +275,9 @@ def test_reshape_and_cache_flash(
y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
return y.contiguous()
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
if kv_cache_dtype != "nvfp4":
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
def convert_fp8_local(output, input, scale, kv_dtype):
fp8_input = input.view(current_platform.fp8_dtype())
@@ -257,7 +293,7 @@ def test_reshape_and_cache_flash(
result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
output.copy_(result)
# Clone the KV caches.
# Clone the KV caches (for non-nvfp4, used as reference baseline).
if kv_cache_dtype == "fp8":
cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
@@ -265,25 +301,27 @@ def test_reshape_and_cache_flash(
convert_fp8_local(
cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype
)
else:
elif kv_cache_dtype != "nvfp4":
cloned_key_cache = key_cache_compact.clone()
cloned_value_cache = value_cache_compact.clone()
# Call the reshape_and_cache kernel.
if implementation == "cuda":
opcheck(
torch.ops._C_cache_ops.reshape_and_cache_flash,
(
key,
value,
key_cache,
value_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
v_scale,
),
cond=(head_size == HEAD_SIZES[0]),
)
if kv_cache_dtype != "nvfp4":
opcheck(
torch.ops._C_cache_ops.reshape_and_cache_flash,
(
key,
value,
key_cache,
value_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
v_scale,
),
cond=(head_size == HEAD_SIZES[0]),
)
ops.reshape_and_cache_flash(
key,
value,
@@ -309,6 +347,46 @@ def test_reshape_and_cache_flash(
k_scale,
v_scale,
)
if kv_cache_dtype == "nvfp4":
# Verify NVFP4 by dequantizing the entire cache and comparing
# the written positions against original bf16 values.
# Same pattern as FP8: dequant whole cache, then extract and compare.
from tests.kernels.quantization.nvfp4_utils import (
dequant_nvfp4_kv_cache,
)
def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
# data_cache: [N, T, H, data_dim] NHD (contiguous inner dims)
# scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims)
# Permute to HND layout for the dequant utility.
data_hnd = data_cache.permute(0, 2, 1, 3)
scale_hnd = scale_cache.permute(0, 2, 1, 3)
result_hnd = dequant_nvfp4_kv_cache(
data_hnd, scale_hnd, global_scale, head_size, block_size
)
return result_hnd.permute(0, 2, 1, 3) # back to [N, T, H, D]
result_key_cache = dequant_nvfp4_cache_nhd(
nvfp4_key_data, key_scale_cache, k_scale.item()
)
result_value_cache = dequant_nvfp4_cache_nhd(
nvfp4_value_data, value_scale_cache, v_scale.item()
)
# Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
num_slots = num_blocks * block_size
result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size)
result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size)
torch.testing.assert_close(
result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
)
torch.testing.assert_close(
result_value_flat[slot_mapping], value.float(), atol=1.5, rtol=0.5
)
return
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
+48 -4
View File
@@ -13,6 +13,11 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
selective_state_update,
)
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupSpec,
MambaSpec,
)
try:
import flashinfer.mamba # noqa: F401
@@ -22,22 +27,40 @@ except ImportError:
HAS_FLASHINFER = False
def _kv_cache_config_with_ssu(mamba_type: str = "mamba2") -> KVCacheConfig:
spec = MambaSpec(
block_size=16,
shapes=((16, 64),),
dtypes=(torch.float16,),
mamba_type=mamba_type,
)
return KVCacheConfig(
num_blocks=1,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(layer_names=["l0"], kv_cache_spec=spec)],
)
def test_default_backend_is_triton():
initialize_mamba_ssu_backend(MambaConfig())
initialize_mamba_ssu_backend(MambaConfig(), _kv_cache_config_with_ssu())
backend = get_mamba_ssu_backend()
assert isinstance(backend, TritonSSUBackend)
assert backend.name == "triton"
def test_explicit_triton_backend():
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
)
backend = get_mamba_ssu_backend()
assert isinstance(backend, TritonSSUBackend)
@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
def test_flashinfer_backend_init():
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.FLASHINFER), _kv_cache_config_with_ssu()
)
backend = get_mamba_ssu_backend()
assert isinstance(backend, FlashInferSSUBackend)
assert backend.name == "flashinfer"
@@ -53,6 +76,25 @@ def test_uninitialized_backend_raises():
mod._mamba_ssu_backend = old
@pytest.mark.parametrize(
"mamba_type", ["linear_attention", "gdn_attention", "short_conv"]
)
def test_init_is_noop_for_non_ssu_mamba_type(mamba_type):
import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
old = mod._mamba_ssu_backend
mod._mamba_ssu_backend = None
try:
initialize_mamba_ssu_backend(
MambaConfig(), _kv_cache_config_with_ssu(mamba_type)
)
assert mod._mamba_ssu_backend is None
with pytest.raises(RuntimeError, match="not been initialized"):
get_mamba_ssu_backend()
finally:
mod._mamba_ssu_backend = old
@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
def test_flashinfer_import_error():
with pytest.raises(ImportError, match="FlashInfer is required"):
@@ -61,7 +103,9 @@ def test_flashinfer_import_error():
def test_triton_basic_call():
set_random_seed(0)
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
)
device = "cuda"
batch_size = 2
dim = 64
+67
View File
@@ -135,3 +135,70 @@ def test_fused_topk_bias(
topk_weights_ref.to(torch.float32), topk_weights, atol=1e-2, rtol=1e-2
)
torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0)
@pytest.mark.skipif(
not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
)
@pytest.mark.parametrize("num_experts", [6, 8, 16])
@pytest.mark.parametrize("topk", [3, 4])
@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
@pytest.mark.parametrize("bad_value", [float("nan"), float("inf")])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
def test_fused_topk_nan_inf_clamp(
num_experts: int,
topk: int,
scoring_func: str,
bad_value: float,
dtype: torch.dtype,
):
"""Regression test for the NaN/Inf clamp in topk_softmax_kernels.cu.
Degenerate hidden states (e.g., from CUDA graph padding) can produce
NaN/Inf gating logits. Without the clamp, softmax/sigmoid outputs are
NaN and the argmax loop picks expert 0 for every top-k slot (since
"NaN > NaN" is false per IEEE 754), yielding duplicate expert IDs that
crash downstream MoE sort kernels. The fix clamps NaN/Inf to 0 before
argmax so index tie-breaking selects unique experts [0, 1, ..., k-1].
"""
torch.manual_seed(0)
num_tokens = 4
hidden_size = 1024
hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
# Row 0: all normal. Rows 1-3: fully poisoned with NaN or Inf.
gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
gating_output[1:, :] = bad_value
topk_weights, topk_ids, _ = fused_topk(
hidden_states=hidden_states,
gating_output=gating_output,
topk=topk,
renormalize=False,
scoring_func=scoring_func,
)
# Normal row must still match the torch reference.
ref_weights, ref_ids = torch_topk(
gating_output=gating_output[:1],
topk=topk,
renormalize=False,
scoring_func=scoring_func,
)
torch.testing.assert_close(
ref_weights.to(torch.float32), topk_weights[:1], atol=1e-2, rtol=1e-2
)
torch.testing.assert_close(ref_ids.to(torch.int32), topk_ids[:1], atol=0, rtol=0)
# Poisoned rows: IDs must be unique (no duplicates) and weights must be
# finite (no NaN/Inf propagation into downstream MoE kernels).
for row in range(1, num_tokens):
row_ids = topk_ids[row]
assert row_ids.unique().numel() == topk, (
f"Row {row} has duplicate expert IDs {row_ids.tolist()} "
f"(bad_value={bad_value}, scoring_func={scoring_func})"
)
assert torch.isfinite(topk_weights[row]).all(), (
f"Row {row} has non-finite weights {topk_weights[row].tolist()} "
f"(bad_value={bad_value}, scoring_func={scoring_func})"
)
+57
View File
@@ -0,0 +1,57 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm.model_executor.models.gemma4 import (
gemma4_fused_routing_kernel_triton,
gemma4_routing_function_torch,
)
def sort_by_id(w, ids):
order = ids.argsort(dim=-1)
return w.gather(1, order), ids.gather(1, order)
# Gemma4 MoE Model has context length of 250K
# the minus 1 is to ensure that edge cases are tested
@pytest.mark.parametrize("num_tokens", [1, 2, 2048, 250000])
@pytest.mark.parametrize("num_experts", [128]) # gemma4 moe experts
@pytest.mark.parametrize("topk", [8]) # gemma4 topk
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
def test_gemma4_routing_kernel_triton(
num_tokens: int,
num_experts: int,
topk: int,
dtype: torch.dtype,
):
torch.manual_seed(0)
gating = torch.randn(num_tokens, num_experts, dtype=dtype, device="cuda")
scales = torch.rand(num_experts, dtype=torch.float32, device="cuda")
ref_w, ref_ids = gemma4_routing_function_torch(gating, topk, scales)
tri_w, tri_ids = gemma4_fused_routing_kernel_triton(gating, topk, scales)
# Sort by expert id — to remove tie-breaking differences
ref_ws, ref_is = sort_by_id(ref_w, ref_ids)
tri_ws, tri_is = sort_by_id(tri_w, tri_ids)
ids_match = (ref_is == tri_is).all().item()
weights_match = torch.allclose(ref_ws, tri_ws, atol=1e-2, rtol=1e-2)
all_match = ids_match and weights_match
max_err = (ref_ws - tri_ws).abs().max().item()
print(
f"T={num_tokens:5d} E={num_experts:4d} K={topk} "
f"{str(dtype).split('.')[-1]:7s} ids={ids_match} max_Δweight={max_err:.2e}"
)
if not all_match:
bad = (ref_is != tri_is).any(dim=-1).nonzero(as_tuple=True)[0]
if len(bad):
r = bad[0].item()
print(
f" first bad row {r}: ref_ids={ref_ids[r].tolist()} "
f"tri_ids={tri_ids[r].tolist()}"
)
assert all_match
+45 -166
View File
@@ -14,8 +14,6 @@ import pytest
import torch
from torch.nn import Parameter
from torch.nn import functional as F
from transformers import MixtralConfig
from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
import vllm.model_executor.layers.fused_moe # noqa
from tests.kernels.moe.utils import (
@@ -24,10 +22,7 @@ from tests.kernels.moe.utils import (
modular_triton_fused_moe,
)
from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
from vllm._aiter_ops import rocm_aiter_ops
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.distributed.parallel_state import init_distributed_environment
from vllm.forward_context import get_forward_context, set_forward_context
from vllm.model_executor.layers.fused_moe import (
MoEActivation,
fused_topk,
@@ -56,12 +51,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
marlin_quantize,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
from vllm.model_executor.models.mixtral import MixtralMoE
from vllm.platforms import current_platform
from vllm.scalar_type import ScalarType, scalar_types
from vllm.utils.math_utils import next_power_of_2
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.worker.workspace import init_workspace_manager
def iterative_moe(
@@ -150,12 +143,14 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [
{
"a_type": [scalar_types.bfloat16],
"b_type": scalar_types.float4_e2m1f,
"c_type": [scalar_types.bfloat16],
"group_blocks": [2],
},
# MXFP8
{
"a_type": [scalar_types.bfloat16],
"b_type": scalar_types.float8_e4m3fn,
"c_type": [scalar_types.bfloat16],
"group_blocks": [2],
},
# AWQ-INT4 with INT8 activation
@@ -681,154 +676,35 @@ def test_fused_moe_wn16(
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("padding", [True, False])
@pytest.mark.parametrize(
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
)
@torch.inference_mode()
def test_mixtral_moe(
default_vllm_config,
dist_init,
dtype: torch.dtype,
padding: bool,
use_rocm_aiter: bool,
monkeypatch,
):
"""Make sure our Mixtral MoE implementation agrees with the one from
huggingface."""
# Explicitly set AITER env var based on test parameter to ensure
# consistent behavior regardless of external environment
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
rocm_aiter_ops.refresh_env_variables()
if use_rocm_aiter and dtype == torch.float32:
pytest.skip("AITER ROCm test skip for float32")
monkeypatch.setenv("RANK", "0")
monkeypatch.setenv("LOCAL_RANK", "0")
monkeypatch.setenv("WORLD_SIZE", "1")
monkeypatch.setenv("MASTER_ADDR", "localhost")
monkeypatch.setenv("MASTER_PORT", "12345")
init_distributed_environment()
init_workspace_manager(torch.accelerator.current_device_index())
# Instantiate our and huggingface's MoE blocks
vllm_config.compilation_config.static_forward_context = dict()
with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
config = MixtralConfig()
hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
vllm_moe = MixtralMoE(
num_experts=config.num_local_experts,
top_k=config.num_experts_per_tok,
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
params_dtype=dtype,
tp_size=1,
dp_size=1,
).cuda()
# Load the weights
vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data
if isinstance(hf_moe.experts, torch.nn.ModuleList):
# Transformers v4
for i in range(config.num_local_experts):
weights = (
hf_moe.experts[i].w1.weight.data,
hf_moe.experts[i].w3.weight.data,
)
vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
else:
# Transformers v5
vllm_moe.experts.w13_weight.data[:] = hf_moe.experts.gate_up_proj.data
vllm_moe.experts.w2_weight.data[:] = hf_moe.experts.down_proj.data
# TODO: remove this line after https://github.com/huggingface/transformers/pull/43622
hf_moe.experts.config._experts_implementation = "eager"
# Generate input batch of dimensions [batch_size, seq_len, hidden_dim]
hf_inputs = torch.randn((1, 64, config.hidden_size)).to(dtype).to("cuda")
# vLLM uses 1D query [num_tokens, hidden_dim]
vllm_inputs = hf_inputs.flatten(0, 1)
# Pad the weight if moe padding is enabled
if padding:
vllm_moe.experts.w13_weight = Parameter(
F.pad(vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[
..., 0:-128
],
requires_grad=False,
)
vllm_moe.experts.w2_weight = Parameter(
F.pad(vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128],
requires_grad=False,
)
torch.accelerator.synchronize()
torch.accelerator.empty_cache()
# FIXME (zyongye) fix this after we move self.kernel
# assignment in FusedMoE.__init__
vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts)
# need to override the forward context for unittests, otherwise it assumes
# we're running the model forward pass (the model specified in vllm_config)
get_forward_context().all_moe_layers = None
# Run forward passes for both MoE blocks
hf_states = hf_moe.forward(hf_inputs)
if isinstance(hf_states, tuple):
# Transformers v4
hf_states = hf_states[0]
vllm_states = vllm_moe.forward(vllm_inputs)
mixtral_moe_tol = {
torch.float32: 1e-3,
torch.float16: 1e-3,
torch.bfloat16: 1e-2,
}
if use_rocm_aiter:
# The values of rtol and atol are set based on the tests in ROCM AITER package.
# https://github.com/ROCm/aiter/blob/dfed377f4be7da96ca2d75ac0761f569676f7240/op_tests/test_moe.py#L174
torch.testing.assert_close(
hf_states.flatten(0, 1), vllm_states, rtol=0.01, atol=100
)
else:
torch.testing.assert_close(
hf_states.flatten(0, 1),
vllm_states,
rtol=mixtral_moe_tol[dtype],
atol=mixtral_moe_tol[dtype],
)
MARLIN_MOE_SCENARIOS = [
# (m, n, k, e, topk, ep_size, act_order, is_k_full)
# No act_order: is_k_full=True matches usual case (marlin_is_k_full).
# N>=256 required for Marlin kernel thread config for MXFP8.
# Single token, small matrices
(1, 128, 256, 5, 2, 1, False, True),
# Single token, large matrices
(1, 1024, 2048, 5, 2, 1, False, True),
# Unaligned m, small matrices
(133, 256, 256, 5, 2, 1, False, True),
# Unaligned m, large matrices
(133, 1024, 2048, 12, 3, 1, False, True),
# Aligned batch, small matrices
(128, 256, 256, 5, 2, 1, False, True),
# Aligned batch, large matrices
(128, 1024, 2048, 12, 3, 1, False, True),
# Expert parallelism
(64, 1024, 2048, 12, 3, 4, False, True),
# Act order with is_k_full=True (no tensor parallelism)
(1, 1024, 2048, 5, 2, 1, True, True),
# Act order with is_k_full=False (tensor parallelism)
(133, 256, 256, 5, 2, 1, True, False),
]
def marlin_moe_generate_valid_test_cases():
import itertools
m_list = [1, 123, 666]
n_list = [128, 1024]
k_list = [256, 2048]
e_list = [5, 12]
topk_list = [2, 3]
ep_size_list = [1, 4]
act_order_list = [True, False]
is_k_full_list = [True, False]
all_combinations = itertools.product(
MOE_MARLIN_QUANT_TEST_CONFIGS,
m_list,
n_list,
k_list,
e_list,
topk_list,
ep_size_list,
act_order_list,
is_k_full_list,
)
def is_invalid(
def is_valid(
a_type,
b_type,
c_type,
@@ -845,39 +721,42 @@ def marlin_moe_generate_valid_test_cases():
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
if group_size > 0 and k % group_size != 0:
return False
if act_order and group_size in [-1, k, n]:
return False
if group_size in [k, n]:
return False
if not act_order and is_k_full:
if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
return False
return a_type.size_bits < 16 or a_type is c_type
cases = []
for case in all_combinations:
quant_test_config, m, n, k, _, _, _, act_order, *_ = case
if act_order and not quant_test_config.get("support_act_order", False):
continue
for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS:
f16_types = [scalar_types.float16]
inner_combinations = itertools.product(
quant_test_config.get("a_type", f16_types),
[quant_test_config["b_type"]],
quant_test_config.get("c_type", f16_types),
quant_test_config["group_blocks"],
inner_combinations = list(
itertools.product(
quant_test_config.get("a_type", f16_types),
[quant_test_config["b_type"]],
quant_test_config.get("c_type", f16_types),
quant_test_config["group_blocks"],
)
)
supports_act_order = quant_test_config.get("support_act_order", False)
for sub_case in inner_combinations:
if (
sub_case[0] == scalar_types.float8_e4m3fn
and current_platform.get_device_capability() not in [89, 120]
):
continue
args = sub_case + (m, n, k) + case[4:]
if is_invalid(*args):
cases.append(args)
for scenario in MARLIN_MOE_SCENARIOS:
m, n, k, e, topk, ep_size, act_order, is_k_full = scenario
if act_order and not supports_act_order:
continue
args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full)
if is_valid(*args):
cases.append(args)
return cases
+8
View File
@@ -465,6 +465,14 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
if config.enable_eplb and config.ep_size == 1:
return False, "EPLB only works with EP+DP"
# Disable fp4 tests until flashinfer is updated or the Dockerfile is
# modified to install cublasLt.h. See #39525.
if (
config.quantization == "modelopt_fp4"
and current_platform.is_device_capability_family(100)
):
return False, "Temporarily skip until #39525 is resolved"
return True, None
+248
View File
@@ -0,0 +1,248 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels."""
import random
import pytest
import torch
from tests.kernels.utils import torch_moe_single
from vllm import _custom_ops as ops
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
random.seed(42)
set_random_seed(42)
MXFP4_BLOCK_SIZE = 32
def align(val: int, alignment: int = 128) -> int:
return int((val + alignment - 1) // alignment * alignment)
def calc_diff(x, y):
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
def is_sm100_supported() -> bool:
return current_platform.is_cuda() and current_platform.is_device_capability_family(
100
)
def compute_ref_output(
input_tensor: torch.Tensor,
weight_list: list[torch.Tensor],
expert_offsets: list[int],
expert_offset: int,
num_experts: int,
) -> torch.Tensor:
"""Reference output using torch_moe_single with top-1 routing."""
score = torch.full(
(expert_offset, num_experts),
-1e9,
device=input_tensor.device,
dtype=torch.float32,
)
for g in range(num_experts):
start = expert_offsets[g]
end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset
score[start:end, g] = 0.0
return torch_moe_single(
input_tensor, torch.stack(weight_list, dim=0), score, topk=1
)
@pytest.mark.skipif(
not is_sm100_supported(),
reason="cutlass_mxfp4_group_mm requires CUDA SM100",
)
@pytest.mark.parametrize("num_experts", [8, 16, 32])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16])
def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype):
"""
Test the MXFP4 grouped GEMM kernel by:
1. Creating random per-expert inputs and weights
2. Quantizing both to MXFP4 using the CUDA kernel
3. Running the CUTLASS grouped GEMM
4. Comparing against BF16 reference
"""
device = "cuda"
alignment = 128
# N and K must be multiples of 128 for clean swizzle layout
n_g = random.randint(1, 16) * alignment
k_g = random.randint(1, 16) * alignment
expert_offset = 0
expert_offsets_input = []
problem_sizes = []
input_list = []
weight_list = []
for g in range(num_experts):
m_g = random.randint(1, 256)
expert_offsets_input.append(expert_offset)
expert_offset += m_g
problem_sizes.append([m_g, n_g, k_g])
input_list.append(
torch.normal(0.0, std=0.5, size=(m_g, k_g), device=device, dtype=out_dtype)
)
weight_list.append(
torch.normal(0.0, std=0.5, size=(n_g, k_g), device=device, dtype=out_dtype)
)
input_tensor = torch.concat(input_list, dim=0) # [M_total, K]
# --- Quantize INPUTS via mxfp4_experts_quant ---
input_bs_offsets = []
tot = 0
for g in range(num_experts):
input_bs_offsets.append(tot)
tot += align(problem_sizes[g][0], 128)
input_bs_offsets.append(tot)
_inp_expert_offsets = torch.tensor(
expert_offsets_input + [expert_offset], device=device, dtype=torch.int32
)
_inp_bs_offsets = torch.tensor(input_bs_offsets, device=device, dtype=torch.int32)
input_quant, input_sf = ops.mxfp4_experts_quant(
input_tensor,
_inp_expert_offsets,
_inp_bs_offsets,
num_experts,
topk=1,
)
# --- Quantize WEIGHTS via mxfp4_experts_quant ---
# Treat each expert's N weight rows as an "expert" with N tokens
weight_tensor = torch.concat(weight_list, dim=0) # [E*N, K]
weight_expert_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
# N is always multiple of 128, so blockscale offsets are clean
weight_bs_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
_wt_expert_offsets = torch.tensor(
weight_expert_offsets, device=device, dtype=torch.int32
)
_wt_bs_offsets = torch.tensor(weight_bs_offsets, device=device, dtype=torch.int32)
weight_quant, weight_sf = ops.mxfp4_experts_quant(
weight_tensor,
_wt_expert_offsets,
_wt_bs_offsets,
num_experts,
topk=1,
)
# Reshape weight quantized data to [E, N, K//2]
weight_quant = weight_quant[: num_experts * n_g].view(num_experts, n_g, k_g // 2)
# Reshape weight scale factors to [E, N, K//32]
# The quant kernel produces uint8 SF buffer. Each row has K//32 SFs.
scales_per_row = k_g // MXFP4_BLOCK_SIZE
weight_sf_flat = weight_sf.view(-1)[: num_experts * n_g * scales_per_row]
weight_sf_3d = weight_sf_flat.view(num_experts, n_g, scales_per_row)
# Output
output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
_problem_sizes = torch.tensor(problem_sizes, device=device, dtype=torch.int32)
_expert_offsets = torch.tensor(
expert_offsets_input, device=device, dtype=torch.int32
)
_input_bs = torch.tensor(input_bs_offsets[:-1], device=device, dtype=torch.int32)
# Run the MXFP4 grouped GEMM
ops.cutlass_mxfp4_moe_mm(
output,
input_quant,
weight_quant,
input_sf,
weight_sf_3d,
_problem_sizes,
_expert_offsets,
_input_bs,
)
# Reference: BF16 matmul
ref_output = compute_ref_output(
input_tensor=input_tensor,
weight_list=weight_list,
expert_offsets=expert_offsets_input,
expert_offset=expert_offset,
num_experts=num_experts,
)
# Compare per-expert
for g in range(num_experts):
start = expert_offsets_input[g]
end = expert_offsets_input[g + 1] if g + 1 < num_experts else expert_offset
if start == end:
continue
baseline = ref_output[start:end]
actual = output[start:end]
diff = calc_diff(actual, baseline)
print(
f"m_g={end - start} n_g={n_g} k_g={k_g} "
f"num_experts={num_experts}, "
f"out_dtype={out_dtype}, diff={diff:.5f}"
)
# FP4 quantization is very lossy (~4 bits precision)
# Comparing quantized vs full-precision gives cosine diff of 0.05-0.15
assert diff < 0.15, f"Expert {g}: diff={diff:.5f} exceeds threshold"
@pytest.mark.skipif(
not is_sm100_supported(),
reason="mxfp4_experts_quant requires CUDA SM100",
)
def test_mxfp4_experts_quant_basic():
"""
Basic smoke test for the MXFP4 experts quantization kernel.
"""
device = "cuda"
num_experts = 4
k = 256
tokens_per_expert = 16
total_tokens = tokens_per_expert * num_experts
input_tensor = torch.randn(total_tokens, k, device=device, dtype=torch.bfloat16) / 5
expert_offsets = [i * tokens_per_expert for i in range(num_experts + 1)]
blockscale_offsets = [
align(i * tokens_per_expert, 128) for i in range(num_experts + 1)
]
_expert_offsets = torch.tensor(expert_offsets, device=device, dtype=torch.int32)
_blockscale_offsets = torch.tensor(
blockscale_offsets, device=device, dtype=torch.int32
)
output, output_sf = ops.mxfp4_experts_quant(
input_tensor,
_expert_offsets,
_blockscale_offsets,
num_experts,
topk=1,
)
assert output.shape == (total_tokens, k // 2)
assert output.dtype == torch.uint8
assert output_sf.dtype == torch.uint8
assert output.any(), "Quantized output is all zeros"
print(
f"MXFP4 experts quant: output shape={output.shape}, sf shape={output_sf.shape}"
)
print("PASSED")
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
@@ -11,6 +11,11 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
)
from vllm.platforms import current_platform
skipif_not_cuda_rocm = pytest.mark.skipif(
not (current_platform.is_cuda() or current_platform.is_rocm()),
reason="Only supported on CUDA/ROCm platforms.",
)
@pytest.mark.parametrize(
"platform_method,expected_backend",
@@ -190,3 +195,83 @@ def test_select_cuda_flashinfer_cutlass_backend(
assert selected_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS
assert experts_cls is not None
@skipif_not_cuda_rocm
def test_select_lora_backend_prefers_triton():
"""LoRA-enabled unquantized MoE should select Triton backend."""
moe_config = make_dummy_moe_config()
moe_config.is_lora_enabled = True
selected_backend, experts_cls = select_unquantized_moe_backend(
moe_config=moe_config
)
assert selected_backend == UnquantizedMoeBackend.TRITON
assert experts_cls is not None
@skipif_not_cuda_rocm
def test_select_lora_explicit_non_triton_backend():
"""LoRA should override explicit non-Triton backend to Triton."""
moe_config = make_dummy_moe_config()
moe_config.is_lora_enabled = True
# Use string from mapping in function map_unquantized_backend()
moe_config.moe_backend = "flashinfer_cutlass"
selected_backend, experts_cls = select_unquantized_moe_backend(
moe_config=moe_config
)
assert selected_backend == UnquantizedMoeBackend.TRITON
assert experts_cls is not None
@skipif_not_cuda_rocm
@pytest.mark.parametrize("is_lora_enabled", [False, True])
def test_select_explicit_triton_backend(is_lora_enabled):
"""Explicit triton backend selection should return Triton."""
moe_config = make_dummy_moe_config()
moe_config.is_lora_enabled = is_lora_enabled
moe_config.moe_backend = "triton"
selected_backend, experts_cls = select_unquantized_moe_backend(
moe_config=moe_config
)
assert selected_backend == UnquantizedMoeBackend.TRITON
assert experts_cls is not None
@skipif_not_cuda_rocm
def test_select_explicit_triton_ignores_flashinfer_env(monkeypatch):
"""Explicit triton backend should override FlashInfer env selection."""
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
moe_config = make_dummy_moe_config()
moe_config.is_lora_enabled = False
moe_config.moe_backend = "triton"
selected_backend, experts_cls = select_unquantized_moe_backend(
moe_config=moe_config
)
assert selected_backend == UnquantizedMoeBackend.TRITON
assert experts_cls is not None
@skipif_not_cuda_rocm
def test_select_lora_ignores_flashinfer_env(monkeypatch):
"""LoRA path should still choose Triton even if FlashInfer env is on."""
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
moe_config = make_dummy_moe_config()
moe_config.is_lora_enabled = True
selected_backend, experts_cls = select_unquantized_moe_backend(
moe_config=moe_config
)
assert selected_backend == UnquantizedMoeBackend.TRITON
assert experts_cls is not None
+54
View File
@@ -88,6 +88,60 @@ def break_fp4_bytes(a, dtype):
return values.reshape(m, n * 2).to(dtype=dtype)
def dequant_nvfp4_kv_cache(
fp4_data: torch.Tensor,
block_scale: torch.Tensor,
global_scale: float,
head_size: int,
block_size: int,
) -> torch.Tensor:
"""Dequantize an NVFP4 KV cache with 4x4-swizzled block scales.
The input must be in HND layout so that the last two dims are
(block_size, last_dim). For NHD caches, permute to HND first.
Args:
fp4_data: [..., num_heads, block_size, head_size//2] uint8 packed fp4.
block_scale: [..., num_heads, block_size, head_size//16] fp8 block
scales (as uint8 or float8_e4m3fn).
global_scale: checkpoint dequant scale (k_scale or v_scale).
head_size: head dimension.
block_size: page size.
Returns:
[..., num_heads, block_size, head_size] float32.
"""
data_dim = head_size // 2
scale_dim = head_size // 16
fp4_packed = fp4_data
sf_swizzled = block_scale.view(torch.uint8)
# Unswizzle 4x4 block scales on (block_size, scale_dim) plane.
# [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S]
batch_shape = sf_swizzled.shape[:-2]
T, S = block_size, scale_dim
sg = S // 4
sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4)
ndim = sf_reshape.ndim
# Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg)
perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32)
# Unpack fp4
shape = fp4_packed.shape # [..., T, data_dim]
fp4_flat = fp4_packed.reshape(-1, data_dim)
fp4_vals = break_fp4_bytes(fp4_flat, torch.float32)
fp4_vals = fp4_vals.reshape(*shape[:-1], head_size)
# Dequant: fp4_val * block_scale * global_scale per 16-element group
return (
fp4_vals.reshape(*shape[:-1], scale_dim, 16)
* (sf_f32 * global_scale).unsqueeze(-1)
).reshape(*shape[:-1], head_size)
def get_nvfp4_global_scale(a: torch.Tensor):
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32)
File diff suppressed because one or more lines are too long
@@ -46,7 +46,7 @@ AITER_MODEL_LIST = [
),
pytest.param(
"openai-community/gpt2", # gpt2
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
marks=[pytest.mark.core_model],
),
pytest.param("Milos/slovak-gpt-j-405M"), # gptj
pytest.param("bigcode/tiny_starcoder_py"), # gpt_bigcode
@@ -143,11 +143,6 @@ def test_models(
# in parts of the operators
pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
if current_platform.is_cpu() and model in ("openai-community/gpt2",):
# These models are sensitive to the rounding error
# Fuse ops to reduce rounding
monkeypatch.setenv("VLLM_CPU_CI_ENV", "0")
with hf_runner(model) as hf_model:
hf_outputs = hf_model.generate_greedy_logprobs_limit(
example_prompts, max_tokens, num_logprobs
@@ -4,7 +4,6 @@
import os
from collections.abc import Sequence
import librosa
import pytest
import regex as re
from huggingface_hub import snapshot_download
@@ -14,6 +13,7 @@ from vllm.assets.image import ImageAsset
from vllm.logprobs import SampleLogprobs
from vllm.lora.request import LoRARequest
from vllm.multimodal.image import convert_image_mode, rescale_image_size
from vllm.multimodal.media.audio import load_audio
from ....conftest import (
IMAGE_ASSETS,
@@ -290,7 +290,7 @@ def test_vision_speech_models(
num_logprobs: int,
) -> None:
# use the example speech question so that the model outputs are reasonable
audio = librosa.load(speech_question, sr=None)
audio = load_audio(speech_question, sr=None)
image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
inputs_vision_speech = [
@@ -25,6 +25,7 @@ if TYPE_CHECKING:
PIXTRAL_ID = "mistralai/Pixtral-12B-2409"
MISTRAL_SMALL_3_1_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
MINISTRAL_3B_ID = "mistralai/Ministral-3-3B-Instruct-2512"
MODELS = [PIXTRAL_ID, MISTRAL_SMALL_3_1_ID]
@@ -116,6 +117,7 @@ assert FIXTURES_PATH.exists()
FIXTURE_LOGPROBS_CHAT = {
PIXTRAL_ID: FIXTURES_PATH / "pixtral_chat.json",
MISTRAL_SMALL_3_1_ID: FIXTURES_PATH / "mistral_small_3_chat.json",
MINISTRAL_3B_ID: FIXTURES_PATH / "ministral_3b_chat.json",
}
OutputsLogprobs = list[tuple[list[int], str, SampleLogprobs | None]]
@@ -209,3 +211,41 @@ def test_chat(
name_0="h100_ref",
name_1="output",
)
@large_gpu_test(min_gb=16)
@pytest.mark.parametrize("dtype", ["bfloat16"])
def test_chat_consolidated(vllm_runner, dtype: str, local_asset_server) -> None:
EXPECTED_CHAT_LOGPROBS = load_outputs_w_logprobs(
FIXTURE_LOGPROBS_CHAT[MINISTRAL_3B_ID]
)
with vllm_runner(
MINISTRAL_3B_ID,
dtype=dtype,
tokenizer_mode="mistral",
load_format="mistral",
config_format="mistral",
max_model_len=8192,
limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
) as vllm_model:
outputs = []
urls_all = [local_asset_server.url_for(u) for u in IMG_URLS]
msgs = [
_create_msg_format(urls_all[:1]),
_create_msg_format(urls_all[:2]),
_create_msg_format(urls_all),
]
for msg in msgs:
output = vllm_model.llm.chat(msg, sampling_params=SAMPLING_PARAMS)
outputs.extend(output)
logprobs = vllm_runner._final_steps_generate_w_logprobs(outputs)
for i in range(len(logprobs)):
assert logprobs[i][-1] is None
logprobs[i] = logprobs[i][:-1]
check_logprobs_close(
outputs_0_lst=EXPECTED_CHAT_LOGPROBS,
outputs_1_lst=logprobs,
name_0="h100_ref",
name_1="output",
)
@@ -4,11 +4,11 @@
from collections.abc import Sequence
from typing import Any
import librosa
import pytest
from transformers import AutoModelForSpeechSeq2Seq
from vllm.assets.audio import AudioAsset
from vllm.multimodal.audio import AudioResampler
from vllm.platforms import current_platform
from ....conftest import HfRunner, PromptAudioInput, VllmRunner
@@ -93,13 +93,12 @@ def run_test(
def resampled_assets() -> list[tuple[Any, int]]:
audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
sampled_assets = []
resampler = AudioResampler(target_sr=WHISPER_SAMPLE_RATE)
for asset in audio_assets:
audio, orig_sr = asset.audio_and_sample_rate
# Resample to Whisper's expected sample rate (16kHz)
if orig_sr != WHISPER_SAMPLE_RATE:
audio = librosa.resample(
audio, orig_sr=orig_sr, target_sr=WHISPER_SAMPLE_RATE
)
audio = resampler.resample(audio, orig_sr=orig_sr)
sampled_assets.append(
(audio, WHISPER_SAMPLE_RATE),
)
@@ -140,88 +140,3 @@ def test_audio_token_count_matches_hf_processor_math():
_count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 0) == 1499
)
assert _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 1) == 375
def test_audio_feature_pipeline_matches_hf_small_config():
from transformers.models.audioflamingo3 import (
modeling_audioflamingo3 as hf_audioflamingo3_modeling,
)
from transformers.models.audioflamingo3.configuration_audioflamingo3 import (
AudioFlamingo3Config,
)
from vllm.model_executor.models.audioflamingo3 import (
AudioFlamingo3Encoder,
AudioFlamingo3MultiModalProjector,
_build_audio_encoder_attention_mask,
_flatten_valid_audio_embeddings,
)
text_config = {
"model_type": "qwen2",
"intermediate_size": 64,
"initializer_range": 0.02,
"hidden_size": 32,
"max_position_embeddings": 1024,
"num_hidden_layers": 2,
"num_attention_heads": 4,
"num_key_value_heads": 2,
"vocab_size": 128,
"pad_token_id": 1,
"use_mrope": False,
}
audio_config = {
"hidden_size": 16,
"num_attention_heads": 4,
"intermediate_size": 32,
"num_hidden_layers": 2,
"num_mel_bins": 80,
"max_source_positions": 1500,
"dropout": 0.0,
"attention_dropout": 0.0,
"activation_dropout": 0.0,
"encoder_layerdrop": 0.0,
}
torch.manual_seed(0)
config = AudioFlamingo3Config(
text_config=text_config,
audio_config=audio_config,
audio_token_id=0,
)
hf_model = hf_audioflamingo3_modeling.AudioFlamingo3ForConditionalGeneration(
config
).eval()
vllm_encoder = AudioFlamingo3Encoder(config.audio_config).eval()
vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
vllm_projector = AudioFlamingo3MultiModalProjector(config).eval()
vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
input_features = torch.randn(3, 80, 3000)
feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
feature_attention_mask[0, :3000] = True
feature_attention_mask[1, :2500] = True
feature_attention_mask[2, :1500] = True
hf_output = hf_model.get_audio_features(
input_features,
feature_attention_mask,
return_dict=True,
).pooler_output
vllm_attention_mask = _build_audio_encoder_attention_mask(
feature_attention_mask,
dtype=vllm_encoder.conv1.weight.dtype,
device=vllm_encoder.conv1.weight.device,
)
vllm_hidden_states = vllm_encoder(
input_features,
attention_mask=vllm_attention_mask,
)
vllm_output, _ = _flatten_valid_audio_embeddings(
vllm_projector(vllm_hidden_states),
feature_attention_mask,
)
torch.testing.assert_close(vllm_output, hf_output)
+2 -2
View File
@@ -3,12 +3,12 @@
from pathlib import Path
from unittest.mock import patch
import librosa
import numpy as np
import pybase64 as base64
import pytest
from vllm.multimodal.media import AudioMediaIO
from vllm.multimodal.media.audio import load_audio
from ...conftest import AudioTestAssets
@@ -73,6 +73,6 @@ def test_audio_media_io_from_video(video_assets):
video_path = video_assets[0].video_path
with open(video_path, "rb") as f:
audio, sr = audio_io.load_bytes(f.read())
audio_ref, sr_ref = librosa.load(video_path, sr=None)
audio_ref, sr_ref = load_audio(video_path, sr=None)
assert sr == sr_ref
np.testing.assert_allclose(audio_ref, audio, atol=1e-4)
+3 -6
View File
@@ -26,11 +26,8 @@ def test_placeholder_range_get_num_embeds(is_embed, expected):
"is_embed,expected",
[
(None, None),
(
torch.tensor([False, True, False, True, True]),
torch.tensor([0, 1, 1, 2, 3]),
),
(torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
(torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]),
(torch.tensor([True, True, True]), [1, 2, 3]),
],
)
def test_placeholder_range_embeds_cumsum(is_embed, expected):
@@ -41,6 +38,6 @@ def test_placeholder_range_embeds_cumsum(is_embed, expected):
assert pr.embeds_cumsum is None
return
assert torch.equal(pr.embeds_cumsum, expected)
assert pr.embeds_cumsum == expected
# cached_property should return the same object on repeated access
assert pr.embeds_cumsum is pr.embeds_cumsum
-1
View File
@@ -38,6 +38,5 @@ def test_model_experts_int8_startup(
dtype=dtype,
enforce_eager=True,
quantization="experts_int8",
allow_deprecated_quantization=True,
) as vllm_model:
vllm_model.generate_greedy(example_prompts, max_tokens)
+28 -54
View File
@@ -18,9 +18,7 @@ from vllm.model_executor.layers.quantization.turboquant.config import (
TQ_PRESETS,
TurboQuantConfig,
)
from vllm.model_executor.layers.quantization.turboquant.quantizer import (
generate_wht_signs,
)
from vllm.platforms import current_platform
from vllm.utils.math_utils import next_power_of_2
# ============================================================================
@@ -345,7 +343,8 @@ class TestLloydMax:
# Rotation matrix tests (GPU required)
# ============================================================================
CUDA_AVAILABLE = torch.cuda.is_available()
GPGPU_AVAILABLE = torch.cuda.is_available() or torch.xpu.is_available()
DEVICE_TYPE = current_platform.device_type
def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
@@ -360,16 +359,16 @@ def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Te
return Q.to(device)
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
class TestRotationMatrix:
"""Tests for the QR-based rotation (standalone benchmarks only)."""
@pytest.mark.parametrize("dim", [64, 96, 128, 256])
def test_rotation_matrix_shape_and_orthogonal(self, dim):
Pi = generate_rotation_matrix(dim, seed=42, device="cuda")
Pi = generate_rotation_matrix(dim, seed=42, device=DEVICE_TYPE)
assert Pi.shape == (dim, dim)
eye = Pi @ Pi.T
assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
f"Pi not orthogonal for dim={dim}"
)
@@ -385,13 +384,13 @@ class TestRotationMatrix:
def test_rotation_matrix_det_is_pm1(self):
"""Orthogonal matrix determinant must be +1 or -1."""
Pi = generate_rotation_matrix(128, seed=42, device="cuda")
Pi = generate_rotation_matrix(128, seed=42, device=DEVICE_TYPE)
det = torch.linalg.det(Pi)
assert abs(abs(det.item()) - 1.0) < 1e-4
# ============================================================================
# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard)
# Hadamard rotation tests (serving path: _build_hadamard)
# ============================================================================
@@ -403,58 +402,34 @@ def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
return (H / math.sqrt(d)).to(torch.device(device))
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
class TestWHTRotation:
"""Tests for the WHT rotation actually used in serving."""
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
class TestHadamardRotation:
"""Tests for the Hadamard rotation used in serving."""
@pytest.mark.parametrize("dim", [64, 128, 256])
def test_wht_orthonormal(self, dim):
"""signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
signs = generate_wht_signs(dim, seed=42, device="cuda")
H = _build_hadamard(dim, "cuda")
PiT = (signs.unsqueeze(1) * H).contiguous()
eye = PiT @ PiT.T
assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
f"WHT rotation not orthonormal for dim={dim}"
def test_hadamard_orthonormal(self, dim):
"""H must be orthonormal: H @ H^T = I."""
H = _build_hadamard(dim, DEVICE_TYPE)
eye = H @ H.T
assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
f"Hadamard not orthonormal for dim={dim}"
)
@pytest.mark.parametrize("dim", [64, 128, 256])
def test_wht_self_inverse(self, dim):
"""PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
signs = generate_wht_signs(dim, seed=42, device="cuda")
H = _build_hadamard(dim, "cuda")
PiT = (signs.unsqueeze(1) * H).contiguous()
Pi = PiT.T.contiguous()
# Pi @ PiT should be identity (rotation then inverse)
result = Pi @ PiT
assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
f"WHT rotation not self-inverse for dim={dim}"
def test_hadamard_symmetric(self, dim):
"""Sylvester Hadamard must be symmetric: H = H^T."""
H = _build_hadamard(dim, DEVICE_TYPE)
assert torch.allclose(H, H.T, atol=1e-6), (
f"Hadamard not symmetric for dim={dim}"
)
def test_wht_signs_deterministic(self):
"""Same seed must produce identical signs."""
s1 = generate_wht_signs(128, seed=42)
s2 = generate_wht_signs(128, seed=42)
assert torch.equal(s1, s2)
def test_wht_signs_different_seeds(self):
"""Different seeds must produce different signs."""
s1 = generate_wht_signs(128, seed=42)
s2 = generate_wht_signs(128, seed=99)
assert not torch.equal(s1, s2)
def test_wht_signs_are_pm1(self):
"""All sign values must be exactly +1 or -1."""
signs = generate_wht_signs(128, seed=42)
assert torch.all(signs.abs() == 1.0)
# ============================================================================
# Store → Decode round-trip test (GPU + Triton required)
# ============================================================================
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
class TestStoreDecodeRoundTrip:
"""End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""
@@ -487,13 +462,12 @@ class TestStoreDecodeRoundTrip:
block_size = 16
num_blocks = 1
device = torch.device("cuda")
device = torch.device(DEVICE_TYPE)
# Generate rotation
signs = generate_wht_signs(D, seed=42, device=device)
H = _build_hadamard(D, "cuda")
PiT = (signs.unsqueeze(1) * H).contiguous().float()
Pi = PiT.T.contiguous()
# Pure Hadamard rotation (symmetric: H = H^T, so Pi = PiT = H)
H = _build_hadamard(D, DEVICE_TYPE)
PiT = H
Pi = H
# Generate centroids
centroids, _ = solve_lloyd_max(D, cfg.centroid_bits)
-16
View File
@@ -17,22 +17,6 @@ from vllm import LLM, SamplingParams
from vllm.platforms import current_platform
@pytest.mark.skip(reason="In V1, we reject tokens > max_seq_len")
def test_duplicated_ignored_sequence_group():
"""https://github.com/vllm-project/vllm/issues/1655"""
sampling_params = SamplingParams(temperature=0.01, top_p=0.1, max_tokens=256)
llm = LLM(
model="distilbert/distilgpt2",
max_num_batched_tokens=4096,
tensor_parallel_size=1,
)
prompts = ["This is a short prompt", "This is a very long prompt " * 1000]
outputs = llm.generate(prompts, sampling_params=sampling_params)
assert len(prompts) == len(outputs)
@pytest.mark.parametrize(
"model",
[
File diff suppressed because it is too large Load Diff
@@ -64,9 +64,6 @@ class TestQwen3xmlToolParser(ToolParserTests):
"test_empty_arguments": "Qwen3XML streaming has systematic issues",
"test_surrounding_text": "Qwen3XML streaming has systematic issues",
"test_escaped_strings": "Qwen3XML streaming has systematic issues",
"test_malformed_input": (
"Qwen3XML parser is lenient with malformed input"
),
"test_streaming_reconstruction": (
"Qwen3XML streaming reconstruction has known issues"
),
@@ -0,0 +1,116 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression tests for Responses API tool-calling request adjustment.
Covers two bugs on the ``/v1/responses`` path that broke streaming tool
calling for parsers relying on special-token delimiters (Gemma4):
1. :class:`Gemma4ToolParser.adjust_request` used an
``isinstance(request, ChatCompletionRequest)`` guard, so a
:class:`ResponsesRequest` with tools never had
``skip_special_tokens`` flipped to ``False``. The default (``True``)
stripped ``<|tool_call>`` / ``<tool_call|>`` delimiters, causing
:meth:`Gemma4ToolParser.extract_tool_calls_streaming` to fall through
to the content branch and leak the raw ``call:fn{...}`` body via
``response.output_text.delta``.
2. :meth:`ToolParser.adjust_request` built
:class:`ResponseTextConfig` in two steps (bare constructor then
``.format = ...``). Under Pydantic v2 the later assignment is not
tracked in ``__fields_set__``, which can drop the nested config from
``model_dump``. It also passed a ``description`` kwarg carrying the
wrong-purpose string ``"Response format for tool calling"``.
"""
from __future__ import annotations
from typing import Any
from openai.types.responses.tool_param import FunctionToolParam
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.tool_parsers.abstract_tool_parser import ToolParser
from vllm.tool_parsers.gemma4_tool_parser import Gemma4ToolParser
def _get_weather_tool() -> FunctionToolParam:
return FunctionToolParam(
type="function",
name="get_weather",
description="Get current weather for a city",
parameters={
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"],
},
strict=True,
)
def _build_responses_request(*, tool_choice: str) -> ResponsesRequest:
return ResponsesRequest(
model="gemma4-test",
input=[{"role": "user", "content": "What is the weather in Hanoi?"}],
tools=[_get_weather_tool()],
tool_choice=tool_choice,
stream=True,
max_output_tokens=200,
)
class _StubTokenizer:
"""Minimal tokenizer stub to satisfy ``Gemma4ToolParser.__init__``."""
def get_vocab(self) -> dict[str, int]:
return {"<|tool_call>": 256_000, "<tool_call|>": 256_001, '<|"|>': 52}
def test_gemma4_adjust_request_sets_skip_special_tokens_on_responses() -> None:
"""``Gemma4ToolParser.adjust_request`` must flip
``skip_special_tokens=False`` for both ``ChatCompletionRequest`` and
``ResponsesRequest`` so that ``<|tool_call>`` delimiters reach the
streaming extractor. The previous
``isinstance(ChatCompletionRequest)`` guard omitted the Responses
path, causing raw ``call:fn{...}`` text to leak via
``response.output_text.delta``.
"""
parser = Gemma4ToolParser.__new__(Gemma4ToolParser)
parser.model_tokenizer = _StubTokenizer()
request = _build_responses_request(tool_choice="auto")
assert request.skip_special_tokens is True, (
"Precondition: ResponsesRequest.skip_special_tokens default is True"
)
Gemma4ToolParser.adjust_request(parser, request)
assert request.skip_special_tokens is False
def test_tool_parser_adjust_request_builds_valid_response_text_config() -> None:
"""``ToolParser.adjust_request`` must produce a ``ResponseTextConfig``
whose dumped form contains the JSON schema under the ``schema`` alias
and does not leak the unrelated ``"Response format for tool calling"``
description string that the previous two-step construction injected.
"""
parser = ToolParser.__new__(ToolParser)
parser.model_tokenizer = None
request = _build_responses_request(tool_choice="required")
ToolParser.adjust_request(parser, request)
assert request.text is not None
assert request.text.format is not None
assert request.text.format.type == "json_schema"
dump: dict[str, Any] = request.text.model_dump(mode="json", by_alias=True)
fmt = dump.get("format") or {}
assert fmt.get("type") == "json_schema"
assert fmt.get("name") == "tool_calling_response"
assert fmt.get("strict") is True
# Nested config must be present under the alias. Two-step Pydantic v2
# construction could drop it from __fields_set__.
assert "schema" in fmt and isinstance(fmt["schema"], dict)
# The old code passed a wrong-purpose string; valid field should now
# either be absent or None (the openai-python default).
assert fmt.get("description") in (None, "")
+41 -2
View File
@@ -7,17 +7,39 @@ from tests.models.utils import check_embeddings_close
from vllm.utils.serial_utils import (
EMBED_DTYPES,
ENDIANNESS,
MM_METADATA_DTYPES,
EmbedDType,
Endianness,
MmMetadataDType,
binary2tensor,
tensor2binary,
)
FLOAT_EMBED_DTYPES = tuple(EMBED_DTYPES.keys())
INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys())
def _build_integer_tensor(
embed_dtype: MmMetadataDType, shape: tuple[int, ...]
) -> torch.Tensor:
torch_dtype = MM_METADATA_DTYPES[embed_dtype].torch_dtype
if torch_dtype is torch.bool:
return torch.randint(0, 2, shape, dtype=torch.int32).to(torch.bool)
if torch_dtype is torch.uint8:
return torch.randint(0, 256, shape, dtype=torch.uint8)
if torch_dtype is torch.int32:
return torch.randint(-(2**20), 2**20, shape, dtype=torch.int32)
if torch_dtype is torch.int64:
return torch.randint(-(2**62), 2**62, shape, dtype=torch.int64)
raise AssertionError(f"Unsupported non-floating embed dtype: {embed_dtype}")
@pytest.mark.parametrize("endianness", ENDIANNESS)
@pytest.mark.parametrize("embed_dtype", EMBED_DTYPES.keys())
@pytest.mark.parametrize("embed_dtype", FLOAT_EMBED_DTYPES)
@torch.inference_mode()
def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
def test_encode_and_decode_floats(embed_dtype: EmbedDType, endianness: Endianness):
for i in range(10):
tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32)
shape = tensor.shape
@@ -40,3 +62,20 @@ def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
name_1="new",
tol=1e-2,
)
@pytest.mark.parametrize("endianness", ENDIANNESS)
@pytest.mark.parametrize("embed_dtype", INTEGER_EMBED_DTYPES)
@torch.inference_mode()
def test_encode_and_decode_integers(
embed_dtype: MmMetadataDType, endianness: Endianness
):
shape = (2, 3, 5, 7, 11, 13)
for i in range(10):
tensor = _build_integer_tensor(embed_dtype, shape)
binary = tensor2binary(tensor, embed_dtype, endianness)
new_tensor = binary2tensor(binary, shape, embed_dtype, endianness)
assert new_tensor.dtype == MM_METADATA_DTYPES[embed_dtype].torch_dtype
torch.testing.assert_close(tensor, new_tensor, atol=0, rtol=0)
@@ -0,0 +1,105 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Test batch-invariant matmul against torch.matmul for various shape combinations.
Tests correctness (matches torch.matmul) and batch invariance (result for one
item doesn't change based on other items in the batch).
"""
import pytest
import torch
from utils import skip_unsupported
from vllm.model_executor.layers.batch_invariant import matmul_batch_invariant
from vllm.platforms import current_platform
DEVICE_TYPE = current_platform.device_type
@skip_unsupported
@pytest.mark.parametrize(
"a_shape,b_shape",
[
# 2D x 2D
((32, 64), (64, 16)),
# 2D x 3D
((64, 16), (4, 16, 32)),
# 3D x 2D
((4, 32, 64), (64, 16)),
# 4D x 2D
((1, 4, 32, 64), (64, 16)),
# 3D x 3D
((4, 32, 64), (4, 64, 16)),
# 3D x 4D
((2, 32, 64), (1, 2, 64, 16)),
# 4D x 3D (Gemma4 pattern)
((1, 2, 32, 64), (2, 64, 16)),
# 4D x 4D
((1, 2, 32, 64), (4, 2, 64, 16)),
# 2D x 4D
((32, 64), (1, 2, 64, 16)),
# 2D x 5D
((32, 64), (1, 2, 2, 64, 16)),
# 5D x 2D
((1, 2, 2, 32, 64), (64, 16)),
# 5D x 5D
((1, 2, 4, 32, 64), (1, 2, 4, 64, 16)),
],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_matmul_correctness(a_shape, b_shape, dtype):
"""
Compare matmul_batch_invariant against torch.matmul for various shapes.
"""
device = torch.device(DEVICE_TYPE)
torch.manual_seed(42)
a = torch.rand(a_shape, dtype=dtype, device=device)
b = torch.rand(b_shape, dtype=dtype, device=device)
# Standard implementation (CUDA ops)
standard_output = torch.matmul(a, b)
# Batch-invariant implementation (Triton)
triton_output = matmul_batch_invariant(a, b)
# Compare outputs
# Use looser tolerance for bfloat16 due to its lower precision
if dtype == torch.bfloat16:
rtol, atol = 1e-1, 1e-1 # 10% relative tolerance for bfloat16
else:
rtol, atol = 1e-2, 1e-2 # 1% for float16/float32
torch.testing.assert_close(
triton_output,
standard_output,
rtol=rtol,
atol=atol,
msg=f"matmul mismatch for a ndim={a.ndim}, b ndim={b.ndim},",
)
@skip_unsupported
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_matmul_batch_invariance(dtype):
"""
Verify that the result for one item is bitwise identical regardless
of what other items are in the batch.
"""
device = torch.device(DEVICE_TYPE)
torch.manual_seed(42)
a_single = torch.rand((1, 64, 32), dtype=dtype, device=device)
b = torch.rand((32, 128), dtype=dtype, device=device)
standard_output = matmul_batch_invariant(a_single, b)
a_batch = torch.rand((8, 64, 32), dtype=dtype, device=device)
a_batch[3] = a_single[0]
batch_output = matmul_batch_invariant(a_batch, b)
batch_output_a = batch_output[3]
assert torch.equal(standard_output[0], batch_output_a)
@@ -32,8 +32,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
# 3 blocks, store just the middle block (skip first and last)
# blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
runner.new_request(token_ids=[0] * offloaded_block_size * 3)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(
list(keys)[1:2]
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output(list(keys)[1:2])
)
runner.run(decoded_tokens=[0])
@@ -45,18 +45,22 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
runner.manager.prepare_store.assert_not_called()
# +1 token -> single block, fail prepare_store
runner.manager.prepare_store.side_effect = lambda keys: None
runner.manager.prepare_store.side_effect = lambda keys, req_context: None
runner.run(decoded_tokens=[0])
runner.manager.prepare_store.assert_called()
# 1 more block (+ token for async scheduling)
# now set block_hashes_to_store = []
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output([])
)
runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
# 1 more block (+ token for kicking off offloading)
# now check touch was called with all 6 blocks
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output(keys)
)
runner.run(
decoded_tokens=[0] * (offloaded_block_size + 1),
expected_stored_gpu_block_indexes=(15, 16, 17),
@@ -89,13 +93,17 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
runner.new_request(
token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output([])
)
runner.run(decoded_tokens=[EOS_TOKEN_ID])
runner.manager.lookup.assert_not_called()
# single block lookup with no hits
runner.new_request(token_ids=[1] * offloaded_block_size)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output([])
)
runner.run(decoded_tokens=[EOS_TOKEN_ID])
runner.manager.lookup.assert_called()
assert len(list(runner.manager.lookup.call_args.args[0])) == 1
@@ -103,7 +111,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
# single block lookup with a hit
runner.scheduler.reset_prefix_cache()
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output([])
)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
@@ -113,7 +123,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
runner.new_request(
token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output([])
)
runner.manager.lookup.return_value = 1
runner.run(
decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
@@ -164,14 +176,18 @@ def test_request_preemption(request_runner, async_scheduling: bool):
# 2 blocks, store all, without flushing
# blocks = [0, 1, 2], [3, 4, 5]
runner.new_request(token_ids=[0] * offloaded_block_size * 2)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output(keys)
)
runner.run(
decoded_tokens=[0],
complete_transfers=False,
)
# decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output(keys)
)
runner.run(
decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
complete_transfers=False,
@@ -195,7 +211,9 @@ def test_request_preemption(request_runner, async_scheduling: bool):
# request should now return from preemption
# re-load [0, ..., 8] from the CPU and store [9, 10, 11]
runner.manager.lookup.return_value = 3
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output(keys)
)
runner.run(
decoded_tokens=[0] * gpu_block_size,
expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
@@ -222,7 +240,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
# store 1 blocks
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output(keys)
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=(0, 1, 2),
@@ -253,7 +273,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
# complete transfers
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output([])
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_loaded_gpu_block_indexes=(0, 1, 2),
@@ -278,7 +300,9 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
# store 1 blocks
runner.new_request(token_ids=[0] * offloaded_block_size)
runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
runner.manager.prepare_store.side_effect = (
lambda keys, req_context: generate_store_output(keys)
)
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored_gpu_block_indexes=(0, 1, 2),
@@ -115,7 +115,7 @@ class MockOffloadingSpec(OffloadingSpec):
self.manager = MagicMock(spec=OffloadingManager)
self.manager.lookup.return_value = 0
self.manager.prepare_load = lambda keys: MockLoadStoreSpec(keys)
self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
self.handler = MockOffloadingHandler()
def get_manager(self) -> OffloadingManager:
+68 -49
View File
@@ -11,6 +11,7 @@ from vllm.v1.kv_offload.abstract import (
OffloadingEvent,
OffloadKey,
PrepareStoreOutput,
ReqContext,
make_offload_key,
)
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
@@ -19,6 +20,14 @@ from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext:
"""Create a ReqContext as production code would, from a request's params."""
return ReqContext(kv_transfer_params=kv_transfer_params)
_EMPTY_REQ_CTX = make_req_context()
@dataclass
class ExpectedPrepareStoreOutput:
keys_to_store: list[int]
@@ -103,7 +112,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
)
# store [1, 2] and complete
manager.prepare_store(to_keys([1, 2]))
manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
manager.complete_store(to_keys([1, 2]))
# touch [1] to make block 2 the LRU candidate
@@ -113,7 +122,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
# - block 2 is already stored -> filtered out of keys_to_store
# - block 2 must NOT be evicted even though it is the LRU candidate
# - block 1 (ID 0) is evicted instead; new blocks [3,4,5] get IDs 2,3,0
prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]))
prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
@@ -127,7 +136,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
manager.complete_store(to_keys([2, 3, 4, 5]))
# block 2 must still be present in the cache
assert manager.lookup(to_keys([2])) == 1
assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1
def test_cpu_manager():
@@ -140,7 +149,7 @@ def test_cpu_manager():
)
# prepare store [1, 2]
prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]))
prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
@@ -151,7 +160,7 @@ def test_cpu_manager():
)
# lookup [1, 2] -> not ready
assert cpu_manager.lookup(to_keys([1, 2])) == 0
assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
# no events so far
assert list(cpu_manager.take_events()) == []
@@ -161,12 +170,14 @@ def test_cpu_manager():
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
# lookup [1, 2]
assert cpu_manager.lookup(to_keys([1])) == 1
assert cpu_manager.lookup(to_keys([1, 2])) == 2
assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2
assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
# prepare store [2, 3, 4, 5] -> evicts [1]
prepare_store_output = cpu_manager.prepare_store(to_keys([2, 3, 4, 5]))
prepare_store_output = cpu_manager.prepare_store(
to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX
)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
@@ -180,23 +191,23 @@ def test_cpu_manager():
verify_events(cpu_manager.take_events(), expected_evictions=({1},))
# prepare store with no space
assert cpu_manager.prepare_store(to_keys([1, 6])) is None
assert cpu_manager.prepare_store(to_keys([1, 6]), _EMPTY_REQ_CTX) is None
# complete store [2, 3, 4, 5]
cpu_manager.complete_store(to_keys([2, 3, 4, 5]))
# prepare load [2, 3]
prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]))
prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
verify_load_output(prepare_load_output, [1, 2])
# prepare store with no space ([2, 3] is being loaded)
assert cpu_manager.prepare_store(to_keys([6, 7, 8])) is None
assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
# complete load [2, 3]
cpu_manager.complete_load(to_keys([2, 3]))
# prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest)
prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]))
prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
@@ -213,7 +224,7 @@ def test_cpu_manager():
cpu_manager.touch(to_keys([5, 6, 7]))
# prepare store [7, 9] -> evicts [8] (oldest following previous touch)
prepare_store_output = cpu_manager.prepare_store(to_keys([9]))
prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
@@ -227,8 +238,8 @@ def test_cpu_manager():
cpu_manager.complete_store(to_keys([7, 9]), success=False)
# assert [7] is still stored, but [9] is not
assert cpu_manager.lookup(to_keys([7])) == 1
assert cpu_manager.lookup(to_keys([9])) == 0
assert cpu_manager.lookup(to_keys([7]), _EMPTY_REQ_CTX) == 1
assert cpu_manager.lookup(to_keys([9]), _EMPTY_REQ_CTX) == 0
verify_events(
cpu_manager.take_events(),
@@ -260,7 +271,9 @@ class TestARCPolicy:
cpu_manager, arc_policy = self._make_manager()
# prepare store [1, 2]
prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]))
prepare_store_output = cpu_manager.prepare_store(
to_keys([1, 2]), _EMPTY_REQ_CTX
)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
@@ -271,7 +284,7 @@ class TestARCPolicy:
)
# lookup [1, 2] -> not ready
assert cpu_manager.lookup(to_keys([1, 2])) == 0
assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
# no events so far
assert list(cpu_manager.take_events()) == []
@@ -281,9 +294,9 @@ class TestARCPolicy:
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
# lookup [1, 2]
assert cpu_manager.lookup(to_keys([1])) == 1
assert cpu_manager.lookup(to_keys([1, 2])) == 2
assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2
assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
# blocks should be in T1 (recent)
assert len(arc_policy.t1) == 2
@@ -297,7 +310,7 @@ class TestARCPolicy:
cpu_manager, arc_policy = self._make_manager(enable_events=False)
# store and complete block 1
cpu_manager.prepare_store(to_keys([1]))
cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1]))
# block 1 starts in T1 (recent)
@@ -319,7 +332,9 @@ class TestARCPolicy:
cpu_manager, _ = self._make_manager()
# prepare and complete store [1, 2, 3, 4]
prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
prepare_store_output = cpu_manager.prepare_store(
to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX
)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
@@ -331,19 +346,21 @@ class TestARCPolicy:
cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
# prepare load [2, 3] (increases ref_cnt)
prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]))
prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
verify_load_output(prepare_load_output, [1, 2])
# prepare store [5, 6, 7] with [2, 3] being loaded
# should fail because [2, 3] have ref_cnt > 0
assert cpu_manager.prepare_store(to_keys([5, 6, 7])) is None
assert cpu_manager.prepare_store(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) is None
# complete load [2, 3]
cpu_manager.complete_load(to_keys([2, 3]))
# now prepare store [5, 6, 7] should succeed
# ARC will evict blocks one at a time from T1 as needed
prepare_store_output = cpu_manager.prepare_store(to_keys([5, 6, 7]))
prepare_store_output = cpu_manager.prepare_store(
to_keys([5, 6, 7]), _EMPTY_REQ_CTX
)
assert prepare_store_output is not None
# Should successfully evict enough blocks to make room (at least 1)
assert len(prepare_store_output.evicted_keys) >= 1
@@ -357,13 +374,13 @@ class TestARCPolicy:
cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
# store blocks 1, 2 (fills cache)
cpu_manager.prepare_store(to_keys([1, 2]))
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2]))
initial_target = arc_policy.target_t1_size
# store block 3, evicting block 1 (moves to B1 ghost list)
cpu_manager.prepare_store(to_keys([3]))
cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([3]))
# block 1 should be in B1 (ghost list)
@@ -384,7 +401,7 @@ class TestARCPolicy:
cpu_manager, arc_policy = self._make_manager(enable_events=False)
# store blocks 1, 2, 3, 4
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
# promote blocks 3, 4 to T2 by touching them
@@ -399,7 +416,7 @@ class TestARCPolicy:
arc_policy.target_t1_size = 1
# store block 5, should evict from T1 (block 1, LRU in T1)
output = cpu_manager.prepare_store(to_keys([5]))
output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
assert output is not None
assert to_keys([1]) == output.evicted_keys
@@ -418,12 +435,12 @@ class TestARCPolicy:
cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
# fill cache with blocks 1, 2
cpu_manager.prepare_store(to_keys([1, 2]))
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2]))
# store many blocks to fill ghost lists
for i in range(3, 20):
cpu_manager.prepare_store(to_keys([i]))
cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([i]))
# ghost lists should not exceed cache_capacity
@@ -438,7 +455,7 @@ class TestARCPolicy:
cpu_manager, arc_policy = self._make_manager()
# store blocks 1, 2, 3, 4
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
# promote 3, 4 to T2
@@ -453,7 +470,7 @@ class TestARCPolicy:
assert len(arc_policy.t2) == 3
# store block 5, should evict from T1 (block 2, only one in T1)
prepare_store_output = cpu_manager.prepare_store(to_keys([5]))
prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
verify_store_output(
prepare_store_output,
ExpectedPrepareStoreOutput(
@@ -471,11 +488,11 @@ class TestARCPolicy:
cpu_manager, arc_policy = self._make_manager()
# store blocks 1, 2, 3, 4
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
# prepare store block 5 (will evict block 1)
prepare_store_output = cpu_manager.prepare_store(to_keys([5]))
prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
assert prepare_store_output is not None
assert len(prepare_store_output.evicted_keys) == 1
@@ -483,7 +500,7 @@ class TestARCPolicy:
cpu_manager.complete_store(to_keys([5]), success=False)
# block 5 should not be in cache
assert cpu_manager.lookup(to_keys([5])) == 0
assert cpu_manager.lookup(to_keys([5]), _EMPTY_REQ_CTX) == 0
# block 5 should not be in T1 or T2
assert to_keys([5])[0] not in arc_policy.t1
assert to_keys([5])[0] not in arc_policy.t2
@@ -500,11 +517,13 @@ class TestARCPolicy:
cpu_manager, arc_policy = self._make_manager()
# store [1, 2]
cpu_manager.prepare_store(to_keys([1, 2]))
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2]))
# store [3, 4, 5] -> evicts [1]
prepare_store_output = cpu_manager.prepare_store(to_keys([3, 4, 5]))
prepare_store_output = cpu_manager.prepare_store(
to_keys([3, 4, 5]), _EMPTY_REQ_CTX
)
assert prepare_store_output is not None
assert len(prepare_store_output.evicted_keys) == 1
cpu_manager.complete_store(to_keys([3, 4, 5]))
@@ -517,13 +536,13 @@ class TestARCPolicy:
assert len(arc_policy.t2) == 2
# store [6] -> should evict from T1 (4 is oldest in T1)
prepare_store_output = cpu_manager.prepare_store(to_keys([6]))
prepare_store_output = cpu_manager.prepare_store(to_keys([6]), _EMPTY_REQ_CTX)
assert prepare_store_output is not None
cpu_manager.complete_store(to_keys([6]))
# verify blocks 2, 3 (in T2) are still present
assert cpu_manager.lookup(to_keys([2])) == 1
assert cpu_manager.lookup(to_keys([3])) == 1
assert cpu_manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1
assert cpu_manager.lookup(to_keys([3]), _EMPTY_REQ_CTX) == 1
# verify events
events = list(cpu_manager.take_events())
@@ -543,34 +562,34 @@ def test_filter_reused_manager():
)
# Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
assert manager.lookup(to_keys([1, 2])) == 0
assert manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
# prepare store [1, 2] -> should be filtered
prepare_store_output = manager.prepare_store(to_keys([1, 2]))
prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
assert prepare_store_output is not None
assert prepare_store_output.keys_to_store == []
# Lookup [1] -> 2nd time, eligible now
assert manager.lookup(to_keys([1])) == 0
assert manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 0
# prepare store [1, 2] -> [1] should be eligible, [2] should be filtered
prepare_store_output = manager.prepare_store(to_keys([1, 2]))
prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
assert prepare_store_output is not None
assert prepare_store_output.keys_to_store == to_keys([1])
# Lookup [3, 4] -> 1st time
# (evicts [2] from tracker since max_size is 3 and tracker has [1])
assert manager.lookup(to_keys([3, 4])) == 0
assert manager.lookup(to_keys([3, 4]), _EMPTY_REQ_CTX) == 0
# Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])
assert to_keys([2])[0] not in manager.counts
# Lookup [2] again -> (this adds [2] back to the tracker as 1st time)
assert manager.lookup(to_keys([2])) == 0
assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 0
# Verify [2] was re-added with count=1 (not eligible yet)
assert manager.counts.get(to_keys([2])[0]) == 1
# prepare store [2] -> should still be filtered out since count was reset
prepare_store_output = manager.prepare_store(to_keys([2]))
prepare_store_output = manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX)
assert prepare_store_output is not None
assert prepare_store_output.keys_to_store == []
+135
View File
@@ -1150,6 +1150,38 @@ def cutlass_fp4_moe_mm(
)
def cutlass_mxfp4_moe_mm(
out_tensors: torch.Tensor,
a_tensors: torch.Tensor,
b_tensors: torch.Tensor,
a_scales: torch.Tensor,
b_scales: torch.Tensor,
problem_sizes: torch.Tensor,
expert_offsets: torch.Tensor,
sf_offsets: torch.Tensor,
):
"""
An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
- a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
- a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout)
- Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales).
- expert_offsets/sf_offsets: expert boundary indices
- problem_sizes: (num_experts, 3) with (M, N, K) per expert
"""
return torch.ops._C.cutlass_mxfp4_group_mm(
out_tensors,
a_tensors,
b_tensors,
a_scales,
b_scales,
problem_sizes,
expert_offsets,
sf_offsets,
)
def mxfp8_experts_quant(
input_tensor: torch.Tensor,
problem_sizes: torch.Tensor,
@@ -1848,6 +1880,109 @@ def silu_and_mul_scaled_fp4_experts_quant(
return output, output_scales
def mxfp4_experts_quant(
input_tensor: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
n_experts: int,
topk: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Quantize input tensor to MXFP4 for packed MoE inputs.
Uses 32-element blocks with E8M0 (power-of-two) scale factors.
MXFP4 has no global scale - only block-level E8M0 scale factors.
Args:
input_tensor: [m_topk, k] BF16/FP16 activations
expert_offsets: [n_experts+1] token boundaries per expert
blockscale_offsets: [n_experts+1] SF row boundaries per expert
n_experts: number of experts
topk: number of top-k experts
Returns:
output: [m_topk, k//2] packed E2M1 values (uint8)
output_scales: E8M0 blockscales in swizzled layout (uint8 view)
"""
assert not current_platform.is_rocm()
assert input_tensor.ndim == 2
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
m_numtopk, k = input_tensor.shape
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
f"{MAX_TOKENS_PER_EXPERT})"
f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use"
f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
)
scales_k = k // 32
padded_k = (scales_k + (4 - 1)) // 4
output = torch.empty(
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
)
output_scales = torch.empty(
MAX_TOKENS_PER_EXPERT * topk,
padded_k,
dtype=torch.int32,
device=input_tensor.device,
)
torch.ops._C.mxfp4_experts_quant(
output,
output_scales,
input_tensor,
expert_offsets,
blockscale_offsets,
n_experts,
)
# E8M0 SFs are stored as uint8
output_scales = output_scales.view(torch.uint8)
return output, output_scales
def silu_and_mul_mxfp4_experts_quant(
input_tensor: torch.Tensor,
expert_offsets: torch.Tensor,
blockscale_offsets: torch.Tensor,
n_experts: int,
topk: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
MXFP4 has no global scale - only block-level E8M0 scale factors.
"""
assert not current_platform.is_rocm()
assert input_tensor.ndim == 2
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
m_numtopk, k_times_2 = input_tensor.shape
assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
k = k_times_2 // 2
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk
scales_k = k // 32
padded_k = (scales_k + (4 - 1)) // 4
output = torch.empty(
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
)
output_scales = torch.empty(
MAX_TOKENS_PER_EXPERT * topk,
padded_k,
dtype=torch.int32,
device=input_tensor.device,
)
torch.ops._C.silu_and_mul_mxfp4_experts_quant(
output,
output_scales,
input_tensor,
expert_offsets,
blockscale_offsets,
n_experts,
)
output_scales = output_scales.view(torch.uint8)
return output, output_scales
# fp8
def scaled_fp8_quant(
input: torch.Tensor,
+17
View File
@@ -22,6 +22,23 @@ else:
except ImportError:
from torch.library import impl_abstract as register_fake
if hasattr(torch.ops._xpu_C, "fp8_gemm"):
@register_fake("_xpu_C::fp8_gemm")
def _fp8_gemm_fake(
q_input: torch.Tensor,
q_weight: torch.Tensor,
out_dtype: torch.dtype,
input_scales: torch.Tensor,
weight_scale: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
input_2d = q_input.view(-1, q_input.shape[-1])
M = input_2d.size(0)
N = q_weight.size(1)
return torch.empty((M, N), dtype=out_dtype, device=q_input.device)
if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"):
@register_fake("_xpu_C::fp8_gemm_w8a16")
+3 -1
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""The request function for API endpoints."""
import codecs
import io
import json
import os
@@ -25,11 +26,12 @@ class StreamedResponseHandler:
def __init__(self):
self.buffer = ""
self._decoder = codecs.getincrementaldecoder("utf-8")()
def add_chunk(self, chunk_bytes: bytes) -> list[str]:
"""Add a chunk of bytes to the buffer and return any complete
messages."""
chunk_str = chunk_bytes.decode("utf-8")
chunk_str = self._decoder.decode(chunk_bytes)
self.buffer += chunk_str
messages = []
+16 -14
View File
@@ -8,7 +8,7 @@ from dataclasses import dataclass
from functools import partial
from pathlib import Path
from types import TracebackType
from typing import ClassVar
from typing import TYPE_CHECKING, ClassVar
from typing_extensions import Self, override
@@ -17,20 +17,8 @@ from vllm.utils.import_utils import PlaceholderModule
from .utils import sanitize_filename
try:
import matplotlib.pyplot as plt
except ImportError:
plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
try:
if TYPE_CHECKING:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
try:
import seaborn as sns
except ImportError:
seaborn = PlaceholderModule("seaborn")
@dataclass
@@ -265,6 +253,20 @@ def _plot_fig(
fig_height: float,
fig_dpi: int,
):
# Lazy-import matplotlib/pandas/seaborn
try:
import matplotlib.pyplot as plt
except ImportError:
plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
try:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
try:
import seaborn as sns
except ImportError:
sns = PlaceholderModule("seaborn")
fig_group, fig_data = fig_group_data
row_groups = full_groupby(
+16 -14
View File
@@ -6,7 +6,7 @@ from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from typing import ClassVar
from typing import TYPE_CHECKING, ClassVar
from vllm.utils.collection_utils import full_groupby
from vllm.utils.import_utils import PlaceholderModule
@@ -14,20 +14,8 @@ from vllm.utils.import_utils import PlaceholderModule
from .plot import DummyExecutor, _json_load_bytes
from .utils import sanitize_filename
try:
import matplotlib.pyplot as plt
except ImportError:
plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
try:
if TYPE_CHECKING:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
try:
import seaborn as sns
except ImportError:
seaborn = PlaceholderModule("seaborn")
def _first_present(run_data: dict[str, object], keys: list[str]):
@@ -195,6 +183,20 @@ def _plot_fig(
print("[END FIGURE]")
return
# Lazy-import matplotlib/pandas/seaborn
try:
import matplotlib.pyplot as plt
except ImportError:
plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
try:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
try:
import seaborn as sns
except ImportError:
sns = PlaceholderModule("seaborn")
df = pd.DataFrame.from_records(fig_data)
df = df.dropna(subset=["tokens_per_user", "tokens_per_gpu"])
+2 -7
View File
@@ -290,14 +290,9 @@ class CUDAGraphWrapper:
# across layers will make the cudagraph capture very slow.
# therefore, we only run gc for the first graph,
# and disable gc for the rest of the graphs.
stack.enter_context(patch("gc.collect", lambda: None))
stack.enter_context(
patch("gc.collect", lambda *args, **kwargs: None)
)
stack.enter_context(
patch(
"torch.accelerator.empty_cache",
lambda *args, **kwargs: None,
)
patch("torch.accelerator.empty_cache", lambda: None)
)
if self.graph_pool is not None:
+1 -3
View File
@@ -30,6 +30,7 @@ CacheDType = Literal[
"turboquant_3bit_nc",
"int8_per_token_head",
"fp8_per_token_head",
"nvfp4",
]
MambaDType = Literal["auto", "float32", "float16"]
MambaCacheMode = Literal["all", "align", "none"]
@@ -101,8 +102,6 @@ class CacheConfig:
kv_cache_dtype_skip_layers: list[str] = field(default_factory=list)
"""Layer patterns to skip KV cache quantization. Accepts layer indices
(e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window')."""
cpu_kvcache_space_bytes: int | None = None
"""(CPU backend only) CPU key-value cache space."""
mamba_page_size_padded: int | None = None
""" Optional override for mamba page size; used by hybrid mamba/attention
models to ensure exact alignment with attention page size."""
@@ -183,7 +182,6 @@ class CacheConfig:
"num_gpu_blocks_override",
"enable_prefix_caching",
"prefix_caching_hash_algo",
"cpu_kvcache_space_bytes",
"mamba_page_size_padded",
"user_specified_block_size",
"user_specified_mamba_block_size",

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