Compare commits

..
Author SHA1 Message Date
Shengqi ChenandGitHub a02155c787 Merge branch 'main' into cuda-arch-fixup 2026-07-21 09:33:48 +08:00
Chris LeonardandGitHub 97a98006b0 Update qutlass cmake for stable abi (#47879)
Signed-off-by: Chris Leonard <chleonar@redhat.com>
2026-07-20 18:31:10 -07:00
0a684ab0c0 [Bugfix] Fix WSL circular import from pin_memory warning_once (#48444)
Signed-off-by: AlejandroParedesLT <alejandroparedeslatorre@gmail.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-20 18:30:55 -07:00
0d9210a502 Fixes non-coalesced HBM access in marlin_int4_fp8_preprocess_kernel_awq (#47268)
Signed-off-by: xjx <493337577@qq.com>
Signed-off-by: flutist-alibaba <30485581+flutist@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-20 18:30:38 -07:00
1d874867ea [Misc][Docs] Fix broken protocol link in speech_to_text doc (#47212)
Signed-off-by: oonyshch <xonyshch@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-21 01:05:05 +00:00
2e2e626b40 [Bugfix] Count per-group blocks in get_max_concurrency_for_kv_cache_config (#48317)
Signed-off-by: David Orman <ormandj@corenode.com>
Co-authored-by: Luke Alonso <lalonso@gmail.com>
Co-authored-by: Martin Vit <martin@voipmonitor.org>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-07-21 00:29:10 +00:00
Nick HillandGitHub af91f4b3e4 [Cleanup] Remove unused StructuredOutputRequest.status field (#49235)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-21 00:06:16 +00:00
2396a61108 [Attention][MLA][DCP] Query replication for MLA decode (DeepSeek-V2/R1 + Kimi-K2.5) (#45964)
Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-07-20 23:51:27 +00:00
97a668152b [RL Infra][FlashInfer] Enable router replay output from FlashInfer monolithic MoE kernel (#44214)
Signed-off-by: Xuanyu Zhang <xuanyu.zhang@mistral.ai>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: OpenAI Codex <noreply@openai.com>
2026-07-20 16:45:10 -07:00
58b2012aa2 [copy of #45208] CuMem slept-L1 fragmentation accounting (#49208)
Signed-off-by: haosdent <haosdent@gmail.com>
Signed-off-by: Justin Wood <justin.m.wood@me.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: haosdent <haosdent@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Justin Wood <jwood@me.com>
2026-07-20 23:08:11 +00:00
Ning XieandGitHub b7c20d0cfa [chore] adjust logo be more friendly to white background terminal (#48938)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
2026-07-20 15:15:28 -07:00
TJianandGitHub a2b1f9fc3b [ROCm] [Release] [Bugfix] Fix the per commit wheel release pipeline. (#49245)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-07-20 22:12:38 +00:00
642076d26c Support loading sample_from_anchor flag from speculators config (#48639)
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
2026-07-20 14:52:43 -07:00
Charlie FuandGitHub 5feb3950e5 [ROCm][CI] fix test_rocm_quick_reduce.py (#49234)
Signed-off-by: charlifu <charlifu@amd.com>
2026-07-20 16:39:33 -05:00
4ec199b66a [Bugfix][Spec-Decode] Populate draft seq_lens_cpu_upper_bound for spec-decode attention metadata (#44492)
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-07-20 20:50:58 +00:00
Roberto L. CastroGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
7ca017778f [Feat][Perf] Add new warmup infrastructure for JITs (#47451)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-20 13:21:55 -07:00
fbfe58133d [Bugfix][KV Offload] Preserve reachable tails for hybrid SWA groups (#48911)
Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
Co-authored-by: Colton Ottley <colton@ottleyengineering.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
2026-07-20 22:12:36 +03:00
9dd62d80ab Cosmos3 FP8 ModelOpt/Diffusers remapping (#48952)
Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>
Signed-off-by: wkutak <wkutak@nvidia.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
2026-07-20 11:31:48 -07:00
f878367898 [Revert][Bugfix] Restore MiniCPM-V 4.6 ViT QKV weight loader (#49193)
Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: wjinxu <1299461899@qq.com>
2026-07-20 18:17:46 +00:00
Matthew BonanniGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
bd091079cb [Attention] FlashAttention 4 SM100 FP8 kv cache support (#42569)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-20 10:53:27 -07:00
b23bd73f54 [XPU]add sycl path for Mhc (#47245)
Signed-off-by: root <xiaolong.guo@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-20 15:32:54 +00:00
Bugen ZhaoandGitHub e2d7adeb64 [Rust Frontend] Bump xgrammar-structural-tag and enable local extension (#49161)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-20 16:22:24 +01:00
Isotr0pyandGitHub 15cb8e140d [Multimodal] Allow keeping original image mode for ImageIO (#49159)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
2026-07-20 13:42:45 +00:00
f007cceb42 [KV Offload] Support self-describing KV events with TieringOffloadingSpec (#48679)
Signed-off-by: Change72 <changg@nvidia.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-20 16:41:58 +03:00
0a5069e4e3 [Bugfix][Gemma4] Fix ModelOpt mixed-precision MoE config mapping (#48563)
Signed-off-by: wangqian <601731555@qq.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-20 06:39:28 -07:00
8ce53a616e [Bugfix] Zero new KV blocks for quantized + sliding-window hybrid caches (#47574)
Signed-off-by: EdalatiAli <aliedalati@cohere.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
2026-07-20 13:18:17 +00:00
Lena OnyshchenkoandGitHub ae10e855ab [Misc][Docs] Remove duplicate CodeGeex4 row in XPU model table (#47210)
Signed-off-by: oonyshch <xonyshch@gmail.com>
2026-07-20 10:05:36 +00:00
hclandGitHub 530ee36a0d fix(openai): reject non-numeric logprobs with 400 instead of 500 (#49144)
Signed-off-by: Chenglun Hu <chenglunhu@gmail.com>
2026-07-20 10:04:50 +00:00
Salt SatoandGitHub d835ad572c [Bugfix][Rust Frontend] Map missing prompt logprobs for single-token prompts in chat and raw generate (#49111)
Signed-off-by: Feathbow <feathbow@gmail.com>
2026-07-20 10:00:06 +00:00
47d0597ca2 [Misc][Docs] Fix broken csrc kernel links in fusions doc (#47211)
Signed-off-by: oonyshch <xonyshch@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-20 09:44:25 +00:00
ReidandGitHub 818cf61e91 [Rust Frontend] Fix macro-based content format detection (#49042)
Signed-off-by: reidliu41 <reid201711@gmail.com>
2026-07-20 09:39:13 +00:00
Bugen ZhaoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
c01618fdc8 [Rust][Benchmark] Integrate vllm-bench to vllm-rs & vllm CLI (#48930)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-20 09:31:25 +00:00
Xiaochang WuGitHubKunshang Jimergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
823eaf667d [XPU] FP8 o_proj with fp8_bmm and load-time scale transpose (#48334)
Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-20 16:32:03 +08:00
SageGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
f1f1259692 [Rust Frontend] Use zero-copy slicing for multimodal tensors (#48781)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
2026-07-20 16:28:25 +08:00
zofiaGitHubmayuyuacemergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Kunshang Ji
df13b5aef5 [XPU] [MoE] add quant input when prepare for fusedmoe (#47122)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
Signed-off-by: zofia <110436990+zufangzhu@users.noreply.github.com>
Co-authored-by: mayuyuace <qiming1.zhang@intel.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-20 15:47:26 +08:00
Sihan ChenGitHubLi, Jiang <jiang1.li@intel.com>
4938d44a3b [CPU] fixes heterogeneous NIXL KV transfer into CPU_ATTN decode workers (#47871)
Signed-off-by: Spycsh <sihan.chen@intel.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-07-20 07:33:13 +00:00
37bf988c2f [XPU][Bugfix] Fix GroupCoordinator device_index (#47295)
Signed-off-by: Michal Ganczarenko <michal.ganczarenko@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-20 15:25:56 +08:00
aoshen02andGitHub 9459fc6471 [Bugfix][RL] Set vLLM config during weight reload (#45989)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
2026-07-20 15:02:56 +08:00
5245c80564 [Doc] Document blocks_per_chunk in the KV offloading guide (#49100)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
2026-07-20 09:48:43 +03:00
9bc266d923 [Bugfix][KV Offload] Propagate EAGLE mode to SimpleCPU coordinator (#49071)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-20 06:39:11 +00:00
5c9f6557d7 [Hardware][CPU] Enable granite-4 model on cpu (#47641)
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Co-authored-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com>
Co-authored-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan>
Co-authored-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Co-authored-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-07-20 06:15:16 +00:00
Shengqi ChenandGitHub 319db65b68 Merge branch 'main' into cuda-arch-fixup 2026-07-09 01:55:33 +08:00
Shengqi ChenandGitHub 83a7669827 Merge branch 'main' into cuda-arch-fixup
Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-07 00:13:53 +08:00
Shengqi ChenandGitHub 401bed48ad Merge branch 'main' into cuda-arch-fixup
Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-03 15:54:56 +08:00
Shengqi ChenandGitHub d492d1e697 Merge branch 'main' into cuda-arch-fixup 2026-06-30 21:30:29 +08:00
Shengqi Chen f37e113590 [Build] Apply ruff format to CUDA arch regex
Keep the compiled CUDA arch regex on one line to match ruff-format output.

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 20:37:59 +08:00
Shengqi Chen a10e369f06 [Build] Address CUDA arch review comments
Fix CUDA arch warning tests so they do not depend on importing the stable libtorch extension, which keeps the warning coverage active in lightweight CI environments and avoids mypy treating a fixture value as a base class.

Use regex for the compiled-arch parser, apply formatter output, and make the CUTLASS grouped GEMM Python support query fall back to false when the op is unavailable or unimplemented in the current build.

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 20:03:28 +08:00
Shengqi Chen aa3f2efe42 [Temp] Cherry-pick #47139 to fix build
Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:38:28 +08:00
Shengqi ChenandCodex a32d1bff95 [Doc] Document CUDA wheel architecture coverage
Explain that pre-built CUDA wheels use the architecture lists selected by the release and build pipelines, which may be narrower than the full set vLLM can build from source.

Call out CUDA 12.9 architecture-specific wheel coverage, CUDA 13 family-specific targets, and the no-kernel-image error users may see when a wheel does not cover their GPU.

Co-authored-by: Codex <codex@openai.com>

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:27:42 +08:00
Shengqi ChenandCodex ee47a21fcd [Build] Warn on uncovered CUDA device architectures
Expose the compiled CUDA arch list from the stable extension and check visible CUDA devices against it during CUDA platform startup. The runtime check distinguishes exact architecture targets from CUDA 13 family targets so users get an early warning before hitting missing kernel images.

The startup warning is limited to the NVML-backed CUDA platform path to preserve the existing no-CUDA-init import behavior for non-NVML environments.

Co-authored-by: Codex <codex@openai.com>

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:27:42 +08:00
Shengqi ChenandCodex 0f471a3088 [Build] Scope stable CUDA kernel feature macros
Keep optional stable CUDA kernel feature macros on the source files that consume them instead of adding them to VLLM_GPU_FLAGS. This avoids perturbing unrelated compile commands and invalidating more cache entries when optional kernel families change.

Also align CUTLASS grouped MoE support with the SM10x/SM11x family so Thor works under both CUDA 12 SM101 and CUDA 13 SM110 reporting, and remove the stale ENABLE_CUTLASS_MLA definition left after the old CUTLASS MLA path was deleted.

Co-authored-by: Codex <codex@openai.com>

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:27:42 +08:00
190 changed files with 7228 additions and 2025 deletions
+5 -1
View File
@@ -18,6 +18,8 @@ steps:
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
- tests/kernels/mamba/test_cpu_short_conv.py
- tests/kernels/mamba/test_causal_conv1d.py
- tests/kernels/mamba/test_mamba_ssm.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -28,7 +30,9 @@ steps:
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
# Note: SDE can't be downloaded from CI host because of AWS WAF
# - label: CPU-Compatibility Tests
-3
View File
@@ -7,9 +7,6 @@
set -euo pipefail
# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
git fetch --tags --force origin
# The Rust frontend build needs protoc.
if ! command -v protoc >/dev/null 2>&1; then
brew install protobuf
@@ -40,7 +40,9 @@ function cpu_tests() {
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
# skip tests requiring model downloads if HF_TOKEN is not set
# due to rate-limits
@@ -97,3 +99,4 @@ function cpu_tests() {
# All of CPU tests are expected to be finished less than 40 mins.
export -f cpu_tests
timeout 2h bash -c cpu_tests
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Basic Correctness
key: basic-correctness
timeout_in_minutes: 68
timeout_in_minutes: 45
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: 45
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -51,7 +51,7 @@ steps:
- label: e2e Scheduling (1 GPU)
key: e2e-scheduling-1-gpu
timeout_in_minutes: 53
timeout_in_minutes: 35
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: 75
timeout_in_minutes: 50
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: 68
timeout_in_minutes: 45
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: 83
timeout_in_minutes: 45
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: 75
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -31,7 +31,7 @@ steps:
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 83
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/config/
+5 -1
View File
@@ -5,12 +5,14 @@ steps:
- label: Model Executor
device: h200_35gb
key: model-executor
timeout_in_minutes: 60
timeout_in_minutes: 45
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands:
- apt-get update && apt-get install -y curl libsodium23
@@ -34,7 +36,9 @@ steps:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
+1 -1
View File
@@ -137,7 +137,7 @@ steps:
- label: Language Models Test (MTEB)
key: language-models-test-mteb
timeout_in_minutes: 68
timeout_in_minutes: 45
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: 68
timeout_in_minutes: 45
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: 75
timeout_in_minutes: 50
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: 75
timeout_in_minutes: 50
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -85,7 +85,7 @@ steps:
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 98
timeout_in_minutes: 65
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: 150
timeout_in_minutes: 90
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
+86 -42
View File
@@ -227,6 +227,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
cuda_archs_loose_intersection(CUDA_ARCHS
"${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
set(VLLM_COMPILED_CUDA_ARCHS)
foreach(_ARCH ${CUDA_ARCHS})
set(_COMPILED_ARCH "${_ARCH}")
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
if(_ARCH MATCHES "^(10|11|12)\\.0$")
set(_COMPILED_ARCH "${_ARCH}f")
endif()
elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
if(_ARCH MATCHES "^(10\\.(0|1|3)|12\\.(0|1))$")
set(_COMPILED_ARCH "${_ARCH}a")
endif()
endif()
list(APPEND VLLM_COMPILED_CUDA_ARCHS "${_COMPILED_ARCH}")
endforeach()
list(JOIN VLLM_COMPILED_CUDA_ARCHS "," VLLM_COMPILED_CUDA_ARCHS_STR)
else()
#
# For other GPU targets override the GPU architectures detected by cmake/torch
@@ -385,8 +401,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
#
# _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
#
# Shared entry sources are part of the base extension source list, but some
# optional kernel families below append source-local feature macros to them.
set(VLLM_STABLE_TORCH_BINDINGS_SRC
"csrc/libtorch_stable/torch_bindings.cpp")
set(CACHE_KERNELS_SRC "csrc/libtorch_stable/cache_kernels.cu")
set(SCALED_MM_ENTRY_SRC
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
set(NVFP4_QUANT_ENTRY_SRC
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu")
set(NVFP4_SCALED_MM_ENTRY_SRC
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
set(VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/torch_bindings.cpp"
"${VLLM_STABLE_TORCH_BINDINGS_SRC}"
"csrc/libtorch_stable/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
"csrc/libtorch_stable/activation_kernels.cu"
@@ -409,7 +437,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/sampler.cu"
"csrc/libtorch_stable/topk.cu"
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
"csrc/libtorch_stable/cache_kernels.cu"
"${CACHE_KERNELS_SRC}"
"csrc/libtorch_stable/cache_kernels_fused.cu"
"csrc/libtorch_stable/custom_all_reduce.cu"
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
@@ -425,11 +453,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
"9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(COOPERATIVE_TOPK_ARCHS)
list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
endif()
endif()
if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -467,11 +490,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
"${SCALED_MM_ENTRY_SRC}"
"${NVFP4_QUANT_ENTRY_SRC}"
"${NVFP4_SCALED_MM_ENTRY_SRC}"
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
set_compile_definitions_for_srcs(
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC}"
DEFINITIONS VLLM_COMPILED_CUDA_ARCHS=\"${VLLM_COMPILED_CUDA_ARCHS_STR}\")
#
# Machete kernels
@@ -547,11 +573,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
CUDA_ARCHS "${CUDA_ARCHS}")
if(COOPERATIVE_TOPK_ARCHS)
set(COOPERATIVE_TOPK_SRC "csrc/libtorch_stable/cooperative_topk.cu")
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cooperative_topk.cu")
"${COOPERATIVE_TOPK_SRC}")
set_gencode_flags_for_srcs(
SRCS "csrc/libtorch_stable/cooperative_topk.cu"
SRCS "${COOPERATIVE_TOPK_SRC}"
CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC};${COOPERATIVE_TOPK_SRC}"
DEFINITIONS VLLM_ENABLE_COOPERATIVE_TOPK=1)
endif()
# Only build Marlin kernels if we are building for at least some compatible archs.
@@ -760,8 +790,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM90_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM90_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM90=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM90_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
@@ -794,8 +826,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM120_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM120_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM120=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM120_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
@@ -828,8 +862,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM100_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM100_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM100=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}")
@@ -858,8 +894,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_C2X_SRCS}"
CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_C2X_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_C2X=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_C2X_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
else()
if (SCALED_MM_3X_ARCHS)
@@ -884,8 +922,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${CUTLASS_MOE_SM90_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM90_SRCS}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM90=1)
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM90_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
else()
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
@@ -908,8 +948,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${CUTLASS_MOE_SM100_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM100_SRCS}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM10X_OR_SM11X=1)
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
# The implementation is named sm100 historically, but it is built for the
# SM10x/SM11x family: CUDA 12 Thor reports SM101, CUDA 13 Thor reports
# SM110. Keep the compile-time macro aligned with runtime dispatch.
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
else()
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
@@ -950,10 +995,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
#
# SM12x FP4 kernels. These share some generic NVFP4 quantization entry
# sources with the SM10x/11x block below; set_gencode_flags_for_srcs appends
# per-source flags, so shared files accumulate both SM12x and SM10x/11x
# gencodes when both families are requested.
# Shared FP4 implementation sources live next to the FP4 arch logic because
# both SM12x and SM10x/11x append family-specific gencodes and feature macros
# to them.
set(FP4_SHARED_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
else()
@@ -961,18 +1012,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS)
set(FP4_SM120_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
${FP4_SHARED_SRCS}
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
)
set_gencode_flags_for_srcs(
SRCS "${FP4_SM120_SRCS}"
CUDA_ARCHS "${FP4_SM120_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM120_SRCS}"
DEFINITIONS ENABLE_NVFP4_SM120=1)
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM120=1)
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
else()
message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.")
@@ -987,14 +1039,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS)
set(FP4_SM100_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
${FP4_SHARED_SRCS}
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
message(STATUS
"Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).")
@@ -1002,9 +1050,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${FP4_SM100_SRCS}"
CUDA_ARCHS "${FP4_SM100_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM100_SRCS}"
DEFINITIONS ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
else()
message(STATUS "Not building SM10x/11x NVFP4/MXFP4 as no compatible archs were found.")
@@ -1058,7 +1107,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${CUTLASS_MLA_SRCS}"
CUDA_ARCHS "${MLA_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MLA_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
# Add MLA-specific include directories only to MLA source files
set_source_files_properties(${CUTLASS_MLA_SRCS}
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
@@ -1106,10 +1154,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# Needed to use cuda/hip APIs from C-shim
if(VLLM_GPU_LANG STREQUAL "CUDA")
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
if(COOPERATIVE_TOPK_ARCHS)
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_COOPERATIVE_TOPK=1)
endif()
# Needed by CUTLASS kernels
target_compile_definitions(_C_stable_libtorch PRIVATE
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+3
View File
@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
"csrc/cpu/layernorm.cpp"
"csrc/cpu/mla_decode.cpp"
"csrc/cpu/pos_encoding.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/torch_bindings.cpp")
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
"csrc/cpu/activation.cpp"
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
+5 -5
View File
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
else()
set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
if(NOT _qutlass_fc_root)
@@ -125,8 +125,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
CUDA_ARCHS "${QUTLASS_ARCHS}"
)
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
# Keep it as its own extension (registers torch.ops._qutlass_C).
define_extension_target(
_qutlass_C
DESTINATION vllm
@@ -139,9 +137,11 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE
QUTLASS_DISABLE_PYBIND=1
QUTLASS_MINIMAL_BUILD=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
TORCH_TARGET_VERSION=0x020B000000000000ULL
USE_CUDA)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+22
View File
@@ -344,6 +344,28 @@ macro(set_gencode_flags_for_srcs)
endif()
endmacro()
#
# For a list of source files append preprocessor definitions to file-specific
# compile options. Use this for optional kernel feature macros so toggling one
# kernel family does not perturb the compile command for every source in the
# extension target.
#
macro(set_compile_definitions_for_srcs)
set(options)
set(oneValueArgs)
set(multiValueArgs SRCS DEFINITIONS)
cmake_parse_arguments(arg "${options}" "${oneValueArgs}"
"${multiValueArgs}" ${ARGN})
foreach(_DEF ${arg_DEFINITIONS})
set_property(
SOURCE ${arg_SRCS}
APPEND PROPERTY
COMPILE_DEFINITIONS "${_DEF}"
)
endforeach()
endmacro()
#
# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form
# `<major>.<minor>[letter]` compute the "loose intersection" with the
+8 -7
View File
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
reg.val[1] = fp16_to_fp32_bits(raw_lo);
}
float reduce_sum() const {
AliasReg ar;
ar.reg = reg;
float result = 0;
unroll_loop<int, VEC_ELEM_NUM>(
[&result, &ar](int i) { result += ar.values[i]; });
return result;
// VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
// Step 1: pairwise sum of the two 4-wide halves
__vector float s = vec_add(reg.val[0], reg.val[1]);
// Step 2: rotate by 8 bytes (2 floats) and add
s = vec_add(s, vec_sld(s, s, 8));
// Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total
s = vec_add(s, vec_sld(s, s, 4));
return vec_extract(s, 0);
}
FP32Vec8 exp() const {
f32x4x2_t out;
+285
View File
@@ -0,0 +1,285 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
// mamba_kernels.hpp.
#include "cpu/mamba_kernels.hpp"
#include <ATen/ATen.h>
#include <torch/library.h>
#include <c10/util/Optional.h>
#include "cpu_types.hpp"
// ---------------------------------------------------------------------------
// causal_conv1d_update
// ---------------------------------------------------------------------------
at::Tensor causal_conv1d_update_cpu_impl(
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
const c10::optional<at::Tensor>& bias,
const c10::optional<std::string>& activation,
const c10::optional<at::Tensor>& conv_state_indices,
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
bool do_silu = false;
if (activation.has_value()) {
const std::string& act = activation.value();
do_silu = (act == "silu" || act == "swish");
}
at::ScalarType dtype = x.scalar_type();
// Input x: contiguous in native dtype.
at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
// conv_state: NEVER copy the full paged tensor just for layout reasons.
// If the dtype matches we work directly on conv_state (contiguous or not)
// by extracting strides and passing them to the kernel.
// Only a dtype-conversion copy is made when types differ (rare for BF16).
bool state_type_ok = (conv_state.scalar_type() == dtype);
at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
// state_c and conv_state may be non-contiguous — that is intentional.
// Weight: coerce to same dtype if needed (should match in practice)
at::Tensor w_c =
(weight.scalar_type() != dtype)
? weight.to(dtype).contiguous()
: (weight.is_contiguous() ? weight : weight.contiguous());
// Bias stays float32 (small scalar, used only for fp32 accumulation)
at::Tensor bias_f32;
if (bias.has_value() && bias.value().defined())
bias_f32 = bias.value().to(at::kFloat).contiguous();
int64_t batch = x_c.size(0);
int64_t dim = x_c.size(1);
int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
int64_t width = w_c.size(1);
int64_t state_len = state_c.size(2);
// Extract strides — works for contiguous AND non-contiguous (transposed)
// state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
// contiguous) stride(1): between conv channels (dim stride) stride(2):
// between state elements (=1 when contiguous, =dim when transposed)
int64_t stride_s_slot = state_c.stride(0);
int64_t stride_s_dim = state_c.stride(1);
int64_t stride_s_state = state_c.stride(2);
at::Tensor out = x_c.clone(); // native dtype, no float32 alloc
const int32_t* cache_idx_ptr = nullptr;
at::Tensor cache_idx_int;
if (conv_state_indices.has_value()) {
cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
}
VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
out.data_ptr<scalar_t>(), cache_idx_ptr,
static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
do_silu);
});
// Write back only when a type-conversion copy was made.
// Layout-only non-contiguity is handled via strides above — no copy needed.
if (!state_type_ok) conv_state.copy_(state_c);
return out;
}
// ---------------------------------------------------------------------------
// selective_state_update
// ---------------------------------------------------------------------------
void selective_state_update_cpu_impl(
at::Tensor& state, // (nstates, nheads, dim, dstate)
const at::Tensor& x, // (N, nheads, dim)
const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
const at::Tensor& C, const c10::optional<at::Tensor>& D,
const c10::optional<at::Tensor>& z,
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
const c10::optional<at::Tensor>& state_batch_indices,
const c10::optional<at::Tensor>& dst_state_batch_indices,
int64_t null_block_id, at::Tensor& out,
const c10::optional<at::Tensor>& num_accepted_tokens,
const c10::optional<at::Tensor>& cu_seqlens) {
at::ScalarType state_type = state.scalar_type();
at::ScalarType input_type = x.scalar_type();
// x, B, C must be contiguous and match input_type
auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor x_in = ensure_input(x);
at::Tensor B_in = ensure_input(B);
at::Tensor C_in = ensure_input(C);
at::Tensor z_in;
if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
// A, D, dt_bias are float32 model parameters that arrive here as expanded
// tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
// We need just the scalar value per head as a (nheads,) 1-D array so that
// A_ptr[h] in the kernel correctly reads head h's value.
//
// Strategy: peel trailing expanded (stride=0) dims via .select(), which is
// a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0)
// → .select(2,0) → (nheads, head_dim) strides (1,0)
// → .select(1,0) → (nheads,) stride (1,) ← contiguous, free.
// No allocation, no type conversion (A is already float32).
auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
at::Tensor r = t;
// Peel trailing dimensions that are broadcast (stride=0 or size=1)
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32
at::Tensor D_f32, dt_bias_f32;
if (D.has_value() && D.value().defined())
D_f32 = to_per_head_1d_f32(D.value());
if (dt_bias.has_value() && dt_bias.value().defined())
dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
// dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
// the type conversion so we convert head_dim x fewer elements.
at::Tensor dt_f32;
{
// If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
// take a zero-copy view of index 0 along that dim first.
at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads)
at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
}
int64_t nheads = state.size(1);
int64_t dim = state.size(2);
int64_t dstate = state.size(3);
int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
? cu_seqlens.value().size(0) - 1
: x_in.size(0);
int64_t ngroups = B_in.size(1);
// Strides
int64_t stride_state_n = state.stride(0);
int64_t stride_state_h = state.stride(1);
int64_t stride_state_d = state.stride(2);
int64_t stride_x_n = x_in.stride(0);
int64_t stride_x_h = x_in.stride(1);
int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads)
int64_t stride_BC_n = B_in.stride(0);
int64_t stride_BC_g = B_in.stride(1);
int64_t stride_out_n = out.stride(0);
int64_t stride_out_h = out.stride(1);
// Optional index pointers
auto get_int32_ptr =
[](const c10::optional<at::Tensor>& opt) -> const int32_t* {
return (opt.has_value() && opt.value().defined())
? opt.value().data_ptr<int32_t>()
: nullptr;
};
const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
// Dispatch on (state_t, input_t, out_t): write directly into `out`
// without any intermediate float32 buffer.
VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
using state_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
using input_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
using out_t = scalar_t;
mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
state.data_ptr<state_t>(), stride_state_n, stride_state_h,
stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
nheads, ngroups, dim, dstate, dt_softplus);
});
});
});
}
// ---------------------------------------------------------------------------
// mamba_chunk_scan_fwd_cpu
// ---------------------------------------------------------------------------
void mamba_chunk_scan_fwd_cpu_impl(
at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller
at::Tensor&
final_states, // [batch, nheads, headdim, dstate] float32 contiguous
const at::Tensor& x, // [seqlen, nheads, headdim]
const at::Tensor&
dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
const at::Tensor& A, // [nheads] float32
const at::Tensor& B, // [seqlen, ngroups, dstate]
const at::Tensor& C, // [seqlen, ngroups, dstate]
const c10::optional<at::Tensor>& D, // [nheads] float32 (optional)
const c10::optional<at::Tensor>& z, // [seqlen, nheads, headdim] (optional)
const at::Tensor& cu_seqlens // [batch+1] int32
) {
const at::ScalarType input_type = x.scalar_type();
auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor x_in = ensure_contig(x);
at::Tensor B_in = ensure_contig(B);
at::Tensor C_in = ensure_contig(C);
at::Tensor z_in;
if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
// A and D are float32 model parameters, potentially broadcast-expanded.
// Strip trailing broadcast dims to get a contiguous (nheads,) array.
auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
at::Tensor r = t;
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor A_f32 = to_per_head_f32(A);
at::Tensor D_f32;
if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
// dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
// Python.
at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
const int64_t batch = final_states.size(0);
const int64_t nheads = final_states.size(1);
const int64_t headdim = final_states.size(2);
const int64_t dstate = final_states.size(3);
const int64_t ngroups = B_in.size(1);
TORCH_CHECK(final_states.is_contiguous(),
"mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
TORCH_CHECK(out.is_contiguous(),
"mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
"raw data_ptr)");
VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
ngroups, headdim, dstate);
});
}
+382
View File
@@ -0,0 +1,382 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Fused CPU vector kernels for Mamba decode-step hotspots:
// - causal_conv1d_update (depthwise 1-D conv state roll + compute)
// - selective_state_update (SSM recurrence, single-step)
#pragma once
#include "cpu_types.hpp"
#include <cmath>
#include <cstring>
#include <cstdint>
#include <algorithm>
namespace mamba_cpu {
// ---------------------------------------------------------------------------
// causal_conv1d_update — templated for native BF16/FP32
//
// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
// Explicit strides are passed so the kernel writes directly into the
// correct memory locations without making a contiguous copy of the full
// paged tensor (which was the source of the 34-41% direct_copy_kernel).
//
// stride_s_slot = state.stride(0) — between cache slots
// stride_s_dim = state.stride(1) — between conv_dim channels
// stride_s_state = state.stride(2) — between state elements
//
// When stride_s_state == 1 (contiguous), the memmove fast path is used.
// ---------------------------------------------------------------------------
template <typename scalar_t>
inline void causal_conv1d_update_kernel(
const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
int64_t width, int64_t state_len, bool do_silu) {
#pragma omp parallel for
for (int64_t b = 0; b < batch; ++b) {
int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
if (cache_idx == pad_slot_id) continue;
for (int64_t t = 0; t < seqlen; ++t) {
const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
// Base of this slot in the (possibly non-contiguous) paged state
scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
for (int64_t d = 0; d < dim; ++d) {
float x_val = static_cast<float>(x_b[d * seqlen]);
scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state
const scalar_t* w = weight_ptr + d * width;
// Accumulate in float32 for precision
float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
for (int64_t k = 0; k < state_len; ++k) {
acc += static_cast<float>(w[k]) *
static_cast<float>(sd[k * stride_s_state]);
}
acc += static_cast<float>(w[state_len]) * x_val;
// Shift state left and append new input.
// Use memmove when contiguous (stride==1); element loop otherwise.
if (stride_s_state == 1) {
if (state_len > 1)
std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
} else {
for (int64_t k = 0; k < state_len - 1; ++k)
sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
if (state_len > 0)
sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
}
if (do_silu) {
float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
: std::exp(acc) / (1.0f + std::exp(acc));
acc *= sigmoid;
}
out_b[d * seqlen] = static_cast<scalar_t>(acc);
}
}
}
}
// ---------------------------------------------------------------------------
// selective_state_update
//
// Template parameters:
// state_t - dtype of ssm_state cache (typically BFloat16)
// input_t - dtype of x, B, C (typically BFloat16)
// out_t - dtype of output tensor (typically BFloat16)
// Write directly — no float32 intermediate buffer needed.
//
// A, D, dt_bias are accepted as const float* (they are always float32
// model parameters in Mamba2). This eliminates the per-call float32→BF16
// conversion and the .contiguous() materialisation of the broadcast-expand.
//
// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
// ---------------------------------------------------------------------------
template <typename state_t, typename input_t, typename out_t = float>
inline void selective_state_update_kernel(
state_t* __restrict__ state_ptr, int64_t stride_state_n,
int64_t stride_state_h, int64_t stride_state_d,
const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
// dt: (N, nheads) — scalar per head, NOT expanded to head_dim
const float* __restrict__ dt_ptr, int64_t stride_dt_n,
// A: (nheads,) float32 — scalar per head
const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
// D: (nheads,) float32 — scalar per head (nullptr if not used)
const float* __restrict__ D_ptr,
// z: same shape as x (optional)
const input_t* __restrict__ z_ptr,
// dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
int64_t stride_out_n, int64_t stride_out_h,
const int32_t* __restrict__ state_batch_indices,
const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
const int32_t* __restrict__ num_accepted_tokens,
const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
using state_vec_t = vec_op::vec_t<state_t>;
using input_vec_t = vec_op::vec_t<input_t>;
constexpr int VEC_ELEM_NUM = 8;
int64_t nheads_per_group = nheads / ngroups;
for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
int64_t bos, seq_len;
if (cu_seqlens != nullptr) {
bos = cu_seqlens[seq_idx];
seq_len = cu_seqlens[seq_idx + 1] - bos;
} else {
bos = seq_idx;
seq_len = 1;
}
int64_t state_read_idx = (state_batch_indices != nullptr)
? state_batch_indices[seq_idx]
: seq_idx;
if (state_read_idx == null_block_id) continue;
int64_t state_write_idx = (num_accepted_tokens == nullptr)
? ((dst_state_batch_indices != nullptr)
? dst_state_batch_indices[seq_idx]
: state_read_idx)
: -1;
state_t* s = state_ptr + state_read_idx * stride_state_n;
for (int64_t t = 0; t < seq_len; ++t) {
int64_t token_idx = bos + t;
const input_t* x_tok = x_ptr + token_idx * stride_x_n;
// dt: (N, nheads) — one float per head per token
const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
out_t* out_tok = out_ptr + token_idx * stride_out_n;
#pragma omp parallel for
for (int64_t h = 0; h < nheads; ++h) {
int64_t g = h / nheads_per_group;
const input_t* x_h = x_tok + h * stride_x_h;
const input_t* B_g = B_tok + g * stride_BC_g;
const input_t* C_g = C_tok + g * stride_BC_g;
out_t* out_h = out_tok + h * stride_out_h;
state_t* s_h = s + h * stride_state_h;
// Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
float dt_val = dt_tok[h];
if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
if (dt_softplus) {
dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
}
const float A_val = A_ptr[h]; // scalar: same for all dim, dstate
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
const input_t* z_h =
(z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
: nullptr;
vec_op::FP32Vec8 dt_vec(dt_val);
// dA = exp(A * dt): A and dt are SCALARS per head, so compute once
// and broadcast. This saves 7 redundant std::exp() calls that
// FP32Vec8::exp() would otherwise make on the broadcast vector.
const float dA_scalar = std::exp(A_val * dt_val);
vec_op::FP32Vec8 dA(dA_scalar); // broadcast
for (int64_t d = 0; d < dim; ++d) {
float x_val = static_cast<float>(x_h[d]);
vec_op::FP32Vec8 out_vec(0.0f);
state_t* s_hd = s_h + d * stride_state_d;
const input_t* B_g_base = B_g;
const input_t* C_g_base = C_g;
vec_op::FP32Vec8 x_vec(x_val);
// dBx = B * x * dt — same dA for all dstate (A is scalar)
// s_new = s * dA + B * x * dt
int64_t n = 0;
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
vec_op::FP32Vec8 s_new = s_v * dA + dBx;
state_vec_t(s_new).save(s_hd + n);
out_vec = out_vec + s_new * C_v;
}
float out_val = out_vec.reduce_sum();
for (; n < dstate; ++n) {
// Reuse dA_scalar computed once per head — no exp() re-call
float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
s_hd[n] = static_cast<state_t>(s_new);
out_val += s_new * static_cast<float>(C_g[n]);
}
if (D_ptr != nullptr) out_val += x_val * D_val;
if (z_h != nullptr) {
float z_val = static_cast<float>(z_h[d]);
float sigmoid = (z_val >= 0)
? 1.0f / (1.0f + std::exp(-z_val))
: std::exp(z_val) / (1.0f + std::exp(z_val));
out_val *= z_val * sigmoid;
}
out_h[d] = static_cast<out_t>(out_val);
}
}
if (num_accepted_tokens != nullptr &&
dst_state_batch_indices != nullptr) {
int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
}
}
}
if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
state_write_idx != state_read_idx) {
state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
}
}
}
// ---------------------------------------------------------------------------
// mamba_chunk_scan_fwd
//
// Prefill SSM recurrence for Mamba2 / SSD models.
//
// Key difference from selective_state_update_kernel (decode path):
// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
// Each thread owns a (batch, head) slice and runs the entire token
// sequence without any per-token OpenMP synchronisation overhead.
// For seqlen=256, this eliminates 256 thread-barrier launches per batch.
//
// `dt` arrives already processed (float32, after bias + softplus + clamp)
// to keep this kernel simple. Preprocessing is done in the Python wrapper.
//
// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
// tensor, pre-initialised by the caller (zero or from initial_states).
// Each (b, h) slice is private to exactly one thread via collapse(2), so
// there are no write conflicts.
//
// D is treated as a scalar per head ([nheads] float32).
// ---------------------------------------------------------------------------
template <typename input_t>
inline void mamba_chunk_scan_fwd_kernel(
float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32
const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim]
const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed)
const float* __restrict__ A_ptr, // [nheads] f32
const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate]
const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate]
const float* __restrict__ D_ptr, // [nheads] f32 (nullable)
const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable)
input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim]
const int32_t* __restrict__ cu_seqlens, // [batch+1] int32
int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
int64_t dstate) {
using input_vec_t = vec_op::vec_t<input_t>;
constexpr int VEC_ELEM_NUM = 8;
const int64_t nheads_per_group = nheads / ngroups;
// states layout: [batch, nheads, headdim, dstate] contiguous (caller
// guarantee)
const int64_t stride_s_b = nheads * headdim * dstate;
const int64_t stride_s_h = headdim * dstate;
// stride_s_d = dstate, stride_s_n = 1
#pragma omp parallel for collapse(2) schedule(static)
for (int64_t b = 0; b < batch; ++b) {
for (int64_t h = 0; h < nheads; ++h) {
const int64_t seq_start = cu_seqlens[b];
const int64_t seq_end = cu_seqlens[b + 1];
const int64_t g = h / nheads_per_group;
const float A_val = A_ptr[h];
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
// Working state slice: states[b, h, :, :] — float32, headdim * dstate.
// Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
for (int64_t t = seq_start; t < seq_end; ++t) {
const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
const float* dt_h = dt_ptr + t * nheads + h;
const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
const input_t* z_h = (z_ptr != nullptr)
? z_ptr + t * nheads * headdim + h * headdim
: nullptr;
input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
const float dt_val = *dt_h;
const float dA_val = std::exp(A_val * dt_val);
const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar
const vec_op::FP32Vec8 dt_vec(dt_val);
for (int64_t d = 0; d < headdim; ++d) {
const float x_val = static_cast<float>(x_h[d]);
float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32
// Vectorised SSM update + readout over dstate:
// s_new = s * dA + x * dt * B
// y += s_new * C
int64_t n = 0;
vec_op::FP32Vec8 y_vec(0.0f);
const vec_op::FP32Vec8 x_vec(x_val);
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
const vec_op::FP32Vec8 s_v(s_bhd + n);
const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
s_new.save(s_bhd + n);
y_vec = y_vec + s_new * C_v;
}
float y_val = y_vec.reduce_sum();
// Scalar tail for remaining dstate elements
for (; n < dstate; ++n) {
const float B_n = static_cast<float>(B_g[n]);
const float C_n = static_cast<float>(C_g[n]);
const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
s_bhd[n] = s_new;
y_val += s_new * C_n;
}
// D skip connection (scalar per head)
if (D_ptr != nullptr) y_val += x_val * D_val;
// z gating: out = y * z * sigmoid(z) (SiLU)
if (z_h != nullptr) {
const float z_val = static_cast<float>(z_h[d]);
const float sigmoid =
(z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
: std::exp(z_val) / (1.0f + std::exp(z_val));
y_val *= z_val * sigmoid;
}
out_h[d] = static_cast<input_t>(y_val);
}
}
}
}
}
} // namespace mamba_cpu
+50
View File
@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping,
const int64_t block_size);
at::Tensor causal_conv1d_update_cpu_impl(
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
const c10::optional<at::Tensor>& bias,
const c10::optional<std::string>& activation,
const c10::optional<at::Tensor>& conv_state_indices,
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
void selective_state_update_cpu_impl(
at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
const c10::optional<at::Tensor>& state_batch_indices,
const c10::optional<at::Tensor>& dst_state_batch_indices,
int64_t null_block_id, at::Tensor& out,
const c10::optional<at::Tensor>& num_accepted_tokens,
const c10::optional<at::Tensor>& cu_seqlens);
void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
const at::Tensor& x, const at::Tensor& dt,
const at::Tensor& A, const at::Tensor& B,
const at::Tensor& C,
const c10::optional<at::Tensor>& D,
const c10::optional<at::Tensor>& z,
const at::Tensor& cu_seqlens);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
@@ -595,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"block_size) -> ()",
&compute_slot_mapping_kernel_impl);
// Mamba CPU kernels
ops.def(
"causal_conv1d_update_cpu_vec("
"Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
"Tensor? bias, str? activation, Tensor? conv_state_indices, "
"Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
&causal_conv1d_update_cpu_impl);
ops.def(
"selective_state_update_cpu("
"Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
"Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
"Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
"SymInt null_block_id, Tensor(a13!) out, "
"Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
&selective_state_update_cpu_impl);
ops.def(
"mamba_chunk_scan_fwd_cpu("
"Tensor(a0!) out, Tensor(a1!) final_states, "
"Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
"Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
&mamba_chunk_scan_fwd_cpu_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels
+2
View File
@@ -47,6 +47,8 @@ torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
torch::stable::Tensor const& perm);
#ifndef USE_ROCM
std::string get_compiled_cuda_archs();
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
@@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
// AWQ zeros: (size_k // group_size, size_n // 8)
const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
int32_t group_size) {
int32_t val =
qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
int32_t zero =
qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
blockIdx.y];
// Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
// blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
// same row (stride 1) instead of striding across rows (stride size_n/8).
int col = blockIdx.y * 32 + threadIdx.x;
if (col >= size_n / 8) return;
(void)size_k;
int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
int32_t new_val = 0;
#pragma unroll
@@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
zero >>= 4;
}
output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
output[blockIdx.x * (size_n / 8) + col] = new_val;
}
torch::stable::Tensor marlin_int4_fp8_preprocess(
@@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
"qweight.size(0) % qzeros.size(0) != 0");
STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
dim3 blocks(size_k / 32, size_n / 8);
dim3 blocks(size_k, (size_n / 8 + 31) / 32);
marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
@@ -51,7 +51,8 @@ void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
#endif
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
torch::stable::Tensor const& a_tensors,
torch::stable::Tensor const& b_tensors,
@@ -83,8 +84,9 @@ void cutlass_scaled_mm_sm100(torch::stable::Tensor& c,
std::optional<torch::stable::Tensor> const& bias);
#endif
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
(defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
(defined(ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
void get_cutlass_moe_mm_data_caller(
const torch::stable::Tensor& topk_ids,
@@ -175,11 +177,14 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
// CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
// or CUDA 12.8 and SM100 (Blackwell). Only report archs that have an
// actual cutlass_moe_mm dispatch compiled into this file.
// or CUDA 12.8 and SM10x/SM11x (Blackwell / Thor). CUDA 12 reports Thor as
// SM101 while CUDA 13 reports it as SM110, but both use this sm100-named
// implementation. Only report archs that have an actual cutlass_moe_mm
// dispatch compiled into this file.
#if defined CUDA_VERSION
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
return CUDA_VERSION >= 12080;
}
@@ -281,8 +286,11 @@ void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
torch::stable::Tensor const& c_strides, bool per_act_token,
bool per_out_ch) {
int32_t version_num = get_sm_version_num();
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
if (version_num >= 100 && version_num < 110) {
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
// Keep runtime dispatch aligned with the CMake arch list and support query:
// CUDA 12 Thor is SM101 and CUDA 13 Thor is SM110.
if (version_num >= 100 && version_num < 120) {
cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
expert_offsets, problem_sizes, a_strides, b_strides,
c_strides, per_act_token, per_out_ch);
@@ -316,8 +324,9 @@ void get_cutlass_moe_mm_data(
// This function currently gets compiled only if we have a valid cutlass moe
// mm to run it for.
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
problem_sizes2, input_permutation,
@@ -338,8 +347,9 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
const bool swap_ab) {
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
@@ -362,8 +372,9 @@ void get_cutlass_batched_moe_mm_data(
// This function currently gets compiled only if we have a valid cutlass moe
// mm to run it for.
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
problem_sizes2, expert_num_tokens,
+6
View File
@@ -4,6 +4,10 @@
#include <torch/csrc/stable/library.h>
#ifndef USE_ROCM
std::string get_compiled_cuda_archs() { return VLLM_COMPILED_CUDA_ARCHS; }
#endif
// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
// Note: We register under namespace "_C" so ops are accessible as
// torch.ops._C.<op_name> for compatibility with existing code.
@@ -30,6 +34,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
ops.def("get_compiled_cuda_archs() -> str");
#ifndef USE_ROCM
@@ -763,6 +768,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
// ops.impl("op_name", &func) without a dispatch key in the non-stable API.
STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
#ifndef USE_ROCM
ops.impl("get_compiled_cuda_archs", TORCH_BOX(&get_compiled_cuda_archs));
ops.impl("cutlass_scaled_mm_supports_fp8",
TORCH_BOX(&cutlass_scaled_mm_supports_fp8));
ops.impl("cutlass_group_gemm_supported",
+2 -2
View File
@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
+2 -2
View File
@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:
- Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu)
- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
### SiLU+Mul + Quantization (`fuse_act_quant`)
@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
- Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
### RMSNorm + Padding (`fuse_act_padding`)
+4 -3
View File
@@ -68,13 +68,14 @@ vllm serve <model> \
| --- | --- | --- | --- | --- |
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
## Secondary Tiers
@@ -179,7 +180,7 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
@@ -46,6 +46,26 @@ export CPU_ARCH=$(uname -m) # x86_64 or aarch64
uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
```
!!! warning "CUDA architecture coverage"
Pre-built CUDA wheels are compiled for the CUDA architectures selected by
vLLM's release and build pipelines. This list is intentionally smaller than
every architecture CMake can build, because each additional architecture
increases wheel size.
In particular, CUDA 12.9 wheels do not use CUDA 13 family-specific targets.
To keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
architecture-specific Blackwell/Thor targets, such as `sm_103` or
`sm_121`, even though vLLM can build them from source. CUDA 13 wheels use
family-specific targets such as `sm_100f`, `sm_110f`, and `sm_120f`, which
cover the corresponding major-version GPU family.
If vLLM logs a warning that your visible CUDA device is not covered by the
wheel's compiled CUDA architectures, or if you see a CUDA error such as
`no kernel image is available for execution on the device`, install a CUDA
13 wheel when possible, or build from source with a `TORCH_CUDA_ARCH_LIST`
that includes your GPU.
#### Install the latest code
LLM inference is a fast-evolving field, and the latest code may contain bug fixes, performance improvements, and new features that are not released yet. To allow users to try the latest code without waiting for the next release, vLLM provides wheels for every commit since `v0.5.3` on <https://wheels.vllm.ai/nightly>. There are multiple indices that could be used:
@@ -31,10 +31,8 @@
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | |
| 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | |
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | |
| meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | |
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | |
| Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | |
| RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | |
@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
- `response_format`: Format of the response ("json", "text") (optional)
- `temperature`: Sampling temperature between 0 and 1 (optional)
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182).
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).
**Response Format:**
+27
View File
@@ -174,6 +174,33 @@ If the script runs successfully, you should see the message `sanity check is suc
If the test script hangs or crashes, usually it means the hardware/drivers are broken in some sense. You should try to contact your system administrator or hardware vendor for further assistance. As a common workaround, you can try to tune some NCCL environment variables, such as `export NCCL_P2P_DISABLE=1` to see if it helps. Please check [their documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) for more information. Please only use these environment variables as a temporary workaround, as they might affect the performance of the system. The best solution is still to fix the hardware/drivers so that the test script can run successfully.
## CUDA architecture not covered by the wheel
If vLLM logs a warning that the current wheel was built for a set of CUDA
architectures but one of your visible CUDA devices is not covered, the installed
wheel may not contain native CUDA kernels for that GPU. The same issue may also
surface later as a CUDA runtime error such as `no kernel image is available for
execution on the device`.
The architecture list in a pre-built wheel is determined by the vLLM release and
build pipelines, then filtered by CMake before individual kernels choose their
own per-kernel architectures. It is not the same as the full set of
architectures that vLLM can build from source.
CUDA 12.9 wheels use architecture-specific targets for newer NVIDIA GPUs. To
keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
Blackwell/Thor targets such as `sm_103` or `sm_121`, even though vLLM can build
them from source. CUDA 13 wheels can use family-specific targets such as
`sm_100f`, `sm_110f`, and `sm_120f`, which cover the corresponding
major-version GPU family.
If you see this warning or error, install a CUDA 13 wheel when possible.
Otherwise, build vLLM from source and set `TORCH_CUDA_ARCH_LIST` to include
your GPU's compute capability. See the CUDA installation guide's
[pre-built wheels](../getting_started/installation/gpu.md#pre-built-wheels) and
[build wheel from source](../getting_started/installation/gpu.md#build-wheel-from-source)
sections for details.
## Python multiprocessing
### `RuntimeError` Exception
+3 -2
View File
@@ -5560,6 +5560,7 @@ dependencies = [
"tracing",
"tracing-subscriber",
"uuid",
"vllm-bench",
"vllm-chat",
"vllm-engine-core-client",
"vllm-managed-engine",
@@ -6320,9 +6321,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
[[package]]
name = "xgrammar-structural-tag"
version = "0.1.0+xgrammar.0.2.2.4d145cc"
version = "0.2.0+xgrammar.0.2.4.dd729e7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6"
checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
dependencies = [
"auto_impl",
"serde",
+2 -1
View File
@@ -132,6 +132,7 @@ trait-set = "0.3.0"
url = "2.5.7"
uuid = { version = "1.22.0", features = ["v4"] }
validator = { version = "0.20.0", features = ["derive"] }
vllm-bench = { path = "src/bench" }
vllm-chat = { path = "src/chat" }
vllm-engine-core-client = { path = "src/engine-core-client" }
vllm-llm = { path = "src/llm" }
@@ -142,7 +143,7 @@ vllm-server = { path = "src/server" }
vllm-text = { path = "src/text" }
vllm-tokenizer = { path = "src/tokenizer" }
winnow = { version = "1.0.2", features = ["simd"] }
xgrammar-structural-tag = "0.1.0"
xgrammar-structural-tag = "0.2.0"
zeromq = { version = "0.6.0", default-features = false, features = [
"tokio-runtime",
"all-transport",
+4 -11
View File
@@ -3,8 +3,6 @@
use std::fmt;
use clap::Parser;
/// Backend type for the benchmark endpoint.
#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendKind {
@@ -77,7 +75,7 @@ pub enum DatasetName {
ShareGpt,
#[value(name = "sonnet")]
Sonnet,
#[value(name = "speed-bench")]
#[value(name = "speed-bench", alias = "speed_bench")]
SpeedBench,
#[value(name = "hf")]
Hf,
@@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig {
}
/// High-performance benchmark client for vLLM serving endpoints.
#[derive(Parser, Debug, Clone)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
pub struct Cli {
#[derive(clap::Args, Debug, Clone)]
pub struct BenchServeArgs {
/// The type of backend or endpoint to use for the benchmark.
#[arg(long, default_value = "openai")]
pub backend: BackendKind,
@@ -659,7 +652,7 @@ pub struct Cli {
pub lora_assignment: LoraAssignment,
}
impl Cli {
impl BenchServeArgs {
/// Resolve the base URL from explicit --base-url or from --host/--port.
pub fn resolve_base_url(&self) -> String {
if let Some(ref base) = self.base_url {
+212 -188
View File
@@ -4,7 +4,9 @@
use std::collections::HashMap;
use std::sync::Arc;
use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
use crate::cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
use crate::error::{BenchError, Result};
@@ -215,63 +217,63 @@ pub struct BenchConfig {
}
impl BenchConfig {
pub fn from_cli(cli: &Cli) -> Result<Self> {
if cli.burstiness <= 0.0 {
pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
if args.burstiness <= 0.0 {
return Err(BenchError::Config("Burstiness must be positive".into()));
}
if cli.num_prompts == 0 {
if args.num_prompts == 0 {
return Err(BenchError::Config(
"--num-prompts must be at least 1".into(),
));
}
if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
return Err(BenchError::Config(
"--request-rate must be positive (or inf)".into(),
));
}
if cli.max_model_len == Some(0) {
if args.max_model_len == Some(0) {
return Err(BenchError::Config(
"--max-model-len must be at least 1".into(),
));
}
let base_url = cli.resolve_base_url();
let api_url = cli.resolve_api_url();
let base_url = args.resolve_base_url();
let api_url = args.resolve_api_url();
let extra_headers = cli.parse_headers()?;
let mut extra_body = cli.parse_extra_body()?;
let extra_headers = args.parse_headers()?;
let mut extra_body = args.parse_extra_body()?;
// Merge sampling parameters into extra_body (matches Python behavior).
// Python collects non-None sampling params and merges them UNDER extra_body,
// meaning extra_body keys take precedence over sampling params.
{
let mut sampling_params = serde_json::Map::new();
if let Some(v) = cli.top_p {
if let Some(v) = args.top_p {
sampling_params.insert("top_p".into(), serde_json::json!(v));
}
if let Some(v) = cli.top_k {
if let Some(v) = args.top_k {
sampling_params.insert("top_k".into(), serde_json::json!(v));
}
if let Some(v) = cli.min_p {
if let Some(v) = args.min_p {
sampling_params.insert("min_p".into(), serde_json::json!(v));
}
if let Some(v) = cli.temperature {
if let Some(v) = args.temperature {
sampling_params.insert("temperature".into(), serde_json::json!(v));
}
if let Some(v) = cli.frequency_penalty {
if let Some(v) = args.frequency_penalty {
sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
}
if let Some(v) = cli.presence_penalty {
if let Some(v) = args.presence_penalty {
sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
}
if let Some(v) = cli.repetition_penalty {
if let Some(v) = args.repetition_penalty {
sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
}
if !sampling_params.is_empty() {
if !cli.backend.is_openai_compatible() {
if !args.backend.is_openai_compatible() {
return Err(BenchError::Config(
"Sampling parameters are only supported by openai-compatible backends."
.into(),
@@ -299,7 +301,7 @@ impl BenchConfig {
}
// Parse metadata
let metadata = match &cli.metadata {
let metadata = match &args.metadata {
None => None,
Some(items) => {
let mut pairs = Vec::new();
@@ -314,24 +316,24 @@ impl BenchConfig {
};
// Parse goodput SLOs
let goodput = parse_goodput(&cli.goodput)?;
let goodput = parse_goodput(&args.goodput)?;
// Parse ramp-up config
let ramp_up = parse_ramp_up(cli)?;
let ramp_up = parse_ramp_up(args)?;
// Default percentile metrics based on backend type
let default_percentile_metrics = if cli.backend.is_pooling() {
let default_percentile_metrics = if args.backend.is_pooling() {
"e2el"
} else {
"ttft,tpot,itl,e2el"
};
let percentile_metrics_str =
cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
let selected_percentile_metrics: Vec<String> =
percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
let sweep_summary_percentiles = cli
let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
let sweep_summary_percentiles = args
.sweep_summary_percentiles
.as_deref()
.map(|raw| parse_percentiles(raw, true))
@@ -344,38 +346,38 @@ impl BenchConfig {
selected_percentiles.push(90.0);
}
let tokenizer_id = if cli.skip_tokenizer_init {
let tokenizer_id = if args.skip_tokenizer_init {
None
} else {
Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
args.tokenizer.clone().or_else(|| args.model.clone())
};
// Resolve input/output lengths
let random_input_len = cli.resolved_random_input_len();
let random_output_len = cli.resolved_random_output_len();
let per_turn_input_len = cli.resolved_per_turn_input_len();
let random_input_len = args.resolved_random_input_len();
let random_output_len = args.resolved_random_output_len();
let per_turn_input_len = args.resolved_per_turn_input_len();
// Normalized multi-turn turn counts (computed in validation block below, defaults
// to num_turns if multi-turn mode is not active)
let mut multi_turn_min_turns = cli.multi_turn_num_turns;
let mut multi_turn_max_turns = cli.multi_turn_num_turns;
let mut multi_turn_min_turns = args.multi_turn_num_turns;
let mut multi_turn_max_turns = args.multi_turn_num_turns;
// For random datasets with openai-compatible backends, default to ignore_eos.
// Exception: multi-turn mode, where ignore_eos causes unbounded context growth
// across turns. Multi-turn uses min_tokens instead for output length control.
// Pooling backends don't generate tokens, so ignore_eos is irrelevant.
let ignore_eos = if cli.backend.is_pooling() {
let ignore_eos = if args.backend.is_pooling() {
false
} else {
cli.ignore_eos
|| ((cli.dataset_name == DatasetName::Random
|| cli.dataset_name == DatasetName::RandomMm)
&& cli.backend.is_openai_compatible()
&& !cli.multi_turn)
args.ignore_eos
|| ((args.dataset_name == DatasetName::Random
|| args.dataset_name == DatasetName::RandomMm)
&& args.backend.is_openai_compatible()
&& !args.multi_turn)
};
// Pooling backends don't support multi-turn
if cli.backend.is_pooling() && cli.multi_turn {
if args.backend.is_pooling() && args.multi_turn {
return Err(BenchError::Config(
"Pooling/embedding backends do not support --multi-turn".into(),
));
@@ -383,7 +385,7 @@ impl BenchConfig {
// LoRA validation. Adapter names must be non-empty after trim; pooling
// backends are out of scope (vLLM LoRA routing is for generative paths).
let lora_modules = match cli.lora_modules.as_ref() {
let lora_modules = match args.lora_modules.as_ref() {
None => None,
Some(names) => {
if names.is_empty() {
@@ -391,7 +393,7 @@ impl BenchConfig {
"--lora-modules requires at least one adapter name".into(),
));
}
if cli.backend.is_pooling() {
if args.backend.is_pooling() {
return Err(BenchError::Config(
"--lora-modules is not supported for pooling/embedding backends".into(),
));
@@ -411,18 +413,18 @@ impl BenchConfig {
};
// Random-MM validation and config parsing
let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
if cli.backend != BackendKind::OpenaiChat {
let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
if args.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config(
"Multi-modal content (images) is only supported on 'openai-chat' backend."
.into(),
));
}
let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
&cli.random_mm_limit_mm_per_prompt,
&args.random_mm_limit_mm_per_prompt,
)?;
let buckets =
crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
(limit, buckets)
} else {
(MmLimitPerPrompt::default(), Vec::new())
@@ -432,18 +434,18 @@ impl BenchConfig {
// sonnet (uses built-in Shakespeare's sonnets).
// Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;
// Batched inputs only make sense for pooling backends (the generation
// backends send one prompt per request).
if cli.random_batch_size == 0 {
if args.random_batch_size == 0 {
return Err(BenchError::Config(
"--random-batch-size must be at least 1".into(),
));
}
if cli.random_batch_size > 1
&& !cli.backend.is_pooling()
&& cli.dataset_name != DatasetName::RandomRerank
if args.random_batch_size > 1
&& !args.backend.is_pooling()
&& args.dataset_name != DatasetName::RandomRerank
{
return Err(BenchError::Config(
"--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -451,16 +453,16 @@ impl BenchConfig {
}
// random-rerank validation (mirrors Python RandomDatasetForReranking)
let is_reranker = !cli.no_reranker;
if cli.dataset_name == DatasetName::RandomRerank {
if !cli.backend.is_pooling() {
let is_reranker = !args.no_reranker;
if args.dataset_name == DatasetName::RandomRerank {
if !args.backend.is_pooling() {
return Err(BenchError::Config(
"--dataset-name random-rerank requires an embeddings/pooling backend \
(e.g. --backend vllm-rerank)"
.into(),
));
}
if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
return Err(BenchError::Config(
"--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
(the query is folded into the first batch slot)"
@@ -470,8 +472,8 @@ impl BenchConfig {
}
// Custom dataset validation
if cli.dataset_name == DatasetName::Custom {
match cli.dataset_path.as_deref() {
if args.dataset_name == DatasetName::Custom {
match args.dataset_path.as_deref() {
None => {
return Err(BenchError::Config(
"--dataset-path is required for --dataset-name custom \
@@ -486,7 +488,7 @@ impl BenchConfig {
}
_ => {}
}
if !cli.skip_chat_template {
if !args.skip_chat_template {
eprintln!(
"NOTE: client-side chat template rendering is not supported; custom \
dataset prompts are sent raw (equivalent to --skip-chat-template)."
@@ -495,29 +497,29 @@ impl BenchConfig {
}
// Prefix repetition validation
if cli.dataset_name == DatasetName::PrefixRepetition {
if cli.prefix_repetition_num_prefixes == 0 {
if args.dataset_name == DatasetName::PrefixRepetition {
if args.prefix_repetition_num_prefixes == 0 {
return Err(BenchError::Config(
"--prefix-repetition-num-prefixes must be at least 1".into(),
));
}
if cli.num_prompts < cli.prefix_repetition_num_prefixes {
if args.num_prompts < args.prefix_repetition_num_prefixes {
return Err(BenchError::Config(format!(
"--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
cli.num_prompts, cli.prefix_repetition_num_prefixes
args.num_prompts, args.prefix_repetition_num_prefixes
)));
}
}
// HF dataset validation
if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
return Err(BenchError::Config(
"--dataset-path is required for --dataset-name hf \
(set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
.into(),
));
}
if let Some(len) = cli.hf_output_len
if let Some(len) = args.hf_output_len
&& len == 0
{
return Err(BenchError::Config(
@@ -526,13 +528,13 @@ impl BenchConfig {
}
// Multi-turn validation
if cli.multi_turn {
if cli.backend != BackendKind::OpenaiChat {
if args.multi_turn {
if args.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config(
"--multi-turn requires --backend openai-chat".into(),
));
}
if cli.multi_turn_num_turns == 0 {
if args.multi_turn_num_turns == 0 {
return Err(BenchError::Config(
"--multi-turn-num-turns must be at least 1".into(),
));
@@ -541,18 +543,18 @@ impl BenchConfig {
// Normalize and validate min/max turns. ShareGPT only consumes max_turns
// (the loader walks all available turns up to the cap), so the
// min/num/max coupling used for synthetic generation does not apply.
if cli.dataset_name == DatasetName::ShareGpt {
if cli.multi_turn_max_turns == 1 {
if args.dataset_name == DatasetName::ShareGpt {
if args.multi_turn_max_turns == 1 {
return Err(BenchError::Config(
"--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
));
}
} else {
(multi_turn_min_turns, multi_turn_max_turns) =
match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
(0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
(m, 0) => (m, cli.multi_turn_num_turns),
(0, x) => (cli.multi_turn_num_turns, x),
match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
(0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
(m, 0) => (m, args.multi_turn_num_turns),
(0, x) => (args.multi_turn_num_turns, x),
(m, x) => (m, x),
};
if multi_turn_min_turns < 1 {
@@ -575,8 +577,8 @@ impl BenchConfig {
}
// Validate prefix sharing ratios
let pg = cli.multi_turn_prefix_global_ratio;
let pc = cli.multi_turn_prefix_conversation_ratio;
let pg = args.multi_turn_prefix_global_ratio;
let pc = args.multi_turn_prefix_conversation_ratio;
if !(0.0..=1.0).contains(&pg) {
return Err(BenchError::Config(
"--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
@@ -592,20 +594,20 @@ impl BenchConfig {
"--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
));
}
if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
return Err(BenchError::Config(
"Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
));
}
}
if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
return Err(BenchError::Config(format!(
"--steady-state-threshold must be in (0.0, 1.0], got {}",
cli.steady_state_threshold
args.steady_state_threshold
)));
}
if let Some(mw) = cli.steady_state_min_window
if let Some(mw) = args.steady_state_min_window
&& mw < 0.0
{
return Err(BenchError::Config(format!(
@@ -613,122 +615,122 @@ impl BenchConfig {
)));
}
if cli.profile_batch_threshold.is_some() && !cli.profile {
if args.profile_batch_threshold.is_some() && !args.profile {
return Err(BenchError::Config(
"--profile-batch-threshold requires --profile".into(),
));
}
if cli.profile_duration <= 0.0 {
if args.profile_duration <= 0.0 {
return Err(BenchError::Config(
"--profile-duration must be positive".into(),
));
}
if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
return Err(BenchError::Config(
"--profile-duration requires --profile-batch-threshold".into(),
));
}
Ok(BenchConfig {
backend: cli.backend,
backend: args.backend,
base_url,
api_url,
model: cli.model.clone(),
model_name: cli.served_model_name.clone(),
model: args.model.clone(),
model_name: args.served_model_name.clone(),
tokenizer_id,
tokenizer_mode: cli.tokenizer_mode.clone(),
trust_remote_code: cli.trust_remote_code,
skip_tokenizer_init: cli.skip_tokenizer_init,
dataset_name: cli.dataset_name,
dataset_path: cli.dataset_path.clone(),
max_model_len: cli.max_model_len,
tokenizer_mode: args.tokenizer_mode.clone(),
trust_remote_code: args.trust_remote_code,
skip_tokenizer_init: args.skip_tokenizer_init,
dataset_name: args.dataset_name,
dataset_path: args.dataset_path.clone(),
max_model_len: args.max_model_len,
random_input_len,
random_output_len,
random_prefix_len: cli.random_prefix_len,
random_prefix_len: args.random_prefix_len,
random_range_ratio,
random_batch_size: cli.random_batch_size,
random_batch_size: args.random_batch_size,
is_reranker,
custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
prefix_repetition_output_len: cli
custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
prefix_repetition_output_len: args
.output_len
.unwrap_or(cli.prefix_repetition_output_len),
random_cache_hit_fraction: cli.random_cache_hit_fraction,
random_cache_ratio: cli.random_cache_ratio,
sharegpt_output_len: cli.sharegpt_output_len,
sonnet_input_len: cli.sonnet_input_len,
sonnet_output_len: cli.sonnet_output_len,
sonnet_prefix_len: cli.sonnet_prefix_len,
no_oversample: cli.no_oversample,
disable_shuffle: cli.disable_shuffle,
num_prompts: cli.num_prompts,
request_rate: cli.request_rate,
burstiness: cli.burstiness,
max_concurrency: cli.max_concurrency,
steady_state_threshold: cli.steady_state_threshold,
steady_state_min_window: cli.steady_state_min_window,
no_steady_state: cli.no_steady_state,
disable_tqdm: cli.disable_tqdm,
num_warmups: cli.num_warmups,
profile: cli.profile,
profile_batch_threshold: cli.profile_batch_threshold,
profile_duration: cli.profile_duration,
save_result: cli.save_result,
save_detailed: cli.save_detailed,
append_result: cli.append_result,
result_dir: cli.result_dir.clone(),
result_filename: cli.result_filename.clone(),
seed: cli.seed,
.unwrap_or(args.prefix_repetition_output_len),
random_cache_hit_fraction: args.random_cache_hit_fraction,
random_cache_ratio: args.random_cache_ratio,
sharegpt_output_len: args.sharegpt_output_len,
sonnet_input_len: args.sonnet_input_len,
sonnet_output_len: args.sonnet_output_len,
sonnet_prefix_len: args.sonnet_prefix_len,
no_oversample: args.no_oversample,
disable_shuffle: args.disable_shuffle,
num_prompts: args.num_prompts,
request_rate: args.request_rate,
burstiness: args.burstiness,
max_concurrency: args.max_concurrency,
steady_state_threshold: args.steady_state_threshold,
steady_state_min_window: args.steady_state_min_window,
no_steady_state: args.no_steady_state,
disable_tqdm: args.disable_tqdm,
num_warmups: args.num_warmups,
profile: args.profile,
profile_batch_threshold: args.profile_batch_threshold,
profile_duration: args.profile_duration,
save_result: args.save_result,
save_detailed: args.save_detailed,
append_result: args.append_result,
result_dir: args.result_dir.clone(),
result_filename: args.result_filename.clone(),
seed: args.seed,
ignore_eos,
insecure: cli.insecure,
insecure: args.insecure,
selected_percentile_metrics,
selected_percentiles,
sweep_summary_percentiles,
label: cli.label.clone(),
logprobs: cli.logprobs,
request_id_prefix: cli.get_request_id_prefix(),
ready_check_timeout_sec: cli.ready_check_timeout_sec,
label: args.label.clone(),
logprobs: args.logprobs,
request_id_prefix: args.get_request_id_prefix(),
ready_check_timeout_sec: args.ready_check_timeout_sec,
extra_headers,
extra_body,
metadata,
dry_run: cli.dry_run,
dry_run: args.dry_run,
goodput,
ramp_up,
multi_turn: cli.multi_turn,
multi_turn_num_turns: cli.multi_turn_num_turns,
multi_turn: args.multi_turn,
multi_turn_num_turns: args.multi_turn_num_turns,
multi_turn_min_turns,
multi_turn_max_turns,
sharegpt_multi_turn_max_turns: if cli.multi_turn
&& cli.dataset_name == DatasetName::ShareGpt
&& cli.multi_turn_max_turns != 0
sharegpt_multi_turn_max_turns: if args.multi_turn
&& args.dataset_name == DatasetName::ShareGpt
&& args.multi_turn_max_turns != 0
{
Some(cli.multi_turn_max_turns)
Some(args.multi_turn_max_turns)
} else {
None
},
per_turn_input_len,
multi_turn_concurrency: cli.multi_turn_concurrency,
multi_turn_delay_ms: cli.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
speed_bench_config: cli.speed_bench_config,
speed_bench_category: cli.speed_bench_category.clone(),
speed_bench_max_input_len: cli.speed_bench_max_input_len,
hf_split: cli.hf_split.clone(),
hf_subset: cli.hf_subset.clone(),
hf_output_len: cli.hf_output_len,
hf_text_column: cli.hf_text_column.clone(),
reset_prefix_cache: cli.reset_prefix_cache,
prompt_token_ids: cli.prompt_token_ids,
random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
multi_turn_concurrency: args.multi_turn_concurrency,
multi_turn_delay_ms: args.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
speed_bench_config: args.speed_bench_config,
speed_bench_category: args.speed_bench_category.clone(),
speed_bench_max_input_len: args.speed_bench_max_input_len,
hf_split: args.hf_split.clone(),
hf_subset: args.hf_subset.clone(),
hf_output_len: args.hf_output_len,
hf_text_column: args.hf_text_column.clone(),
reset_prefix_cache: args.reset_prefix_cache,
prompt_token_ids: args.prompt_token_ids,
random_mm_base_items_per_request: args.random_mm_base_items_per_request,
random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
random_mm_limit,
random_mm_buckets,
enable_multimodal_chat: cli.enable_multimodal_chat,
enable_multimodal_chat: args.enable_multimodal_chat,
lora_modules,
lora_assignment: cli.lora_assignment,
lora_assignment: args.lora_assignment,
})
}
}
@@ -811,17 +813,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
Ok(config)
}
fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
let strategy = match cli.ramp_up_strategy {
fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
let strategy = match args.ramp_up_strategy {
None => return Ok(None),
Some(s) => s,
};
let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
})?;
let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
})?;
@@ -843,7 +845,21 @@ mod tests {
use clap::Parser;
use super::*;
use crate::cli::Cli;
use crate::cli::BenchServeArgs;
#[derive(Parser)]
struct TestCli {
#[command(flatten)]
args: BenchServeArgs,
}
fn parse_args<I, T>(args: I) -> BenchServeArgs
where
I: IntoIterator<Item = T>,
T: Into<std::ffi::OsString> + Clone,
{
TestCli::parse_from(args).args
}
fn base_multi_turn_args() -> Vec<&'static str> {
vec![
@@ -859,8 +875,8 @@ mod tests {
#[test]
fn test_prefix_sharing_defaults_to_zero() {
let args = base_multi_turn_args();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
}
@@ -874,8 +890,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.8",
]);
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
}
@@ -889,8 +905,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.6",
]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -902,16 +918,16 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.5",
]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
fn test_prefix_sharing_out_of_range_fails() {
let mut args = base_multi_turn_args();
args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -928,8 +944,8 @@ mod tests {
"--multi-turn-prefix-global-ratio",
"0.1",
];
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -944,8 +960,8 @@ mod tests {
"--dataset-name",
"sharegpt",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.multi_turn_max_turns, 3);
assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -968,8 +984,8 @@ mod tests {
"--multi-turn-max-turns",
"2",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
}
@@ -987,8 +1003,8 @@ mod tests {
"--multi-turn-max-turns",
"1",
];
let cli = Cli::parse_from(args);
let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
let args = parse_args(args);
let err = BenchConfig::from_args(&args).unwrap_err().to_string();
assert!(
err.contains("at least 2 for ShareGPT"),
"expected ShareGPT-specific error, got: {err}"
@@ -1009,8 +1025,8 @@ mod tests {
"--multi-turn-max-turns",
"20",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
}
@@ -1018,8 +1034,8 @@ mod tests {
#[test]
fn test_sweep_summary_percentiles_default_empty() {
let args = base_multi_turn_args();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert!(config.sweep_summary_percentiles.is_empty());
assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1034,8 +1050,8 @@ mod tests {
"--sweep-summary-percentiles",
"90,95,90",
]);
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1045,8 +1061,8 @@ mod tests {
fn test_invalid_sweep_summary_percentile_fails() {
let mut args = base_multi_turn_args();
args.extend(["--sweep-summary-percentiles", "101"]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -1058,12 +1074,20 @@ mod tests {
"--max-model-len",
"4096",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.max_model_len, Some(4096));
}
#[test]
fn test_tokenizer_id_deferred_when_model_is_unspecified() {
let args = parse_args(["vllm-bench"]);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.tokenizer_id, None);
}
#[test]
fn test_zero_max_model_len_fails() {
let args = vec![
@@ -1073,9 +1097,9 @@ mod tests {
"--max-model-len",
"0",
];
let cli = Cli::parse_from(args);
let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err());
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
fn test_range_ratio_parse_float() {
+5 -2
View File
@@ -40,8 +40,11 @@ impl HubRepo {
.build()
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
rt.block_on(async move {
let api = hf_hub::api::tokio::Api::new()
.map_err(|e| format!("Failed to init HF API: {e}"))?;
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
if let Ok(token) = std::env::var("HF_TOKEN") {
builder = builder.with_token(Some(token));
}
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
})
})
+86
View File
@@ -0,0 +1,86 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
use anyhow::Context;
pub use cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use config::BenchConfig;
/// Prepare process-wide resources for a benchmark run.
pub fn prepare_process() {
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
}
}
/// Run the online serving benchmark.
pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = args.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_args(&args).context("Configuration error")?;
async {
if config.multi_turn {
if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
args.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
} else if let Some(ref sweep_rate) = args.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if args.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, args.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
}
.await
.context("Benchmark failed")
}
+14 -74
View File
@@ -1,92 +1,32 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
use anyhow::Context;
use clap::Parser;
use cli::Cli;
use config::BenchConfig;
#[derive(Parser)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
struct Cli {
#[command(flatten)]
args: vllm_bench::BenchServeArgs,
}
fn main() -> anyhow::Result<()> {
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
}
let cli = Cli::parse();
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = cli.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
vllm_bench::prepare_process();
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.expect("Failed to build tokio runtime");
.context("Failed to build tokio runtime")?;
runtime
.block_on(async {
if config.multi_turn {
if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
cli.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
.await?;
} else if let Some(ref sweep_rate) = cli.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if cli.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, cli.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
})
.context("Benchmark failed")
runtime.block_on(vllm_bench::run(cli.args))
}
+4 -4
View File
@@ -38,7 +38,7 @@ pub(super) fn build_batched_items(
let keep_on_cpu = spec.keep_on_cpu_keys.contains(key);
let (value, field) = match spec.field_layout_for(key) {
Some(FieldLayout::Batched) => (
tensor.batched_value_at(index)?,
tensor.batched_wire_value_at(index)?,
MmField::Batched(MmBatchedField { keep_on_cpu }),
),
Some(FieldLayout::Flat { sizes_key }) => {
@@ -47,7 +47,7 @@ pub(super) fn build_batched_items(
})?;
let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
(
tensor.flat_value_range(start, end)?,
tensor.flat_wire_value_range(start, end)?,
MmField::Flat(MmFlatField {
slices: vec![MmSlice::Slice(SliceSpec {
start: Some(0),
@@ -60,7 +60,7 @@ pub(super) fn build_batched_items(
)
}
None => (
tensor.clone(),
tensor.try_into()?,
MmField::Shared(MmSharedField {
batch_size: len,
keep_on_cpu,
@@ -71,7 +71,7 @@ pub(super) fn build_batched_items(
data.insert(
key.clone(),
MmFieldElem {
data: Some(value.try_into()?),
data: Some(value),
field,
},
);
+68 -81
View File
@@ -12,7 +12,7 @@ use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
use crate::error::{Error, Result, bail_multimodal, multimodal};
/// Representation for multimodal kwarg values for transformation.
#[derive(Debug, Clone)]
#[derive(Debug)]
pub(super) enum KwargValue {
/// Float tensor with row-major flat data and shape.
F32Tensor { data: Vec<f32>, shape: Vec<usize> },
@@ -107,28 +107,19 @@ impl KwargValue {
}
}
impl TryFrom<KwargValue> for ProtocolKwargValue {
impl TryFrom<&KwargValue> for ProtocolKwargValue {
type Error = Error;
fn try_from(value: KwargValue) -> Result<Self> {
match value {
KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::Passthrough(value) => Ok(value),
}
fn try_from(value: &KwargValue) -> Result<Self> {
let tensor = match value {
KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
KwargValue::Passthrough(value) => return Ok(value.clone()),
};
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
}
}
@@ -145,63 +136,55 @@ impl KwargValue {
}
}
/// Extract one media item from a batched tensor field.
/// Convert one media item from a batched tensor field to wire bytes.
///
/// Batched fields use their first axis as media-item index and drop that
/// axis in the per-feature value, matching vLLM's batched-field semantics.
pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> {
match self {
Self::F32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::F32Tensor { data, shape })
}
Self::F16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::F16Tensor { data, shape })
}
Self::Bf16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::Bf16Tensor { data, shape })
}
Self::I64Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::I64Tensor { data, shape })
}
Self::U32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::U32Tensor { data, shape })
}
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
}
pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
self.wire_value_range(index, index + 1, true)
}
/// Extract one media item's variable-length range from a flat tensor field.
/// Convert one media item's flat tensor range directly to wire bytes.
///
/// Flat fields keep the first axis as the sliced length for this item.
pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> {
match self {
pub(super) fn flat_wire_value_range(
&self,
start: usize,
end: usize,
) -> Result<ProtocolKwargValue> {
self.wire_value_range(start, end, false)
}
fn wire_value_range(
&self,
start: usize,
end: usize,
drop_axis: bool,
) -> Result<ProtocolKwargValue> {
let tensor = match self {
Self::F32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::F32Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_f32(shape, data)
}
Self::F16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::F16Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_f16(shape, data)
}
Self::Bf16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::Bf16Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_bf16(shape, data)
}
Self::I64Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::I64Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_i64(shape, data)
}
Self::U32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::U32Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_u32(shape, data)
}
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
}
Self::Passthrough(value) => return Ok(value.clone()),
};
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
}
}
@@ -240,13 +223,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
}
/// Slice a flat row-major tensor along its first axis.
fn slice_first_axis_range<T: Clone>(
fn slice_first_axis_range<'a, T>(
shape: &[usize],
data: &[T],
data: &'a [T],
start: usize,
end: usize,
drop_axis: bool,
) -> Result<(Vec<usize>, Vec<T>)> {
) -> Result<(Vec<usize>, &'a [T])> {
let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
if start > end || end > first_dim {
bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
@@ -270,7 +253,7 @@ fn slice_first_axis_range<T: Clone>(
shape[0] = end - start;
shape
};
Ok((out_shape, data[data_start..data_end].to_vec()))
Ok((out_shape, &data[data_start..data_end]))
}
#[cfg(test)]
@@ -278,35 +261,39 @@ mod tests {
use super::*;
#[test]
fn batched_value_at_drops_first_axis() {
fn batched_wire_value_at_drops_first_axis() {
let value = KwargValue::F32Tensor {
data: vec![1.0, 2.0, 3.0, 4.0],
shape: vec![2, 2],
};
let value = value.batched_value_at(1).unwrap();
let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
panic!("expected tensor");
};
assert!(matches!(
value,
KwargValue::F32Tensor { data, shape }
if shape == vec![2] && data == vec![3.0, 4.0]
));
assert_eq!(tensor.shape, vec![2]);
assert_eq!(
tensor.data.into_raw_view().unwrap(),
[3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
);
}
#[test]
fn flat_value_range_keeps_first_axis() {
fn flat_wire_value_range_keeps_first_axis() {
let value = KwargValue::U32Tensor {
data: (0..10).collect(),
shape: vec![5, 2],
};
let value = value.flat_value_range(1, 3).unwrap();
let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
panic!("expected tensor");
};
assert!(matches!(
value,
KwargValue::U32Tensor { data, shape }
if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
));
assert_eq!(tensor.shape, vec![2, 2]);
assert_eq!(
tensor.data.into_raw_view().unwrap(),
[2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::<Vec<_>>()
);
}
#[test]
@@ -336,7 +323,7 @@ mod tests {
let value =
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
else {
panic!("expected tensor");
};
@@ -351,7 +338,7 @@ mod tests {
let value =
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
else {
panic!("expected tensor");
};
+4 -4
View File
@@ -130,7 +130,7 @@ fn build_video_item(
let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key);
let (value, field) = match support.spec.field_layout_for(&key) {
Some(FieldLayout::Batched) => (
tensor.batched_value_at(0)?,
tensor.batched_wire_value_at(0)?,
MmField::Batched(MmBatchedField { keep_on_cpu }),
),
Some(FieldLayout::Flat { .. }) => {
@@ -138,7 +138,7 @@ fn build_video_item(
.first_dim()
.ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?;
(
tensor,
(&tensor).try_into()?,
MmField::Flat(MmFlatField {
slices: vec![MmSlice::Slice(SliceSpec {
start: Some(0),
@@ -151,7 +151,7 @@ fn build_video_item(
)
}
None => (
tensor,
(&tensor).try_into()?,
MmField::Shared(MmSharedField {
batch_size: 1,
keep_on_cpu,
@@ -162,7 +162,7 @@ fn build_video_item(
data.insert(
key,
MmFieldElem {
data: Some(value.try_into()?),
data: Some(value),
field,
},
);
+1 -1
View File
@@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor {
Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
};
apply_structural_tag_constraint(request, parser.structural_tag_model())?;
apply_structural_tag_constraint(request, parser.structural_tag_builder())?;
if parser.preserve_special_tokens() {
request.decode_options.skip_special_tokens = false;
@@ -7,7 +7,8 @@ use thiserror_ext::AsReport;
use vllm_engine_core_client::protocol::structured_outputs::{
StructuredOutputBackend, StructuredOutputsParams,
};
use vllm_parser::tool::StructuralTagModel;
use vllm_parser::tool::StructuralTagBuilder;
use xgrammar_structural_tag::builders::StructuralTagOptions;
use xgrammar_structural_tag::{
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
build_structural_tag,
@@ -20,9 +21,9 @@ use crate::{Error, Result as ChatResult};
/// support and the request's tool choice.
pub(super) fn apply_structural_tag_constraint(
request: &mut ChatRequest,
model: Option<StructuralTagModel>,
builder: Option<&dyn StructuralTagBuilder>,
) -> ChatResult<()> {
let Some(model) = model else {
let Some(builder) = builder else {
return Ok(());
};
let Some(tool_choice) = structural_tag_tool_choice(request) else {
@@ -42,11 +43,16 @@ pub(super) fn apply_structural_tag_constraint(
})
.collect::<Vec<_>>();
let structural_tag = build_structural_tag(model, &tools, tool_choice, false)
.and_then(|tag| tag.to_json_string())
.map_err(|error| Error::StructuralTag {
message: error.to_report_string(),
})?;
let structural_tag = build_structural_tag(
builder,
&tools,
tool_choice,
StructuralTagOptions::default().with_reasoning(false),
)
.and_then(|tag| tag.to_json_string())
.map_err(|error| Error::StructuralTag {
message: error.to_report_string(),
})?;
// Overwrite any existing structured output settings with the structural tag constraint.
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
@@ -141,7 +147,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let tag = structural_tag_value(&request);
@@ -154,7 +160,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none());
@@ -169,7 +175,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let params = structured_outputs(&request);
@@ -184,7 +190,7 @@ mod tests {
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let tag = structural_tag_value(&request);
@@ -201,7 +207,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let params = structured_outputs(&request);
@@ -221,7 +227,7 @@ mod tests {
);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let tag = structural_tag_value(&request).to_string();
@@ -234,7 +240,7 @@ mod tests {
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none());
@@ -249,7 +255,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed");
let params = structured_outputs(&request);
+14 -3
View File
@@ -236,9 +236,10 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool {
loops.into_iter().any(|loop_ast| {
matches!(loop_ast.target, Expr::Var(_))
&& message_varnames
.iter()
.any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content")))
&& (is_var_access(&loop_ast.iter, "content")
|| message_varnames.iter().any(|varname| {
is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))
}))
})
}
@@ -315,6 +316,16 @@ mod tests {
);
}
#[test]
fn detects_openai_template_with_content_parameter_loop() {
assert_eq!(
detect(
"{% macro render(content) %}{% for item in content %}{{ item }}{% endfor %}{% endmacro %}{% for message in messages %}{{ render(message.content) }}{% endfor %}"
),
ChatTemplateContentFormat::OpenAi
);
}
#[test]
fn detects_openai_template_with_messages_alias() {
assert_eq!(
+20
View File
@@ -1309,6 +1309,26 @@ mod tests {
.assert_eq(&rendered);
}
#[test]
fn qwen35_template_auto_detects_openai_multimodal_content() {
let mut request = image_request();
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_mm(
QWEN3_5_0_8B_TEMPLATE,
&request,
ChatTemplateContentFormatOption::Auto,
)
.unwrap();
expect![[r#"
Text(
"<|im_start|>user\na<|vision_start|><|image_pad|><|vision_end|>b<|im_end|>\n",
)
"#]]
.assert_debug_eq(&rendered.prompt);
}
#[test]
fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() {
let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
+1
View File
@@ -29,6 +29,7 @@ tokio-util.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
uuid.workspace = true
vllm-bench.workspace = true
vllm-chat.workspace = true
vllm-engine-core-client.workspace = true
vllm-managed-engine.workspace = true
+11 -1
View File
@@ -79,13 +79,23 @@ impl Cli {
}
/// Supported top-level CLI commands.
#[derive(Debug, Subcommand, PartialEq, Eq)]
#[derive(Debug, Subcommand)]
pub enum Command {
/// Run the Rust OpenAI frontend as a Python-supervised worker.
Frontend(FrontendArgs),
/// Launch a managed Python headless engine, then run the Rust OpenAI
/// frontend.
Serve(ServeArgs),
/// Run vLLM benchmarks.
#[command(subcommand)]
Bench(BenchCommand),
}
/// Supported benchmark commands.
#[derive(Debug, Subcommand)]
pub enum BenchCommand {
/// Benchmark online serving throughput.
Serve(vllm_bench::BenchServeArgs),
}
/// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
+21 -1
View File
@@ -5,7 +5,27 @@ use expect_test::expect;
use vllm_engine_core_client::TransportMode;
use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
use super::{Cli, Command};
use super::{BenchCommand, Cli, Command};
#[test]
fn bench_serve_args_parse_without_managed_engine_repartition() {
let cli = Cli::try_parse_from([
"vllm-rs",
"bench",
"serve",
"--backend",
"openai-chat",
"--request-rate",
"inf",
])
.unwrap();
let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
panic!("expected bench serve args");
};
assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
assert!(args.request_rate.is_infinite());
}
#[test]
fn serve_args_forward_python_flags_with_separator() {
+5 -1
View File
@@ -12,7 +12,7 @@ use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use vllm_managed_engine::ManagedEngineHandle;
use crate::cli::{Cli, Command};
use crate::cli::{BenchCommand, Cli, Command};
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -100,6 +100,10 @@ fn main() -> Result<()> {
async fn async_main(cli: Cli) -> Result<()> {
match cli.command {
Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
Command::Bench(BenchCommand::Serve(bench_args)) => {
vllm_bench::prepare_process();
vllm_bench::run(bench_args).await
}
Command::Serve(args) => {
let handshake_port = args.managed_engine.resolve_handshake_port()?;
@@ -55,52 +55,57 @@ pub struct WireNdArray {
impl WireNdArray {
/// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "float32".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
})
}
/// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "float16".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
})
}
/// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "bfloat16".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
})
}
/// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "int64".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
})
}
/// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "uint32".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
})
}
+3 -3
View File
@@ -4,7 +4,7 @@
use std::sync::Arc;
use vllm_parser::tool::{
Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput,
Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput,
};
use vllm_parser::unified::{
UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
@@ -85,8 +85,8 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
self.inner.preserve_special_tokens()
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
self.inner.structural_tag_model()
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
self.inner.structural_tag_builder()
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekDsmlToolParser, DsmlTokens};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3.2 models.
///
@@ -47,8 +47,8 @@ impl ToolParser for DeepSeekV32ToolParser {
true
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::DeepSeekV32)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::DeepSeekV32.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekDsmlToolParser, DsmlTokens};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V4 models.
///
@@ -50,8 +50,8 @@ impl ToolParser for DeepSeekV4ToolParser {
true
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::DeepSeekV4)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::DeepSeekV4.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -73,7 +73,7 @@ mod tests {
use super::DeepSeekV4ToolParser;
use crate::tool::test_utils::{collect_stream, test_tools};
use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _};
use crate::tool::{ToolParser, ToolParserTestExt as _};
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
let params = params
@@ -91,13 +91,10 @@ mod tests {
}
#[test]
fn deepseek_v4_exposes_structural_tag_model() {
fn deepseek_v4_exposes_structural_tag_builder() {
let parser = DeepSeekV4ToolParser::new(&test_tools());
assert_eq!(
parser.structural_tag_model(),
Some(StructuralTagModel::DeepSeekV4)
);
assert!(parser.structural_tag_builder().is_some());
}
#[test]
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3 JSON-fenced tool calls.
///
@@ -35,8 +35,8 @@ impl ToolParser for DeepSeekV3ToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::DeepSeekR1)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::DeepSeekR1.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3.1 raw JSON tool calls.
///
@@ -31,8 +31,8 @@ impl ToolParser for DeepSeekV31ToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::DeepSeekV31)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::DeepSeekV31.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{GlmXmlToolParser, Separator};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for GLM-4.7 MoE XML-style tool calls.
///
@@ -25,8 +25,8 @@ impl ToolParser for Glm47MoeToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Glm47)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Glm47.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool};
use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALLS_START: &str = "<tool_calls>";
const TOOL_CALLS_END: &str = "</tool_calls>";
@@ -116,8 +116,8 @@ impl ToolParser for HyV3ToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::HyV3)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::HyV3.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
parser_name: "Hermes",
@@ -48,8 +48,8 @@ impl ToolParser for HermesToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Hermes)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Hermes.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+5 -3
View File
@@ -12,7 +12,9 @@ use super::{
argument_delta_event, tool_call_header_event,
};
use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{
Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput,
};
#[derive(Debug, Clone, PartialEq, Eq)]
enum LlamaJsonMode {
@@ -136,8 +138,8 @@ impl ToolParser for Llama3JsonToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Llama)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Llama.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
parser_name: "Qwen XML",
@@ -50,8 +50,8 @@ impl ToolParser for Qwen3XmlToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Qwen3)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Qwen3.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -11,7 +11,7 @@ use winnow::token::{literal, rest, take_until, take_while};
use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool};
use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
@@ -150,8 +150,8 @@ impl ToolParser for KimiK2ToolParser {
true
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Kimi)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Kimi.builder())
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
+3 -3
View File
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool};
use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALL_START: &str = "<minimax:tool_call>";
const TOOL_CALL_END: &str = "</minimax:tool_call>";
@@ -115,8 +115,8 @@ impl ToolParser for MinimaxM2ToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Minimax)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Minimax.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -36,7 +36,7 @@ pub use qwen_coder::Qwen3CoderToolParser;
pub use seed_oss::SeedOssToolParser;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub use xgrammar_structural_tag::Model as StructuralTagModel;
pub use xgrammar_structural_tag::builders::StructuralTagBuilder;
use crate::utils;
@@ -187,8 +187,8 @@ pub trait ToolParser: Send {
false
}
/// Return the xgrammar structural-tag model used for strict tool calling.
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
/// Return the xgrammar structural-tag builder used for strict tool calling.
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
None
}
+7 -10
View File
@@ -9,8 +9,8 @@ use winnow::token::{literal, take_until};
use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::Tool;
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALL_START: &str = "<tool_call>";
const TOOL_CALL_END: &str = "</tool_call>";
@@ -146,8 +146,8 @@ impl ToolParser for Qwen3CoderToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Qwen3Coder)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Qwen3Coder.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -294,7 +294,7 @@ mod tests {
use serde_json::{Value, json};
use thiserror_ext::AsReport;
use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser};
use super::{Qwen3CoderToolParser, ToolParser};
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
@@ -308,13 +308,10 @@ mod tests {
}
#[test]
fn qwen_coder_exposes_structural_tag_model() {
fn qwen_coder_exposes_structural_tag_builder() {
let parser = Qwen3CoderToolParser::new(&test_tools());
assert_eq!(
parser.structural_tag_model(),
Some(StructuralTagModel::Qwen3Coder)
);
assert!(parser.structural_tag_builder().is_some());
}
#[test]
+4 -7
View File
@@ -7,7 +7,7 @@ use vllm_tokenizer::DynTokenizer;
use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
use crate::reasoning::ReasoningParser;
use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Unified parser that composes existing reasoning and tool parsers.
pub struct CombinedParser {
@@ -79,8 +79,8 @@ impl UnifiedParser for CombinedParser {
|| self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens())
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
self.tool.as_ref().and_then(|parser| parser.structural_tag_model())
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
self.tool.as_ref().and_then(|parser| parser.structural_tag_builder())
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -269,10 +269,7 @@ mod tests {
fn combined_parser_emits_tool_calls_from_visible_content() {
let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap();
let mut parser = CombinedParser::new(None, Some(tool));
assert!(matches!(
parser.structural_tag_model(),
Some(crate::tool::StructuralTagModel::Qwen3)
));
assert!(parser.structural_tag_builder().is_some());
let output = collect(
&mut parser,
+3 -3
View File
@@ -16,7 +16,7 @@ use vllm_tokenizer::DynTokenizer;
use crate::reasoning::ReasoningError;
use crate::tool::{
StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
};
/// Result alias for unified parser operations.
@@ -171,8 +171,8 @@ pub trait UnifiedParser: Send {
false
}
/// Return the xgrammar structural-tag model used for strict tool calling.
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
/// Return the xgrammar structural-tag builder used for strict tool calling.
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
None
}
@@ -238,13 +238,18 @@ fn collect_generate(
None
};
let prompt_logprobs = if include_prompt_logprobs {
let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
ApiError::server_error(
"raw generate response requested prompt_logprobs but generation returned none"
.to_string(),
)
})?;
Some(raw_prompt_logprobs_to_maps(prompt_logprobs))
match collected.prompt_logprobs.as_ref() {
Some(prompt_logprobs) => Some(raw_prompt_logprobs_to_maps(prompt_logprobs)),
// A single-token prompt has no scored positions; same mapping
// as /v1/completions.
None if collected.prompt_token_ids.len() == 1 => Some(vec![None]),
None => {
return Err(ApiError::server_error(
"raw generate response requested prompt_logprobs but generation returned none"
.to_string(),
));
}
}
} else {
None
};
@@ -472,4 +477,48 @@ mod tests {
Some(2)
);
}
#[test]
fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
request_id: "raw-1".to_string(),
prompt_logprobs: None,
token_ids: vec![3],
logprobs: None,
finish_reason: FinishReason::stop_eos(),
usage: vllm_llm::TokenUsage {
prompt_token_count: prompt_token_ids.len(),
output_token_count: 1,
cached_token_count: 0,
},
kv_transfer_params: None,
ec_transfer_params: None,
prompt_token_ids,
};
let response = collect_generate(
output_without_payload(vec![9707]),
"raw-1".to_string(),
ApiServerOptions::default(),
ResponseOptions {
include_prompt_logprobs: true,
..Default::default()
},
)
.expect("single-token prompt without payload maps to [None]");
let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
assert_eq!(prompt_logprobs.len(), 1);
assert!(prompt_logprobs[0].is_none());
collect_generate(
output_without_payload(vec![9707, 11]),
"raw-2".to_string(),
ApiServerOptions::default(),
ResponseOptions {
include_prompt_logprobs: true,
..Default::default()
},
)
.expect_err("multi-token prompt without payload is an engine failure");
}
}
@@ -35,7 +35,7 @@ use crate::routes::openai::chat_completions::types::{
ChatMessageDelta,
};
use crate::routes::openai::utils::logprobs::{
decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps,
decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
};
use crate::routes::openai::utils::types::{
ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
@@ -181,14 +181,11 @@ async fn collect_chat_completion(
None
};
let prompt_logprobs = if include_prompt_logprobs {
Some(decoded_prompt_logprobs_to_maps(
prompt_logprobs.as_ref().ok_or_else(|| {
server_error!(
"chat response requested prompt_logprobs but generation returned none"
)
})?,
Some(prompt_logprobs_to_maps(
prompt_logprobs.as_ref(),
&prompt_token_ids,
return_tokens_as_token_ids,
))
)?)
} else {
None
};
@@ -5,7 +5,6 @@ mod convert;
mod types;
mod validate;
use std::collections::HashMap;
use std::convert::Infallible;
use std::result::Result;
use std::sync::Arc;
@@ -29,8 +28,8 @@ use vllm_text::{
use self::convert::{ResponseOptions, prepare_completion_request};
use super::utils::logprobs::{
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
decoded_prompt_logprobs_to_openai, text_len,
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai,
prompt_logprobs_to_maps, text_len,
};
use super::utils::types::Usage;
use crate::config::ApiServerOptions;
@@ -505,27 +504,6 @@ fn prompt_only_logprobs_to_openai(
))
}
fn prompt_logprobs_to_maps(
prompt_logprobs: Option<&DecodedPromptLogprobs>,
prompt_token_ids: &[u32],
return_tokens_as_token_ids: bool,
) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
if let Some(prompt_logprobs) = prompt_logprobs {
return Ok(decoded_prompt_logprobs_to_maps(
prompt_logprobs,
return_tokens_as_token_ids,
));
}
if let [_token_id] = prompt_token_ids {
return Ok(vec![None]);
}
Err(server_error!(
"completion response requested prompt_logprobs but generation returned none"
))
}
fn usage_chunk(
request_id: &str,
response_model: &str,
@@ -100,20 +100,31 @@ pub fn decoded_prompt_logprobs_to_openai(
})
}
/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response
/// shape.
pub fn decoded_prompt_logprobs_to_maps(
prompt_logprobs: &DecodedPromptLogprobs,
/// Map decoded prompt logprobs into vLLM-style per-position maps, treating a
/// missing single-token payload as `[None]`.
pub fn prompt_logprobs_to_maps(
prompt_logprobs: Option<&DecodedPromptLogprobs>,
prompt_token_ids: &[u32],
return_tokens_as_token_ids: bool,
) -> Vec<Option<HashMap<String, f32>>> {
std::iter::once(None)
.chain(prompt_logprobs.scored_positions.iter().map(|position| {
Some(position_top_logprobs_map(
position,
return_tokens_as_token_ids,
))
}))
.collect()
) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
if let Some(prompt_logprobs) = prompt_logprobs {
return Ok(std::iter::once(None)
.chain(prompt_logprobs.scored_positions.iter().map(|position| {
Some(position_top_logprobs_map(
position,
return_tokens_as_token_ids,
))
}))
.collect());
}
if let [_token_id] = prompt_token_ids {
return Ok(vec![None]);
}
Err(server_error!(
"prompt_logprobs were requested but generation returned none"
))
}
/// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
@@ -275,7 +286,13 @@ pub fn clamp_logprob(logprob: f32) -> f32 {
mod tests {
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
use super::decoded_logprobs_to_openai_chat;
use super::{decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps};
#[test]
fn prompt_logprobs_maps_reject_missing_multi_token_payload() {
prompt_logprobs_to_maps(None, &[9707, 11], false)
.expect_err("multi-token prompt without payload is an engine failure");
}
fn sample_logprobs() -> DecodedLogprobs {
DecodedLogprobs {
+131
View File
@@ -0,0 +1,131 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
from types import ModuleType, SimpleNamespace
from typing import Any
import pytest
from vllm.platforms.interface import DeviceCapability
@pytest.fixture
def cuda_platform_base(monkeypatch: pytest.MonkeyPatch) -> Any:
stable_libtorch_module = ModuleType("vllm._C_stable_libtorch")
monkeypatch.setitem(
sys.modules,
"vllm._C_stable_libtorch",
stable_libtorch_module,
)
from vllm.platforms.cuda import CudaPlatformBase
return CudaPlatformBase
def test_compiled_arch_covers_device(cuda_platform_base: Any) -> None:
assert cuda_platform_base._compiled_arch_covers_device(
"12.1a", DeviceCapability(12, 1)
)
assert not cuda_platform_base._compiled_arch_covers_device(
"12.0a", DeviceCapability(12, 1)
)
assert cuda_platform_base._compiled_arch_covers_device(
"12.0f", DeviceCapability(12, 1)
)
assert not cuda_platform_base._compiled_arch_covers_device(
"10.0f", DeviceCapability(12, 1)
)
def test_warn_if_device_arch_not_compiled(
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
) -> None:
def device_count(cls: type[Any]) -> int:
return 2
def get_device_capability(
cls: type[Any], device_id: int = 0
) -> DeviceCapability | None:
capabilities = {
0: DeviceCapability(12, 1),
1: DeviceCapability(10, 3),
}
return capabilities[device_id]
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
return f"GPU {device_id}"
warnings: list[tuple[str, str, str]] = []
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
warnings.append((message, compiled_archs, devices))
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
monkeypatch.setattr(
cuda_platform_base,
"get_device_capability",
classmethod(get_device_capability),
)
monkeypatch.setattr(
cuda_platform_base, "get_device_name", classmethod(get_device_name)
)
monkeypatch.setattr(
"vllm.platforms.cuda.torch.ops",
SimpleNamespace(
_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "12.0f,10.0a")
),
)
monkeypatch.setattr(
"vllm.platforms.cuda.logger",
SimpleNamespace(warning_once=warning_once),
)
cuda_platform_base._warn_if_device_arch_not_compiled()
assert len(warnings) == 1
assert warnings[0][1] == "12.0f, 10.0a"
assert "1: GPU 1 (compute capability 10.3)" in warnings[0][2]
def test_warn_if_device_arch_not_compiled_no_warning(
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
) -> None:
def device_count(cls: type[Any]) -> int:
return 1
def get_device_capability(
cls: type[Any], device_id: int = 0
) -> DeviceCapability | None:
return DeviceCapability(10, 3)
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
raise AssertionError("get_device_name should not be called for covered devices")
warnings: list[tuple[str, str, str]] = []
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
warnings.append((message, compiled_archs, devices))
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
monkeypatch.setattr(
cuda_platform_base,
"get_device_capability",
classmethod(get_device_capability),
)
monkeypatch.setattr(
cuda_platform_base, "get_device_name", classmethod(get_device_name)
)
monkeypatch.setattr(
"vllm.platforms.cuda.torch.ops",
SimpleNamespace(_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "10.0f")),
)
monkeypatch.setattr(
"vllm.platforms.cuda.logger",
SimpleNamespace(warning_once=warning_once),
)
cuda_platform_base._warn_if_device_arch_not_compiled()
assert warnings == []
+13 -3
View File
@@ -60,6 +60,8 @@ def _make_quick_allreduce(
qar.use_fp16_kernels = use_fp16_kernels
qar.qr_quant_level = QuickReduceRegime[quant_level]
qar.qr_max_size = qr_max_size
qar.qr_min_size = None
qar.qr_quantization_min_size = None
return qar
@@ -511,13 +513,21 @@ def test_quick_reduce_regime_values():
assert QuickReduceRegime.INT8.value == 1
assert QuickReduceRegime.INT6.value == 2
assert QuickReduceRegime.INT4.value == 3
assert QuickReduceRegime.NONE.value == 4
assert QuickReduceRegime.INT3.value == 4
assert QuickReduceRegime.NONE.value == 5
def test_quick_reduce_regime_names():
from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
assert set(QuickReduceRegime.__members__) == {
"FP",
"INT8",
"INT6",
"INT4",
"INT3",
"NONE",
}
@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
@@ -693,7 +703,7 @@ def test_quick_allreduce_min_size_table():
for dtype in [torch.float16, torch.bfloat16]:
for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
assert len(min_sizes) == 4
assert len(min_sizes) == 5
assert all(size > 0 for size in min_sizes)
@@ -515,3 +515,15 @@ def test_structured_outputs_structural_tag_invalid(structural_tag):
messages=[{"role": "user", "content": "hello"}],
structured_outputs={"structural_tag": structural_tag},
)
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"])
def test_non_numeric_logprobs_rejected(field_name):
"""A non-numeric logprobs value must be a clean 400 validation error, not a
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
**{field_name: "2"},
)
@@ -610,3 +610,16 @@ class TestCompletionPromptListLimit:
max_tokens=1,
)
assert len(request.prompt_embeds) == 5
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"])
def test_non_numeric_logprobs_rejected(field_name):
"""A non-numeric logprobs value must be a clean 400 validation error, not a
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
CompletionRequest(
model=MODEL_NAME,
prompt="Test prompt",
max_tokens=10,
**{field_name: "2"},
)
@@ -6,7 +6,6 @@ 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,7 +6,6 @@ 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
@@ -11,9 +11,6 @@ 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,
)
@@ -76,59 +73,6 @@ 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 = (
+1 -1
View File
@@ -425,7 +425,7 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No
match the single-call result.
"""
from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
causal_conv1d_torch,
causal_conv1d_fn_cpu as causal_conv1d_torch,
)
x, weight, bias = _conv_inputs(total_tokens)
+8 -3
View File
@@ -18,8 +18,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
DEVICE = current_platform.device_type
pytestmark = pytest.mark.skipif(
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
not (
current_platform.is_cuda_alike()
or current_platform.is_xpu()
or current_platform.is_cpu()
),
reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU",
)
@@ -284,7 +288,8 @@ def test_causal_conv1d_varlen(
batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
):
device = DEVICE
torch.accelerator.empty_cache()
if not current_platform.is_cpu():
torch.accelerator.empty_cache()
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
if itype == torch.bfloat16:
rtol, atol = 1e-2, 5e-2
+48 -2
View File
@@ -20,8 +20,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
DEVICE = current_platform.device_type
pytestmark = pytest.mark.skipif(
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
reason="mamba_ssm kernels require CUDA-alike or XPU",
not (
current_platform.is_cuda_alike()
or current_platform.is_xpu()
or current_platform.is_cpu()
),
reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU",
)
# selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op,
@@ -342,6 +346,13 @@ def test_selective_scan(
@pytest.mark.parametrize("has_z", [False, True])
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update(dim, dstate, has_z, itype):
device = DEVICE
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -436,6 +447,13 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
@pytest.mark.parametrize("max_seq_len", [1, 2, 4])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
device = DEVICE
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -697,6 +715,13 @@ def test_selective_scan_varlen(
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
# tests correctness in case subset of the sequences are padded
@pytest.mark.parametrize("with_padding", [True, False])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_with_batch_indices(
with_padding, dim, dstate, has_z, itype
):
@@ -789,6 +814,13 @@ def test_selective_state_update_with_batch_indices(
@pytest.mark.parametrize("ngroups", [1, 4])
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_with_heads_with_batch_indices(
dim, dstate, ngroups, has_z, tie_hdim, itype
):
@@ -862,6 +894,13 @@ def test_selective_state_update_with_heads_with_batch_indices(
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
@pytest.mark.parametrize("max_seq_len", [2, 4])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_with_num_accepted_tokens(
dim, dstate, has_z, itype, max_seq_len
):
@@ -988,6 +1027,13 @@ def test_selective_state_update_with_num_accepted_tokens(
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
@pytest.mark.parametrize("max_seq_len", [2, 4])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_varlen_with_num_accepted(
dim, dstate, has_z, itype, max_seq_len
):
@@ -0,0 +1,880 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end tests for routed-expert capture on the monolithic MoE path.
These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the
expert IDs picked by FlashInfer's fused router-and-experts kernels (the
"monolithic" path). When ``set_capture_fn`` is installed on
a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call
should:
* allocate an int16 ``(num_tokens, top_k)`` buffer,
* pass it to FlashInfer as ``routing_replay_out``,
* invoke the callback after the kernel returns.
"""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import patch
import pytest
import torch
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
FusedMoEQuantConfig,
RoutingMethodType,
fp8_w8a8_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
TrtLlmBf16ExpertsMonolithic,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
TrtLlmFp8ExpertsMonolithic,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
TrtLlmNvFp4ExpertsMonolithic,
)
from vllm.platforms import current_platform
try:
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
except ImportError:
pytest.skip("flashinfer not available", allow_module_level=True)
if not has_flashinfer_trtllm_fused_moe() or not current_platform.is_cuda():
pytest.skip(
"Requires FlashInfer TRT-LLM fused MoE on CUDA",
allow_module_level=True,
)
if not current_platform.has_device_capability(100):
pytest.skip(
"TRT-LLM fused MoE kernels require SM100+",
allow_module_level=True,
)
def _shuffle_bf16_weights_block_major_k(
w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128
) -> torch.Tensor:
"""Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by
``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert
row shuffle.
"""
from flashinfer import shuffle_matrix_a
from flashinfer.fused_moe import convert_to_block_layout
num_experts = w.shape[0]
shuffled = []
for i in range(num_experts):
t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m)
shuffled.append(convert_to_block_layout(t, block_k))
return torch.stack(shuffled).view(torch.bfloat16)
def _make_bf16_monolithic_experts(
num_experts: int,
top_k: int,
hidden_size: int,
intermediate_size: int,
routing_method: RoutingMethodType,
device: torch.device,
) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]:
"""Construct the monolithic BF16 experts plus the BlockMajorK weights
expected by ``trtllm_bf16_moe``.
"""
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
moe_config = FusedMoEConfig(
num_experts=num_experts,
experts_per_token=top_k,
hidden_dim=hidden_size,
intermediate_size=intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
moe_parallel_config=parallel_cfg,
in_dtype=torch.bfloat16,
activation=MoEActivation.SILU,
device=device,
routing_method=routing_method,
max_num_tokens=max(8, 1),
)
quant_config = FusedMoEQuantConfig.make(
quant_dtype=None,
per_act_token_quant=False,
per_out_ch_quant=False,
block_shape=None,
)
experts = TrtLlmBf16ExpertsMonolithic(
moe_config=moe_config, quant_config=quant_config
)
gemm1 = (
torch.randn(
num_experts,
2 * intermediate_size,
hidden_size,
device=device,
dtype=torch.bfloat16,
)
* 0.1
)
gemm2 = (
torch.randn(
num_experts,
hidden_size,
intermediate_size,
device=device,
dtype=torch.bfloat16,
)
* 0.1
)
w13 = _shuffle_bf16_weights_block_major_k(gemm1)
w2 = _shuffle_bf16_weights_block_major_k(gemm2)
return experts, w13, w2
def _run_bf16_monolithic(
experts: TrtLlmBf16ExpertsMonolithic,
hidden_states: torch.Tensor,
w13: torch.Tensor,
w2: torch.Tensor,
router_logits: torch.Tensor,
num_experts: int,
*,
n_group: int | None = None,
topk_group: int | None = None,
routed_scaling_factor: float | None = None,
e_score_correction_bias: torch.Tensor | None = None,
) -> torch.Tensor:
return experts.apply(
hidden_states=hidden_states,
w1=w13,
w2=w2,
router_logits=router_logits,
activation=MoEActivation.SILU,
global_num_experts=num_experts,
expert_map=None,
a1q_scale=None,
apply_router_weight_on_input=False,
num_expert_group=n_group,
topk_group=topk_group,
e_score_correction_bias=e_score_correction_bias,
routed_scaling_factor=routed_scaling_factor,
)
_DSV3_NUM_EXPERTS = 32
_DSV3_N_GROUP = 4
_DSV3_TOPK_GROUP = 2
def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor:
return torch.randn(num_experts, device=device, dtype=torch.bfloat16)
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
@pytest.mark.parametrize("top_k", [2, 4])
def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts(
num_tokens: int,
top_k: int,
) -> None:
"""The capture callback should receive the int16 routed-expert IDs the
kernel actually used, the values should be valid expert indices, and
each token should pick ``top_k`` distinct experts."""
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
pytest.skip(
f"DSV3 requires top_k <= n_group * topk_group "
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
)
torch.manual_seed(0)
device = torch.device("cuda:0")
num_experts = _DSV3_NUM_EXPERTS
hidden_size = 1024
intermediate_size = 1024
experts, w13, w2 = _make_bf16_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
routing_method=RoutingMethodType.DeepSeekV3,
device=device,
)
captured: list[torch.Tensor] = []
def capture_fn(replay_out: torch.Tensor) -> None:
captured.append(replay_out.clone())
assert experts.supports_routing_replay_capture()
experts.set_capture_fn(capture_fn)
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
)
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
routing_bias = _make_dsv3_routing_bias(num_experts, device)
_ = _run_bf16_monolithic(
experts,
hidden_states=hidden_states,
w13=w13,
w2=w2,
router_logits=router_logits,
num_experts=num_experts,
n_group=_DSV3_N_GROUP,
topk_group=_DSV3_TOPK_GROUP,
routed_scaling_factor=1.0,
e_score_correction_bias=routing_bias,
)
assert len(captured) == 1
replay = captured[0]
assert replay.dtype == torch.int16
assert replay.shape == (num_tokens, top_k)
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
for t in range(num_tokens):
unique = replay[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts, "
f"got {unique.numel()} ({replay[t].tolist()})"
)
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
@pytest.mark.parametrize(
"routing_method",
[
RoutingMethodType.Renormalize,
RoutingMethodType.RenormalizeNaive,
],
)
def test_trtllm_bf16_monolithic_routing_replay_non_dsv3(
num_tokens: int,
routing_method: RoutingMethodType,
) -> None:
"""Routing replay works for non-DeepSeekV3 routing methods too.
FlashInfer's ``routing_replay_out`` is routing-method-agnostic."""
torch.manual_seed(0)
device = torch.device("cuda:0")
num_experts = 8
top_k = 2
hidden_size = 1024
intermediate_size = 1024
experts, w13, w2 = _make_bf16_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
routing_method=routing_method,
device=device,
)
captured: list[torch.Tensor] = []
experts.set_capture_fn(lambda r: captured.append(r.clone()))
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
)
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
_ = _run_bf16_monolithic(
experts,
hidden_states=hidden_states,
w13=w13,
w2=w2,
router_logits=router_logits,
num_experts=num_experts,
)
assert len(captured) == 1
replay = captured[0]
assert replay.dtype == torch.int16
assert replay.shape == (num_tokens, top_k)
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
for t in range(num_tokens):
unique = replay[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts, "
f"got {unique.numel()} ({replay[t].tolist()})"
)
def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None:
"""With no callback installed the kernel should not see a
``routing_replay_out`` tensor verify the helper short-circuits."""
torch.manual_seed(0)
device = torch.device("cuda:0")
experts, _, _ = _make_bf16_monolithic_experts(
num_experts=_DSV3_NUM_EXPERTS,
top_k=2,
hidden_size=1024,
intermediate_size=1024,
routing_method=RoutingMethodType.DeepSeekV3,
device=device,
)
# No callback installed.
buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device)
assert buf is None
# Dispatch is also a no-op.
experts._maybe_dispatch_routing_replay(buf, num_tokens=4)
def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None:
"""FlashInfer's ``routing_replay_out`` is supported by all routing
methods, so ``supports_routing_replay_capture`` should be True
regardless of routing method."""
device = torch.device("cuda:0")
for routing_method in (
RoutingMethodType.DeepSeekV3,
RoutingMethodType.Renormalize,
RoutingMethodType.RenormalizeNaive,
):
experts, _, _ = _make_bf16_monolithic_experts(
num_experts=_DSV3_NUM_EXPERTS,
top_k=2,
hidden_size=1024,
intermediate_size=1024,
routing_method=routing_method,
device=device,
)
assert experts.supports_routing_replay_capture() is True, (
f"{routing_method!r} should support routing replay capture"
)
def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None:
"""When capture is installed, the allocated buffer is int16 and shaped
``(num_tokens, experts_per_token)``."""
device = torch.device("cuda:0")
experts, _, _ = _make_bf16_monolithic_experts(
num_experts=_DSV3_NUM_EXPERTS,
top_k=4,
hidden_size=1024,
intermediate_size=1024,
routing_method=RoutingMethodType.DeepSeekV3,
device=device,
)
experts.set_capture_fn(lambda r: None)
buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device)
assert buf is not None
assert buf.dtype == torch.int16
assert buf.shape[0] >= 11
assert buf.shape[1] == 4
assert buf.device.type == "cuda"
def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
"""End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback
on the monolithic experts and verify the captured rows land in the
capturer's device buffer at the correct layer slot.
Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
for the monolithic path: a single closure is installed on the monolithic
``fused_experts`` (in addition to ``router.set_capture_fn`` on the
non-monolithic path) and the capturer routes per-layer based on the
closed-over ``layer_id``.
"""
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
RoutedExpertsCapturer,
)
torch.manual_seed(7)
device = torch.device("cuda:0")
num_tokens = 4
top_k = 2
num_experts = _DSV3_NUM_EXPERTS
hidden_size = 1024
intermediate_size = 1024
experts, w13, w2 = _make_bf16_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
routing_method=RoutingMethodType.DeepSeekV3,
device=device,
)
num_layers = 3
layer_id = 1
capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
capturer.dp_rank = 0
capturer.tp_size = 1
capturer.device_buffer = torch.full(
(num_tokens + 4, num_layers, top_k),
-1,
dtype=torch.int32,
device=device,
)
def capture_fn(replay_out: torch.Tensor) -> None:
capturer.capture(layer_id, replay_out)
experts.set_capture_fn(capture_fn)
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
)
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
routing_bias = _make_dsv3_routing_bias(num_experts, device)
# Patch get_forward_context to return a dp_metadata=None context so the
# capturer takes the single-DP branch.
import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec
with patch.object(
rec,
"get_forward_context",
return_value=SimpleNamespace(dp_metadata=None),
):
_ = _run_bf16_monolithic(
experts,
hidden_states=hidden_states,
w13=w13,
w2=w2,
router_logits=router_logits,
num_experts=num_experts,
n_group=_DSV3_N_GROUP,
topk_group=_DSV3_TOPK_GROUP,
routed_scaling_factor=1.0,
e_score_correction_bias=routing_bias,
)
captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu()
# Valid expert IDs at this layer.
assert (captured >= 0).all()
assert (captured < num_experts).all()
for t in range(num_tokens):
unique = captured[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts at layer "
f"{layer_id}, got {unique.numel()}"
)
# Other layers / trailing token rows untouched.
for other_layer in range(num_layers):
if other_layer == layer_id:
continue
assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), (
f"layer {other_layer} should be untouched, got writes"
)
assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), (
"tail rows beyond num_tokens should remain sentinel"
)
# ----------------------------------------------------------------------------
# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic``
# ----------------------------------------------------------------------------
def _make_fp8_block_scale_monolithic_experts(
num_experts: int,
top_k: int,
hidden_size: int,
intermediate_size: int,
device: torch.device,
) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]:
"""Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale
code path with DSV3 routing.
Weights are shuffled into the BlockMajorK layout the kernel expects
(same helper the vLLM weight loader uses for DeepSeek-FP8 models).
"""
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
_shuffle_deepseek_fp8_moe_weights,
)
block_k = 128
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
moe_config = FusedMoEConfig(
num_experts=num_experts,
experts_per_token=top_k,
hidden_dim=hidden_size,
intermediate_size=intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
moe_parallel_config=parallel_cfg,
in_dtype=torch.bfloat16,
activation=MoEActivation.SILU,
device=device,
routing_method=RoutingMethodType.DeepSeekV3,
max_num_tokens=max(8, 1),
)
# Random fp8 weights + ones-block scales (the kernel decoder only cares
# that the per-block scales are present and finite for routing/replay).
gemm1 = torch.randn(
num_experts, 2 * intermediate_size, hidden_size, device=device
).to(torch.float8_e4m3fn)
gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
torch.float8_e4m3fn
)
w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2)
w1_scale = torch.ones(
num_experts,
2 * intermediate_size // block_k,
hidden_size // block_k,
device=device,
dtype=torch.float32,
)
w2_scale = torch.ones(
num_experts,
hidden_size // block_k,
intermediate_size // block_k,
device=device,
dtype=torch.float32,
)
quant_config = fp8_w8a8_moe_quant_config(
w1_scale=w1_scale,
w2_scale=w2_scale,
block_shape=[block_k, block_k],
per_act_token_quant=False,
)
experts = TrtLlmFp8ExpertsMonolithic(
moe_config=moe_config, quant_config=quant_config
)
return experts, w13_shuffled, w2_shuffled
def _run_fp8_block_scale_monolithic(
experts: TrtLlmFp8ExpertsMonolithic,
hidden_states_fp8: torch.Tensor,
hidden_states_scale: torch.Tensor,
w13: torch.Tensor,
w2: torch.Tensor,
router_logits: torch.Tensor,
num_experts: int,
routing_bias: torch.Tensor,
) -> torch.Tensor:
return experts.apply(
hidden_states=hidden_states_fp8,
w1=w13,
w2=w2,
router_logits=router_logits,
activation=MoEActivation.SILU,
global_num_experts=num_experts,
expert_map=None,
# The block-scale apply path reads ``a1q_scale`` and transposes it
# to ``(hidden_size/128, num_tokens)`` for the kernel call.
a1q_scale=hidden_states_scale,
apply_router_weight_on_input=False,
num_expert_group=_DSV3_N_GROUP,
topk_group=_DSV3_TOPK_GROUP,
e_score_correction_bias=routing_bias,
routed_scaling_factor=1.0,
)
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
@pytest.mark.parametrize("top_k", [2, 4])
def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts(
num_tokens: int,
top_k: int,
) -> None:
"""End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale
path, DSV3 routing) captures valid expert IDs."""
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
pytest.skip(
f"DSV3 requires top_k <= n_group * topk_group "
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
)
torch.manual_seed(0)
device = torch.device("cuda:0")
num_experts = _DSV3_NUM_EXPERTS
hidden_size = 1024
intermediate_size = 1024
experts, w13, w2 = _make_fp8_block_scale_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
device=device,
)
assert experts.supports_routing_replay_capture()
captured: list[torch.Tensor] = []
experts.set_capture_fn(lambda r: captured.append(r.clone()))
# Per-token / per-block hidden scales (ones is fine for the routing
# path; the GEMM output isn't being asserted on).
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
).to(torch.float8_e4m3fn)
hidden_states_scale = torch.ones(
num_tokens, hidden_size // 128, device=device, dtype=torch.float32
)
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
routing_bias = _make_dsv3_routing_bias(num_experts, device)
_ = _run_fp8_block_scale_monolithic(
experts,
hidden_states_fp8=hidden_states,
hidden_states_scale=hidden_states_scale,
w13=w13,
w2=w2,
router_logits=router_logits,
num_experts=num_experts,
routing_bias=routing_bias,
)
assert len(captured) == 1
replay = captured[0]
assert replay.dtype == torch.int16
assert replay.shape == (num_tokens, top_k)
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
for t in range(num_tokens):
unique = replay[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts, "
f"got {unique.numel()} ({replay[t].tolist()})"
)
# ----------------------------------------------------------------------------
# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic``
# ----------------------------------------------------------------------------
def _make_nvfp4_monolithic_experts(
num_experts: int,
top_k: int,
hidden_size: int,
intermediate_size: int,
device: torch.device,
) -> tuple[
TrtLlmNvFp4ExpertsMonolithic,
torch.Tensor, # w13 (packed nvfp4 uint8)
torch.Tensor, # w13 block-scale (fp8)
torch.Tensor, # w2 (packed nvfp4 uint8)
torch.Tensor, # w2 block-scale (fp8)
torch.Tensor, # input global scale (per-tensor float32)
]:
"""Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights
and DSV3 routing.
NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor
"global" scale. We follow the layout in ``test_ocp_mx_moe.py`` /
``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``:
* weights: uint8 (packed fp4) ``(E, M, K//2)``
* weight scales: fp8 ``(E, M, K//16)``
* hidden states: uint8 (packed fp4) ``(N, K//2)``
* hidden state scales: fp8 ``(N, K//16)``
"""
from flashinfer import fp4_quantize
block_size = 16
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
moe_config = FusedMoEConfig(
num_experts=num_experts,
experts_per_token=top_k,
hidden_dim=hidden_size,
intermediate_size=intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
moe_parallel_config=parallel_cfg,
in_dtype=torch.bfloat16,
activation=MoEActivation.SILU,
device=device,
routing_method=RoutingMethodType.DeepSeekV3,
max_num_tokens=max(8, 1),
)
gemm1 = torch.randn(
num_experts,
2 * intermediate_size,
hidden_size,
device=device,
dtype=torch.bfloat16,
)
gemm2 = torch.randn(
num_experts,
hidden_size,
intermediate_size,
device=device,
dtype=torch.bfloat16,
)
# Per-tensor weight scaling factor (used to build ``g1_alphas`` /
# ``g2_alphas`` below).
w_global_scale = torch.tensor(1.0, device=device)
# Per-tensor input scaling factor.
a_global_scale = torch.tensor(1.0, device=device)
w13_q, w13_scale = fp4_quantize(
gemm1,
w_global_scale,
block_size,
sf_use_ue8m0=False,
is_sf_swizzled_layout=False,
)
w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
num_experts, 2 * intermediate_size, hidden_size // block_size
)
w2_q, w2_scale = fp4_quantize(
gemm2,
w_global_scale,
block_size,
sf_use_ue8m0=False,
is_sf_swizzled_layout=False,
)
w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
num_experts, hidden_size, intermediate_size // block_size
)
# NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2,
# g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these.
g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
quant_config = FusedMoEQuantConfig.make(
quant_dtype="nvfp4",
per_act_token_quant=False,
per_out_ch_quant=False,
block_shape=None,
w1_scale=w13_scale,
w2_scale=w2_scale,
g1_alphas=g_alphas,
g2_alphas=g_alphas,
a1_gscale=a_global_scale,
a2_gscale=a2_gscale,
)
experts = TrtLlmNvFp4ExpertsMonolithic(
moe_config=moe_config, quant_config=quant_config
)
return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale
def _run_nvfp4_monolithic(
experts: TrtLlmNvFp4ExpertsMonolithic,
hidden_states_q: torch.Tensor,
hidden_states_scale: torch.Tensor,
router_logits: torch.Tensor,
num_experts: int,
routing_bias: torch.Tensor,
) -> torch.Tensor:
"""The monolithic NVFP4 apply expects packed fp4 hidden states + the
matching fp8 per-block scale stored in the ``a1q_scale`` slot."""
# Stash the weight tensors on the experts in the locations the apply()
# implementation reads from (it pulls them from quant_config / scales
# already; w1/w2 come in as args).
return experts.apply(
hidden_states=hidden_states_q,
w1=experts._w13_packed,
w2=experts._w2_packed,
router_logits=router_logits,
activation=MoEActivation.SILU,
global_num_experts=num_experts,
expert_map=None,
a1q_scale=hidden_states_scale,
apply_router_weight_on_input=False,
num_expert_group=_DSV3_N_GROUP,
topk_group=_DSV3_TOPK_GROUP,
e_score_correction_bias=routing_bias,
routed_scaling_factor=1.0,
)
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
@pytest.mark.parametrize("top_k", [2, 4])
def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts(
num_tokens: int,
top_k: int,
) -> None:
"""End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs
on the DSV3 routing path."""
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
pytest.skip(
f"DSV3 requires top_k <= n_group * topk_group "
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
)
from flashinfer import fp4_quantize
torch.manual_seed(0)
device = torch.device("cuda:0")
num_experts = _DSV3_NUM_EXPERTS
hidden_size = 1024
intermediate_size = 1024
block_size = 16
experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
device=device,
)
# The apply() reads w1/w2 from its args, but we keep them on the experts
# for convenience of the helper.
experts._w13_packed = w13_q
experts._w2_packed = w2_q
assert experts.supports_routing_replay_capture()
captured: list[torch.Tensor] = []
experts.set_capture_fn(lambda r: captured.append(r.clone()))
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
)
hidden_states_q, hidden_states_scale = fp4_quantize(
hidden_states,
a_gs,
block_size,
sf_use_ue8m0=False,
is_sf_swizzled_layout=False,
)
# The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave
# ``hidden_states_scale`` in its native (uint8 packed) form.
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
routing_bias = _make_dsv3_routing_bias(num_experts, device)
_ = _run_nvfp4_monolithic(
experts,
hidden_states_q=hidden_states_q,
hidden_states_scale=hidden_states_scale,
router_logits=router_logits,
num_experts=num_experts,
routing_bias=routing_bias,
)
assert len(captured) == 1
replay = captured[0]
assert replay.dtype == torch.int16
assert replay.shape == (num_tokens, top_k)
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
for t in range(num_tokens):
unique = replay[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts, "
f"got {unique.numel()} ({replay[t].tolist()})"
)
@@ -0,0 +1,31 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import torch
from vllm import _custom_ops as ops
def test_cutlass_group_gemm_python_guard_allows_thor(monkeypatch):
seen_capabilities: list[int] = []
def fake_cutlass_group_gemm_supported(capability: int) -> bool:
seen_capabilities.append(capability)
return True
monkeypatch.setattr(
torch.ops,
"_C",
SimpleNamespace(cutlass_group_gemm_supported=fake_cutlass_group_gemm_supported),
)
# CUDA 12 reports Thor as SM101 and CUDA 13 reports it as SM110. Both
# should reach the C++ query for the SM10x/SM11x CUTLASS MoE kernel.
assert ops.cutlass_group_gemm_supported(101)
assert ops.cutlass_group_gemm_supported(110)
# SM120 uses separate kernels and must not be advertised by this path.
assert not ops.cutlass_group_gemm_supported(120)
assert seen_capabilities == [101, 110]
+306
View File
@@ -0,0 +1,306 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import ast
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any, cast
import pytest
from vllm.model_executor.warmup.jit_warmup import (
VllmJitKernel,
WarmupIntRange,
get_ast_full_name,
zip_inputs,
)
def _next_power_of_2(value: int) -> int:
return 1 << max(0, value - 1).bit_length()
def _round_up(value: int, *, multiple: int) -> int:
return ((value + multiple - 1) // multiple) * multiple
def _config(
*,
bias: int = 0,
disabled: bool = False,
name: str = "base",
vectorized: bool = False,
) -> SimpleNamespace:
return SimpleNamespace(
bias=bias,
disabled=disabled,
name=name,
vectorized=vectorized,
)
class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
block_size: int
work: int
vector_width: int
descriptor: tuple[object, ...]
enabled: bool
def dispatch( # type: ignore[override]
self,
*,
tokens: int,
cfg: Any,
lanes: int = 1,
mode: str = "default",
debug: int = 0,
) -> CompileKey:
block_size = _next_power_of_2(tokens)
work: int = block_size * lanes + cfg.bias
return self.CompileKey(
block_size=block_size,
work=work,
vector_width=4 if cfg.vectorized and block_size >= 4 else 1,
descriptor=(
cfg.name,
mode,
-block_size,
block_size % 3,
block_size**2,
),
enabled=not cfg.disabled,
)
def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]:
return self._trace_dispatch(self.dispatch)(
tokens=WarmupIntRange(1, max_tokens + 1),
cfg=cfg,
# This argument is intentionally unused by dispatch expressions.
debug=WarmupIntRange(0, 100),
)
def compile(self, compile_key: CompileKey) -> None:
pass
class RecordingToyKernel(ToyKernel):
def __init__(self) -> None:
self.compiled: list[ToyKernel.CompileKey] = []
super().__init__()
def compile(self, compile_key: ToyKernel.CompileKey) -> None:
self.compiled.append(compile_key)
def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None:
cfg = _config()
assert ToyKernel().get_warmup_keys(5, cfg) == [
ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
]
def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
assert ToyKernel().compile_key(
{
"tokens": 4,
"cfg": cfg,
"lanes": 2,
}
) == ToyKernel.CompileKey(
block_size=4,
work=11,
vector_width=4,
descriptor=("cfg", "default", -4, 1, 16),
enabled=False,
)
def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
cfg = _config(vectorized=True)
keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
zip_inputs(
dict(tokens=1, mode="small"),
dict(tokens=4, mode="wide"),
),
cfg=cfg,
lanes=(1, 2),
)
assert keys == [
ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True),
ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True),
ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True),
ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True),
]
def test_zip_inputs_validates_input_rows() -> None:
with pytest.raises(ValueError, match="requires at least one"):
zip_inputs()
with pytest.raises(ValueError, match="rows must be mappings"):
zip_inputs(cast(Any, ("tokens", 1)))
with pytest.raises(ValueError, match="at least one dispatch input name"):
zip_inputs({})
with pytest.raises(ValueError, match="dispatch input names must be strings"):
zip_inputs(cast(Any, {1: 2}))
with pytest.raises(ValueError, match="same dispatch input names"):
zip_inputs({"tokens": 1}, {"mode": "small"})
def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None:
kernel = ToyKernel()
with pytest.raises(TypeError, match="zip_inputs"):
kernel._trace_dispatch(kernel.dispatch)(
cast(Any, {"tokens": 1}),
cfg=_config(),
)
with pytest.raises(ValueError, match="specified more than once"):
kernel._trace_dispatch(kernel.dispatch)(
zip_inputs(dict(tokens=1, mode="small")),
tokens=2,
cfg=_config(),
)
def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def dispatch( # type: ignore[override]
self,
*,
tokens: int,
block_size: int,
) -> CompileKey:
return self.CompileKey(value=_round_up(tokens, multiple=block_size))
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def dispatch( # type: ignore[override]
self,
*,
tokens: int,
block_size: int,
) -> CompileKey:
return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size}))
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
assert HelperKernel().compile_key(
{
"tokens": 5,
"block_size": 4,
}
) == HelperKernel.CompileKey(value=8)
with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"):
StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
if value > 0:
value = 1
return self.CompileKey(value=value)
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
return self.CompileKey(**{"value": value})
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
with pytest.raises(ValueError, match="local assignments"):
BranchKernel()
with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"):
KwargsReturnKernel()
def test_dispatch_reports_unsupported_expression_with_context() -> None:
class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: object
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
return self.CompileKey(value={value})
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
with pytest.raises(ValueError) as exc_info:
UnsupportedKernel().compile_key({"value": 1})
message = str(exc_info.value)
assert "Unsupported dispatch expression" in message
assert "{value}" in message
assert "Supported dispatch expressions" in message
def test_warmup_compiles_all_returned_keys_in_order() -> None:
kernel = RecordingToyKernel()
cfg = _config()
kernel.warmup(3, cfg)
assert kernel.compiled == [
ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
]
def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
dotted_expr = ast.parse("foo.bar.baz").body[0]
call_expr = ast.parse("foo()").body[0]
assert isinstance(dotted_expr, ast.Expr)
assert isinstance(call_expr, ast.Expr)
assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz"
assert get_ast_full_name(call_expr.value) is None
@@ -66,6 +66,12 @@ def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
)
def _make_modular_routed_experts():
return types.SimpleNamespace(
quant_method=types.SimpleNamespace(is_monolithic=False),
)
def test_base_router_capture_pre_eplb_mapping():
router = _make_router()
captured = []
@@ -122,6 +128,8 @@ def test_gpu_model_runner_binds_router_capture(monkeypatch):
def __init__(self):
self.layer_id = 7
self.router = _make_router()
self.routed_experts = _make_modular_routed_experts()
self._quant_method = self.routed_experts.quant_method
class DummyCapturer:
def __init__(self):
@@ -160,6 +168,8 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
def __init__(self):
self.layer_id = 11
self.router = _make_router()
self.routed_experts = _make_modular_routed_experts()
self._quant_method = self.routed_experts.quant_method
class DummyCapturer:
def __init__(self):
@@ -197,6 +207,8 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
def __init__(self, layer_id):
self.layer_id = layer_id
self.router = _make_router()
self.routed_experts = _make_modular_routed_experts()
self._quant_method = self.routed_experts.quant_method
target_module = DummyFusedMoE(layer_id=7)
draft_module = DummyFusedMoE(layer_id=0)
@@ -222,6 +234,49 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
assert draft_module.router.capture_fn is None
def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
from vllm.v1.worker import gpu_model_runner as gmr
class DummyFusedMoE:
def __init__(self):
self.layer_id = 3
self.router = _make_router()
# Use a concrete monolithic expert and override its capability
# instead of instantiating the abstract base class directly.
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
CPUExpertsFp8,
)
fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8)
self.routed_experts = types.SimpleNamespace(
quant_method=types.SimpleNamespace(
is_monolithic=True,
moe_kernel=types.SimpleNamespace(
impl=types.SimpleNamespace(fused_experts=fused_experts)
),
)
)
self._quant_method = self.routed_experts.quant_method
self._quant_method.moe_kernel.impl.fused_experts = fused_experts
fused_experts.supports_routing_replay_capture = lambda: False
class DummyCapturer:
def capture(self, layer_id, topk_ids):
pass
dummy_module = DummyFusedMoE()
import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
dummy_self = types.SimpleNamespace(
model=types.SimpleNamespace(modules=lambda: [dummy_module])
)
with pytest.raises(ValueError, match="monolithic MoE kernel"):
gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
def test_routed_experts_capturer_single_dp_no_metadata():
"""dp_metadata is None: capture writes the full topk_ids rows."""
capturer = _capturer_with_buffer(dp_rank=0)
+13 -2
View File
@@ -9,6 +9,7 @@ import torch.nn.functional as F
from transformers import AutoModel
from vllm.platforms import current_platform
from vllm.utils.mem_constants import MiB_bytes
from ....conftest import HfRunner
from ....utils import VLLM_PATH
@@ -106,7 +107,12 @@ def test_prm_models(
if current_platform.is_cpu():
pytest.skip("CPU only supports V1")
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
with vllm_runner(
model,
max_model_len=1024,
dtype=dtype,
kv_cache_memory_bytes=64 * MiB_bytes,
) as vllm_model:
vllm_outputs = vllm_model.token_classify(math_step_prompts)
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
@@ -145,7 +151,12 @@ def test_prm_models_with_golden_outputs(
if not FIXTURE_REWARD_RESULT.get(model):
pytest.skip(f"No available golden outputs for {model}.")
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
with vllm_runner(
model,
max_model_len=1024,
dtype=dtype,
kv_cache_memory_bytes=64 * MiB_bytes,
) as vllm_model:
vllm_outputs = vllm_model.token_classify(math_step_prompts)
golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])
+38
View File
@@ -74,6 +74,44 @@ def test_cosmos3_new_checkpoint_weights_mapper():
)
def test_cosmos3_modelopt_quantizer_weights_mapper():
"""ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers
(``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable
``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former
(it has no parameter for them) while keeping the latter."""
from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration
mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper
# Native ModelOpt quantizer buffers are dropped.
assert (
mapper.apply_list(
[
"layers.0.self_attn.to_q.input_quantizer._amax",
"layers.0.self_attn.to_q.weight_quantizer._amax",
"layers.0.self_attn.to_q.weight_quantizer._scale",
"layers.0.mlp.down_proj.output_quantizer._amax",
]
)
== []
)
# The FP8 scale sidecars vLLM actually consumes are kept and remapped.
assert mapper.apply_list(
[
"layers.0.self_attn.to_q.weight",
"layers.0.self_attn.to_q.weight_scale",
"layers.0.self_attn.to_q.input_scale",
"layers.0.mlp.down_proj.input_scale",
]
) == [
"language_model.model.layers.0.self_attn.q_proj.weight",
"language_model.model.layers.0.self_attn.q_proj.weight_scale",
"language_model.model.layers.0.self_attn.q_proj.input_scale",
"language_model.model.layers.0.mlp.down_proj.input_scale",
]
def test_cosmos3_edge_checkpoint_weights_mapper():
from vllm.model_executor.models.cosmos3_edge import (
Cosmos3EdgeForConditionalGeneration,
+2 -8
View File
@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.
import pytest
from tests.quantization.utils import is_quant_method_supported
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype
from vllm.platforms import current_platform
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
from ..utils import check_logprobs_close
@@ -70,13 +70,7 @@ def test_models(
if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():
pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")
if not (
current_platform.is_xpu()
or (
get_flash_attn_version() == 3
and current_platform.is_device_capability_family(90)
)
):
if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype):
pytest.skip(
f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
)
+29
View File
@@ -6,6 +6,7 @@ import mimetypes
import os
import shutil
import time
from io import BytesIO
from tempfile import NamedTemporaryFile, TemporaryDirectory
import aiohttp
@@ -111,6 +112,34 @@ async def test_fetch_image_base64(
assert _image_equals(data_image_sync, data_image_async)
@pytest.mark.asyncio
async def test_fetch_image_keep_original_mode():
"""media_io_kwargs can disable the default RGB conversion."""
# RGBA image: opaque black pixel on a fully transparent background
rgba_image = Image.new("RGBA", (4, 4), (0, 0, 0, 0))
rgba_image.putpixel((2, 2), (0, 0, 0, 255))
buffer = BytesIO()
rgba_image.save(buffer, "PNG")
data_url = (
f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
)
# Default behavior: RGBA is composited onto a white background
default_image = MediaConnector().fetch_image(data_url)
assert default_image.mode == "RGB"
assert default_image.getpixel((0, 0)) == (255, 255, 255)
assert default_image.getpixel((2, 2)) == (0, 0, 0)
# image_mode=None via media_io_kwargs: original mode is preserved
connector = MediaConnector(media_io_kwargs={"image": {"image_mode": None}})
image_sync = connector.fetch_image(data_url)
image_async = await connector.fetch_image_async(data_url)
for image in (image_sync, image_async):
assert image.mode == "RGBA"
assert image.getpixel((0, 0)) == (0, 0, 0, 0)
assert image.getpixel((2, 2)) == (0, 0, 0, 255)
@pytest.mark.asyncio
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
async def test_fetch_image_local_files(image_url: str):
+23
View File
@@ -80,6 +80,29 @@ def test_image_media_io_rgba_custom_background(tmp_path):
assert green_numpy[0][0][2] == 0 # B
def test_image_media_io_no_mode_conversion(tmp_path):
"""image_mode=None skips conversion and preserves the original mode."""
# RGBA image: opaque black pixel on a fully transparent background
rgba_image = Image.new("RGBA", (10, 10), (0, 0, 0, 0))
rgba_image.putpixel((5, 5), (0, 0, 0, 255))
test_image_path = tmp_path / "test_rgba.png"
rgba_image.save(test_image_path)
# Default behavior: RGBA is composited onto a white background
image_io_default = ImageMediaIO()
converted_default = image_io_default.load_file(test_image_path)
assert converted_default.media.mode == "RGB"
assert converted_default.media.getpixel((0, 0)) == (255, 255, 255)
assert converted_default.media.getpixel((5, 5)) == (0, 0, 0)
# image_mode=None: original mode and alpha channel are preserved
image_io_keep = ImageMediaIO(image_mode=None)
converted_keep = image_io_keep.load_file(test_image_path)
assert converted_keep.media.mode == "RGBA"
assert converted_keep.media.getpixel((0, 0)) == (0, 0, 0, 0)
assert converted_keep.media.getpixel((5, 5)) == (0, 0, 0, 255)
def test_image_media_io_rgba_background_color_validation():
"""Test that invalid rgba_background_color values are properly rejected."""
@@ -48,12 +48,12 @@ def _check_dense_embedding(data, index=0):
def _check_sparse_embedding(data, check_tokens=False):
expected_weights = [
{"token_id": 32, "weight": 0.0552978515625, "token": "?"},
{"token_id": 70, "weight": 0.09808349609375, "token": " the"},
{"token_id": 83, "weight": 0.08154296875, "token": " is"},
{"token_id": 111, "weight": 0.11810302734375, "token": " of"},
{"token_id": 4865, "weight": 0.1171875, "token": " What"},
{"token_id": 9942, "weight": 0.292236328125, "token": " France"},
{"token_id": 10323, "weight": 0.2802734375, "token": " capital"},
{"token_id": 70, "weight": 0.09808349609375, "token": "the"},
{"token_id": 83, "weight": 0.08154296875, "token": "is"},
{"token_id": 111, "weight": 0.11810302734375, "token": "of"},
{"token_id": 4865, "weight": 0.1171875, "token": "What"},
{"token_id": 9942, "weight": 0.292236328125, "token": "France"},
{"token_id": 10323, "weight": 0.2802734375, "token": "capital"},
]
expected_embed = {x["token_id"]: x for x in expected_weights}
+32
View File
@@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
def test_modelopt_mixed_precision_composes_gemma4_mappers():
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
from vllm.model_executor.models.gemma4_mm import (
Gemma4ForConditionalGeneration,
)
config = _mixed_precision_config(
{
"model.language_model.layers.0.experts": {
"quant_algo": "NVFP4",
"group_size": 16,
},
"model.language_model.layers.1.moe.experts.gate_up_proj": {
"quant_algo": "NVFP4",
"group_size": 16,
},
}
)
config.apply_vllm_mapper(
Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
)
config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
expected_prefix = "language_model.model.layers.0.moe.experts"
assert set(config.quantized_layers) == {
expected_prefix,
"language_model.model.layers.1.moe.gate_up_proj",
}
assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
from vllm.model_executor.layers.linear import LinearBase
+17
View File
@@ -67,6 +67,23 @@ def test_memory_profiling():
non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa
assert abs(non_torch_ratio - 1) <= 0.05
assert result.torch_peak_increase == 1024 * 1024 * 1024
expected_total_consumed = (256 + 512) * 1024 * 1024
total_consumed_ratio = result.total_consumed / expected_total_consumed
assert abs(total_consumed_ratio - 1) <= 0.05, (
f"total_consumed={result.total_consumed}, "
f"expected={expected_total_consumed}, "
f"ratio={total_consumed_ratio}"
)
expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
assert abs(non_kv_ratio - 1) <= 0.05, (
f"non_kv_cache_memory={result.non_kv_cache_memory}, "
f"expected={expected_non_kv}, "
f"ratio={non_kv_ratio}"
)
del weights
lib.cudaFree(handle1)
lib.cudaFree(handle2)
+64 -7
View File
@@ -21,6 +21,7 @@ from vllm.platforms import current_platform
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import (
STR_DTYPE_TO_TORCH_DTYPE,
is_quantized_kv_cache,
is_torch_equal_or_newer,
set_random_seed,
)
@@ -45,6 +46,11 @@ BACKENDS_TO_TEST = [
DEVICE_TYPE = current_platform.device_type
FP8_KV_CACHE_DTYPES = {
"fp8": torch.float8_e4m3fn,
"fp8_e4m3": torch.float8_e4m3fn,
}
# Remove flashinfer from the list if it's not available
try:
import flashinfer # noqa: F401
@@ -110,6 +116,7 @@ def create_and_prepopulate_kv_cache(
num_blocks: int,
common_attn_metadata: CommonAttentionMetadata,
randomize_blocks: bool = True,
kv_cache_dtype: str = "auto",
) -> torch.Tensor:
"""Create and prepopulate a KV cache with context data.
@@ -140,8 +147,18 @@ def create_and_prepopulate_kv_cache(
block_table = common_attn_metadata.block_table_tensor
slot_mapping = common_attn_metadata.slot_mapping
# For an fp8 kv cache, store the cache in the fp8 dtype so that assigning
# the higher-precision context tensors quantizes them, mirroring runtime.
fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype
kv_cache = torch.zeros(
num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device
num_blocks,
block_size,
num_kv_heads,
2 * head_size,
dtype=storage_dtype,
device=device,
)
kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
@@ -195,7 +212,12 @@ def create_and_prepopulate_kv_cache(
] * block_size + token_inter_block_offsets.to(device)
# Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
return kv_cache.transpose(1, 2).contiguous()
kv_cache = kv_cache.transpose(1, 2).contiguous()
if fp8_kv_cache:
kv_cache = kv_cache.view(torch.uint8)
return kv_cache
class MockAttentionLayer:
@@ -224,6 +246,7 @@ def run_attention_backend(
kv_cache: torch.Tensor,
attn_type: AttentionType = AttentionType.DECODER,
sliding_window: int | None = None,
kv_cache_dtype: str = "auto",
) -> torch.Tensor:
"""Run attention computation using the specified backend's AttentionImpl."""
@@ -291,13 +314,16 @@ def run_attention_backend(
alibi_slopes=None,
sliding_window=sliding_window,
attn_type=attn_type,
kv_cache_dtype="auto",
kv_cache_dtype=kv_cache_dtype,
)
# Create mock layer and output buffer
mock_layer = MockAttentionLayer(device)
output = torch.empty_like(query)
if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:
query = query.to(current_platform.fp8_dtype())
# Run forward pass
# NOTE: The query, key, and value are already shaped correctly
# in the calling test function.
@@ -324,6 +350,7 @@ def _test_backend_correctness(
atol: float = 1e-2,
rtol: float = 1e-2,
tensor_parallel_size: int = 1,
kv_cache_dtype: str = "auto",
):
"""
Test that all backends produce similar outputs to a reference implementation
@@ -372,6 +399,7 @@ def _test_backend_correctness(
num_gpu_blocks=8192,
hf_config_override=hf_config_override,
)
vllm_config.cache_config.cache_dtype = kv_cache_dtype
device = torch.device(f"{DEVICE_TYPE}:0")
kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type)
@@ -392,6 +420,13 @@ def _test_backend_correctness(
block_size = vllm_config.cache_config.block_size
scale = 1.0 / (head_size**0.5)
fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
if fp8_kv_cache:
query_fp8_dtype = current_platform.fp8_dtype()
kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype]
atol = max(atol, 6e-2)
rtol = max(rtol, 1e-1)
# 2. Generate data and compute SDPA reference output
all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
all_sdpa_outputs = []
@@ -407,10 +442,17 @@ def _test_backend_correctness(
k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
if fp8_kv_cache:
q_ref = q.to(query_fp8_dtype).to(dtype)
k_ref = k_full.to(kv_fp8_dtype).to(dtype)
v_ref = v_full.to(kv_fp8_dtype).to(dtype)
else:
q_ref, k_ref, v_ref = q, k_full, v_full
# SDPA expects (N, H, L, D), so unsqueeze batch and permute
q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2)
k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2)
v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2)
if num_q_heads != num_kv_heads:
assert num_q_heads % num_kv_heads == 0, (
@@ -471,6 +513,7 @@ def _test_backend_correctness(
num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
common_attn_metadata=common_attn_metadata,
randomize_blocks=True,
kv_cache_dtype=kv_cache_dtype,
)
# 4. Run vLLM backends and compare
@@ -488,6 +531,12 @@ def _test_backend_correctness(
else:
backend_cls = None
if is_quantized_kv_cache(kv_cache_dtype) and (
backend_cls is None
or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
):
continue
if backend_name == AttentionBackendEnum.FLASHINFER:
set_kv_cache_layout("HND")
reset_kv_cache_layout = True
@@ -521,6 +570,7 @@ def _test_backend_correctness(
kv_cache_for_backend,
sliding_window=sliding_window,
attn_type=attn_type,
kv_cache_dtype=kv_cache_dtype,
)
finally:
if reset_kv_cache_layout:
@@ -570,8 +620,13 @@ def _test_backend_correctness(
)
@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
def test_causal_backend_correctness(
default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
default_vllm_config,
batch_spec_name: str,
model: str,
tensor_parallel_size: int,
kv_cache_dtype: str,
):
"""Test backend's correctness with causal attention."""
@@ -612,6 +667,7 @@ def test_causal_backend_correctness(
SMALL_BLOCK_BACKENDS,
causal_mask_mod,
tensor_parallel_size=tensor_parallel_size,
kv_cache_dtype=kv_cache_dtype,
)
# Fast FlexAttention needs to run with block_size=128
@@ -623,6 +679,7 @@ def test_causal_backend_correctness(
causal_mask_mod,
block_size=128,
tensor_parallel_size=tensor_parallel_size,
kv_cache_dtype=kv_cache_dtype,
)
+1
View File
@@ -119,6 +119,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
layer.kv_b_proj.quant_method = None
layer.is_aiter_triton_fp4_bmm_enabled = False
layer.is_aiter_triton_fp8_bmm_enabled = False
layer.dcp_q_replicate = False
layer.quant_config = None
layer.layer_name = "test"
+2
View File
@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (
EncoderOnlyAttentionSpec,
FullAttentionSpec,
MambaSpec,
get_kv_quant_mode,
)
@@ -178,6 +179,7 @@ def create_standard_kv_cache_spec(
head_size=vllm_config.model_config.get_head_size(),
dtype=vllm_config.model_config.dtype,
sliding_window=vllm_config.model_config.get_sliding_window(),
kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype),
)
+1 -3
View File
@@ -9,7 +9,6 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import RequestStatus
from vllm.v1.structured_output import StructuredOutputGrammar
from vllm.v1.utils import ConstantList
from .utils import create_requests, create_scheduler
@@ -263,7 +262,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler = object.__new__(AsyncScheduler)
request = create_requests(num_requests=1, num_tokens=1)[0]
request.structured_output_request = Mock()
request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
request.structured_output_request.grammar = Mock()
request.structured_output_request.grammar.accept_tokens.return_value = False
request.status = RequestStatus.RUNNING
request.num_computed_tokens = request.num_tokens
@@ -285,7 +284,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler.kv_event_publisher = Mock()
scheduler.finished_req_ids = set()
scheduler.finished_req_ids_dict = None
scheduler.grammar_compile_error_reqs = set()
scheduler.vllm_config = Mock()
scheduler.vllm_config.model_config.enable_return_routed_experts = False
scheduler.enable_return_routed_experts = False
+123
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import copy
import hashlib
import importlib
from collections.abc import Callable
@@ -1477,6 +1478,128 @@ def test_get_max_concurrency_for_kv_cache_config():
assert num_tokens == max_concurrency_hybrid_model * max_model_len
assert max_concurrency == max_concurrency_hybrid_model
# Unequal group sizes in the standard layout: each group's pages cost
# whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the
# same as the equal-hybrid case above, regardless of the second group
# holding only 2 layers.
kv_cache_config_unequal_groups = KVCacheConfig(
num_blocks=1153 * 3,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec),
],
)
assert (
get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_unequal_groups
)
== 3
)
# UniformTypeKVCacheSpecs group (worker config shape): the aggregated
# spec's memory/page ratio equals a single layer's page count, so the
# group needs 1024 blocks and the request 1153 in total. The previous
# formula normalized both groups' memory by the first group's page size,
# reporting 3459/1057 = 3.27 here instead of 3 — and a different value
# again for the scheduler-config shape below.
uniform_full_spec = UniformTypeKVCacheSpecs(
block_size=full_attention_spec.block_size,
kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)},
)
kv_cache_config_uniform_group = KVCacheConfig(
num_blocks=1153 * 3,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec),
KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec),
],
)
assert (
get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_uniform_group
)
== 3
)
# Scheduler-config shape: generate_scheduler_kv_cache_config replaces the
# uniform-type group's spec with a representative per-layer spec.
# Capacity must not change between the two shapes (the engine computes
# on the scheduler config, the worker loop on the worker config).
kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config(
[copy.deepcopy(kv_cache_config_uniform_group)]
)
assert get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_scheduler_shape
) == get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_uniform_group
)
def test_get_max_concurrency_packed_kv_cache_config():
from vllm.v1.core.kv_cache_utils import (
_get_kv_cache_config_packed,
_use_packed_kv_cache_config,
)
model_config = ModelConfig(
"Qwen/Qwen1.5-7B",
runner="generate",
dtype="float16",
max_model_len=16384,
)
scheduler_config = SchedulerConfig(
max_num_batched_tokens=1024,
enable_chunked_prefill=True,
max_model_len=model_config.max_model_len,
is_encoder_decoder=model_config.is_encoder_decoder,
async_scheduling=False,
)
vllm_config = VllmConfig(
model_config=model_config,
scheduler_config=scheduler_config,
)
# All-UniformTypeKVCacheSpecs groups select the packed layout.
mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
swa_specs = {
f"layer_{i}": SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
)
for i in range(4, 6)
}
kv_cache_groups = [
KVCacheGroupSpec(
list(mla_specs),
UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
),
KVCacheGroupSpec(
list(swa_specs),
UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
),
]
assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
vllm_config, kv_cache_groups, 2 * GiB_bytes
)
assert num_blocks > 0
kv_cache_config_packed = KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=kv_cache_tensors,
kv_cache_groups=kv_cache_groups,
)
# Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
# the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
# previous formula normalized by the first group's page size and gave
# 1061 blocks per request instead of 1097.
assert get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_packed
) == num_blocks / (1024 + 73)
def test_allocate_with_lookahead():
"""Verify that lookahead tokens correctly affect block allocation"""
+2 -56
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import dataclasses
from concurrent.futures import Future
from unittest.mock import Mock
import pytest
@@ -37,7 +36,7 @@ from vllm.v1.kv_cache_interface import (
)
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager
from vllm.v1.structured_output import StructuredOutputManager
from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
@@ -3007,58 +3006,6 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
assert len(scheduler.skipped_waiting) == 1
@pytest.mark.parametrize("async_grammar", [True, False])
def test_grammar_compile_error_finishes_only_request(async_grammar: bool):
scheduler = create_scheduler()
manager = scheduler.structured_output_manager
manager.backend = Mock()
manager.backend.compile_grammar.side_effect = RuntimeError(
"forced FSM compilation error"
)
manager._use_async_grammar_compilation = async_grammar
sampling_params = SamplingParams(
max_tokens=16,
structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),
)
sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
request = Request(
request_id="grammar-error",
prompt_token_ids=[0, 1],
sampling_params=sampling_params,
pooling_params=None,
)
manager.grammar_init(request)
assert request.structured_output_request is not None
grammar_future = request.structured_output_request._grammar
assert isinstance(grammar_future, Future)
assert isinstance(grammar_future.exception(timeout=5), RuntimeError)
scheduler.add_request(request)
scheduler_output = scheduler.schedule()
assert not scheduler_output.num_scheduled_tokens
engine_core_outputs = scheduler.update_from_output(
scheduler_output,
ModelRunnerOutput(req_ids=[], req_id_to_index={}),
)
assert request.status == RequestStatus.FINISHED_ERROR
assert request.request_id not in scheduler.requests
output = engine_core_outputs[0].outputs[0]
assert output.request_id == request.request_id
assert output.finish_reason == FinishReason.ERROR
assert output.stop_reason is None
healthy_request = create_requests(num_requests=1, req_ids=["healthy-request"])[0]
scheduler.add_request(healthy_request)
next_output = scheduler.schedule()
assert [req.req_id for req in next_output.scheduled_new_reqs] == [
healthy_request.request_id
]
def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler = object.__new__(Scheduler)
sampling_params = SamplingParams(ignore_eos=True, max_tokens=4)
@@ -3072,7 +3019,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
pooling_params=None,
)
request.structured_output_request = Mock()
request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
request.structured_output_request.grammar = Mock()
request.structured_output_request.grammar.accept_tokens.return_value = False
request.status = RequestStatus.RUNNING
request.num_computed_tokens = request.num_tokens
@@ -3093,7 +3040,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler.kv_event_publisher = Mock()
scheduler.finished_req_ids = set()
scheduler.finished_req_ids_dict = None
scheduler.grammar_compile_error_reqs = set()
scheduler.vllm_config = Mock()
scheduler.vllm_config.model_config.enable_return_routed_experts = False
scheduler.enable_return_routed_experts = False
@@ -7,7 +7,13 @@ import torch
from tests.v1.kv_connector.unit.utils import create_vllm_config
from vllm.config import KVEventsConfig, KVTransferConfig
from vllm.distributed.kv_events import MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored
from vllm.distributed.kv_events import (
MEDIUM_CPU,
MEDIUM_FS,
MEDIUM_OBJ,
BlockRemoved,
BlockStored,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
build_offloading_config,
)
@@ -62,8 +68,9 @@ def _wire_hash(block_hash: BlockHash):
return maybe_convert_block_hash(block_hash)
def _request(*, block_hashes: list[BlockHash], token_count: int):
def _request(*, block_hashes: list[BlockHash], token_count: int, req_id: str = "req"):
req = MagicMock()
req.request_id = req_id
req.block_hashes = block_hashes
req.all_token_ids = list(range(1, token_count + 1))
req.lora_request = None
@@ -104,10 +111,32 @@ def _record_chunks(
return keys
def _record_lookup_chunks(
tracker: OffloadingEventsTracker,
req,
group_config: GroupOffloadConfig,
num_chunks: int,
) -> list[OffloadKey]:
keys: list[OffloadKey] = []
hbf = group_config.hashes_per_chunk
for chunk_idx in range(num_chunks):
tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
assert tail_hash is not None
key = make_offload_key(tail_hash, group_config.group_idx)
tracker.record_lookup(
req,
group_config,
chunk_idx,
key,
)
keys.append(key)
return keys
def _stored_event(
keys: list[OffloadKey],
locality: Locality | None = None,
medium: str = _CPU_MEDIUM,
locality: Locality | None = None,
) -> OffloadingEvent:
return OffloadingEvent(
keys=keys,
@@ -119,8 +148,8 @@ def _stored_event(
def _removed_event(
keys: list[OffloadKey],
locality: Locality | None = None,
medium: str = _CPU_MEDIUM,
locality: Locality | None = None,
) -> OffloadingEvent:
return OffloadingEvent(
keys=keys,
@@ -130,6 +159,21 @@ def _removed_event(
)
def _lookup_chunk() -> tuple[
OffloadingEventsTracker, MagicMock, GroupOffloadConfig, OffloadKey
]:
tracker = _tracker()
req = _request(block_hashes=[_hash(0)], token_count=4)
group_config = _group_config()
key = _record_lookup_chunks(
tracker,
req,
group_config,
num_chunks=1,
)[0]
return tracker, req, group_config, key
def test_take_events_forwards_locality_to_rich_store():
tracker = _tracker()
req = _request(block_hashes=[_hash(0)], token_count=4)
@@ -220,18 +264,37 @@ def test_take_events_publishes_routable_block_stored():
assert len(tracker._pending_event_metadata) == 6
def test_take_events_factor_gt_1_chunk_store_and_remove():
def test_promotion_emits_full_cpu_stored_event():
tracker, _, _, key = _lookup_chunk()
[event] = tracker.take_events([_stored_event([key])])
assert isinstance(event, BlockStored)
assert event.medium == MEDIUM_CPU
assert event.block_hashes == [_wire_hash(_hash(0))]
assert event.parent_block_hash is None
assert event.token_ids == [1, 2, 3, 4]
assert event.block_size == 4
assert event.lora_id is None
assert event.lora_name is None
assert event.extra_keys is None
assert event.group_idx == 0
assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
assert event.kv_cache_spec_sliding_window is None
def test_lookup_promotion_factor_gt_1_store_and_remove():
block_size = 4
blocks_per_chunk = 3
blocks_per_chunk = 2
tracker = _tracker()
group_config = _group_config(
block_size=block_size, blocks_per_chunk=blocks_per_chunk
)
req = _request(
block_hashes=[_hash(i) for i in range(6)],
block_hashes=[_hash(i) for i in range(4)],
token_count=block_size * blocks_per_chunk * 2,
)
keys = _record_chunks(tracker, req, group_config, num_chunks=2)
keys = _record_lookup_chunks(tracker, req, group_config, num_chunks=2)
stored = list(tracker.take_events([_stored_event(keys)]))
assert len(stored) == 2
@@ -293,6 +356,7 @@ def test_take_events_opt_out_keeps_placeholders():
group_config = _group_config()
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
keys = _record_chunks(tracker, req, group_config, num_chunks=3)
_record_lookup_chunks(tracker, req, group_config, num_chunks=3)
assert not tracker.self_describing_enabled
assert not tracker._pending_event_metadata
@@ -315,11 +379,21 @@ def test_take_events_opt_out_keeps_placeholders():
assert len(events[3].block_hashes) == 3
def test_record_store_skips_sliding_window_group():
@pytest.mark.parametrize(
"sliding_window_size_in_chunks",
[1, 2],
ids=["ssm", "sliding-window"],
)
def test_event_metadata_skips_non_full_attention_group(
sliding_window_size_in_chunks: int,
):
tracker = _tracker()
group_config = _group_config(sliding_window_size_in_chunks=2)
group_config = _group_config(
sliding_window_size_in_chunks=sliding_window_size_in_chunks
)
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
keys = _record_chunks(tracker, req, group_config, num_chunks=3)
_record_lookup_chunks(tracker, req, group_config, num_chunks=3)
assert not tracker._pending_event_metadata
@@ -329,6 +403,57 @@ def test_record_store_skips_sliding_window_group():
assert events[0].block_size == 0
def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
tracker = _tracker()
block_hashes = [_hash(0), _hash(1)]
req = _request(block_hashes=block_hashes, token_count=8)
group_config = _group_config(blocks_per_chunk=2)
key = _record_chunks(tracker, req, group_config, num_chunks=1)[0]
confirmed_meta = tracker._pending_event_metadata[key]
lookup_req = _request(
block_hashes=block_hashes,
token_count=8,
req_id="new-request",
)
tracker.record_lookup(
lookup_req,
group_config,
0,
key,
)
assert tracker._pending_event_metadata[key] is confirmed_meta
removed = list(tracker.take_events([_removed_event([key])]))
assert len(removed) == 1
assert removed[0].block_hashes == [
_wire_hash(_hash(0)),
_wire_hash(_hash(1)),
]
stored = list(tracker.take_events([_stored_event([key])]))
assert len(stored) == 1
assert stored[0].block_size == 0
assert stored[0].token_ids == []
tracker.record_lookup(lookup_req, group_config, 0, key)
removed = list(tracker.take_events([_removed_event([key])]))
assert removed[0].block_hashes == [
_wire_hash(_hash(0)),
_wire_hash(_hash(1)),
]
@pytest.mark.parametrize("medium", [MEDIUM_FS, MEDIUM_OBJ])
def test_secondary_stored_event_does_not_mutate_cpu_metadata(medium: str):
tracker, _, _, key = _lookup_chunk()
expected_metadata = dict(tracker._pending_event_metadata)
stored = list(tracker.take_events([_stored_event([key], medium)]))
assert stored[0].token_ids == [1, 2, 3, 4]
assert tracker._pending_event_metadata == expected_metadata
def test_take_events_groups_removed_hashes_by_kv_group():
tracker = _tracker()
group0_config = _group_config(group_idx=0, blocks_per_chunk=2)
@@ -378,7 +503,7 @@ def test_reset_cache_clears_side_table():
tracker = _tracker()
group_config = _group_config()
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
_record_chunks(tracker, req, group_config, num_chunks=3)
_record_lookup_chunks(tracker, req, group_config, num_chunks=3)
assert tracker._pending_event_metadata
@@ -387,7 +512,7 @@ def test_reset_cache_clears_side_table():
assert not tracker._pending_event_metadata
def test_tiering_rejects_self_describing_kv_events():
def test_tiering_accepts_self_describing_kv_events():
vllm_config = create_vllm_config(
block_size=4,
max_num_batched_tokens=16,
@@ -423,5 +548,9 @@ def test_tiering_rejects_self_describing_kv_events():
],
)
with pytest.raises(ValueError, match="TieringOffloadingSpec"):
TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
spec = TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
tracker = OffloadingEventsTracker(spec.kv_events_config)
assert spec.kv_events_config.enable_kv_cache_events
assert spec.kv_events_config.self_describing_kv_events
assert tracker.self_describing_enabled

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