Compare commits

...
Author SHA1 Message Date
Jee Jee Li 42b812edf5 FIX
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-05-13 12:48:02 +00:00
hissu-hyvarinenandGitHub 0a62f5eec9 [AMD] skip machete tests for rocm (#42326)
Signed-off-by: Hissu Hyvarinen <hissu.hyvarinen@amd.com>
2026-05-13 12:11:03 +00:00
PikaPikachuandGitHub 3b1ef03be4 [Bugfix][Quark] Fix W8A8 INT8 garbage outputs on Step-3.5-Flash (and other 3-key fused-MoE Quark exports) (#41892)
Signed-off-by: kangletian <kangletian@hotmail.com>
2026-05-13 11:59:49 +00:00
3c413a5481 Triton attention: add USE_TD constexpr for tensor descriptor Q/K/V load/store (#40327)
Signed-off-by: Artur Fierka <artur.fierka@intel.com>
Co-authored-by: quinnlp <quinnlp@users.noreply.github.com>
2026-05-13 13:57:41 +02:00
Ronen SchafferandGitHub 79fd1bc7ed [kv_offload] Add req_id to ReqContext for per-request tracking (#42507)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
2026-05-13 11:11:10 +00:00
SILONG ZENGandGitHub cee6751e54 [Bugfix][Qwen3-VL] Fix pipeline-parallel deepstack initialization (#42394)
Signed-off-by: MrZ20 <2609716663@qq.com>
2026-05-13 10:58:42 +00:00
16863072ca [Bugfix] Fix scipy audio resampling ratio (#42233)
Signed-off-by: JooHo Lee <BWAAEEEK@users.noreply.github.com>
Co-authored-by: JooHo Lee <BWAAEEEK@users.noreply.github.com>
2026-05-13 18:52:41 +08:00
Andreas KaratzasandGitHub d628a3c5cb [ROCm][CI] Skip ROCm batch invalid-input test pending torch fix (#41572)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-05-13 18:50:47 +08:00
akii96andGitHub 74dffae666 [ROCm] Run AITER RMSNorm pad fusion before AR RMS fusion (#42411)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
2026-05-13 18:35:12 +08:00
97c4317bf5 [Bugfix][Frontend] Default max_tokens server-side on /inference/v1/generate (#42329)
Signed-off-by: hallerite <git@hallerite.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-13 11:16:46 +02:00
f6e868fbdf [CI] Use uv with Python 3.12 for PyPI wheel upload (#42470)
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-13 02:12:06 -07:00
Yifan QiaoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
13bf242100 [Feat][KVConnector] Add bind_gpu_block_pool() to KVConnectorBase_V1 (#39654)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-05-13 02:10:29 -07:00
Jiangyun ZhuandGitHub 140dc2ec30 [Bugfix] Install nvidia-cutlass-dsl[cu13] extra on CUDA 13 platforms (#42438)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-05-13 01:57:21 -07:00
Yifan QiaoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
9ce74042d3 [Bugfix][SimpleCPUOffloadBackend] Dedup in-flight CPU offload stores across scheduler steps (#41289)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-05-13 01:53:32 -07:00
sychen52andGitHub a8c13d2837 Patch SlidingWindowSpec.real_page_size_bytes for nvfp4 kv (#42464)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
2026-05-13 01:46:30 -07:00
Jee Jee Li 446bbb6ad3 Move
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-05-13 08:27:59 +00:00
Jee Jee Li 2e6f2b451e Done
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-05-13 08:15:45 +00:00
Shanshan ShenandGitHub 92def124bc [MM][Perf][CG] Support ViT full CUDA graph for Qwen3.5 (#42151)
Signed-off-by: shen-shanshan <467638484@qq.com>
2026-05-13 16:00:32 +08:00
85b2fecab7 [5/n] Migrate CUTLASS MLA, hadamard, awq, allspark and DSV3 fused a gemm to torch stable ABI (continued) (#42339)
Signed-off-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>
Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>
2026-05-13 07:24:39 +00:00
Taneem IbrahimGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>Wentao Ye
503697c9ce [chore] Refactor pooling metadata token ID accessors (#42368)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
2026-05-13 06:08:01 +00:00
Nicolò LucchesiandGitHub 71bcd02ef3 [Bugfix][PD] Fix multi-node TP (TP>8) (#39907)
Signed-off-by: NickLucche <nlucches@redhat.com>
2026-05-12 22:20:57 -07:00
Matthew BonanniandGitHub dcacdf9a88 [Attention] Sync FA with upstream (#41052) 2026-05-12 23:34:18 -04:00
bnellnmandGitHub 18f6bf5a21 [MoE Refactor] Add sequence parallel tests to test_moe_layer.py (#41299)
Signed-off-by: Bill Nell <bnell@redhat.com>
2026-05-12 21:52:19 -04:00
AlecandGitHub 07534b8782 [PD] Bump NIXL connector dependency to 1.x (#42364)
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
2026-05-12 18:05:01 -07:00
Wentao YeandGitHub 3d635c58c0 [Perf] Optimize MLA compute_prefill_context memory allocation (#42460)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-05-12 16:23:46 -07:00
+4 ebeb09d822 [KV Transfer] Add MooncakeStoreConnector for KV cache offloading via Mooncake distributed store (#40900)
Signed-off-by: leichao.lc <leichao.lc@antgroup.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: leichao.lc <leichao.lc@antgroup.com>
Co-authored-by: ivanium <yifanqiao@inferact.ai>
Co-authored-by: aoshen524 <aoshen@inferact.ai>
Co-authored-by: Dao007forever <daole@inferact.ai>
Co-authored-by: Teng Ma <sima.mt@alibaba-inc.com>
Co-authored-by: Pz1116 <zpbzpb123123@gmail.com>
Co-authored-by: foraxe <1055696449@qq.com>
Co-authored-by: Skywalker-EP <173423846@qq.com>
Co-authored-by: fems14 <1804143737@qq.com>
Co-authored-by: jianzs <zheng.shoujian@outlook.com>
Co-authored-by: baxingpiaochong <771405853@qq.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-05-12 16:09:10 -07:00
Michael GoinandGitHub 184577ae46 [Build] DeepGEMM: trim comments, add integration notes + TODOs (#42429)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-05-12 15:57:58 -07:00
Kevin H. LuuandGitHub 8c4fc4202a [CI] Inline build artifact annotations in release pipeline (#42357)
Signed-off-by: khluu <khluu000@gmail.com>
2026-05-12 15:57:43 -07:00
Nick HillandGitHub fe8b42e80c [CI] Fix test_async_scheduling.py flakiness (#42455)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-05-12 21:38:32 +00:00
Giancarlo DelfinandGitHub fe5b4e0fe7 [Model Runner V2] Apply synthetic mode to probabilistic rejection sampler (#41035) 2026-05-12 13:37:03 -07:00
0ce6613b9c platforms: add uses_cpu_device() hook to Platform for DeviceConfig (#42313)
Signed-off-by: Viktor Pus <viktorpus@tenstorrent.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
2026-05-12 12:39:17 -07:00
379f0ec369 [CI] Migrate 6 verified jobs from gpu_1_queue to h200_18gb MIG (#42446)
Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-12 11:52:01 -07:00
KaivalyaMDabhadkarandGitHub 67c89fe40a [Model][Bugfix] Fix Step3-VL image_embeds input path (#42333)
Signed-off-by: Kaivalya Dabhadkar <kdabhadkar@nvidia.com>
2026-05-12 18:47:55 +00:00
d9b4990783 [MoE Refactor] EPLB refactoring for FusedMoE (#41055)
Signed-off-by: Bill Nell <bnell@redhat.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
2026-05-12 14:16:31 -04:00
4d591db470 [MoE Refactor] Introduce RoutedExperts alias for FusedMoE and don't store SharedExperts in MK (#40735)
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
2026-05-12 13:37:44 -04:00
yzong-rhandGitHub 6ff7405b81 [Bugfix] [Frontend] Responses API, fix merging of messages (#42189)
Signed-off-by: Yifan Zong <yzong@redhat.com>
Signed-off-by: Yifan <yzong@redhat.com>
2026-05-12 16:09:59 +00:00
Yan Ru PeiandGitHub bcb9c133ba feat(kv-events): emit KV cache metadata (#40984)
Signed-off-by: PeaBrane <yanrpei@gmail.com>
2026-05-12 15:58:48 +00:00
c8a6e272e0 [CPU] Fix rotary embedding for CPU without flash-attn ops (#42225)
Signed-off-by: jmamou <jonathan.mamou@intel.com>
Signed-off-by: Jonathan Mamou <jonathan.mamou@intel.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
2026-05-12 15:05:35 +00:00
Wentao YeandGitHub a1b2d87498 [Refactor] Clean up pooling models build_tok_params logic (#42341)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-05-12 15:05:05 +00:00
Martin HickeyandGitHub 418ba8ef14 [kv_offload][BugFix] Fix store deferral (#41945)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>
2026-05-12 18:04:44 +03:00
Juan Pérez de AlgabaGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
5a6a9fc6f6 [docs] Added one new contact to the Vulnerability Management team (#42145)
Signed-off-by: jperezde <jperezde@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-05-12 10:59:59 -04:00
289cee0473 [vLLM IR] Minor improvements (#39362) (#39558)
Signed-off-by: Avishek Goswami <avishek.goswami@ibm.com>
Co-authored-by: Avishek Goswami <avishek.goswami@ibm.com>
2026-05-12 10:58:36 -04:00
shanjiazandGitHub 6ccb10d794 Added peagle speculators support (#41826)
Signed-off-by: shanjiaz <zsjwpianpian@gmail.com>
2026-05-12 07:55:57 -07:00
7a9cc5e7f0 [Model] Support MiniCPM-V 4.6 (#41254)
Signed-off-by: caitianchi <caitianchi@tc-mb.com>
Signed-off-by: tc-mb <157115220+tc-mb@users.noreply.github.com>
Co-authored-by: caitianchi <caitianchi@tc-mb.com>
2026-05-12 14:28:10 +00:00
d077622d60 [Build] Build bundled DeepGEMM _C per-Python so the wheel imports on every CPython (#41516)
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-12 10:27:29 -04:00
dd6b3a5ef5 [Perf] Use 2D-grid to eliminate divmod in W8W8 group quant (#42153)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-12 10:01:30 -04:00
593d5a4033 [Bugfix] Fix mismatched kernel-per-logical blocks in NIXL HMA transfer (#42097)
Signed-off-by: ZhanqiuHu <zhu@redhat.com>
Signed-off-by: Zhanqiu Hu <zhu@redhat.com>
Signed-off-by: NickLucche <nlucches@redhat.com>
Co-authored-by: NickLucche <nlucches@redhat.com>
2026-05-12 15:53:30 +02:00
bnellnmandGitHub 6427603ae8 [MoE Refactor] Move remaining experts classes to experts directory (#42334)
Signed-off-by: Bill Nell <bnell@redhat.com>
2026-05-12 09:19:46 -04:00
206eaed08d [MoE Refactor] Move expert map related code into ExpertMapManager class (#41046)
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
2026-05-12 09:18:27 -04:00
8f89381fc6 [Hybrid] Warmup Mamba2 SSD kernel (#39822)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-05-12 12:46:22 +00:00
Dipika SikkaandGitHub a7b801e26d [MXFP4] Support for linear layers + compressed-tensors integration (#41664) 2026-05-12 07:49:33 -04:00
Kunshang JiandGitHub 4df1be9547 [XPU] bump up vllm-xpu-kernels to v0.1.8 (#42410)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
2026-05-12 11:47:37 +00:00
bc03f280c8 [XPU] keep generator state of sycl kernel align with pytorch (#41771)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Co-authored-by: Qiming Zhang <qiming1.zhang@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-05-12 11:44:47 +00:00
997132911e [Doc] Fix typo in llm-d documentation link (#42397)
Signed-off-by: Florian Woerner <florian.woerner@onmyown.io>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
2026-05-12 04:26:46 -07:00
haosdentandGitHub fc8bf6eedb [CI] De-flake Language Models Test (Extended Generation) test_models(False-False-5-32-bigcode/starcoder2-3b) (#42392)
Signed-off-by: haosdent <haosdent@gmail.com>
2026-05-12 10:46:48 +00:00
liuzhenweiandGitHub 07a40ede19 [UT][XPU] fix test_parallel_sampling due to global random state (#42388)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
2026-05-12 18:03:23 +08:00
Kevin H. LuuandGitHub e1c8776e90 [CI] Move DockerHub and PyPI publish steps to end of release pipeline (#42355)
Signed-off-by: khluu <khluu000@gmail.com>
2026-05-12 09:17:42 +00:00
243 changed files with 9830 additions and 2474 deletions
+1
View File
@@ -8,6 +8,7 @@ run_all_patterns:
- "CMakeLists.txt"
- "requirements/common.txt"
- "requirements/cuda.txt"
- "requirements/kv_connectors.txt"
- "requirements/build/cuda.txt"
- "requirements/test/cuda.txt"
- "setup.py"
+79 -64
View File
@@ -28,6 +28,7 @@ steps:
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
env:
DOCKER_BUILDKIT: "1"
@@ -41,6 +42,7 @@ steps:
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
env:
DOCKER_BUILDKIT: "1"
@@ -54,6 +56,7 @@ steps:
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
env:
DOCKER_BUILDKIT: "1"
@@ -67,6 +70,7 @@ steps:
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
env:
DOCKER_BUILDKIT: "1"
@@ -80,6 +84,7 @@ steps:
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
env:
DOCKER_BUILDKIT: "1"
@@ -93,6 +98,7 @@ steps:
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
env:
DOCKER_BUILDKIT: "1"
@@ -138,6 +144,7 @@ steps:
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'
- label: "Build release image - aarch64 - CUDA 13.0"
depends_on: ~
@@ -160,6 +167,7 @@ steps:
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'
- label: "Build release image - x86_64 - CUDA 12.9"
depends_on: ~
@@ -184,6 +192,7 @@ steps:
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'
- label: "Build release image - aarch64 - CUDA 12.9"
depends_on: ~
@@ -205,6 +214,7 @@ steps:
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
depends_on: ~
@@ -231,6 +241,7 @@ steps:
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
depends_on: ~
@@ -255,6 +266,7 @@ steps:
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
depends_on: ~
@@ -280,6 +292,7 @@ steps:
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
depends_on: ~
@@ -303,6 +316,7 @@ steps:
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'
- block: "Build release image for x86_64 CPU"
key: block-cpu-release-image-build
@@ -320,6 +334,7 @@ steps:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
- "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
env:
DOCKER_BUILDKIT: "1"
@@ -339,6 +354,7 @@ steps:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
- "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
env:
DOCKER_BUILDKIT: "1"
@@ -356,15 +372,7 @@ steps:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
- label: "Annotate release workflow - CUDA 13.0"
depends_on:
- create-multi-arch-manifest
id: annotate-release-workflow
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/annotate-release.sh"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"'
- label: "Create multi-arch manifest - CUDA 12.9"
depends_on:
@@ -377,6 +385,7 @@ steps:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"'
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
depends_on:
@@ -389,6 +398,7 @@ steps:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"'
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
depends_on:
@@ -401,6 +411,7 @@ steps:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"'
- label: "Publish nightly multi-arch image to DockerHub"
depends_on:
@@ -438,59 +449,6 @@ steps:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- block: "Publish release images to DockerHub"
key: block-publish-release-images
depends_on:
- create-multi-arch-manifest
- create-multi-arch-manifest-cuda-12-9
- create-multi-arch-manifest-ubuntu2404
- create-multi-arch-manifest-cuda-12-9-ubuntu2404
- build-rocm-release-image
- input-release-version
# Wait for CPU builds if their block steps were unblocked, so publish
# doesn't race the in-progress CPU build. allow_failure lets publish
# proceed when the operator legitimately leaves the CPU block steps
# unblocked or the CPU build fails.
- step: build-cpu-release-image-x86
allow_failure: true
- step: build-cpu-release-image-arm64
allow_failure: true
if: build.env("NIGHTLY") != "1"
- label: "Publish release images to DockerHub"
depends_on:
- block-publish-release-images
key: publish-release-images-dockerhub
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/publish-release-images.sh"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- group: "Publish wheels"
key: "publish-wheels"
steps:
- block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
key: block-upload-release-wheels
depends_on:
- input-release-version
- build-wheels
- label: "Upload release wheels to PyPI"
depends_on:
- block-upload-release-wheels
id: upload-release-wheels
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
# =============================================================================
# ROCm Release Pipeline (x86_64 only)
# =============================================================================
@@ -604,7 +562,7 @@ steps:
echo ""
echo " Build complete - Image and wheels cached"
fi
artifact_paths:
- "artifacts/rocm-base-wheels/*.whl"
env:
@@ -820,7 +778,7 @@ steps:
# Push to ECR
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
echo ""
echo " Successfully built and pushed ROCm release image"
echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
@@ -847,3 +805,60 @@ steps:
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
# =============================================================================
# Publish to DockerHub and PyPI (at the end so all builds complete first)
# =============================================================================
- block: "Publish release images to DockerHub"
key: block-publish-release-images
depends_on:
- create-multi-arch-manifest
- create-multi-arch-manifest-cuda-12-9
- create-multi-arch-manifest-ubuntu2404
- create-multi-arch-manifest-cuda-12-9-ubuntu2404
- build-rocm-release-image
- input-release-version
# Wait for CPU builds if their block steps were unblocked, so publish
# doesn't race the in-progress CPU build. allow_failure lets publish
# proceed when the operator legitimately leaves the CPU block steps
# unblocked or the CPU build fails.
- step: build-cpu-release-image-x86
allow_failure: true
- step: build-cpu-release-image-arm64
allow_failure: true
if: build.env("NIGHTLY") != "1"
- label: "Publish release images to DockerHub"
depends_on:
- block-publish-release-images
key: publish-release-images-dockerhub
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/publish-release-images.sh"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- group: "Publish wheels"
key: "publish-wheels"
steps:
- block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
key: block-upload-release-wheels
depends_on:
- input-release-version
- build-wheels
- label: "Upload release wheels to PyPI"
depends_on:
- block-upload-release-wheels
id: upload-release-wheels
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
+9
View File
@@ -0,0 +1,9 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Append a build artifact line to the Buildkite annotation.
# Usage: annotate-build-artifact.sh <label> <value>
set -e
echo "- **${1}**: \`${2}\`" | \
buildkite-agent annotate --append --style 'info' --context 'release-artifacts'
-27
View File
@@ -1,27 +0,0 @@
#!/bin/bash
set -ex
# Get release version, default to 1.0.0.dev for nightly/per-commit builds
RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null | sed 's/^v//')
if [ -z "${RELEASE_VERSION}" ]; then
RELEASE_VERSION="1.0.0.dev"
fi
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
To download the wheel (by commit):
\`\`\`
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
(Optional) For CUDA 12.9:
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
(Optional) For CPU:
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
\`\`\`
Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
EOF
@@ -39,11 +39,17 @@ fi
set -x # avoid printing secrets above
# install twine and sdist build prerequisites from pypi
python3 -m venv /tmp/vllm-release-env
# install uv if not already available
if ! command -v uv &> /dev/null; then
curl -LsSf https://astral.sh/uv/install.sh | UV_VERSION=0.11.14 sh
export PATH="$HOME/.local/bin:$PATH"
fi
# install twine and sdist build prerequisites using uv with Python 3.12
uv venv --python 3.12 /tmp/vllm-release-env
source /tmp/vllm-release-env/bin/activate
pip install twine
pip install -r requirements/build/cuda.txt
uv pip install twine
uv pip install -r requirements/build/cuda.txt
python3 -m twine --version
# copy release wheels to local directory
+6
View File
@@ -91,6 +91,7 @@ steps:
- label: Kernels KDA Test
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
- vllm/model_executor/layers/fla/ops/kda.py
- vllm/model_executor/layers/fla/ops/chunk_delta_h.py
@@ -105,7 +106,11 @@ steps:
device: h100
num_devices: 1
source_file_dependencies:
- cmake/external_projects/deepgemm.cmake
- tools/install_deepgemm.sh
- tools/build_deepgemm_C.py
- tools/setup_deepgemm_pythons.sh
- tools/check_wheel_deepgemm.py
- vllm/utils/deep_gemm.py
- vllm/model_executor/layers/fused_moe
- vllm/model_executor/layers/quantization
@@ -115,6 +120,7 @@ steps:
- tests/kernels/attention/test_deepgemm_attention.py
- tests/quantization/test_cutlass_w4a16.py
commands:
- python3 ../tools/check_wheel_deepgemm.py
- pytest -v -s kernels/quantization/test_block_fp8.py
- pytest -v -s kernels/moe/test_deepgemm.py
- pytest -v -s kernels/moe/test_batched_deepgemm.py
+1
View File
@@ -114,6 +114,7 @@ steps:
- label: LM Eval TurboQuant KV Cache
key: lm-eval-turboquant-kv-cache
timeout_in_minutes: 75
device: h200_18gb
source_file_dependencies:
- vllm/model_executor/layers/quantization/turboquant/
- vllm/v1/attention/backends/turboquant_attn.py
+3
View File
@@ -79,6 +79,7 @@ steps:
- tests/v1/core
- tests/v1/executor
- tests/v1/kv_offload
- tests/v1/simple_kv_offload
- tests/v1/worker
- tests/v1/kv_connector/unit
- tests/v1/metrics
@@ -90,6 +91,7 @@ steps:
- pytest -v -s -m 'not cpu_test' v1/core
- pytest -v -s v1/executor
- pytest -v -s v1/kv_offload
- pytest -v -s v1/simple_kv_offload
- pytest -v -s v1/worker
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
@@ -212,6 +214,7 @@ steps:
depends_on: ~
optional: true
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
- tests/standalone_tests/python_only_compile.sh
- setup.py
+2 -3
View File
@@ -106,13 +106,12 @@ steps:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/v1/spec_decode/test_max_len.py
- tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
- tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- tests/v1/spec_decode/test_rejection_sampler_utils.py
- tests/v1/e2e/spec_decode/test_spec_decode.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
- pytest -v -s v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
+1
View File
@@ -5,6 +5,7 @@ steps:
- label: Basic Models Tests (Initialization)
key: basic-models-tests-initialization
timeout_in_minutes: 45
device: h200_18gb
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -5,6 +5,7 @@ steps:
- label: Language Models Tests (Standard)
key: language-models-tests-standard
timeout_in_minutes: 25
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/language
+5 -5
View File
@@ -84,19 +84,19 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: DFlash Speculators Correctness
key: dflash-speculators-correctness
timeout_in_minutes: 30
- label: Speculators Correctness
key: speculators-correctness
timeout_in_minutes: 60
device: h100
optional: true
num_devices: 1
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/model_executor/models/qwen3_dflash.py
- tests/v1/spec_decode/test_speculators_dflash.py
- tests/v1/spec_decode/test_speculators_correctness.py
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
- label: Spec Decode MTP hybrid (B200)
timeout_in_minutes: 30
+3
View File
@@ -9,6 +9,9 @@ PATH=${cuda_home}/bin:$PATH
LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH
# Install requirements
if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then
sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt
fi
$python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt
# Limit the number of parallel jobs to avoid OOM
+71 -70
View File
@@ -369,7 +369,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
FetchContent_MakeAvailable(cutlass)
list(APPEND VLLM_EXT_SRC
"csrc/quantization/awq/gemm_kernels.cu"
"csrc/cutlass_extensions/common.cpp")
set_gencode_flags_for_srcs(
@@ -501,46 +500,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
" in CUDA target architectures")
endif()
# Only build AllSpark kernels if we are building for at least some compatible archs.
cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
if (ALLSPARK_ARCHS)
set(ALLSPARK_SRCS
"csrc/quantization/gptq_allspark/allspark_repack.cu"
"csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
set_gencode_flags_for_srcs(
SRCS "${ALLSPARK_SRCS}"
CUDA_ARCHS "${ALLSPARK_ARCHS}")
list(APPEND VLLM_EXT_SRC "${ALLSPARK_SRCS}")
message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
else()
message(STATUS "Not building AllSpark kernels as no compatible archs found"
" in CUDA target architectures")
endif()
# CUTLASS MLA Archs and flags
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
set(SRCS
"csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${MLA_ARCHS}")
list(APPEND VLLM_EXT_SRC "${SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
# Add MLA-specific include directories only to MLA source files
set_source_files_properties(${SRCS}
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
else()
message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
# clear MLA_ARCHS
set(MLA_ARCHS)
endif()
# Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
@@ -568,24 +527,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif()
endif()
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
set_gencode_flags_for_srcs(
SRCS "${DSV3_FUSED_A_GEMM_SRC}"
CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
else()
message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
"in CUDA target architectures.")
endif()
#
# Machete kernels
@@ -657,16 +598,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif()
# Hadacore kernels
cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
if(HADACORE_ARCHS)
set(SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${HADACORE_ARCHS}")
list(APPEND VLLM_EXT_SRC "${SRCS}")
message(STATUS "Building hadacore")
endif()
# if CUDA endif
endif()
@@ -716,7 +647,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/permute_cols.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu")
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
endif()
if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -725,6 +657,40 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
CUDA_ARCHS "${CUDA_ARCHS}")
endif()
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
set(SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
else()
message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
"in CUDA target architectures.")
endif()
# Only build AllSpark kernels if we are building for at least some compatible archs.
cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
if (ALLSPARK_ARCHS)
set(SRCS
"csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu"
"csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${ALLSPARK_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
else()
message(STATUS "Not building AllSpark kernels as no compatible archs found"
" in CUDA target architectures")
endif()
#
# CUTLASS scaled_mm kernels (moved from _C to _C_stable_libtorch)
#
@@ -1034,6 +1000,41 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif()
endif()
# CUTLASS MLA Archs and flags
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
set(SRCS
"csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${MLA_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
# Add MLA-specific include directories only to MLA source files
set_source_files_properties(${SRCS}
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
else()
message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
# clear MLA_ARCHS
set(MLA_ARCHS)
endif()
# Hadacore kernels
cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
if(HADACORE_ARCHS)
set(SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
set_gencode_flags_for_srcs(
SRCS "${SRCS}"
CUDA_ARCHS "${HADACORE_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
message(STATUS "Building hadacore")
endif()
message(STATUS "Enabling C_stable extension.")
define_extension_target(
_C_stable_libtorch
+2 -2
View File
@@ -27,10 +27,10 @@ from vllm.model_executor.layers.fused_moe.config import (
RoutingMethodType,
_get_config_dtype_str,
)
from vllm.model_executor.layers.fused_moe.fused_moe import *
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
TritonOrDeepGemmExperts,
)
from vllm.model_executor.layers.fused_moe.fused_moe import *
from vllm.transformers_utils.config import get_config
from vllm.triton_utils import triton
from vllm.utils.argparse_utils import FlexibleArgumentParser
+69 -37
View File
@@ -53,48 +53,80 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
if(DEEPGEMM_ARCHS)
message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")
find_package(CUDAToolkit REQUIRED)
#
# Build the _C pybind11 extension from DeepGEMM's C++ source.
# This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
# DeepGEMM integration notes
# --------------------------
# We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle a
# `_C.cpython-X.Y-*.so` for every CPython in `requires-python`. The
# per-Python build is delegated to tools/build_deepgemm_C.py.
#
# Why per-Python: DeepGEMM's binding uses PYBIND11_MODULE, which links
# private CPython symbols — a single `_C.abi3.so` is not viable today
# (see #41476 / #41512 for the failed attempt).
#
# TODOs (tracked in vllm-project/vllm#42431):
# - Replace DeepGEMM's pybind11 binding with a TORCH_LIBRARY + shim
# binding (cf. vllm-flash-attention/csrc/common/pytorch_shim.h) to
# collapse to one `_C.abi3.so`. Needs either an upstream change or
# a maintained binding fork in vLLM.
# - AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop
# the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime
# requirement.
#
Python_add_library(_deep_gemm_C MODULE WITH_SOABI
"${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")
# The pybind11 module name must be _C to match DeepGEMM's Python imports.
set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C")
# DEEPGEMM_PYTHON_INTERPRETERS: ":"-separated target Python paths.
# Empty/unset → fall back to the build interpreter (editable installs).
# (Empty-but-set env vars test as DEFINED in cmake — treat as unset.)
if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "")
string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}")
else()
set(_dg_pythons "${Python_EXECUTABLE}")
endif()
message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
target_compile_definitions(_deep_gemm_C PRIVATE
"-DTORCH_EXTENSION_NAME=_C")
# add_custom_command does no implicit header scanning; glob explicitly so
# header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild.
file(GLOB_RECURSE _dg_headers
"${deepgemm_SOURCE_DIR}/csrc/*.h"
"${deepgemm_SOURCE_DIR}/csrc/*.hpp"
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")
target_include_directories(_deep_gemm_C PRIVATE
"${deepgemm_SOURCE_DIR}/csrc"
"${deepgemm_SOURCE_DIR}/deep_gemm/include"
"${deepgemm_SOURCE_DIR}/third-party/cutlass/include"
"${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include"
"${deepgemm_SOURCE_DIR}/third-party/fmt/include")
target_compile_options(_deep_gemm_C PRIVATE
$<$<COMPILE_LANGUAGE:CXX>:-O3>
$<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
$<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
# torch_python is required because DeepGEMM uses pybind11 type casters
# for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which
# use torch::Library custom ops.
find_library(TORCH_PYTHON_LIBRARY torch_python
PATHS "${TORCH_INSTALL_PREFIX}/lib"
REQUIRED)
target_link_libraries(_deep_gemm_C PRIVATE
torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}"
CUDA::cudart CUDA::nvrtc)
# Install the shared library into the vendored package directory
install(TARGETS _deep_gemm_C
LIBRARY DESTINATION vllm/third_party/deep_gemm
COMPONENT _deep_gemm_C)
set(_dg_markers)
set(_dg_seen_soabis)
foreach(_pybin IN LISTS _dg_pythons)
execute_process(
COMMAND "${_pybin}" -c
"import sysconfig; print(sysconfig.get_config_var('SOABI'))"
OUTPUT_VARIABLE _dg_soabi
OUTPUT_STRIP_TRAILING_WHITESPACE
COMMAND_ERROR_IS_FATAL ANY)
# Dedup interpreters that resolve to the same CPython.
if(_dg_soabi IN_LIST _dg_seen_soabis)
continue()
endif()
list(APPEND _dg_seen_soabis "${_dg_soabi}")
set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}")
set(_dg_marker "${_dg_dir}/.built")
add_custom_command(
OUTPUT "${_dg_marker}"
COMMAND "${Python_EXECUTABLE}"
"${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
"${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}"
COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}"
DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
"${deepgemm_SOURCE_DIR}/csrc/python_api.cpp"
${_dg_headers}
COMMENT "Building DeepGEMM _C for ${_pybin}"
VERBATIM)
list(APPEND _dg_markers "${_dg_marker}")
install(DIRECTORY "${_dg_dir}/"
DESTINATION vllm/third_party/deep_gemm
COMPONENT _deep_gemm_C
FILES_MATCHING PATTERN "_C.cpython-*.so")
endforeach()
add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers})
#
# Vendor DeepGEMM Python package files
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG f5bc33cfc02c744d24a2e9d50e6db656de40611c
GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+26 -18
View File
@@ -1,7 +1,13 @@
#pragma once
// For TORCH_CHECK
#include <torch/library.h>
#include <cstdint>
#include <string>
#include <tuple>
#include <utility>
#include <variant>
// For STD_TORCH_CHECK
#include <torch/headeronly/util/Exception.h>
namespace vllm {
@@ -45,7 +51,7 @@ class ScalarType {
// IEEE 754 compliant floating point type
static constexpr ScalarType float_IEEE754(uint8_t exponent,
uint8_t mantissa) {
TORCH_CHECK(mantissa > 0 && exponent > 0);
STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754);
}
@@ -53,11 +59,12 @@ class ScalarType {
static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
bool finite_values_only,
NanRepr nan_repr) {
TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
TORCH_CHECK(mantissa > 0 && exponent > 0);
TORCH_CHECK(nan_repr != NAN_IEEE_754,
"use `float_IEEE754` constructor for floating point types that "
"follow IEEE 754 conventions");
STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
STD_TORCH_CHECK(
nan_repr != NAN_IEEE_754,
"use `float_IEEE754` constructor for floating point types that "
"follow IEEE 754 conventions");
return ScalarType(exponent, mantissa, true, 0, finite_values_only,
nan_repr);
}
@@ -176,8 +183,8 @@ class ScalarType {
private:
double _floating_point_max() const {
TORCH_CHECK(mantissa <= 52 && exponent <= 11,
"Cannot represent max/min as a double for type ", str());
STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11,
"Cannot represent max/min as a double for type ", str());
uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
@@ -186,8 +193,8 @@ class ScalarType {
uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
TORCH_CHECK(exponent < 11,
"Cannot represent max/min as a double for type ", str());
STD_TORCH_CHECK(exponent < 11,
"Cannot represent max/min as a double for type ", str());
max_exponent += 1;
}
@@ -216,16 +223,17 @@ class ScalarType {
if (is_floating_point()) {
return {_floating_point_max()};
} else {
TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
"Cannot represent max as a int64_t");
STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
"Cannot represent max as a int64_t");
return {(int64_t(1) << mantissa) - 1};
}
}
constexpr std::variant<int64_t, double> _raw_min() const {
if (is_floating_point()) {
TORCH_CHECK(is_signed(),
"We currently assume all floating point types are signed");
STD_TORCH_CHECK(
is_signed(),
"We currently assume all floating point types are signed");
constexpr uint64_t sign_bit_double = (uint64_t(1) << 63);
double max = _floating_point_max();
@@ -233,8 +241,8 @@ class ScalarType {
uint64_t min_raw = max_raw | sign_bit_double;
return {*reinterpret_cast<double*>(&min_raw)};
} else {
TORCH_CHECK(!is_signed() || size_bits() <= 64,
"Cannot represent min as a int64_t");
STD_TORCH_CHECK(!is_signed() || size_bits() <= 64,
"Cannot represent min as a int64_t");
if (is_signed()) {
// set the top bit to 1 (i.e. INT64_MIN) and the rest to 0
// then perform an arithmetic shift right to set all the bits above
@@ -18,13 +18,12 @@ limitations under the License.
* Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
* by Alcanderian JieXin Liang
*/
#include "core/registration.h"
#include "libtorch_stable/torch_utils.h"
#include <torch/csrc/stable/library.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cutlass/cutlass.h>
#include <cutlass/kernel_hardware_info.h>
#include <torch/all.h>
#include <cute/tensor.hpp>
#include <iostream>
@@ -35,27 +34,27 @@ limitations under the License.
// clang-format off
#if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
void sm100_cutlass_mla_decode(
torch::Tensor const& out,
torch::Tensor const& lse,
torch::Tensor const& q_nope,
torch::Tensor const& q_pe,
torch::Tensor const& kv_c_and_k_pe_cache,
torch::Tensor const& seq_lens,
torch::Tensor const& page_table,
torch::Tensor const& workspace,
torch::stable::Tensor const& out,
torch::stable::Tensor const& lse,
torch::stable::Tensor const& q_nope,
torch::stable::Tensor const& q_pe,
torch::stable::Tensor const& kv_c_and_k_pe_cache,
torch::stable::Tensor const& seq_lens,
torch::stable::Tensor const& page_table,
torch::stable::Tensor const& workspace,
double sm_scale,
int64_t num_kv_splits) {
TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
}
int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
}
#else
#define CUTLASS_CHECK(status) \
{ \
cutlass::Status error = status; \
TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
STD_TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
}
using namespace cute;
@@ -100,23 +99,23 @@ struct MlaSm100 {
template <typename T>
typename T::Fmha::Arguments args_from_options(
at::Tensor const& out,
at::Tensor const& lse,
at::Tensor const& q_nope,
at::Tensor const& q_pe,
at::Tensor const& kv_c_and_k_pe_cache,
at::Tensor const& seq_lens,
at::Tensor const& page_table,
torch::stable::Tensor const& out,
torch::stable::Tensor const& lse,
torch::stable::Tensor const& q_nope,
torch::stable::Tensor const& q_pe,
torch::stable::Tensor const& kv_c_and_k_pe_cache,
torch::stable::Tensor const& seq_lens,
torch::stable::Tensor const& page_table,
double sm_scale,
int64_t num_kv_splits) {
cutlass::KernelHardwareInfo hw_info;
hw_info.device_id = q_nope.device().index();
hw_info.device_id = q_nope.get_device_index();
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
int batches = q_nope.sizes()[0];
int page_count_per_seq = page_table.sizes()[1];
int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
int page_size = kv_c_and_k_pe_cache.sizes()[1];
int batches = q_nope.size(0);
int page_count_per_seq = page_table.size(1);
int page_count_total = kv_c_and_k_pe_cache.size(0);
int page_size = kv_c_and_k_pe_cache.size(1);
int max_seq_len = page_size * page_count_per_seq;
using TileShapeH = typename T::TileShapeH;
using TileShapeD = typename T::TileShapeD;
@@ -186,14 +185,14 @@ typename T::Fmha::Arguments args_from_options(
template <typename Element, typename ElementOut, bool IsPaged128, typename PersistenceOption>
void runMla(
at::Tensor const& out,
at::Tensor const& lse,
at::Tensor const& q_nope,
at::Tensor const& q_pe,
at::Tensor const& kv_c_and_k_pe_cache,
at::Tensor const& seq_lens,
at::Tensor const& page_table,
at::Tensor const& workspace,
torch::stable::Tensor const& out,
torch::stable::Tensor const& lse,
torch::stable::Tensor const& q_nope,
torch::stable::Tensor const& q_pe,
torch::stable::Tensor const& kv_c_and_k_pe_cache,
torch::stable::Tensor const& seq_lens,
torch::stable::Tensor const& page_table,
torch::stable::Tensor const& workspace,
double sm_scale,
int64_t num_kv_splits,
cudaStream_t stream) {
@@ -220,37 +219,37 @@ void runMla(
}()
void sm100_cutlass_mla_decode(
torch::Tensor const& out,
torch::Tensor const& lse,
torch::Tensor const& q_nope,
torch::Tensor const& q_pe,
torch::Tensor const& kv_c_and_k_pe_cache,
torch::Tensor const& seq_lens,
torch::Tensor const& page_table,
torch::Tensor const& workspace,
torch::stable::Tensor const& out,
torch::stable::Tensor const& lse,
torch::stable::Tensor const& q_nope,
torch::stable::Tensor const& q_pe,
torch::stable::Tensor const& kv_c_and_k_pe_cache,
torch::stable::Tensor const& seq_lens,
torch::stable::Tensor const& page_table,
torch::stable::Tensor const& workspace,
double sm_scale,
int64_t num_kv_splits) {
auto in_dtype = q_nope.dtype();
at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device());
const int page_size = kv_c_and_k_pe_cache.sizes()[1];
auto in_dtype = q_nope.scalar_type();
torch::stable::accelerator::DeviceGuard device_guard(q_nope.get_device_index());
const cudaStream_t stream = get_current_cuda_stream(q_nope.get_device_index());
const int page_size = kv_c_and_k_pe_cache.size(1);
// NOTE(alcanderian): IsPersistent has bug with manual split_kv.
// Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8)
// Maybe per batch split kv will fix this.
DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
if (in_dtype == at::ScalarType::Half) {
if (in_dtype == torch::headeronly::ScalarType::Half) {
runMla<cutlass::half_t, cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
} else if (in_dtype == at::ScalarType::BFloat16) {
} else if (in_dtype == torch::headeronly::ScalarType::BFloat16) {
runMla<cutlass::bfloat16_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
} else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
} else if (in_dtype == torch::headeronly::ScalarType::Float8_e4m3fn) {
runMla<cutlass::float_e4m3_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
} else {
TORCH_CHECK(false, "Unsupported input data type of MLA");
STD_TORCH_CHECK(false, "Unsupported input data type of MLA");
}
return true;
});
@@ -280,12 +279,12 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
#endif
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
m.impl("sm100_cutlass_mla_decode", &sm100_cutlass_mla_decode);
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("sm100_cutlass_mla_decode", TORCH_BOX(&sm100_cutlass_mla_decode));
}
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CatchAll, m) {
m.impl("sm100_cutlass_mla_get_workspace_size", &sm100_cutlass_mla_get_workspace_size);
STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
m.impl("sm100_cutlass_mla_get_workspace_size", TORCH_BOX(&sm100_cutlass_mla_get_workspace_size));
}
// clang-format on
@@ -20,13 +20,15 @@
* limitations under the License.
*/
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <torch/all.h>
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>
#include "core/registration.h"
#include "libtorch_stable/torch_utils.h"
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cstdlib>
#include <mutex>
@@ -34,7 +36,7 @@
namespace {
inline int getSMVersion() {
auto* props = at::cuda::getCurrentDeviceProperties();
auto* props = get_device_prop();
return props->major * 10 + props->minor;
}
@@ -700,37 +702,40 @@ template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
cudaStream_t);
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
torch::Tensor const& mat_b) {
TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
void dsv3_fused_a_gemm(torch::stable::Tensor& output,
torch::stable::Tensor const& mat_a,
torch::stable::Tensor const& mat_b) {
STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
int const num_tokens = mat_a.size(0);
int const hd_in = mat_a.size(1);
int const hd_out = mat_b.size(1);
constexpr int kHdIn = 7168;
constexpr int kHdOut = 2112;
TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
"required 1 <= mat_a.shape[0] <= 16")
TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
TORCH_CHECK(output.size(0) == num_tokens,
"required output.shape[0] == mat_a.shape[0]")
TORCH_CHECK(output.size(1) == hd_out,
"required output.shape[1] == mat_b.shape[1]")
STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
"required 1 <= mat_a.shape[0] <= 16");
STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
STD_TORCH_CHECK(output.size(0) == num_tokens,
"required output.shape[0] == mat_a.shape[0]");
STD_TORCH_CHECK(output.size(1) == hd_out,
"required output.shape[1] == mat_b.shape[1]");
TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
mat_b.scalar_type() == torch::kBFloat16,
"Only BFloat16 input dtype is supported")
TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
"Only BFloat16 output dtype is supported")
STD_TORCH_CHECK(
mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"Only BFloat16 input dtype is supported");
STD_TORCH_CHECK(
output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"Only BFloat16 output dtype is supported");
TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
auto stream = get_current_cuda_stream(mat_a.get_device_index());
if (num_tokens <= 8) {
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
@@ -746,6 +751,6 @@ void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
}
}
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
m.impl("dsv3_fused_a_gemm", &dsv3_fused_a_gemm);
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("dsv3_fused_a_gemm", TORCH_BOX(&dsv3_fused_a_gemm));
}
+22
View File
@@ -143,4 +143,26 @@ void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
const torch::stable::Tensor& expert_offsets,
const torch::stable::Tensor& sf_offsets);
// AWQ ops
torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
torch::stable::Tensor _kernel,
torch::stable::Tensor _scaling_factors,
torch::stable::Tensor _zeros,
int64_t split_k_iters);
torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
torch::stable::Tensor _scaling_factors,
torch::stable::Tensor _zeros,
int64_t split_k_iters, int64_t thx,
int64_t thy);
// DSV3 fused A GEMM: conditionally compiled so declaration and impl
// registration are in the source file (dsv3_fused_a_gemm.cu)
// AllSpark ops: declarations are in the source files
// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
#endif
torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
bool inplace);
@@ -7,10 +7,11 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
}
*/
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/csrc/stable/ops.h>
#include "libtorch_stable/torch_utils.h"
#include "dequantize.cuh"
#include "libtorch_stable/quantization/awq/dequantize.cuh"
#include <cuda_fp16.h>
@@ -410,10 +411,11 @@ __global__ void __launch_bounds__(64)
} // namespace awq
} // namespace vllm
torch::Tensor awq_dequantize(torch::Tensor _kernel,
torch::Tensor _scaling_factors,
torch::Tensor _zeros, int64_t split_k_iters,
int64_t thx, int64_t thy) {
torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
torch::stable::Tensor _scaling_factors,
torch::stable::Tensor _zeros,
int64_t split_k_iters, int64_t thx,
int64_t thy) {
int in_c = _kernel.size(0);
int qout_c = _kernel.size(1);
int out_c = qout_c * 8;
@@ -437,23 +439,24 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
y_blocks = (int)(in_c / 8);
}
const at::cuda::OptionalCUDAGuard device_guard(device_of(_scaling_factors));
const torch::stable::accelerator::DeviceGuard device_guard(
_scaling_factors.get_device_index());
auto options = torch::TensorOptions()
.dtype(_scaling_factors.dtype())
.device(_scaling_factors.device());
at::Tensor _de_kernel = torch::empty({in_c, out_c}, options);
auto _de_kernel =
torch::stable::empty({in_c, out_c}, _scaling_factors.scalar_type(),
std::nullopt, _scaling_factors.device());
auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
auto de_kernel = reinterpret_cast<half*>(_de_kernel.data_ptr<at::Half>());
auto scaling_factors =
reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
auto de_kernel = reinterpret_cast<half*>(
_de_kernel.mutable_data_ptr<torch::headeronly::Half>());
auto scaling_factors = reinterpret_cast<half*>(
_scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
dim3 num_blocks(x_blocks, y_blocks);
dim3 threads_per_block(x_thread, y_thread);
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const cudaStream_t stream = get_current_cuda_stream();
vllm::awq::dequantize_weights<<<num_blocks, threads_per_block, 0, stream>>>(
kernel, scaling_factors, zeros, de_kernel, G);
@@ -466,27 +469,30 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
// zeros: IC // G, OC // 8 [int32] -> cast to IC // G, OC [uint4b]
// assume that batch_size < 16 for now
torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
torch::Tensor _scaling_factors, torch::Tensor _zeros,
int64_t split_k_iters) {
torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
torch::stable::Tensor _kernel,
torch::stable::Tensor _scaling_factors,
torch::stable::Tensor _zeros,
int64_t split_k_iters) {
int num_in_feats = _in_feats.size(0);
int num_in_channels = _in_feats.size(1);
const at::cuda::OptionalCUDAGuard device_guard(device_of(_in_feats));
const torch::stable::accelerator::DeviceGuard device_guard(
_in_feats.get_device_index());
auto options = torch::TensorOptions()
.dtype(_in_feats.dtype())
.device(_in_feats.device());
at::Tensor _out_feats =
torch::empty({split_k_iters, num_in_feats, _kernel.size(1) * 8}, options);
auto _out_feats = torch::stable::empty(
{split_k_iters, num_in_feats, _kernel.size(1) * 8},
_in_feats.scalar_type(), std::nullopt, _in_feats.device());
int num_out_feats = _out_feats.size(-2);
int num_out_channels = _out_feats.size(-1);
auto in_feats = reinterpret_cast<half*>(_in_feats.data_ptr<at::Half>());
auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
auto scaling_factors =
reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
auto in_feats = reinterpret_cast<half*>(
_in_feats.mutable_data_ptr<torch::headeronly::Half>());
auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
auto out_feats = reinterpret_cast<half*>(
_out_feats.mutable_data_ptr<torch::headeronly::Half>());
auto scaling_factors = reinterpret_cast<half*>(
_scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
int group_size = num_in_channels / _scaling_factors.size(0);
if (num_out_channels % 64 != 0)
@@ -498,7 +504,7 @@ torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
if (num_out_channels % group_size != 0)
throw std::invalid_argument("OC is not multiple of Group size");
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const cudaStream_t stream = get_current_cuda_stream();
if (num_out_channels % 128 == 0) {
int j_factors1 = num_out_channels / 128 / 1;
dim3 num_blocks((num_out_feats + 16 - 1) / 16 * j_factors1 * split_k_iters);
@@ -522,5 +528,5 @@ torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
num_in_feats, num_in_channels, num_out_channels, out_feats);
}
return _out_feats.sum(0);
return torch::stable::sum(_out_feats, 0);
}
@@ -1,20 +1,28 @@
#include "allspark_utils.cuh"
#include <torch/all.h>
#include "core/registration.h"
#include <torch/csrc/stable/c/shim.h>
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>
#include <cublas_v2.h>
at::Tensor as_g_workspace;
#include "core/registration.h"
#include "libtorch_stable/torch_utils.h"
torch::stable::Tensor as_g_workspace;
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
torch::Tensor allspark_w8a16_gemm(
torch::Tensor const& a, torch::Tensor const& b_qweight,
torch::Tensor const& b_scales, std::optional<torch::Tensor> const& b_qzeros,
int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
torch::stable::Tensor allspark_w8a16_gemm(
torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
torch::stable::Tensor const& b_scales,
std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
int64_t group_size, int64_t sm_count, int64_t sm_version,
int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
TORCH_CHECK_NOT_IMPLEMENTED(
STD_TORCH_CHECK_NOT_IMPLEMENTED(
false, "allspark_w8a16_gemm(..) requires CUDA_ARCH >= 8.0");
return torch::empty({1, 1});
return torch::stable::empty({1, 1});
}
#else
@@ -848,8 +856,8 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
const int N_32align, const int N,
const int K, const int GroupSize,
cudaStream_t stream) {
TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
"Unsupported shape");
STD_TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
"Unsupported shape");
if (GroupSize == -1) {
const int BLOCK = 128;
dim3 grid(N_32align / 32, ((K / 16) + 3) / 4);
@@ -859,7 +867,7 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
}
// TODO: Support SubChannel
else {
TORCH_CHECK(false, "Now only support PerChannel");
STD_TORCH_CHECK(false, "Now only support PerChannel");
}
}
@@ -916,24 +924,27 @@ void allspark_qgemm_w8a16_perc_ampere(
} // namespace allspark
torch::Tensor allspark_w8a16_gemm(
torch::Tensor const& a, torch::Tensor const& b_qweight,
torch::Tensor const& b_scales, std::optional<torch::Tensor> const& b_qzeros,
int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
torch::stable::Tensor allspark_w8a16_gemm(
torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
torch::stable::Tensor const& b_scales,
std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
int64_t group_size, int64_t sm_count, int64_t sm_version,
int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
// Verify device and strides
TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
if (has_zp) {
TORCH_CHECK(b_qzeros.value().device().is_cuda(), "b_qzeros is not on GPU");
TORCH_CHECK(b_qzeros.value().is_contiguous(), "b_qzeros is not contiguous");
STD_TORCH_CHECK(b_qzeros.value().device().is_cuda(),
"b_qzeros is not on GPU");
STD_TORCH_CHECK(b_qzeros.value().is_contiguous(),
"b_qzeros is not contiguous");
}
int m = a.size(0);
@@ -941,16 +952,17 @@ torch::Tensor allspark_w8a16_gemm(
int k = a.size(1);
// Verify shape
TORCH_CHECK(b_qweight.size(0) == n_32align,
"Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
", n_32align = ", n_32align);
TORCH_CHECK(b_qweight.size(1) == k,
"Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
", k = ", k);
STD_TORCH_CHECK(b_qweight.size(0) == n_32align,
"Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
", n_32align = ", n_32align);
STD_TORCH_CHECK(b_qweight.size(1) == k,
"Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
", k = ", k);
TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
STD_TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
const torch::stable::accelerator::DeviceGuard device_guard(
a.get_device_index());
const void* a_ptr = reinterpret_cast<const void*>(a.data_ptr());
const uint8_t* b_ptr = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
const void* b_scale_ptr = reinterpret_cast<const void*>(b_scales.data_ptr());
@@ -959,12 +971,12 @@ torch::Tensor allspark_w8a16_gemm(
b_zero_ptr = reinterpret_cast<const void*>(b_qzeros.value().data_ptr());
}
auto c_options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
torch::Tensor c = torch::empty({m, n}, c_options);
void* c_ptr = reinterpret_cast<void*>(c.data_ptr());
auto c =
torch::stable::empty({m, n}, a.scalar_type(), std::nullopt, a.device());
void* c_ptr = reinterpret_cast<void*>(c.mutable_data_ptr());
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
cudaStream_t stream = get_current_cuda_stream();
cublasHandle_t handle = get_current_cuda_blas_handle();
allspark::BlockTileSplitkParams fused_gemm_params;
@@ -976,14 +988,15 @@ torch::Tensor allspark_w8a16_gemm(
m, n, k, sm_count, fused_gemm_params);
}
auto ws_options = torch::TensorOptions().dtype(at::kChar).device(a.device());
if (as_g_workspace.numel() <
ws_size) { // ws_options: kChar, so numel() is bytes
as_g_workspace = torch::empty({long(ws_size)}, ws_options);
as_g_workspace = torch::stable::empty({static_cast<int64_t>(ws_size)},
torch::headeronly::ScalarType::Char,
std::nullopt, a.device());
}
void* ws = reinterpret_cast<void*>(as_g_workspace.data_ptr());
if (a.dtype() == at::ScalarType::Half) {
if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
allspark::allspark_qgemm_w8a16_perc_ampere<__half, uint8_t>(
reinterpret_cast<const __half*>(a_ptr), b_ptr,
reinterpret_cast<const __half*>(b_scale_ptr),
@@ -991,7 +1004,7 @@ torch::Tensor allspark_w8a16_gemm(
reinterpret_cast<__half*>(c_ptr), m, n_32align, n, k, ws,
fused_gemm_params, group_size, CUBLAS_M_THRESHOLD, sm_version, stream,
handle);
} else if (a.dtype() == at::ScalarType::BFloat16) {
} else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
allspark::allspark_qgemm_w8a16_perc_ampere<__nv_bfloat16, uint8_t>(
reinterpret_cast<const __nv_bfloat16*>(a_ptr), b_ptr,
reinterpret_cast<const __nv_bfloat16*>(b_scale_ptr),
@@ -1006,6 +1019,6 @@ torch::Tensor allspark_w8a16_gemm(
#endif
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
m.impl("allspark_w8a16_gemm", &allspark_w8a16_gemm);
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("allspark_w8a16_gemm", TORCH_BOX(&allspark_w8a16_gemm));
}
@@ -1,6 +1,11 @@
#include "allspark_utils.cuh"
#include <torch/all.h>
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>
#include "core/registration.h"
#include "libtorch_stable/torch_utils.h"
namespace allspark {
@@ -99,36 +104,40 @@ void rearrange_kn_weight_as_n32k16_order_ldg16(
} // namespace allspark
void rearrange_kn_weight_as_n32k16_order(
torch::Tensor const& b_qweight, torch::Tensor const& b_scales,
std::optional<torch::Tensor> const& b_zeros, bool has_zp,
torch::Tensor& b_qweight_reorder, torch::Tensor& b_scales_reorder,
std::optional<torch::Tensor> const& b_zeros_reorder, const int64_t K,
const int64_t N, const int64_t N_32align) {
torch::stable::Tensor const& b_qweight,
torch::stable::Tensor const& b_scales,
std::optional<torch::stable::Tensor> const& b_zeros, bool has_zp,
torch::stable::Tensor& b_qweight_reorder,
torch::stable::Tensor& b_scales_reorder,
std::optional<torch::stable::Tensor> const& b_zeros_reorder,
const int64_t K, const int64_t N, const int64_t N_32align) {
// Verify device and strides
TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
"b_qweight_reorder is not on GPU");
TORCH_CHECK(b_qweight_reorder.is_contiguous(),
"b_qweight_reorder is not contiguous");
STD_TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
"b_qweight_reorder is not on GPU");
STD_TORCH_CHECK(b_qweight_reorder.is_contiguous(),
"b_qweight_reorder is not contiguous");
TORCH_CHECK(b_scales_reorder.device().is_cuda(),
"b_scales_reorder is not on GPU");
TORCH_CHECK(b_scales_reorder.is_contiguous(),
"b_scales_reorder is not contiguous");
STD_TORCH_CHECK(b_scales_reorder.device().is_cuda(),
"b_scales_reorder is not on GPU");
STD_TORCH_CHECK(b_scales_reorder.is_contiguous(),
"b_scales_reorder is not contiguous");
if (has_zp) {
TORCH_CHECK(b_zeros.value().device().is_cuda(), "b_zeros is not on GPU");
TORCH_CHECK(b_zeros.value().is_contiguous(), "b_zeros is not contiguous");
STD_TORCH_CHECK(b_zeros.value().device().is_cuda(),
"b_zeros is not on GPU");
STD_TORCH_CHECK(b_zeros.value().is_contiguous(),
"b_zeros is not contiguous");
TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
"b_zeros_reorder is not on GPU");
TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
"b_zeros_reorder is not contiguous");
STD_TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
"b_zeros_reorder is not on GPU");
STD_TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
"b_zeros_reorder is not contiguous");
}
const uint8_t* matB = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
@@ -136,18 +145,20 @@ void rearrange_kn_weight_as_n32k16_order(
const void* b_zero = has_zp ? b_zeros.value().data_ptr() : nullptr;
uint8_t* matB_reorder =
reinterpret_cast<uint8_t*>(b_qweight_reorder.data_ptr());
void* b_scale_reorder = b_scales_reorder.data_ptr();
void* b_zero_reorder = has_zp ? b_zeros_reorder.value().data_ptr() : nullptr;
reinterpret_cast<uint8_t*>(b_qweight_reorder.mutable_data_ptr());
void* b_scale_reorder = b_scales_reorder.mutable_data_ptr();
void* b_zero_reorder =
has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr;
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (b_scales.dtype() == at::ScalarType::Half) {
cudaStream_t stream = get_current_cuda_stream();
if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(
matB, reinterpret_cast<const __half*>(b_scale),
reinterpret_cast<const __half*>(b_zero), matB_reorder,
reinterpret_cast<__half*>(b_scale_reorder),
reinterpret_cast<__half*>(b_zero_reorder), K, N, N_32align, stream);
} else if (b_scales.dtype() == at::ScalarType::BFloat16) {
} else if (b_scales.scalar_type() ==
torch::headeronly::ScalarType::BFloat16) {
allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__nv_bfloat16>(
matB, reinterpret_cast<const __nv_bfloat16*>(b_scale),
reinterpret_cast<const __nv_bfloat16*>(b_zero), matB_reorder,
@@ -157,7 +168,7 @@ void rearrange_kn_weight_as_n32k16_order(
}
}
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("rearrange_kn_weight_as_n32k16_order",
&rearrange_kn_weight_as_n32k16_order);
TORCH_BOX(&rearrange_kn_weight_as_n32k16_order));
}
@@ -1,13 +1,12 @@
#pragma once
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <iostream>
#include "../marlin/marlin_dtypes.cuh"
#include "quantization/marlin/marlin_dtypes.cuh"
using marlin::MarlinScalarType2;
namespace allspark {
@@ -11,18 +11,16 @@ Redistribution and use in source and binary forms, with or without modification,
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
***********/
#include <torch/all.h>
#include "libtorch_stable/torch_utils.h"
#include "libtorch_stable/dispatch_utils.h"
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/macros.h>
#include <stdint.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <cuda/annotated_ptr>
#include <c10/cuda/CUDAException.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "core/registration.h"
#include "dispatch_utils.h"
namespace hadacore {
@@ -65,12 +63,12 @@ constexpr int launch_configs_big[7][3] = {
};
// a 4x2, b 2x2, c 2x2
template <torch::ScalarType dtype>
template <torch::headeronly::ScalarType dtype>
__device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32& c0, b32& c1){
static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16);
static_assert(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16);
// d, a, b, c
b32 zero = 0;
if constexpr(dtype == torch::ScalarType::Half) {
if constexpr(dtype == torch::headeronly::ScalarType::Half) {
asm (
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};\n\t"
@@ -89,7 +87,7 @@ __device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1,
}
// a 4x2, b 4x2, c 4x2
template <torch::ScalarType dtype>
template <torch::headeronly::ScalarType dtype>
__device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32 b2, b32 b3, b32& c0, b32& c1, b32& c2, b32& c3){
mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b0, b1, c0, c1);
mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b2, b3, c2, c3);
@@ -108,11 +106,11 @@ __device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(b32& a0) {
#define n_p(i) ((val_1n[i] & 0x0000FFFF) | val_1p[i] << 16)
#define n_n(i) ((val_1n[i] & 0x0000FFFF) | val_1n[i] << 16)
template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::ScalarType dtype>
template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::headeronly::ScalarType dtype>
__global__ void __launch_bounds__(32 * warps_per_block, blocks_per_sm)
// a is column major, b is row major
hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
static_assert(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
b32 b_frag_all[num_chunks][4]; // for all chunks, holds matrix fragment (which takes 4 regs of b16x2 * 32 threads)
@@ -162,8 +160,8 @@ hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
constexpr b16 bf16_1p[4] = {0b0011111100110101, 0b0011111100000000, 0b0011111010110101, 0b0011111010000000};
constexpr b16 bf16_1n[4] = {0b1011111100110101, 0b1011111100000000, 0b1011111010110101, 0b1011111010000000};
#define val_type_1p(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
#define val_type_1n(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
#define val_type_1p(i) (((dtype) == torch::headeronly::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
#define val_type_1n(i) (((dtype) == torch::headeronly::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
constexpr b16 val_1p[4] = {val_type_1p(0), val_type_1p(1), val_type_1p(2), val_type_1p(3)};
constexpr b16 val_1n[4] = {val_type_1n(0), val_type_1n(1), val_type_1n(2), val_type_1n(3)};
@@ -684,14 +682,14 @@ constexpr int64_t ceil_div(int64_t a, int64_t b) {
return (a + b - 1) / b;
}
template <torch::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
template <torch::headeronly::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaStream_t stream) {
int64_t shared_size = chunks_per_warp * warps_per_block * 128 * 4;
dim3 block_size = 32 * warps_per_block;
#define CHECK_SHARED_LIM() { \
if (shared_size > 48 * 1024) { \
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
STD_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
} \
} \
@@ -714,10 +712,10 @@ void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaSt
kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
STD_CUDA_KERNEL_LAUNCH_CHECK();
}
template <torch::ScalarType dtype>
template <torch::headeronly::ScalarType dtype>
void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream) {
int64_t num_chunks = numel / 256; // caller required to ensure divisible by 256
// for size 256, use (2, 1)
@@ -764,54 +762,54 @@ void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cu
}
}
template void run_fht<torch::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
template void run_fht<torch::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
template void run_fht<torch::headeronly::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
template void run_fht<torch::headeronly::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
} // namespace hadacore
constexpr bool is_power_of_two(int x) { return x && !(x & (x - 1)); }
torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace) {
torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, bool inplace) {
auto dtype = x.scalar_type();
TORCH_CHECK(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
TORCH_CHECK(x.is_cuda());
STD_TORCH_CHECK(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
STD_TORCH_CHECK(x.is_cuda());
const int had_size = x.size(-1);
TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
STD_TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
"Only power of two Hadamard sizes up to 2^15 are supported, got ", had_size);
const auto res_shape = x.sizes();
x = x.reshape({-1, had_size});
x = torch::stable::reshape(x, {-1, had_size});
auto numel = x.numel();
if (numel % 256 != 0) {
x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 0, 0, (256 - numel % 256) / had_size}));
x = torch::stable::pad(x, {0, 0, 0, (256 - numel % 256) / had_size});
}
if (x.stride(-1) != 1) {
x = x.contiguous();
x = torch::stable::contiguous(x);
}
torch::Tensor out = inplace ? x : torch::empty_like(x);
torch::stable::Tensor out = inplace ? x : torch::stable::empty_like(x);
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index());
auto stream = get_current_cuda_stream();
VLLM_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
auto constexpr SCALAR_TYPE = c10::CppTypeToScalarType<scalar_t>::value;
VLLM_STABLE_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
auto constexpr SCALAR_TYPE = torch::headeronly::CppTypeToScalarType<scalar_t>::value;
hadacore::run_fht<SCALAR_TYPE>(x.data_ptr(), x.data_ptr(), x.numel(), had_size, stream);
});
if (numel % 256 != 0) {
out = out.narrow(0, 0, numel / had_size);
out = torch::stable::narrow(out, 0, 0, numel / had_size);
}
if (inplace && out.data_ptr() != x.data_ptr()) {
x.copy_(out.view(res_shape));
torch::stable::copy_(x, torch::stable::view(out, res_shape));
return x;
}
return out.reshape(res_shape);
return torch::stable::reshape(out, res_shape);
}
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
m.impl("hadacore_transform", &hadacore_transform);
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
m.impl("hadacore_transform", TORCH_BOX(&hadacore_transform));
}
@@ -156,6 +156,17 @@ inline int GetGroupsPerBlock(int64_t num_groups) {
return 1;
}
// Largest divisor of padded_groups_per_row that is <= 16. ry = 16 / kx.
inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) {
if (padded_groups_per_row % 16 == 0) {
return 16;
}
if (padded_groups_per_row % 8 == 0) {
return 8;
}
return 4;
}
void per_token_group_quant_8bit(const torch::stable::Tensor& input,
torch::stable::Tensor& output_q,
torch::stable::Tensor& output_s,
@@ -247,11 +258,11 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
//
// Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
// THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128.
template <typename T, typename DST_DTYPE, int GROUP_SIZE>
template <typename T, typename DST_DTYPE, int GROUP_SIZE, int kGroupsPerBlockX,
int kRowsPerBlock>
__global__ void per_token_group_quant_8bit_packed_register_kernel(
const T* __restrict__ input, void* __restrict__ output_q,
unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
const int groups_per_block, const int padded_groups_per_row,
unsigned int* __restrict__ output_s_packed, const int padded_groups_per_row,
const int groups_per_row, const int mn, const int output_q_mn_extent,
const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
const float min_8bit, const float max_8bit) {
@@ -260,27 +271,25 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
constexpr int VEC_SIZE = 32 / sizeof(T); // 16 for bf16/fp16
static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
"GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
// Each group's 8 threads must live in a single warp octet so the
// 0xffu << (threadIdx.x & 24u) shuffle mask selects exactly the lanes
// that share a group. Requires 32 % THREADS_PER_GROUP == 0 and the host
// to launch num_threads as a multiple of THREADS_PER_GROUP (which it does
// via num_threads = groups_per_block * THREADS_PER_GROUP).
static_assert(32 % THREADS_PER_GROUP == 0,
"THREADS_PER_GROUP must divide warp size for the shuffle "
"mask to be valid");
static_assert(
kGroupsPerBlockX > 0 && (kGroupsPerBlockX & (kGroupsPerBlockX - 1)) == 0,
"kGroupsPerBlockX must be a positive power of 2");
static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
const int lane_id = threadIdx.x % THREADS_PER_GROUP;
const int64_t block_group_id = blockIdx.x * groups_per_block;
const int64_t global_group_id = block_group_id + local_group_id;
if (global_group_id >= num_groups_padded) {
const int sf_k_local = local_group_id % kGroupsPerBlockX;
const int row_local = local_group_id / kGroupsPerBlockX;
const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
if (mn_idx >= tma_aligned_mn) {
return;
}
const int sf_k_idx =
static_cast<int>(global_group_id % padded_groups_per_row);
const int mn_idx = static_cast<int>(global_group_id / padded_groups_per_row);
const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
// Load 16 input elements (32 B) into registers as two adjacent uint4
@@ -443,34 +452,53 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
constexpr int THREADS_PER_GROUP = 8;
const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row;
const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
STD_TORCH_CHECK(padded_groups_per_row % 4 == 0,
"padded_groups_per_row=", padded_groups_per_row,
" is not a multiple of 4.");
const int kx = GetGroupsPerBlockX(padded_groups_per_row);
const int ry = 16 / kx;
const int64_t blocks_x = padded_groups_per_row / kx;
const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
const int num_threads = (kx * ry) * THREADS_PER_GROUP;
// CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
blocks_y <= static_cast<int64_t>(INT32_MAX),
"per_token_group_quant_8bit_packed grid too large: (",
blocks_x, ", ", blocks_y, ").");
auto dst_type = output_q.scalar_type();
const int64_t num_blocks = num_groups_padded / groups_per_block;
const int num_threads = groups_per_block * THREADS_PER_GROUP;
// CUDA caps grid.x at 2^31 - 1; this fits any realistic shape but guard
// against pathological inputs.
STD_TORCH_CHECK(num_blocks <= static_cast<int64_t>(INT32_MAX),
"per_token_group_quant_8bit_packed grid too large: ",
num_blocks, " blocks (max ", INT32_MAX, ").");
#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
do { \
dim3 grid(static_cast<unsigned int>(num_blocks)); \
dim3 block(num_threads); \
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128> \
<<<grid, block, 0, stream>>>( \
static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
num_groups_padded, groups_per_block, \
static_cast<int>(padded_groups_per_row), \
static_cast<int>(groups_per_row), static_cast<int>(mn), \
static_cast<int>(output_q_mn_extent), \
static_cast<int>(tma_aligned_mn), num_scale_elems, \
static_cast<float>(eps), static_cast<float>(min_8bit), \
static_cast<float>(max_8bit)); \
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
do { \
dim3 grid(static_cast<unsigned int>(blocks_x), \
static_cast<unsigned int>(blocks_y)); \
dim3 block(num_threads); \
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
RY> \
<<<grid, block, 0, stream>>>( \
static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
static_cast<int>(padded_groups_per_row), \
static_cast<int>(groups_per_row), static_cast<int>(mn), \
static_cast<int>(output_q_mn_extent), \
static_cast<int>(tma_aligned_mn), num_scale_elems, \
static_cast<float>(eps), static_cast<float>(min_8bit), \
static_cast<float>(max_8bit)); \
} while (0)
#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
do { \
if (kx == 16) { \
LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 16, 1); \
} else if (kx == 8) { \
LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 8, 2); \
} else if (kx == 4) { \
LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 4, 4); \
} else { \
STD_TORCH_CHECK(false, "Unsupported kx value ", kx); \
} \
} while (0)
VLLM_STABLE_DISPATCH_HALF_TYPES(
@@ -488,6 +516,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
}));
#undef LAUNCH_REG_KERNEL
#undef LAUNCH_REG_KERNEL_INST
}
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
+57
View File
@@ -218,7 +218,54 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def(
"cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
"Tensor)");
// SM100 CUTLASS MLA decode
// conditionally compiled so impl registrations are in source file
ops.def(
"sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
" Tensor q_pe, Tensor kv_c_and_k_pe_cache,"
" Tensor seq_lens, Tensor page_table,"
" Tensor workspace, float scale,"
" int num_kv_splits) -> ()");
ops.def(
"sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
" int sm_count, int num_kv_splits) "
"-> int");
// Quantized GEMM for AWQ.
ops.def(
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
"Tensor _zeros, SymInt split_k_iters) -> Tensor");
// Dequantization for AWQ.
ops.def(
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
"Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
// conditionally compiled so impl registration is in source file
ops.def(
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
// reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
ops.def(
"rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
"Tensor? b_zeros, "
"bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
"Tensor!? b_zeros_reorder, "
"int K, int N, int N_32align) -> ()");
// AllSpark quantization ops
ops.def(
"allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
"Tensor? b_qzeros, "
"SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
"CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
#endif
// Hadamard transforms
// conditionally compiled so impl registration is in source file
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
}
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -254,6 +301,16 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
// mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
// W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
// AWQ ops
ops.impl("awq_gemm", TORCH_BOX(&awq_gemm));
ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize));
// DSV3 fused A GEMM: conditionally compiled so impl registration is in
// source file (dsv3_fused_a_gemm.cu)
// AllSpark ops: conditionally compiled so impl registrations are in source
// files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
#endif
}
+66
View File
@@ -6,12 +6,71 @@
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/util/shim_utils.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <deque>
#include <mutex>
#include <string>
#include <vector>
// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED.
#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \
STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__)
// Device properties cache for stable ABI compatibility.
// Uses raw CUDA/HIP APIs instead of ATen functions.
// Using inline ensures a single instance across all translation units.
inline std::deque<std::once_flag> device_flags;
inline std::vector<cudaDeviceProp> device_properties;
inline std::once_flag vectors_init_flag;
inline void do_init_device_vectors() {
int device_count;
cudaError_t err = cudaGetDeviceCount(&device_count);
if (err != cudaSuccess) {
STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " +
std::string(cudaGetErrorString(err)));
}
device_flags.resize(device_count);
device_properties.resize(device_count);
}
inline void initDeviceVectors() {
std::call_once(vectors_init_flag, do_init_device_vectors);
}
inline void initDeviceProperty(int device_index) {
cudaDeviceProp device_prop{};
cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index);
if (err != cudaSuccess) {
STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " +
std::string(cudaGetErrorString(err)));
}
device_properties[device_index] = device_prop;
}
// Get device properties using raw CUDA/HIP APIs (stable ABI compatible).
// Caches results per device so cudaGetDeviceProperties is called at most once
// per device.
inline cudaDeviceProp* get_device_prop() {
initDeviceVectors();
int device_index;
cudaError_t err = cudaGetDevice(&device_index);
if (err != cudaSuccess) {
STD_TORCH_CHECK(
false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err)));
}
STD_TORCH_CHECK(device_index >= 0 && static_cast<size_t>(device_index) <
device_properties.size(),
"CUDA device index " + std::to_string(device_index) +
" out of range [0, " +
std::to_string(device_properties.size()) + ")");
std::call_once(device_flags[device_index], initDeviceProperty, device_index);
return &device_properties[device_index];
}
// Utility to get the current CUDA stream for a given device using stable APIs.
// Returns a cudaStream_t for use in kernel launches.
inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
@@ -20,3 +79,10 @@ inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
return reinterpret_cast<cudaStream_t>(stream_ptr);
}
// Utility to get the current cuBLAS handle using stable APIs.
inline cublasHandle_t get_current_cuda_blas_handle() {
void* blas_handle_ptr = nullptr;
TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr));
return reinterpret_cast<cublasHandle_t>(blas_handle_ptr);
}
-20
View File
@@ -200,19 +200,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
#ifndef USE_ROCM
torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
torch::Tensor _scaling_factors, torch::Tensor _zeros,
int64_t split_k_iters);
torch::Tensor awq_dequantize(torch::Tensor _kernel,
torch::Tensor _scaling_factors,
torch::Tensor _zeros, int64_t split_k_iters,
int64_t thx, int64_t thy);
#endif
torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
int64_t n,
std::optional<at::ScalarType> const& dtype);
@@ -302,8 +289,6 @@ std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
int64_t open_mem_handle(torch::Tensor& mem_handle);
void free_shared_buffer(int64_t buffer);
torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace);
#ifdef USE_ROCM
fptr_t init_custom_qr(int64_t rank, int64_t world_size,
std::optional<int64_t> qr_max_size = std::nullopt);
@@ -315,11 +300,6 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
int64_t qr_max_size();
#endif
#ifndef USE_ROCM
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
torch::Tensor const& mat_b);
#endif
#ifndef USE_ROCM
torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
torch::Tensor const& norm_weight,
+8 -4
View File
@@ -2,10 +2,14 @@
#ifndef _marlin_cuh
#define _marlin_cuh
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
// These torch headers are only needed by non-stable callers (e.g. ops.cu).
// Guard them so that stable ABI targets can still include marlin.cuh
// for Vec, constants, and cp_async helpers without pulling in torch/all.h.
#ifndef TORCH_TARGET_VERSION
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#endif
#include <cuda.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
-51
View File
@@ -263,22 +263,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// Quantization ops
#ifndef USE_ROCM
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
ops.def(
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
// conditionally compiled so impl registration is in source file
// Quantized GEMM for AWQ.
ops.def(
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
"Tensor _zeros, SymInt split_k_iters) -> Tensor");
ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);
// Dequantization for AWQ.
ops.def(
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
"Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
@@ -408,22 +392,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
" -> ()");
// conditionally compiled so impl registration is in source file
// SM100 CUTLASS MLA decode
ops.def(
"sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
" Tensor q_pe, Tensor kv_c_and_k_pe_cache,"
" Tensor seq_lens, Tensor page_table,"
" Tensor workspace, float scale,"
" int num_kv_splits) -> ()");
// conditionally compiled so impl in source file
// SM100 CUTLASS MLA workspace
ops.def(
"sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
" int sm_count, int num_kv_splits) "
"-> int");
// conditionally compiled so impl in source file
#endif
// Quantized GEMM for GPTQ.
@@ -496,26 +464,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"Tensor? last_chunk_indices) -> ()");
ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
// Hadamard transforms
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
#ifndef USE_ROCM
// reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
ops.def(
"rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
"Tensor? b_zeros, "
"bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
"Tensor!? b_zeros_reorder, "
"int K, int N, int N_32align) -> ()");
// conditionally compiled so impl in source file
// AllSpark quantization ops
ops.def(
"allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
"Tensor? b_qzeros, "
"SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
"CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
ops.def(
"minimax_allreduce_rms("
"Tensor input,"
+17 -1
View File
@@ -199,7 +199,10 @@ COPY requirements/cuda.txt requirements/cuda.txt
COPY use_existing_torch.py use_existing_torch.py
COPY pyproject.toml pyproject.toml
RUN --mount=type=cache,target=/root/.cache/uv \
if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
fi \
&& if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
echo "Installing torch nightly..." \
&& uv pip install --python /opt/venv/bin/python3 torch torchaudio torchvision --pre \
--index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
@@ -301,6 +304,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
python3 use_existing_torch.py --prefix; \
fi
# Provision one bare Python per `requires-python` entry; cmake reads
# DEEPGEMM_PYTHON_INTERPRETERS to build DeepGEMM `_C` for each. See
# cmake/external_projects/deepgemm.cmake for the full picture.
COPY tools/setup_deepgemm_pythons.sh tools/build_deepgemm_C.py tools/
ENV DEEPGEMM_VENV_PREFIX=/opt/dgenv
RUN --mount=type=cache,target=/root/.cache/uv \
tools/setup_deepgemm_pythons.sh > /tmp/dg_pythons.txt
# Build the vLLM wheel
# if USE_SCCACHE is set, use sccache to speed up compilation
# AWS credentials mounted at ~/.aws/credentials for sccache S3 auth (optional)
@@ -328,6 +339,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
&& export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" \
&& export VLLM_MAIN_CUDA_VERSION="${VLLM_MAIN_CUDA_VERSION}" \
&& export VLLM_DOCKER_BUILD_CONTEXT=1 \
&& export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) \
&& sccache --show-stats \
&& python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
&& sccache --show-stats; \
@@ -345,6 +357,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" && \
export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" && \
export VLLM_DOCKER_BUILD_CONTEXT=1 && \
export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) && \
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
fi
@@ -616,6 +629,9 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL
COPY requirements/common.txt /tmp/common.txt
COPY requirements/cuda.txt /tmp/requirements-cuda.txt
RUN --mount=type=cache,target=/root/.cache/uv \
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /tmp/requirements-cuda.txt; \
fi && \
uv pip install --system -r /tmp/requirements-cuda.txt \
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
rm /tmp/requirements-cuda.txt /tmp/common.txt
@@ -34,6 +34,7 @@ you may contact the following individuals:
- Simon Mo - <simon.mo@hey.com>
- Russell Bryant - <rbryant@redhat.com>
- Juan Pérez de Algaba - <jperezde@redhat.com>
- Huzaifa Sidhpurwala - <huzaifas@redhat.com>
## Slack Discussion
+1 -1
View File
@@ -2,4 +2,4 @@
vLLM can be deployed with [llm-d](https://github.com/llm-d/llm-d), a Kubernetes-native distributed inference serving stack providing well-lit paths for anyone to serve large generative AI models at scale. It helps achieve the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators and infrastructure providers.
You can use vLLM with llm-d directly by following [this guide](https://llm-d.ai/docs/guide) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
You can use vLLM with llm-d directly by following [the official guides](https://llm-d.ai/docs/guides) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
+2 -1
View File
@@ -85,8 +85,9 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
| Architecture | Models | CG for Image | CG for Video |
| ------------ | ------ | ------------ | ------------ |
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
!!! note
Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
+2 -2
View File
@@ -81,7 +81,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
| Kernel | Input act. format | Quant. types | Quant. format | Activation function | Apply Weight On Input | Modular | Source |
| ------ | ----------------- | ------------ | ------------- | ------------------- | --------------------- | ------- | ------ |
| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.experts.triton_moe.TritonExperts] |
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.BatchedTritonExperts] |
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
@@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
| rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
| cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
| naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.NaiveBatchedExperts] |
| naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] |
!!! info "Table key"
1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -0,0 +1,161 @@
# MooncakeStoreConnector Usage Guide
MooncakeStoreConnector is a KV cache connector that uses [MooncakeDistributedStore](https://github.com/kvcache-ai/Mooncake) as a shared KV cache pool. Unlike `MooncakeConnector` which does direct point-to-point KV transfer between prefiller and decoder, MooncakeStoreConnector enables KV cache offloading to an external distributed store, supporting:
- **CPU offloading**: Extend effective KV cache capacity by offloading to CPU memory via Mooncake's transfer engine.
- **Prefix caching across instances**: Hash-based deduplication allows multiple vLLM instances to share cached KV blocks through the store.
- **Single-node and multi-node deployment**: Works both as a standalone KV cache extension and in disaggregated prefill-decode setups.
## Prerequisites
### Install Mooncake
Install mooncake through pip:
```bash
uv pip install mooncake-transfer-engine
```
Refer to the [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions and building from source.
### Start the Mooncake Master Server
The Mooncake master manages metadata and coordinates the distributed store. Start it before launching vLLM:
```bash
mooncake_master --port 50051
```
Default ports:
- RPC: 50051
Multiple vLLM instances can share the same master server.
### Configure Mooncake
Create a JSON configuration file (e.g., `mooncake_config.json`):
```json
{
"metadata_server": "P2PHANDSHAKE",
"master_server_address": "127.0.0.1:50051",
"global_segment_size": "80GB",
"local_buffer_size": "4GB",
"protocol": "rdma",
"device_name": ""
}
```
- `protocol`: Use `"rdma"` for best performance. `"tcp"` works as a fallback.
- `global_segment_size`: CPU memory contributed to the distributed pool (per GPU).
- `local_buffer_size`: Private buffer for this node's own operations (per GPU).
Set the config path via environment variable:
```bash
export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
```
## Usage
### Single-Node KV Cache Offloading
Use MooncakeStoreConnector to offload KV cache to CPU memory, extending the effective cache size:
```bash
MOONCAKE_CONFIG_PATH=mooncake_config.json \
vllm serve meta-llama/Llama-3.1-8B-Instruct \
--kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}'
```
### Disaggregated Prefill-Decode (XpYd)
In disaggregated prefill-decode mode, use `MultiConnector` to combine `MooncakeConnector` (point-to-point KV transfer) with `MooncakeStoreConnector` (shared KV cache pool). This enables both direct P2P transfer between prefiller and decoder, and cross-instance prefix cache sharing via the distributed store.
**Prefiller Node:**
```bash
MOONCAKE_CONFIG_PATH=mooncake_config.json \
VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \
vllm serve meta-llama/Llama-3.1-8B-Instruct \
--port 8100 \
--kv-transfer-config '{
"kv_connector": "MultiConnector",
"kv_role": "kv_producer",
"kv_connector_extra_config": {
"connectors": [
{
"kv_connector": "MooncakeConnector",
"kv_role": "kv_producer"
},
{
"kv_connector": "MooncakeStoreConnector",
"kv_role": "kv_producer"
}
]
}
}'
```
**Decoder Node:**
```bash
MOONCAKE_CONFIG_PATH=mooncake_config.json \
VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \
vllm serve meta-llama/Llama-3.1-8B-Instruct \
--port 8200 \
--kv-transfer-config '{
"kv_connector": "MultiConnector",
"kv_role": "kv_consumer",
"kv_connector_extra_config": {
"connectors": [
{
"kv_connector": "MooncakeConnector",
"kv_role": "kv_consumer"
},
{
"kv_connector": "MooncakeStoreConnector",
"kv_role": "kv_consumer"
}
]
}
}'
```
**Proxy:**
A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details.
## Environment Variables
| Variable | Description | Default |
| --- | --- | --- |
| `MOONCAKE_CONFIG_PATH` | Path to Mooncake JSON config file | (required) |
| `VLLM_MOONCAKE_BOOTSTRAP_PORT` | Bootstrap port for MooncakeConnector P2P transfer (disagg mode only) | 8998 |
## KV Transfer Config
### KV Role Options
- **kv_producer**: For prefiller instances that store KV caches to the pool.
- **kv_consumer**: For decoder instances that load KV caches from the pool.
- **kv_both**: The instance both stores and loads KV caches. Use this for single-node CPU offloading.
### kv_connector_extra_config
- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
- `discard_partial_chunks` (bool): Discard partial block chunks during store. Default: `true`.
- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
## Notes
### Cross-DP Prefix Cache Hits
When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks:
```bash
PYTHONHASHSEED=0 vllm serve ...
```
Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits.
@@ -2179,6 +2179,92 @@ def run_qwen3_vl_moe(questions: list[str], modality: str) -> ModelRequestData:
)
# Qwen3.5-Dense
def run_qwen3_5(questions: list[str], modality: str) -> ModelRequestData:
model_name = "Qwen/Qwen3.5-4B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
max_num_seqs=5,
mm_processor_kwargs={
"min_pixels": 28 * 28,
"max_pixels": 1280 * 28 * 28,
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
prompts = [
(
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
for question in questions
]
return ModelRequestData(
engine_args=engine_args,
prompts=prompts,
)
# Qwen3.5-MoE
def run_qwen3_5_moe(questions: list[str], modality: str) -> ModelRequestData:
model_name = "Qwen/Qwen3.5-35B-A3B"
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
engine_args = EngineArgs(
model=model_name,
max_model_len=4096,
max_num_seqs=5,
mm_processor_kwargs={
"min_pixels": 28 * 28,
"max_pixels": 1280 * 28 * 28,
"fps": 1,
},
limit_mm_per_prompt=mm_limit,
)
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
if modality == "image":
placeholder = image_placeholder
elif modality == "video":
placeholder = video_placeholder
elif modality == "image+video":
placeholder = image_placeholder + video_placeholder
prompts = [
(
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{placeholder}"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
for question in questions
]
return ModelRequestData(
engine_args=engine_args,
prompts=prompts,
)
# R-4B
def run_r_vl(questions: list[str], modality: str) -> ModelRequestData:
assert modality == "image"
@@ -2442,6 +2528,8 @@ model_example_map = {
"qwen2_5_omni": run_qwen2_5_omni,
"qwen3_vl": run_qwen3_vl,
"qwen3_vl_moe": run_qwen3_vl_moe,
"qwen3_5": run_qwen3_5,
"qwen3_5_moe": run_qwen3_5_moe,
"rvl": run_r_vl,
"skywork_chat": run_skyworkr1v,
"smolvlm": run_smolvlm,
@@ -2460,13 +2548,17 @@ MODELS_NEED_VIDEO_METADATA = [
"molmo2",
"qwen3_vl",
"qwen3_vl_moe",
"qwen3_5",
"qwen3_5_moe",
]
MODELS_SUPPORT_VIT_CUDA_GRAPH = [
"qwen2_5_vl",
"qwen3_vl",
"qwen3_vl_moe",
"qwen2_5_vl",
"qwen3_5",
"qwen3_5_moe",
]
+2 -2
View File
@@ -20,6 +20,6 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
# Required for faster safetensors model loading
fastsafetensors >= 0.2.2
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
nvidia-cutlass-dsl>=4.4.2
# QuACK and Cutlass DSL for FA4 and DeepSeek-V4 kernels
nvidia-cutlass-dsl[cu13]>=4.5.0
quack-kernels>=0.3.3
+1 -3
View File
@@ -1,5 +1,3 @@
lmcache >= 0.3.9
nixl[cu13] >= 0.7.1, <= 0.10.1 # Required for disaggregated prefill
nixl-cu12 >= 0.7.1, <= 0.10.1
nixl-cu13 >= 0.7.1, <= 0.10.1
nixl >= 1.1.0 # Required for disaggregated prefill
mooncake-transfer-engine >= 0.3.8
+1 -1
View File
@@ -15,4 +15,4 @@ torch==2.11.0+xpu
torchaudio
torchvision
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.7/vllm_xpu_kernels-0.1.7-cp38-abi3-manylinux_2_28_x86_64.whl
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.8/vllm_xpu_kernels-0.1.8-cp38-abi3-manylinux_2_28_x86_64.whl
+3
View File
@@ -970,6 +970,9 @@ def get_requirements() -> list[str]:
# vllm-flash-attn is built only for CUDA 12.x.
# Skip for other versions.
continue
if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
# [cu13] extra is the default; strip it on CUDA 12 builds.
req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl")
modified_requirements.append(req)
requirements = modified_requirements
elif _is_hip():
+41
View File
@@ -1686,3 +1686,44 @@ def disable_log_dedup(monkeypatch):
logger._print_warning_once = original_print_warning_once
logger._print_info_once = original_print_info_once
logger._print_debug_once = original_print_debug_once
@pytest.fixture(scope="function")
def fake_vllm_ir(monkeypatch):
"""
Pytest fixture to allow isolated IR op registration in tests.
Replaces IrOp.registry with an empty dict and swaps ``vllm_ir_torch_lib`` for a
fresh ``Library`` with a unique namespace per test (see ``Library.ns``).
Torch keeps registrations for the process lifetime; reusing the fragment
name ``vllm_ir`` and defining the same op string again can segfault. A
random library name keeps each fixture run on a disjoint namespace.
The test Library is kept alive until after monkeypatch teardown so PyTorch's
C++ state is not freed while references may still exist.
Usage:
def test_my_ir_op(fake_vllm_ir):
@vllm.ir.register_op
def my_test_op(x: torch.Tensor) -> torch.Tensor:
return x * 2
result = my_test_op(torch.tensor([1, 2, 3]))
# Registry and library cleaned up automatically after the test
"""
import secrets
from torch.library import Library
from vllm.ir.op import IrOp
monkeypatch.setattr(IrOp, "registry", {})
# Keep a local reference so the Library is not GC'd before monkeypatch
# teardown restores the original reference.
test_lib = Library(f"vllm_ir_{secrets.token_hex(8)}", "FRAGMENT")
monkeypatch.setattr("vllm.ir.op.vllm_ir_torch_lib", test_lib)
yield
del test_lib
@@ -10,6 +10,7 @@ import torch
from tests.kernels.moe.utils import make_test_quant_config
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.distributed.eplb.eplb_state import EplbLayerState
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
from vllm.distributed.parallel_state import (
ensure_model_parallel_initialized,
@@ -201,7 +202,7 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig):
dtype=torch.int32,
device=device,
)
fml.enable_eplb = True
fml.eplb_state = EplbLayerState()
fml.set_eplb_state(
lidx,
torch.zeros(
+3 -1
View File
@@ -3,7 +3,9 @@
import pytest
from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
from vllm.model_executor.layers.fused_moe.expert_map_manager import (
determine_expert_map,
)
def verify_round_robin_pattern(expert_map, ep_rank, ep_size, global_num_experts):
+14 -2
View File
@@ -9,7 +9,10 @@ from vllm.distributed.kv_events import BlockRemoved, BlockStored
_FAKE_HASH: bytes = b"\xab" * 32
def _make_block_stored(group_idx: int | None = None) -> BlockStored:
def _make_block_stored(
group_idx: int | None = None,
kv_cache_spec_sliding_window: int | None = None,
) -> BlockStored:
return BlockStored(
block_hashes=[_FAKE_HASH],
parent_block_hash=None,
@@ -19,10 +22,13 @@ def _make_block_stored(group_idx: int | None = None) -> BlockStored:
medium="GPU",
lora_name=None,
group_idx=group_idx,
kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,
)
def _make_block_removed(group_idx: int | None = None) -> BlockRemoved:
def _make_block_removed(
group_idx: int | None = None,
) -> BlockRemoved:
return BlockRemoved(
block_hashes=[_FAKE_HASH],
medium="GPU",
@@ -72,3 +78,9 @@ def test_block_removed_hash_same_for_equal_group_idx():
event_a = _make_block_removed(group_idx=1)
event_b = _make_block_removed(group_idx=1)
assert hash(event_a) == hash(event_b)
def test_block_stored_hash_differs_by_sliding_window():
event_a = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=128)
event_b = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=256)
assert hash(event_a) != hash(event_b)
@@ -4,7 +4,6 @@
from unittest.mock import patch
import pytest
from openai.types.chat import ChatCompletionMessageParam
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from openai.types.responses.response_function_tool_call_output_item import (
ResponseFunctionToolCallOutputItem,
@@ -17,10 +16,8 @@ from openai.types.responses.response_reasoning_item import (
Summary,
)
from vllm.entrypoints.constants import MCP_PREFIX
from vllm.entrypoints.openai.responses.utils import (
_construct_single_message_from_response_item,
_maybe_combine_reasoning_and_tool_call,
_construct_message_from_response_item,
construct_chat_messages_with_tool_call,
construct_input_messages,
convert_tool_responses_to_completions_format,
@@ -28,6 +25,96 @@ from vllm.entrypoints.openai.responses.utils import (
)
def _single_chat_message(item):
message = _construct_message_from_response_item(item)
assert message is not None
return message
def make_output_message(
text: str,
*,
id: str = "msg_1",
status: str = "completed",
) -> ResponseOutputMessage:
return ResponseOutputMessage(
id=id,
content=[
ResponseOutputText(
annotations=[],
text=text,
type="output_text",
logprobs=None,
)
],
role="assistant",
status=status,
type="message",
)
def make_reasoning_item(
*,
content_text: str | None = None,
summary_text: str | None = None,
content: list[Content] | None = None,
summary: list[Summary] | None = None,
encrypted_content: str | None = None,
id: str = "reasoning_1",
status: str | None = None,
) -> ResponseReasoningItem:
if content is None and content_text is not None:
content = [Content(text=content_text, type="reasoning_text")]
if summary is None and summary_text is not None:
summary = [Summary(text=summary_text, type="summary_text")]
return ResponseReasoningItem(
id=id,
summary=[] if summary is None else summary,
type="reasoning",
content=content,
encrypted_content=encrypted_content,
status=status,
)
def make_function_call(
*,
call_id: str,
name: str = "test_function",
arguments: str = "{}",
id: str = "tool_id",
status: str | None = None,
) -> ResponseFunctionToolCall:
kwargs = {
"type": "function_call",
"id": id,
"call_id": call_id,
"name": name,
"arguments": arguments,
}
if status is not None:
kwargs["status"] = status
return ResponseFunctionToolCall(**kwargs)
def make_function_call_output(
*,
call_id: str,
output: str = "42",
id: str = "output_1",
status: str = "completed",
) -> ResponseFunctionToolCallOutputItem:
return ResponseFunctionToolCallOutputItem(
id=id,
type="function_call_output",
call_id=call_id,
output=output,
status=status,
)
class TestResponsesUtils:
"""Tests for convert_tool_responses_to_completions_format function."""
@@ -87,7 +174,7 @@ class TestResponsesUtils:
message["tool_calls"][0]["function"]["arguments"] == '{"code": "123+456"}'
)
def test_construct_single_message_from_response_item(self):
def test_construct_chat_messages_preserves_single_item_conversions(self):
item = ResponseReasoningItem(
id="lol",
summary=[],
@@ -101,7 +188,7 @@ class TestResponsesUtils:
encrypted_content=None,
status=None,
)
formatted_item = _construct_single_message_from_response_item(item)
formatted_item = _single_chat_message(item)
assert formatted_item["role"] == "assistant"
assert formatted_item["reasoning"] == "Leroy Jenkins"
@@ -119,7 +206,7 @@ class TestResponsesUtils:
status=None,
)
formatted_item = _construct_single_message_from_response_item(item)
formatted_item = _single_chat_message(item)
assert formatted_item["role"] == "assistant"
assert (
formatted_item["reasoning"]
@@ -133,11 +220,22 @@ class TestResponsesUtils:
output="1234",
status="completed",
)
formatted_item = _construct_single_message_from_response_item(tool_call_output)
formatted_item = _single_chat_message(tool_call_output)
assert formatted_item["role"] == "tool"
assert formatted_item["content"] == "1234"
assert formatted_item["tool_call_id"] == "temp"
formatted_item = _single_chat_message(
{
"type": "function_call_output",
"call_id": "temp_dict",
"output": "5678",
}
)
assert formatted_item["role"] == "tool"
assert formatted_item["content"] == "5678"
assert formatted_item["tool_call_id"] == "temp_dict"
item = ResponseReasoningItem(
id="lol",
summary=[],
@@ -147,7 +245,7 @@ class TestResponsesUtils:
status=None,
)
with pytest.raises(ValueError):
_construct_single_message_from_response_item(item)
construct_chat_messages_with_tool_call([item])
output_item = ResponseOutputMessage(
id="msg_bf585bbbe3d500e0",
@@ -164,7 +262,7 @@ class TestResponsesUtils:
type="message",
)
formatted_item = _construct_single_message_from_response_item(output_item)
formatted_item = _single_chat_message(output_item)
assert formatted_item["role"] == "assistant"
assert formatted_item["content"] == "dongyi"
@@ -192,7 +290,7 @@ class TestReasoningItemContentPriority:
encrypted_content=None,
status=None,
)
formatted = _construct_single_message_from_response_item(item)
formatted = _single_chat_message(item)
assert formatted["reasoning"] == "This is the actual content"
def test_content_only(self):
@@ -210,7 +308,7 @@ class TestReasoningItemContentPriority:
encrypted_content=None,
status=None,
)
formatted = _construct_single_message_from_response_item(item)
formatted = _single_chat_message(item)
assert formatted["reasoning"] == "Content without summary"
@patch("vllm.entrypoints.openai.responses.utils.logger")
@@ -229,7 +327,7 @@ class TestReasoningItemContentPriority:
encrypted_content=None,
status=None,
)
formatted = _construct_single_message_from_response_item(item)
formatted = _single_chat_message(item)
assert formatted["reasoning"] == "Fallback summary text"
mock_logger.warning.assert_called_once()
assert (
@@ -252,7 +350,7 @@ class TestReasoningItemContentPriority:
encrypted_content=None,
status=None,
)
formatted = _construct_single_message_from_response_item(item)
formatted = _single_chat_message(item)
assert formatted["reasoning"] == "Summary when content empty"
mock_logger.warning.assert_called_once()
assert (
@@ -269,7 +367,7 @@ class TestReasoningItemContentPriority:
encrypted_content=None,
status=None,
)
formatted = _construct_single_message_from_response_item(item)
formatted = _single_chat_message(item)
assert formatted["reasoning"] == ""
def test_encrypted_content_raises(self):
@@ -293,7 +391,7 @@ class TestReasoningItemContentPriority:
status=None,
)
with pytest.raises(ValueError):
_construct_single_message_from_response_item(item)
construct_chat_messages_with_tool_call([item])
@patch("vllm.entrypoints.openai.responses.utils.logger")
def test_summary_with_multiple_entries_uses_first(self, mock_logger):
@@ -315,7 +413,7 @@ class TestReasoningItemContentPriority:
encrypted_content=None,
status=None,
)
formatted = _construct_single_message_from_response_item(item)
formatted = _single_chat_message(item)
assert formatted["reasoning"] == "First summary"
mock_logger.warning.assert_called_once()
assert (
@@ -343,7 +441,7 @@ class TestReasoningItemContentPriority:
encrypted_content=None,
status=None,
)
_construct_single_message_from_response_item(item)
construct_chat_messages_with_tool_call([item])
mock_logger.warning.assert_not_called()
@@ -626,119 +724,125 @@ class TestShouldContinueFinalMessage:
assert should_continue_final_message([dict_item]) is False
class TestMaybeCombineReasoningAndToolCall:
"""Tests for _maybe_combine_reasoning_and_tool_call function."""
class TestConstructChatMessagesCombinePolicy:
"""Tests for contiguous assistant-side merging."""
def test_returns_none_when_item_id_is_none(self):
"""
Test fix from PR #31999: when item.id is None, should return None
instead of raising TypeError on startswith().
"""
item = ResponseFunctionToolCall(
type="function_call",
id=None, # This was causing TypeError before the fix
call_id="call_123",
name="test_function",
arguments="{}",
)
messages: list[ChatCompletionMessageParam] = []
@pytest.mark.parametrize(
("items", "expected_content", "expected_reasoning", "expected_tool_call_ids"),
[
pytest.param(
[
make_reasoning_item(content_text="Let me think"),
make_output_message("Hello"),
],
"Hello",
"Let me think",
None,
id="reasoning-output-messages",
),
pytest.param(
[
make_function_call(call_id="call_123"),
make_function_call(call_id="call_456"),
],
None,
None,
["call_123", "call_456"],
id="consecutive-tool-calls",
),
pytest.param(
[
make_reasoning_item(content_text="Let me think"),
make_function_call(call_id="call_123"),
],
None,
"Let me think",
["call_123"],
id="reasoning-tool-call",
),
pytest.param(
[
make_output_message("Hello"),
make_function_call(call_id="call_123"),
],
"Hello",
None,
["call_123"],
id="output-tool-call",
),
pytest.param(
[
make_reasoning_item(content_text="Thinking"),
make_output_message("Hello"),
make_function_call(call_id="call_123"),
make_function_call(call_id="call_456"),
],
"Hello",
"Thinking",
["call_123", "call_456"],
id="reasoning-output-tool-call",
),
],
)
def test_assistant_side_items_merge_until_tool_output(
self,
items,
expected_content,
expected_reasoning,
expected_tool_call_ids,
):
messages = construct_chat_messages_with_tool_call(items)
result = _maybe_combine_reasoning_and_tool_call(item, messages)
assert len(messages) == 1
assert messages[0]["role"] == "assistant"
if expected_content is None:
assert "content" not in messages[0]
else:
assert messages[0]["content"] == expected_content
if expected_reasoning is None:
assert "reasoning" not in messages[0]
else:
assert messages[0]["reasoning"] == expected_reasoning
if expected_tool_call_ids is None:
assert "tool_calls" not in messages[0]
else:
assert [tool_call["id"] for tool_call in messages[0]["tool_calls"]] == (
expected_tool_call_ids
)
assert result is None
def test_returns_none_when_id_does_not_start_with_mcp_prefix(self):
"""Test that non-MCP tool calls are not combined."""
item = ResponseFunctionToolCall(
type="function_call",
id="regular_id", # Does not start with MCP_PREFIX
call_id="call_123",
name="test_function",
arguments="{}",
)
messages = [{"role": "assistant", "reasoning": "some reasoning"}]
result = _maybe_combine_reasoning_and_tool_call(item, messages)
assert result is None
def test_returns_none_when_last_message_is_not_assistant(self):
"""Test that non-assistant last message returns None."""
item = ResponseFunctionToolCall(
type="function_call",
id=f"{MCP_PREFIX}tool_id",
call_id="call_123",
name="test_function",
arguments="{}",
)
messages = [{"role": "user", "content": "hello"}]
result = _maybe_combine_reasoning_and_tool_call(item, messages)
assert result is None
def test_returns_none_when_last_message_has_no_reasoning(self):
"""Test that assistant message without reasoning returns None."""
item = ResponseFunctionToolCall(
type="function_call",
id=f"{MCP_PREFIX}tool_id",
call_id="call_123",
name="test_function",
arguments="{}",
)
messages = [{"role": "assistant", "content": "some content"}]
result = _maybe_combine_reasoning_and_tool_call(item, messages)
assert result is None
def test_combines_reasoning_and_mcp_tool_call(self):
"""Test successful combination of reasoning message and MCP tool call."""
item = ResponseFunctionToolCall(
type="function_call",
id=f"{MCP_PREFIX}tool_id",
call_id="call_123",
name="test_function",
arguments='{"arg": "value"}',
)
messages = [{"role": "assistant", "reasoning": "I need to call this tool"}]
result = _maybe_combine_reasoning_and_tool_call(item, messages)
assert result is not None
assert result["role"] == "assistant"
assert result["reasoning"] == "I need to call this tool"
assert "tool_calls" in result
assert len(result["tool_calls"]) == 1
assert result["tool_calls"][0]["id"] == "call_123"
assert result["tool_calls"][0]["function"]["name"] == "test_function"
assert result["tool_calls"][0]["function"]["arguments"] == '{"arg": "value"}'
assert result["tool_calls"][0]["type"] == "function"
def test_returns_none_for_non_function_tool_call_type(self):
"""Test that non-ResponseFunctionToolCall items return None."""
# Pass a dict instead of ResponseFunctionToolCall
item = {"type": "message", "content": "hello"}
messages = [{"role": "assistant", "reasoning": "some reasoning"}]
result = _maybe_combine_reasoning_and_tool_call(item, messages)
assert result is None
def test_returns_none_when_id_is_empty_string(self):
"""Test that empty string id returns None (falsy check)."""
item = ResponseFunctionToolCall(
type="function_call",
id="", # Empty string is falsy
call_id="call_123",
name="test_function",
arguments="{}",
)
messages = [{"role": "assistant", "reasoning": "some reasoning"}]
result = _maybe_combine_reasoning_and_tool_call(item, messages)
assert result is None
@pytest.mark.parametrize(
("items", "num_expected_messages"),
[
pytest.param(
[
make_output_message("Hello"),
make_output_message("World"),
],
2,
id="consecutive-output-messages",
),
pytest.param(
[
make_reasoning_item(content_text="Let me think"),
make_reasoning_item(content_text="Let me think more"),
],
2,
id="consecutive-reasoning-messages",
),
pytest.param(
[
make_function_call(call_id="call_123"),
make_function_call_output(call_id="call_123", output="42"),
make_function_call(call_id="call_456"),
],
3,
id="interrupted-by-non-assistant-item",
),
],
)
def test_merge_chain_breaks(self, items, num_expected_messages):
messages = construct_chat_messages_with_tool_call(items)
assert len(messages) == num_expected_messages
class TestConstructInputMessagesInstructionsLeak:
@@ -145,6 +145,7 @@ def test_openapi_stateless(case: Case):
("POST", "/v1/chat/completions/batch"): LONG_TIMEOUT_SECONDS,
("POST", "/v1/completions"): LONG_TIMEOUT_SECONDS,
("POST", "/v1/messages"): LONG_TIMEOUT_SECONDS,
("POST", "/inference/v1/generate"): LONG_TIMEOUT_SECONDS,
}.get(key, DEFAULT_TIMEOUT_SECONDS)
# No need to verify SSL certificate for localhost
@@ -0,0 +1,70 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the disagg request/response protocol.
These tests intentionally avoid spinning up a server they exercise the
pydantic validators on ``GenerateRequest`` directly so they run fast and
fail loudly if the validator semantics ever drift.
"""
import json
from vllm.entrypoints.serve.disagg.protocol import GenerateRequest
from vllm.sampling_params import SamplingParams
def _base_payload() -> dict:
return {"token_ids": [1, 2, 3], "sampling_params": {}}
def test_omitted_max_tokens_is_not_provided():
"""Body without ``max_tokens`` must surface as 'not provided' so the
server can apply its own default instead of the dataclass 16."""
req = GenerateRequest.model_validate(_base_payload())
# SamplingParams' dataclass default leaks through the parsed instance —
# this is exactly the bug the server-side defaulting works around.
assert req.sampling_params.max_tokens == 16
assert not req.is_sampling_param_provided("max_tokens")
def test_explicit_max_tokens_is_provided():
"""Even when the client picks the same value as the dataclass default,
it must register as explicitly set so the server won't override it."""
payload = _base_payload()
payload["sampling_params"] = {"max_tokens": 16}
req = GenerateRequest.model_validate(payload)
assert req.sampling_params.max_tokens == 16
assert req.is_sampling_param_provided("max_tokens")
payload["sampling_params"] = {"max_tokens": 256}
req = GenerateRequest.model_validate(payload)
assert req.sampling_params.max_tokens == 256
assert req.is_sampling_param_provided("max_tokens")
def test_other_fields_tracked_independently():
payload = _base_payload()
payload["sampling_params"] = {"temperature": 0.0}
req = GenerateRequest.model_validate(payload)
assert not req.is_sampling_param_provided("max_tokens")
assert req.is_sampling_param_provided("temperature")
def test_json_roundtrip_preserves_provided_keys():
payload = _base_payload()
payload["sampling_params"] = {"temperature": 0.5}
req = GenerateRequest.model_validate_json(json.dumps(payload))
assert not req.is_sampling_param_provided("max_tokens")
assert req.is_sampling_param_provided("temperature")
def test_internal_instance_construction_treats_all_as_provided():
"""When internal callers build ``GenerateRequest`` from a pre-resolved
``SamplingParams`` instance, every field is considered explicitly set
so server-side defaulting can't clobber values resolved upstream."""
sp = SamplingParams(max_tokens=500, temperature=0.0)
req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=sp)
assert req.is_sampling_param_provided("max_tokens")
assert req.is_sampling_param_provided("temperature")
# And keys we never touched should also count as provided in this path.
assert req.is_sampling_param_provided("top_p")
@@ -114,6 +114,36 @@ async def test_generate_endpoint(client):
assert "choices" in data
@pytest.mark.asyncio
async def test_generate_defaults_max_tokens_when_omitted(client):
"""Regression: omitting ``max_tokens`` must not silently cap at 16.
``SamplingParams.max_tokens`` defaults to 16. Before the server-side
defaulting was wired up, every request that didn't set ``max_tokens``
truncated mid-generation. The server should now fill it in from
``max_model_len - prompt_len`` (matching ``/v1/chat/completions``).
"""
payload = {
"model": MODEL_NAME,
"token_ids": [1, 2, 3],
"sampling_params": {
"temperature": 0.0,
"ignore_eos": True,
},
"stream": False,
}
resp = await client.post(GEN_ENDPOINT, json=payload)
resp.raise_for_status()
data = resp.json()
completion_tokens = len(data["choices"][0]["token_ids"])
# max_model_len=1024 in the test fixture, prompt is 3 tokens, so we
# should get ~1021 tokens of output (capped at max_model_len boundary).
assert completion_tokens > 16, (
f"expected server-side default to exceed the legacy 16-token cap, "
f"got {completion_tokens}"
)
@pytest.mark.asyncio
async def test_generate_stream(client):
payload = {
+214 -66
View File
@@ -13,20 +13,40 @@ from torch.fx.experimental.proxy_tensor import make_fx
import vllm.ir.op
from vllm.ir.op import RESERVED_PROVIDERS, IrOp, IrOpImpl
# This should not exist
assert "_custom_add" not in IrOp.registry
class CustomError(Exception):
pass
@vllm.ir.register_op(allow_inplace=True)
def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return x + y
@pytest.fixture
def custom_add_op(fake_vllm_ir):
"""Register ``_custom_add`` plus impl_a, impl_b, impl_even for this test."""
@vllm.ir.register_op(allow_inplace=True)
def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return x + y
@_custom_add.register_impl("impl_a")
def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return x + y + 10
@_custom_add.register_impl("impl_b", inplace=True)
def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
"""Computes x+y+20"""
x.add_(y)
x.add_(20)
return x
@_custom_add.register_impl(
"impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0
)
def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return x + y + 50
return _custom_add
def test_registration_overloads():
def test_registration_overloads(fake_vllm_ir):
assert all(
n not in IrOp.registry for n in ["_custom_sub", "_custom_mul", "_custom_div"]
)
@@ -63,7 +83,7 @@ def test_registration_overloads():
return x * y - 100
def test_no_kw_only_args():
def test_no_kw_only_args(fake_vllm_ir):
# kw-only args not supported
with pytest.raises(ValueError, match="keyword-only arguments"):
@@ -78,18 +98,23 @@ def test_no_kw_only_args():
class TestIrOpCustomAdd:
# Registration invariants
def test_decorated_object(self):
def test_decorated_object(self, custom_add_op):
"""Make sure that referring directly to an op is correct"""
_custom_add = custom_add_op
assert isinstance(_custom_add, IrOp)
assert "_custom_add" in IrOp.registry
assert _custom_add is IrOp.registry["_custom_add"]
def test_torch_op_is_registered(self):
assert hasattr(torch.ops.vllm_ir, "_custom_add")
assert callable(torch.ops.vllm_ir._custom_add.default)
def test_torch_op_is_registered(self, custom_add_op):
_custom_add = custom_add_op
torch_ops = getattr(torch.ops, vllm.ir.op.vllm_ir_torch_lib.ns)
assert hasattr(torch_ops, "_custom_add")
assert callable(torch_ops._custom_add.default)
assert _custom_add.torch_op is torch_ops._custom_add.default
# Semantic correctness
def test_semantics_match_native(self):
def test_semantics_match_native(self, custom_add_op):
_custom_add = custom_add_op
x = torch.randn(4, 5)
y = torch.randn(4, 5)
@@ -103,7 +128,9 @@ class TestIrOpCustomAdd:
# Implementation registration
# -------------------------
def test_register_impl_is_non_intrusive(self):
def test_register_impl_is_non_intrusive(self, custom_add_op):
_custom_add = custom_add_op
@_custom_add.register_impl("dummy_provider")
def dummy_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return x + y + 123
@@ -117,7 +144,8 @@ class TestIrOpCustomAdd:
# Native semantics must still hold
torch.testing.assert_close(_custom_add(x, y), x + y)
def test_schema_contains_tensor_signature(self):
def test_schema_contains_tensor_signature(self, custom_add_op):
_custom_add = custom_add_op
schema = _custom_add._schema_str
assert "Tensor" in schema
@@ -131,10 +159,19 @@ class TestIrOpCustomAdd:
@pytest.mark.parametrize("symbolic_trace", [True, False])
@pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
def test_trace_sees_single_custom_op(
self, symbolic_trace: bool, enable_torch_wrap: bool, overload: str
self,
custom_add_op,
symbolic_trace: bool,
enable_torch_wrap: bool,
overload: str,
):
_custom_add = custom_add_op
op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)
torch_op = (
_custom_add.torch_op
if overload == "default"
else _custom_add.maybe_inplace.torch_op
)
def fn(x, y):
return op_fn(x, y)
@@ -175,30 +212,13 @@ class TestIrOpCustomAdd:
assert len(ir_nodes) == 1, gm.code
@_custom_add.register_impl("impl_a")
def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return x + y + 10
@_custom_add.register_impl("impl_b", inplace=True)
def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
"""Computes x+y+20"""
x.add_(y)
x.add_(20)
return x
@_custom_add.register_impl("impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0)
def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return x + y + 50
class TestIrOpImplDispatch:
def test_register_impl(self):
def test_register_impl(self, custom_add_op):
_custom_add = custom_add_op
assert "impl_a" in _custom_add.impls
impl = _custom_add.impls["impl_a"]
assert impl is impl_a
assert impl is _custom_add.impls["impl_a"]
assert impl.op is _custom_add
assert impl.provider == "impl_a"
assert callable(impl.impl_fn)
@@ -211,14 +231,15 @@ class TestIrOpImplDispatch:
return x + y + 30
# Check the original impl is still intact
assert _custom_add.impls["impl_a"] is impl_a
assert _custom_add.impls["impl_a"] is impl
# Check support all args
assert impl_a.supports_all_args
assert impl_b.supports_all_args
assert not impl_even.supports_all_args
assert _custom_add.impls["impl_a"].supports_all_args
assert _custom_add.impls["impl_b"].supports_all_args
assert not _custom_add.impls["impl_even"].supports_all_args
def test_reserved_provider_rejected(self):
def test_reserved_provider_rejected(self, custom_add_op):
_custom_add = custom_add_op
for provider in RESERVED_PROVIDERS:
with pytest.raises(AssertionError):
@@ -226,7 +247,8 @@ class TestIrOpImplDispatch:
def bad_impl(x, y):
return x + y
def test_set_priority_scoped(self):
def test_set_priority_scoped(self, custom_add_op):
_custom_add = custom_add_op
assert _custom_add.get_priority() == []
with _custom_add.set_priority(["impl_even", "impl_b"]):
@@ -251,20 +273,25 @@ class TestIrOpImplDispatch:
assert _custom_add.get_priority() == []
@pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
def test_dispatch_priority_order(self, overload: str):
def test_dispatch_priority_order(self, custom_add_op, overload: str):
_custom_add = custom_add_op
op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)
torch_op = (
_custom_add.torch_op
if overload == "default"
else _custom_add.maybe_inplace.torch_op
)
x = torch.tensor(1, dtype=torch.int32)
y = torch.tensor(2, dtype=torch.int32)
with _custom_add.set_priority(["impl_b", "impl_a"]):
assert _custom_add.dispatch(x, y) is impl_b
assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_b"]
out1 = op_fn(x.clone(), y)
out2 = torch_op(x.clone(), y)
with _custom_add.set_priority(["impl_a"]):
assert _custom_add.dispatch(x, y) is impl_a
assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_a"]
out3 = op_fn(x.clone(), y)
out4 = torch_op(x.clone(), y)
@@ -275,7 +302,9 @@ class TestIrOpImplDispatch:
assert out3.item() == 1 + 2 + 10
assert out4.item() == 1 + 2 + 10
def test_unsupported_impl_filtered(self):
def test_unsupported_impl_filtered(self, custom_add_op):
_custom_add = custom_add_op
@_custom_add.register_impl("impl_unsupported", supported=False)
def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return x + y + 999
@@ -291,8 +320,9 @@ class TestIrOpImplDispatch:
assert out.item() == 1 + 2 + 10
def test_supports_args_runtime_dispatch_and_warning(
self, caplog_vllm: pytest.LogCaptureFixture
self, custom_add_op, caplog_vllm: pytest.LogCaptureFixture
):
_custom_add = custom_add_op
x1 = torch.ones((2, 2), dtype=torch.int32)
y1 = torch.full((2, 2), 2, dtype=torch.int32)
@@ -312,7 +342,7 @@ class TestIrOpImplDispatch:
# Check dispatching
assert _custom_add.get_priority() == ["impl_even", "native"]
assert _custom_add.dispatch(x1, y1) is impl_even
assert _custom_add.dispatch(x1, y1) is _custom_add.impls["impl_even"]
assert _custom_add.dispatch(x2, y2) is _custom_add.impls["native"]
out1 = _custom_add(x1, y1) # size(1) == 2 → impl_even
@@ -324,8 +354,12 @@ class TestIrOpImplDispatch:
assert torch.all(out2 == 1 + 2)
def test_default_priority(
self, caplog_vllm: pytest.LogCaptureFixture, disable_log_dedup
self,
custom_add_op,
caplog_vllm: pytest.LogCaptureFixture,
disable_log_dedup,
):
_custom_add = custom_add_op
# Make sure logs are not deduplicated to properly test the warning
x = torch.tensor([3], dtype=torch.int32)
y = torch.tensor([4], dtype=torch.int32)
@@ -347,15 +381,23 @@ class TestIrOpImplDispatch:
assert "priority not set" in message
@vllm.ir.register_op
def _custom_mm(
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
) -> torch.Tensor:
tmp = x @ y
return tmp if bias is None else tmp + bias
@pytest.fixture
def custom_mm_op(fake_vllm_ir):
"""Fixture that registers ``_custom_mm`` (isolated by ``fake_vllm_ir``)."""
@vllm.ir.register_op
def _custom_mm(
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
) -> torch.Tensor:
tmp = x @ y
return tmp if bias is None else tmp + bias
return _custom_mm
def test_default_args():
def test_default_args(custom_mm_op):
_custom_mm = custom_mm_op
# Test that default args are properly applied when dispatching and calling
@_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True)
def impl_mm(
@@ -373,7 +415,8 @@ def test_default_args():
assert _custom_mm.dispatch(x1, x2) is impl_mm
def test_bad_impl_registrations():
def test_bad_impl_registrations(custom_mm_op):
_custom_mm = custom_mm_op
# Check bad schema
with pytest.raises(ValueError, match="does not match native schema"):
@@ -446,7 +489,8 @@ def test_bad_impl_registrations():
) -> torch.Tensor:
return x @ y + 40
assert set(_custom_mm.impls.keys()) == {"impl_mm", "native"}
# With fixture, each test gets a fresh op with only "native" impl
assert set(_custom_mm.impls.keys()) == {"native"}
IMPL_OOT_SRC = """
@@ -460,14 +504,14 @@ def impl_mm_oot(
"""
def load_custom_mm_module(file_path: Path):
def load_custom_mm_module(file_path: Path, custom_mm_op):
spec = importlib.util.spec_from_file_location("_custom_mm_oot", file_path)
assert spec is not None
module = importlib.util.module_from_spec(spec)
# Inject the variable into the module's global namespace
# This allows the @_custom_mm.register_impl decorator to work
module._custom_mm = _custom_mm # type: ignore[attr-defined]
module._custom_mm = custom_mm_op # type: ignore[attr-defined]
# Execute the file; this triggers the decorator
assert spec.loader is not None
@@ -475,12 +519,13 @@ def load_custom_mm_module(file_path: Path):
return module
def test_uuid_and_oot(tmp_path: Path):
def test_uuid_and_oot(custom_mm_op, tmp_path: Path):
_custom_mm = custom_mm_op
file_path = tmp_path / "_custom_mm_oot.py"
file_path.write_text(IMPL_OOT_SRC)
assert "impl_mm_oot" not in _custom_mm.impls
_ = load_custom_mm_module(file_path)
_ = load_custom_mm_module(file_path, _custom_mm)
assert "impl_mm_oot" in _custom_mm.impls
uuid = _custom_mm.impls["impl_mm_oot"].uuid()
@@ -489,7 +534,7 @@ def test_uuid_and_oot(tmp_path: Path):
# Replace file source
file_path.write_text(IMPL_OOT_SRC + " # added file source")
assert "impl_mm_oot" not in _custom_mm.impls
_ = load_custom_mm_module(file_path)
_ = load_custom_mm_module(file_path, _custom_mm)
assert "impl_mm_oot" in _custom_mm.impls
uuid1 = _custom_mm.impls["impl_mm_oot"].uuid()
@@ -499,7 +544,7 @@ def test_uuid_and_oot(tmp_path: Path):
# Back to original
file_path.write_text(IMPL_OOT_SRC)
assert "impl_mm_oot" not in _custom_mm.impls
_ = load_custom_mm_module(file_path)
_ = load_custom_mm_module(file_path, _custom_mm)
assert "impl_mm_oot" in _custom_mm.impls
uuid2 = _custom_mm.impls["impl_mm_oot"].uuid()
@@ -571,3 +616,106 @@ class TestTolerance:
op = IrOp("_tol_test_unknown", _test_native)
with pytest.raises(ValueError, match="No tolerance defined"):
op.get_tolerance(torch.complex64)
def test_naming_validation(fake_vllm_ir):
"""Test that op and provider names are validated ([a-z_][a-z_0-9]*)."""
# Valid op and provider names
@vllm.ir.register_op
def _valid_name_123(x: torch.Tensor) -> torch.Tensor:
return x
@_valid_name_123.register_impl("valid_provider_123")
def valid_impl(x: torch.Tensor) -> torch.Tensor:
return x + 1
# Invalid op names should fail
with pytest.raises(ValueError, match="name.*invalid"):
@vllm.ir.register_op
def InvalidName(x: torch.Tensor) -> torch.Tensor:
return x
with pytest.raises(ValueError, match="name.*invalid"):
@vllm.ir.register_op(name="123invalid")
def some_func(x: torch.Tensor) -> torch.Tensor:
return x
# Invalid provider names should fail
with pytest.raises(ValueError, match="name.*invalid"):
@_valid_name_123.register_impl("Invalid-Provider")
def invalid_impl(x: torch.Tensor) -> torch.Tensor:
return x + 1
def test_registration_stack_traces(fake_vllm_ir):
"""Test that stack traces are captured for ops and impls."""
@vllm.ir.register_op
def _test_stack(x: torch.Tensor) -> torch.Tensor:
return x
@_test_stack.register_impl("test_provider")
def test_impl(x: torch.Tensor) -> torch.Tensor:
return x + 1
# Verify op stack trace
assert hasattr(_test_stack, "_registration_stack")
assert len(_test_stack._registration_stack) > 0
op_stack_str = "".join(_test_stack._registration_stack)
assert "test_op.py" in op_stack_str
# Last frame should be the decorator in user code, not internal decorator logic
assert "@vllm.ir.register_op" in _test_stack._registration_stack[-1]
assert "return decorator(f)" not in op_stack_str
# Verify impl stack trace
impl = _test_stack.impls["test_provider"]
assert hasattr(impl, "_registration_stack")
assert len(impl._registration_stack) > 0
impl_stack_str = "".join(impl._registration_stack)
assert "test_op.py" in impl_stack_str
# Last frame should be the decorator in user code
assert '@_test_stack.register_impl("test_provider")' in impl._registration_stack[-1]
def test_op_repr_uses_docstring(fake_vllm_ir):
"""Test that __str__ uses the function's docstring and __repr__ is simple."""
@vllm.ir.register_op
def _test_repr_with_doc(x: torch.Tensor) -> torch.Tensor:
"""First line of docstring.
Additional details here.
"""
return x
@vllm.ir.register_op
def _test_repr_no_doc(x: torch.Tensor) -> torch.Tensor:
return x
# __str__ with docstring: uses first line only
str_with = str(_test_repr_with_doc)
assert "IrOp('_test_repr_with_doc')" in str_with
assert "First line of docstring." in str_with
assert "Additional details" not in str_with
# __str__ without docstring: simple format
assert str(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')"
# __repr__ should be simple for both
assert repr(_test_repr_with_doc) == "IrOp('_test_repr_with_doc')"
assert repr(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')"
def test_vllm_ir_fixture(fake_vllm_ir):
"""Test that the fake_vllm_ir fixture provides test isolation."""
@vllm.ir.register_op
def _test_fixture(x: torch.Tensor) -> torch.Tensor:
return x
assert "_test_fixture" in IrOp.registry
# Fixture will automatically clean up after test
@@ -345,3 +345,182 @@ def test_triton_unified_attn_fp16_input_fp8_output(
torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
f"{torch.max(torch.abs(output_fp16 - ref_output))}",
)
# USE_TD path covers two head-size regimes:
# - pow2 (HEAD_SIZE == HEAD_SIZE_PADDED): full TD path including Q/O.
# - non-pow2 (96, HEAD_SIZE_PADDED=128): gates USE_TD_QO off — Q load
# and output store fall back to pointer path, KV tile TD load remains.
# The non-pow2 case mirrors real models like Phi-3-mini (head_size=96).
HEAD_SIZES_USE_TD = [128, 256, 96]
def _run_use_td_case(
seq_lens: list[tuple[int, int]],
num_heads: tuple[int, int],
head_size: int,
block_size: int,
sliding_window: int | None,
soft_cap: float | None,
seq_threshold_3D: int,
dtype: torch.dtype = torch.bfloat16,
num_blocks: int = 2048,
) -> None:
"""Shared driver for the USE_TD test cases.
Runs ``unified_attention(..., use_td=True)`` and compares against the
reference paged-attention implementation that the sibling non-TD
tests use.
"""
torch.set_default_device(DEVICE_TYPE)
set_random_seed(0)
num_seqs = len(seq_lens)
query_lens = [x[0] for x in seq_lens]
kv_lens = [x[1] for x in seq_lens]
num_query_heads, num_kv_heads = num_heads
assert num_query_heads % num_kv_heads == 0
max_query_len = max(query_lens)
max_kv_len = max(kv_lens)
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
scale = head_size**-0.5
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
key_cache = torch.randn(
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
)
value_cache = torch.randn_like(key_cache)
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
dim=0, dtype=torch.int32
)
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
block_tables = torch.randint(
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
)
output = torch.empty_like(query)
num_par_softmax_segments = 16
head_size_padded = next_power_of_2(head_size)
softmax_segm_output = torch.empty(
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
dtype=torch.float32,
)
softmax_segm_max = torch.empty(
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
dtype=torch.float32,
)
softmax_segm_expsum = torch.empty(
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
dtype=torch.float32,
)
unified_attention(
q=query,
k=key_cache,
v=value_cache,
out=output,
cu_seqlens_q=cu_query_lens,
seqused_k=kv_lens_tensor,
max_seqlen_q=max_query_len,
max_seqlen_k=max_kv_len,
softmax_scale=scale,
causal=True,
window_size=window_size,
block_table=block_tables,
softcap=soft_cap if soft_cap is not None else 0,
q_descale=None,
k_descale=None,
v_descale=None,
seq_threshold_3D=seq_threshold_3D,
num_par_softmax_segments=num_par_softmax_segments,
softmax_segm_output=softmax_segm_output,
softmax_segm_max=softmax_segm_max,
softmax_segm_expsum=softmax_segm_expsum,
use_td=True,
)
ref_output = ref_paged_attn(
query=query,
key_cache=key_cache,
value_cache=value_cache,
query_lens=query_lens,
kv_lens=kv_lens,
block_tables=block_tables,
scale=scale,
sliding_window=sliding_window,
soft_cap=soft_cap,
)
torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
@pytest.mark.parametrize(
"seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES_USE_TD)
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
@pytest.mark.parametrize("sliding_window", [None, 128])
@pytest.mark.parametrize("soft_cap", [None, 50.0])
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
@torch.inference_mode()
def test_triton_unified_attn_use_td(
seq_lens: list[tuple[int, int]],
num_heads: tuple[int, int],
head_size: int,
sliding_window: int | None,
block_size: int,
soft_cap: float | None,
num_blocks: int,
seq_threshold_3D: int,
) -> None:
"""Exercise the USE_TD (tensor-descriptor) Q/K/V load/store path.
Covers both 2D and 3D kernels via ``seq_threshold_3D``. Two routes
to the USE_TD_QO=False fallback (pointer path for Q/O with TD still
active for KV tile loads):
- non-pow2 ``num_queries_per_kv`` via ``NUM_HEADS`` entry ``(5, 1)``,
- non-pow2 ``head_size`` via ``HEAD_SIZES_USE_TD`` entry ``96``.
"""
_run_use_td_case(
seq_lens=seq_lens,
num_heads=num_heads,
head_size=head_size,
block_size=block_size,
sliding_window=sliding_window,
soft_cap=soft_cap,
seq_threshold_3D=seq_threshold_3D,
num_blocks=num_blocks,
)
# Prefill-heavy shape: long query drives the prefill kernel path where
# ``_get_tile_size`` returns 32, which exceeds block_size=16 and must be
# clamped by the fix in 'clamp TILE_SIZE to block_size when USE_TD'.
# Only the prefill launch exercises the clamp, so parameterize only over
# the (num_heads, seq_threshold_3D=0) combinations needed to cover it.
@pytest.mark.parametrize("num_heads", [(4, 4), (5, 1)])
@torch.inference_mode()
def test_triton_unified_attn_use_td_tile_clamp(
num_heads: tuple[int, int],
) -> None:
"""Regression guard: ``USE_TD`` needs ``BLOCK_SIZE % TILE_SIZE == 0``.
With ``block_size=16`` and ``head_size=128`` (non-Gemma3),
``_get_tile_size`` returns 32 for prefill, which violates the
``USE_TD`` constraint unless clamped to ``block_size``. Without
the clamp the triton kernel ``static_assert`` fires at compile time.
"""
_run_use_td_case(
seq_lens=[(256, 256), (128, 128)],
num_heads=num_heads,
head_size=128,
block_size=16,
sliding_window=None,
soft_cap=None,
seq_threshold_3D=0,
)
@@ -15,16 +15,16 @@ from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
BatchedDeepGemmExperts,
)
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import DeepGemmExperts
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
BatchedTritonExperts,
NaiveBatchedExperts,
)
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
TritonOrDeepGemmExperts,
)
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
MoEPrepareAndFinalizeNoDPEPModular,
)
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
TritonOrDeepGemmExperts,
)
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
cutlass_fp4_supported,
)
+1 -1
View File
@@ -9,7 +9,7 @@ from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_confi
from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
BatchedDeepGemmExperts,
)
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
BatchedTritonExperts,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+1 -1
View File
@@ -16,7 +16,7 @@ from tests.kernels.quant_utils import native_batched_masked_quant_matmul
from tests.kernels.utils import torch_experts
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe import fused_topk
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
invoke_moe_batched_triton_kernel,
)
from vllm.platforms import current_platform
+1 -1
View File
@@ -31,7 +31,7 @@ from vllm.model_executor.layers.fused_moe.config import (
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
_valid_deep_gemm_shape,
)
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
TritonOrDeepGemmExperts,
)
from vllm.platforms import current_platform
+3 -1
View File
@@ -19,7 +19,9 @@ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
)
from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
BatchedTritonExperts,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
per_token_group_quant_fp8,
+2 -2
View File
@@ -25,10 +25,10 @@ from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantDesc,
fp8_w8a8_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
TritonOrDeepGemmExperts,
)
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
per_token_group_quant_fp8,
)
+1 -1
View File
@@ -1588,7 +1588,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
layer.apply_router_weight_on_input = False
layer.routed_scaling_factor = None
layer.shared_experts = None
layer._maybe_init_expert_routing_tables = lambda: None
layer._expert_routing_tables = lambda: None
quant_method.process_weights_after_loading(layer)
+174 -66
View File
@@ -31,12 +31,13 @@ from vllm.config import (
VllmConfig,
set_current_vllm_config,
)
from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
from vllm.distributed.parallel_state import (
from vllm.distributed import (
get_ep_group,
get_eplb_group,
tensor_model_parallel_all_gather,
)
from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.fused_moe import FusedMoE, fused_experts
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -49,6 +50,7 @@ from vllm.model_executor.layers.quantization.modelopt import (
ModelOptFp8Config,
ModelOptNvFp4Config,
)
from vllm.model_executor.models.utils import sequence_parallel_chunk
from vllm.platforms import current_platform
from vllm.utils.flashinfer import (
has_flashinfer_nvlink_one_sided,
@@ -81,6 +83,9 @@ PARALLEL_COMBOS = [
[1, 4, False],
[2, 1, True],
[4, 1, True],
# This combination indicates sequence parallel.
# See ParallelConfig.use_sequence_parallel.
[2, 2, True],
]
# TODO: should this even be set manually? let oracles handle this
@@ -112,24 +117,24 @@ QUANT_METHODS = [
# Which quantization methods each backend supports.
# fmt: off
BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
"allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
"mori": {None, "fp8", "modelopt_fp8"},
"flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"},
"flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"},
"deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"},
"allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
"mori": {None, "fp8", "modelopt_fp8"},
"flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
"flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501
"deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
"deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
"nixl_ep": {None, "fp8", "modelopt_fp8"},
"nixl_ep": {None, "fp8_blocked", "modelopt_fp8"},
}
# Map from backend -> (DP/EP support, DP support, TP support)
BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool]] = {
"allgather_reducescatter": (True, True, True),
"mori": (True, False, False),
"flashinfer_nvlink_two_sided": (False, True, False),
"flashinfer_nvlink_one_sided": (False, True, False),
"deepep_low_latency": (True, False, False),
"deepep_high_throughput": (True, False, False),
"nixl_ep": (True, False, False),
# Map from backend -> (DP/EP support, DP support, TP support, SP support)
BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool, bool]] = {
"allgather_reducescatter": (True, True, True, True),
"mori": (True, False, False, True),
"flashinfer_nvlink_two_sided": (False, True, False, False),
"flashinfer_nvlink_one_sided": (False, True, False, False),
"deepep_low_latency": (True, False, False, True),
"deepep_high_throughput": (True, False, False, True),
"nixl_ep": (True, False, False, True),
}
# fmt: on
@@ -163,6 +168,45 @@ def override_normalize_e4m3fn_to_e4m3fnuz():
vllm.model_executor.layers.quantization.utils.w8a8_utils.normalize_e4m3fn_to_e4m3fnuz = mock_normalize_e4m3fn_to_e4m3fnuz # noqa: E501
def sp_wrapper(
fn: Callable | FusedMoE, is_sequence_parallel: bool | None = None
) -> Callable:
"""Wrapper to handle sequence parallelism chunking and gathering.
For SP with EP:
- The TP group is created with the original tensor_parallel_size (e.g., 2)
- get_tp_group() has the correct world_size for SP operations
- sequence_parallel_chunk() uses get_tensor_model_parallel_world_size()
- tensor_model_parallel_all_gather() uses get_tp_group()
- Both should work correctly even when EP is enabled
"""
if isinstance(fn, FusedMoE):
assert is_sequence_parallel is None
is_sequence_parallel = fn.is_sequence_parallel
else:
assert is_sequence_parallel is not None
if is_sequence_parallel:
def wrapper(
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
) -> torch.Tensor:
# Split sequence across TP ranks
# Both hidden_states and router_logits have [num_tokens, ...] shape
hidden_states = sequence_parallel_chunk(hidden_states)
router_logits = sequence_parallel_chunk(router_logits)
# Run MoE on local chunk
result = fn(hidden_states, router_logits)
# Gather results from all TP ranks
result = tensor_model_parallel_all_gather(result, 0)
# Remove any padding added by SP.
return result[: hidden_states.shape[0]]
return wrapper
return fn
def maybe_roundup_layer_hidden_size(
hidden_size: int,
act_dtype: torch.dtype,
@@ -272,6 +316,15 @@ class MoETestConfig:
dp_size: int = 1
tp_size: int = 1
@property
def is_sequence_parallel(self) -> bool:
# Sequence parallelism: EP enabled + TP dimension used for sequence splitting
# In test config: ep_size represents total expert parallel size
# tp_size represents the original TP dimension (becomes sp_size in FusedMoE)
# dp_size represents data parallel size
# For SP: we need EP enabled (ep_size > 1) and sequence splitting (tp_size > 1)
return self.ep_size > 1 and self.tp_size > 1
# TODO: add more error messages
def id(self) -> str:
def proc(s: str) -> str:
@@ -404,11 +457,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
"leads to large differences.",
)
# gate requires shared_experts (use_overlapped mode)
# TODO: also not sure this is true
if config.use_gate and not config.use_shared_experts:
return False, "gate requires shared_experts (use_overlapped mode)"
# Skip modelopt_fp4 if not on B100+ (compute capability 10.0+)
if (
config.quantization == "modelopt_fp4"
@@ -445,7 +493,7 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
)
if config.backend == "nixl_ep":
from vllm.model_executor.layers.fused_moe.nixl_ep_prepare_finalize import ( # noqa: E501
from vllm.model_executor.layers.fused_moe.prepare_finalize.nixl_ep import ( # noqa: E501
NixlEPPrepareAndFinalize,
)
@@ -456,11 +504,11 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
)
if config.backend is not None:
supports_ep_dp, supports_dp, supports_tp = BACKEND_EP_DP_TP_SUPPORT[
config.backend
]
supports_ep_dp, supports_dp, supports_tp, supports_sp = (
BACKEND_EP_DP_TP_SUPPORT[config.backend]
)
if config.tp_size > 1 and not supports_tp:
if config.tp_size > 1 and not supports_tp and not config.is_sequence_parallel:
return False, f"{config.backend} does not support TP."
if config.dp_size > 1 and config.ep_size == 1 and not supports_dp:
@@ -468,10 +516,34 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
if config.dp_size > 1 and config.ep_size > 1 and not supports_ep_dp:
return False, f"{config.backend} does not support EP/DP."
if config.is_sequence_parallel and not supports_sp:
return False, f"{config.backend} does not support SP."
else:
if config.tp_size > 1 or config.ep_size > 1 or config.dp_size > 1:
return False, "An all2all backend is required for parallelism."
# Sequence parallelism specific validations
if config.is_sequence_parallel:
if config.ep_size == 1:
return False, "Sequence parallelism requires EP to be enabled (ep_size > 1)"
if config.tp_size == 1:
return (
False,
"Sequence parallelism requires tp_size > 1 for sequence splitting",
)
# SP is essentially EP + sequence splitting
# Verify the relationship: ep_size should equal dp_size * tp_size
# (when pcp_size=1).
expected_ep_size = config.dp_size * config.tp_size
if config.ep_size != expected_ep_size:
return False, (
f"For sequence parallelism: ep_size ({config.ep_size}) should equal "
f"dp_size * tp_size ({expected_ep_size})"
)
if config.enable_eplb:
if config.ep_size == 1:
return False, "EPLB requires EP."
@@ -485,14 +557,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
if config.num_experts % config.dp_size != 0:
return False, "EPLB requires num_experts divisible by ep_size"
# 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
@@ -751,7 +815,8 @@ def create_shared_experts_from_config(
in_dtype: torch.dtype,
tp_size: int = 1,
tp_rank: int = 0,
device: torch.device | str | None = None,
is_sequence_parallel: bool = False,
device: torch.device | str | None = "cuda",
) -> TestMLP | None:
"""Create TestMLP for shared experts from config.
@@ -772,7 +837,7 @@ def create_shared_experts_from_config(
s_w2 = shared_experts_config.w2
# Apply TP chunking if needed
if tp_size > 1:
if tp_size > 1 and not is_sequence_parallel:
s_w1 = tp_chunk_gate_up(s_w1, tp_rank, tp_size, dim=1, device=device)
s_w2 = chunk_by_rank(s_w2, tp_rank, tp_size, dim=0, device=device)
else:
@@ -920,6 +985,7 @@ def make_fused_moe_layer(
routed_input_transform: torch.nn.Module | None = None,
routed_output_transform: torch.nn.Module | None = None,
pcp_size: int | None = 1,
is_sequence_parallel: bool = False,
) -> FusedMoE:
quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts)
@@ -959,6 +1025,7 @@ def make_fused_moe_layer(
enable_eplb=enable_eplb,
num_redundant_experts=num_redundant_experts,
has_bias=has_bias,
is_sequence_parallel=is_sequence_parallel,
**kwargs,
)
@@ -1004,7 +1071,6 @@ def make_fake_moe_layer(
activation: str = "silu",
indices_type: torch.dtype | None = None,
expert_map: torch.Tensor | None = None,
enable_eplb: bool = False,
expert_load_view: torch.Tensor | None = None,
logical_to_physical_map: torch.Tensor | None = None,
logical_replica_count: torch.Tensor | None = None,
@@ -1015,6 +1081,7 @@ def make_fake_moe_layer(
tp_size: int = 1,
dp_size: int = 1,
ep_size: int = 1,
is_sequence_parallel: bool = False,
) -> Callable:
quant_dtype = None
activation = MoEActivation.from_str(activation)
@@ -1022,7 +1089,6 @@ def make_fake_moe_layer(
router = create_fused_moe_router(
top_k=top_k,
global_num_experts=global_num_experts,
# eplb_state=None, # TODO
renormalize=renormalize,
use_grouped_topk=use_grouped_topk,
num_expert_group=num_expert_group,
@@ -1032,7 +1098,6 @@ def make_fake_moe_layer(
routed_scaling_factor=routed_scaling_factor,
e_score_correction_bias=e_score_correction_bias,
num_fused_shared_experts=0, # TODO
enable_eplb=enable_eplb,
# TODO(bnell): once we can construct the MK at init time, we
# can make this a value.
indices_type_getter=lambda: indices_type,
@@ -1046,7 +1111,8 @@ def make_fake_moe_layer(
w2_s = None
shared_experts = create_shared_experts_from_config(
shared_experts_config, in_dtype, 1, 0, "cuda"
shared_experts_config,
in_dtype,
)
quant_config = FusedMoEQuantConfig.make(
@@ -1117,7 +1183,7 @@ def make_fake_moe_layer(
def _test_body_regular(
moe_layer: Callable,
moe_layer: FusedMoE,
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
vllm_config: VllmConfig,
@@ -1134,7 +1200,7 @@ def _test_body_regular(
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp,
):
output = moe_layer(hidden_states, router_logits)
output = sp_wrapper(moe_layer)(hidden_states, router_logits)
return baseline_output, output
@@ -1167,6 +1233,8 @@ def _test_body_eplb(
) -> tuple[torch.Tensor, torch.Tensor]:
device = torch.accelerator.current_accelerator()
is_sequence_parallel = moe_layer.is_sequence_parallel
"""EPLB test body: compare output before and after expert weight rearrangement."""
# Get "before" output with original weight arrangement
with set_forward_context(
@@ -1175,7 +1243,7 @@ def _test_body_eplb(
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp,
):
output_before = moe_layer(hidden_states, router_logits)
output_before = sp_wrapper(moe_layer)(hidden_states, router_logits)
# Create a fresh FusedMoE layer with enable_eplb=True
# Delete the original layer's registration so the constructor can
@@ -1206,6 +1274,7 @@ def _test_body_eplb(
gate=gate,
routed_input_transform=routed_input_transform,
routed_output_transform=routed_output_transform,
is_sequence_parallel=is_sequence_parallel,
)
if eplb_moe_layer._expert_map is not None:
@@ -1267,7 +1336,7 @@ def _test_body_eplb(
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp,
):
output_after = eplb_moe_layer(hidden_states, router_logits)
output_after = sp_wrapper(eplb_moe_layer)(hidden_states, router_logits)
return output_before, output_after
@@ -1275,11 +1344,12 @@ def _test_body_eplb(
# TODO: make this take a MoETestConfig
def _run_one_config(
vllm_config: VllmConfig,
ep_size: int,
dp_size: int,
tp_size: int,
dp_rank: int,
tp_rank: int,
ep_size: int, # Expert parallel size (total across all ranks)
dp_size: int, # Data parallel size (number of DP groups)
tp_size: int, # Tensor parallel size OR sequence parallel size (when use_ep=True)
dp_rank: int, # Current rank in data parallel dimension
tp_rank: int, # Current rank in tensor/sequence parallel dimension
is_sequence_parallel: bool, # Whether to use sequence parallelism
m: int,
n: int,
k: int,
@@ -1293,15 +1363,22 @@ def _run_one_config(
use_routed_input_transform: bool,
**kwargs,
) -> None:
set_random_seed(7)
"""Generic test loop that sets up environment and delegates to test_body_fn.
This function is called directly by test_moe_layer and test_moe_layer_eplb
via parallel_launch_with_config, passing either _test_body_regular or
_test_body_eplb as the test_body_fn parameter.
Parameter Interpretation:
- When is_sequence_parallel=False (standard TP or EP):
* ep_size: Number of expert parallel ranks (or 1 if no EP)
* tp_size: Number of tensor parallel ranks (or 1 if no TP)
* Weights are chunked by ep_size (experts) and tp_size (tensors)
- When is_sequence_parallel=True (EP + sequence splitting):
* ep_size: Number of expert parallel ranks (equals dp_size * tp_size)
* tp_size: Number of ranks to split sequence across (becomes sp_size in FusedMoE)
* Weights are chunked by ep_size (experts) but NOT by tp_size
* Input sequences are chunked by tp_size (via sp_wrapper)
"""
world_size = tp_size * dp_size
set_random_seed(7)
use_ep = ep_size > 1
assert vllm_config.parallel_config.enable_expert_parallel == use_ep
@@ -1337,6 +1414,8 @@ def _run_one_config(
routed_output_transform = test_data.routed_output_transform
activation = "silu"
# Create baseline layer with FULL weights (no EP chunking)
# Baseline represents the expected output using full model
baseline_layer = make_fake_moe_layer(
w1=w1,
w2=w2,
@@ -1354,26 +1433,45 @@ def _run_one_config(
ep_size=ep_size,
dp_size=dp_size,
activation=activation,
is_sequence_parallel=is_sequence_parallel,
)
baseline_output = baseline_layer(hidden_states, router_logits)
with set_current_vllm_config(vllm_config):
# Compute baseline output with SP wrapper if needed
# sp_wrapper handles sequence chunking/gathering for SP
baseline_output = sp_wrapper(baseline_layer, is_sequence_parallel)(
hidden_states, router_logits
)
del baseline_layer
torch.accelerator.empty_cache()
with set_current_vllm_config(vllm_config):
# Chunk weights for EP/TP (after baseline is created)
# Chunk weights for EP BEFORE creating FusedMoE
# FusedMoE uses EP-chunked weights and handles reductions internally
if ep_size > 1:
w1 = chunk_by_rank(w1, dp_rank, dp_size, dim=0, device=device)
w2 = chunk_by_rank(w2, dp_rank, dp_size, dim=0, device=device)
# Split experts across ranks (dimension 0 is the expert dimension)
# When EP is enabled, use EP group rank and ep_size for chunking
ep_rank = get_ep_group().rank_in_group
w1 = chunk_by_rank(w1, ep_rank, ep_size, dim=0, device=device)
w2 = chunk_by_rank(w2, ep_rank, ep_size, dim=0, device=device)
if tp_size > 1:
# Chunk weights for TP (only if NOT doing sequence parallelism)
# Sequence parallelism splits tokens/sequences, not weight tensors
if tp_size > 1 and not is_sequence_parallel:
w1 = tp_chunk_gate_up(w1, tp_rank, tp_size, dim=1, device=device)
w2 = chunk_by_rank(w2, tp_rank, tp_size, dim=2, device=device)
# Setup shared experts if needed
# In SP mode, shared experts should NOT be TP-chunked (same as routed experts)
# tp_size is used for sequence splitting, not weight splitting
shared_experts = create_shared_experts_from_config(
shared_experts_config, in_dtype, tp_size, tp_rank, device
shared_experts_config,
in_dtype,
tp_size,
tp_rank,
is_sequence_parallel,
device,
)
# Determine hidden size for MoE layer
@@ -1399,14 +1497,17 @@ def _run_one_config(
routed_input_transform=routed_input_transform,
routed_output_transform=routed_output_transform,
activation=activation,
is_sequence_parallel=is_sequence_parallel,
)
if moe_layer._expert_map is not None:
moe_layer._expert_map = moe_layer._expert_map.to(device)
num_tokens = m
# num_tokens_across_dp should have one entry per DP group, not per total rank
# When EP is enabled, dp_size represents the number of DP groups
num_tokens_across_dp = torch.tensor(
[num_tokens] * world_size,
[num_tokens] * dp_size,
device=device,
dtype=torch.int,
)
@@ -1448,7 +1549,7 @@ def _run_one_config(
else:
atol, rtol = 3.5e-2, 3.5e-2
elif quantization in ("fp8", "fp8_blocked", "modelopt_fp8"):
atol, rtol = 6e-2, 6e-2
atol, rtol = 6.5e-2, 6.5e-2
elif quantization == "modelopt_fp4":
if k >= 2048:
atol = rtol = 1e-1 + (k * 1e-4)
@@ -1532,6 +1633,7 @@ def test_moe_layer_no_parallel(
test_config.tp_size,
0,
0,
False,
test_config.m,
test_config.n,
test_config.k,
@@ -1592,6 +1694,7 @@ def _parallel_worker(
test_config.tp_size,
dp_rank,
tp_rank,
test_config.is_sequence_parallel,
test_config.m,
test_config.n,
test_config.k,
@@ -1676,7 +1779,12 @@ def test_moe_layer(
"""
num_gpus = current_platform.device_count()
world_size = tp_size * dp_size
ep_size = 1 if not use_ep else world_size # or dp_size?
# When use_ep=True: FusedMoEParallelConfig flattens tp_size across dp ranks
# Result: ep_size = dp_size * pcp_size * tp_size
# Since pcp_size=1 in these tests: ep_size = dp_size * tp_size = world_size
# When use_ep=False: no expert parallelism, ep_size = 1
ep_size = 1 if not use_ep else world_size
assert world_size > 1
# Check if enough GPUs available
@@ -10,7 +10,9 @@ import pytest
import torch
from vllm.model_executor.layers.fused_moe import fused_topk
from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
from vllm.model_executor.layers.fused_moe.expert_map_manager import (
determine_expert_map,
)
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
moe_permute,
moe_permute_unpermute_supported,
+4 -7
View File
@@ -36,9 +36,11 @@ TOP_KS = [2, 4, 6]
NUM_EXPERTS = [8, 16, 64]
def setup_eplb_state(enable_eplb: bool, global_num_experts: int) -> EplbLayerState:
def setup_eplb_state(
enable_eplb: bool, global_num_experts: int
) -> EplbLayerState | None:
if not enable_eplb:
return EplbLayerState()
return None
# Initialize EPLB state with proper tensors for testing
# For testing purposes, we use a simple 1:1 mapping (no redundant experts)
@@ -349,7 +351,6 @@ def test_fused_topk(
top_k=top_k,
global_num_experts=global_num_experts,
renormalize=renormalize,
enable_eplb=enable_eplb,
eplb_state=eplb_state,
)
@@ -400,7 +401,6 @@ def test_fused_topk_bias(
top_k=top_k,
global_num_experts=global_num_experts,
renormalize=renormalize,
enable_eplb=enable_eplb,
eplb_state=eplb_state,
)
@@ -469,7 +469,6 @@ def test_grouped_topk(
top_k=top_k,
global_num_experts=global_num_experts,
renormalize=renormalize,
enable_eplb=enable_eplb,
eplb_state=eplb_state,
)
@@ -540,7 +539,6 @@ def test_custom(
global_num_experts=global_num_experts,
custom_routing_function=custom_routing_function,
renormalize=renormalize,
enable_eplb=enable_eplb,
eplb_state=eplb_state,
)
@@ -580,7 +578,6 @@ def test_custom(
# router = create_fused_moe_router(
# top_k=top_k,
# global_num_experts=global_num_experts,
# enable_eplb=enable_eplb,
# eplb_state=eplb_state,
# )
@@ -183,6 +183,7 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
x=hidden_states,
topk_weights=topk_weights,
topk_ids=topk_ids,
shared_experts=None,
shared_experts_input=None,
)
+4 -4
View File
@@ -17,13 +17,13 @@ from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
TritonExperts,
)
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
BatchedTritonExperts,
NaiveBatchedExperts,
)
from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
TritonExperts,
)
from vllm.model_executor.layers.fused_moe.fused_moe import (
fused_experts,
)
@@ -143,6 +143,15 @@ def test_models(
# in parts of the operators
pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
if model == "bigcode/starcoder2-3b":
# Replace example.txt's Test1 (an NL prompt) with a code prompt:
# starcoder2-3b is a code model, so NL prompts give near-uniform
# digit logits where HF<->vLLM bf16 drift can reorder top-K.
example_prompts = list(example_prompts)
example_prompts[1] = (
"def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - "
)
with hf_runner(model) as hf_model:
hf_outputs = hf_model.generate_greedy_logprobs_limit(
example_prompts, max_tokens, num_logprobs
@@ -42,6 +42,18 @@ def qwen_vl_chat_template(content: str) -> str:
MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
"qwen2_5_vl": VitCudagraphTestConfig(
model="Qwen/Qwen2.5-VL-3B-Instruct",
image_prompt=qwen_vl_chat_template(
"<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
),
video_prompt=qwen_vl_chat_template(
"<|vision_start|><|video_pad|><|vision_end|>"
"Describe this video in one sentence."
),
needs_video_metadata=False,
marks=[pytest.mark.core_model],
),
"qwen3_vl": VitCudagraphTestConfig(
model="Qwen/Qwen3-VL-2B-Instruct",
image_prompt=qwen_vl_chat_template(
@@ -54,8 +66,8 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
needs_video_metadata=True,
marks=[pytest.mark.core_model],
),
"qwen2_5_vl": VitCudagraphTestConfig(
model="Qwen/Qwen2.5-VL-3B-Instruct",
"qwen3_5": VitCudagraphTestConfig(
model="Qwen/Qwen3.5-0.8B",
image_prompt=qwen_vl_chat_template(
"<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
),
@@ -63,7 +75,7 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
"<|vision_start|><|video_pad|><|vision_end|>"
"Describe this video in one sentence."
),
needs_video_metadata=False,
needs_video_metadata=True,
marks=[pytest.mark.core_model],
),
}
@@ -0,0 +1,58 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for Step3-VL precomputed image embedding inputs."""
import pytest
import torch
from vllm.model_executor.models.step3_vl import (
Step3VLForConditionalGeneration,
Step3VLImageEmbeddingInputs,
)
class _FakeStep3VL:
@staticmethod
def _process_image_features(image_features: torch.Tensor) -> torch.Tensor:
return image_features
def test_image_embedding_inputs_construction():
"""Step3VLImageEmbeddingInputs should store embeddings in the data field."""
image_embeds = torch.randn(2, 16, 64)
inputs = Step3VLImageEmbeddingInputs(
type="image_embeds",
data=image_embeds,
)
assert inputs["type"] == "image_embeds"
assert torch.equal(inputs["data"], image_embeds)
assert torch.equal(inputs.data, image_embeds)
def test_image_embedding_inputs_validation_rejects_wrong_rank():
"""Validation should reject tensors with wrong rank."""
with pytest.raises(ValueError, match="rank"):
Step3VLImageEmbeddingInputs(
type="image_embeds",
data=torch.randn(16, 64),
)
def test_process_image_embeds_does_not_require_pixel_input_fields():
"""The image_embeds branch should not reference patch pixel metadata."""
image_embeds = torch.randn(2, 4, 8)
image_input = Step3VLImageEmbeddingInputs(
type="image_embeds",
data=image_embeds,
)
outputs = Step3VLForConditionalGeneration._process_image_input(
_FakeStep3VL(),
image_input,
)
assert len(outputs) == 2
assert torch.equal(outputs[0], image_embeds[0])
assert torch.equal(outputs[1], image_embeds[1])
+19
View File
@@ -1098,6 +1098,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
},
trust_remote_code=True,
),
"MiniCPMV4_6ForConditionalGeneration": _HfExamplesInfo(
"openbmb/MiniCPM-V-4_6",
min_transformers_version="5.7.0",
),
"MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo(
"MiniMaxAI/MiniMax-VL-01",
trust_remote_code=True,
@@ -1489,6 +1493,21 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
"Qwen/Qwen3-VL-8B-Instruct",
speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
),
# [PEagle]
"PEagleDraftModel": _HfExamplesInfo(
"Qwen/Qwen3-8B",
trust_remote_code=True,
speculative_model="nm-testing/qwen3-8b-peagle-speculators",
tokenizer="Qwen/Qwen3-8B",
use_original_num_layers=True,
),
"PeagleLlamaForCausalLM": _HfExamplesInfo(
"Qwen/Qwen3-8B",
trust_remote_code=True,
speculative_model="nm-testing/qwen3-8b-peagle-speculators",
tokenizer="Qwen/Qwen3-8B",
use_original_num_layers=True,
),
# [MTP]
"DeepSeekMTPModel": _HfExamplesInfo(
"luccafong/deepseek_mtp_main_random",
+21 -1
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# test_audio.py
import math
from unittest.mock import patch
import numpy as np
@@ -45,7 +46,6 @@ def test_resample_audio_scipy(dummy_audio):
assert np.all(out_same == dummy_audio)
@pytest.mark.xfail(reason="resample_audio_scipy is buggy for non-integer ratios")
def test_resample_audio_scipy_non_integer_ratio(dummy_audio):
out = resample_audio_scipy(dummy_audio, orig_sr=5, target_sr=3)
@@ -56,6 +56,26 @@ def test_resample_audio_scipy_non_integer_ratio(dummy_audio):
assert np.isfinite(out).all()
def test_resample_audio_scipy_non_divisible_sample_rates():
audio = np.arange(441, dtype=float)
out = resample_audio_scipy(audio, orig_sr=44100, target_sr=16000)
expected_len = math.ceil(len(audio) * 16000 / 44100)
assert len(out) == expected_len
assert isinstance(out, np.ndarray)
assert np.isfinite(out).all()
def test_resample_audio_scipy_resamples_last_axis_for_multichannel():
audio = np.arange(2 * 441, dtype=float).reshape(2, 441)
out = resample_audio_scipy(audio, orig_sr=44100, target_sr=16000)
expected_len = math.ceil(audio.shape[-1] * 16000 / 44100)
assert out.shape == (2, expected_len)
assert np.isfinite(out).all()
def test_audio_resampler_pyav_calls_resample(dummy_audio):
resampler = AudioResampler(target_sr=22050, method="pyav")
with patch("vllm.multimodal.audio.resample_audio_pyav") as mock_resample:
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
CompressedTensorsConfig,
CompressedTensorsLinearMethod,
CompressedTensorsW4A4Fp4,
CompressedTensorsW4A4Mxfp4,
CompressedTensorsW4A8Fp8,
CompressedTensorsW4A16Fp4,
CompressedTensorsW8A8Fp8,
@@ -689,3 +690,31 @@ def test_compressed_tensors_mxfp8_moe_setup(vllm_runner):
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=4)
assert output
@pytest.mark.skipif(
not current_platform.is_cuda() or not current_platform.has_device_capability(80),
reason="MXFP4 requires ampere or newer",
)
def test_compressed_tensors_mxfp4(vllm_runner):
model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-MXFP4"
with vllm_runner(model_path, enforce_eager=True) as llm:
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
o_proj = layer.self_attn.o_proj
gate_up_proj = layer.mlp.gate_up_proj
down_proj = layer.mlp.down_proj
for proj in (qkv_proj, o_proj, gate_up_proj, down_proj):
assert isinstance(proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(proj.scheme, CompressedTensorsW4A4Mxfp4)
# Verify group size
assert proj.scheme.group_size == 32
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=4)
assert output
+1 -1
View File
@@ -14,7 +14,7 @@ import torch
from vllm.platforms import current_platform
if not current_platform.has_device_capability(90):
if not current_platform.has_device_capability(90) or current_platform.is_rocm():
pytest.skip(
"Machete W4A16 requires Hopper (sm_90).",
allow_module_level=True,
+148
View File
@@ -43,11 +43,16 @@ from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupSpec,
KVCacheSpec,
KVCacheSpecKind,
KVCacheTensor,
MambaSpec,
MLAAttentionSpec,
SinkFullAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
UniformTypeKVCacheSpecs,
get_kv_cache_spec_kind,
get_kv_cache_spec_sliding_window,
)
from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
from vllm.v1.request import Request
@@ -1865,6 +1870,149 @@ def new_mla_spec(cache_dtype_str=None):
)
def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():
assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION
sliding_window_mla_spec = SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
)
assert (
get_kv_cache_spec_kind(sliding_window_mla_spec)
== KVCacheSpecKind.SLIDING_WINDOW_MLA
)
sink_full_attention_spec = SinkFullAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sink_len=4,
)
assert (
get_kv_cache_spec_kind(sink_full_attention_spec)
== KVCacheSpecKind.SINK_FULL_ATTENTION
)
def test_get_kv_cache_spec_kind_unwraps_uniform_type_specs():
uniform_mla_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": new_mla_spec(),
"layer_2": new_mla_spec(cache_dtype_str="fp8"),
},
)
assert get_kv_cache_spec_kind(uniform_mla_spec) == KVCacheSpecKind.MLA_ATTENTION
uniform_swa_mla_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
),
"layer_2": SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=1024,
dtype=torch.float32,
sliding_window=128,
),
},
)
assert (
get_kv_cache_spec_kind(uniform_swa_mla_spec)
== KVCacheSpecKind.SLIDING_WINDOW_MLA
)
def test_get_kv_cache_spec_kind_unknown_for_mixed_uniform_type_specs():
uniform_mixed_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": new_mla_spec(),
"layer_2": SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
),
},
)
assert get_kv_cache_spec_kind(uniform_mixed_spec) == KVCacheSpecKind.UNKNOWN
def test_get_kv_cache_spec_sliding_window_reads_windowed_specs():
full_attention_spec = FullAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
)
sliding_window_spec = SlidingWindowSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sliding_window=128,
)
assert get_kv_cache_spec_sliding_window(full_attention_spec) is None
assert get_kv_cache_spec_sliding_window(sliding_window_spec) == 128
def test_get_kv_cache_spec_sliding_window_unwraps_uniform_type_specs():
uniform_window_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": SlidingWindowSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sliding_window=128,
),
"layer_2": SlidingWindowSpec(
block_size=16,
num_kv_heads=2,
head_size=64,
dtype=torch.float32,
sliding_window=128,
),
},
)
mixed_window_spec = UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer_1": SlidingWindowSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sliding_window=128,
),
"layer_2": SlidingWindowSpec(
block_size=16,
num_kv_heads=1,
head_size=64,
dtype=torch.float32,
sliding_window=256,
),
},
)
assert get_kv_cache_spec_sliding_window(uniform_window_spec) == 128
assert get_kv_cache_spec_sliding_window(mixed_window_spec) is None
def test_merge_mla_spec():
kv_cache_specs = [
new_mla_spec(),
+98 -9
View File
@@ -8,6 +8,7 @@ from collections.abc import Callable
import pytest
import torch
import vllm.v1.core.kv_cache_manager as kv_cache_manager
import vllm.v1.core.kv_cache_utils as kv_cache_utils
from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored
from vllm.lora.request import LoRARequest
@@ -35,6 +36,7 @@ from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheSpecKind,
MambaSpec,
SlidingWindowSpec,
)
@@ -1933,6 +1935,7 @@ def test_kv_cache_events(blocks_to_cache: int):
== len(manager.block_pool.cached_block_hash_to_block)
)
assert len(block.token_ids) == block.block_size * len(block.block_hashes)
assert block.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
assert len(manager.block_pool.kv_event_queue) == 0
stored_block_hash = block.block_hashes
@@ -1946,6 +1949,7 @@ def test_kv_cache_events(blocks_to_cache: int):
events = manager.take_events()
for blocks in events[:-1]:
assert isinstance(blocks, BlockRemoved)
assert blocks.block_hashes[0] in stored_block_hash
assert len(events) == blocks_to_cache + 1
assert isinstance(events[-2], BlockRemoved)
@@ -2022,6 +2026,8 @@ def test_null_parent_block_hash():
]
assert event.block_hashes == expected_new_hashes
assert event.group_idx == kv_cache_group_id
assert event.kv_cache_spec_kind is None
assert event.kv_cache_spec_sliding_window is None
# Ensure we didn't accidentally assign a hash to the null block.
assert pool.null_block.block_hash is None
@@ -2095,12 +2101,14 @@ def test_block_stored_event_group_idx(group_id: int):
block_size = 4
num_tokens = block_size * 2
pool = BlockPool(
num_gpu_blocks=5,
manager = KVCacheManager(
make_kv_cache_config_three_types(block_size, num_blocks=5),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
enable_kv_cache_events=True,
hash_block_size=block_size,
)
pool = manager.block_pool
req = make_request(
"req_grp_idx",
@@ -2119,10 +2127,26 @@ def test_block_stored_event_group_idx(group_id: int):
kv_cache_group_id=group_id,
)
events = pool.take_events()
events = manager.take_events()
assert len(events) == 1
assert isinstance(events[0], BlockStored)
assert events[0].group_idx == group_id
assert (
events[0].kv_cache_spec_kind
== [
KVCacheSpecKind.FULL_ATTENTION.value,
KVCacheSpecKind.SLIDING_WINDOW.value,
KVCacheSpecKind.MAMBA.value,
][group_id]
)
assert (
events[0].kv_cache_spec_sliding_window
== [
None,
2 * block_size,
None,
][group_id]
)
def test_block_stored_event_group_idx_multiple_groups():
@@ -2137,13 +2161,38 @@ def test_block_stored_event_group_idx_multiple_groups():
block_size = 4
num_tokens = block_size * 2
# null block + 4 usable (2 per group)
pool = BlockPool(
num_gpu_blocks=5,
manager = KVCacheManager(
KVCacheConfig(
num_blocks=5,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["layer1"],
FullAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
),
KVCacheGroupSpec(
["layer2"],
SlidingWindowSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
sliding_window=128,
),
),
],
),
max_model_len=8192,
enable_caching=True,
hash_block_size=block_size,
enable_kv_cache_events=True,
hash_block_size=block_size,
)
pool = manager.block_pool
req = make_request(
"req_multi_grp",
@@ -2174,12 +2223,52 @@ def test_block_stored_event_group_idx_multiple_groups():
kv_cache_group_id=1,
)
events = pool.take_events()
events = manager.take_events()
assert len(events) == 2
assert isinstance(events[0], BlockStored)
assert events[0].group_idx == 0
assert events[0].kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
assert events[0].kv_cache_spec_sliding_window is None
assert isinstance(events[1], BlockStored)
assert events[1].group_idx == 1
assert events[1].kv_cache_spec_kind == KVCacheSpecKind.SLIDING_WINDOW.value
assert events[1].kv_cache_spec_sliding_window == 128
def test_block_stored_event_group_idx_out_of_bounds(monkeypatch):
"""Out-of-range group_idx events are returned without metadata annotation."""
block_size = 4
manager = KVCacheManager(
make_kv_cache_config(block_size, num_blocks=5),
max_model_len=8192,
enable_caching=True,
enable_kv_cache_events=True,
hash_block_size=block_size,
)
event = BlockStored(
block_hashes=[1],
parent_block_hash=None,
token_ids=list(range(block_size)),
block_size=block_size,
lora_id=None,
medium=None,
lora_name=None,
group_idx=1,
)
manager.block_pool.kv_event_queue.append(event)
warnings = []
def collect_warning(message, *args, **kwargs):
del kwargs
warnings.append(message % args if args else message)
monkeypatch.setattr(kv_cache_manager.logger, "warning", collect_warning)
events = manager.take_events()
assert events == [event]
assert event.kv_cache_spec_kind is None
assert event.kv_cache_spec_sliding_window is None
assert warnings == ["Group index `1` not in KV cache metadata"]
@pytest.mark.parametrize("group_id", [0, 1, 2])
@@ -429,7 +429,7 @@ def _logprobs_match(
and lps_a.keys() == lps_b.keys()
and all(
a.decoded_token == b.decoded_token
and a.rank == b.rank
and a.rank == pytest.approx(b.rank, rel=0.005)
and a.logprob == pytest.approx(b.logprob, rel=rel_tol, abs=abs_tol)
for a, b in ((lps_a[x], lps_b[x]) for x in lps_a)
)
+2 -1
View File
@@ -79,10 +79,11 @@ def _get_test_sampling_params(
structured_outputs: bool = False,
) -> tuple[list[SamplingParams], list[int]]:
"""Generate random sampling params for a batch."""
rng = random.Random(seed)
def get_mostly_n_gt1() -> int:
r"""Mostly n \in [2,20], ~1/3 n=1"""
x = random.randint(0, 28)
x = rng.randint(0, 28)
if x < 10:
return 1
else:
@@ -627,7 +627,7 @@ def _make_scheduler_with_lookup(
return scheduler
_EMPTY_REQ_CTX = ReqContext()
_EMPTY_REQ_CTX = ReqContext(req_id="")
class TestMaximalPrefixLookup:
@@ -459,9 +459,6 @@ class RequestRunner:
self.worker_connector.bind_connector_metadata(kv_connector_metadata)
self.worker_connector.start_load_kv(self._dummy_ctx)
if scheduler_output.total_num_scheduled_tokens > 0:
self.worker_connector.wait_for_save()
if complete_transfers:
self.offloading_spec.complete_transfers()
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock, patch
from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( # noqa: E501
ExampleConnectorMetadata,
)
@@ -38,11 +40,20 @@ def test_kv_connector_mixin_clears_metadata():
vllm_config.kv_transfer_config.kv_role = "kv_both"
vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = "unit"
# Initialize the global connector instance
kv_cache_config = KVCacheConfig(
num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
)
ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
# Initialize the global connector instance.
# kv_transfer init now syncs engine_id across TP, so unit tests need
# a minimal mocked TP group.
mock_tp_group = MagicMock()
mock_tp_group.broadcast_object.side_effect = lambda value, src=0: value
with patch(
"vllm.distributed.parallel_state.get_tp_group",
return_value=mock_tp_group,
):
ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
try:
# Minimal scheduler output with empty metadata; mixin should still
@@ -0,0 +1,258 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock, patch
from vllm.config import set_current_vllm_config
from vllm.distributed.kv_events import BlockStored
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
KVConnectorRole,
)
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
connector,
worker,
)
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
MooncakeStoreConnectorMetadata,
)
from vllm.v1.outputs import KVConnectorOutput
from .utils import create_vllm_config
def _make_vllm_config():
return create_vllm_config(
kv_connector="MooncakeStoreConnector",
kv_role="kv_both",
)
def _make_block_stored() -> BlockStored:
return BlockStored(
block_hashes=[b"hash"],
parent_block_hash=None,
token_ids=[1, 2, 3],
block_size=16,
lora_id=None,
medium="cpu",
lora_name=None,
)
def test_scheduler_role_initializes_store_scheduler_only():
vllm_config = _make_vllm_config()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreScheduler"
) as mock_scheduler,
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreWorker"
) as mock_worker,
):
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
mock_scheduler.assert_called_once_with(vllm_config)
mock_worker.assert_not_called()
assert conn.connector_scheduler is mock_scheduler.return_value
assert conn.connector_worker is None
def test_worker_role_initializes_store_worker_on_rank0():
vllm_config = _make_vllm_config()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreScheduler"
) as mock_scheduler,
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreWorker"
) as mock_worker,
):
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
mock_scheduler.assert_not_called()
mock_worker.assert_called_once_with(vllm_config)
assert conn.connector_scheduler is None
assert conn.connector_worker is mock_worker.return_value
def test_worker_role_initializes_on_nonzero_rank():
vllm_config = _make_vllm_config()
vllm_config.parallel_config.rank = 1
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreWorker"
) as mock_worker,
):
connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
mock_worker.assert_called_once_with(vllm_config)
def test_lookup_rpc_path_uses_data_parallel_index_in_dense_dp():
vllm_config = _make_vllm_config()
vllm_config.parallel_config.data_parallel_rank = 0
vllm_config.parallel_config.data_parallel_index = 3
path = worker.get_zmq_rpc_path_lookup(vllm_config)
assert path.endswith("_dp_rank3")
def test_lookup_rpc_path_uses_local_rank_when_local_engines_only():
vllm_config = _make_vllm_config()
vllm_config.parallel_config.data_parallel_index = 7
vllm_config.parallel_config.data_parallel_rank_local = 1
vllm_config.parallel_config.data_parallel_hybrid_lb = True
path = worker.get_zmq_rpc_path_lookup(vllm_config)
assert path.endswith("_dp_rank1")
def test_worker_methods_delegate_to_store_worker():
vllm_config = _make_vllm_config()
kv_caches = {"layer0": MagicMock()}
metadata = MooncakeStoreConnectorMetadata(set(), set())
finished_req_ids = {"req-1"}
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreWorker"
) as mock_worker_cls,
):
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
worker_inst = mock_worker_cls.return_value
worker_inst.get_finished.return_value = ({"req-1"}, {"req-2"})
conn.bind_connector_metadata(metadata)
conn.register_kv_caches(kv_caches)
result = conn.get_finished(finished_req_ids)
worker_inst.register_kv_caches.assert_called_once_with(kv_caches)
worker_inst.get_finished.assert_called_once_with(finished_req_ids, metadata)
assert result == ({"req-1"}, {"req-2"})
def test_get_kv_connector_kv_cache_events_returns_none_when_empty():
vllm_config = _make_vllm_config()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreWorker"
) as mock_worker_cls,
):
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
mock_worker_cls.return_value.get_kv_events.return_value = []
assert conn.get_kv_connector_kv_cache_events() is None
def test_get_kv_connector_kv_cache_events_wraps_worker_events():
vllm_config = _make_vllm_config()
event = _make_block_stored()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreWorker"
) as mock_worker_cls,
):
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
mock_worker_cls.return_value.get_kv_events.return_value = [event]
kv_events = conn.get_kv_connector_kv_cache_events()
assert isinstance(kv_events, connector.MooncakeStoreKVEvents)
assert kv_events.get_number_of_workers() == 1
assert kv_events.get_all_events() == [event]
def test_prefer_cross_layer_blocks_from_config():
# Default: disabled
vllm_config = _make_vllm_config()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreScheduler"
),
):
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
assert conn.prefer_cross_layer_blocks is False
# Enabled via config
vllm_config_enabled = create_vllm_config(
kv_connector="MooncakeStoreConnector",
kv_role="kv_both",
kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
)
with (
set_current_vllm_config(vllm_config_enabled),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreScheduler"
),
):
conn_enabled = connector.MooncakeStoreConnector(
vllm_config_enabled, KVConnectorRole.SCHEDULER
)
assert conn_enabled.prefer_cross_layer_blocks is True
def test_register_cross_layers_kv_cache_delegates_to_worker():
vllm_config = _make_vllm_config()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreWorker"
) as mock_worker_cls,
):
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
fake_tensor = MagicMock()
fake_backend = MagicMock()
conn.register_cross_layers_kv_cache(fake_tensor, fake_backend)
worker_inst = mock_worker_cls.return_value
worker_inst.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
def test_update_connector_output_and_take_events():
vllm_config = _make_vllm_config()
event = _make_block_stored()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"connector.MooncakeStoreScheduler"
),
):
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
kv_events = connector.MooncakeStoreKVEvents(num_workers=1)
kv_events.add_events([event])
conn.update_connector_output(KVConnectorOutput(kv_cache_events=kv_events))
assert conn._kv_cache_events is kv_events
assert list(conn.take_events()) == [event]
assert conn._kv_cache_events is None
@@ -0,0 +1,300 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
from unittest.mock import MagicMock, patch
import torch
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
worker,
)
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
ChunkedTokenDatabase,
KeyMetadata,
ReqMeta,
)
def _make_store_sending_thread(
store: MagicMock,
) -> worker.KVCacheStoreSendingThread:
token_database = ChunkedTokenDatabase(
KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
)
token_database.set_kv_caches_base_addr([0x1000])
token_database.set_block_len([256])
thread = worker.KVCacheStoreSendingThread(
store=store,
token_database=token_database,
block_size=16,
tp_rank=0,
put_step=1,
kv_role="kv_producer",
ready_event=threading.Event(),
)
thread.request_queue.task_done = MagicMock()
return thread
def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
return ReqMeta(
req_id=req_id,
token_len_chunk=32,
block_ids=[0, 1],
block_hashes=block_hashes,
can_save=True,
original_block_size=16,
)
def test_store_sending_thread_skips_request_during_cpu_pressure():
store = MagicMock()
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
store.batch_put_from_multi_buffers.side_effect = [
[-200, -200],
[256, 256],
[256, 256],
]
thread = _make_store_sending_thread(store)
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
assert thread._store_pressure_active is True
assert "req-a" in thread._skip_store_requests
assert store.batch_put_from_multi_buffers.call_count == 1
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))
assert store.batch_put_from_multi_buffers.call_count == 1
thread.add_stored_request("req-b")
thread._handle_request(_make_store_req("req-b", [b"b0", b"b1"]))
assert thread._store_pressure_active is False
assert "req-a" not in thread._skip_store_requests
assert store.batch_put_from_multi_buffers.call_count == 2
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a4", b"a5"]))
assert store.batch_put_from_multi_buffers.call_count == 3
def test_store_sending_thread_only_skips_on_no_available_handle():
store = MagicMock()
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
store.batch_put_from_multi_buffers.side_effect = [
[-500, -500],
[256, 256],
]
thread = _make_store_sending_thread(store)
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
assert thread._store_pressure_active is False
assert "req-a" not in thread._skip_store_requests
assert store.batch_put_from_multi_buffers.call_count == 1
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))
assert store.batch_put_from_multi_buffers.call_count == 2
# ---------------------------------------------------------------------------
# Helpers for register_kv_caches tests
# ---------------------------------------------------------------------------
def _auto_set_ready_event(*args, **kwargs):
"""Side effect for mocked thread constructors that auto-sets ready_event."""
for arg in args:
if isinstance(arg, threading.Event):
arg.set()
for val in kwargs.values():
if isinstance(val, threading.Event):
val.set()
return MagicMock()
def _make_bare_worker(
*,
num_gpu_blocks: int = 10,
block_size: int = 16,
kv_role: str = "kv_both",
) -> worker.MooncakeStoreWorker:
"""Construct a MooncakeStoreWorker via __new__, bypassing __init__.
Sets only the attributes that register_kv_caches() reads so we can
test the stride-based layout detection without a real
MooncakeDistributedStore.
"""
w = object.__new__(worker.MooncakeStoreWorker)
w.cache_config = MagicMock()
w.cache_config.num_gpu_blocks = num_gpu_blocks
w.store = MagicMock()
w.store.register_buffer.return_value = 0
w.use_mla = False
w.token_database = ChunkedTokenDatabase(
KeyMetadata("test-model", 0, 0, 0, 0), block_size=block_size
)
w.kv_role = kv_role
w.block_size = block_size
w.tp_rank = 0
w.put_step = 1
w.enable_kv_events = False
w.kv_send_thread = None
w.kv_recv_thread = None
return w
# ---------------------------------------------------------------------------
# register_kv_caches tests
# ---------------------------------------------------------------------------
def test_register_kv_caches_blocks_first_single_segment():
"""Blocks-first layout (FlashInfer/MLA): one segment per layer."""
num_blocks = 10
page_size_elements = 64 # elements per block
w = _make_bare_worker(num_gpu_blocks=num_blocks)
# Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims
tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16)
with (
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.KVCacheStoreSendingThread",
side_effect=_auto_set_ready_event,
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.KVCacheStoreRecvingThread",
side_effect=_auto_set_ready_event,
),
):
w.register_kv_caches({"layer0": tensor})
assert len(w.kv_caches_base_addr) == 1
assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
assert len(w.block_len) == 1
assert w.block_len[0] == expected_block_len
w.store.register_buffer.assert_called_once_with(
tensor.untyped_storage().data_ptr(),
tensor.untyped_storage().nbytes(),
)
def test_register_kv_caches_kv_first_two_segments():
"""K/V-first layout (FlashAttn): two segments (K, V) per layer."""
num_blocks = 10
block_size_tokens = 16
num_kv_heads = 4
head_size = 8
w = _make_bare_worker(num_gpu_blocks=num_blocks)
# Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost
tensor = torch.zeros(
2,
num_blocks,
block_size_tokens,
num_kv_heads,
head_size,
dtype=torch.float16,
)
with (
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.KVCacheStoreSendingThread",
side_effect=_auto_set_ready_event,
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.KVCacheStoreRecvingThread",
side_effect=_auto_set_ready_event,
),
):
w.register_kv_caches({"layer0": tensor})
# K/V-first: dim 0 has stride > page_size, so 2 segments
assert len(w.kv_caches_base_addr) == 2
assert len(w.block_len) == 2
el = tensor.element_size()
seg_stride = tensor.stride(0) * el # stride of the K/V dim in bytes
base = tensor.untyped_storage().data_ptr()
assert w.kv_caches_base_addr[0] == base
assert w.kv_caches_base_addr[1] == base + seg_stride
assert w.block_len[0] == seg_stride // num_blocks
assert w.block_len[1] == seg_stride // num_blocks
def test_register_kv_caches_cross_layer_single_segment():
"""Cross-layer tensor: single segment with block_len = page_size * num_layers."""
num_blocks = 10
num_layers = 4
per_layer_page_elements = 64 # elements per layer per block
w = _make_bare_worker(num_gpu_blocks=num_blocks)
# Cross-layer blocks-first tensor: all layers packed into a single
# contiguous block. Shape (num_blocks, num_layers * per_layer_page)
# mimics the physical layout after stride reordering.
total_page_elements = num_layers * per_layer_page_elements
tensor = torch.zeros(num_blocks, total_page_elements, dtype=torch.float16)
with (
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.KVCacheStoreSendingThread",
side_effect=_auto_set_ready_event,
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.KVCacheStoreRecvingThread",
side_effect=_auto_set_ready_event,
),
):
# Use the cross-layer wrapper key, same as register_cross_layers_kv_caches
w.register_kv_caches({"__cross_layer__": tensor})
assert len(w.kv_caches_base_addr) == 1
assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
# block_len should be per_layer_page_size * num_layers
assert (
expected_block_len
== num_layers * per_layer_page_elements * tensor.element_size()
)
assert len(w.block_len) == 1
assert w.block_len[0] == expected_block_len
# Also verify via register_cross_layers_kv_caches wrapper
w2 = _make_bare_worker(num_gpu_blocks=num_blocks)
with (
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.KVCacheStoreSendingThread",
side_effect=_auto_set_ready_event,
),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.KVCacheStoreRecvingThread",
side_effect=_auto_set_ready_event,
),
):
w2.register_cross_layers_kv_caches(tensor)
assert w2.kv_caches_base_addr == w.kv_caches_base_addr
assert w2.block_len == w.block_len
@@ -261,10 +261,11 @@ def test_multi_example_connector_consistency():
)
events = get_connector_events()
# First event is set_xfer_handshake_metadata from initialization, then
# on_new_request when the request is enqueued, then
# get_num_new_matched_tokens and update_state_after_alloc from generate().
assert events["storage1-SCHEDULER"][:5] == [
# First event is bind_gpu_block_pool from initialization, then
# set_xfer_handshake_metadata, then on_new_request when the request is enqueued,
# then get_num_new_matched_tokens and update_state_after_alloc from generate().
assert events["storage1-SCHEDULER"][:6] == [
"bind_gpu_block_pool",
"set_xfer_handshake_metadata",
"on_new_request",
"get_num_new_matched_tokens 0",
@@ -283,7 +284,8 @@ def test_multi_example_connector_consistency():
"wait_for_layer_load",
"save_kv_layer",
]
assert events["storage2-SCHEDULER"][:5] == [
assert events["storage2-SCHEDULER"][:6] == [
"bind_gpu_block_pool",
"set_xfer_handshake_metadata",
"on_new_request",
"get_num_new_matched_tokens 0",
@@ -93,27 +93,39 @@ def test_logical_to_kernel_block_ids_with_hma():
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"group_spec_types,expansion_stride,remote_block_ids,expected_remote_block_ids",
"group_spec_types,remote_physical_per_logical,"
"local_physical_per_logical,tp_ratio,remote_block_ids,"
"expected_remote_block_ids",
[
pytest.param(
("FullAttentionSpec", "SlidingWindowSpec"),
2,
2,
1,
([0, 1, 2], [3, 4]),
[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]],
id="dense_fa_swa",
),
# Nemotron-3-Nano-30B-A3B 4p1d (P_TP=4, D_TP=1):
# remote_physical_per_logical=34, local_physical_per_logical=66.
# FA logical block 5 → kernel [170..203], block 6 → [204..237].
# Mamba block unchanged.
pytest.param(
("FullAttentionSpec", "MambaSpec"),
261,
([0, 1, 2], [10, 11]),
[[0, 1, 261, 262, 522, 523], [10, 11]],
34,
66,
-4,
([5, 6], [2]),
[list(range(170, 238)), [2]],
id="mamba_fa_ssm",
),
],
)
def test_read_blocks_for_req_expands_remote_ids(
group_spec_types,
expansion_stride,
remote_physical_per_logical,
local_physical_per_logical,
tp_ratio,
remote_block_ids,
expected_remote_block_ids,
):
@@ -148,7 +160,7 @@ def test_read_blocks_for_req_expands_remote_ids(
resolved_types = tuple(spec_name_to_type[n] for n in group_spec_types)
worker = object.__new__(NixlConnectorWorker)
worker._physical_blocks_per_logical_kv_block = 2
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
has_mamba = any(t is MambaSpec for t in resolved_types)
has_swa = any(t is SlidingWindowSpec for t in resolved_types)
@@ -159,9 +171,11 @@ def test_read_blocks_for_req_expands_remote_ids(
remote_engine_id = "remote-engine"
worker.transfer_topo = MagicMock()
worker.transfer_topo.tp_ratio.return_value = 1
# tp_ratio not exercised (all_source_ranks is empty so no reads run),
# but set for realism.
worker.transfer_topo.tp_ratio.return_value = tp_ratio
remote_info = MagicMock()
remote_info.remote_physical_blocks_per_logical = expansion_stride
remote_info.remote_physical_blocks_per_logical = remote_physical_per_logical
worker.transfer_topo.get_engine_info.return_value = remote_info
worker.use_mla = False
@@ -192,6 +206,168 @@ def test_read_blocks_for_req_expands_remote_ids(
)
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"local_physical_per_logical,remote_physical_per_logical,"
"local_block_ids,remote_block_ids,"
"expected_local,expected_remote",
[
# 10 kernel blocks of data, local has more logical blocks.
# remote physical_per_logical=10 → 1 logical → 10 kernel blocks
# local physical_per_logical=6 → 2 logical → 12 kernel blocks
# Trim local from 12 to 10.
pytest.param(
6,
10,
[list(range(12)), [42]],
[list(range(10)), [42]],
[list(range(10)), [42]],
[list(range(10)), [42]],
id="align_local6_remote10",
),
# 10 kernel blocks of data, remote has more logical blocks.
# remote physical_per_logical=6 → 2 logical → 12 kernel blocks
# local physical_per_logical=10 → 1 logical → 10 kernel blocks
# Trim remote from 12 to 10.
pytest.param(
10,
6,
[list(range(10)), [42]],
[list(range(12)), [42]],
[list(range(10)), [42]],
[list(range(10)), [42]],
id="align_local10_remote6",
),
],
)
def test_apply_prefix_caching_mamba_hybrid(
local_physical_per_logical,
remote_physical_per_logical,
local_block_ids,
remote_block_ids,
expected_local,
expected_remote,
):
"""_apply_prefix_caching front-trims FA groups to
min(local, remote) for Mamba hybrid models with heterogeneous TP.
"""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
)
from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
worker = object.__new__(NixlConnectorWorker)
worker._has_mamba = True
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
worker._group_spec_types = (FullAttentionSpec, MambaSpec)
worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)
aligned_local, aligned_remote = worker._apply_prefix_caching(
local_block_ids, remote_block_ids, remote_physical_per_logical
)
assert aligned_local == expected_local, (
f"Expected local {expected_local}, got {aligned_local}"
)
assert aligned_remote == expected_remote, (
f"Expected remote {expected_remote}, got {aligned_remote}"
)
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"local_physical_per_logical,remote_physical_per_logical,"
"remote_fa_blocks,local_fa_blocks,ssm_blocks,"
"correct_remote_fa,correct_local_fa",
[
# 10 kernel blocks of data (640 tokens).
# remote physical_per_logical=10 → 1 logical → 10 kernel [0..9]
# local physical_per_logical=6 → 2 logical → 12 kernel [0..11]
# 1st local logical block cached → suffix [6..11]
# Correct: transfer only uncached suffix tokens (384-639)
# = remote [6,7,8,9] → local [6,7,8,9].
# Actual (front-trim): remote[:6]=[0..5] → local [6..11]. Wrong.
pytest.param(
6,
10,
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
[6, 7, 8, 9, 10, 11],
[42],
[6, 7, 8, 9],
[6, 7, 8, 9],
id="local6_remote10_fail",
),
# 15 kernel blocks of data (960 tokens).
# remote physical_per_logical=6 → 3 logical → 18 kernel [0..17]
# local physical_per_logical=10 → 2 logical → 20 kernel [0..19]
# 1st local logical block cached → suffix [10..19]
# Correct: transfer only uncached suffix tokens (640-959)
# = remote [10,11,12,13,14] → local [10,11,12,13,14].
# Actual (front-trim): remote[:10]=[0..9] → local [10..19]. Wrong.
pytest.param(
10,
6,
list(range(18)),
list(range(10, 20)),
[42],
[10, 11, 12, 13, 14],
[10, 11, 12, 13, 14],
id="local10_remote6_fail",
),
],
)
def test_mismatched_physical_per_logical_fails_with_prefix_caching(
local_physical_per_logical,
remote_physical_per_logical,
remote_fa_blocks,
local_fa_blocks,
ssm_blocks,
correct_remote_fa,
correct_local_fa,
):
"""Demonstrate that _apply_prefix_caching front-trims ([:N])
in the Mamba hybrid path, which fails when prefix caching produces
suffix-only local blocks.
Prefix caching operates at logical block granularity. When a logical
block is cached locally, the decode side only allocates kernel blocks
for the uncached suffix. The front-trim pairs remote prefix blocks
with local suffix slots a silent data corruption.
"""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
)
worker = object.__new__(NixlConnectorWorker)
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
worker.kv_cache_config = make_kv_cache_config(
block_size=16,
mamba_enabled=True,
)
worker._has_mamba = True
worker._group_spec_types = tuple(
type(g.kv_cache_spec) for g in worker.kv_cache_config.kv_cache_groups
)
local_block_ids = (local_fa_blocks, ssm_blocks)
remote_block_ids = (remote_fa_blocks, ssm_blocks)
aligned_local, aligned_remote = worker._apply_prefix_caching(
local_block_ids,
remote_block_ids,
remote_physical_per_logical,
)
assert (
aligned_remote[0] != correct_remote_fa or aligned_local[0] != correct_local_fa
), (
f"Prefix caching with mismatched physical_per_logical should not "
f"produce correct transfer ids: "
f"remote={aligned_remote[0]}, local={aligned_local[0]}, "
f"correct_remote={correct_remote_fa}, correct_local={correct_local_fa}"
)
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
"""Test that a prefill instance returns fewer "remote blocks" for the SWA groups
@@ -564,3 +740,119 @@ def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_rati
)
assert compute_physical_blocks_per_logical(ssm_sizes, block_len) == expected_ratio
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"mamba_enabled,swa_enabled,"
"local_physical_per_logical,remote_physical_per_logical,"
"logical_block_ids,expected_kernel_block_ids",
[
# Qwen3.5-0.8B 4P2D (kernel_block_size=64):
# prefill TP=4: logical_block_size=384 → physical_per_logical=6
# decode TP=2: logical_block_size=640 → physical_per_logical=10
# FA logical [0] → remote kernel [0..9] (1 * 10)
# SSM logical [10] → unchanged [10]
pytest.param(
True,
False,
6,
10,
([0], [10]),
[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10]],
id="qwen35_4p2d",
),
# Qwen3.5-0.8B 2P4D (kernel_block_size=64):
# prefill TP=2: logical_block_size=640 → physical_per_logical=10
# decode TP=4: logical_block_size=384 → physical_per_logical=6
# FA logical [0, 1] → remote kernel [0..5, 6..11] (2 * 6)
# SSM logical [10] → unchanged [10]
pytest.param(
True,
False,
10,
6,
([0, 1], [10]),
[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], [10]],
id="qwen35_2p4d",
),
# Homogeneous TP (kernel_block_size=64):
# both sides: logical_block_size=640 → physical_per_logical=10
# FA logical [0] → kernel [0..9], SSM unchanged
pytest.param(
True,
False,
10,
10,
([0], [10]),
[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10]],
id="homo_tp",
),
# remote physical_per_logical=1: early return, no expansion
pytest.param(
True,
False,
10,
1,
([0, 1, 2], [5]),
[[0, 1, 2], [5]],
id="mamba_remote_physical_per_logical_1",
),
# Pure FA (no mamba): single group expanded with remote stride
pytest.param(
False,
False,
2,
4,
([0, 1],),
[[0, 1, 2, 3, 4, 5, 6, 7]],
id="pure_fa",
),
# FA + SWA (no mamba): both groups expanded
pytest.param(
False,
True,
2,
3,
([0, 1], [2, 3]),
[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]],
id="fa_swa",
),
],
)
def test_logical_to_remote_kernel_block_ids(
mamba_enabled,
swa_enabled,
local_physical_per_logical,
remote_physical_per_logical,
logical_block_ids,
expected_kernel_block_ids,
):
"""Verify _logical_to_remote_kernel_block_ids uses the remote
physical_per_logical for FA expansion, not the local one.
This was the root cause of silent accuracy corruption in Qwen3.5
heterogeneous TP (e.g. 4P2D): the old code used local physical_per_logical
for the expansion arange, producing wrong kernel block indices.
Qwen3.5-0.8B values verified by verify_conv_split.py (issue #13).
"""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
)
worker = object.__new__(NixlConnectorWorker)
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
worker.kv_cache_config = make_kv_cache_config(
block_size=16,
mamba_enabled=mamba_enabled,
swa_enabled=swa_enabled,
)
result = worker._logical_to_remote_kernel_block_ids(
logical_block_ids,
remote_physical_per_logical,
)
assert list(result) == expected_kernel_block_ids, (
f"Expected {expected_kernel_block_ids}, got {result}"
)
+4 -2
View File
@@ -19,9 +19,11 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext:
def make_req_context(
req_id: str = "", 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)
return ReqContext(req_id=req_id, kv_transfer_params=kv_transfer_params)
_EMPTY_REQ_CTX = make_req_context()
+116 -16
View File
@@ -179,6 +179,7 @@ _req_counter = 0
def make_request(
num_blocks: int = 2,
request_id: str | None = None,
extra_tokens: int = 1,
) -> Request:
"""Create a Request with deterministic block hashes."""
global _req_counter
@@ -186,13 +187,7 @@ def make_request(
if request_id is None:
request_id = f"req-{_req_counter}"
# Add one extra token beyond the last full block so that
# ``max_cache_hit_length = num_tokens - 1`` (see
# KVCacheManager.get_computed_blocks) does not truncate the final
# full block: ``find_longest_cache_hit`` uses
# ``max_length // block_size`` and would otherwise drop one block
# when the prompt is an exact multiple of block_size.
num_tokens = num_blocks * BLOCK_SIZE + 1
num_tokens = num_blocks * BLOCK_SIZE + extra_tokens
start = _req_counter * 10000
prompt_token_ids = list(range(start, start + num_tokens))
sampling_params = SamplingParams(max_tokens=1)
@@ -386,6 +381,8 @@ def test_eager_store_and_load_roundtrip() -> None:
block_hasher=req._block_hasher,
)
hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
# make_request pads num_tokens by +1 beyond the last full block, so the
# manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact.
assert hit_tokens == num_blocks * BLOCK_SIZE
assert is_async is True
@@ -405,7 +402,44 @@ def test_eager_store_and_load_roundtrip() -> None:
# ---------------------------------------------------------------------------
# Test 1b: Lazy store-and-load roundtrip
# Test 1b: Boundary — max_hit_len cap drops the last full block when the
# prompt is an exact multiple of BLOCK_SIZE.
# ---------------------------------------------------------------------------
def test_max_hit_len_cap_drops_last_full_block() -> None:
"""When num_tokens is an exact multiple of BLOCK_SIZE, the manager's
``max_hit_len = num_tokens - 1`` cap forces ``find_longest_cache_hit`` to
drop the final block (since ``max_length // block_size`` rounds down).
"""
fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
sched = fix.scheduler
num_blocks = 2
req = make_request(num_blocks=num_blocks, extra_tokens=0)
assert req.num_tokens == num_blocks * BLOCK_SIZE
kv_blocks = _alloc_and_register(fix, req, num_blocks)
sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
sched_out = make_scheduler_output(
{req.request_id: num_blocks * BLOCK_SIZE},
new_reqs={req.request_id: kv_blocks.get_block_ids()},
)
meta = sched.build_connector_meta(sched_out)
simulate_store_completion(sched, meta.store_event)
req2 = Request(
request_id="req-cap-boundary",
prompt_token_ids=req.prompt_token_ids,
sampling_params=req.sampling_params,
pooling_params=None,
mm_features=None,
block_hasher=req._block_hasher,
)
hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
assert hit_tokens == (num_blocks - 1) * BLOCK_SIZE
# ---------------------------------------------------------------------------
# Test 1c: Lazy store-and-load roundtrip
# ---------------------------------------------------------------------------
def _flush_old_blocks_to_lru_head(
gpu_pool: BlockPool,
@@ -469,8 +503,11 @@ def test_lazy_store_and_load_roundtrip() -> None:
hit_tokens, is_async = sched.get_num_new_matched_tokens(
req_old2, num_computed_tokens=0
)
assert hit_tokens == num_blocks * BLOCK_SIZE, (
f"Expected {num_blocks * BLOCK_SIZE} hit tokens, got {hit_tokens}"
# make_request pads num_tokens by +1 beyond the last full block, so the
# manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact.
expected_hit = num_blocks * BLOCK_SIZE
assert hit_tokens == expected_hit, (
f"Expected {expected_hit} hit tokens, got {hit_tokens}"
)
assert is_async is True
@@ -537,7 +574,72 @@ def test_eager_duplicate_store_skipped() -> None:
# ---------------------------------------------------------------------------
# Test 2b: Lazy duplicate store is skipped
# Test 2b: Eager dedup of in-flight stores across consecutive steps
# ---------------------------------------------------------------------------
def test_eager_in_flight_store_dedup_across_steps() -> None:
"""Eager: a second request sharing a prefix with an in-flight store
must not re-offload the same GPU blocks before completion lands.
Simulates a GPU prefix-cache hit by reusing the first request's
GPU block IDs in the second scheduler step, which is the path the
real scheduler takes when two requests share a prefix.
"""
fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
sched = fix.scheduler
num_blocks = 2
req = make_request(num_blocks=num_blocks)
kv_blocks = _alloc_and_register(fix, req, num_blocks)
sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
block_ids = kv_blocks.get_block_ids()
sched_out = make_scheduler_output(
{req.request_id: num_blocks * BLOCK_SIZE},
new_reqs={req.request_id: block_ids},
)
meta1 = sched.build_connector_meta(sched_out)
assert meta1.store_event >= 0
assert len(meta1.store_cpu_blocks) == num_blocks
# In-flight set tracks the scheduled GPU blocks until completion.
assert sched._in_flight_store_gpu_blocks == set(meta1.store_gpu_blocks)
cpu_free_after_first = get_cpu_free_blocks(sched)
# Second request shares the prefix and reuses the same GPU block IDs
# (the real scheduler path: GPU prefix cache returns the same blocks).
# Do NOT simulate completion — the first store is still in-flight.
req2 = Request(
request_id="req-dup-eager-inflight",
prompt_token_ids=req.prompt_token_ids,
sampling_params=req.sampling_params,
pooling_params=None,
mm_features=None,
block_hasher=req._block_hasher,
)
req2.num_computed_tokens = num_blocks * BLOCK_SIZE
sched.update_state_after_alloc(req2, kv_blocks, num_external_tokens=0)
sched_out2 = make_scheduler_output(
{req2.request_id: num_blocks * BLOCK_SIZE},
new_reqs={req2.request_id: block_ids},
)
meta2 = sched.build_connector_meta(sched_out2)
if meta2.store_event >= 0:
assert len(meta2.store_cpu_blocks) == 0, (
"Expected no new CPU blocks for in-flight duplicate hashes"
)
assert get_cpu_free_blocks(sched) == cpu_free_after_first, (
"Second request should not consume CPU blocks while the first "
"store is still in-flight"
)
# After completion, the in-flight set is cleared.
simulate_store_completion(sched, meta1.store_event)
assert sched._in_flight_store_gpu_blocks == set()
# ---------------------------------------------------------------------------
# Test 2c: Lazy duplicate store is skipped
# ---------------------------------------------------------------------------
def test_lazy_duplicate_store_skipped() -> None:
"""Lazy: blocks already offloaded to CPU should not be offloaded again.
@@ -1092,11 +1194,9 @@ def test_partial_gpu_prefix_plus_cpu_load() -> None:
hit_tokens, is_async = sched.get_num_new_matched_tokens(
req2, num_computed_tokens=gpu_local_computed
)
# CPU should hit blocks 2,3 (not 4,5 — those are beyond the CPU range).
num_cpu_hit_blocks = 2
# Actually CPU has all 6 stored; it returns hits starting from position 2.
# The number of CPU hit blocks = min(remaining request blocks, CPU cached).
# Here remaining = 6 - 2 = 4 blocks are in CPU, so hit = 4 * BLOCK_SIZE.
# CPU has all 6 blocks stored. make_request pads num_tokens by +1, so
# the manager's num_tokens - 1 cap leaves all full blocks intact:
# remaining hashable range = 6 - 2 = 4 blocks, all hit.
num_cpu_hit_blocks = 4
assert hit_tokens == num_cpu_hit_blocks * BLOCK_SIZE, (
f"Expected {num_cpu_hit_blocks * BLOCK_SIZE} CPU hit tokens, got {hit_tokens}"
@@ -6,8 +6,8 @@ import math
import pytest
import torch
from vllm.v1.worker.gpu.spec_decode.probabilistic_rejection_sampler_utils import (
probabilistic_rejection_sample,
from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
rejection_sample,
)
VOCAB_SIZE = 4096
@@ -167,7 +167,7 @@ def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: fl
num_trials=num_trials,
)
sampled, num_sampled = probabilistic_rejection_sample(
sampled, num_sampled = rejection_sample(
**inputs, num_speculative_steps=num_speculative_steps
)
@@ -201,7 +201,7 @@ def test_greedy_rejection_sample(num_speculative_steps: int):
num_trials=num_trials,
)
sampled, num_sampled = probabilistic_rejection_sample(
sampled, num_sampled = rejection_sample(
**inputs, num_speculative_steps=num_speculative_steps
)
@@ -213,3 +213,70 @@ def test_greedy_rejection_sample(num_speculative_steps: int):
assert (sampled[accepted_mask] == target_argmax).all(), (
"Greedy sampling produced tokens that are not the target argmax"
)
@pytest.mark.parametrize(
"num_speculative_steps,temperature,unconditional_rates",
[
(3, 1.0, [0.9, 0.5, 0.2]),
(3, 0.0, [0.9, 0.5, 0.2]),
(3, 1.0, [1.0, 1.0, 1.0]),
(3, 0.0, [1.0, 1.0, 1.0]),
(3, 1.0, [0.0, 0.0, 0.0]),
(3, 0.0, [0.0, 0.0, 0.0]),
(1, 1.0, [0.7]),
(1, 0.0, [0.7]),
],
)
def test_synthetic_rejection_sample(
num_speculative_steps: int,
temperature: float,
unconditional_rates: list[float],
):
"""
Verify that synthetic rejection sampling produces the expected
per-position acceptance rates. The unconditional rate at position i
is P(all draft steps 0..i accepted) = product(conditional_rates[0:i+1]).
This is approximately mean(num accepted >= i + 1) over many trials.
"""
from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
torch.manual_seed(42)
device = "cuda"
num_trials = 10 * VOCAB_SIZE
deviation_tol = 1e-2
target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
if temperature > 0:
target_logits_1d /= temperature
draft_logits_1d /= temperature
inputs = _build_rejection_sample_inputs(
target_logits_1d,
draft_logits_1d,
num_speculative_steps,
temperature=temperature,
num_trials=num_trials,
)
conditional_rates = unconditional_to_conditional_rates(unconditional_rates)
synthetic_conditional_rates = torch.tensor(
conditional_rates, dtype=torch.float32, device=device
)
_, num_sampled = rejection_sample(
**inputs,
num_speculative_steps=num_speculative_steps,
synthetic_conditional_rates=synthetic_conditional_rates,
)
# num_sampled includes the resampled/bonus token.
num_accepted = num_sampled - 1
for i, expected_rate in enumerate(unconditional_rates):
observed_rate = (num_accepted >= i + 1).float().mean().item()
assert abs(observed_rate - expected_rate) < deviation_tol, (
f"Step {i}: observed rate {observed_rate:.4f} deviates from "
f"expected rate {expected_rate:.4f} by more than {deviation_tol}."
)
@@ -1,5 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import dataclasses
import pytest
import torch
@@ -9,22 +11,55 @@ from vllm import LLM
from vllm.config import SpeculativeConfig
from vllm.distributed import cleanup_dist_env_and_memory
MODEL_PATH = "nm-testing/dflash-qwen3-8b-speculators"
EXPECTED_GSM8K_ACCURACY = 0.885
ACCURACY_RTOL = 0.03
EXPECTED_ACCEPTANCE_LEN = 3.45
ACCEPTANCE_LEN_RTOL = 0.15
# Expected per-position acceptance rates (accepted_at_pos / num_drafts)
# Based on GSM8K evaluation with Qwen3-8B dflash speculators.
EXPECTED_PER_POS_ACCEPTANCE_RATES = [0.795, 0.611, 0.429, 0.282]
PER_POS_RTOL = 0.15
@dataclasses.dataclass
class SpeculatorTestConfig:
model_path: str
method: str
display_name: str
expected_gsm8k_accuracy: float
accuracy_rtol: float
expected_acceptance_len: float
acceptance_len_rtol: float
expected_per_pos_acceptance_rates: tuple[float, ...]
per_pos_rtol: float
quantization: str | None = None
parallel_drafting: bool | None = None
def compute_spec_decode_stats(
metrics,
) -> dict:
DFLASH_CONFIG = SpeculatorTestConfig(
model_path="nm-testing/dflash-qwen3-8b-speculators",
method="dflash",
display_name="DFlash",
expected_gsm8k_accuracy=0.885,
accuracy_rtol=0.03,
expected_acceptance_len=3.45,
acceptance_len_rtol=0.15,
expected_per_pos_acceptance_rates=(0.795, 0.611, 0.429, 0.282),
per_pos_rtol=0.15,
quantization="fp8",
)
PEAGLE_CONFIG = SpeculatorTestConfig(
model_path="nm-testing/qwen3-8b-peagle-speculators",
method="eagle3",
display_name="PEagle",
expected_gsm8k_accuracy=0.88,
accuracy_rtol=0.05,
expected_acceptance_len=2.27,
acceptance_len_rtol=0.20,
expected_per_pos_acceptance_rates=(0.66, 0.36, 0.18, 0.09),
per_pos_rtol=0.20,
parallel_drafting=True,
)
SPECULATOR_CONFIGS = [
pytest.param(DFLASH_CONFIG, id="dflash"),
pytest.param(PEAGLE_CONFIG, id="peagle"),
]
def compute_spec_decode_stats(metrics) -> dict:
"""Extract all spec-decode metrics and compute derived stats."""
name2metric = {m.name: m for m in metrics}
@@ -67,25 +102,26 @@ def print_spec_decode_stats(stats: dict) -> None:
print("===============================\n")
def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
@pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
def test_speculators_model(vllm_runner, example_prompts, monkeypatch, config):
"""
Test DFlash speculators model properly initializes speculative decoding.
Test speculators model properly initializes speculative decoding.
Verifies:
1. Speculative config is automatically initialized from speculators config
2. Method is detected as 'dflash'
3. The draft model path is correctly set
4. Speculative tokens count is valid (num_speculative_tokens=8)
5. Text generation works with speculative decoding enabled
2. Method is detected correctly
3. parallel_drafting is set correctly (if applicable)
4. The draft model path is correctly set
5. Speculative tokens count is valid
6. Text generation works with speculative decoding enabled
"""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
with vllm_runner(
MODEL_PATH,
dtype=torch.bfloat16,
enforce_eager=True,
quantization="fp8",
) as vllm_model:
runner_kwargs = dict(dtype=torch.bfloat16, enforce_eager=True)
if config.quantization:
runner_kwargs["quantization"] = config.quantization
with vllm_runner(config.model_path, **runner_kwargs) as vllm_model:
vllm_config = vllm_model.llm.llm_engine.vllm_config
assert isinstance(vllm_config.speculative_config, SpeculativeConfig), (
@@ -93,40 +129,43 @@ def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
)
spec_config = vllm_config.speculative_config
assert spec_config.method == "dflash", (
f"Expected method='dflash', got '{spec_config.method}'"
assert spec_config.method == config.method, (
f"Expected method='{config.method}', got '{spec_config.method}'"
)
if config.parallel_drafting is not None:
assert spec_config.parallel_drafting is config.parallel_drafting, (
f"Expected parallel_drafting={config.parallel_drafting} "
f"for {config.display_name} model"
)
assert spec_config.num_speculative_tokens > 0, (
f"Expected positive speculative tokens, "
f"got {spec_config.num_speculative_tokens}"
)
assert spec_config.model == MODEL_PATH, (
f"Draft model should be {MODEL_PATH}, got {spec_config.model}"
assert spec_config.model == config.model_path, (
f"Draft model should be {config.model_path}, got {spec_config.model}"
)
vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens=20)
assert vllm_outputs, f"No outputs generated for speculators model {MODEL_PATH}"
assert vllm_outputs, (
f"No outputs generated for speculators model {config.model_path}"
)
@pytest.mark.slow_test
@large_gpu_mark(min_gb=40)
def test_dflash_speculators_correctness(monkeypatch):
@pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
def test_speculators_correctness(monkeypatch, config):
"""
E2E correctness test for DFlash via the speculators auto-detect path.
E2E correctness test via the speculators auto-detect path.
Evaluates GSM8k accuracy to ensure the speculators-format model produces
correct outputs, and checks that acceptance length does not collapse under
batched inference (lm-eval style).
Observed per-position acceptance rates on GSM8K (1319 prompts):
pos 0: 0.795, pos 1: 0.611, pos 2: 0.429, pos 3: 0.282,
pos 4: 0.169, pos 5: 0.093, pos 6: 0.048, pos 7: 0.023
Observed mean AL: 3.45 (GSM8K dataset, max_num_seqs=128)
"""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
spec_llm = LLM(
model=MODEL_PATH,
model=config.model_path,
trust_remote_code=True,
max_model_len=4096,
max_num_seqs=128,
@@ -137,7 +176,7 @@ def test_dflash_speculators_correctness(monkeypatch):
results = evaluate_gsm8k_offline(spec_llm)
accuracy = results["accuracy"]
accuracy_threshold = EXPECTED_GSM8K_ACCURACY * (1 - ACCURACY_RTOL)
accuracy_threshold = config.expected_gsm8k_accuracy * (1 - config.accuracy_rtol)
assert accuracy >= accuracy_threshold, (
f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}"
)
@@ -147,19 +186,18 @@ def test_dflash_speculators_correctness(monkeypatch):
print_spec_decode_stats(stats)
acceptance_len = stats["acceptance_len"]
al_threshold = EXPECTED_ACCEPTANCE_LEN * (1 - ACCEPTANCE_LEN_RTOL)
al_threshold = config.expected_acceptance_len * (1 - config.acceptance_len_rtol)
assert acceptance_len >= al_threshold, (
f"DFlash speculators acceptance length too low: "
f"{config.display_name} speculators acceptance length too low: "
f"{acceptance_len:.2f} < {al_threshold:.2f}"
)
# Check per-position acceptance rates for the first few positions.
per_pos_rates = stats["per_pos_acceptance_rates"]
for i, expected_rate in enumerate(EXPECTED_PER_POS_ACCEPTANCE_RATES):
for i, expected_rate in enumerate(config.expected_per_pos_acceptance_rates):
assert i < len(per_pos_rates), (
f"Missing per-position acceptance rate for position {i}"
)
threshold = expected_rate * (1 - PER_POS_RTOL)
threshold = expected_rate * (1 - config.per_pos_rtol)
assert per_pos_rates[i] >= threshold, (
f"Per-position acceptance rate at pos {i} too low: "
f"{per_pos_rates[i]:.4f} < {threshold:.4f} "
+85
View File
@@ -0,0 +1,85 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Build DeepGEMM's `_C` pybind11 extension for <TARGET_PY>.
Driven from cmake/external_projects/deepgemm.cmake. The driver runs against
the build interpreter's torch; <TARGET_PY> is only consulted for INCLUDEPY
and SOABI, so target venvs don't need torch installed.
Usage: python build_deepgemm_C.py <DEEPGEMM_SRC_DIR> <OUTPUT_DIR> <TARGET_PY>
"""
import json
import os
import subprocess
import sys
from pathlib import Path
import torch
from torch.utils import cpp_extension
if len(sys.argv) != 4:
sys.exit(f"usage: {sys.argv[0]} <SRC> <OUT> <TARGET_PY>")
src = Path(sys.argv[1]).resolve()
out = Path(sys.argv[2]).resolve()
target_py = sys.argv[3]
out.mkdir(parents=True, exist_ok=True)
info = json.loads(
subprocess.check_output(
[
target_py,
"-c",
"import sysconfig, json; "
"print(json.dumps({k: sysconfig.get_config_var(k) "
"for k in ('EXT_SUFFIX', 'INCLUDEPY')}))",
]
).decode()
)
cuda_home = cpp_extension.CUDA_HOME
if cuda_home is None:
sys.exit("CUDA_HOME not found; cannot build DeepGEMM _C")
# CCCL lives outside the standard CUDAToolkit search (mirrors DeepGEMM's setup.py).
includes = [
info["INCLUDEPY"],
f"{cuda_home}/include",
f"{cuda_home}/include/cccl",
str(src / "csrc"),
str(src / "deep_gemm/include"),
str(src / "third-party/cutlass/include"),
str(src / "third-party/cutlass/tools/util/include"),
str(src / "third-party/fmt/include"),
*cpp_extension.include_paths(device_type="cuda"),
]
cmd = [
os.environ.get("CXX", "g++"),
"-shared",
"-fPIC",
"-std=c++20",
"-O3",
"-g0",
"-Wno-psabi",
"-Wno-deprecated-declarations",
"-DTORCH_API_INCLUDE_EXTENSION_H",
"-DTORCH_EXTENSION_NAME=_C",
f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}",
*(f"-I{p}" for p in includes),
str(src / "csrc/python_api.cpp"),
*(f"-L{p}" for p in cpp_extension.library_paths(device_type="cuda")),
f"-L{cuda_home}/lib64",
"-ltorch",
"-ltorch_python",
"-ltorch_cpu",
"-ltorch_cuda",
"-lc10",
"-lc10_cuda",
"-lcudart",
"-lnvrtc",
"-o",
str(out / f"_C{info['EXT_SUFFIX']}"),
]
print("[build_deepgemm_C] " + " ".join(cmd), flush=True)
subprocess.check_call(cmd)
+41
View File
@@ -0,0 +1,41 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Assert the installed vLLM has a `_C.cpython-X.Y-*.so` for every CPython
covered by `requires-python`. Fails closed if a Python's `.so` is missing
from the wheel i.e. the regression that surfaced in #41476/#41512.
Run from a CI test job after vLLM is installed, e.g. the H100 deepgemm
kernel tests in .buildkite/test_areas/kernels.yaml.
"""
import importlib.util
import os
import sys
from pathlib import Path
import regex as re
import tomllib
SO_RE = re.compile(r"^_C\.cpython-(\d)(\d+)-")
def required_pythons() -> list[str]:
pyproject = Path(__file__).resolve().parent.parent / "pyproject.toml"
spec = tomllib.loads(pyproject.read_text())["project"]["requires-python"]
m = re.match(r">=3\.(\d+),<3\.(\d+)", spec)
if not m:
sys.exit(f"unexpected requires-python format: {spec!r}")
return [f"3.{v}" for v in range(int(m[1]), int(m[2]))]
spec = importlib.util.find_spec("vllm.third_party.deep_gemm")
if spec is None or spec.origin is None:
sys.exit("vllm.third_party.deep_gemm not importable; is vllm installed?")
pkg_dir = Path(spec.origin).parent
found = {f"{m[1]}.{m[2]}" for f in os.listdir(pkg_dir) if (m := SO_RE.match(f))}
required = required_pythons()
missing = [v for v in required if v not in found]
print(f"deepgemm _C: found {sorted(found)}, required {required}, missing {missing}")
sys.exit(1 if missing else 0)
+36
View File
@@ -0,0 +1,36 @@
#!/usr/bin/env bash
# Provision one bare Python per `requires-python` entry (or per argument) and
# print their paths as ":"-separated DEEPGEMM_PYTHON_INTERPRETERS. Skip this
# entirely if you already have interpreter paths.
#
# Usage:
# export DEEPGEMM_PYTHON_INTERPRETERS=$(tools/setup_deepgemm_pythons.sh)
# python setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
#
# Optional: DEEPGEMM_VENV_PREFIX (default: /tmp/dgenv).
set -euo pipefail
if [ "$#" -eq 0 ]; then
# Derive the matrix from `requires-python = ">=3.X,<3.Y"` in pyproject.toml.
pyproject="$(dirname "$0")/../pyproject.toml"
spec=$(grep -E '^requires-python' "$pyproject" \
| grep -oE '>=3\.[0-9]+,<3\.[0-9]+')
lo=${spec#>=3.}; lo=${lo%%,*}
hi=${spec##*<3.}
set -- $(seq "$lo" $((hi - 1)) | sed 's/^/3./')
fi
prefix="${DEEPGEMM_VENV_PREFIX:-/tmp/dgenv}"
mkdir -p "$prefix"
paths=""
for V in "$@"; do
venv="$prefix/$V"
# uv-managed Python ensures Python.h is present; system 3.X-dev packages
# on the manylinux / Ubuntu build bases are not always installed.
[ -x "$venv/bin/python" ] || \
uv venv --python "$V" "$venv" --python-preference only-managed --seed \
>/dev/null
paths="$paths:$venv/bin/python"
done
echo "${paths#:}"
+5 -3
View File
@@ -143,6 +143,11 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
if self.pass_config.fuse_gemm_comms:
self.passes += [AsyncTPPass(config)]
if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
# Run the more specific RMSNorm+router-pad fusion before
# AR+RMS, since both consume fused_add_rms_norm.
self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
if self.pass_config.fuse_allreduce_rms:
if rocm_aiter_ops.is_enabled():
self.passes += [RocmAiterAllReduceFusionPass(config)]
@@ -164,9 +169,6 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
if rocm_aiter_ops.is_enabled():
self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():
self.passes += [MLADualRMSNormFusionPass(config)]
+7 -2
View File
@@ -65,8 +65,13 @@ class DeviceConfig:
elif isinstance(self.device, torch.device):
self.device_type = self.device.type
# Some device types require processing inputs on CPU
if self.device_type in ["tpu"]:
# Some platforms require processing inputs on CPU.
from vllm.platforms import current_platform
if (
current_platform.uses_host_device_handling()
and self.device_type == current_platform.device_type
):
self.device = None
else:
# Set device with device type
+2 -2
View File
@@ -1345,7 +1345,7 @@ class ModelConfig:
)
raise AssertionError(f"Unsupported block type: {block_type}")
def get_mamba_chunk_size(self) -> int | None:
def get_mamba_chunk_size(self) -> int:
"""
Returns the mamba chunk size if it exists
"""
@@ -1356,7 +1356,7 @@ class ModelConfig:
chunk_size = getattr(self.hf_text_config, "chunk_size", None)
# Since Mamba1 does not have a chunk notion
# we use a default chunk size of 1024.
# we use a default chunk size of 2048.
if chunk_size is None:
chunk_size = 2048
@@ -335,6 +335,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):
_lock = threading.Lock()
def __init__(self, cpu_group, tcp_store_group=None):
assert tcp_store_group is not None
super().__init__(cpu_group, tcp_store_group)
self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS
@@ -270,7 +270,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
input_tensor = input_.movedim(0, dim).contiguous()
if sizes is not None:
assert len(sizes) == world_size
assert len(sizes) == world_size, f"{len(sizes)} == {world_size}"
assert input_tensor.shape[0] == sum(sizes)
chunk_size = sizes[self.rank_in_group]
else:

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