Compare commits

..
Author SHA1 Message Date
Andreas Karatzas d7cbffd5f8 [CI][AMD] Retry Hugging Face tests online
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-07-23 03:30:30 +00:00
jcotant-inferactandGitHub fc5fda105f [Docs] Re-add Reo.dev analytics beacon (#49474) 2026-07-23 03:03:32 +00:00
Matej SirovatkaandGitHub b07ec92faa [Bugfix] Make shared NVFP4 MoE scales writable (#49489)
Signed-off-by: S1ro1 <matej.sirovatka@gmail.com>
2026-07-22 19:17:53 -07:00
Colin ZGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
27ffbfde8d Fused Shared Expert Support for AMD Quark DeepSeek-V4 Model Checkpoints (#48044)
Signed-off-by: Colin Zeng <Colin.Zeng@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-23 00:34:17 +00:00
Nick HillandGitHub 229e01e9e1 [BugFix] Handle per-group prefix-hit divergence for hybrid models with KV connector (#48425) 2026-07-22 17:19:11 -07:00
191146dba5 Add quantization label automation (#49492)
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 19:59:15 -04:00
Summer YangandGitHub f3a920a076 [Core][DSV4] Compact MXFP4 indexer KV cache and packed group overlays (#48993) 2026-07-22 16:58:44 -07:00
Nils MattesonGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
149daf0d72 [Bugfix] Exclude location-derived path vars from torch.compile cache factors (#47573)
Signed-off-by: Nils Matteson <nilsmatteson@icloud.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-22 16:56:09 -07:00
Michael GoinandGitHub 917fdb5bf7 [Bugfix] Fix DeepGEMM warmup when using FlashInferFp8DeepGEMMDynamicBlockScaledKernel (#49467)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-07-22 16:28:49 -07:00
stefankoncarevicandGitHub 4b594b4aa1 [Bugfix][CI] Fix topk_softplus_sqrt no-op on non-XPU platforms (#49452)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
2026-07-22 15:36:17 -07:00
7d10a4cfce [Bugfix] Retry config read to survive concurrent HF cache refresh (#49001)
Signed-off-by: pei.zhang <pei.zhang@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-22 15:35:38 -07:00
Nick HillandGitHub 910cc8543a [Bugfix] Restore gather_and_maybe_dequant_cache OOB guard (#49427)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 13:47:37 -07:00
Nick HillandGitHub 431934522b [CI] Fix stale/fragile untethered kernels-root tests (#49423)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 14:43:07 -06:00
61a09532f2 Bump Flashinfer version to 0.6.15 (#48914)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Signed-off-by: Wei Zhao <weizha@oci-aga-slurm-1-vscode-02.cm.cluster>
Co-authored-by: Wei Zhao <weizha@oci-aga-slurm-1-vscode-02.cm.cluster>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-07-22 13:32:05 -07:00
Ben BrowningGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
3de4b2bf3c [Bugfix][Parser] Fix special tokens (EOS/BOS) leaking into reasoning content (#48748)
Signed-off-by: Ben Browning <56071+bbrowning@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-22 16:10:55 -04:00
b44311b6ef [CI] stabilize GDN prefill CuTeDSL test (#49388)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: Codex <noreply@openai.com>
2026-07-22 09:03:46 -07:00
Nick HillandGitHub b0d7875180 [CI] Increase timeout of pytorch-compilation-unit-tests (#49450)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 15:39:09 +00:00
Divakar VermaandGitHub 53c2f20dd9 [ROCm][CI] skip moe weight padding for eplb (#49350)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
2026-07-22 10:18:01 -05:00
Wentao YeandGitHub 37e370fe93 [DSv4 Perf] Skip empty c128 kernel launch, around 2x kernel performance improvement. (#48957)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-07-22 10:55:20 -04:00
Guan-Ming ChiuandGitHub 2dc5a72e7e [Bugfix][Renderer] Rebuild vision chunk UUIDs in async render path (#49400)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
2026-07-22 14:11:24 +00:00
Andrey TalmanandGitHub c79ff5f918 [Build] Bump vllm-flash-attn to C++20-compatible commit for torch-nightly (#49326)
Signed-off-by: Andrey Talman <atalman@fb.com>
2026-07-22 13:51:59 +00:00
Teresa ChenandGitHub 1a659a0c37 Upgrade tpu-inference to v0.25.0 (#49431) 2026-07-22 11:52:56 +00:00
SageandGitHub 0f6cf7f628 [Rust Frontend] Extract request preparation from the inference path (#49045)
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
2026-07-22 11:31:36 +00:00
c79ad3ae21 [Rust Frontend][gRPC] Add abort control RPC (#49255)
Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Connor Carpenter <connorc@nvidia.com>
2026-07-22 11:31:01 +00:00
wang.yuqiandGitHub 61c9ef986a [Frontend] Parallelize preprocessing within the same request for pooling models online serving. (#49153)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-07-22 10:56:23 +00:00
LiangqiusongandGitHub d6dbdb9b0d [XPU] WA of topk_softplus_sqrt arg mismatch on XPU (#49408)
Signed-off-by: xiaolong <xiaolong.guo@intel.com>
2026-07-22 16:16:13 +08:00
liuzhenweiandGitHub 06da482fb4 [XPU] WA of topk_softmax arg mismatch on XPU (#49395)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
2026-07-22 01:05:01 -07:00
2f75e7f712 [CI] Increase timeouts for jobs exceeding current limits (#49374)
Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 00:16:16 -07:00
Ziming HuangandGitHub 7c21548ce3 [PD][Bugfix] Fix NIXL hybrid MLA+mamba heterogeneous TP (#49297)
Signed-off-by: ZeldaHuang <zelda.huanghuang@gmail.com>
2026-07-22 07:11:39 +00:00
Guan-Ming ChiuandGitHub 9df2f91232 [Renderer] Offload derender CPU work to renderer thread pool (#49396)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
2026-07-22 06:55:03 +00:00
387189c429 [ROCm] Remove redundant AITER fused_qk_rmsnorm probe (avoids config-time HIP init) (#47992)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-21 22:58:14 -05:00
Kunshang JiandGitHub 75576c63be Add auto label for xpu relate issue (#49398)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-22 03:33:20 +00:00
Andreas KaratzasandGitHub 16aca639b7 [ROCm] Upgrade NIXL and UCX (#49251)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
2026-07-21 22:30:02 -05:00
Woosuk KwonandGitHub 6049424b7e [MRV2] Always build attn metadata at capture time (#49364)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-07-21 19:27:16 -07:00
060b5f61dc [Bugfix][Attention] Ignore empty MLA context chunks during merge (#49294)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 02:23:48 +00:00
ec59c1579f [MoE Refactor] Migrate MoeWNA16Method quantization method over to using the new MK oracle scheme. (#44120)
Signed-off-by: Bill Nell <bnell@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: mgoin <mgoin64@gmail.com>
2026-07-21 19:20:30 -07:00
Isotr0pyandGitHub 1750e443f2 [Misc] Move PyNvVideoCodec stuff out of gpu worker (#49322)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
2026-07-22 10:11:56 +08:00
ba18929079 [Bugfix][SpecDecode] Scope MTP completeness checks outside bucketed updates (#49178)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
2026-07-21 18:56:39 -07:00
0500ca6a58 [CI][Bugfix] Fix ROCm FP8 KV cache dtype in attention backend test (#49380)
Signed-off-by: pei.zhang <pei.zhang@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-22 01:34:49 +00:00
a1c15bcb0f [CI][Bugfix] Fix and wire streaming-input tests (#49356)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-22 00:35:23 +00:00
Nick HillandGitHub 4809de7317 [Misc] Fix terminal output logo coloring (#49344)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 00:21:29 +00:00
stefankoncarevicandGitHub 05781e21dd [ROCm][CI] Fix order-dependent failure in test_flash_attn_accepts_handled_fp8_variants (MI355) (#49329)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
2026-07-21 18:54:33 -05:00
gnovackandGitHub 85f638a2b8 skip cudagraph/DP padding in topk (#48979)
Signed-off-by: gnovack <novackgm@gmail.com>
2026-07-21 15:20:30 -07:00
Nick HillandGitHub 08e5067561 [CI] Bump timeout of entrypoints-integration-api-server-openai-part-2 (#49359)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-21 14:54:03 -07:00
a7d00ec051 [Bugfix] DFlash fc sized wrong when num_target_layers != num_hidden_layers (#48524)
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
2026-07-21 14:42:52 -07:00
Michael GoinandGitHub b8fb56d970 [CI] Add gemma-4-E4B-it-assistant to CI gsm8k for GemmaMTP (#49243)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-07-21 16:51:04 -04:00
96a739289e [Bugfix] fix cutalss version upgrade bug, need update MSG new commit (#49016)
Signed-off-by: lengrongfu <lenronfu@gmail.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
2026-07-21 13:42:14 -07:00
60d443f738 [CI/Build][The Rock][BugFix] Use fork method in test_multiproc_executor_multi_node for py 3.14 compat and fix test_multiproc_executor_shutdown_cleanup (#48655)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-21 15:40:29 -05:00
1dca300653 [CI] Fix and wire encoder/manager cudagraph unit tests (#49339)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 20:41:16 +01:00
Flora FengandGitHub fca252d59e [CI][Bugfix] Reduce max_model_len in OOT embedding test to fix KV-cache OOM on small GPUs (#49351)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-07-21 15:00:06 -04:00
33178f9006 Fix Qwen3-VL M-RoPE on the Transformers modeling backend (grids + compile) (#49292)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-21 18:28:39 +00:00
b2b8f679d0 [Bugfix][Spec Decode] Restrict embedding-width share guard to EAGLE drafts (#47953)
Signed-off-by: Evan Takahashi <evan10takahashi@gmail.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
2026-07-21 10:53:37 -07:00
de6ec294ef [Bugfix] Fix DSA crash under breakable piecewise cudagraphs (#49302)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 18:28:13 +01:00
stefankoncarevicandGitHub 61e10f0116 [ROCm][CI] Fix AITER MLA fp8 decode metadata regression test (#48845)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
2026-07-21 12:19:45 -05:00
6e96891ba0 [ROCm] Bump AITER to v0.1.16.post5 (#48683)
Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com>
Signed-off-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-21 10:39:11 -05:00
47f1b47a73 Ci/add laguna xs gsm8k (#49241)
Signed-off-by: Joe Rowell <joerowell4@gmail.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Joe Rowell <joerowell4@gmail.com>
2026-07-21 11:30:42 -04:00
5aab491bc9 [CI] Wire tests/models/inkling into a B200 job (#49325)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 15:43:20 +01:00
5812e1a66b [Test] Add DeepSeek MTP parallel-load tests (#41653)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
2026-07-21 10:38:20 -04:00
8950394e0a [Bugfix] Prefix-cache metrics double-counted when a KV connector defers requests (#48860)
Signed-off-by: Seiji Eicher <58963096+eicherseiji@users.noreply.github.com>
Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Lixi Zhang <muxixibbb@gmail.com>
Co-authored-by: harjothkhara <harjoth.khara@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-07-21 15:02:03 +01:00
Roberto L. CastroandGitHub 7bb49be4d1 [Bugfix] Handle MLA fallback during FA4 JIT warmup (#49306) 2026-07-21 13:58:05 +00:00
c67650f04b [XPU][DeepSeekV4]Add DeepSeek-V4 fuse_index_q SYCL kernel path (#45991)
Signed-off-by: Avery Miao <avery.miao@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-21 21:22:34 +08:00
f890e1dbe2 [BugFix] Set graph_pool_id before FULL CUDA graph capture in ModelRunner V2 (#48843)
Signed-off-by: Markov Ilya <markovilya19@gmail.com>
Co-authored-by: Markov Ilya <markovilya19@gmail.com>
2026-07-21 08:59:26 -04:00
181 changed files with 5057 additions and 2364 deletions
+28 -28
View File
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_METADATA_VERSION="1"
IMAGE_EXISTED_BEFORE_BUILD=0
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
return 0
fi
for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
if [[ -n "${!var:-}" ]]; then
echo "Using provided ${var}: ${!var}"
continue
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
compute_dependency_cache_keys() {
local bake_dir=""
local dockerfile_rocm=""
local rixl_branch=""
local nixl_branch=""
local ucx_branch=""
local rocshmem_branch=""
local deepep_branch=""
local rixl_material=""
local nixl_material=""
local rocshmem_material=""
local deepep_material=""
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
RIXL_CACHE_KEY=$(
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
NIXL_CACHE_KEY=$(
compose_dependency_cache_key \
"${rixl_branch}-ucx-${ucx_branch}" \
"${rixl_material}"
"${nixl_branch}-ucx-${ucx_branch}" \
"${nixl_material}"
)
export RIXL_CACHE_KEY
echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
export NIXL_CACHE_KEY
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
fi
if [[ -n "${rocshmem_branch}" ]]; then
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
case "${target}" in
rixl-rocm-ci)
if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
elif [[ -n "${RIXL_BRANCH:-}" ]]; then
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
nixl-rocm-ci)
if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
elif [[ -n "${NIXL_BRANCH:-}" ]]; then
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
fi
;;
rocshmem-rocm-ci)
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
resolve_ci_base_dependency_targets() {
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
local rixl_ref=""
local nixl_ref=""
local rocshmem_ref=""
local deepep_ref=""
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
case "${mode}" in
always)
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
add_dependency_cache_target "${target}"
fi
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
;;
esac
if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
if dependency_cache_ref_exists "${rixl_ref}"; then
echo "RIXL dependency cache exists: ${rixl_ref}"
if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
if dependency_cache_ref_exists "${nixl_ref}"; then
echo "NIXL dependency cache exists: ${nixl_ref}"
else
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
add_dependency_cache_target "rixl-rocm-ci"
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
add_dependency_cache_target "nixl-rocm-ci"
fi
fi
+58 -2
View File
@@ -56,6 +56,15 @@ fi
export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS PYTEST_TIMEOUT ROCM_DOCKER_TTY
export PYTHONFAULTHANDLER
# The AMD pipeline template sets this for eligible pilot jobs. Capture it
# before clear_ci_orchestration_env removes CI-only controls from test processes.
hf_offline_retry_enabled="${VLLM_CI_HF_OFFLINE_RETRY:-0}"
if [[ "${hf_offline_retry_enabled}" != "0" && "${hf_offline_retry_enabled}" != "1" ]]; then
echo "VLLM_CI_HF_OFFLINE_RETRY must be 0 or 1" >&2
exit 2
fi
hf_retry_command_file=""
# Export Python path for commands that run directly on the host. Containerized
# tests set this to /vllm-workspace below so spawned Python processes do not
# depend on their current working directory.
@@ -73,6 +82,7 @@ report_docker_usage() {
clear_ci_orchestration_env() {
unset -v \
VLLM_TEST_GROUP_NAME \
VLLM_CI_HF_OFFLINE_RETRY \
VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE \
VLLM_CI_ARTIFACT_STEP \
VLLM_TEST_CACHE \
@@ -92,6 +102,29 @@ clear_ci_orchestration_env() {
VLLM_ALLOW_DEPRECATED_BEAM_SEARCH
}
prepare_hf_retry_command_file() {
local command_text=$1
if ! hf_retry_command_file=$(mktemp "${TMPDIR:-/tmp}/vllm-hf-offline-command.XXXXXX"); then
echo "Failed to create the Hugging Face retry command file" >&2
return 1
fi
if ! printf '%s\n' "${command_text}" >"${hf_retry_command_file}"; then
echo "Failed to write the Hugging Face retry command file" >&2
rm -f -- "${hf_retry_command_file}"
hf_retry_command_file=""
return 1
fi
}
# shellcheck disable=SC2329 # Called by cleanup functions registered as traps.
cleanup_hf_retry_command_file() {
if [[ -n "${hf_retry_command_file}" ]]; then
rm -f -- "${hf_retry_command_file}"
hf_retry_command_file=""
fi
}
cleanup_network() {
local max_nodes=${NUM_NODES:-2}
for node in $(seq 0 $((max_nodes - 1))); do
@@ -649,6 +682,7 @@ if is_native_runtime; then
artifact_work_dir=""
cleanup_native_workspace() {
cleanup_hf_retry_command_file
if [[ -n "${artifact_work_dir}" ]]; then
rm -rf "${artifact_work_dir}"
fi
@@ -691,7 +725,13 @@ if is_native_runtime; then
run_native_preflight || exit 1
# Keep AMD CI orchestration variables out of vLLM's runtime environment.
clear_ci_orchestration_env
/bin/bash -o pipefail -c "${commands}"
if [[ "${hf_offline_retry_enabled}" == "1" ]]; then
prepare_hf_retry_command_file "${commands}" || exit 1
bash "${VLLM_CI_WORKSPACE:-/vllm-workspace}/.buildkite/scripts/hf-offline-retry.sh" \
"${hf_retry_command_file}"
else
/bin/bash -o pipefail -c "${commands}"
fi
handle_pytest_exit "$?"
fi
@@ -709,6 +749,7 @@ artifact_work_dir=""
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
remove_docker_container() {
cleanup_hf_retry_command_file
if docker container inspect "${container_name}" >/dev/null 2>&1; then
docker rm -f "${container_name}" || true
fi
@@ -850,6 +891,10 @@ fi
# --- Route: multi-node vs single-node ---
clear_ci_orchestration_env
if is_multi_node "$commands"; then
if [[ "${hf_offline_retry_enabled}" == "1" ]]; then
echo "Hugging Face offline retry is not enabled for AMD multi-node jobs yet" >&2
exit 2
fi
echo "--- Multi-node job detected"
export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
@@ -894,6 +939,16 @@ if is_multi_node "$commands"; then
fi
else
echo "--- Single-node job"
hf_retry_container_args=()
container_test_command="${commands}"
if [[ "${hf_offline_retry_enabled}" == "1" ]]; then
prepare_hf_retry_command_file "${commands}" || exit 1
hf_retry_container_path="/tmp/vllm-hf-offline-command.sh"
hf_retry_container_args=(
-v "${hf_retry_command_file}:${hf_retry_container_path}:ro"
)
container_test_command="bash /vllm-workspace/.buildkite/scripts/hf-offline-retry.sh ${hf_retry_container_path}"
fi
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
docker_run_terminal_args=(-i)
if [[ "${ROCM_DOCKER_TTY}" == "1" ]]; then
@@ -950,9 +1005,10 @@ else
-e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
-e "PYTORCH_ROCM_ARCH=" \
"${standalone_merge_base_env[@]}" \
"${hf_retry_container_args[@]}" \
--name "${container_name}" \
"${image_name}" \
/bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
/bin/bash -c "${CONTAINER_PREFLIGHT} && ${container_test_command}"
exit_code=$?
handle_pytest_exit "$exit_code"
@@ -0,0 +1,122 @@
#!/usr/bin/env bash
set -euo pipefail
script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
helper="${script_dir}/hf-offline-retry.sh"
test_root=$(mktemp -d "${TMPDIR:-/tmp}/vllm-hf-offline-retry-test.XXXXXX")
cleanup() {
rm -rf -- "${test_root}"
}
trap cleanup EXIT
assert_equal() {
local expected=$1
local actual=$2
local description=$3
if [[ "${actual}" != "${expected}" ]]; then
printf 'FAIL: %s\n expected: %q\n actual: %q\n' \
"${description}" "${expected}" "${actual}" >&2
exit 1
fi
}
status_command_file="${test_root}/status-command.sh"
cat >"${status_command_file}" <<'VLLM_TEST_COMMAND'
if [[ "${HF_HUB_OFFLINE-}" == "1" &&
"${TRANSFORMERS_OFFLINE-}" == "1" &&
"${HF_DATASETS_OFFLINE-}" == "1" ]]; then
printf '%s\n' offline >>"${ATTEMPTS_FILE}"
exit "${OFFLINE_STATUS}"
fi
if [[ -z "${HF_HUB_OFFLINE+x}" &&
-z "${TRANSFORMERS_OFFLINE+x}" &&
-z "${HF_DATASETS_OFFLINE+x}" ]]; then
printf '%s\n' online >>"${ATTEMPTS_FILE}"
exit "${ONLINE_STATUS}"
fi
exit 91
VLLM_TEST_COMMAND
run_status_case() {
local name=$1
local offline_status=$2
local online_status=$3
local expected_status=$4
local expected_attempts=$5
local attempts_file="${test_root}/${name}.attempts"
local actual_status
local actual_attempts
set +e
ATTEMPTS_FILE="${attempts_file}" \
OFFLINE_STATUS="${offline_status}" \
ONLINE_STATUS="${online_status}" \
bash "${helper}" "${status_command_file}"
actual_status=$?
set -e
actual_attempts=$(<"${attempts_file}")
assert_equal "${expected_status}" "${actual_status}" "${name} exit status"
assert_equal "${expected_attempts}" "${actual_attempts}" "${name} attempts"
}
run_status_case offline_success 0 99 0 "offline"
run_status_case retry_status_1 1 0 0 $'offline\nonline'
run_status_case retry_status_2 2 0 0 $'offline\nonline'
run_status_case retry_status_123 123 0 0 $'offline\nonline'
run_status_case no_retry_status_42 42 0 42 "offline"
run_status_case online_failure 1 7 7 $'offline\nonline'
transport_file="${test_root}/transport.out"
transport_command_file="${test_root}/transport-command.sh"
cat >"${transport_command_file}" <<'VLLM_TRANSPORT_COMMAND'
cat >"${OUTPUT_FILE}" <<'INNER_PAYLOAD'
literal:$HOME
quotes:'single' "double"
two command-file lines
INNER_PAYLOAD
printf 'expanded:%s\n' "${TRANSPORT_VALUE}" >>"${OUTPUT_FILE}"
VLLM_TRANSPORT_COMMAND
OUTPUT_FILE="${transport_file}" \
TRANSPORT_VALUE='value with spaces, $dollars, and "quotes"' \
bash "${helper}" "${transport_command_file}"
transport_output=$(<"${transport_file}")
assert_equal \
$'literal:$HOME\nquotes:\'single\' "double"\ntwo command-file lines\nexpanded:value with spaces, $dollars, and "quotes"' \
"${transport_output}" \
"quotes, dollars, and newlines"
stdin_file="${test_root}/stdin.out"
stdin_payload="stdin payload with \$dollars and \"quotes\""
stdin_command_file="${test_root}/stdin-command.sh"
cat >"${stdin_command_file}" <<'VLLM_STDIN_COMMAND'
IFS= read -r input
printf '%s\n' "${input}" >"${OUTPUT_FILE}"
VLLM_STDIN_COMMAND
printf '%s\n' "${stdin_payload}" |
OUTPUT_FILE="${stdin_file}" bash "${helper}" "${stdin_command_file}"
stdin_output=$(<"${stdin_file}")
assert_equal \
"${stdin_payload}" \
"${stdin_output}" \
"stdin preservation"
set +e
bash "${helper}" >/dev/null 2>&1
usage_status=$?
bash "${helper}" "${test_root}/missing-command" >/dev/null 2>&1
missing_status=$?
set -e
assert_equal "64" "${usage_status}" "missing argument exit status"
assert_equal "66" "${missing_status}" "unreadable command file exit status"
echo "PASS: hf-offline-retry"
+48
View File
@@ -0,0 +1,48 @@
#!/usr/bin/env bash
set -uo pipefail
if [[ $# -ne 1 ]]; then
echo "Usage: $0 COMMAND_FILE" >&2
exit 64
fi
command_file=$1
if [[ ! -f "${command_file}" || ! -r "${command_file}" ]]; then
echo "Command file is not readable: ${command_file}" >&2
exit 66
fi
run_attempt() (
if [[ "$1" == "offline" ]]; then
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export HF_DATASETS_OFFLINE=1
else
unset HF_HUB_OFFLINE
unset TRANSFORMERS_OFFLINE
unset HF_DATASETS_OFFLINE
fi
bash -e -o pipefail -- "${command_file}"
)
echo "--- :package: Hugging Face offline attempt"
run_attempt offline
offline_status=$?
if [[ ${offline_status} -eq 0 ]]; then
exit 0
fi
case "${offline_status}" in
1 | 2 | 123)
echo "--- :globe_with_meridians: Hugging Face offline attempt failed; retrying online"
run_attempt online
online_status=$?
exit "${online_status}"
;;
*)
exit "${offline_status}"
;;
esac
@@ -29,7 +29,6 @@ PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
CARGO_LLVM_COV_VERSION="${CARGO_LLVM_COV_VERSION:-0.8.7}"
log_section() {
echo "--- $*"
@@ -107,18 +106,6 @@ install_cargo_nextest() {
"cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
}
install_cargo_llvm_cov() {
log_section "Installing cargo-llvm-cov ${CARGO_LLVM_COV_VERSION}"
local toolchain
toolchain="$(rust_toolchain)"
rustup component add --toolchain "$toolchain" llvm-tools-preview
cargo binstall \
--no-confirm \
--force \
--secure \
"cargo-llvm-cov@${CARGO_LLVM_COV_VERSION}"
}
install_uv() {
log_section "Installing uv ${UV_VERSION}"
curl -L --proto '=https' --tlsv1.2 -sSf \
@@ -189,41 +176,14 @@ run_tests() {
setup_pyo3_python
install_cargo_binstall
install_cargo_nextest
install_cargo_llvm_cov
log_section "Running cargo nextest with Rust coverage"
mkdir -p artifacts
export LLVM_PROFILE_FILE_NAME="vllm-rust-unit-%4m.profraw"
cargo llvm-cov clean \
--manifest-path rust/Cargo.toml \
--profraw-only
set +e
cargo llvm-cov nextest \
log_section "Running cargo nextest"
cargo nextest run \
--manifest-path rust/Cargo.toml \
--workspace \
--all-features \
--locked \
--no-fail-fast \
--no-clean \
--lcov \
--output-path artifacts/rust-unit.lcov \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/'
local coverage_rc=$?
local upload_rc=0
if [[ $coverage_rc -eq 0 ]]; then
# shellcheck source=.buildkite/scripts/rust-coverage.sh
source .buildkite/scripts/rust-coverage.sh
rust_coverage_upload artifacts/rust-unit.lcov rust-unit
upload_rc=$?
fi
set -e
if [[ $coverage_rc -ne 0 ]]; then
return "$coverage_rc"
fi
return "$upload_rc"
--no-fail-fast
}
install_protoc
-182
View File
@@ -1,182 +0,0 @@
#!/bin/sh
RUST_CODECOV_VERSION="v11.3.1"
RUST_CODECOV_SHA256="ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d"
rust_coverage_repo_root() {
if [ -f /vllm-workspace/.buildkite/scripts/rust-coverage.sh ]; then
printf '%s\n' /vllm-workspace
elif [ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ] \
&& [ -d "$BUILDKITE_BUILD_CHECKOUT_PATH" ]; then
printf '%s\n' "$BUILDKITE_BUILD_CHECKOUT_PATH"
else
git rev-parse --show-toplevel
fi
}
rust_coverage_start() {
RUST_COVERAGE_FLAG=${1:?coverage flag is required}
RUST_COVERAGE_DIR="/tmp/vllm-rust-coverage/${BUILDKITE_JOB_ID:-local}"
export RUST_COVERAGE_FLAG RUST_COVERAGE_DIR
mkdir -p "$RUST_COVERAGE_DIR"
LLVM_PROFILE_FILE="$RUST_COVERAGE_DIR/rust-%4m.profraw"
export LLVM_PROFILE_FILE
trap rust_coverage_finalize 0
}
rust_coverage_objects() {
rust_cov_objects_manifest="$(dirname "$(command -v llvm-cov)")/../objects"
python3 - "$rust_cov_objects_manifest" <<'PY'
from pathlib import Path
import sys
for relative in Path(sys.argv[1]).read_text().splitlines():
for entry in sys.path:
path = Path(entry or ".").resolve() / relative
if path.is_file():
print(path)
break
else:
raise RuntimeError(f"installed Rust coverage object was not found: {relative}")
PY
}
rust_coverage_collect() {
rust_cov_collect_flag=${1:?coverage flag is required}
rust_cov_collect_lcov="$RUST_COVERAGE_DIR/$rust_cov_collect_flag.lcov"
rust_cov_collect_objects=$(rust_coverage_objects) || return 1
rust_cov_collect_primary=
set --
while IFS= read -r rust_cov_collect_object; do
if [ -z "$rust_cov_collect_primary" ]; then
rust_cov_collect_primary=$rust_cov_collect_object
else
set -- "$@" "--object=$rust_cov_collect_object"
fi
done <<EOF
$rust_cov_collect_objects
EOF
llvm-profdata merge \
-sparse \
"$RUST_COVERAGE_DIR"/*.profraw \
-o "$RUST_COVERAGE_DIR/merged.profdata" || return 1
llvm-cov export \
"$rust_cov_collect_primary" \
"$@" \
--format=lcov \
--instr-profile="$RUST_COVERAGE_DIR/merged.profdata" \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/' \
> "$rust_cov_collect_lcov" || return 1
RUST_COVERAGE_LCOV=$rust_cov_collect_lcov
export RUST_COVERAGE_LCOV
}
rust_coverage_upload() {
rust_cov_upload_lcov=${1:?LCOV path is required}
rust_cov_upload_flag=${2:?coverage flag is required}
rust_cov_upload_repo_root=$(rust_coverage_repo_root) || return 1
if [ "$(uname -m)" != "x86_64" ]; then
echo "Rust coverage upload currently supports x86_64 CI agents" >&2
return 1
fi
rust_cov_upload_codecov_dir=$(mktemp -d /tmp/codecov-bin.XXXXXX) \
|| return 1
curl -fsSL \
"https://github.com/codecov/codecov-cli/releases/download/${RUST_CODECOV_VERSION}/codecovcli_linux" \
-o "$rust_cov_upload_codecov_dir/codecov" || return 1
echo "$RUST_CODECOV_SHA256 $rust_cov_upload_codecov_dir/codecov" \
| sha256sum -c - || return 1
chmod +x "$rust_cov_upload_codecov_dir/codecov" || return 1
rust_cov_upload_slug="vllm-project/vllm"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ] \
&& [ -n "${BUILDKITE_PULL_REQUEST_REPO:-}" ]; then
rust_cov_upload_slug=$(echo "$BUILDKITE_PULL_REQUEST_REPO" \
| sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
case "$rust_cov_upload_slug" in
*/*) ;;
*) rust_cov_upload_slug="vllm-project/vllm" ;;
esac
fi
rust_cov_upload_branch=${BUILDKITE_BRANCH:?BUILDKITE_BRANCH is required}
if [ -z "${CODECOV_TOKEN:-}" ]; then
# Codecov accepts tokenless public uploads on unprotected branch names.
# A colon-separated prefix keeps feature-branch and fork uploads from
# requiring a repository secret.
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
rust_cov_upload_branch="pr${BUILDKITE_PULL_REQUEST}:$rust_cov_upload_branch"
else
rust_cov_upload_branch="buildkite:$rust_cov_upload_branch"
fi
fi
set --
set -- "$@" upload-process
set -- "$@" --file "$rust_cov_upload_lcov"
# LCOV paths are mapped server-side by codecov.yml. Skip the CLI's local
# source-line fix scanning, which is unrelated to path mapping.
set -- "$@" --disable-search --disable-file-fixes
set -- "$@" --fail-on-error --git-service github
set -- "$@" --build "${BUILDKITE_BUILD_NUMBER:?BUILDKITE_BUILD_NUMBER is required}"
set -- "$@" --branch "$rust_cov_upload_branch"
set -- "$@" --sha "${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}"
set -- "$@" --slug "$rust_cov_upload_slug"
set -- "$@" --flag "$rust_cov_upload_flag"
set -- "$@" --name "${rust_cov_upload_flag}-${BUILDKITE_JOB_ID:?BUILDKITE_JOB_ID is required}"
set -- "$@" --dir "$rust_cov_upload_repo_root"
set -- "$@" --network-root-folder "$rust_cov_upload_repo_root"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
set -- "$@" --pr "$BUILDKITE_PULL_REQUEST"
fi
rust_cov_upload_log="$rust_cov_upload_codecov_dir/codecov.log"
# E2E steps run from tests/, so execute from the repository root to resolve
# codecov.yml and repository paths consistently.
(
cd "$rust_cov_upload_repo_root" || exit 1
"$rust_cov_upload_codecov_dir/codecov" "$@"
) >"$rust_cov_upload_log" 2>&1
rust_cov_upload_rc=$?
cat "$rust_cov_upload_log"
# v11.3.1 can log API failures while returning zero even with
# --fail-on-error. Preserve the strict CI contract explicitly.
if grep -aEq 'error.* -- ' "$rust_cov_upload_log"; then
echo "Codecov CLI reported an upload error" >&2
rust_cov_upload_rc=1
fi
rm -rf "$rust_cov_upload_codecov_dir"
return "$rust_cov_upload_rc"
}
rust_coverage_finalize() {
rust_cov_finalize_test_rc=$?
trap - 0
set +e
rust_coverage_collect "$RUST_COVERAGE_FLAG"
rust_cov_finalize_collect_rc=$?
rust_cov_finalize_upload_rc=0
if [ "$rust_cov_finalize_collect_rc" -eq 0 ]; then
rust_coverage_upload "$RUST_COVERAGE_LCOV" "$RUST_COVERAGE_FLAG"
rust_cov_finalize_upload_rc=$?
fi
find "$RUST_COVERAGE_DIR" -type f -name '*.profraw' -delete
if [ "$rust_cov_finalize_test_rc" -ne 0 ]; then
exit "$rust_cov_finalize_test_rc"
fi
if [ "$rust_cov_finalize_collect_rc" -ne 0 ]; then
exit "$rust_cov_finalize_collect_rc"
fi
exit "$rust_cov_finalize_upload_rc"
}
+9
View File
@@ -25,6 +25,9 @@
# and $$BUILDKITE_PARALLEL_JOB_COUNT environment variables.
# working_dir(str): specify the place where the command should execute, default to /vllm-workspace/tests
# source_file_dependencies(list): the list of prefixes to opt-in the test for, if empty, the test will always run.
# hf_offline_retry(bool): for single-node run-amd-test jobs, run offline first
# and retry eligible failures online. May be set as the file default or a
# per-step override.
# When adding a test
# - If the test belongs to an existing group, add it there
@@ -105,6 +108,8 @@
#####################################################################################################################################
# AMD-only pilot: test jobs inherit this; non-test jobs opt out below.
hf_offline_retry: true
steps:
#########################################################################################################################################
@@ -411,6 +416,7 @@ steps:
#------------------------------------------------------------ mi250 · rust -----------------------------------------------------------#
- label: Rust Frontend Cargo Style + Clippy # TBD
hf_offline_retry: false
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
@@ -425,6 +431,7 @@ steps:
- bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
- label: Rust Frontend Cargo Tests # TBD
hf_offline_retry: false
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
@@ -441,6 +448,7 @@ steps:
#----------------------------------------------------------- mi250 · docker ----------------------------------------------------------#
- label: Docker Build Metadata (ROCm) # TBD
hf_offline_retry: false
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
@@ -2264,6 +2272,7 @@ steps:
#--------------------------------------------------------- mi300 · ray_compat ---------------------------------------------------------#
- label: Ray Dependency Compatibility Check # TBD
hf_offline_retry: false
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Basic Correctness
key: basic-correctness
timeout_in_minutes: 45
timeout_in_minutes: 68
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Benchmarks CLI Test
key: benchmarks-cli-test
timeout_in_minutes: 30
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
+3
View File
@@ -26,7 +26,10 @@ steps:
- vllm/v1/cudagraph_dispatcher.py
- vllm/config/compilation.py
- vllm/compilation
- vllm/v1/worker/encoder_cudagraph.py
- vllm/v1/worker/encoder_cudagraph_defs.py
commands:
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
- pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
+1 -1
View File
@@ -51,7 +51,7 @@ steps:
- label: e2e Scheduling (1 GPU)
key: e2e-scheduling-1-gpu
timeout_in_minutes: 35
timeout_in_minutes: 53
device: h200_18gb
source_file_dependencies:
- vllm/v1/
+4 -4
View File
@@ -39,7 +39,7 @@ steps:
- label: Entrypoints Integration (API Server)
key: entrypoints-integration-api-server
device: h200_35gb
timeout_in_minutes: 50
timeout_in_minutes: 75
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -59,7 +59,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 1)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 45
timeout_in_minutes: 68
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -78,7 +78,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 2)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 45
timeout_in_minutes: 83
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -156,7 +156,7 @@ steps:
- label: Entrypoints Integration (Pooling)
device: h200_35gb
key: entrypoints-integration-pooling
timeout_in_minutes: 50
timeout_in_minutes: 75
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
+4 -1
View File
@@ -31,7 +31,7 @@ steps:
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 45
timeout_in_minutes: 83
device: h200_18gb
source_file_dependencies:
- vllm/config/
@@ -90,6 +90,7 @@ steps:
- tests/v1/kv_offload
- tests/v1/simple_kv_offload
- tests/v1/worker
- tests/v1/streaming_input
- tests/v1/kv_connector/unit
- tests/v1/ec_connector/unit
- tests/v1/metrics
@@ -103,6 +104,7 @@ steps:
- pytest -v -s v1/kv_offload
- pytest -v -s v1/simple_kv_offload
- pytest -v -s v1/worker
- pytest -v -s v1/streaming_input
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
@@ -143,6 +145,7 @@ steps:
- pytest -v -s -m 'cpu_test' v1/core
- pytest -v -s v1/structured_output
- pytest -v -s v1/test_serial_utils.py
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'cpu_test' v1/metrics
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: Model Executor
device: h200_35gb
key: model-executor
timeout_in_minutes: 45
timeout_in_minutes: 60
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
+13
View File
@@ -46,6 +46,19 @@ steps:
depends_on:
- image-build-amd
- label: Inkling Unit Tests (B200)
key: inkling-unit-tests-b200
timeout_in_minutes: 40
device: b200-k8s
source_file_dependencies:
- vllm/models/inkling/
- vllm/cute_utils/
- cmake/external_projects/tml_fa4.cmake
- tests/models/inkling/
commands:
# FA4 kernel tests require SM100; the suite skips them elsewhere.
- pytest -v -s models/inkling
- label: Basic Models Test (Other CPU) # 5min
key: basic-models-test-other-cpu
depends_on:
+1 -1
View File
@@ -137,7 +137,7 @@ steps:
- label: Language Models Test (MTEB)
key: language-models-test-mteb
timeout_in_minutes: 45
timeout_in_minutes: 68
device: h200_18gb
optional: true
source_file_dependencies:
+4 -4
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 45
timeout_in_minutes: 68
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -20,7 +20,7 @@ steps:
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 50
timeout_in_minutes: 75
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -54,7 +54,7 @@ steps:
- label: "Multi-Modal Models (Standard) 4: other + whisper"
device: h200_35gb
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 50
timeout_in_minutes: 75
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -85,7 +85,7 @@ steps:
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 65
timeout_in_minutes: 98
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: PyTorch Compilation Unit Tests
device: h200_35gb
key: pytorch-compilation-unit-tests
timeout_in_minutes: 90
timeout_in_minutes: 110
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
-30
View File
@@ -8,11 +8,6 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/benchmarks/
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
@@ -28,7 +23,6 @@ steps:
- tests/entrypoints/openai/test_uds.py
- tests/v1/sample/test_logprobs_e2e.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
@@ -49,11 +43,6 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
- vllm/v1/engine/
@@ -65,7 +54,6 @@ steps:
# - tests/entrypoints/serve/dev/test_sleep.py
- tests/entrypoints/serve/tokenize/test_tokenization.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
@@ -84,16 +72,10 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- tests/utils.py
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
@@ -104,17 +86,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- vllm/tool_parsers/
- tests/utils.py
- tests/tool_use/
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
@@ -125,11 +101,6 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/distributed/
- vllm/engine/
- vllm/executor/
@@ -140,7 +111,6 @@ steps:
- tests/v1/distributed/test_hybrid_lb_dp.py
- tests/v1/distributed/test_internal_lb_dp.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export NCCL_CUMEM_HOST_ENABLE=0
@@ -26,7 +26,5 @@ steps:
- rust-toolchain.toml
- .buildkite/test_areas/rust_frontend_cargo.yaml
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
commands:
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
+16
View File
@@ -170,3 +170,19 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
key: spec-decode-deepseek-mtp-parallel-load-b200
timeout_in_minutes: 30
device: b200-k8s
optional: true
num_devices: 2
source_file_dependencies:
- vllm/v1/spec_decode/llm_base_proposer.py
- vllm/v1/spec_decode/eagle.py
- vllm/v1/worker/gpu/spec_decode/eagle/
- vllm/model_executor/models/deepseek_mtp.py
- vllm/model_executor/models/deepseek_v2.py
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
commands:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
+12
View File
@@ -181,6 +181,18 @@ pull_request_rules:
add:
- performance
- name: label-quantization
description: Automatically apply quantization label
conditions:
- label != stale
- or:
- files~=^vllm/model_executor/layers/quantization/
- title~=(?i)quant
actions:
label:
add:
- quantization
- name: label-qwen
description: Automatically apply qwen label
conditions:
+42 -1
View File
@@ -130,6 +130,47 @@ jobs:
},
],
},
quantization: {
keywords: [
{
term: "quantization",
searchIn: "both"
},
{
term: "quantized",
searchIn: "both"
},
],
},
"intel-gpu": {
// Keyword search - matches whole words only (with word boundaries)
keywords: [
{
term: "B50",
searchIn: "both"
},
{
term: "B60",
searchIn: "both"
},
{
term: "B70",
searchIn: "both"
},
{
term: "intel gpu",
searchIn: "both"
},
{
term: "Arc GPU",
searchIn: "both"
},
{
term: "BMG",
searchIn: "both"
},
],
},
// Add more label configurations here as needed
// example: {
// keywords: [...],
@@ -491,4 +532,4 @@ jobs:
issue_number: context.issue.number,
body: message,
});
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
-1
View File
@@ -257,4 +257,3 @@ vllm/grpc/vllm_engine_pb2.pyi
# Ignore generated cpu headers
csrc/cpu/cpu_attn_dispatch_generated.h
rust-coverage-tools/
-37
View File
@@ -8,8 +8,6 @@
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
CARGO_LLVM_COV_VERSION="0.8.7"
COVERAGE_TOOLS_DIR="$REPO_ROOT/rust-coverage-tools"
# Read the required toolchain from rust-toolchain.toml.
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -32,39 +30,4 @@ else
PROFILE_ARG="--release"
fi
rm -rf "$COVERAGE_TOOLS_DIR"
mkdir -p "$COVERAGE_TOOLS_DIR/bin" "$COVERAGE_TOOLS_DIR/lib"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
# rustc wrapper flags are invisible to Cargo's normal fingerprinting.
# Keep instrumented intermediates isolated when local builds switch modes.
export CARGO_TARGET_DIR="$REPO_ROOT/rust/target/coverage"
rustup component add --toolchain "$TOOLCHAIN" llvm-tools-preview
cargo +"$TOOLCHAIN" install \
--locked \
--version "$CARGO_LLVM_COV_VERSION" \
cargo-llvm-cov
eval "$(
cargo +"$TOOLCHAIN" llvm-cov show-env \
--manifest-path "$REPO_ROOT/rust/Cargo.toml" \
--sh
)"
# Build scripts and proc macros can run during compilation. Their profiles
# are unrelated to runtime coverage and would otherwise pollute the tree.
export LLVM_PROFILE_FILE=/dev/null
export VLLM_RUST_COVERAGE_OBJECTS="$COVERAGE_TOOLS_DIR/objects"
fi
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
LLVM_BIN_DIR="$(dirname "$(rustup run "$TOOLCHAIN" rustc \
--print target-libdir)")/bin"
cp "$LLVM_BIN_DIR"/{llvm-cov,llvm-profdata} "$COVERAGE_TOOLS_DIR/bin/"
chmod 0755 "$COVERAGE_TOOLS_DIR/bin/"*
cp -L "$LLVM_BIN_DIR"/../lib/libLLVM.so* "$COVERAGE_TOOLS_DIR/lib/"
chmod 0644 "$COVERAGE_TOOLS_DIR/lib/"*
fi
+1 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
-13
View File
@@ -10,16 +10,3 @@ fixes:
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
- "/usr/lib/python3.*/site-packages/vllm/::vllm/"
# Map Rust sources built in the E2E image and on Buildkite agents.
- "/workspace/rust/::rust/"
- "/var/lib/buildkite-agent/.*/rust/::rust/"
flags:
rust-unit:
paths:
- rust/
carryforward: false
rust-e2e:
paths:
- rust/
carryforward: false
+3
View File
@@ -1025,6 +1025,9 @@ __global__ void gather_and_maybe_dequant_cache(
batch_offset += offset;
int32_t block_table_id = batch_offset / block_size;
int32_t slot_id = batch_offset % block_size;
// seq_starts may push the block index past the end of the batch's block
// table row.
if (block_table_id >= block_table_stride) continue;
int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
int32_t block_id = block_table[block_table_offset];
int64_t cache_offset =
+6 -3
View File
@@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias);
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> is_padding);
void topk_sigmoid(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor);
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding);
void topk_softplus_sqrt(
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
@@ -25,7 +27,8 @@ void topk_softplus_sqrt(
double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid);
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding);
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
std::optional<torch::stable::Tensor> topk_ids,
@@ -174,7 +174,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int end_expert,
const bool renormalize,
const float* bias,
const double routed_scaling_factor)
const double routed_scaling_factor,
const bool* is_padding)
{
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int expert = result_kvp.key;
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool is_pad_row = is_padding != nullptr && is_padding[block_row];
const int idx = k * block_row + k_idx;
// Return the unbiased scores for output weights
output[idx] = inputs_after_softmax[thread_read_offset + expert];
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : num_experts);
assert(is_pad_row || indices[idx] >= 0);
source_rows[idx] = k_idx * num_rows + block_row;
if (renormalize) {
selected_sum += inputs_after_softmax[thread_read_offset + expert];
@@ -277,7 +280,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias, const double routed_scaling_factor)
const float* bias, const double routed_scaling_factor, const bool* is_padding)
{
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Add a guard to ignore experts not included by this node
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
// The lead thread from each sub-group will write out the final results to global memory. (This will be a
// single) thread per row of the input/output matrices.
const int idx = k * thread_row + k_idx;
output[idx] = max_val;
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -605,7 +610,7 @@ struct TopkConstants
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias, const double routed_scaling_factor, cudaStream_t stream)
const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)
{
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
@@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding);
}
#ifndef USE_ROCM
@@ -627,7 +632,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream);
bias, routed_scaling_factor, stream, is_padding);
#else
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -635,13 +640,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream); \
bias, routed_scaling_factor, stream, is_padding); \
} else if (WARP_SIZE == 32) { \
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, routed_scaling_factor, stream); \
bias, routed_scaling_factor, stream, is_padding); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -661,7 +666,8 @@ void topkGatingKernelLauncher(
const bool renormalize,
const float* bias,
const double routed_scaling_factor,
cudaStream_t stream) {
cudaStream_t stream,
const bool* is_padding) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
#ifndef USE_ROCM
@@ -736,7 +742,7 @@ void topkGatingKernelLauncher(
}
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding);
}
}
}
@@ -755,7 +761,8 @@ void dispatch_topk_launch(
int num_tokens, int num_experts, int topk, bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
cudaStream_t stream)
cudaStream_t stream,
std::optional<torch::stable::Tensor> is_padding)
{
const float* bias_ptr = nullptr;
if (bias.has_value()) {
@@ -769,6 +776,18 @@ void dispatch_topk_launch(
bias_ptr = bias_tensor.const_data_ptr<float>();
}
const bool* is_padding_ptr = nullptr;
if (is_padding.has_value()) {
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,
"is_padding tensor must be bool");
STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
"is_padding size mismatch, expected: ", num_tokens);
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
}
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -777,7 +796,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream);
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
} else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -786,7 +805,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream);
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
} else {
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
@@ -796,7 +815,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream);
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
}
}
@@ -806,7 +825,8 @@ void topk_softmax(
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias)
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> is_padding)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -825,15 +845,15 @@ void topk_softmax(
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream);
bias, 1.0, stream, is_padding);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream);
bias, 1.0, stream, is_padding);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, 1.0, stream);
bias, 1.0, stream, is_padding);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -846,7 +866,8 @@ void topk_sigmoid(
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor)
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -865,15 +886,15 @@ void topk_sigmoid(
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream);
bias, routed_scaling_factor, stream, is_padding);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream);
bias, routed_scaling_factor, stream, is_padding);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, routed_scaling_factor, stream);
bias, routed_scaling_factor, stream, is_padding);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -80,22 +80,27 @@ __launch_bounds__(128) __global__
OutIndType* indices, int num_rows,
int num_experts, float routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid) {
const HashIndType* tid2eid,
const bool* is_padding) {
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int lane = threadIdx.x % 32;
if (warp >= num_rows) return;
const int64_t token_id = load_index_as_int64(input_ids, warp);
const bool is_pad_row = is_padding != nullptr && is_padding[warp];
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
int expert = 0;
float weight = 0.f;
if (lane < 6) {
if (lane < 6 && !is_pad_row) {
// only load and calculate for 6 experts
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
const float x = input[warp * num_experts + expert];
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
if (isnan(weight)) {
weight = 0.f;
}
}
float weight_sum = weight;
#pragma unroll
@@ -111,7 +116,8 @@ __launch_bounds__(128) __global__
const int offset = warp * 6 + lane;
output[offset] =
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
indices[offset] = static_cast<OutIndType>(expert);
indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
: static_cast<OutIndType>(-1);
}
}
@@ -120,7 +126,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
int num_rows, int num_experts,
double routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream) {
const HashIndType* tid2eid, cudaStream_t stream,
const bool* is_padding) {
if (num_rows == 0) return;
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
cudaLaunchConfig_t config = {};
@@ -134,7 +141,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
config.numAttrs = 1;
const float scale = static_cast<float>(routed_scaling_factor);
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
num_experts, scale, input_ids, tid2eid);
num_experts, scale, input_ids, tid2eid, is_padding);
}
#endif
@@ -166,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
const int num_rows, IndType* indices, int* source_rows, const int k,
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const HashIndType* input_ids, const HashIndType* tid2eid) {
const HashIndType* input_ids, const HashIndType* tid2eid,
const bool* is_padding) {
static_assert(std::is_same_v<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -231,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
return;
}
const bool row_is_active = finished ? !finished[thread_row] : true;
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
// We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read.
@@ -249,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
cudaGridDependencySynchronize();
#endif
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
// to float
if constexpr (std::is_same_v<InputType, float>) {
if (is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = 0.f;
}
} else if constexpr (std::is_same_v<InputType, float>) {
using VecType = AlignedArray<float, ELTS_PER_LDG>;
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
const VecType* vec_thread_read_ptr =
@@ -315,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
if constexpr (USE_HASH) {
const int64_t token_id = load_index_as_int64(input_ids, thread_row);
const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
if (!is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
row_chunk[ii] = sqrtf(val);
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
// Dummy/padding tokens can result in NaN values, so
// clamp them to 0.0. Note: this clamp could likely be removed if
// 'is_padding' is made mandatory
if (isnan(val)) {
val = 0.f;
}
row_chunk[ii] = val;
}
}
float selected_sum = 0.f;
#pragma unroll
@@ -335,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
if (expert == expert_idx) {
indices[idx] = static_cast<IndType>(expert);
indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
: static_cast<IndType>(-1);
selected_sum += row_chunk[ii];
break;
}
@@ -379,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
#endif
return;
} else {
if (!is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
// Dummy/padding tokens can result in NaN values, so
// clamp them to 0.0. Note: this clamp could likely be removed if
// 'is_padding' is made mandatory
if (isnan(val)) {
val = 0.f;
}
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
}
row_chunk[ii] = val;
}
// Original TopK path: find top-k experts by score
@@ -450,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Add a guard to ignore experts not included by this node
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool should_process_row =
row_is_active && node_uses_expert && !is_pad_row;
// The lead thread from each sub-group will write out the final results
// to global memory. (This will be a single) thread per row of the
// input/output matrices.
const int idx = k * thread_row + k_idx;
if (correction_bias != nullptr) {
if (correction_bias != nullptr && should_process_row) {
max_val -= correction_bias[expert];
}
output[idx] = max_val;
indices[idx] =
should_process_row ? (expert - start_expert) : NUM_EXPERTS;
!is_pad_row ? expert - start_expert : static_cast<IndType>(-1);
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -544,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const bool use_hash, const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream) {
const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
static constexpr int BYTES_PER_LDG =
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants =
@@ -573,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper(
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
indices, source_row, k, start_expert, end_expert,
renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid);
input_ids, tid2eid, is_padding);
#else
kernel<<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert,
end_expert, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid);
input_ids, tid2eid, is_padding);
#endif
})
}
@@ -592,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
stream);
stream, is_padding);
#else
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -601,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream); \
tid2eid, stream, is_padding); \
} else if (WARP_SIZE == 32) { \
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
MAX_BYTES>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream); \
tid2eid, stream, is_padding); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -622,14 +654,14 @@ void topkGatingSoftplusSqrtKernelLauncher(
const int topk, const bool renormalize, double routed_scaling_factor,
const float* correction_bias, const bool use_hash,
const HashIndType* input_ids, const HashIndType* tid2eid,
cudaStream_t stream) {
cudaStream_t stream, const bool* is_padding) {
#ifndef USE_ROCM
if constexpr (std::is_same_v<InputType, float>) {
if (use_hash && topk == 6 && renormalize &&
(num_experts == 256 || num_experts == 384)) {
launchDsv4HashTopk<IndType, HashIndType>(
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
routed_scaling_factor, input_ids, tid2eid, stream);
routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
return;
}
}
@@ -728,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch(
int num_experts, int topk, bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
const std::optional<torch::stable::Tensor>& is_padding) {
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
bias_ptr = correction_bias.value().const_data_ptr<float>();
@@ -737,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch(
auto launch = [&](auto* topk_indices_ptr) {
using OutIndType =
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
const bool* is_padding_ptr = nullptr;
if (is_padding.has_value()) {
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==
torch::headeronly::ScalarType::Bool,
"is_padding tensor must be bool");
STD_TORCH_CHECK(is_padding_tensor.dim() == 1,
"is_padding tensor must be 1D");
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
"is_padding size mismatch, expected: ", num_tokens);
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),
"is_padding tensor must be contiguous");
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
}
if (tid2eid.has_value()) {
STD_TORCH_CHECK(input_ids.has_value(),
"input_ids is required for hash MoE");
@@ -751,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
tid2eid.value().const_data_ptr<int64_t>(), stream);
tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
} else {
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
torch::headeronly::ScalarType::Int);
@@ -761,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
tid2eid.value().const_data_ptr<int>(), stream);
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
}
} else {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -769,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
static_cast<const OutIndType*>(nullptr), stream);
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
}
};
@@ -793,7 +842,8 @@ void topk_softplus_sqrt(
bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid) {
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding) {
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
const int topk = topk_weights.size(-1);
@@ -806,21 +856,22 @@ void topk_softplus_sqrt(
dispatch_topk_softplus_sqrt_launch<float>(
gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
token_expert_indices, num_tokens, num_experts, topk, renormalize,
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
is_padding);
} else if (gating_output.scalar_type() ==
torch::headeronly::ScalarType::Half) {
dispatch_topk_softplus_sqrt_launch<__half>(
reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream);
input_ids, tid2eid, stream, is_padding);
} else if (gating_output.scalar_type() ==
torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream);
input_ids, tid2eid, stream, is_padding);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
gating_output.scalar_type());
+3 -3
View File
@@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
"bias) -> ()");
"bias, Tensor? is_padding) -> ()");
// Apply topk sigmoid to the gating outputs.
m.def(
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, "
"Tensor? bias, float routed_scaling_factor) -> ()");
"Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()");
m.def(
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, float "
"routed_scaling_factor, Tensor? "
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
"bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
// Calculate the result of moe by summing up the partial results
// from all selected experts. topk_ids/expert_map are optional and, when
+1 -12
View File
@@ -294,9 +294,6 @@ FROM base AS rust-build
ARG BUILD_OS
ARG USE_SCCACHE
ARG SCCACHE_ENDPOINT
# Temporary default for the initial CI validation. Set this back to 0 when
# ci-infra passes VLLM_RUST_COVERAGE=1 explicitly.
ARG VLLM_RUST_COVERAGE=1
# Install native tools needed only for Rust/protoc builds.
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
@@ -796,7 +793,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
# https://docs.flashinfer.ai/installation.html
# From versions.json: .flashinfer.version
ARG FLASHINFER_VERSION=0.6.14
ARG FLASHINFER_VERSION=0.6.15.post1
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -905,14 +902,6 @@ COPY ./vllm/collect_env.py .
# note that this uses vllm installed by `pip`
FROM vllm-base AS test
COPY --from=rust-build \
/workspace/rust-coverage-tools/ \
/opt/vllm-rust-coverage/
ENV PATH=/opt/vllm-rust-coverage/bin:${PATH}
ENV LD_LIBRARY_PATH=/opt/vllm-rust-coverage/lib:${LD_LIBRARY_PATH}
ENV LLVM_PROFILE_FILE=/dev/null
ADD . /vllm-workspace/
ARG PYTHON_VERSION
+53 -31
View File
@@ -339,18 +339,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
# RIXL/UCX build stages
FROM base AS build_rixl
ARG RIXL_BRANCH="39be1de8"
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
ARG UCX_BRANCH="bfb51733"
# NIXL/UCX build stages
FROM base AS build_nixl
ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ENV ROCM_PATH=/opt/rocm
ENV UCX_HOME=/usr/local/ucx
ENV RIXL_HOME=/usr/local/rixl
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
ENV NIXL_HOME=/usr/local/nixl
# RIXL build system dependences and RDMA support
# NIXL build system dependencies and RDMA support
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
libgrpc-dev \
libgrpc++-dev \
@@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
&& rm -rf /var/lib/apt/lists/*
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system meson auditwheel patchelf tomlkit
uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
RUN --mount=type=cache,target=/root/.cache/ccache \
cd /usr/local/src && \
@@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
RUN --mount=type=cache,target=/root/.cache/ccache \
git clone ${RIXL_REPO} /opt/rixl && \
cd /opt/rixl && \
git checkout ${RIXL_BRANCH} && \
git clone ${NIXL_REPO} /opt/nixl && \
cd /opt/nixl && \
git checkout ${NIXL_BRANCH} && \
CC="ccache gcc" CXX="ccache g++" \
meson setup build --prefix=${RIXL_HOME} \
meson setup build --prefix=${NIXL_HOME} \
-Ducx_path=${UCX_HOME} \
-Drocm_path=${ROCM_PATH} && \
-Dwheel_variant=rocm \
-Dbuild_tests=false \
-Dbuild_examples=false && \
cd build && \
ninja -j$(nproc) && \
ninja install
ninja install && \
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
> /etc/ld.so.conf.d/nixl.conf && \
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
>> /etc/ld.so.conf.d/nixl.conf && \
ldconfig
# Generate RIXL wheel
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
# so configure the ROCm wheel variant directly through Meson.
# Exclude libcore and libpull from auditwheel: transitive dependencies
# that are not shipped in the wheel and vary across base images.
RUN cd /opt/rixl && \
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
contrib/build-wheel.sh && \
mkdir -p /app/install && \
_ucx_install_dir=${UCX_HOME} \
./contrib/build-wheel.sh \
--output-dir /app/install \
--rocm-dir ${ROCM_PATH} \
RUN cd /opt/nixl && \
./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
CC="ccache gcc" CXX="ccache g++" \
uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
--python ${PYTHON_VERSION} \
-Csetup-args=-Ducx_path=${UCX_HOME} \
-Csetup-args=-Dwheel_variant=rocm \
-Csetup-args=-Dbuild_tests=false \
-Csetup-args=-Dbuild_examples=false && \
mkdir -p /tmp/nixl_wheels/repaired /app/install && \
auditwheel repair \
--exclude 'libamdhip64*' \
--exclude 'libcore*' \
--exclude 'libpull*' \
/tmp/nixl_wheels/nixl_rocm*.whl \
--plat manylinux_2_34_$(uname -m) \
--wheel-dir /tmp/nixl_wheels/repaired && \
./contrib/wheel_add_ucx_plugins.py \
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
/tmp/nixl_wheels/repaired/*.whl && \
cp /tmp/nixl_wheels/repaired/*.whl /app/install
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
# invalidate the slow ROCShmem build.
@@ -660,10 +680,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
fi
# Install RIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
# Install NIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
# Copy ROCShmem runtime libraries.
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
@@ -724,6 +744,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
# ROCm profiler limits workaround.
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -796,9 +817,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
&& pip uninstall -y vllm \
&& uv pip install --system *.whl
# Install RIXL wheel
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
uv pip install --system /rixl_install/*.whl
# Install NIXL ROCm wheel
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
uv pip install --system /nixl_install/*.whl
ARG COMMON_WORKDIR
ARG BASE_IMAGE
@@ -813,6 +834,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV TOKENIZERS_PARALLELISM=false
+1 -1
View File
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
ARG FA_BRANCH="0e60e394"
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
ARG AITER_BRANCH="v0.1.16.post3"
ARG AITER_BRANCH="v0.1.16.post5"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
ARG MORI_BRANCH="v1.1.0"
ARG MORI_REPO="https://github.com/ROCm/mori.git"
+13 -13
View File
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
}
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
# stage and is irrelevant when building --target ci_base itself.
variable "CI_BASE_IMAGE" {
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
# functions produce no entries when the variable is empty.
variable "RIXL_BRANCH" {
variable "NIXL_BRANCH" {
default = ""
}
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
default = ""
}
variable "RIXL_CACHE_KEY" {
variable "NIXL_CACHE_KEY" {
default = ""
}
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
])
}
# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
# These stages are pinned to specific upstream commit hashes, so cache keys use
# those hashes rather than the Buildkite commit. This means the cache persists
# across all vLLM commits as long as the upstream dependency pins don't change.
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
function "get_cache_from_rocm_deps" {
params = []
result = compact([
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
])
}
function "get_cache_to_rocm_rixl" {
function "get_cache_to_rocm_nixl" {
params = []
result = compact([
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
])
}
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
# (e.g., requirements change), these stages are cache hits if their upstream
# pins haven't changed -- saving ~35min of compilation.
target "rixl-rocm-ci" {
target "nixl-rocm-ci" {
inherits = ["_common-rocm", "_ci-rocm"]
target = "build_rixl"
target = "build_nixl"
cache-from = get_cache_from_rocm_deps()
cache-to = get_cache_to_rocm_rixl()
cache-to = get_cache_to_rocm_nixl()
output = ["type=cacheonly"]
}
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
output = ["type=cacheonly"]
}
# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
# files drifts from the remote image label. Per-PR builds then pull the result
# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
]),
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
# Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
# are cache hits even when ci_base itself needs rebuilding.
get_cache_from_rocm_deps(),
)
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
# Group for ci_base builds -- exports dependency stage caches alongside the
# ci_base image so future rebuilds can reuse them independently.
group "ci-base-rocm-ci-with-deps" {
targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
}
+2 -2
View File
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
# overrides) and writes a small HCL override before invoking CI targets.
variable "RIXL_BRANCH" {
variable "NIXL_BRANCH" {
default = ""
}
@@ -106,7 +106,7 @@ target "test-rocm" {
output = ["type=docker"]
}
# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
# torchcodec, requirements, etc.). Used by the weekly scheduled build and
# the auto-rebuild trigger when requirements change in a PR.
target "ci-base-rocm" {
+1 -4
View File
@@ -46,9 +46,6 @@
"TORCH_CUDA_ARCH_LIST": {
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
},
"VLLM_RUST_COVERAGE": {
"default": "1"
},
"MAX_JOBS": {
"default": "2"
},
@@ -71,7 +68,7 @@
"default": "true"
},
"FLASHINFER_VERSION": {
"default": "0.6.14"
"default": "0.6.15.post1"
},
"GDRCOPY_CUDA_VERSION": {
"default": "12.8"
+1 -5
View File
@@ -13,11 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files
For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
+3
View File
@@ -0,0 +1,3 @@
// Reo.Dev documentation tracking
// https://docs.reo.dev/integrations/input-sources/developer-insights/documentation
!function(){var e,t,n;e="d5c4337961ef0ac",t=function(){Reo.init({clientID:"d5c4337961ef0ac", enableThirdPartyTracking: true})},(n=document.createElement("script")).src="https://static.reo.dev/"+e+"/reo.js",n.defer=!0,n.onload=t,document.head.appendChild(n)}();
+1
View File
@@ -160,3 +160,4 @@ extra_javascript:
- https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
- mkdocs/javascript/edit_and_feedback.js
- mkdocs/javascript/slack_and_forum.js
- mkdocs/javascript/reo.js
+3 -3
View File
@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
# install_requires so the published wheel does not carry an unresolvable pin
--extra-index-url https://flashinfer.ai/whl/
flashinfer-python==0.6.14
flashinfer-cubin==0.6.14
flashinfer-python==0.6.15.post1
flashinfer-cubin==0.6.15.post1
apache-tvm-ffi==0.1.10
tilelang==0.1.9
nvidia-cudnn-frontend>=1.19.1
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
nvidia-cutlass-dsl[cu13]==4.6.0
quack-kernels>=0.4.0 # Required for tml-fa4
quack-kernels>=0.6.1 # Required for CUTLASS DSL 4.6 by MSA
# Tokenspeed_MLA for faster mla with spec decode
tokenspeed-mla==0.1.8; platform_system == "Linux"
+1 -1
View File
@@ -12,4 +12,4 @@ ray[data]
setuptools==78.1.0
setuptools-rust>=1.9.0
nixl==0.3.0
tpu-inference==0.24.0
tpu-inference==0.25.0
+13
View File
@@ -14,6 +14,10 @@ service Generate {
rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
}
service Control {
rpc Abort (AbortRequest) returns (AbortResponse) {}
}
// ======================================================================================
// Generate Request
// ======================================================================================
@@ -201,3 +205,12 @@ message TokenIds {
repeated uint32 ids = 1;
}
// ======================================================================================
// Control
// ======================================================================================
message AbortRequest {
repeated string request_ids = 1;
}
message AbortResponse {}
+106 -61
View File
@@ -54,6 +54,7 @@ mod stream;
use vllm_engine_core_client::EngineCoreClient;
use vllm_engine_core_client::protocol::dtype::ModelDtype;
use vllm_engine_core_client::protocol::multimodal::MmFeatures;
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
use vllm_llm::Llm;
use vllm_text::{Prompt, TextLlm, TextRequest};
@@ -88,6 +89,92 @@ pub fn validate_parser_overrides(
Ok(())
}
/// Chat request preparation shared by inference and render-only frontends.
pub struct ChatRequestProcessor {
backend: DynChatBackend,
/// Effective model dtype reported by the engine.
/// Absent for text-only frontends without an engine handshake.
model_dtype: Option<ModelDtype>,
}
impl ChatRequestProcessor {
/// Create a processor with multimodal support using the effective model
/// dtype reported by the engine.
fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self {
Self {
backend,
model_dtype: Some(model_dtype),
}
}
/// Create a render-only processor that rejects multimodal requests.
pub fn render_only(backend: DynChatBackend) -> Self {
Self {
backend,
model_dtype: None,
}
}
async fn finalize_rendered_prompt(
&self,
request: &ChatRequest,
rendered: RenderedPrompt,
) -> Result<(Prompt, Option<MmFeatures>)> {
match self.model_dtype {
Some(model_dtype) => {
multimodal::finalize_rendered_prompt(
request,
rendered,
self.backend.multimodal_model_info(),
model_dtype,
)
.await
}
None if !request.has_multimodal() => Ok((rendered.prompt, None)),
None => Err(Error::UnsupportedMultimodalRenderer),
}
}
/// Prepare one chat request without submitting it to an engine.
pub async fn prepare(
&self,
mut request: ChatRequest,
options: NewChatOutputProcessorOptions<'_>,
) -> Result<(TextRequest, DynChatOutputProcessor)> {
request.validate()?;
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
let arrival_time = vllm_llm::current_unix_timestamp_secs();
let output_processor = self.backend.new_chat_output_processor(&mut request, options)?;
let rendered = self.backend.chat_renderer().render(&request)?;
let reasoning_parser_kwargs =
request
.sampling_params
.structured_outputs
.is_some()
.then(|| ReasoningParserKwargs {
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
});
let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
let text_request = TextRequest {
request_id: request.request_id,
prompt,
mm_features,
sampling_params: request.sampling_params,
decode_options: request.decode_options,
intermediate: request.intermediate,
priority: request.priority,
cache_salt: request.cache_salt,
add_special_tokens: request.add_special_tokens,
data_parallel_rank: request.data_parallel_rank,
reasoning_parser_kwargs,
lora_request: request.lora_request,
arrival_time: Some(arrival_time),
};
Ok((text_request, output_processor))
}
}
/// Structured chat facade above [`TextLlm`].
///
/// This layer stays above raw text semantics: it takes care of chat-template
@@ -95,9 +182,7 @@ pub fn validate_parser_overrides(
/// request semantics such as tool calls.
pub struct ChatLlm {
text: TextLlm,
backend: DynChatBackend,
/// Effective model dtype reported by the engine.
model_dtype: ModelDtype,
processor: ChatRequestProcessor,
/// Tool-call parser selection.
tool_call_parser: ParserSelection,
/// Reasoning parser selection.
@@ -112,8 +197,7 @@ impl ChatLlm {
Self {
text,
backend,
model_dtype,
processor: ChatRequestProcessor::new(backend, model_dtype),
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
}
@@ -140,7 +224,7 @@ impl ChatLlm {
/// Override the effective model dtype used for multimodal tensor encoding.
pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
self.model_dtype = model_dtype;
self.processor.model_dtype = Some(model_dtype);
self
}
@@ -172,57 +256,23 @@ impl ChatLlm {
}
/// Render, tokenize, and submit one chat request.
pub async fn chat(&self, mut request: ChatRequest) -> Result<ChatEventStream> {
request.validate()?;
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
let arrival_time = vllm_llm::current_unix_timestamp_secs();
let output_processor = self.backend.new_chat_output_processor(
&mut request,
NewChatOutputProcessorOptions {
tool_call_parser: &self.tool_call_parser,
reasoning_parser: &self.reasoning_parser,
},
)?;
let rendered = self.backend.chat_renderer().render(&request)?;
let reasoning_parser_kwargs =
request
.sampling_params
.structured_outputs
.is_some()
.then(|| ReasoningParserKwargs {
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
});
let (prompt, mm_features) = multimodal::finalize_rendered_prompt(
&request,
rendered,
self.backend.multimodal_model_info(),
self.model_dtype,
)
.await?;
let text_request = TextRequest {
request_id: request.request_id.clone(),
prompt,
mm_features,
sampling_params: request.sampling_params,
decode_options: request.decode_options,
intermediate: request.intermediate,
priority: request.priority,
cache_salt: request.cache_salt,
add_special_tokens: request.add_special_tokens,
data_parallel_rank: request.data_parallel_rank,
reasoning_parser_kwargs,
lora_request: request.lora_request,
arrival_time: Some(arrival_time),
};
pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
let (text_request, output_processor) = self
.processor
.prepare(
request,
NewChatOutputProcessorOptions {
tool_call_parser: &self.tool_call_parser,
reasoning_parser: &self.reasoning_parser,
},
)
.await?;
let request_id = text_request.request_id.clone();
let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
let structured_stream = output_processor.process(decoded_stream)?;
Ok(ChatEventStream::new(request.request_id, structured_stream))
Ok(ChatEventStream::new(request_id, structured_stream))
}
/// Render through the chat template and tokenize, without submitting to the engine.
@@ -233,14 +283,9 @@ impl ChatLlm {
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
request.validate()?;
let rendered = self.backend.chat_renderer().render(&request)?;
let (prompt, _mm_features) = multimodal::finalize_rendered_prompt(
&request,
rendered,
self.backend.multimodal_model_info(),
self.model_dtype,
)
.await?;
let rendered = self.processor.backend.chat_renderer().render(&request)?;
let (prompt, _mm_features) =
self.processor.finalize_rendered_prompt(&request, rendered).await?;
let tokenizer = self.text.tokenizer();
let token_ids = match prompt {
+6 -11
View File
@@ -8,7 +8,7 @@ use tonic_health::ServingStatus;
use tonic_health::server::HealthReporter;
use tracing::{info, warn};
use super::GenerateGrpcService;
use super::{ControlGrpcService, GenerateGrpcService};
pub(crate) async fn monitor_health(
mut health_reporter: HealthReporter,
@@ -16,21 +16,18 @@ pub(crate) async fn monitor_health(
shutdown: CancellationToken,
) {
let generate_service = GenerateGrpcService::NAME;
let control_service = ControlGrpcService::NAME;
let status = ServingStatus::NotServing;
let health_event_first = tokio::select! {
result = engine_health.wait_for(|healthy| !*healthy) => {
match result {
Ok(_) => warn!(
generate_service,
overall_service = true,
status = ?status,
reason = "engine_unhealthy",
"marking gRPC health services as not serving"
),
Err(error) => warn!(
%error,
generate_service,
overall_service = true,
status = ?status,
reason = "health_channel_closed",
"engine health channel closed; marking gRPC health services as not serving"
@@ -40,8 +37,6 @@ pub(crate) async fn monitor_health(
}
_ = shutdown.cancelled() => {
info!(
generate_service,
overall_service = true,
status = ?status,
reason = "server_shutdown",
"server shutting down; marking gRPC health services as not serving"
@@ -51,20 +46,20 @@ pub(crate) async fn monitor_health(
};
health_reporter.set_not_serving::<GenerateGrpcService>().await;
// Generate is currently the only engine-backed gRPC service, so overall
// server health intentionally mirrors it.
health_reporter.set_not_serving::<ControlGrpcService>().await;
// Both gRPC services use the same engine client, so overall server health
// mirrors their shared engine health.
health_reporter.set_service_status("", status).await;
if health_event_first {
shutdown.cancelled().await;
info!(
generate_service,
overall_service = true,
reason = "server_shutdown",
"server shutting down; closing gRPC health watches"
);
}
health_reporter.clear_service_status(generate_service).await;
health_reporter.clear_service_status(control_service).await;
health_reporter.clear_service_status("").await;
}
+32
View File
@@ -26,8 +26,10 @@ pub mod pb {
}
pub(crate) use health::monitor_health;
pub use pb::control_server::ControlServer;
pub use pb::generate_server::GenerateServer;
pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
#[cfg(test)]
@@ -44,6 +46,36 @@ impl GenerateServiceImpl {
}
}
/// gRPC control service backed by the shared application state.
pub struct ControlServiceImpl {
state: Arc<AppState>,
}
impl ControlServiceImpl {
pub fn new(state: Arc<AppState>) -> Self {
Self { state }
}
}
#[tonic::async_trait]
impl pb::control_server::Control for ControlServiceImpl {
async fn abort(
&self,
request: Request<pb::AbortRequest>,
) -> Result<Response<pb::AbortResponse>, Status> {
let request_ids = request.into_inner().request_ids;
if request_ids.is_empty() {
return Ok(Response::new(pb::AbortResponse {}));
}
self.state
.chat
.abort(&request_ids)
.await
.map_err(|error| Status::internal(error.to_report_string()))?;
Ok(Response::new(pb::AbortResponse {}))
}
}
#[tonic::async_trait]
impl pb::generate_server::Generate for GenerateServiceImpl {
type GenerateStreamStream =
+122 -18
View File
@@ -38,8 +38,9 @@ use vllm_tokenizer::test_utils::TestTokenizer;
use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
use super::pb::control_client::ControlClient;
use super::pb::generate_client::GenerateClient;
use super::{GenerateServer, GenerateServiceImpl, pb};
use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb};
use crate::listener::{Listener, MaybeTlsListener};
use crate::state::AppState;
use crate::tls;
@@ -153,10 +154,6 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
dealer.recv().await.expect("recv engine message").into_vec()
}
fn test_llm(client: EngineCoreClient) -> Llm {
Llm::new(client).with_request_id_randomization(false)
}
#[derive(Clone, Debug)]
struct FakeTextBackend;
@@ -206,6 +203,7 @@ async fn setup_grpc_service(
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> (
GenerateServer<GenerateServiceImpl>,
ControlServer<ControlServiceImpl>,
tokio::sync::watch::Receiver<bool>,
MockEngineTask,
) {
@@ -243,12 +241,13 @@ async fn setup_grpc_service(
let engine_health = client.subscribe_health();
let chat = ChatLlm::from_shared_backend(
test_llm(client),
Llm::new(client),
Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
);
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
(
GenerateServer::new(GenerateServiceImpl::new(state)),
GenerateServer::new(GenerateServiceImpl::new(state.clone())),
ControlServer::new(ControlServiceImpl::new(state)),
engine_health,
engine_task,
)
@@ -264,9 +263,11 @@ async fn grpc_test_server(
tokio::task::JoinHandle<()>,
MockEngineTask,
) {
let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (generate_service, control_service, engine_health, engine_task) =
setup_grpc_service(engine_id, output_specs).await;
let (channel, server_task) = start_grpc_test_server(
svc,
generate_service,
control_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
@@ -276,11 +277,13 @@ async fn grpc_test_server(
async fn start_grpc_test_server(
generate_service: GenerateServer<GenerateServiceImpl>,
control_service: ControlServer<ControlServiceImpl>,
engine_health: tokio::sync::watch::Receiver<bool>,
shutdown: tokio_util::sync::CancellationToken,
) -> (Channel, tokio::task::JoinHandle<()>) {
let (health_reporter, health_service) = health_reporter();
health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr");
@@ -289,6 +292,7 @@ async fn start_grpc_test_server(
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
let server = TonicServer::builder()
.add_service(health_service)
.add_service(control_service)
.add_service(generate_service)
.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
let health_monitor =
@@ -319,7 +323,8 @@ async fn grpc_tls_test_server(
certs: &TestCerts,
cert_reqs: i32,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (generate_service, control_service, _engine_health, engine_task) =
setup_grpc_service(engine_id, output_specs).await;
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
.expect("build grpc tls config");
@@ -329,7 +334,8 @@ async fn grpc_tls_test_server(
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
TonicServer::builder()
.add_service(svc)
.add_service(control_service)
.add_service(generate_service)
.serve_with_incoming(incoming)
.await
.expect("grpc tls server");
@@ -409,7 +415,7 @@ async fn grpc_server_with_keepalive(
engine_id: impl Into<EngineId>,
keepalive: Option<Duration>,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, _engine_health, engine_task) =
let (generate_service, control_service, _engine_health, engine_task) =
setup_grpc_service(engine_id, default_stream_output_specs()).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
@@ -425,7 +431,8 @@ async fn grpc_server_with_keepalive(
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
builder
.add_service(svc)
.add_service(control_service)
.add_service(generate_service)
.serve_with_incoming(incoming)
.await
.expect("grpc server");
@@ -1073,14 +1080,106 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn control_abort_resolves_external_id_and_empty_is_noop() {
let (generate_service, control_service, engine_health, engine_task) =
setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await;
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
let mut generate_client = GenerateClient::new(channel.clone());
let mut control_client = ControlClient::new(channel);
let request_id = "test-abort-active";
let mut stream = generate_client
.generate_stream(pb::GenerateRequest {
request_id: request_id.to_string(),
model: "test-model".to_string(),
prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
stopping: Some(pb::StoppingCriteria {
max_new_tokens: 10,
..Default::default()
}),
..Default::default()
})
.await
.expect("start generation")
.into_inner();
loop {
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for active generation output")
.expect("read active generation output")
.expect("generation ended before producing output");
if let Some(output) = response.outputs {
assert!(
output.finish_info.is_none(),
"generation finished before abort behavior was exercised"
);
break;
}
}
control_client
.abort(pb::AbortRequest::default())
.await
.expect("empty abort should be a no-op");
assert!(
tokio::time::timeout(Duration::from_millis(100), stream.message())
.await
.is_err(),
"empty abort unexpectedly ended the active generation"
);
control_client
.abort(pb::AbortRequest {
request_ids: vec![
request_id.to_string(),
request_id.to_string(),
"unknown".to_string(),
],
})
.await
.expect("abort active generation");
let finish_reason = loop {
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for aborted generation")
.expect("read aborted generation")
.expect("generation ended without an aborted response");
if let Some(finish_info) = response.outputs.and_then(|output| output.finish_info) {
break finish_info.finish_reason;
}
};
assert_eq!(finish_reason, pb::finish_info::FinishReason::Aborted as i32);
control_client
.abort(pb::AbortRequest {
request_ids: vec![request_id.to_string()],
})
.await
.expect("repeated abort should be idempotent");
engine_task.await.expect("mock engine task");
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
let (generate_service, _connected_engine_health, _engine_task) =
let (generate_service, control_service, _connected_engine_health, _engine_task) =
setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
@@ -1088,7 +1187,7 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
let mut health_client = HealthClient::new(channel);
let mut health_streams = Vec::new();
for service in ["vllm.Generate", ""] {
for service in ["vllm.Generate", "vllm.Control", ""] {
let service_label = if service.is_empty() {
"overall"
} else {
@@ -1143,14 +1242,19 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_watch_closes_on_graceful_shutdown() {
let (generate_service, engine_health, _engine_task) = setup_grpc_service(
let (generate_service, control_service, engine_health, _engine_task) = setup_grpc_service(
b"engine-grpc-health-shutdown",
default_stream_output_specs(),
)
.await;
let shutdown = tokio_util::sync::CancellationToken::new();
let (channel, server_task) =
start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
shutdown.clone(),
)
.await;
let mut health_client = HealthClient::new(channel);
let mut stream = health_client
.watch(HealthCheckRequest {
+4
View File
@@ -207,6 +207,9 @@ where
let (health_reporter, health_service) = health_reporter();
let engine_health = state.engine_core_client().subscribe_health();
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
health_reporter.set_serving::<grpc::ControlGrpcService>().await;
let control_service =
grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()));
let generate_service =
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
let svc = TonicServer::builder()
@@ -214,6 +217,7 @@ where
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
.layer(middleware::request_runtime_layer(state.clone()))
.add_service(health_service)
.add_service(control_service)
.add_service(generate_service);
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
+108 -72
View File
@@ -38,32 +38,22 @@ trait_set! {
pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
}
/// Raw text facade above [`Llm`].
///
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
/// in, decoded text deltas and terminal metadata flow out.
pub struct TextLlm {
/// Generate-only client owned by this text facade.
llm: Llm,
/// Text request preparation shared by inference and render-only frontends.
pub struct TextRequestProcessor {
/// Tokenizer/model metadata backend responsible for prompt encode/decode
/// and sampling hints.
backend: DynTextBackend,
/// Runtime context window size reported by the engine startup handshake.
/// Render-only frontends supply the downstream engine's effective value.
max_model_len: u32,
/// Maximum number of top log probabilities accepted by this text facade.
max_logprobs: i32,
}
impl TextLlm {
/// Create a new text-generation facade from a shared LLM client plus a text
/// backend.
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
// The engine-reported value reflects the post-profiling, auto-fitted
// KV cache limit used at runtime.
let max_model_len = llm.engine_core_client().max_model_len();
impl TextRequestProcessor {
/// Create a processor with the effective model context length.
pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self {
Self {
llm,
backend,
max_model_len,
max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
@@ -78,61 +68,18 @@ impl TextLlm {
self
}
/// Return the backend model ID.
pub fn model_id(&self) -> &str {
self.backend.model_id()
}
/// Expose the underlying engine-core client for low-level utility/admin
/// calls.
pub fn engine_core_client(&self) -> &EngineCoreClient {
self.llm.engine_core_client()
}
/// Return the tokenizer used by this text backend.
/// Return the tokenizer used by this processor.
pub fn tokenizer(&self) -> DynTokenizer {
self.backend.tokenizer()
}
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
pub fn tokenizer_vocab_size(&self) -> usize {
self.backend.tokenizer_vocab_size()
/// Return the effective model context length.
pub fn max_model_len(&self) -> u32 {
self.max_model_len
}
/// Model vocabulary size from the model config, used to bound generated
/// token IDs and logits-domain sampling controls.
pub fn model_vocab_size(&self) -> usize {
self.backend.model_vocab_size()
}
/// Tokenize if needed, lower to a generate request, and return the raw
/// token stream.
pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
let (_, raw_stream) = self.generate_inner(request).await?;
Ok(raw_stream)
}
/// Tokenize if needed, lower to a generate request, and stream
/// incrementally decoded text.
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
let (text_request, raw_stream) = self.generate_inner(request).await?;
let tokenizer = self.backend.tokenizer();
let decoded_stream = output::decoded_text_event_stream(
text_request.request_id,
tokenizer,
raw_stream,
text_request.decode_options,
text_request.intermediate,
);
Ok(decoded_stream)
}
async fn generate_inner(
&self,
mut request: TextRequest,
) -> Result<(TextRequest, GenerateOutputStream)> {
/// Tokenize and lower one request without submitting it to an engine.
pub fn prepare(&self, mut request: TextRequest) -> Result<PreparedTextRequest> {
request.validate()?;
if request.arrival_time.is_none() {
@@ -146,7 +93,6 @@ impl TextLlm {
// and infra workloads bypass chat rendering and tokenizer overhead entirely.
Prompt::TokenIds(token_ids) => token_ids,
};
let sampling_hints = self.backend.sampling_hints()?;
let sampling_limits = SamplingLimits {
max_model_len: self.max_model_len,
@@ -155,16 +101,106 @@ impl TextLlm {
tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
};
let PreparedTextRequest {
text_request,
generate_request,
} = lower_text_request(
lower_text_request(
request,
prompt_token_ids,
sampling_hints,
sampling_limits,
&*tokenizer,
)?;
tokenizer.as_ref(),
)
}
}
/// Raw text facade above [`Llm`].
///
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
/// in, decoded text deltas and terminal metadata flow out.
pub struct TextLlm {
/// Generate-only client owned by this text facade.
llm: Llm,
/// Shared engine-free request preparation.
processor: TextRequestProcessor,
}
impl TextLlm {
/// Create a new text-generation facade from a shared LLM client plus a text
/// backend.
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
// The engine-reported value reflects the post-profiling, auto-fitted
// KV cache limit used at runtime.
let max_model_len = llm.engine_core_client().max_model_len();
Self {
llm,
processor: TextRequestProcessor::new(backend, max_model_len),
}
}
/// Override the maximum accepted logprobs count.
pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
self.processor = self.processor.with_max_logprobs(max_logprobs);
self
}
/// Return the backend model ID.
pub fn model_id(&self) -> &str {
self.processor.backend.model_id()
}
/// Expose the underlying engine-core client for low-level utility/admin
/// calls.
pub fn engine_core_client(&self) -> &EngineCoreClient {
self.llm.engine_core_client()
}
/// Return the tokenizer used by this text backend.
pub fn tokenizer(&self) -> DynTokenizer {
self.processor.tokenizer()
}
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
pub fn tokenizer_vocab_size(&self) -> usize {
self.processor.backend.tokenizer_vocab_size()
}
/// Model vocabulary size from the model config, used to bound generated
/// token IDs and logits-domain sampling controls.
pub fn model_vocab_size(&self) -> usize {
self.processor.backend.model_vocab_size()
}
/// Tokenize if needed, lower to a generate request, and return the raw
/// token stream.
pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
let (_, raw_stream) = self.generate_inner(request).await?;
Ok(raw_stream)
}
/// Tokenize if needed, lower to a generate request, and stream
/// incrementally decoded text.
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
let (text_request, raw_stream) = self.generate_inner(request).await?;
let tokenizer = self.processor.tokenizer();
let decoded_stream = output::decoded_text_event_stream(
text_request.request_id,
tokenizer,
raw_stream,
text_request.decode_options,
text_request.intermediate,
);
Ok(decoded_stream)
}
async fn generate_inner(
&self,
request: TextRequest,
) -> Result<(TextRequest, GenerateOutputStream)> {
let PreparedTextRequest {
text_request,
generate_request,
} = self.processor.prepare(request)?;
let raw_stream = self.llm.generate(generate_request).await?;
Ok((text_request, raw_stream))
+62
View File
@@ -214,3 +214,65 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs():
base_hash = CacheConfig().compute_hash()
assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash
assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash
def test_envs_compile_factors_relocation_invariant(tmp_path):
"""Relocating HOME or the XDG roots must not change the compile-cache
env hash.
Location-derived env vars (VLLM_XLA_CACHE_PATH from XDG_CACHE_HOME,
VLLM_CONFIG_ROOT from XDG_CONFIG_HOME/HOME) carry no information about
compiled artifacts, only about where directories live. When they leak
into compile_factors(), a cache produced under one HOME/XDG layout
silently misses under another - which defeats copying or pre-baking a
compile cache into a container image.
"""
import os
import subprocess
import sys
code = """
import sys
import logging
logging.disable(logging.CRITICAL)
from vllm import envs
from vllm.config.utils import hash_factors
print(hash_factors(envs.compile_factors()))
"""
def hash_with(extra_env):
env = {**dict(os.environ), "VLLM_LOGGING_LEVEL": "ERROR"}
# Drop explicit overrides so the derived defaults are what is
# exercised, then apply the relocation under test.
for key in ("VLLM_XLA_CACHE_PATH", "VLLM_CONFIG_ROOT", "VLLM_CACHE_ROOT"):
env.pop(key, None)
env.update(extra_env)
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
check=True,
env=env,
)
return result.stdout.strip()
xdg_cache = tmp_path / "relocated-xdg-cache"
xdg_config = tmp_path / "relocated-xdg-config"
new_home = tmp_path / "relocated-home"
for d in (xdg_cache, xdg_config, new_home):
d.mkdir()
base = hash_with({})
relocated_xdg = hash_with(
{"XDG_CACHE_HOME": str(xdg_cache), "XDG_CONFIG_HOME": str(xdg_config)}
)
relocated_home = hash_with({"HOME": str(new_home)})
assert relocated_xdg == base, (
"XDG_CACHE_HOME/XDG_CONFIG_HOME relocation changed the compile-cache "
"env hash - a location-only derived var is leaking into the key"
)
assert relocated_home == base, (
"HOME relocation changed the compile-cache env hash - a "
"location-only derived var is leaking into the key"
)
+9
View File
@@ -43,6 +43,15 @@ def test_language_model_only_affects_model_hash():
assert base_hash != lm_only_hash
@pytest.mark.parametrize("backend_arg", ["video_backend", "backend"])
def test_use_gpu_video_backend_from_media_io_kwargs(backend_arg: str):
config = MultiModalConfig(
media_io_kwargs={"video": {backend_arg: "pynvvideocodec"}}
)
assert config.use_gpu_video_backend()
def test_mm_encoder_fp8_scale_path_requires_fp8():
with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
+8 -5
View File
@@ -254,8 +254,8 @@ def test_multiproc_executor_shutdown_cleanup():
for worker in executor.workers:
assert not worker.proc.is_alive(), "Worker processes should be terminated"
# Verify shutdown event is set
assert executor.shutdown_event.is_set(), "Shutdown event should be set"
# Verify shutdown flag is set
assert executor.shutting_down, "Shutdown flag should be set"
# Multiple shutdowns should be safe (idempotent)
executor.shutdown()
@@ -292,7 +292,6 @@ def test_multiproc_executor_pipeline_parallel():
"Max concurrent batches should follow the configured PP/async "
"scheduling policy"
)
finally:
# Clean up
executor.shutdown()
@@ -338,6 +337,10 @@ def test_multiproc_executor_multi_node():
- Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
Total world_size = 4, nnodes = 2
"""
# Python 3.14+ changed default multiprocessing start method to 'forkserver'
# which cannot pickle nested functions. Use 'fork' for this test.
mp_ctx = multiprocessing.get_context("fork")
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
port = s.getsockname()[1]
@@ -405,12 +408,12 @@ def test_multiproc_executor_multi_node():
executor.shutdown()
# Create a queue to collect results from both processes
result_queue: multiprocessing.Queue[dict[str, int | bool]] = multiprocessing.Queue()
result_queue: multiprocessing.Queue[dict[str, int | bool]] = mp_ctx.Queue()
# Start both node processes
processes = []
for node_rank in range(2):
p = multiprocessing.Process(
p = mp_ctx.Process(
target=run_node,
args=(node_rank, result_queue, port),
name=f"Node{node_rank}",
@@ -9,8 +9,6 @@ from pydantic import TypeAdapter, ValidationError
from vllm import PoolingParams
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
from vllm.entrypoints.pooling.embed.protocol import (
CohereEmbedContent,
CohereEmbedInput,
CohereEmbedRequest,
EmbeddingBatchChatInputRequest,
EmbeddingBatchChatRequest,
@@ -19,7 +17,10 @@ from vllm.entrypoints.pooling.embed.protocol import (
EmbeddingCompletionRequest,
EmbeddingRequest,
)
from vllm.entrypoints.pooling.typing import PoolingServeContext
from vllm.entrypoints.pooling.typing import (
PoolingEngineInput,
PoolingServeContext,
)
from vllm.outputs import PoolingOutput, PoolingRequestOutput
@@ -410,6 +411,7 @@ class TestChunkedEmbeddingProcessing:
def _make_handler(cls):
handler = object.__new__(EmbedIOProcessor)
handler.model_config = cls._FakeModelConfig()
handler.enable_chunked_processing = True
return handler
@staticmethod
@@ -421,15 +423,29 @@ class TestChunkedEmbeddingProcessing:
}
)
assert isinstance(request, EmbeddingCompletionRequest)
pooling_params = PoolingParams()
return PoolingServeContext(
request=request,
pooling_params=PoolingParams(),
pooling_params=pooling_params,
model_name="test",
request_id="embd-client-prompt-999-chunk-888",
engine_inputs=[
{"prompt_token_ids": [0, 1, 2, 3, 4]},
{"prompt_token_ids": [10, 11]},
PoolingEngineInput(
prompts={"prompt_token_ids": [0, 1, 2, 3, 4]},
params=pooling_params,
lora_requests=None,
priorities=0,
),
PoolingEngineInput(
prompts={"prompt_token_ids": [10, 11]},
params=pooling_params,
lora_requests=None,
priorities=0,
),
],
lora_request=None,
priorities=0,
prompt_extras=None,
)
@staticmethod
@@ -450,7 +466,7 @@ class TestChunkedEmbeddingProcessing:
handler = self._make_handler()
ctx = self._make_context()
handler._pre_process_chunked(ctx)
handler.maybe_pre_process_chunked(ctx)
assert ctx.prompt_request_ids == [
"embd-client-prompt-999-chunk-888-prompt-0-chunk-0",
@@ -488,227 +504,3 @@ class TestChunkedEmbeddingProcessing:
ctx.final_res_batch[1].outputs.data,
torch.tensor([9.0, 9.0]),
)
class TestPreProcessCohereOnline:
"""Unit tests for EmbedIOProcessor._pre_process_cohere_online."""
@staticmethod
def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
return PoolingServeContext(
request=CohereEmbedRequest(model="test", **request_kwargs),
pooling_params=PoolingParams(),
model_name="test",
request_id="embd-test",
)
@staticmethod
def _make_handler():
handler = object.__new__(EmbedIOProcessor)
handler._validate_input_type = lambda _input_type: None
return handler
def test_text_only_without_task_prefix_uses_completion_path(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"])
calls: list[tuple[str, object]] = []
def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
calls.append(("completion", prompt_input))
return ["completion"]
handler._get_task_instruction_prefix = lambda _input_type: None
handler._has_chat_template = lambda: False
handler._preprocess_cmpl_online = preprocess_cmpl_online
handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
"text-only request should not require chat rendering"
)
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["completion"]
assert calls == [("completion", ["hello"])]
def test_text_only_falls_back_to_prefixed_completion_without_template(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"], input_type="query")
calls: list[tuple[str, object]] = []
def preprocess_cmpl(request, prompt_input, prompt_embeds):
calls.append(("completion", prompt_input))
return ["fallback"]
handler._get_task_instruction_prefix = lambda _input_type: "query: "
handler._has_chat_template = lambda: False
handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
"chat rendering should be skipped without a template"
)
handler._preprocess_cmpl_online = preprocess_cmpl
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["fallback"]
assert calls == [("completion", ["query: hello"])]
def test_text_only_with_template_uses_chat_path(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"], input_type="query")
calls: list[tuple[str, object]] = []
def batch_render_chat(
request,
all_messages,
truncate_prompt_tokens,
truncation_side,
):
calls.append(
(
"chat",
{
"request": request,
"all_messages": all_messages,
"truncate_prompt_tokens": truncate_prompt_tokens,
"truncation_side": truncation_side,
},
)
)
return ["chat"]
handler._get_task_instruction_prefix = lambda _input_type: "query: "
handler._has_chat_template = lambda: True
handler._batch_render_chat = batch_render_chat
handler._preprocess_cmpl_online = lambda *_args, **_kwargs: pytest.fail(
"completion path should be skipped when a template exists"
)
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["chat"]
assert calls == [
(
"chat",
{
"request": ctx.request,
"all_messages": [
handler._mixed_input_to_messages(
CohereEmbedInput(
content=[CohereEmbedContent(type="text", text="hello")]
),
task_prefix="query: ",
)
],
"truncate_prompt_tokens": -1,
"truncation_side": None,
},
)
]
class TestPreProcessOpenAIEmbeddingChatOnline:
"""Unit tests for OpenAI embedding chat preprocessing."""
class _FakeModelConfig:
max_model_len = 128
encoder_config: dict[str, object] = {}
pooler_config = None
multimodal_config = None
is_encoder_decoder = False
class _FakeRenderer:
tokenizer = object()
def __init__(self):
self.calls = []
def render_chat(
self,
all_messages,
chat_params,
tok_params,
prompt_extras=None,
):
self.calls.append(
{
"all_messages": all_messages,
"chat_params": chat_params,
"tok_params": tok_params,
"prompt_extras": prompt_extras,
}
)
return all_messages, [
{"prompt_token_ids": [index]} for index, _ in enumerate(all_messages)
]
@classmethod
def _make_handler(cls, renderer):
handler = object.__new__(EmbedIOProcessor)
handler.renderer = renderer
handler.model_config = cls._FakeModelConfig()
handler.chat_template = "template"
handler.chat_template_content_format = "auto"
handler.trust_request_chat_template = False
handler.enable_chunked_processing = False
return handler
@staticmethod
def _make_context(
request: (
EmbeddingChatRequest
| EmbeddingBatchChatRequest
| EmbeddingChatInputRequest
| EmbeddingBatchChatInputRequest
),
) -> PoolingServeContext[
EmbeddingChatRequest
| EmbeddingBatchChatRequest
| EmbeddingChatInputRequest
| EmbeddingBatchChatInputRequest
]:
return PoolingServeContext(
request=request,
pooling_params=PoolingParams(),
model_name="test",
request_id="embd-test",
)
def test_chat_template_kwargs_forwarded_for_batched_input_messages(self):
request = TypeAdapter(EmbeddingRequest).validate_python(
{
"model": "test",
"input": [
[{"role": "user", "content": "hello"}],
[{"role": "user", "content": "goodbye"}],
],
"add_generation_prompt": True,
"chat_template_kwargs": {"instruction": "Represent the query: "},
"mm_processor_kwargs": {"max_pixels": 1},
"cache_salt": "salt",
}
)
assert isinstance(request, EmbeddingBatchChatInputRequest)
renderer = self._FakeRenderer()
handler = self._make_handler(renderer)
ctx = self._make_context(request)
handler.pre_process_online(ctx)
assert ctx.engine_inputs == [
{"prompt_token_ids": [0]},
{"prompt_token_ids": [1]},
]
assert len(renderer.calls) == 1
call = renderer.calls[0]
assert call["all_messages"] == request.messages
assert call["prompt_extras"] == {
"mm_processor_kwargs": {"max_pixels": 1},
"cache_salt": "salt",
}
chat_template_kwargs = call["chat_params"].chat_template_kwargs
assert chat_template_kwargs["instruction"] == "Represent the query: "
assert chat_template_kwargs["add_generation_prompt"] is True
assert chat_template_kwargs["continue_final_message"] is False
assert "tools" not in chat_template_kwargs
assert chat_template_kwargs["tokenize"] is False
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
from types import SimpleNamespace
import pytest
import torch
@@ -10,10 +9,7 @@ import torch
from tests.models.utils import softmax
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
from vllm.entrypoints.pooling.scoring.typing import ScoringData
from vllm.platforms import current_platform
from vllm.renderers import TokenizeParams
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
PROMPT = "The chef prepared a delicious meal."
@@ -145,45 +141,6 @@ def test_max_tokens_per_doc(llm: LLM):
assert with_limit_tokens < no_limit_tokens
def test_token_type_ids_follow_post_tokenization():
processor = object.__new__(CrossEncoderIOProcessor)
processor.tokenizer = SimpleNamespace(truncation_side="right", pad_token_id=-1)
processor.renderer = SimpleNamespace(process_for_engine=lambda prompt, _: prompt)
processor.model_config = None
processor.get_score_prompt = lambda **_: (
"",
{
"prompt_token_ids": list(range(32)),
"token_type_ids": [0] * 16 + [1] * 16,
},
)
engine_inputs, pooling_params = processor._pre_process(
ScoringData(data_1=["query"], data_2=["document"]),
TokenizeParams(
max_total_tokens=None,
truncate_prompt_tokens=16,
truncation_side="left",
),
PoolingParams(task="classify", extra_kwargs={"cache_salt": "salt"}),
)
assert engine_inputs[0]["prompt_token_ids"] == list(range(16, 32))
assert pooling_params[0].extra_kwargs == {
"cache_salt": "salt",
"compressed_token_type_ids": 0,
}
engine_inputs, pooling_params = processor._pre_process(
ScoringData(data_1=["query"], data_2=["document"]),
TokenizeParams(max_total_tokens=None, pad_prompt_tokens=40),
PoolingParams(task="classify"),
)
assert engine_inputs[0]["prompt_token_ids"] == list(range(32)) + [-1] * 8
assert pooling_params[0].extra_kwargs == {"compressed_token_type_ids": 16}
def test_pooling_params(llm: LLM):
def get_outputs(use_activation):
outputs = llm.score(
@@ -161,7 +161,9 @@ def _make_pooling_serving(lora_name: str) -> _ConcretePoolingServing:
return serving
def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
def _make_pooling_ctx(
model_name: str, serving: PoolingBaseServing
) -> PoolingServeContext:
mock_request = MagicMock()
mock_request.model = model_name
return PoolingServeContext(
@@ -169,6 +171,9 @@ def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
model_name=MODEL_NAME,
request_id="test-id",
pooling_params=PoolingParams(),
lora_request=serving._maybe_get_adapters(mock_request),
priorities=0,
prompt_extras=None,
)
@@ -176,9 +181,7 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
"""LoRA adapter name must populate ctx.lora_request without raising."""
lora_name = "bot-embed-lora"
serving = _make_pooling_serving(lora_name)
ctx = _make_pooling_ctx(lora_name)
ctx.lora_request = serving._maybe_get_adapters(ctx.request)
ctx = _make_pooling_ctx(lora_name, serving)
assert ctx.lora_request is not None
assert ctx.lora_request.lora_name == lora_name
@@ -187,7 +190,6 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
def test_pooling_maybe_get_adapters_unknown_model_raises():
"""An unrecognised model name must still raise VLLMNotFoundError."""
serving = _make_pooling_serving("some-lora")
ctx = _make_pooling_ctx("unknown-model")
with pytest.raises(VLLMNotFoundError):
serving._maybe_get_adapters(ctx.request)
_make_pooling_ctx("unknown-model", serving)
@@ -6,6 +6,7 @@ max_concurrency: 100
server_args: >-
--enforce-eager
--max-model-len 4096
--max-num-batched-tokens 32768
--safetensors-load-strategy prefetch
--moe-backend flashinfer_cutlass
--prefill-context-parallel-size 4
@@ -6,6 +6,7 @@ max_concurrency: 100
server_args: >-
--enforce-eager
--max-model-len 4096
--max-num-batched-tokens 32768
--safetensors-load-strategy prefetch
--moe-backend flashinfer_cutlass
--tensor-parallel-size 2
@@ -0,0 +1,9 @@
model_name: "poolside/Laguna-XS.2-NVFP4"
accuracy_threshold: 0.86
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
server_args: >-
--enforce-eager
--max-model-len 4096
--trust-remote-code
@@ -2,4 +2,7 @@ model_name: "google/gemma-4-E4B-it-qat-mobile-ct"
accuracy_threshold: 0.50
num_questions: 1319
num_fewshot: 5
server_args: "--enforce-eager --max-model-len 4096"
server_args: >-
--enforce-eager
--max-model-len 4096
--speculative-config '{"method":"mtp","model":"google/gemma-4-E4B-it-assistant","num_speculative_tokens":4}'
@@ -3,3 +3,4 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
Qwen1.5-MoE-W4A16-CT.yaml
DeepSeek-V2-Lite-Instruct-FP8.yaml
Qwen3-30B-A3B-NVFP4.yaml
Laguna-XS.2-NVFP4.yaml
@@ -546,8 +546,11 @@ def test_flash_attn_accepts_handled_fp8_variants(
):
"""FlashAttentionBackend must accept the two fp8 dtypes it can actually
handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""
import vllm.v1.attention.backends.flash_attn as fa_mod
import vllm.v1.attention.backends.fa_utils as fa_utils_mod
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
monkeypatch.setattr(fa_mod.current_platform, "is_xpu", lambda: True)
# The fp8 decision is made in fa_utils, using its own current_platform
# binding, so patch is_xpu there (not on flash_attn's) to stay robust to
# import order across earlier tests that patch vllm.platforms.current_platform.
monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True)
assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
@@ -11,6 +11,9 @@ from vllm._custom_ops import (
scaled_fp8_quant,
)
from vllm.platforms import current_platform
from vllm.v1.attention.ops.triton_merge_attn_states import (
mask_empty_context,
)
from vllm.v1.attention.ops.triton_merge_attn_states import (
merge_attn_states as merge_attn_states_triton,
)
@@ -73,6 +76,59 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]
all_case_info: list[tuple] = []
def test_mask_empty_context() -> None:
query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
query_start_loc = torch.cat(
(torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
).cuda()
context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
context_start_loc = torch.cat(
(torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
).cuda()
num_heads, num_tokens, head_dim = 4, 165, 16
lse = torch.randn(num_heads, num_tokens, device="cuda")
output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
# Empty-context rows carry undefined (possibly non-finite) attention output.
output[33:164] = float("nan")
expected_lse = lse.clone()
expected_lse[:, 33:164] = float("-inf")
expected_output = output.clone()
expected_output[33:164] = 0.0
mask_empty_context(lse, output, query_start_loc, context_start_loc)
torch.testing.assert_close(lse, expected_lse)
torch.testing.assert_close(output, expected_output)
@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
"""When a token is empty on both sides (both LSE -inf), the 0/0 softmax
scales must not surface as NaN in the merged output."""
num_tokens, num_heads, head_size = 6, 8, 128
prefix_output = torch.zeros(
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
)
prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")
suffix_output = torch.zeros(
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
)
suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
# Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
# their outputs and set both LSEs to -inf).
empty = slice(2, 4)
prefix_lse[:, empty] = float("-inf")
suffix_lse[:, empty] = float("-inf")
output = torch.empty_like(prefix_output)
merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
assert not output.isnan().any()
def generate_markdown_table():
global all_case_info
table_header = (
@@ -40,10 +40,11 @@ BATCH_SIZE = 128
CONTEXT_LEN = 8192
PAGE_SIZE = 1
# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8
# KV-cache -> fp8_e4m3 kv.
EXPECTED_Q_DTYPE = torch.bfloat16
EXPECTED_KV_DTYPE = torch.float8_e4m3fn
# On the fp8 KV-cache path the builder forwards the platform fp8 dtype (aiter
# dtypes.fp8) for both q and kv. Mirror it via current_platform.fp8_dtype()
# instead of hardcoding a literal (see #47276).
EXPECTED_Q_DTYPE = current_platform.fp8_dtype()
EXPECTED_KV_DTYPE = current_platform.fp8_dtype()
# The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
# is excluded: it holds raw device pointers, never equal across allocations.
+15
View File
@@ -0,0 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
@pytest.fixture(autouse=True)
def reset_default_torch_device():
"""Several kernel tests call torch.set_default_device without restoring
it, which poisons subsequent tests in the same pytest run (e.g. CPU
tensors silently created on CUDA). Restore the factory default after
every test.
"""
yield
torch.set_default_device(None)
+14 -19
View File
@@ -33,12 +33,10 @@ from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402
@pytest.mark.parametrize("num_seqs", [1, 5, 257])
@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
seq_lens = torch.randint(
1,
130,
(num_seqs,),
dtype=torch.int32,
)
rng_cpu = torch.Generator("cpu").manual_seed(1234)
rng = torch.Generator("cuda").manual_seed(2345)
seq_lens = torch.randint(1, 130, (num_seqs,), dtype=torch.int32, generator=rng_cpu)
cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32)
cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0)
total_tokens = int(cu_seqlens[-1].item())
@@ -56,8 +54,9 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_k_dim,
device="cuda",
dtype=dtype,
generator=rng,
)
k = torch.randn_like(q)
k = torch.randn_like(q, generator=rng)
v = torch.randn(
1,
total_tokens,
@@ -65,29 +64,24 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_v_dim,
device="cuda",
dtype=dtype,
generator=rng,
)
q = F.normalize(q.float(), p=2, dim=-1).to(dtype)
k = F.normalize(k.float(), p=2, dim=-1).to(dtype)
a = torch.randn(
1,
total_tokens,
num_v_heads,
device="cuda",
dtype=dtype,
1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
)
b = torch.randn(
1,
total_tokens,
num_v_heads,
device="cuda",
dtype=dtype,
1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
)
# Match upstream FLA GatedDeltaNet synthetic initialization:
# https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16)
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(
0, 16, generator=rng
)
A_log = torch.log(A)
dt = torch.exp(
torch.rand(num_v_heads, device="cuda", dtype=torch.float32)
torch.rand(num_v_heads, device="cuda", dtype=torch.float32, generator=rng)
* (math.log(0.1) - math.log(0.001))
+ math.log(0.001)
)
@@ -105,6 +99,7 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_k_dim,
device="cuda",
dtype=state_dtype,
generator=rng,
)
* 0.05
)
@@ -6,6 +6,7 @@ import pytest
import torch
import torch.nn.functional as F
import vllm._custom_ops as ops
from vllm.model_executor.layers.fused_moe.config import (
RoutingMethodType,
get_routing_method_type,
@@ -231,3 +232,119 @@ def test_dsv4_fast_topk(
atol=2e-5,
rtol=2e-5,
)
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="This test is skipped on non-CUDA platform.",
)
@pytest.mark.parametrize("use_hash", [False, True])
@pytest.mark.parametrize("use_bias", [False, True])
@pytest.mark.parametrize("use_padding_mask", [False, True])
@pytest.mark.parametrize("pad_with_nan", [False, True])
@pytest.mark.parametrize("num_experts", [128, 256, 384])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
def test_fused_topk_softplus_sqrt_padding(
use_hash: bool,
use_bias: bool,
use_padding_mask: bool,
pad_with_nan: bool,
num_experts: int,
dtype: torch.dtype,
):
"""Verify explicit padding and NaN-padded rows do not affect real rows."""
torch.manual_seed(0)
num_tokens = 8
topk = 6
indices_dtype = torch.int32
gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
padding_rows = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
padding_rows[1::2] = True
if pad_with_nan:
gating_output[padding_rows] = float("nan")
is_padding = padding_rows if use_padding_mask else None
# A negative correction bias makes explicit pad rows look selectable unless
# the kernel uses the is_padding guard.
e_score_correction_bias = None
if use_bias:
e_score_correction_bias = (
-torch.rand((num_experts,), dtype=torch.float32, device="cuda") - 1.0
)
input_ids = None
hash_indices_table = None
if use_hash:
vocab_size = 64
hash_indices_table = torch.stack(
[torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
).to(device="cuda", dtype=indices_dtype)
input_ids = torch.randint(
0, vocab_size, (num_tokens,), dtype=indices_dtype, device="cuda"
)
topk_weights = torch.empty(num_tokens, topk, dtype=torch.float32, device="cuda")
topk_ids = torch.empty(num_tokens, topk, dtype=indices_dtype, device="cuda")
token_expert_indices = torch.empty(
num_tokens, topk, dtype=torch.int32, device="cuda"
)
ops.topk_hash_softplus_sqrt(
topk_weights,
topk_ids,
token_expert_indices,
gating_output,
renormalize=True,
routed_scaling_factor=1.0,
e_score_correction_bias=e_score_correction_bias,
input_tokens=input_ids,
hash_indices_table=hash_indices_table,
is_padding=is_padding,
)
if use_padding_mask:
pad_ids = topk_ids[padding_rows]
pad_weights = topk_weights[padding_rows]
assert torch.equal(pad_ids, torch.full_like(pad_ids, -1)), (
f"Explicit pad rows should contain only -1 ids, got {pad_ids.tolist()}"
)
assert (pad_weights == 0).all(), (
"Explicit pad rows should have all-zero weights, "
f"got {pad_weights.tolist()}"
)
if pad_with_nan:
nan_pad_weights = topk_weights[padding_rows]
assert torch.isfinite(nan_pad_weights).all(), (
f"NaN-padded rows have non-finite weights, got {nan_pad_weights.tolist()}"
)
assert (nan_pad_weights == 0).all(), (
"NaN-padded rows should have all-zero weights, "
f"got {nan_pad_weights.tolist()}"
)
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
gating_output=gating_output,
topk=topk,
renormalize=True,
routed_scaling_factor=1.0,
e_score_correction_bias=e_score_correction_bias,
input_ids=input_ids,
hash_indices_table=hash_indices_table,
)
rows_to_compare = torch.ones(num_tokens, dtype=torch.bool, device="cuda")
if use_padding_mask or pad_with_nan:
rows_to_compare = ~padding_rows
sorted_ref_ids, idx_ref = topk_ids_ref[rows_to_compare].sort(dim=-1)
sorted_ids, idx_ops = topk_ids[rows_to_compare].sort(dim=-1)
torch.testing.assert_close(
sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0
)
sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref)
sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops)
torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
+5 -3
View File
@@ -21,9 +21,9 @@ def test_gather_cache_oob():
seq_starts causes the block_table offset to read out of bounds.
"""
batch_size = 1
block_size = 64
entry_size = 128
# The kernel only supports the MLA entry sizes.
entry_size = 576
block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
@@ -34,6 +34,7 @@ def test_gather_cache_oob():
seq_len = 65
cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
token_to_seq = torch.zeros(seq_len, dtype=torch.int32, device="cuda")
# src_cache: [num_blocks, block_size, entry_size]
num_blocks = 5
@@ -51,7 +52,8 @@ def test_gather_cache_oob():
dst,
block_table,
cu_seq_lens,
batch_size,
token_to_seq,
seq_len,
"auto", # kv_cache_dtype
scale,
seq_starts,
+21
View File
@@ -13,6 +13,7 @@ These tests cover:
"""
import math
from types import SimpleNamespace
import pytest
import torch
@@ -27,6 +28,7 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
_launch_two_stage_sparse_attn_compressor,
)
from vllm.models.deepseek_v4.compressor import _get_c128_boundary
from vllm.platforms import current_platform
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
@@ -58,6 +60,25 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
return x_fp8, scales
@pytest.mark.parametrize(
("starts", "query_start_loc", "expected"),
[
([0], [0, 127], False),
([0], [0, 128], True),
([127], [0, 1], True),
([128], [0, 127], False),
([1, 255], [0, 1, 2], True),
(None, [0, 1], None),
],
)
def test_get_c128_boundary(starts, query_start_loc, expected):
metadata = SimpleNamespace(
_num_computed_tokens_cpu=None if starts is None else torch.tensor(starts),
query_start_loc_cpu=torch.tensor(query_start_loc),
)
assert _get_c128_boundary(metadata) is expected
# ── Test A: DeepseekV4 Attention path ──────────────────────────────────────────────
+1 -1
View File
@@ -264,7 +264,7 @@ def test_block_mask_direct_vs_slow_path():
device = torch.device("cuda")
vllm_config = create_vllm_config(
model_name="meta-llama/Meta-Llama-3-8B", block_size=16, max_model_len=1024
model_name="Qwen/Qwen2.5-1.5B-Instruct", block_size=16, max_model_len=1024
)
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
+13 -4
View File
@@ -726,14 +726,19 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
This catches stride/layout bugs that only manifest when the einsum
kernel actually consumes the quantized activations.
"""
from deep_gemm.utils.math import ceil_div
from vllm.utils.deep_gemm import (
fp8_einsum,
is_deep_gemm_supported,
per_block_cast_to_fp8,
transform_sf_into_required_layout,
)
if not is_deep_gemm_supported():
pytest.skip("DeepGEMM not supported on this platform")
def ceil_div(a: int, b: int) -> int:
return (a + b - 1) // b
heads_per_group = num_heads // n_groups
d = heads_per_group * HEAD_DIM
o_lora_rank = 1024
@@ -809,8 +814,12 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
# -- Checks --
# Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
# come from fp32 ordering and UE8M0 boundary shifts only.
# Use relative diff (same metric as test_fp8_einsum.py).
from deep_gemm.testing import calc_diff
# Use relative diff (same metric as deep_gemm.testing.calc_diff).
def calc_diff(x, y):
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
z_diff = calc_diff(z_fused, z_ref)
assert z_diff < 0.01, (
@@ -84,6 +84,22 @@ def norm_rope_ref(x, weight, positions, cos_sin_cache, eps):
return roped
def assert_fp8_cache_close(kv_cache, expected_kv_cache):
"""Compare two e4m3 caches allowing 1 ulp.
On CUDA the fused kernel quantizes K from its fp32 intermediate, while the
reshape_and_cache_flash reference quantizes the bf16-materialized value, so
rounding-boundary values may differ by one e4m3 code.
"""
byte_diff = (kv_cache.int() - expected_kv_cache.int()).abs()
got = kv_cache.view(torch.float8_e4m3fn).float()
exp = expected_kv_cache.view(torch.float8_e4m3fn).float()
ok = (byte_diff <= 1) | ((got == 0) & (exp == 0))
assert bool(ok.all()), (
f"fp8 cache differs by more than 1 ulp in {int((~ok).sum())} elements"
)
# ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
@@ -265,7 +281,7 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
scale,
scale,
)
torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
assert_fp8_cache_close(kv_cache, expected_kv_cache)
else:
for t in range(num_tokens):
s = slot_mapping[t].item()
@@ -383,7 +399,7 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
scale,
scale,
)
torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
assert_fp8_cache_close(kv_cache, expected_kv_cache)
else:
k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)
@@ -41,11 +41,12 @@ def test_fused_recurrent_packed_decode_matches_reference(
A_log = torch.randn((HV,), device=device, dtype=dtype)
dt_bias = torch.randn((HV,), device=device, dtype=dtype)
# Continuous batching indices (include PAD_SLOT_ID=-1 cases).
ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32)
# Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is
# reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1.
ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32)
ssm_state_indices[-3:] = -1
state0 = torch.randn((B, HV, V, K), device=device, dtype=dtype)
state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype)
state_ref = state0.clone()
state_packed = state0.clone()
@@ -94,5 +95,8 @@ def test_fused_recurrent_packed_decode_matches_reference(
atol = 2e-2 if dtype != torch.float32 else 1e-4
rtol = 1e-2 if dtype != torch.float32 else 1e-4
torch.testing.assert_close(out_packed, out_ref, rtol=rtol, atol=atol)
# Output rows for PAD_SLOT_ID entries are never written (uninitialized in
# both paths), so compare only the valid rows.
valid = ssm_state_indices > 0
torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol)
torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)
@@ -58,10 +58,12 @@ def test_fused_sigmoid_gating_delta_rule_update_non_spec(
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
# Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
# state indices start at 1.
ssm_state = torch.rand(
total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
)
state_indices = torch.randperm(total_entries, dtype=torch.int32)[:num_tokens]
state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[:num_tokens]
cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32)
beta = b.sigmoid()
@@ -144,13 +146,14 @@ def test_fused_sigmoid_gating_delta_rule_update_spec(
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
# Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
# state indices start at 1.
ssm_state = torch.rand(
total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
)
state_indices = torch.randperm(
total_entries,
dtype=torch.int32,
)[:num_tokens].view(num_reqs, num_speculative_tokens + 1)
state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[
:num_tokens
].view(num_reqs, num_speculative_tokens + 1)
num_accepted_tokens = torch.randint(
1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32
)
@@ -0,0 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.model_executor.model_loader.mtp_validation import (
disable_mtp_completeness_check,
is_mtp_completeness_check_enabled,
)
def test_disable_mtp_completeness_check_is_scoped():
assert is_mtp_completeness_check_enabled()
with pytest.raises(RuntimeError), disable_mtp_completeness_check():
assert not is_mtp_completeness_check_enabled()
raise RuntimeError
assert is_mtp_completeness_check_enabled()
+97
View File
@@ -6,15 +6,43 @@ import time
import pytest
import vllm.config.multimodal as multimodal_config_module
from vllm.config.multimodal import MultiModalConfig
from vllm.multimodal.gpu_ipc_memory import (
MultiModalGPUMemoryPool,
get_mm_gpu_ipc_pool,
maybe_init_mm_gpu_ipc_pool,
reserve_mm_ipc_gpu_memory,
set_mm_gpu_ipc_pool,
)
from vllm.multimodal.video import (
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
PYNVVIDEOCODEC_MAX_RETAINED_DECODERS,
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
from vllm.utils.mem_constants import GiB_bytes
def _mm_config(
*,
mm_ipc_gpu_memory_gb: float = 0,
video_backend: str | None = None,
) -> MultiModalConfig:
video_kwargs = {} if video_backend is None else {"video_backend": video_backend}
return MultiModalConfig(
mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb,
media_io_kwargs={"video": video_kwargs} if video_kwargs else {},
)
def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int:
return api_process_count * (
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS
+ PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES
)
def test_acquire_release_accounting():
pool = MultiModalGPUMemoryPool(total_bytes=100)
assert pool.available_bytes == 100
@@ -143,3 +171,72 @@ def test_global_pool_splits_budget_across_api_processes():
def test_global_pool_rejects_invalid_api_process_count():
with pytest.raises(ValueError):
maybe_init_mm_gpu_ipc_pool(2, api_process_count=0)
@pytest.mark.parametrize("video_backend", [None, "opencv"])
def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
monkeypatch: pytest.MonkeyPatch,
video_backend: str | None,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
mm_config = _mm_config(
mm_ipc_gpu_memory_gb=0.25,
video_backend=video_backend,
)
assert reserve_mm_ipc_gpu_memory(GiB_bytes, mm_config) == int(0.75 * GiB_bytes)
def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
mm_config = _mm_config(
mm_ipc_gpu_memory_gb=0.25,
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
)
available_bytes = 4 * GiB_bytes
assert reserve_mm_ipc_gpu_memory(available_bytes, mm_config) == (
available_bytes - int(0.25 * GiB_bytes) - _pynvvideocodec_decoder_budget()
)
def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
available_bytes = 4 * GiB_bytes
assert reserve_mm_ipc_gpu_memory(available_bytes, _mm_config()) == (
available_bytes - _pynvvideocodec_decoder_budget()
)
def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
multimodal_config_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
available_bytes = 8 * GiB_bytes
assert reserve_mm_ipc_gpu_memory(
available_bytes,
_mm_config(),
api_process_count=3,
) == available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3)
+42
View File
@@ -920,3 +920,45 @@ class TestDelegatingParserLargeDelta:
assert output.tool_calls[0]["name"] == "get_weather"
args = json.loads(output.tool_calls[0]["arguments"])
assert args == {"location": "Berlin"}
@pytest.mark.parametrize(
"chunk_size",
[1, 2, 3, 5, None],
ids=lambda c: f"chunk={c}",
)
def test_eos_not_leaked_when_reasoning_never_ends(self, chunk_size):
"""EOS must not leak into reasoning_content when the model never
emits </think> (generation ends while still in REASONING state)."""
eos_text = "<end▁of▁sentence>"
eos_id = 128801
vocab = {
**_DSV4_FULL_VOCAB,
eos_text: eos_id,
}
reasoning_text = "Good morning! How can I help you today?"
tokens: list[tuple[int, str]] = []
tid = 100
for word in reasoning_text.split(" "):
prefix = " " if tokens else ""
tokens.append((tid, prefix + word))
tid += 1
tokens.append((eos_id, eos_text))
tokenizer = MockTokenizer(vocab=vocab, tokens=tokens)
parser = _DeepSeekV4Delegating(
tokenizer,
chat_template_kwargs={"thinking": True},
)
deltas = replay_streaming(
parser,
tokens,
chunk_size=chunk_size,
finished_on_last=True,
)
output = collect_output(deltas)
assert reasoning_text in output.reasoning
assert eos_text not in output.reasoning
assert output.content == ""
assert output.tool_calls == []
+29 -7
View File
@@ -1624,19 +1624,41 @@ class TestDropSpecialTokens:
assert delta is not None
assert "<bos>" in delta.reasoning
def test_drops_suppressed_with_skip_tool_parsing(self):
"""When skip_tool_parsing is active, drop tokens are preserved
as content so a later tool-call pass can see them."""
def test_drops_applied_with_skip_tool_parsing(self):
"""Drop tokens are always dropped, even with skip_tool_parsing.
DROP_TERMINALs have no transitions by construction, so no parser
pass can use them."""
for initial_state in (ParserState.REASONING, ParserState.CONTENT):
engine = _make_engine(
vocab=_DROP_VOCAB,
special_tokens=list(_DROP_VOCAB.keys()),
)
engine._engine.skip_tool_parsing = True
engine._engine.reset(initial_state=initial_state)
events = engine._engine.feed("hello<bos>world", [72, 204, 73])
delta = engine._events_to_delta(events)
assert delta is not None
output = (delta.reasoning or "") + (delta.content or "")
assert "<bos>" not in output, f"<bos> leaked in state {initial_state}"
def test_transitions_unaffected_by_drop_in_reasoning_with_skip_tool_parsing(self):
"""With skip_tool_parsing in REASONING state, drop tokens are
removed but configured terminals still fire their transitions."""
engine = _make_engine(
vocab=_DROP_VOCAB,
special_tokens=list(_DROP_VOCAB.keys()),
)
engine._engine.skip_tool_parsing = True
engine._engine.reset()
events = engine._engine.feed("hello<bos>world", [72, 204, 73])
delta = engine._events_to_delta(events)
assert delta is not None
assert "<bos>" in delta.reasoning
events = engine._engine.feed("thought<bos></think>answer", [72, 204, 201, 73])
types = [e.type for e in events]
assert EventType.REASONING_CHUNK in types
assert EventType.REASONING_END in types
assert EventType.TEXT_CHUNK in types
reasoning_text = "".join(
e.value for e in events if e.type == EventType.REASONING_CHUNK
)
assert "<bos>" not in reasoning_text
def test_drops_in_tool_args_state(self):
"""Drop tokens in TOOL_ARGS state are silently discarded."""
@@ -50,7 +50,7 @@ def test_oot_registration_embedding(
m.setenv("VLLM_PLUGINS", "register_dummy_model")
prompts = ["Hello, my name is", "The text does not matter"]
llm = LLM(
model=dummy_gemma2_embedding_path, load_format="dummy", max_model_len=2048
model=dummy_gemma2_embedding_path, load_format="dummy", max_model_len=512
)
outputs = llm.embed(prompts)
+79
View File
@@ -5,13 +5,17 @@
Run `pytest tests/quantization/test_auto_gptq.py -v -s`.
"""
from types import SimpleNamespace
import pytest
import torch
from tests.quantization.utils import is_quant_method_supported
from vllm.model_executor.layers.fused_moe import RoutedExperts
from vllm.model_executor.layers.quantization.auto_gptq import (
AutoGPTQConfig,
AutoGPTQLinearMethod,
AutoGPTQMoEMethod,
)
PROMPT = "On the surface of Mars, we found"
@@ -54,3 +58,78 @@ def test_auto_gptq_quantization_method(vllm_runner, model_id: str, monkeypatch):
def test_auto_gptq_config_get_name():
"""Test that AutoGPTQConfig.get_name() returns 'auto_gptq'."""
assert AutoGPTQConfig.get_name() == "auto_gptq"
def test_auto_gptq_moe_creates_zero_initialized_expert_biases():
method = object.__new__(AutoGPTQMoEMethod)
method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
method.input_dtype = None
method.experts_cls = None
layer = torch.nn.Module()
method.create_weights(
layer=layer,
num_experts=2,
hidden_size=8,
intermediate_size_per_partition=4,
params_dtype=torch.float16,
intermediate_size_full=4,
weight_loader=lambda *args, **kwargs: None,
)
assert layer.w13_bias.shape == (2, 8)
assert layer.w2_bias.shape == (2, 8)
assert torch.count_nonzero(layer.w13_bias) == 0
assert torch.count_nonzero(layer.w2_bias) == 0
def test_routed_experts_loads_per_expert_biases():
class Loader:
quant_config = None
quant_method = object()
moe_config = SimpleNamespace(
is_act_and_mul=True,
tp_rank=0,
moe_parallel_config=SimpleNamespace(tp_size=1),
)
_get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim)
_narrow_expert_data_for_padding = staticmethod(
RoutedExperts._narrow_expert_data_for_padding
)
_load_w13 = RoutedExperts._load_w13
_loaded_expert_biases = set()
@staticmethod
def _map_global_expert_id_to_local_expert_id(expert_id):
return expert_id
loader = Loader()
w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False)
w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False)
for shard_id, loaded in (
("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])),
("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])),
):
assert RoutedExperts.weight_loader(
loader,
w13_bias,
loaded,
weight_name="model.layers.0.mlp.experts.w13_bias",
shard_id=shard_id,
expert_id=0,
return_success=True,
)
assert RoutedExperts.weight_loader(
loader,
w2_bias,
torch.tensor([9.0, 10.0, 11.0, 12.0]),
weight_name="model.layers.0.mlp.experts.w2_bias",
shard_id="w2",
expert_id=0,
return_success=True,
)
assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8))
assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4))
assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"}
+181 -32
View File
@@ -5,45 +5,194 @@ from types import SimpleNamespace
import pytest
import torch
from compressed_tensors.quantization import (
ActivationOrdering,
QuantizationArgs,
QuantizationStrategy,
QuantizationType,
)
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method
from vllm.platforms import current_platform
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
WNA16MoEBackend,
_backend_incompatibility_reason,
_convert_moe_wna16_humming_tensors,
convert_to_wna16_moe_kernel_format,
map_wna16_backend,
)
from vllm.model_executor.layers.quantization import moe_wna16
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
from vllm.model_executor.layers.quantization.moe_wna16 import (
MoeWNA16Config,
MoeWNA16Method,
)
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
def test_moe_wna16_apply_passes_layer_activation(monkeypatch):
captured_kwargs = {}
def test_map_wna16_backend_supports_triton():
assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON
def fake_fused_experts(*args, **kwargs):
captured_kwargs.update(kwargs)
return torch.empty(1, 2)
monkeypatch.setattr(
"vllm.model_executor.layers.fused_moe.fused_experts",
fake_fused_experts,
@pytest.mark.parametrize(
("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"),
[
(
WNA16MoEBackend.TRITON,
AutoAWQConfig(4, 128, True, False),
True,
False,
"AutoAWQ weight layout",
),
(
WNA16MoEBackend.TRITON,
AutoGPTQConfig(4, 128, True, True, False, {}, {}),
False,
False,
"activation ordering",
),
(
WNA16MoEBackend.TRITON,
QuantizationArgs(
num_bits=4,
type=QuantizationType.INT,
strategy=QuantizationStrategy.GROUP,
symmetric=True,
dynamic=False,
group_size=128,
actorder=ActivationOrdering.GROUP,
),
False,
False,
"activation ordering",
),
(
WNA16MoEBackend.TRITON,
AutoGPTQConfig(4, 128, False, True, False, {}, {}),
False,
True,
"bias",
),
(
WNA16MoEBackend.MARLIN,
MoeWNA16Config(
linear_quant_method="gptq",
weight_bits=4,
group_size=128,
has_zp=False,
lm_head_quantized=False,
modules_to_not_convert=None,
full_config={},
),
False,
False,
"MoeWNA16 checkpoint layout",
),
],
)
def test_wna16_oracle_rejects_incompatible_quant_structures(
backend, quant_config, may_have_zp, may_have_bias, expected
):
reason = _backend_incompatibility_reason(
backend=backend,
quant_config=quant_config,
may_have_zp=may_have_zp,
may_have_bias=may_have_bias,
)
assert reason is not None
assert expected in reason
def test_compressed_tensors_weights_are_transposed_for_triton():
quant_config = QuantizationArgs(
num_bits=4,
type=QuantizationType.INT,
strategy=QuantizationStrategy.GROUP,
symmetric=True,
dynamic=False,
group_size=32,
)
w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8)
w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6)
w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8)
w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6)
converted = convert_to_wna16_moe_kernel_format(
backend=WNA16MoEBackend.TRITON,
layer=torch.nn.Module(),
quant_config=quant_config,
input_dtype=None,
w13=w13,
w2=w2,
w13_scale=w13_scale,
w2_scale=w2_scale,
)
assert converted is not None
assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8))
assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8))
assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous())
assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous())
def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):
method = object.__new__(MoeWNA16Method)
method.experts_cls = object
method.wna16_backend = WNA16MoEBackend.HUMMING
method.moe = object()
quant_config = object()
method.get_fused_moe_quant_config = lambda layer: quant_config
layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None))
captured = {}
kernel = object()
def fake_make_wna16_moe_kernel(**kwargs):
captured.update(kwargs)
return kernel
monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel)
method._setup_kernel(layer)
assert method.moe_kernel is kernel
assert captured["backend"] == WNA16MoEBackend.HUMMING
assert captured["layer"] is layer
def test_moe_wna16_humming_adapter_repacks_uint8_tensors():
qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8)
scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4)
qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2)
converted = _convert_moe_wna16_humming_tensors(
{"qweight": qweight, "scales": scales, "qzeros": qzeros},
has_zero_point=True,
)
assert torch.equal(converted["weight"], qweight.view(torch.int32))
assert converted["weight"].shape == (1, 4, 2)
assert torch.equal(converted["weight_scale"], scales)
expected_qzeros = (
qzeros.transpose(-1, -2)
.contiguous()
.view(torch.int32)
.transpose(-1, -2)
.contiguous()
)
assert torch.equal(converted["zero_point"], expected_qzeros)
assert converted["zero_point"].shape == (1, 2, 2)
def test_moe_wna16_uses_humming_quant_config(monkeypatch):
from vllm.model_executor.layers.quantization.utils import humming_utils
method = object.__new__(MoeWNA16Method)
method.moe = SimpleNamespace(disable_inplace=False)
method.moe_quant_config = object()
layer = SimpleNamespace(
w13_qweight=torch.empty(1, 2),
w2_qweight=torch.empty(1, 2),
activation=MoEActivation.GELU_TANH,
apply_router_weight_on_input=False,
global_num_experts=1,
expert_map=None,
method.wna16_backend = WNA16MoEBackend.HUMMING
layer = object()
quant_config = object()
monkeypatch.setattr(
humming_utils,
"get_humming_moe_quant_config",
lambda actual_layer: quant_config if actual_layer is layer else None,
)
output = method.apply(
layer,
x=torch.empty(1, 2),
topk_weights=torch.empty(1, 1),
topk_ids=torch.empty(1, 1, dtype=torch.int32),
shared_experts=None,
shared_experts_input=None,
)
assert output.shape == (1, 2)
assert captured_kwargs["activation"] is MoEActivation.GELU_TANH
assert method.get_fused_moe_quant_config(layer) is quant_config
@@ -1,13 +1,55 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import torch
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
from vllm.model_executor.layers.quantization.utils import flashinfer_fp4_moe
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
prepare_nvfp4_moe_layer_for_fi_or_cutlass,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
align_trtllm_fp4_moe_hidden_dim_for_fi,
)
def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):
monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x)
num_experts = 3
layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False))
w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
w2_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
weight_scale = torch.ones(num_experts)
outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass(
backend=NvFp4MoeBackend.FLASHINFER_CUTLASS,
layer=layer,
w13=w13,
w13_scale=w13_scale,
w13_scale_2=weight_scale,
a13_scale=torch.tensor([1.0, 2.0, 3.0]),
w2=w2,
w2_scale=w2_scale,
w2_scale_2=weight_scale,
a2_scale=torch.tensor([4.0, 5.0, 6.0]),
is_act_and_mul=False,
)
a13_scale, a2_scale = outputs[3], outputs[7]
torch.testing.assert_close(a13_scale, torch.full((num_experts,), 3.0))
torch.testing.assert_close(a2_scale, torch.full((num_experts,), 6.0))
distinct_values = torch.arange(num_experts, dtype=torch.float32)
a13_scale.copy_(distinct_values)
a2_scale.copy_(distinct_values)
torch.testing.assert_close(a13_scale, distinct_values)
torch.testing.assert_close(a2_scale, distinct_values)
def test_align_trtllm_fp4_moe_hidden_dim_noop():
w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)
@@ -46,9 +46,12 @@ BACKENDS_TO_TEST = [
DEVICE_TYPE = current_platform.device_type
# Use the platform's preferred FP8 type so the stored cache matches what the
# backends reinterpret at runtime. On ROCm gfx94x this is e4m3fnuz, not e4m3fn;
# storing e4m3fn bytes there would be re-read as fnuz and produce NaNs.
FP8_KV_CACHE_DTYPES = {
"fp8": torch.float8_e4m3fn,
"fp8_e4m3": torch.float8_e4m3fn,
"fp8": current_platform.fp8_dtype(),
"fp8_e4m3": current_platform.fp8_dtype(),
}
# Remove flashinfer from the list if it's not available
+168 -20
View File
@@ -9,6 +9,7 @@ import torch
from vllm.v1.core.kv_cache_utils import (
_get_kv_cache_config_packed,
_get_kv_cache_groups_uniform_groups,
get_kv_cache_config_from_groups,
)
from vllm.v1.kv_cache_interface import (
@@ -16,6 +17,7 @@ from vllm.v1.kv_cache_interface import (
KVCacheGroupSpec,
KVCacheTensor,
MLAAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
UniformTypeKVCacheSpecs,
)
@@ -109,7 +111,132 @@ def _page_sizes_by_layer(
return page_sizes
def _packing_by_layer(
tensors: list[KVCacheTensor],
) -> dict[str, tuple[int, int]]:
return {
layer_name: (tensor.offset, tensor.block_stride)
for tensor in tensors
for layer_name in tensor.shared_by
}
def _make_views(
groups: list[KVCacheGroupSpec],
num_blocks: int,
tensors: list[KVCacheTensor],
) -> dict[str, torch.Tensor]:
page_sizes = _page_sizes_by_layer(groups)
packing = _packing_by_layer(tensors)
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
return {
layer_name: torch.as_strided(
backing,
size=(num_blocks, page_size),
stride=(packing[layer_name][1], 1),
storage_offset=packing[layer_name][0],
)
for layer_name, page_size in page_sizes.items()
}
def _make_page_group(prefix: str, page_sizes: list[int]) -> KVCacheGroupSpec:
specs = {
f"{prefix}.{i}": MagicMock(page_size_bytes=page_size)
for i, page_size in enumerate(page_sizes)
}
return KVCacheGroupSpec(
layer_names=list(specs),
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=specs),
)
class TestInterleavedPacking:
def test_compact_cache_overlays_fp32_state_group(self):
full_specs = {}
state_specs = {}
for i in range(2):
full_specs[f"mla.{i}"] = MLAAttentionSpec(
block_size=256,
num_kv_heads=1,
head_size=512,
dtype=torch.uint8,
page_size_padded=32768,
indexes_kv_by_block_stride=True,
compress_ratio=4,
)
full_specs[f"indexer.{i}"] = MLAAttentionSpec(
block_size=256,
num_kv_heads=1,
head_size=68,
dtype=torch.uint8,
page_size_padded=4608,
compress_ratio=4,
)
state_specs[f"mla_state.{i}"] = SlidingWindowMLASpec(
block_size=4,
num_kv_heads=1,
head_size=2048,
dtype=torch.float32,
sliding_window=8,
indexes_kv_by_block_stride=True,
)
state_specs[f"indexer_state.{i}"] = SlidingWindowMLASpec(
block_size=4,
num_kv_heads=1,
head_size=512,
dtype=torch.float32,
sliding_window=8,
indexes_kv_by_block_stride=True,
)
grouped_specs = [
UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=full_specs),
UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=state_specs),
]
groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
assert len(groups) == 2
assert {full_specs[f"indexer.{i}"].page_size_bytes for i in range(2)} == {4608}
assert {full_specs[f"indexer.{i}"].real_page_size_bytes for i in range(2)} == {
4352
}
assert {
state_specs[f"indexer_state.{i}"].page_size_bytes for i in range(2)
} == {8192}
full_group_bytes = 2 * (32768 + 4608)
state_group_bytes = 2 * (32768 + 8192)
bytes_per_block = max(full_group_bytes, state_group_bytes)
num_blocks, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, bytes_per_block * 32
)
assert num_blocks == 32
assert {tensor.block_stride for tensor in tensors} == {bytes_per_block}
packing = _packing_by_layer(tensors)
assert packing["mla.0"][0] == packing["mla_state.0"][0] == 0
assert packing["indexer.0"][0] == 32768
assert packing["indexer_state.0"][0] == 32768
def test_deepseek_v4_pro_stride(self):
groups = [
_make_page_group("full", [32768, 4608] * 30 + [1024] * 31),
_make_page_group("c4_state", [32768, 8192] * 30),
_make_page_group("c128_state", [32768] * 31),
_make_page_group("swa.0", [32768] * 31),
_make_page_group("swa.1", [32768] * 30),
]
expected_stride = 1_228_800
num_blocks, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, expected_stride * 32
)
assert num_blocks == 32
assert {tensor.block_stride for tensor in tensors} == {expected_stride}
assert {tensor.size for tensor in tensors} == {expected_stride * 32}
def test_all_tensors_have_block_stride(self):
_, tensors = _run()
for t in tensors:
@@ -122,9 +249,30 @@ class TestInterleavedPacking:
assert sizes.pop() > 0
def test_offsets_within_one_block(self):
_, tensors = _run()
for t in tensors:
assert t.offset < t.block_stride
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
_, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
page_sizes = _page_sizes_by_layer(groups)
packing = _packing_by_layer(tensors)
for layer_name, page_size in page_sizes.items():
offset, block_stride = packing[layer_name]
assert offset + page_size <= block_stride
def test_layouts_are_disjoint_within_each_group(self):
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
_, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
page_sizes = _page_sizes_by_layer(groups)
packing = _packing_by_layer(tensors)
for group in groups:
ranges = sorted(
(packing[name][0], packing[name][0] + page_sizes[name])
for name in group.layer_names
)
assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
def test_all_layers_accounted_for(self):
n_c4, n_c128, n_swa = 5, 4, 7
@@ -135,29 +283,29 @@ class TestInterleavedPacking:
expected = n_c4 * 2 + n_c128 + n_swa
assert len(all_names) == expected
def test_strided_views_are_independent(self):
def test_group_owned_blocks_do_not_alias(self):
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
page_sizes = _page_sizes_by_layer(groups)
num_blocks, tensors = _get_kv_cache_config_packed(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
views = []
for t in tensors:
page_size = page_sizes[t.shared_by[0]]
v = torch.as_strided(
backing,
size=(num_blocks, page_size),
stride=(t.block_stride, 1),
storage_offset=t.offset,
)
views.append(v)
views = _make_views(groups, num_blocks, tensors)
for i, v in enumerate(views):
v.fill_(i + 1)
expected = {}
value = 1
for block_id, group in enumerate(groups):
for layer_name in group.layer_names:
views[layer_name][block_id].fill_(value)
expected[layer_name] = (block_id, value)
value += 1
for i, v in enumerate(views):
assert (v == i + 1).all(), f"View {i} was corrupted"
for layer_name, (block_id, value) in expected.items():
assert (views[layer_name][block_id] == value).all()
# Once the first group releases its block, another group may reuse it.
for layer_name in groups[1].layer_names:
views[layer_name][0].fill_(255)
for layer_name in groups[1].layer_names:
assert (views[layer_name][0] == 255).all()
def test_hma_attention_groups_keep_default_backing(self):
full = _make_full_spec()
+332
View File
@@ -25,14 +25,17 @@ from vllm.multimodal.inputs import (
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
from vllm.utils.hashing import sha256
from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
from vllm.v1.engine import FinishReason
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
MambaSpec,
)
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
@@ -1068,6 +1071,141 @@ def test_preempt_during_execution():
assert requests[1].output_token_ids[0] == 42
def test_prefix_cache_query_not_inflated_by_connector_defer():
"""The GPU prefix-cache query is recorded at admission, so a request the
connector defers several times is counted once, not once per retry."""
num_defers_before_matching = 3
scheduler = create_scheduler(
enable_prefix_caching=True,
use_kv_connector=mock_kv(
matched_tokens=0,
is_async=False,
num_defers_before_matching=num_defers_before_matching,
),
)
request = create_requests(num_requests=1, num_tokens=32, block_size=16)[0]
scheduler.add_request(request)
# Each deferred step re-runs the lookup but records nothing.
for _ in range(num_defers_before_matching):
assert not scheduler.schedule().scheduled_new_reqs
output = scheduler.schedule()
assert any(r.req_id == request.request_id for r in output.scheduled_new_reqs)
stats = scheduler.kv_cache_manager.prefix_cache_stats
assert stats is not None
assert stats.requests == 1
assert stats.queries == request.num_tokens
def test_preemption_re_records_prefix_cache_query():
"""A preempted request re-enters the lookup on resume, so its recomputation
is counted again into the preempted stats."""
scheduler = create_scheduler(enable_prefix_caching=True)
request = create_requests(num_requests=1)[0]
scheduler.add_request(request)
scheduler.schedule()
stats = scheduler.kv_cache_manager.prefix_cache_stats
assert stats is not None
assert (stats.requests, stats.preempted_requests) == (1, 0)
scheduler.running.remove(request)
scheduler._preempt_request(request, 0.0)
assert request.status == RequestStatus.PREEMPTED
scheduler.schedule()
assert request.status == RequestStatus.RUNNING
assert stats.preempted_requests == 1
def test_prefix_cache_stats_not_recorded_when_caching_disabled():
"""With prefix caching off there is no local lookup, so admitting a request
records no phantom miss."""
scheduler = create_scheduler(enable_prefix_caching=False)
for request in create_requests(num_requests=2):
scheduler.add_request(request)
scheduler.schedule()
stats = scheduler.kv_cache_manager.prefix_cache_stats
assert stats is not None
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request():
"""A real cache hit rejected once by allocate_slots and admitted on the next
step is counted exactly once, hits included."""
block_size = 16
scheduler = create_scheduler(
enable_prefix_caching=True,
enable_chunked_prefill=False,
block_size=block_size,
)
# Seed the cache so the next request with the same prompt hits it.
seed = create_requests(
num_requests=1,
num_tokens=block_size * 2,
max_tokens=2,
same_prompt=True,
block_size=block_size,
req_ids=["seed"],
)[0]
scheduler.add_request(seed)
_step_until_done(
scheduler,
scheduler.schedule(),
ModelRunnerOutput(
req_ids=["seed"],
req_id_to_index={"seed": 0},
sampled_token_ids=[[1000]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
),
)
# The seeding step swapped in a fresh accumulator, so re-read it; the
# retried request must be the only thing recorded from here on.
stats = scheduler.kv_cache_manager.prefix_cache_stats
assert stats is not None
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
retried = create_requests(
num_requests=1,
num_tokens=block_size * 3,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["retried"],
)[0]
scheduler.add_request(retried)
# Reject the first allocation attempt, then delegate to the real one.
orig_allocate_slots = scheduler.kv_cache_manager.allocate_slots
allocate_results: list = []
def spy_allocate_slots(*args, **kwargs):
result = None if not allocate_results else orig_allocate_slots(*args, **kwargs)
allocate_results.append(result)
return result
scheduler.kv_cache_manager.allocate_slots = spy_allocate_slots
assert not scheduler.schedule().scheduled_new_reqs
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
assert "retried" in scheduler.schedule().num_scheduled_tokens
assert allocate_results[0] is None and allocate_results[1] is not None
assert (stats.requests, stats.queries, stats.hits) == (
1,
retried.num_tokens,
block_size * 2,
)
def test_scheduler_reset_prefix_cache():
scheduler = create_scheduler(enable_prefix_caching=True)
requests = create_requests(num_requests=10)
@@ -5329,3 +5467,197 @@ def test_async_load_reservation_prevents_wedge_e2e():
assert b.status == RequestStatus.WAITING
assert b.num_preemptions == 0
assert b.request_id not in req_to_blocks
def _create_hybrid_mamba_connector_scheduler(
matched_tokens: int,
block_size: int = 16,
num_blocks: int = 100,
) -> Scheduler:
"""FA + Mamba ("all" cache mode) scheduler with a MockKVConnector."""
model_config = ModelConfig(
model="facebook/opt-125m",
trust_remote_code=True,
dtype="float16",
seed=42,
skip_tokenizer_init=True,
)
vllm_config = VllmConfig(
scheduler_config=SchedulerConfig(
max_num_seqs=4,
max_num_batched_tokens=8192,
max_model_len=8192,
enable_chunked_prefill=True,
is_encoder_decoder=False,
watermark=0.0,
),
model_config=model_config,
cache_config=CacheConfig(
block_size=block_size,
enable_prefix_caching=True,
mamba_cache_mode="all",
),
kv_transfer_config=KVTransferConfig(
kv_connector="MockKVConnector",
kv_role="kv_both",
kv_connector_extra_config={
"matched_tokens": matched_tokens,
"is_async": False,
},
),
)
vllm_config.cache_config.num_gpu_blocks = num_blocks
kv_cache_config = KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["fa"],
FullAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
),
KVCacheGroupSpec(
["mamba"],
MambaSpec(
block_size=block_size,
shapes=((1, 1),),
dtypes=(torch.float32,),
mamba_cache_mode="all",
),
),
],
)
register_all_kvcache_specs(vllm_config)
return Scheduler(
vllm_config=vllm_config,
kv_cache_config=kv_cache_config,
structured_output_manager=StructuredOutputManager(vllm_config),
block_size=block_size,
hash_block_size=block_size,
log_stats=True,
)
@pytest.mark.parametrize(
"matched_tokens,expected_num_computed",
[
# No external hit: resume on the deepest locally-consistent boundary
# (block 0's state survives for both groups).
(0, 16),
# One external block on top of the reconciled local boundary.
(16, 32),
],
)
def test_hybrid_per_group_hit_divergence_with_connector(
matched_tokens: int, expected_num_computed: int
):
"""Per-group prefix hits can diverge for hybrid models with a connector
(#46453): under block pressure the FA prefix tail is evicted while a
deeper Mamba state block survives. The scheduler must not report the
deeper hit as locally computed (evicted FA blocks are not resident ->
engine crash / dirty KV); it falls back to the reconciled boundary that
every group is consistent at.
"""
block_size = 16
scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens)
manager = scheduler.kv_cache_manager
assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
# Seed a 4-block prefix so both groups cache all four boundaries
# (mamba cache mode "all" caches every block's state densely).
[fill] = create_requests(
num_requests=1,
num_tokens=4 * block_size,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["fill"],
)
computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
blocks = manager.allocate_slots(
fill, fill.num_tokens, num_computed, computed_blocks
)
fa_ids = [b.block_id for b in blocks.blocks[0]]
mamba_ids = [b.block_id for b in blocks.blocks[1]]
manager.free(fill)
# Evict the FA tail and the middle mamba states; block 0 (both groups)
# and the deep mamba state at block 3 survive.
manager.block_pool.evict_blocks({fa_ids[2], fa_ids[3], mamba_ids[1], mamba_ids[2]})
# A replay of the prefix plus one extra block now sees diverged
# per-group hits: FA stops at the evicted tail, while the mamba lookup
# finds the deeper surviving state.
[replay] = create_requests(
num_requests=1,
num_tokens=5 * block_size,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["replay"],
)
_, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
replay.block_hashes, replay.num_tokens - 1
)
assert per_group_hits == (2 * block_size, 4 * block_size) # diverged
scheduler.add_request(replay)
output = scheduler.schedule()
num_scheduled = output.num_scheduled_tokens[replay.request_id]
assert replay.num_tokens - num_scheduled == expected_num_computed
def test_hybrid_per_group_hit_divergence_fa_deeper_no_external():
"""The opposite divergence: the FA prefix survives deeper than the Mamba
state and the connector supplies nothing (ext == 0). Reporting the deep FA
hit as locally computed would resume with no valid Mamba state at that
boundary (silent bad output). The scheduler must fall back to the
convergent boundary that every group agrees on (block 0's surviving state).
"""
block_size = 16
scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens=0)
manager = scheduler.kv_cache_manager
assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
# Seed a 4-block prefix in both groups.
[fill] = create_requests(
num_requests=1,
num_tokens=4 * block_size,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["fill"],
)
computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
blocks = manager.allocate_slots(
fill, fill.num_tokens, num_computed, computed_blocks
)
mamba_ids = [b.block_id for b in blocks.blocks[1]]
manager.free(fill)
# Keep all FA blocks; evict every mamba state but block 0. FA reaches 4
# blocks, the mamba hit only reaches 1 -> diverged (FA > Mamba).
manager.block_pool.evict_blocks({mamba_ids[1], mamba_ids[2], mamba_ids[3]})
[replay] = create_requests(
num_requests=1,
num_tokens=5 * block_size,
max_tokens=1,
same_prompt=True,
block_size=block_size,
req_ids=["replay"],
)
_, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
replay.block_hashes, replay.num_tokens - 1
)
assert per_group_hits == (4 * block_size, 1 * block_size) # FA deeper
scheduler.add_request(replay)
output = scheduler.schedule()
num_scheduled = output.num_scheduled_tokens[replay.request_id]
# Must resume at the convergent boundary (block 0), not the deep FA hit.
assert replay.num_tokens - num_scheduled == block_size
+9 -2
View File
@@ -38,8 +38,12 @@ from vllm.v1.structured_output import StructuredOutputManager
EOS_TOKEN_ID = 50256
def mock_kv(matched_tokens: int, is_async: bool):
return MockKVConfig(matched_tokens=matched_tokens, is_async=is_async)
def mock_kv(matched_tokens: int, is_async: bool, num_defers_before_matching: int = 0):
return MockKVConfig(
matched_tokens=matched_tokens,
is_async=is_async,
num_defers_before_matching=num_defers_before_matching,
)
def create_scheduler(
@@ -115,6 +119,9 @@ def create_scheduler(
kv_connector_extra_config={
"matched_tokens": use_kv_connector.matched_tokens,
"is_async": use_kv_connector.is_async,
"num_defers_before_matching": (
use_kv_connector.num_defers_before_matching
),
},
)
elif isinstance(use_kv_connector, str):
+24 -2
View File
@@ -6,7 +6,6 @@ Unit tests for the breakable cudagraph primitives.
from __future__ import annotations
import os
import threading
from contextlib import nullcontext
from unittest.mock import patch
@@ -14,7 +13,21 @@ from unittest.mock import patch
import pytest
import torch
os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
@pytest.fixture(autouse=True)
def _enable_breakable_cudagraph(monkeypatch: pytest.MonkeyPatch):
"""Enable breakable cudagraphs for this module's tests only.
eager_break_during_capture reads the env at decoration time, which
happens inside the test bodies, so a per-test fixture suffices.
monkeypatch restores the env so other test files running in the same
pytest process are unaffected (a module-level os.environ assignment
used to leak into test_cudagraph_dispatch.py and break it).
"""
import vllm.envs as envs
monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
envs.disable_envs_cache()
def test_piecewise_capture_builds_fresh_metadata_for_both_passes():
@@ -86,10 +99,19 @@ def cuda_capture_stream():
"""
if not torch.cuda.is_available():
pytest.skip("CUDA required")
from vllm.utils.torch_utils import _current_stream_tls
prev_stream = getattr(_current_stream_tls, "value", None)
stream = torch.cuda.Stream()
with torch.cuda.stream(stream):
yield stream
torch.cuda.current_stream().wait_stream(stream)
# Exiting torch.cuda.stream() records the default stream in vllm's
# patched set_stream cache. A later CUDAGraphWrapper capture in the same
# process would then run on the default stream, which cannot capture
# (this broke test_cudagraph_dispatch.py when run in-process after this
# file). Restore the pre-fixture value so this module leaves no trace.
_current_stream_tls.value = prev_stream
# ---------------------------------------------------------------------------
@@ -0,0 +1,111 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm.config import (
CompilationConfig,
CUDAGraphMode,
ParallelConfig,
SchedulerConfig,
VllmConfig,
)
from vllm.distributed.device_communicators import pynccl_allocator
from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
pytestmark = pytest.mark.cpu_test
@pytest.fixture(autouse=True)
def _reset_graph_pool_id():
pynccl_allocator._graph_pool_id = None
yield
pynccl_allocator._graph_pool_id = None
def _create_vllm_config() -> MagicMock:
compilation_config = CompilationConfig(
cudagraph_mode="FULL",
cudagraph_capture_sizes=[4],
)
compilation_config.max_cudagraph_capture_size = 4
compilation_config.post_init_cudagraph_sizes()
vllm_config = MagicMock(spec=VllmConfig)
vllm_config.compilation_config = compilation_config
vllm_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=4)
vllm_config.parallel_config = ParallelConfig()
vllm_config.speculative_config = None
vllm_config.num_speculative_tokens = 0
return vllm_config
def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
"""FULL capture must set graph_pool_id before entering torch.cuda.graph().
NCCL symmetric memory checks this global during graph capture; without
it, capture fails with:
AssertionError: graph_pool_id is not set under graph capture
"""
graph_pool = object()
monkeypatch.setattr(
gpu_cudagraph_utils,
"get_pp_group",
lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
)
monkeypatch.setattr(
gpu_cudagraph_utils.current_platform,
"get_global_graph_pool",
lambda: graph_pool,
)
manager = gpu_cudagraph_utils.CudaGraphManager(
vllm_config=_create_vllm_config(),
device=torch.device("cpu"),
cudagraph_mode=CUDAGraphMode.FULL,
decode_query_len=1,
)
desc = BatchExecutionDescriptor(
cg_mode=CUDAGraphMode.FULL,
num_tokens=4,
num_reqs=4,
uniform_token_count=1,
)
manager._capture_descs[CUDAGraphMode.FULL] = [desc]
def create_forward_fn(desc, warmup):
return lambda _mode: None
@contextmanager
def fake_graph_capture(*args, **kwargs):
yield SimpleNamespace(stream=MagicMock())
fake_offloader = MagicMock()
def cuda_graph_enter(*args, **kwargs):
assert pynccl_allocator._graph_pool_id is graph_pool
mock_cuda_graph_ctx = MagicMock()
mock_cuda_graph_ctx.__enter__ = cuda_graph_enter
mock_cuda_graph_ctx.__exit__ = MagicMock(return_value=False)
with (
patch.object(gpu_cudagraph_utils, "graph_capture", fake_graph_capture),
patch.object(gpu_cudagraph_utils, "get_offloader", lambda: fake_offloader),
patch.object(gpu_cudagraph_utils.torch.cuda, "CUDAGraph"),
patch.object(
gpu_cudagraph_utils.torch.cuda,
"graph",
return_value=mock_cuda_graph_ctx,
) as mock_cuda_graph,
):
manager.capture(create_forward_fn)
mock_cuda_graph.assert_called_once()
@@ -108,6 +108,11 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
mgr.token_budgets = sorted(budgets)
mgr.max_batch_size = 16
mgr.use_dp = False
mgr.config = EncoderCudaGraphConfig(
modalities=["image"],
buffer_keys=[],
out_hidden_size=32,
)
mgr.budget_graphs = {"default": {}}
mgr.graph_pool = None
mgr.graph_hits = 0
@@ -0,0 +1,249 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
from contextlib import AsyncExitStack
from dataclasses import dataclass, replace
from typing import Any
import pytest
import torch
from tests.utils import multi_gpu_marks
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.sampling_params import RequestOutputKind
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.metrics.reader import Metric
DEEPSEEK_MTP_MAIN_RANDOM = "luccafong/deepseek_mtp_main_random"
DEEPSEEK_MTP_DRAFT_RANDOM = "luccafong/deepseek_mtp_draft_random"
PROMPT = "The capital of France is"
MAX_TOKENS = 8
MAX_MODEL_LEN = 2048
GPU_MEM_UTIL = 0.85
# Spec / no-spec greedy output should match, but exact match is not guaranteed
# across parallelism layouts (reduction order can flip near-tie argmaxes), so
# gate on a similarity ratio like the spec-decode E2E tests.
MIN_MATCH_RATIO = 0.8
def _token_match_ratio(a: tuple[int, ...], b: tuple[int, ...]) -> float:
n = min(len(a), len(b))
return sum(x == y for x, y in zip(a, b)) / n if n else 0.0
@dataclass(frozen=True)
class InlineConfig:
id: str
tp_size: int = 1
pp_size: int = 1
enable_expert_parallel: bool = False
enable_eplb: bool = False
skip_reason: str | None = None
INLINE_CONFIGS = [
InlineConfig(id="tp2", tp_size=2),
InlineConfig(id="ep2", tp_size=2, enable_expert_parallel=True),
InlineConfig(
id="ep2_eplb",
tp_size=2,
enable_expert_parallel=True,
enable_eplb=True,
),
InlineConfig(
id="pp2",
pp_size=2,
skip_reason=(
"DeepSeek MTP pipeline-parallel support is in flight upstream; "
"see https://github.com/vllm-project/vllm/pull/38104"
),
),
]
def _inline_kwargs(config: InlineConfig, *, with_spec: bool) -> dict[str, Any]:
kwargs: dict[str, Any] = {
"model": DEEPSEEK_MTP_MAIN_RANDOM,
"tensor_parallel_size": config.tp_size,
"pipeline_parallel_size": config.pp_size,
"enable_expert_parallel": config.enable_expert_parallel,
"max_model_len": MAX_MODEL_LEN,
"gpu_memory_utilization": GPU_MEM_UTIL,
"enforce_eager": True,
"trust_remote_code": True,
"disable_log_stats": False,
# MLA + batch invariance currently disables prefix caching at runtime
# (see vllm/v1/attention/backends/mla/common.py); set explicitly so the
# configuration is unambiguous.
"enable_prefix_caching": False,
}
if config.enable_eplb:
kwargs["enable_eplb"] = True
# Rearrangement first fires after step_interval // 4 forward steps
# (the step counter starts at 3/4 of step_interval, see
# vllm/distributed/eplb/eplb_state.py). Keep these small so the EPLB
# routine actually runs within the short MAX_TOKENS generation instead
# of never triggering.
kwargs["eplb_config"] = {
"num_redundant_experts": config.tp_size,
"window_size": 2,
"step_interval": 4,
"log_balancedness": False,
}
if with_spec:
kwargs["speculative_config"] = {
"method": "mtp",
"model": DEEPSEEK_MTP_DRAFT_RANDOM,
"num_speculative_tokens": 1,
}
return kwargs
def _generate_token_ids_inline(llm: LLM) -> tuple[int, ...]:
outputs = llm.generate(
[PROMPT],
SamplingParams(temperature=0.0, max_tokens=MAX_TOKENS, ignore_eos=True),
)
assert outputs and outputs[0].outputs, "expected one completion"
return tuple(outputs[0].outputs[0].token_ids)
def _spec_decode_num_drafts(metrics: list[Metric]) -> int:
name2metric = {m.name: m for m in metrics}
counter = name2metric.get("vllm:spec_decode_num_drafts")
assert counter is not None, (
"spec_decode_num_drafts metric missing; check disable_log_stats=False"
)
return int(counter.value)
@pytest.mark.parametrize(
"config",
[
pytest.param(c, id=c.id, marks=multi_gpu_marks(num_gpus=2))
for c in INLINE_CONFIGS
],
)
def test_deepseek_mtp_load_inline(
monkeypatch: pytest.MonkeyPatch,
config: InlineConfig,
):
"""MTP loads and drafts under TP/EP/EPLB; spec output matches no-spec greedy."""
if config.skip_reason is not None:
pytest.skip(config.skip_reason)
# Reduces run-to-run nondeterminism so spec / no-spec stay close;
# see tests/v1/distributed/test_eagle_dp.py for the same pattern.
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
spec_llm = LLM(**_inline_kwargs(config, with_spec=True))
try:
spec_tokens = _generate_token_ids_inline(spec_llm)
n_drafts = _spec_decode_num_drafts(spec_llm.get_metrics())
finally:
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
no_spec_llm = LLM(**_inline_kwargs(config, with_spec=False))
try:
no_spec_tokens = _generate_token_ids_inline(no_spec_llm)
finally:
del no_spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
# Non-vacuity: a silently-broken MTP drafter that falls back to
# verifier-only decoding would still produce matching output below.
assert n_drafts > 0, (
f"MTP drafter never fired under {config.id}: vllm:spec_decode_num_drafts == 0"
)
match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens)
print(f"\n{config.id}: spec/no-spec match_ratio={match_ratio:.3f}")
assert match_ratio >= MIN_MATCH_RATIO, (
f"Spec / no-spec output divergence under {config.id}: "
f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n"
f" spec_tokens = {spec_tokens}\n"
f" no_spec_tokens= {no_spec_tokens}"
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"dp_size",
[
pytest.param(2, marks=pytest.mark.distributed(num_gpus=2), id="dp2"),
pytest.param(1, id="dp1"),
],
)
async def test_deepseek_mtp_load_dp(monkeypatch: pytest.MonkeyPatch, dp_size: int):
"""MTP loads under DP (on and off) via AsyncLLM; spec matches no-spec greedy."""
if torch.accelerator.device_count() < dp_size:
pytest.skip(f"dp{dp_size} requires at least {dp_size} GPUs")
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
base_args = AsyncEngineArgs(
model=DEEPSEEK_MTP_MAIN_RANDOM,
tensor_parallel_size=1,
data_parallel_size=dp_size,
data_parallel_backend="mp",
max_model_len=MAX_MODEL_LEN,
gpu_memory_utilization=GPU_MEM_UTIL,
enforce_eager=True,
trust_remote_code=True,
enable_prefix_caching=False,
)
spec_args = replace(
base_args,
speculative_config={
"method": "mtp",
"model": DEEPSEEK_MTP_DRAFT_RANDOM,
"num_speculative_tokens": 1,
},
)
sampling_params = SamplingParams(
max_tokens=MAX_TOKENS,
ignore_eos=True,
output_kind=RequestOutputKind.FINAL_ONLY,
temperature=0.0,
)
async def _generate(args: AsyncEngineArgs, request_id: str) -> tuple[int, ...]:
try:
async with AsyncExitStack() as stack:
engine = AsyncLLM.from_engine_args(args)
stack.callback(engine.shutdown)
async for out in engine.generate(
request_id=request_id,
prompt=PROMPT,
sampling_params=sampling_params,
):
token_ids = tuple(out.outputs[0].token_ids)
assert len(token_ids) == MAX_TOKENS
return token_ids
raise AssertionError("AsyncLLM produced no output")
finally:
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
spec_tokens = await asyncio.wait_for(
_generate(spec_args, f"deepseek-mtp-dp{dp_size}-spec"), timeout=600
)
no_spec_tokens = await asyncio.wait_for(
_generate(base_args, f"deepseek-mtp-dp{dp_size}-no-spec"), timeout=600
)
match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens)
print(f"\ndp{dp_size}: spec/no-spec match_ratio={match_ratio:.3f}")
assert match_ratio >= MIN_MATCH_RATIO, (
f"Spec / no-spec output divergence under dp{dp_size}: "
f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n"
f" spec_tokens = {spec_tokens}\n"
f" no_spec_tokens= {no_spec_tokens}"
)
@@ -204,13 +204,12 @@ class FakeNixlWrapper:
def _make_fake_nixl_pkg():
"""Context manager that creates a temporary package making
`from nixl._api import nixl_agent` resolve to our FakeNixlWrapper.
Also creates rixl package for ROCm compatibility.
Also creates the ROCm NIXL packages.
Automatically cleans up the temporary directory when done.
"""
with tempfile.TemporaryDirectory() as td:
# Create both nixl and rixl packages for cross-platform compatibility
for pkg_name in ["nixl", "rixl"]:
for pkg_name in ["nixl", "nixl_rocm"]:
pkg_root = os.path.join(td, pkg_name, "_api")
os.makedirs(pkg_root, exist_ok=True)
@@ -3130,10 +3129,12 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
)
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
FakeNixlWrapper,
)
def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
"""kv_role='kv_both' should emit a deprecation log warning."""
from unittest.mock import patch
from vllm.logger import _print_warning_once
_print_warning_once.cache_clear()
@@ -3156,10 +3157,12 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
assert "deprecated" in msg
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
FakeNixlWrapper,
)
def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init):
"""kv_role='kv_consumer' or 'kv_producer' should NOT emit a warning."""
from unittest.mock import patch
for role in ("kv_consumer", "kv_producer"):
vllm_config = create_vllm_config(kv_role=role)
with patch(
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Verify that GPU memory is fully released after RixlConnector shutdown on ROCm.
"""Verify that GPU memory is released after NixlConnector shutdown on ROCm.
Regression test for ROCm/ucx#33: UCX rocm_ipc transport permanently pinned
GPU memory via hsa_amd_ipc_memory_create during ucp_mem_map, causing
@@ -62,7 +62,7 @@ def _full_gpu_cleanup():
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
def test_gpu_memory_rixl_hma(model_name, sw_size):
def test_gpu_memory_nixl_hma(model_name, sw_size):
"""Track GPU memory through NixlConnector create/infer/shutdown cycle."""
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
@@ -84,7 +84,7 @@ def test_gpu_memory_rixl_hma(model_name, sw_size):
}
print("\n" + "=" * 90)
print("GPU MEMORY -- RIXL NixlConnector HMA (ROCm)")
print("GPU MEMORY -- NIXL NixlConnector HMA (ROCm)")
print("=" * 90)
gc.collect()
torch.accelerator.empty_cache()
@@ -169,14 +169,14 @@ def test_gpu_memory_rixl_hma(model_name, sw_size):
@pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"])
def test_gpu_memory_no_rixl_baseline(model_name):
def test_gpu_memory_no_nixl_baseline(model_name):
"""Same workload without NixlConnector. Comparing driver-level memory
between this and test_gpu_memory_rixl_hma isolates UCX/RIXL impact."""
between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact."""
from vllm import LLM, SamplingParams
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
print("\n" + "=" * 90)
print("CONTROL -- same model, no RIXL connector")
print("CONTROL -- same model, no NIXL connector")
print("=" * 90)
gc.collect()
torch.accelerator.empty_cache()
@@ -209,7 +209,7 @@ def test_gpu_memory_no_rixl_baseline(model_name):
drv_base = snap0["drv_used_mb"]
drv_leaked = snap_final["drv_used_mb"] - drv_base
drv_peak = snap_peak["drv_used_mb"] - drv_base
print(f"\n Driver leaked (no rixl): {drv_leaked:.0f} MB")
print(f"\n Driver leaked (no NIXL): {drv_leaked:.0f} MB")
print("=" * 90)
leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0
+10
View File
@@ -362,6 +362,7 @@ class TestExampleConnector(ExampleConnector):
class MockKVConfig:
matched_tokens: int = 0
is_async: bool = False
num_defers_before_matching: int = 0
class MockKVConnectorMetadata(KVConnectorMetadata):
@@ -384,6 +385,12 @@ class MockKVConnector(KVConnectorBase_V1):
self.config = MockKVConfig(
matched_tokens=extra_config["matched_tokens"],
is_async=extra_config["is_async"],
num_defers_before_matching=extra_config.get(
"num_defers_before_matching", 0
),
)
self._defers_left: defaultdict[str, int] = defaultdict(
lambda: self.config.num_defers_before_matching
)
def get_num_new_matched_tokens(
@@ -391,6 +398,9 @@ class MockKVConnector(KVConnectorBase_V1):
request: Request,
num_computed_tokens: int,
) -> tuple[int | None, bool]:
if self._defers_left[request.request_id] > 0:
self._defers_left[request.request_id] -= 1
return (None, False)
return (self.config.matched_tokens, self.config.is_async)
def update_state_after_alloc(
+37 -1
View File
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Config-only resolution of DFlash draft attention causality.
"""Config-only DFlash behavior.
``dflash_has_any_non_causal`` decides pre-build whether the draft needs a
non-causal-capable backend, so its branch table (explicit override, SWA-derived
@@ -13,8 +13,12 @@ import pytest
from vllm.model_executor.models.qwen3_dflash import (
_dflash_layer_causal,
_get_dflash_fc_input_size,
dflash_has_any_non_causal,
)
from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
get_eagle3_aux_layers_from_config,
)
def _config(num_hidden_layers, layer_types=None, causal_override=None):
@@ -53,3 +57,35 @@ def test_dflash_layer_causal_is_per_layer():
config = _config(2, layer_types=["sliding_attention", "full_attention"])
assert _dflash_layer_causal(config, 0) is True
assert _dflash_layer_causal(config, 1) is False
def _vllm_config(**draft_config):
config = SimpleNamespace(**draft_config)
return SimpleNamespace(
speculative_config=SimpleNamespace(
draft_model_config=SimpleNamespace(hf_config=config)
)
)
def test_dflash_fc_uses_aux_layer_count():
vllm_config = _vllm_config(
num_hidden_layers=5,
hidden_size=4096,
target_hidden_size=None,
target_layer_ids=[1, 17, 32],
)
assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096
@pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"])
def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name):
layer_ids = [1, 17, 32]
vllm_config = _vllm_config(
**{config_name: {"layer_ids": layer_ids}},
)
assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == tuple(
layer_ids
)
@@ -16,7 +16,8 @@ from vllm.sampling_params import SamplingParams
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
pytestmark = pytest.mark.cpu_test
# Not cpu_test: InputBatch allocates pinned (UVA) memory, which requires a
# CUDA device even though the batch tensors live on the CPU.
@pytest.fixture
@@ -28,6 +29,7 @@ def mock_model_runner_with_input_batch():
runner.requests = {}
runner.max_num_reqs = 10
runner.max_model_len = 1024
runner.late_interaction_runner = Mock()
# Create a real InputBatch for e2e testing
runner.input_batch = InputBatch(
@@ -16,7 +16,8 @@ from vllm.v1.core.sched.output import (
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
from vllm.v1.worker.gpu.states import RequestState
pytestmark = pytest.mark.cpu_test
# Not cpu_test: RequestState allocates pinned (UVA) memory, which requires a
# CUDA device even though the request state itself lives on the CPU.
@pytest.fixture
@@ -31,13 +32,12 @@ def mock_model_runner_with_req_states():
num_speculative_steps=0,
vocab_size=32000,
device=torch.device("cpu"),
model_dtype=torch.float32,
cache_draft_logits=False,
)
runner.encoder_cache = None
runner.model_state = Mock()
runner.block_tables = Mock()
runner.lora_state = Mock()
runner.pp_handler = None
runner.sampler = None
runner.prompt_logprobs_worker = None
runner.is_last_pp_rank = False
@@ -53,6 +53,8 @@ def create_scheduler() -> Scheduler:
vllm_config.model_config = MagicMock()
vllm_config.model_config.skip_tokenizer_init = True
vllm_config.model_config.is_multimodal_model = False
vllm_config.model_config.is_encoder_decoder = False
vllm_config.model_config.is_diffusion = False
vllm_config.model_config.max_model_len = 1024
vllm_config.model_config.enable_return_routed_experts = False
vllm_config.cache_config = MagicMock()
@@ -496,7 +498,9 @@ class TestStreamingScheduler(unittest.TestCase):
eco_cycle2 = eco_dict_cycle2[session.client_index].outputs[0]
assert eco_cycle2.finish_reason == FinishReason.STOP
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
assert session in scheduler.waiting
# Sessions paused for streaming input are blocked-waiting, so they
# live in the skipped_waiting queue rather than the main waiting queue.
assert session in scheduler.skipped_waiting
assert session._all_token_ids == [1, 2, 3, 10, STOP_TOKEN]
# CRITICAL ASSERTION: Cached prompt_token_ids STILL must not have changed
-109
View File
@@ -6,122 +6,13 @@ from unittest.mock import patch
import pytest
import vllm.v1.worker.gpu_worker as gpu_worker_module
from vllm.multimodal.video import (
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
PYNVVIDEOCODEC_MAX_RETAINED_DECODERS,
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
from vllm.utils.mem_constants import GiB_bytes
from vllm.v1.worker import startup_plan
from vllm.v1.worker.gpu_worker import Worker
from vllm.v1.worker.startup_plan import (
maybe_apply_startup_plan,
maybe_save_startup_plan,
)
def _worker_with_mm_config(
mm_config: SimpleNamespace,
*,
api_process_count: int = 1,
) -> Worker:
worker = object.__new__(Worker)
worker.model_config = SimpleNamespace(multimodal_config=mm_config)
worker.parallel_config = SimpleNamespace(_api_process_count=api_process_count)
return worker
def _mm_config(
*,
mm_ipc_gpu_memory_gb: float = 0,
video_backend: str | None = None,
) -> SimpleNamespace:
video_kwargs = {} if video_backend is None else {"video_backend": video_backend}
return SimpleNamespace(
mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb,
media_io_kwargs={"video": video_kwargs} if video_kwargs else {},
)
def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int:
return api_process_count * (
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS
+ PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES
)
@pytest.mark.parametrize("video_backend", [None, "opencv"])
def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
monkeypatch: pytest.MonkeyPatch,
video_backend: str | None,
):
monkeypatch.setattr(
gpu_worker_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
worker = _worker_with_mm_config(
_mm_config(mm_ipc_gpu_memory_gb=0.25, video_backend=video_backend)
)
assert worker._reserve_mm_ipc_gpu_memory(GiB_bytes) == int(0.75 * GiB_bytes)
def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
gpu_worker_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
"opencv",
)
worker = _worker_with_mm_config(
_mm_config(
mm_ipc_gpu_memory_gb=0.25,
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
)
)
available_bytes = 4 * GiB_bytes
assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == (
available_bytes - int(0.25 * GiB_bytes) - _pynvvideocodec_decoder_budget()
)
def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
gpu_worker_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
worker = _worker_with_mm_config(_mm_config())
available_bytes = 4 * GiB_bytes
assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == (
available_bytes - _pynvvideocodec_decoder_budget()
)
def test_reserve_mm_ipc_gpu_memory_scales_pynvvideocodec_budget_by_api_servers(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(
gpu_worker_module.envs,
"VLLM_VIDEO_LOADER_BACKEND",
PYNVVIDEOCODEC_VIDEO_BACKEND,
)
worker = _worker_with_mm_config(_mm_config(), api_process_count=3)
available_bytes = 8 * GiB_bytes
assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == (
available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3)
)
# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and
# saved by Worker.determine_available_memory / compile_or_warm_up_model.
+1 -20
View File
@@ -36,22 +36,6 @@ def rust_extensions(*, optional: bool = False) -> list[RustExtension]:
]
def write_coverage_objects(extensions: list[RustExtension], output: Path) -> None:
artifacts = []
for extension in extensions:
for target in sorted(set(extension.target.values())):
target_path = ROOT_DIR.joinpath(*target.split("."))
if extension.binding == Binding.Exec:
matches = [target_path]
else:
matches = sorted(target_path.parent.glob(f"{target_path.name}*.so"))
if len(matches) != 1 or not matches[0].is_file():
raise RuntimeError(f"unable to locate Rust artifact for {target}")
artifacts.append(matches[0].relative_to(ROOT_DIR).as_posix())
output.write_text("\n".join(artifacts) + "\n")
def rust_py_extension_module_names() -> list[str]:
module_names = []
for extension in rust_extensions():
@@ -68,15 +52,12 @@ def rust_py_extension_module_names() -> list[str]:
def build_binary(build_rust_args: list[str]) -> None:
os.chdir(ROOT_DIR)
(ROOT_DIR / "vllm").mkdir(exist_ok=True)
extensions = rust_extensions(optional=False)
setup(
name="vllm-rust-frontend-build",
packages=[],
rust_extensions=extensions,
rust_extensions=rust_extensions(optional=False),
script_args=["build_rust", "--quiet", "--inplace", *build_rust_args],
)
if output := os.getenv("VLLM_RUST_COVERAGE_OBJECTS"):
write_coverage_objects(extensions, Path(output))
def main() -> None:

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