Compare commits

..
Author SHA1 Message Date
Bugen Zhao 7fc60bb26d m3 engine parser (reasoning-only) stage 2026-06-22 20:55:24 +08:00
Palaiologos1453andGitHub 1f2c614c27 Merge branch 'main' into minimax-m3-streaming-reasoning-text-markers-main 2026-06-21 13:07:53 +08:00
Ting SUNandGitHub 183a430c13 [Bugfix][Model Runner V2] Fix min_tokens off-by-one in the V2 GPU sampler (#46243)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
2026-06-21 05:06:49 +00:00
MattandGitHub a346d589f5 [Bugfix] Fix NVFP4/OCP MX MoE emulation (#46254)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
2026-06-20 23:13:10 -05:00
Nick HillandGitHub 7df3d7dada [Core] Ensure memory is pinned prior to async h2d copy (#45424)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-06-20 20:02:24 -07:00
8dd1b702f2 [Misc] Fix stale doc URL and docstring module path (#35530)
Signed-off-by: umut-polat <52835619+umut-polat@users.noreply.github.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>
2026-06-20 23:57:01 +00:00
f57ac274b2 [Render] Add reasoning/tool parsing to /derender + fix byte-fallback FFFD (#45919)
Signed-off-by: aoshen524 <aoshen524@gmail.com>
Co-authored-by: Martin Hickey <martin.hickey@ie.ibm.com>
2026-06-20 19:43:32 -04:00
6e919960af [Perf] Skip/shrink all_token_ids copy in scheduler for non-async and V2 runner (#45840)
Signed-off-by: amanchugh89 <amanchugh.89@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-06-20 22:36:57 +00:00
Jonathan ChenandGitHub c88d3d4775 [SimpleCPUOffloadConnector] PCP + DCP support (#39831)
Signed-off-by: Jonathan Chen <chenleejonathan@gmail.com>
2026-06-20 15:01:06 -07:00
Yifan QiaoandGitHub ab7fcbdd5d [Perf][KVConnector][Mooncake] Compact chunk-hash keys and zero-copy lookup wire format (#45969) 2026-06-20 15:00:11 -07:00
3b4a76b63f [KV-Offloading] : Expose CPU cache usage metric (#45737)
Signed-off-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
Signed-off-by: <>
Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
2026-06-20 21:21:55 +00:00
cc22621b51 [KV Offload] Support packed HMA KV cache layout (#46205)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-06-20 21:19:40 +00:00
77148992cf [Bugfix] Move extract_layer_index back inside is_v32 guard (#46199)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-06-20 21:19:10 +00:00
891cc4b9c5 [Frontend] Report cache usage in Anthropic /v1/messages API (#40912)
Signed-off-by: mistral0105 <zhangshuoming17@mails.ucas.ac.cn>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-06-20 21:12:48 +00:00
TJianandGitHub 1bdf9810aa [ROCm] [Bugfix] Bugfix ROCm Sparse Indexer (#46222)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-06-20 13:38:42 -07:00
Tyler Michael SmithGitHubClaudeCodexmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>kourosh hakhamaneshi
ebfbcfe46a Stop setting CUDA_VISIBLE_DEVICES internally in vLLM, add device_ids arg (#45026)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: kourosh hakhamaneshi <kouroshHakha@users.noreply.github.com>
2026-06-20 13:38:10 -07:00
e9de72fe6c [Bugfix] Guard model_config access in _log_compilation_config (#46198)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-06-20 19:26:38 +00:00
L丶GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
d272418f45 [Perf] Optimize Qwen3-VL multi-video prompt processing (#46026)
Signed-off-by: Sirius29 <422058530@qq.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-20 07:09:18 -07:00
Sumanth R HegdeandGitHub 7ff7f5c8eb Revert "Fix Stale Encoder Cache After Weight Update" (#46125) 2026-06-20 07:09:09 -07:00
Palaiologos1453andGitHub f24d8d5bb4 Merge branch 'main' into minimax-m3-streaming-reasoning-text-markers-main 2026-06-20 20:44:14 +08:00
MattGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
dced290769 [Hardware][AMD][CI] Fix e2e core test group (#46024)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-20 02:04:35 -05:00
JasonLi314andGitHub 93bad11912 [Bugfix] Fix gridDim.y overflow for large row counts (#45255)
Signed-off-by: Jason Li <li.jason.cs@gmail.com>
2026-06-19 23:27:45 -04:00
djramicandGitHub 0fbf42af84 [ROCm] Fix VRAM not freed in test_phi3v (#46046)
Signed-off-by: Djordje Ramic <djoramic@amd.com>
2026-06-19 17:20:59 -05:00
Charlie FuandGitHub e6cd8913dd [ROCm][CI] Skip Qwen3.5-35B-A3B-MXFP4-AITER-TP2 for non gfx950 (#46109)
Signed-off-by: charlifu <charlifu@amd.com>
2026-06-19 17:20:10 -05:00
Ben BrowningandGitHub 859e4d436b [Bugfix][Parser] Fix U+FFFD leak at reasoning-to-content transition in engine parsers (#46159)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-06-19 22:09:28 +00:00
Micah WilliamsonandGitHub 4a083cc858 [ROCm][CI] Pin test_rocm_compressed_tensors_w8a8 to TRITON_ATTN (#46180)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-19 15:20:06 -05:00
Vadim GimpelsonandGitHub ca7e1f2c43 Move CI failure diagnosis docs into ci-fails-buildkite skill (#45975)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
2026-06-19 20:12:40 +00:00
djramicandGitHub dec860fb19 [ROCm] Use vLLM's fp8 quant max in AITER hipBLASLt accuracy test (#46176)
Signed-off-by: Djordje Ramic <djoramic@amd.com>
2026-06-19 13:24:02 -05:00
Harry MellorandGitHub 0a49fb2b13 Fix dead link in docs (#46181)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-19 18:16:09 +00:00
Ben BrowningandGitHub 4a8abf37c7 [Test] Migrate test_openai_schema.py to schemathesis 4.x (#46173)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-06-19 18:05:18 +00:00
01192139bf [DSv4] Pack KV caches into contiguous per-block allocations for DeepSeek V4 (#44577)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-06-19 12:55:42 -04:00
Chris LeonardandGitHub b9a7cd464c [12/n] final _C library kernel migration (#45415) 2026-06-19 06:57:26 -07:00
69bdd34542 [Bugfix] Fall back to Pydantic loc for param in validation errors (#46038)
Signed-off-by: professorsab <135441198+professorsab@users.noreply.github.com>
Co-authored-by: Mahad Durrani <114791389+mahadrehmann@users.noreply.github.com>
2026-06-19 19:11:11 +08:00
Kunshang JiandGitHub ec67d7ae61 [xpu] bump up vllm-xpu-kernels v0.1.10 and upgrade 2618 umd (#40367)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-19 15:37:20 +08:00
ecf9d83520 [AMD][CI] Fix Language Models Test (Extended Generation) failures (#45509)
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-19 12:06:56 +08:00
Samuel ShenandGitHub c9135db27c [Docs] Update stale LMCache examples (#45762)
Signed-off-by: Samuel Shen <slshen@tensormesh.ai>
2026-06-19 03:21:36 +00:00
Jeff (Junze) MaGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2a6c6b9429 [DeepSeek-V4] Support TEP=16 for the block-FP8 shared expert (#46001)
Signed-off-by: Jeff Ma <jeffjma@umich.edu>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-18 20:10:12 -07:00
Jared WenandGitHub ab66606993 [bugfix]Indexer init skip and MTP TopK share for iteration (#45895)
Signed-off-by: JaredforReal <w13431838023@gmail.com>
2026-06-19 09:57:51 +08:00
9ea3a4015b [Bugfix] Fix corrupt outputs in MoE FP8 LoRA responses and MoE base model responses when LoRAs are loaded (#42120)
Signed-off-by: Nicholas Edelman <nedelman@nvidia.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
2026-06-18 18:26:09 -07:00
Flora FengandGitHub 560fb8b867 [Cohere] Remove dead prepare_structured_tag override in Cohere parser (#46099)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-19 01:02:11 +00:00
Wentao YeandGitHub 675cd5d228 [Model Runner V2] Fix MRv2 memory leak test (#46095)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-19 00:36:40 +00:00
7f616c327d [Bugfix] [Parser] Fix empty tool block silently dropping subsequent content (#46091)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>
2026-06-18 23:17:18 +00:00
Ivy XuandGitHub c3c6d723fd [Perf] Remove unused loggers in reasoning/ (#45988)
Signed-off-by: Ivy <fakeshadow1337@gmail.com>
2026-06-18 22:24:29 +00:00
Yifan QiaoGitHubJingyi Yangmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
41dcf49ca5 [Bugfix][KV Connector] Disable Mooncake TP put-striding when DCP > 1 (#45371)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Jingyi Yang <girasoleyang@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-18 15:13:44 -07:00
35e4dd4a69 [KV Connector][Mooncake] Async lookup to reduce scheduler overhead (#45659)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-06-18 21:44:02 +00:00
4ce2d01453 fix(anthropic): auto-detect template support for mid-conversation system messages (#46025)
Signed-off-by: felix0080 <felix0080@users.noreply.github.com>
Signed-off-by: Ben Browning <bbrownin@redhat.com>
Co-authored-by: felix0080 <felix0080@users.noreply.github.com>
Co-authored-by: Ben Browning <bbrownin@redhat.com>
2026-06-18 16:19:11 -04:00
Woosuk KwonandGitHub 16908e132e [MRV2] Make FP32 Gumbel sampling more accurate (#45996)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-06-18 19:42:09 +00:00
Wentao YeandGitHub 225936a1dd [CI Bug] Revert #42379 to fix CI Multi-Modal Models (Extended Generation 1) (#46070)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-18 12:37:39 -07:00
f6ba720963 (security) Upgrade Starlette to >= 1.0.1 to fix CVE-2026-48710 (#45675)
Signed-off-by: jperezde <jperezde@redhat.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-06-18 12:35:13 -07:00
test test 928e13af5f Fix MiniMax M3 streaming reasoning markers
Signed-off-by: test test <2260891073@qq.com>
2026-06-16 01:29:52 +08:00
211 changed files with 5243 additions and 1600 deletions
+1 -1
View File
@@ -67,7 +67,7 @@ steps:
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
pytest -v -s v1/structured_output &&
pytest -v -s v1/test_serial_utils.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
- label: "XPU server test"
depends_on:
+1 -14
View File
@@ -647,7 +647,7 @@ steps:
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- label: e2e Core (1 GPU) # TBD
timeout_in_minutes: 180
timeout_in_minutes: 35
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
optional: true
@@ -2075,19 +2075,6 @@ steps:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
- label: e2e Core (1 GPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
- label: e2e Scheduling (1 GPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+10
View File
@@ -74,6 +74,16 @@ steps:
- tests/v1/e2e/general/
commands:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
mirror:
amd:
device: mi250_1
timeout_in_minutes: 35
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
- vllm/platforms/rocm.py
- label: V1 e2e (2 GPUs)
key: v1-e2e-2-gpus
+1
View File
@@ -104,6 +104,7 @@ steps:
source_file_dependencies:
- csrc/quantization/
- vllm/model_executor/layers/quantization
- vllm/config/
- tests/kernels/quantization
- tests/kernels/quantization/test_rocm_skinny_gemms.py
- vllm/_aiter_ops.py
+10
View File
@@ -101,6 +101,16 @@ steps:
num_devices: 8
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
mirror:
amd:
device: mi300_8
timeout_in_minutes: 180
depends_on:
- image-build-amd
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
key: moe-refactor-integration-test-h100-temporary
@@ -68,7 +68,6 @@ steps:
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
mirror:
amd:
soft_fail: true
device: mi325_1
depends_on:
- image-build-amd
@@ -0,0 +1,35 @@
---
name: ci-fails-buildkite
description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs.
---
# Diagnosing vLLM Buildkite CI Failures
Buildkite logs are public; no login needed.
`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch.
## Fetching logs
```bash
# All failed jobs in a PR's latest build (current branch's PR if omitted):
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
# All failed jobs in a build (--soft also includes soft-failed jobs;
# --all fetches every finished job):
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
# work; pass "-" as a second argument to stream to stdout:
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
```
To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`:
```bash
./ci-clean-log.sh ci.log
```
## Reference
See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage.
+2 -3
View File
@@ -2,15 +2,14 @@
# for more info about CODEOWNERS file
# This lists cover the "core" components of vLLM that require careful review
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
/vllm/lora @jeejeelee
/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
/vllm/model_executor/layers/mamba @tdoublep @tomeras91
/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
/vllm/model_executor/model_loader @22quinn
/vllm/model_executor/layers/batch_invariant.py @yewentao256
/vllm/ir @ProExpertProg
+3 -1
View File
@@ -199,7 +199,9 @@ cython_debug/
.vscode/
# Claude
.claude/
.claude/*
!.claude/skills/
!.claude/skills/**
# Codex
.codex/
-11
View File
@@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository:
- Keep comments and docstrings minimal and concise.
- Assume the reader is familiar with vLLM.
### Diagnosing CI failures
Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
```bash
# All failed-job logs for a PR's latest build (current branch's PR if omitted):
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
# Any Buildkite build or job URL also works:
.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
```
### Commit messages
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
+56 -72
View File
@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
#
# _C extension
# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
#
set(VLLM_EXT_SRC
"csrc/quantization/activation_kernels.cu"
"csrc/torch_bindings.cpp")
if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
set(CUTLASS_REVISION "v4.4.2")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
endif()
if(VLLM_CUTLASS_SRC_DIR)
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
endif()
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
# Please keep this in sync with CUTLASS_REVISION line above.
GIT_TAG ${CUTLASS_REVISION}
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
GIT_SHALLOW TRUE
)
endif()
FetchContent_MakeAvailable(cutlass)
set_gencode_flags_for_srcs(
SRCS "${VLLM_EXT_SRC}"
CUDA_ARCHS "${CUDA_ARCHS}")
# if CUDA endif
endif()
if (VLLM_GPU_LANG STREQUAL "HIP")
# Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
# TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
list(APPEND VLLM_EXT_SRC
if(VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_EXT_SRC
"csrc/torch_bindings.cpp"
"csrc/custom_quickreduce.cu"
"csrc/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
)
# if ROCM endif
endif()
"csrc/libtorch_stable/cuda_utils_kernels.cu")
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
# driver API. This causes problems when linking with earlier versions of CUDA.
# Setting this variable sidesteps the issue by calling the driver directly.
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
# driver API. This causes problems when linking with earlier versions of CUDA.
# Setting this variable sidesteps the issue by calling the driver directly.
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
endif() # _C HIP endif
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
#
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/torch_bindings.cpp"
"csrc/libtorch_stable/activation_kernels.cu"
"csrc/libtorch_stable/quantization/activation_kernels.cu"
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
@@ -429,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
set(CUTLASS_REVISION "v4.4.2")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
endif()
if(VLLM_CUTLASS_SRC_DIR)
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
endif()
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
# Please keep this in sync with CUTLASS_REVISION line above.
GIT_TAG ${CUTLASS_REVISION}
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
GIT_SHALLOW TRUE
)
endif()
FetchContent_MakeAvailable(cutlass)
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${FP4_SM120_SRCS}"
CUDA_ARCHS "${FP4_SM120_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
@@ -962,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${FP4_SM100_SRCS}"
CUDA_ARCHS "${FP4_SM100_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
+29 -5
View File
@@ -60,6 +60,7 @@ endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
set(QUTLASS_SOURCES
csrc/qutlass_registration.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
else()
get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
endif()
endif()
elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
list(APPEND QUTLASS_INCLUDES
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
else()
message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
CUDA_ARCHS "${QUTLASS_ARCHS}"
)
target_sources(_C PRIVATE ${QUTLASS_SOURCES})
target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
target_compile_definitions(_C PRIVATE
# 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
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${QUTLASS_SOURCES}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
USE_SABI 3
WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE
QUTLASS_DISABLE_PYBIND=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
)
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +134,5 @@ else()
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
"CUDA_ARCHS='${CUDA_ARCHS}'.")
endif()
add_custom_target(_qutlass_C)
endif()
@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
using TileShapeD = typename MlaSm100Type::TileShapeD;
arguments.problem_shape =
cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
// Assumes device 0 when getting sm_count.
arguments.hw_info.sm_count =
sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
if (sm_count <= 0) {
int current_device = 0;
cudaGetDevice(&current_device);
arguments.hw_info.sm_count =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
} else {
arguments.hw_info.sm_count = sm_count;
}
arguments.split_kv = static_cast<int>(num_kv_splits);
MlaSm100Type::Fmha::set_split_kv(arguments);
+7 -6
View File
@@ -81,11 +81,11 @@ __global__ void rms_norm_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
scalar_t normalized = static_cast<scalar_t>(x * s_variance);
if constexpr (HasWeight) {
dst.val[j] = normalized * src2.val[j];
float w = static_cast<float>(src2.val[j]);
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
} else {
dst.val[j] = normalized;
dst.val[j] = static_cast<scalar_t>(x * s_variance);
}
}
v_out[i] = dst;
@@ -151,7 +151,8 @@ fused_add_rms_norm_kernel(
#pragma unroll
for (int j = 0; j < width; ++j) {
float x = Converter::convert(res.data[j]);
out.data[j] = Converter::convert(x * s_variance) * w.data[j];
float wf = Converter::convert(w.data[j]);
out.data[j] = Converter::convert(x * s_variance * wf);
}
} else {
#pragma unroll
@@ -198,8 +199,8 @@ fused_add_rms_norm_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
if constexpr (HasWeight) {
input[blockIdx.x * input_stride + idx] =
(scalar_t)(x * s_variance) * weight[idx];
float w = (float)weight[idx];
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
} else {
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance);
}
@@ -66,8 +66,13 @@ __global__ void rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
// Multiply in weight's native dtype to match rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
float w = static_cast<float>(src2.val[j]);
// Round normalized result through scalar_t to match the precision of the
// unfused composite (rms_norm writes scalar_t, then
// static_scaled_fp8_quant re-loads it as float before FP8 conversion).
// Without this round, the fused path is strictly more accurate and
// disagrees with the composite at exact E4M3 quantization tie boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
scale_inv);
@@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int i = 0; i < width; ++i) {
float x = Converter::convert(res.data[i]);
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
float wf = Converter::convert(w.data[i]);
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries. We use the
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
HipT out_norm_h = Converter::convert(x * s_variance * wf);
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
Converter::convert(out_norm_h), scale_inv);
}
@@ -183,8 +192,10 @@ fused_add_rms_norm_static_fp8_quant_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
float w = (float)weight[idx];
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
static_cast<float>(out_norm), scale_inv);
}
@@ -9,7 +9,7 @@
#include <torch/headeronly/core/ScalarType.h>
#include "../../cuda_compat.h"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/dispatch_utils.h"
#include "libtorch_stable/torch_utils.h"
+28
View File
@@ -2,9 +2,25 @@
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/util/Exception.h>
#include <optional>
#include <string>
#include <vector>
#include <torch/csrc/stable/ops.h>
inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
// Ensure tensor is on CUDA
STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
// Get the raw data pointer
void* data_ptr = tensor.mutable_data_ptr();
/// Create a new tensor from the raw data pointer
return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
tensor.device(), tensor.scalar_type());
}
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
torch::stable::Tensor& output_q,
@@ -371,6 +387,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void silu_and_mul_clamp(torch::stable::Tensor& out,
torch::stable::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::stable::Tensor& out,
torch::stable::Tensor& input,
torch::stable::Tensor& scale);
void persistent_masked_m_silu_mul_quant(
const torch::stable::Tensor& input, // (E, T, 2*H)
const torch::stable::Tensor& tokens_per_expert, // (E)
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool use_ue8m0);
void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_tanh_and_mul(torch::stable::Tensor& out,
@@ -1,16 +1,12 @@
#include <ATen/cuda/CUDAContext.h>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include "libtorch_stable/torch_utils.h"
#include <cmath>
#include "core/math.hpp"
#include "../cuda_compat.h"
#include "dispatch_utils.h"
#include "libtorch_stable/core/math.hpp"
#include "cuda_compat.h"
#include "libtorch_stable/dispatch_utils.h"
#include "quantization/w8a8/fp8/common.cuh"
#include <c10/util/Float8_e4m3fn.h>
#ifndef USE_ROCM
#include <cuda_bf16.h>
#include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
#endif
#endif
#include "core/registration.h"
namespace vllm {
template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
} // namespace vllm
// Launch activation, gating, and quantize kernel.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
dim3 block(std::min(d, 512)); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
VLLM_DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "act_and_mul_kernel", [&] { \
VLLM_DISPATCH_FP8_TYPES( \
out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
fp8_t> \
<<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(), \
input.data_ptr<scalar_t>(), \
scale.data_ptr<float>(), d); \
}); \
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
dim3 block(std::min(d, 512)); \
const torch::stable::accelerator::DeviceGuard device_guard( \
input.get_device_index()); \
const cudaStream_t stream = \
get_current_cuda_stream(input.get_device_index()); \
VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "act_and_mul_kernel", [&] { \
VLLM_STABLE_DISPATCH_FP8_TYPES( \
out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
fp8_t> \
<<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<fp8_t>(), \
input.const_data_ptr<scalar_t>(), \
scale.const_data_ptr<float>(), d); \
}); \
});
void silu_and_mul_quant(torch::Tensor& out, // [..., d]
torch::Tensor& input, // [..., 2 * d]
torch::Tensor& scale) {
TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
out.dtype() == torch::kFloat8_e4m3fnuz);
TORCH_CHECK(input.dtype() == torch::kFloat16 ||
input.dtype() == torch::kBFloat16);
TORCH_CHECK(input.size(-1) % 2 == 0);
void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input, // [..., 2 * d]
torch::stable::Tensor& scale) {
STD_TORCH_CHECK(
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
STD_TORCH_CHECK(
input.scalar_type() == torch::headeronly::ScalarType::Half ||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"Input must be FP16 or BF16");
STD_TORCH_CHECK(input.size(-1) % 2 == 0);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
}
void persistent_masked_m_silu_mul_quant(
const at::Tensor& input, // (E, T, 2*H)
const at::Tensor& tokens_per_expert, // (E)
at::Tensor& y_q, // (E, T, H) [OUT]
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
const torch::stable::Tensor& input, // (E, T, 2*H)
const torch::stable::Tensor& tokens_per_expert, // (E)
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool cast_scale_ue8m0) {
#ifndef USE_ROCM
@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
// fixed GROUP_SIZE of 128.
static constexpr int GROUP_SIZE = 128;
TORCH_CHECK(input.dtype() == torch::kBFloat16);
TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
y_q.dtype() == torch::kFloat8_e4m3fnuz);
TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
STD_TORCH_CHECK(input.scalar_type() ==
torch::headeronly::ScalarType::BFloat16);
STD_TORCH_CHECK(
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
bool const is_packed_ue8m0 =
(y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
(y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
cast_scale_ue8m0);
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
is_packed_ue8m0);
using Idx_t = int64_t;
@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(
int const NUM_GROUPS = H / GROUP_SIZE;
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
// TODO: Get this from cuda_arch ?
static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
static constexpr int max_shared_mem_bytes = \
GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \
dim3 grid(sms), block(THREAD_COUNT); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
VLLM_DISPATCH_FP8_TYPES( \
const torch::stable::accelerator::DeviceGuard device_guard( \
input.get_device_index()); \
VLLM_STABLE_DISPATCH_FP8_TYPES( \
y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \
vllm::silu_mul_fp8_quant_deep_gemm_kernel< \
BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \
<<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>( \
reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
(fp8_t*)y_q.data_ptr(), \
reinterpret_cast<scale_t*>(y_s.data_ptr()), \
reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
reinterpret_cast<const __nv_bfloat16*>( \
input.const_data_ptr()), \
y_q.mutable_data_ptr<fp8_t>(), \
reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()), \
reinterpret_cast<const int32_t*>( \
tokens_per_expert.const_data_ptr()), \
E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \
STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \
});
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
Idx_t stride_ys_g = y_s.stride(2);
Idx_t stride_ys_p = 0;
if (!cast_scale_ue8m0) {
TORCH_CHECK(!is_packed_ue8m0);
STD_TORCH_CHECK(!is_packed_ue8m0);
LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
false);
return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
return;
}
TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
TORCH_CHECK(y_s.dtype() == torch::kInt32);
STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);
// Int32 packed ue8m0 scales tensor.
// Let E, T, G be the number to experts, number of tokens and number of groups
@@ -31,7 +31,7 @@
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "core/batch_invariant.hpp"
using namespace cute;
@@ -31,7 +31,7 @@
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "core/batch_invariant.hpp"
using namespace cute;
@@ -19,7 +19,7 @@
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -14,7 +14,7 @@
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -22,7 +22,7 @@
#include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
const int sf_k_local = local_group_id % kGroupsPerBlockX;
const int row_local = local_group_id / kGroupsPerBlockX;
const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
// Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;");
@@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
" is not a multiple of 4.");
const int kx = GetGroupsPerBlockX(padded_groups_per_row);
const int ry = 16 / kx;
const int64_t blocks_x = padded_groups_per_row / kx;
const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
const int64_t sf_k_blocks = padded_groups_per_row / kx;
const int num_threads = (kx * ry) * THREADS_PER_GROUP;
// CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
blocks_y <= static_cast<int64_t>(INT32_MAX),
// CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
constexpr int64_t kMaxGridDimYZ = 65535;
STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
sf_k_blocks <= kMaxGridDimYZ,
"per_token_group_quant_8bit_packed grid too large: (",
blocks_x, ", ", blocks_y, ").");
row_blocks, ", ", sf_k_blocks, ").");
auto dst_type = output_q.scalar_type();
@@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
do { \
cudaLaunchConfig_t config = {}; \
config.gridDim = dim3(static_cast<unsigned int>(blocks_x), \
static_cast<unsigned int>(blocks_y)); \
config.gridDim = dim3(static_cast<unsigned int>(row_blocks), \
static_cast<unsigned int>(sf_k_blocks)); \
config.blockDim = dim3(num_threads); \
config.dynamicSmemBytes = 0; \
config.stream = stream; \
@@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
#else
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
do { \
dim3 grid(static_cast<unsigned int>(blocks_x), \
static_cast<unsigned int>(blocks_y)); \
dim3 grid(static_cast<unsigned int>(row_blocks), \
static_cast<unsigned int>(sf_k_blocks)); \
dim3 block(num_threads); \
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
RY> \
+27
View File
@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
// TODO: Remove this once ROCm upgrade to torch 2.11.
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
// But since they are set back to zero once the kernel is finished we can
// hand wave and say that they have no net effect.
//
// The reason to mark 'workspace' as immutable is so that they don't interfere
// with using ScalarType arguments in the ops. If they are marked as mutable,
// pytorch throws an assert in
// 'torch._higher_order_ops._register_effectful_op' that prevents these
// kernels from being torch.compile'd.
// See the following document for more info on custom types and ops that use
// custom types:
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
ops.def(
"machete_supported_schedules("
@@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor workspace, int k, int max_seq_len) -> ()");
// Activation ops
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s, bool use_ue8m0) -> ()");
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
// Activation function used in SwiGLU.
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
@@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
"float alpha=1.0, float beta=0.0) -> ()");
// SwiGLU activation with FP8 quantization.
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
// Activation function used in GeGLU with `none` approximation.
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
// Activation kernels (shared CUDA/ROCm)
ops.impl("persistent_masked_m_silu_mul_quant",
TORCH_BOX(&persistent_masked_m_silu_mul_quant));
ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
-32
View File
@@ -9,28 +9,6 @@
#include <vector>
torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
// Ensure tensor is on CUDA
if (!tensor.is_cuda()) {
throw std::runtime_error("Tensor must be on CUDA device");
}
// Get the raw data pointer
void* data_ptr = tensor.data_ptr();
// Get tensor sizes and strides
std::vector<int64_t> sizes = tensor.sizes().vec();
std::vector<int64_t> strides = tensor.strides().vec();
// Get tensor options (dtype, device)
auto options = tensor.options();
// Create a new tensor from the raw data pointer
auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
return new_tensor;
}
// rms_norm and fused_add_rms_norm declarations also exist in
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
// because the CPU build still uses these torch::Tensor declarations.
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
torch::Tensor& scale);
void persistent_masked_m_silu_mul_quant(
const at::Tensor& input, // (E, T, 2*H)
const at::Tensor& counts, // (E)
at::Tensor& y_q, // (E, T, H) [OUT]
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool use_ue8m0);
void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
+5
View File
@@ -0,0 +1,5 @@
#include "core/registration.h"
// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
// This stub lets Python import vllm._qutlass_C to trigger op registration.
REGISTER_EXTENSION(_qutlass_C)
-40
View File
@@ -20,17 +20,6 @@
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// vLLM custom ops
//
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s,"
"bool use_ue8m0) -> ()");
ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
&persistent_masked_m_silu_mul_quant);
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
#ifdef USE_ROCM
// TODO: Remove this once we upgrade to torch 2.11.
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
&get_cuda_view_from_cpu_tensor);
#endif
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
// kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
// full-cache bf16/fp8 variants).
// Quantization ops
#ifndef USE_ROCM
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
// But since they are set back to zero once the kernel is finished we can
// hand wave and say that they have no net effect.
//
// The reason to mark 'workspace' as immutable is so that they don't interfere
// with using ScalarType arguments in the ops. If they are marked as mutable,
// pytorch throws an assert in
// 'torch._higher_order_ops._register_effectful_op' that prevents these
// kernels from being torch.compile'd.
// See the following document for more info on custom types and ops that use
// custom types:
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
#endif
}
+7 -7
View File
@@ -75,13 +75,13 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
# Install UMD
RUN mkdir neo && \
cd neo && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-core-2_2.34.4+21428_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-opencl-2_2.34.4+21428_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-ocloc_26.18.38308.1-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-opencl-icd_26.18.38308.1-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libigdgmm12_22.10.0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libze-intel-gpu1_26.18.38308.1-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero_1.28.2+u24.04_amd64.deb && \
dpkg -i *.deb && \
cd .. && \
rm -rf neo
+1 -1
View File
@@ -133,7 +133,7 @@ The model should inherit protocol `IsAttentionFree` and also implement class met
For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).
For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
@@ -4,7 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
* **Upstream vLLM compatibility** It wraps around upstream vLLM without modifying its code.
* **Ease of use** Simplified deployment via Helm charts and observability through Grafana dashboards.
* **High performance** Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache), among others.
* **High performance** Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache) (wired up in vLLM via `--kv-offloading-backend lmcache`; see the [LMCache examples](https://github.com/vllm-project/vllm/tree/main/examples/disaggregated/lmcache) and [docs.lmcache.ai](https://docs.lmcache.ai)), among others.
If you are new to Kubernetes, don't worry: in the vLLM production stack [repo](https://github.com/vllm-project/production-stack), we provide a step-by-step [guide](https://github.com/vllm-project/production-stack/blob/main/tutorials/00-install-kubernetes-env.md) and a [short video](https://www.youtube.com/watch?v=EsTJbQtzj0g) to set up everything and get started in **4 minutes**!
+1 -1
View File
@@ -20,7 +20,7 @@ Two main reasons:
Now supports 9 types of connectors:
- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission. LMCache also offers a multi-process (MP) mode via `LMCacheMPConnector`, where a standalone `lmcache server` holds the KV cache shared by one or more vLLM instances; see the [LMCache examples](../../examples/disaggregated/lmcache/README.md) and the [LMCache docs](https://docs.lmcache.ai) for setup.
- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as:
```bash
@@ -203,6 +203,7 @@ the vLLM JSON config.
### kv_connector_extra_config
- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
- `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`.
- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
- `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
@@ -27,6 +27,7 @@ Currently, there are no pre-built XPU wheels.
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
```bash
git clone https://github.com/vllm-project/vllm.git
+34 -14
View File
@@ -1,10 +1,38 @@
# LMCache Examples
This folder demonstrates how to use LMCache for disaggregated prefilling, CPU offloading and KV cache sharing.
This folder demonstrates how to use LMCache with vLLM v1 for KV cache
offloading, disaggregated prefilling, and KV cache sharing.
## 1. Disaggregated Prefill in vLLM v1
## Integration modes
This example demonstrates how to run LMCache with disaggregated prefill using NIXL on a single node.
LMCache integrates with vLLM v1 in two ways:
- **In-process mode** (`LMCacheConnectorV1`): LMCache runs inside the vLLM
process and is configured through environment variables or a YAML config
file (`LMCACHE_CONFIG_FILE`). This is the simplest way to add single-node
CPU/disk offloading.
- **Multi-process (MP) mode** (`LMCacheMPConnector`): LMCache runs as a
standalone server (`lmcache server`) that owns the KV cache storage; one or
more vLLM instances connect to it. This is the recommended mode for
distributed KV storage and for sharing KV cache across instances. See the
[LMCache docs](https://docs.lmcache.ai) for the full MP setup.
## 1. CPU offload (in-process)
- `python cpu_offload_lmcache.py` - CPU offloading with `LMCacheConnectorV1`
for vLLM v1.
## 2. CPU offload (multi-process)
- `bash cpu_offload_lmcache_mp.sh` - CPU offloading with `LMCacheMPConnector`,
using a standalone `lmcache server`. vLLM provides a built-in shortcut for
this setup via `--kv-offloading-backend lmcache` and
`--kv-offloading-size <GiB>`.
## 3. Disaggregated Prefill in vLLM v1
This example demonstrates how to run LMCache with disaggregated prefill using
NIXL on a single node.
### Prerequisites
@@ -46,15 +74,7 @@ The main script generates several log files:
- `decoder.log` - Logs from the decode server
- `proxy.log` - Logs from the proxy server
## 2. CPU Offload Examples
## 4. KV Cache Sharing
- `python cpu_offload_lmcache.py -v v0` - CPU offloading implementation for vLLM v0
- `python cpu_offload_lmcache.py -v v1` - CPU offloading implementation for vLLM v1
## 3. KV Cache Sharing
The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV caches between vLLM v1 instances.
## 4. Disaggregated Prefill in vLLM v0
The `disaggregated_prefill_lmcache_v0.py` provides an example of how to run disaggregated prefill in vLLM v0.
The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV
caches between vLLM v1 instances through a centralized LMCache server.
@@ -1,20 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of cpu offloading
with LMCache in vLLM v1 or v0.
Usage:
Specify vLLM version
-v v0 : Use LMCacheConnector
model = mistralai/Mistral-7B-Instruct-v0.2
(Includes enable_chunked_prefill = True)
-v v1 : Use LMCacheConnectorV1 (default)
model = meta-llama/Meta-Llama-3.1-8B-Instruct
(Without enable_chunked_prefill)
This file demonstrates the example usage of CPU offloading
with LMCache in vLLM v1.
Note that `lmcache` is needed to run this example.
Requirements:
@@ -23,7 +11,6 @@ Learn more about LMCache environment setup, please refer to:
https://docs.lmcache.ai/getting_started/installation.html
"""
import argparse
import contextlib
import os
import time
@@ -39,8 +26,6 @@ from vllm.engine.arg_utils import EngineArgs
def setup_environment_variables():
# LMCache-related environment variables
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Enable local CPU backend in LMCache
@@ -50,9 +35,9 @@ def setup_environment_variables():
@contextlib.contextmanager
def build_llm_with_lmcache(lmcache_connector: str, model: str):
def build_llm_with_lmcache(model: str):
ktc = KVTransferConfig(
kv_connector=lmcache_connector,
kv_connector="LMCacheConnectorV1",
kv_role="kv_both",
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
@@ -92,23 +77,10 @@ def print_output(
print("-" * 50)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"-v",
"--version",
choices=["v0", "v1"],
default="v1",
help="Specify vLLM version (default: v1)",
)
return parser.parse_args()
def main():
lmcache_connector = "LMCacheConnectorV1"
model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
setup_environment_variables()
with build_llm_with_lmcache(lmcache_connector, model) as llm:
with build_llm_with_lmcache(model) as llm:
# This example script runs two requests with a shared prefix.
# Define the shared prompt and specific prompts
shared_prompt = "Hello, how are you?" * 1000
+43
View File
@@ -0,0 +1,43 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# CPU offloading with LMCache in multi-process (MP) mode.
#
# In MP mode, LMCache runs as a standalone server process (`lmcache server`)
# that owns the KV cache storage. One or more vLLM instances connect to it via
# the `LMCacheMPConnector`. This is the recommended way to run LMCache for
# distributed KV storage and for sharing KV cache across vLLM instances.
#
# vLLM ships a built-in shortcut for this setup: pass `--kv-offloading-backend
# lmcache` together with `--kv-offloading-size <GiB>` and vLLM wires up the
# `LMCacheMPConnector` for you (it defaults to the LMCache server at
# tcp://localhost:5555, matching the `lmcache server` default).
#
# Requires `lmcache` to be installed (`pip install lmcache`).
# Learn more: https://docs.lmcache.ai
set -euo pipefail
MODEL=${MODEL:-meta-llama/Meta-Llama-3.1-8B-Instruct}
# 1. Launch the standalone LMCache server (binds tcp://localhost:5555 by
# default). `--l1-size-gb` sets the CPU memory budget for the L1 cache.
echo "Starting LMCache server..."
lmcache server --host localhost --port 5555 --l1-size-gb 5 &
LMCACHE_SERVER_PID=$!
trap 'kill $LMCACHE_SERVER_PID 2>/dev/null || true' EXIT
# 2. Launch vLLM and offload KV cache to the LMCache server.
# The MP connector currently requires the non-hybrid KV cache manager.
echo "Starting vLLM server with LMCache MP offloading..."
vllm serve "$MODEL" \
--port 8000 \
--kv-offloading-size 5 \
--kv-offloading-backend lmcache \
--disable-hybrid-kv-cache-manager
# Equivalent explicit configuration (instead of the two flags above):
# --kv-transfer-config \
# '{"kv_connector":"LMCacheMPConnector","kv_role":"kv_both",
# "kv_connector_extra_config":{"lmcache.mp.host":"tcp://localhost",
# "lmcache.mp.port":5555}}'
@@ -1,144 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of disaggregated prefilling
with LMCache.
We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
and launch an additional LMCache server.
KV cache is transferred in the following manner:
vLLM prefill node -> LMCache server -> vLLM decode node.
Note that `pip install lmcache` is needed to run this example.
Learn more about LMCache in https://github.com/LMCache/LMCache.
"""
import os
import subprocess
import time
from multiprocessing import Event, Process
from lmcache.experimental.cache_engine import LMCacheEngineBuilder
from lmcache.integration.vllm.utils import ENGINE_NAME
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
# LMCache-related environment variables
# The port to start LMCache server
port = 8100
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Disable local CPU backend in LMCache
os.environ["LMCACHE_LOCAL_CPU"] = "False"
# Set local CPU memory buffer limit to 5.0 GB
os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0"
# Set the remote URL for LMCache server
os.environ["LMCACHE_REMOTE_URL"] = f"lm://localhost:{port}"
# Set the serializer/deserializer between vllm and LMCache server
# `naive` indicates using raw bytes of the tensor without any compression
os.environ["LMCACHE_REMOTE_SERDE"] = "naive"
prompts = [
"Hello, how are you?" * 1000,
]
def run_prefill(prefill_done, prompts):
# We use GPU 0 for prefill node.
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
ktc = KVTransferConfig(
kv_connector="LMCacheConnector",
kv_role="kv_producer",
kv_rank=0,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
# memory. Reduce the value if your GPU has less memory.
llm = LLM(
model="mistralai/Mistral-7B-Instruct-v0.2",
kv_transfer_config=ktc,
max_model_len=8000,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
# llm.generate(prompts, sampling_params)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
generated_text = output.outputs[0].text
print(f"Generated text: {generated_text!r}")
print("Prefill node is finished.")
prefill_done.set()
# Clean up lmcache backend
LMCacheEngineBuilder.destroy(ENGINE_NAME)
def run_decode(prefill_done, prompts, timeout=1):
# We use GPU 1 for decode node.
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
ktc = KVTransferConfig(
kv_connector="LMCacheConnector",
kv_role="kv_consumer",
kv_rank=1,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
# of memory. Reduce the value if your GPU has less memory.
llm = LLM(
model="mistralai/Mistral-7B-Instruct-v0.2",
kv_transfer_config=ktc,
max_model_len=8000,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
print("Waiting for prefill node to finish...")
prefill_done.wait()
time.sleep(timeout)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
generated_text = output.outputs[0].text
print(f"Generated text: {generated_text!r}")
# Clean up lmcache backend
LMCacheEngineBuilder.destroy(ENGINE_NAME)
def run_lmcache_server(port):
server_proc = subprocess.Popen(
["python", "-m", "lmcache.experimental.server", "localhost", str(port)]
)
return server_proc
def main():
prefill_done = Event()
prefill_process = Process(target=run_prefill, args=(prefill_done, prompts))
decode_process = Process(target=run_decode, args=(prefill_done, prompts))
lmcache_server_process = run_lmcache_server(port)
# Start prefill node
prefill_process.start()
# Start decode node
decode_process.start()
# Clean up the processes
decode_process.join()
prefill_process.terminate()
lmcache_server_process.terminate()
lmcache_server_process.wait()
if __name__ == "__main__":
main()
@@ -30,7 +30,6 @@ if [[ $1 == "prefiller" ]]; then
UCX_TLS=cuda_ipc,cuda_copy,tcp \
LMCACHE_CONFIG_FILE=$prefill_config_file \
LMCACHE_USE_EXPERIMENTAL=True \
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES=0 \
@@ -47,7 +46,6 @@ elif [[ $1 == "decoder" ]]; then
UCX_TLS=cuda_ipc,cuda_copy,tcp \
LMCACHE_CONFIG_FILE=$decode_config_file \
LMCACHE_USE_EXPERIMENTAL=True \
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES=1 \
@@ -26,8 +26,6 @@ from vllm.config import KVTransferConfig
# LMCache-related environment variables
# The port to start LMCache server
port = 8100
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Disable local CPU backend in LMCache
+3 -2
View File
@@ -11,13 +11,14 @@ transformers >= 5.5.3
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides.
starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1
aiohttp >= 3.13.3
openai >= 2.0.0 # For Responses API with reasoning content
pydantic >= 2.12.0
prometheus_client >= 0.18.0
pillow # Required for image processing
prometheus-fastapi-instrumentator >= 7.0.0
prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0
tiktoken >= 0.6.0 # Required for DBRX tokenizer
lm-format-enforcer == 0.11.3
llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
+1 -1
View File
@@ -11,7 +11,7 @@ numba == 0.65.0 # Required for N-gram speculative decoding
datasets
peft
pytest-asyncio
tensorizer==2.12.1
tensorizer==2.10.1
packaging>=24.2
setuptools>=77.0.3,<80.0.0
setuptools-scm>=8
+2 -2
View File
@@ -1,6 +1,6 @@
# testing
pytest
tensorizer==2.12.1
tensorizer==2.10.1
pytest-forked
pytest-asyncio
pytest-rerunfailures
@@ -40,7 +40,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test.
schemathesis>=4.0.0 # Required for openai schema test.
# quantization
bitsandbytes==0.49.2
buildkite-test-collector==0.1.9
+24 -49
View File
@@ -35,14 +35,11 @@ arctic-inference==0.1.1
# via -r requirements/test/cuda.in
argcomplete==3.5.1
# via datamodel-code-generator
arrow==1.3.0
# via isoduration
attrs==24.2.0
# via
# aiohttp
# hypothesis
# jsonschema
# pytest-subtests
# referencing
audioread==3.0.1
# via librosa
@@ -57,9 +54,7 @@ azure-identity==1.25.2
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test/cuda.in
# schemathesis
# via -r requirements/test/cuda.in
bitsandbytes==0.49.2
# via -r requirements/test/cuda.in
black==24.10.0
@@ -110,7 +105,6 @@ colorama==0.4.6
# via
# perceptron
# sacrebleu
# schemathesis
colorful==0.5.6
# via ray
colorlog==6.10.1
@@ -183,7 +177,7 @@ et-xmlfile==2.0.0
# via openpyxl
evaluate==0.4.3
# via lm-eval
fastapi==0.128.0
fastapi==0.136.3
# via
# -c requirements/common.txt
# gpt-oss
@@ -206,8 +200,6 @@ filelock==3.16.1
# virtualenv
fonttools==4.55.0
# via matplotlib
fqdn==1.5.1
# via jsonschema
frozendict==2.4.6
# via einx
frozenlist==1.5.0
@@ -269,7 +261,7 @@ h11==0.14.0
# uvicorn
h2==4.3.0
# via httpx
harfile==0.3.0
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
# via huggingface-hub
@@ -309,7 +301,7 @@ hypothesis==6.131.0
# hypothesis-graphql
# hypothesis-jsonschema
# schemathesis
hypothesis-graphql==0.11.1
hypothesis-graphql==0.13.0
# via schemathesis
hypothesis-jsonschema==0.23.1
# via schemathesis
@@ -318,7 +310,6 @@ idna==3.10
# anyio
# email-validator
# httpx
# jsonschema
# requests
# yarl
imagehash==4.3.2
@@ -335,8 +326,6 @@ instanttensor==0.1.5
# via -r requirements/test/cuda.in
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==5.13.2
# via datamodel-code-generator
jinja2==3.1.6
@@ -356,15 +345,14 @@ joblib==1.4.2
# librosa
# nltk
# scikit-learn
jsonpointer==3.0.0
# via jsonschema
jsonschema==4.23.0
# via
# -c requirements/common.txt
# hypothesis-jsonschema
# mistral-common
# ray
# schemathesis
jsonschema-rs==0.46.5
# via schemathesis
jsonschema-specifications==2024.10.1
# via jsonschema
junit-xml==1.9
@@ -715,18 +703,20 @@ pydantic-core==2.41.1
pydantic-extra-types==2.10.5
# via mistral-common
pygments==2.18.0
# via rich
# via
# pytest
# rich
pyjwt==2.11.0
# via msal
pyparsing==3.2.0
# via matplotlib
pyrate-limiter==3.7.0
pyrate-limiter==4.4.0
# via schemathesis
pystemmer==3.0.0
# via mteb
pytablewriter==1.2.0
# via lm-eval
pytest==8.3.5
pytest==9.1.0
# via
# -r requirements/test/cuda.in
# buildkite-test-collector
@@ -737,10 +727,9 @@ pytest==8.3.5
# pytest-mock
# pytest-rerunfailures
# pytest-shard
# pytest-subtests
# pytest-timeout
# schemathesis
pytest-asyncio==0.24.0
pytest-asyncio==1.4.0
# via -r requirements/test/cuda.in
pytest-cov==6.3.0
# via -r requirements/test/cuda.in
@@ -752,13 +741,10 @@ pytest-rerunfailures==14.0
# via -r requirements/test/cuda.in
pytest-shard==0.1.2
# via -r requirements/test/cuda.in
pytest-subtests==0.14.1
# via schemathesis
pytest-timeout==2.3.1
# via -r requirements/test/cuda.in
python-dateutil==2.9.0.post0
# via
# arrow
# botocore
# matplotlib
# pandas
@@ -829,15 +815,12 @@ requests==2.32.3
# tiktoken
responses==0.25.3
# via genai-perf
rfc3339-validator==0.1.4
# via jsonschema
rfc3987==1.3.8
# via jsonschema
rich==13.9.4
# via
# genai-perf
# mteb
# perceptron
# schemathesis
# typer
rouge-score==0.1.2
# via lm-eval
@@ -868,7 +851,7 @@ safetensors==0.7.0
# segmentation-models-pytorch
# timm
# transformers
schemathesis==3.39.15
schemathesis==4.21.6
# via -r requirements/test/cuda.in
scikit-image==0.25.2
# via albumentations
@@ -912,7 +895,6 @@ six==1.16.0
# junit-xml
# opencensus
# python-dateutil
# rfc3339-validator
# rouge-score
smart-open==7.1.0
# via ray
@@ -938,10 +920,10 @@ sqlalchemy==2.0.41
# optuna
sqlitedict==2.1.0
# via lm-eval
starlette==0.50.0
starlette==1.3.1
# via
# -c requirements/common.txt
# fastapi
# schemathesis
# starlette-testclient
starlette-testclient==0.4.1
# via schemathesis
@@ -966,7 +948,8 @@ tenacity==9.1.2
# gpt-oss
# lm-eval
# plotly
tensorizer==2.12.1
# schemathesis
tensorizer==2.10.1
# via -r requirements/test/cuda.in
termcolor==3.1.0
# via gpt-oss
@@ -990,10 +973,6 @@ tokenizers==0.22.2
# -c requirements/common.txt
# -r requirements/test/cuda.in
# transformers
tomli==2.2.1
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch==2.11.0+cu130
# via
# -c requirements/cuda.txt
@@ -1066,8 +1045,6 @@ typer==0.15.2
# huggingface-hub
# perceptron
# transformers
types-python-dateutil==2.9.0.20241206
# via arrow
typing-extensions==4.15.0
# via
# -c requirements/common.txt
@@ -1092,6 +1069,8 @@ typing-extensions==4.15.0
# pydantic
# pydantic-core
# pydantic-extra-types
# pytest-asyncio
# schemathesis
# sentence-transformers
# sqlalchemy
# starlette
@@ -1099,11 +1078,11 @@ typing-extensions==4.15.0
# typer
# typing-inspection
typing-inspection==0.4.2
# via pydantic
# via
# fastapi
# pydantic
tzdata==2024.2
# via pandas
uri-template==1.3.0
# via jsonschema
urllib3==2.2.3
# via
# blobfile
@@ -1122,8 +1101,6 @@ vocos==0.1.0
# via -r requirements/test/cuda.in
wcwidth==0.2.13
# via ftfy
webcolors==24.11.1
# via jsonschema
werkzeug==3.1.3
# via schemathesis
word2number==1.1
@@ -1135,8 +1112,6 @@ xxhash==3.5.0
# datasets
# evaluate
yarl==1.17.1
# via
# aiohttp
# schemathesis
# via aiohttp
zipp==3.23.0
# via importlib-metadata
+2 -2
View File
@@ -1,6 +1,6 @@
# testing
pytest
tensorizer==2.12.1
tensorizer==2.10.1
pytest-forked
pytest-asyncio
pytest-rerunfailures
@@ -31,7 +31,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test.
schemathesis>=4.0.0 # Required for openai schema test.
# quantization
bitsandbytes>=0.49.2
buildkite-test-collector==0.1.9
+2 -2
View File
@@ -2,7 +2,7 @@
# testing
pytest
tensorizer==2.12.1
tensorizer==2.10.1
pytest-forked
pytest-asyncio
pytest-rerunfailures
@@ -39,7 +39,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test
schemathesis>=4.0.0 # Required for openai schema test
# quantization
bitsandbytes==0.49.2
buildkite-test-collector==0.1.9
+22 -48
View File
@@ -51,15 +51,12 @@ arctic-inference==0.1.1
# via -r requirements/test/rocm.in
argcomplete==3.6.3
# via datamodel-code-generator
arrow==1.4.0
# via isoduration
astor==0.8.1
# via depyf
attrs==26.1.0
# via
# aiohttp
# jsonschema
# pytest-subtests
# referencing
audioread==3.0.1
# via librosa
@@ -74,9 +71,7 @@ azure-identity==1.25.3
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test/rocm.in
# schemathesis
# via -r requirements/test/rocm.in
bitsandbytes==0.49.2
# via -r requirements/test/rocm.in
black==26.3.1
@@ -139,7 +134,6 @@ colorama==0.4.6
# via
# perceptron
# sacrebleu
# schemathesis
colorful==0.5.8
# via ray
colorlog==6.10.1
@@ -258,8 +252,6 @@ filelock==3.25.2
# virtualenv
fonttools==4.62.1
# via matplotlib
fqdn==1.5.1
# via jsonschema
frozendict==2.4.7
# via einx
frozenlist==1.8.0
@@ -328,7 +320,7 @@ h11==0.16.0
# uvicorn
h2==4.3.0
# via httpx
harfile==0.4.0
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
# via huggingface-hub
@@ -378,7 +370,7 @@ hypothesis==6.151.9
# hypothesis-graphql
# hypothesis-jsonschema
# schemathesis
hypothesis-graphql==0.12.0
hypothesis-graphql==0.13.0
# via schemathesis
hypothesis-jsonschema==0.23.1
# via schemathesis
@@ -387,7 +379,6 @@ idna==3.11
# anyio
# email-validator
# httpx
# jsonschema
# requests
# yarl
ijson==3.5.0
@@ -408,8 +399,6 @@ interegular==0.3.3
# via lm-format-enforcer
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==8.0.1
# via datamodel-code-generator
jinja2==3.1.6
@@ -435,8 +424,6 @@ joblib==1.5.3
# librosa
# nltk
# scikit-learn
jsonpointer==3.1.0
# via jsonschema
jsonschema==4.26.0
# via
# -c requirements/common.txt
@@ -445,7 +432,8 @@ jsonschema==4.26.0
# mcp
# mistral-common
# ray
# schemathesis
jsonschema-rs==0.46.5
# via schemathesis
jsonschema-specifications==2025.9.1
# via jsonschema
junit-xml==1.9
@@ -792,7 +780,7 @@ prometheus-client==0.24.1
# opentelemetry-exporter-prometheus
# prometheus-fastapi-instrumentator
# ray
prometheus-fastapi-instrumentator==7.1.0
prometheus-fastapi-instrumentator==8.0.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -876,20 +864,22 @@ pydantic-settings==2.13.1
# fastapi
# mcp
pygments==2.19.2
# via rich
# via
# pytest
# rich
pyjwt==2.12.1
# via
# mcp
# msal
pyparsing==3.3.2
# via matplotlib
pyrate-limiter==3.9.0
pyrate-limiter==4.4.0
# via schemathesis
pystemmer==3.0.0
# via mteb
pytablewriter==1.2.1
# via lm-eval
pytest==8.3.5
pytest==9.1.0
# via
# -r requirements/test/rocm.in
# buildkite-test-collector
@@ -900,10 +890,9 @@ pytest==8.3.5
# pytest-mock
# pytest-rerunfailures
# pytest-shard
# pytest-subtests
# pytest-timeout
# schemathesis
pytest-asyncio==0.24.0
pytest-asyncio==1.4.0
# via -r requirements/test/rocm.in
pytest-cov==6.3.0
# via -r requirements/test/rocm.in
@@ -915,13 +904,10 @@ pytest-rerunfailures==14.0
# via -r requirements/test/rocm.in
pytest-shard==0.1.2
# via -r requirements/test/rocm.in
pytest-subtests==0.14.2
# via schemathesis
pytest-timeout==2.3.1
# via -r requirements/test/rocm.in
python-dateutil==2.9.0.post0
# via
# arrow
# botocore
# matplotlib
# pandas
@@ -1016,16 +1002,13 @@ requests==2.32.5
# tiktoken
responses==0.26.0
# via genai-perf
rfc3339-validator==0.1.4
# via jsonschema
rfc3987==1.3.8
# via jsonschema
rich==14.3.3
# via
# genai-perf
# mteb
# perceptron
# rich-toolkit
# schemathesis
# typer
rich-toolkit==0.19.7
# via
@@ -1063,7 +1046,7 @@ safetensors==0.7.0
# segmentation-models-pytorch
# timm
# transformers
schemathesis==3.39.15
schemathesis==4.21.6
# via -r requirements/test/rocm.in
scikit-image==0.26.0
# via albumentations
@@ -1120,7 +1103,6 @@ six==1.17.0
# junit-xml
# opencensus
# python-dateutil
# rfc3339-validator
# rouge-score
smart-open==7.5.1
# via ray
@@ -1149,13 +1131,14 @@ sqlitedict==2.1.0
# via lm-eval
sse-starlette==3.3.4
# via mcp
starlette==0.52.1
starlette==1.3.1
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
# fastapi
# mcp
# model-hosting-container-standards
# prometheus-fastapi-instrumentator
# schemathesis
# sse-starlette
# starlette-testclient
starlette-testclient==0.4.1
@@ -1182,7 +1165,8 @@ tenacity==9.1.4
# via
# gpt-oss
# lm-eval
tensorizer==2.12.1
# schemathesis
tensorizer==2.10.1
# via
# -c requirements/rocm.txt
# -r requirements/test/rocm.in
@@ -1215,10 +1199,6 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
# transformers
tomli==2.4.0
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch-c-dlpack-ext==0.1.5
# via tilelang
tqdm==4.67.3
@@ -1301,8 +1281,10 @@ typing-extensions==4.15.0
# pydantic
# pydantic-core
# pydantic-extra-types
# pytest-asyncio
# referencing
# rich-toolkit
# schemathesis
# sentence-transformers
# sqlalchemy
# starlette
@@ -1317,10 +1299,6 @@ typing-inspection==0.4.2
# mcp
# pydantic
# pydantic-settings
tzdata==2025.3
# via arrow
uri-template==1.3.0
# via jsonschema
urllib3==2.6.3
# via
# blobfile
@@ -1351,8 +1329,6 @@ watchfiles==1.1.1
# uvicorn
wcwidth==0.6.0
# via ftfy
webcolors==25.10.0
# via jsonschema
websockets==16.0
# via uvicorn
werkzeug==3.1.6
@@ -1370,9 +1346,7 @@ xxhash==3.6.0
# datasets
# evaluate
yarl==1.23.0
# via
# aiohttp
# schemathesis
# via aiohttp
z3-solver==4.15.4.0
# via tilelang
zipp==3.23.0
+2 -1
View File
@@ -593,8 +593,9 @@ soxr==0.5.0.post1
# mistral-common
sqlitedict==2.1.0
# via lm-eval
starlette==1.0.0
starlette==1.3.1
# via
# -c requirements/common.txt
# fastapi
# starlette-testclient
starlette-testclient==0.4.1
+1 -1
View File
@@ -17,4 +17,4 @@ torchaudio
torchvision
auto_round_lib>=0.13.3
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10/vllm_xpu_kernels-0.1.10-cp38-abi3-manylinux_2_28_x86_64.whl
+4 -1
View File
@@ -769,6 +769,7 @@ class precompiled_wheel_utils:
"vllm/_C.abi3.so",
"vllm/_C_stable_libtorch.abi3.so",
"vllm/_moe_C_stable_libtorch.abi3.so",
"vllm/_qutlass_C.abi3.so",
"vllm/_flashmla_C.abi3.so",
"vllm/_flashmla_extension_C.abi3.so",
"vllm/_sparse_flashmla_C.abi3.so",
@@ -1135,6 +1136,7 @@ if _is_cuda():
# DeepGEMM requires CUDA 12.3+ (SM90/SM100)
# Optional since it won't build on unsupported architectures
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
# fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
@@ -1149,7 +1151,8 @@ if _is_cpu():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _build_custom_ops():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
+193
View File
@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")
assert args.model == "s3://bucket/model"
assert args.tokenizer == "s3://bucket/tokenizer"
class TestDeviceIds:
def test_device_ids_with_cvd_out_of_range(self, monkeypatch):
"""--device-ids index beyond the CVD set raises ValueError."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "4,5")
args = EngineArgs(model="m", device_ids=[0, 2])
with pytest.raises(ValueError, match="out of range"):
args._resolve_device_ids()
def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):
"""--device-ids are CVD-local indices resolved to physical ids."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "4,5")
args = EngineArgs(model="m", device_ids=[0, 1])
assert args._resolve_device_ids() == [4, 5]
def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):
"""--device-ids support UUID CVD values resolved by the platform."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
),
)
args = EngineArgs(model="m", device_ids=[0, 1])
assert args._resolve_device_ids() == [4, 5]
def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):
"""UUID --device-ids are resolved to physical IDs immediately."""
from vllm.platforms import current_platform
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),
)
args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])
assert args._resolve_device_ids() == [4]
def test_device_ids_reject_mixed_integer_and_uuid_args(self):
"""--device-ids must not mix CVD indices and UUIDs."""
args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])
with pytest.raises(ValueError, match="must not mix"):
args._resolve_device_ids()
def test_no_device_ids(self):
"""No --device-ids returns None."""
args = EngineArgs(model="m")
assert args._resolve_device_ids() is None
def test_cli_parsing(self):
"""--device-ids parses comma-separated string from CLI."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])
assert parsed.device_ids == [0, 2, 4]
def test_cli_parsing_uuid(self):
"""--device-ids parses comma-separated UUID strings from CLI."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(
["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]
)
assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]
def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):
"""assigned_physical_gpu_ids are already physical and not composed with CVD."""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
assert current_platform.device_id_to_physical_device_id(0) == 4
assert current_platform.device_id_to_physical_device_id(1) == 5
assert current_platform.logical_device_id_to_visible_device_id(0) == 0
assert current_platform.logical_device_id_to_visible_device_id(1) == 1
def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):
"""Physical IDs map back to visible ordinals when CVD uses UUIDs."""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])
monkeypatch.setenv(
current_platform.device_control_env_var,
"GPU-abcd1234,GPU-ef567890",
)
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
),
)
assert current_platform.logical_device_id_to_visible_device_id(0) == 1
def test_device_ids_reject_duplicates(self):
"""--device-ids must not contain duplicate entries."""
args = EngineArgs(model="m", device_ids=[2, 2])
with pytest.raises(ValueError, match="duplicates"):
args._resolve_device_ids()
def test_cli_parsing_strips_whitespace(self):
"""--device-ids tolerates whitespace around commas."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])
assert parsed.device_ids == [0, 2, 4]
def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):
"""visible_device_id_to_physical_device_id maps torch device ordinals,
independent of the logical-to-physical mapping.
Regression test: CustomAllreduce passes device.index (a visible
ordinal) and must not index into assigned_physical_gpu_ids, which
raised IndexError for non-identity --device-ids like [2, 3].
"""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])
monkeypatch.delenv(current_platform.device_control_env_var, raising=False)
# CVD unset: visible ordinal == physical ID, even beyond the
# assigned list's length.
assert current_platform.visible_device_id_to_physical_device_id(2) == 2
assert current_platform.visible_device_id_to_physical_device_id(3) == 3
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
assert current_platform.visible_device_id_to_physical_device_id(1) == 5
with pytest.raises(IndexError, match="out of range"):
current_platform.visible_device_id_to_physical_device_id(2)
class TestDpDeviceIdSharding:
def test_dp_supervisor_device_ids_stay_env_relative(self):
"""Regression test: the DP supervisor must pass env-relative indices,
not physical IDs, because each child re-resolves --device-ids
against its inherited device-control env var."""
import argparse
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
args = argparse.Namespace(
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None
)
assert _build_device_ids(args, local_rank=0) == [0, 1]
assert _build_device_ids(args, local_rank=1) == [2, 3]
def test_dp_supervisor_shards_user_device_ids(self):
"""User-provided --device-ids are sharded across DP children."""
import argparse
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
args = argparse.Namespace(
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]
)
assert _build_device_ids(args, local_rank=0) == [4, 5]
assert _build_device_ids(args, local_rank=1) == [6, 7]
with pytest.raises(ValueError, match="needs devices"):
_build_device_ids(args, local_rank=2)
def test_dp_rank_shards_user_assigned_gpu_ids(self):
"""get_physical_gpu_ids_for_local_dp_rank slices the user-provided
--device-ids list instead of recomputing from the env var."""
from vllm.platforms import current_platform
from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
evar = current_platform.device_control_env_var
assert get_physical_gpu_ids_for_local_dp_rank(
evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
) == [6, 7]
with pytest.raises(ValueError, match="needs devices"):
get_physical_gpu_ids_for_local_dp_rank(
evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
)
@@ -8,6 +8,8 @@ AnthropicServingMessages._convert_anthropic_to_openai_request().
Also covers extended-thinking edge cases such as ``redacted_thinking``
blocks echoed back by Anthropic clients, and streaming conversion in
``message_stream_converter``.
Also covers cache usage computation in ``_build_anthropic_usage``.
"""
import json
@@ -18,7 +20,11 @@ import pytest
from vllm.entrypoints.anthropic.protocol import (
AnthropicMessagesRequest,
)
from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
from vllm.entrypoints.anthropic.serving import (
AnthropicServingMessages,
_build_anthropic_usage,
_get_cached_tokens,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionResponseStreamChoice,
ChatCompletionStreamResponse,
@@ -27,6 +33,7 @@ from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
PromptTokenUsageInfo,
UsageInfo,
)
@@ -653,6 +660,108 @@ class TestThinkingBlockConversion:
assert asst.get("content") == "Hi!"
# ======================================================================
# Cache usage computation
# ======================================================================
class TestGetCachedTokens:
"""Tests for _get_cached_tokens helper."""
def test_none_usage(self):
assert _get_cached_tokens(None) is None
def test_no_prompt_tokens_details(self):
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
assert _get_cached_tokens(usage) is None
def test_cached_tokens_present(self):
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
)
assert _get_cached_tokens(usage) == 80
def test_cached_tokens_zero(self):
"""Zero cached tokens should return 0, not None."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
)
assert _get_cached_tokens(usage) == 0
def test_cached_tokens_none_in_details(self):
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
)
assert _get_cached_tokens(usage) is None
class TestBuildAnthropicUsage:
"""Tests for _build_anthropic_usage helper.
Anthropic defines: total_input = input_tokens + cache_read + cache_creation
vLLM's prompt_tokens is the total.
"""
def test_no_cache_info(self):
"""When cache info is unavailable, return raw prompt_tokens."""
result = _build_anthropic_usage(100, 10, None)
assert result.input_tokens == 100
assert result.output_tokens == 10
assert result.cache_read_input_tokens is None
assert result.cache_creation_input_tokens is None
def test_cache_hit(self):
"""When cache is hit, input_tokens excludes cached tokens."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 20 # 100 - 80
assert result.output_tokens == 10
assert result.cache_read_input_tokens == 80
assert result.cache_creation_input_tokens == 0
def test_zero_cached_tokens(self):
"""Zero cached tokens should still set cache_creation to 0."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 100 # 100 - 0
assert result.cache_read_input_tokens == 0
assert result.cache_creation_input_tokens == 0
def test_all_tokens_cached(self):
"""When all tokens are cached, input_tokens should be 0."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 0
assert result.cache_read_input_tokens == 100
assert result.cache_creation_input_tokens == 0
def test_no_prompt_tokens_details(self):
"""UsageInfo without prompt_tokens_details returns no cache info."""
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 100
assert result.cache_read_input_tokens is None
assert result.cache_creation_input_tokens is None
class TestInlineSystemMessageInMessagesArray:
"""Verify that ``role: system`` messages embedded inside the ``messages``
array are preserved in their original position.
@@ -1096,3 +1205,179 @@ class TestMessageStartIncludesTypeAndRole:
message = events[0][1]["message"]
assert message["type"] == "message"
assert message["role"] == "assistant"
class TestStreamingCacheUsageSemantics:
"""Locks in the documented streaming behavior of cache usage fields.
vLLM's OpenAI chat completion streaming only attaches
``prompt_tokens_details`` to the terminal usage chunk. The Anthropic layer
mirrors that contract: cache fields are omitted on ``message_start`` (key
absence signals "unknown") and populated on ``message_delta`` (the final
cumulative count). This is intentionally consistent with vLLM's OpenAI
behavior, even though Anthropic's upstream API populates cache fields on
``message_start``; closing that gap requires plumbing cache info into the
first chunk at the OpenAI layer, which is out of scope here.
"""
@pytest.mark.asyncio
async def test_streaming_cache_fields_absent_then_populated(self):
"""First chunk lacks prompt_tokens_details (vLLM contract);
message_start omits cache fields. The final chunk carries
prompt_tokens_details, so message_delta carries resolved values."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant", content="hi"),
usage=UsageInfo(prompt_tokens=100, total_tokens=100),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(
prompt_tokens=100,
completion_tokens=5,
total_tokens=105,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
output.append(event)
events = _parse_sse_events(output)
# message_start: cache fields unknown → omitted from JSON entirely.
start_usage = events[0][1]["message"]["usage"]
assert events[0][0] == "message_start"
assert start_usage["input_tokens"] == 100
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
# message_delta: authoritative usage with cache fields populated.
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert delta_usage["input_tokens"] == 20 # 100 - 80
assert delta_usage["cache_read_input_tokens"] == 80
assert delta_usage["cache_creation_input_tokens"] == 0
@pytest.mark.asyncio
async def test_streaming_no_cache_hit(self):
"""When the final chunk reports cached_tokens=0, message_delta carries
cache fields = 0 (cache miss); message_start still omits them."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=50, total_tokens=50),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(
prompt_tokens=50,
completion_tokens=5,
total_tokens=55,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
output.append(event)
events = _parse_sse_events(output)
start_usage = events[0][1]["message"]["usage"]
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert start_usage["input_tokens"] == 50
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
assert delta_usage["input_tokens"] == 50 # 50 - 0
assert delta_usage["cache_read_input_tokens"] == 0
assert delta_usage["cache_creation_input_tokens"] == 0
@pytest.mark.asyncio
async def test_streaming_no_prompt_tokens_details_at_all(self):
"""If --enable-prompt-tokens-details is off, no chunk carries cache
info; both message_start and message_delta omit cache fields."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=30, total_tokens=30),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
output.append(event)
events = _parse_sse_events(output)
start_usage = events[0][1]["message"]["usage"]
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
assert "cache_read_input_tokens" not in delta_usage
assert "cache_creation_input_tokens" not in delta_usage
# ======================================================================
# Auto-detection of system-first template requirement
# ======================================================================
Q35_TEMPLATE = (
"{%- for message in messages %}"
"{%- if message.role == 'system' %}"
"{%- if not loop.first %}"
"{{- raise_exception('System message must be at the beginning.') }}"
"{%- endif %}"
"{%- endif %}"
"{%- endfor %}"
)
class TestDetectMergeInlineSystem:
"""Verify _detect_merge_inline_system auto-detection.
Tests three scenarios:
1. Template with system-first guard (e.g. Qwen) → merge needed
2. Template without restrictions → no merge, cache-friendly
3. No template provided → safe default: merge
"""
def test_qwen_template_requires_merge(self):
"""Template with loop.first guard rejects mid-conversation system."""
assert (
AnthropicServingMessages._detect_merge_inline_system(Q35_TEMPLATE) is True
)
def test_no_restriction_no_merge(self):
"""Template without restriction accepts mid-conversation system."""
assert (
AnthropicServingMessages._detect_merge_inline_system(
"{%- for message in messages %}"
"{{- message.role }}: {{ message.content }}\n"
"{%- endfor %}"
)
is False
)
def test_no_template_defaults_merge(self):
"""No chat_template → conservative default: merge."""
assert AnthropicServingMessages._detect_merge_inline_system(None) is True
@@ -364,7 +364,7 @@ class MockVLLMServer:
await self._serve_task
def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
def launch_mock_vllm(child_args: argparse.Namespace):
logger.info("Launching mock vLLM on port %s", child_args.port)
mock_vllm = MockVLLMServer(
port=child_args.port,
@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]
def launch_mock_vllm_with_drain(
child_args: argparse.Namespace, env_updates: dict[str, str]
child_args: argparse.Namespace,
):
logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
mock_vllm = MockVLLMServer(
+57 -44
View File
@@ -6,15 +6,22 @@ from typing import Final
import pytest
import schemathesis
from hypothesis import HealthCheck, settings
from schemathesis import GenerationConfig
from schemathesis.models import Case
from schemathesis import GenerationMode
from schemathesis.config import (
ChecksConfig,
CoveragePhaseConfig,
GenerationConfig,
PhasesConfig,
PositiveDataAcceptanceConfig,
ProjectConfig,
ProjectsConfig,
SchemathesisConfig,
)
from vllm.platforms import current_platform
from ...utils import RemoteOpenAIServer
schemathesis.experimental.OPEN_API_3_1.enable()
MODEL_NAME = "HuggingFaceTB/SmolVLM-256M-Instruct"
MAXIMUM_IMAGES = 2
_ROCM_TIMEOUT_MULTIPLIER = 3 if current_platform.is_rocm() else 1
@@ -44,21 +51,38 @@ def server():
@pytest.fixture(scope="module")
def get_schema(server):
# avoid generating null (\x00) bytes in strings during test case generation
return schemathesis.openapi.from_uri(
return schemathesis.openapi.from_url(
f"{server.url_root}/openapi.json",
generation_config=GenerationConfig(allow_x00=False),
config=SchemathesisConfig(
projects=ProjectsConfig(
default=ProjectConfig(
generation=GenerationConfig(
allow_x00=False,
modes=[GenerationMode.POSITIVE],
),
checks=ChecksConfig(
positive_data_acceptance=PositiveDataAcceptanceConfig(
enabled=False,
),
),
phases=PhasesConfig(
coverage=CoveragePhaseConfig(enabled=False),
),
),
),
),
)
schema = schemathesis.from_pytest_fixture("get_schema")
schema = schemathesis.pytest.from_fixture("get_schema")
@schemathesis.hook
def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
def before_generate_case(context: schemathesis.HookContext, strategy):
op = context.operation
assert op is not None
def no_invalid_types(case: schemathesis.models.Case):
def no_invalid_types(case: schemathesis.Case):
"""
Skips tool_calls with `"type": "custom"` which schemathesis incorrectly
generates instead of the valid `"type": "function"`.
@@ -68,39 +92,25 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
-d '{"messages": [{"role": "assistant", "tool_calls": [{"custom": {"input": "", "name": ""}, "id": "", "type": "custom"}]}]}' \
http://localhost:8000/v1/chat/completions
""" # noqa: E501
if hasattr(case, "body") and isinstance(case.body, dict):
if (
"messages" in case.body
and isinstance(case.body["messages"], list)
and len(case.body["messages"]) > 0
):
for message in case.body["messages"]:
if not isinstance(message, dict):
continue
if (
hasattr(case, "body")
and isinstance(case.body, dict)
and "messages" in case.body
and isinstance(case.body["messages"], list)
and len(case.body["messages"]) > 0
):
for message in case.body["messages"]:
if not isinstance(message, dict):
continue
tool_calls = message.get("tool_calls", [])
if isinstance(tool_calls, list):
for tool_call in tool_calls:
if isinstance(tool_call, dict):
if tool_call.get("type") != "function":
return False
if "custom" in tool_call:
return False
# Sometimes structured_outputs.grammar is generated to be empty
# Causing a server error in EBNF grammar parsing
# https://github.com/vllm-project/vllm/pull/22587#issuecomment-3195253421
structured_outputs = case.body.get("structured_outputs", {})
grammar = (
structured_outputs.get("grammar")
if isinstance(structured_outputs, dict)
else None
)
if grammar == "":
# Allow None (will be handled as no grammar)
# But skip empty strings
return False
tool_calls = message.get("tool_calls", [])
if isinstance(tool_calls, list):
for tool_call in tool_calls:
if isinstance(tool_call, dict):
if tool_call.get("type") != "function":
return False
if "custom" in tool_call:
return False
return True
@@ -108,7 +118,6 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
@schema.parametrize()
@schema.override(headers={"Content-Type": "application/json"})
@settings(
deadline=LONG_TIMEOUT_SECONDS * 1000,
max_examples=50,
@@ -122,7 +131,7 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
# generating large-but-valid request bodies before vLLM is called.
suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.data_too_large],
)
def test_openapi_stateless(case: Case):
def test_openapi_stateless(case: schemathesis.Case):
key = (
case.operation.method.upper(),
case.operation.path,
@@ -151,4 +160,8 @@ def test_openapi_stateless(case: Case):
}.get(key, DEFAULT_TIMEOUT_SECONDS)
# No need to verify SSL certificate for localhost
case.call_and_validate(verify=False, timeout=timeout)
case.call_and_validate(
verify=False,
timeout=timeout,
headers={"Content-Type": "application/json"},
)
@@ -8,6 +8,7 @@ import pytest
import pytest_asyncio
from tests.utils import RemoteLaunchRenderServer
from vllm.tokenizers import get_tokenizer
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
@@ -486,3 +487,438 @@ async def test_derender_completion_kv_transfer_params_passthrough(client):
)
assert response.status_code == 200
assert response.json()["kv_transfer_params"] == kv
# ---------------------------------------------------------------------------
# E2E: render -> derender roundtrip with parser (reasoning + tool calls)
# ---------------------------------------------------------------------------
PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
_E2E_TOOLS = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a city",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
},
},
}
]
@pytest.fixture(scope="module")
def parser_server():
args = [
"--enable-auto-tool-choice",
"--tool-call-parser",
"hermes",
"--reasoning-parser",
"deepseek_r1",
]
with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def parser_client(parser_server):
async with httpx.AsyncClient(
base_url=parser_server.url_for(""), timeout=60.0
) as http_client:
yield http_client
@pytest.fixture(scope="module")
def parser_tokenizer():
return get_tokenizer(PARSER_MODEL)
def _encode(tokenizer, text: str) -> list[int]:
return tokenizer.encode(text, add_special_tokens=False)
def _decoded(tokenizer, token_ids: list[int]) -> str:
return tokenizer.decode(token_ids, skip_special_tokens=True)
def _require_markers_survive(tokenizer, text: str, *markers: str) -> list[int]:
"""Encode text and skip the test if any marker is lost in roundtrip."""
ids = _encode(tokenizer, text)
decoded = tokenizer.decode(ids, skip_special_tokens=False)
for m in markers:
if m not in decoded:
pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip")
return ids
async def _e2e_render_chat(
client: httpx.AsyncClient,
model: str,
messages: list[dict],
) -> dict:
resp = await client.post(
"/v1/chat/completions/render",
json={"model": model, "messages": messages},
)
assert resp.status_code == 200, resp.text
return resp.json()
def _e2e_generate_response(
token_ids: list[int],
request_id: str = "chatcmpl-e2e-test",
) -> dict:
return {
"request_id": request_id,
"choices": [
{
"index": 0,
"token_ids": token_ids,
"finish_reason": "stop",
}
],
}
@pytest.mark.asyncio
async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer):
"""Plain text without reasoning markers roundtrips correctly."""
messages = [{"role": "user", "content": "What is 2+2?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
answer = "The answer is four."
output_ids = _encode(parser_tokenizer, answer)
expected = _decoded(parser_tokenizer, output_ids)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
},
)
assert resp.status_code == 200, resp.text
content = resp.json()["choices"][0]["message"]["content"]
assert content == expected
@pytest.mark.asyncio
async def test_e2e_token_identity(parser_client, parser_tokenizer):
"""encode(derender(token_ids)) == token_ids (RL invariant)."""
messages = [{"role": "user", "content": "Hi"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
answer = "Hello! How can I help?"
output_ids = _encode(parser_tokenizer, answer)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
},
)
assert resp.status_code == 200
content = resp.json()["choices"][0]["message"]["content"]
re_encoded = _encode(parser_tokenizer, content)
assert output_ids == re_encoded
@pytest.mark.asyncio
async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer):
"""CJK + emoji roundtrip without U+FFFD."""
messages = [{"role": "user", "content": "Reply in Chinese"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
answer = "你好世界 😀"
output_ids = _encode(parser_tokenizer, answer)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
},
)
assert resp.status_code == 200
content = resp.json()["choices"][0]["message"]["content"]
assert "" not in content
@pytest.mark.asyncio
async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer):
"""<think>...</think> splits into reasoning + content."""
messages = [{"role": "user", "content": "What is 2+3?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
reasoning_text = "The user wants 2 plus 3. That is 5."
answer_text = "The answer is 5."
output_text = f"<think>{reasoning_text}</think>{answer_text}"
output_ids = _require_markers_survive(parser_tokenizer, output_text, "</think>")
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": PARSER_MODEL,
"messages": messages,
"include_reasoning": True,
},
},
)
assert resp.status_code == 200, resp.text
msg = resp.json()["choices"][0]["message"]
assert msg["reasoning"] is not None
assert reasoning_text in msg["reasoning"]
assert answer_text in msg["content"]
assert "<think>" not in msg["content"]
@pytest.mark.asyncio
async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer):
"""<tool_call> extracted into tool_calls field."""
messages = [{"role": "user", "content": "Weather in Paris?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
output_text = (
"<think>Let me check the weather.</think>"
'<tool_call>\n{"name": "get_weather", '
'"arguments": {"city": "Paris"}}\n</tool_call>'
)
output_ids = _require_markers_survive(
parser_tokenizer,
output_text,
"</think>",
"<tool_call>",
"</tool_call>",
)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": PARSER_MODEL,
"messages": messages,
"tools": _E2E_TOOLS,
"tool_choice": "auto",
},
},
)
assert resp.status_code == 200, resp.text
choice = resp.json()["choices"][0]
assert choice["message"]["tool_calls"]
assert choice["message"]["tool_calls"][0]["function"]["name"] == "get_weather"
@pytest.mark.asyncio
async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenizer):
"""Reasoning + tool call in the same output."""
messages = [{"role": "user", "content": "Weather in Paris?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
reasoning_text = "I should look up the weather."
tool_text = (
'<tool_call>\n{"name": "get_weather", '
'"arguments": {"city": "Paris"}}\n</tool_call>'
)
output_text = f"<think>{reasoning_text}</think>{tool_text}"
output_ids = _require_markers_survive(
parser_tokenizer, output_text, "</think>", "<tool_call>"
)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": PARSER_MODEL,
"messages": messages,
"tools": _E2E_TOOLS,
"tool_choice": "auto",
"include_reasoning": True,
},
},
)
assert resp.status_code == 200, resp.text
choice = resp.json()["choices"][0]
assert choice["message"]["reasoning"] is not None
assert reasoning_text in choice["message"]["reasoning"]
assert choice["message"]["tool_calls"]
@pytest.mark.asyncio
async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer):
"""Without chat_request, derender falls back to plain detokenization."""
messages = [{"role": "user", "content": "Hello"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
answer = "Hi there!"
output_ids = _encode(parser_tokenizer, answer)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
},
)
assert resp.status_code == 200
content = resp.json()["choices"][0]["message"]["content"]
assert "Hi" in content
# ---------------------------------------------------------------------------
# E2E: HarmonyParser + GPT-OSS
# ---------------------------------------------------------------------------
HARMONY_MODEL = "openai/gpt-oss-20b"
def _ensure_harmony_vocab():
"""Pre-cache the o200k_base BPE file needed by openai-harmony.
The Rust tiktoken-rs backend downloads from Azure Blob Storage, which
may be unreachable in some environments. When the cache is cold we
fetch the file ourselves and place it in ``/tmp/tiktoken-rs-cache/``
using the SHA-1(URL) filename that tiktoken-rs expects.
"""
import hashlib
import urllib.request
from pathlib import Path
url = "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
cache_dir = Path("/tmp/tiktoken-rs-cache")
cache_key = hashlib.sha1(url.encode()).hexdigest()
cache_file = cache_dir / cache_key
if not cache_file.exists():
cache_dir.mkdir(parents=True, exist_ok=True)
urllib.request.urlretrieve(url, cache_file)
@pytest.fixture(scope="module")
def harmony_server():
_ensure_harmony_vocab()
args = [
"--trust-remote-code",
"--enable-auto-tool-choice",
"--tool-call-parser",
"openai",
"--reasoning-parser",
"openai_gptoss",
]
with RemoteLaunchRenderServer(HARMONY_MODEL, args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def harmony_client(harmony_server):
async with httpx.AsyncClient(
base_url=harmony_server.url_for(""), timeout=60.0
) as http_client:
yield http_client
@pytest.fixture(scope="module")
def harmony_tokenizer():
return get_tokenizer(HARMONY_MODEL, trust_remote_code=True)
def _harmony_extract_assistant_ids(
tokenizer, assistant_msg: dict, user_content: str = "test"
) -> list[int]:
"""Extract assistant token IDs via apply_chat_template diff."""
prompt = [{"role": "user", "content": user_content}]
full = prompt + [assistant_msg]
text_prompt = tokenizer.apply_chat_template(
prompt, add_generation_prompt=True, tokenize=False
)
text_full = tokenizer.apply_chat_template(
full, add_generation_prompt=False, tokenize=False
)
prompt_ids = tokenizer.encode(text_prompt)
full_ids = tokenizer.encode(text_full)
assistant_ids = list(full_ids[len(prompt_ids) :])
if not assistant_ids:
pytest.skip("Could not extract assistant tokens for Harmony")
return assistant_ids
@pytest.mark.asyncio
async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer):
"""GPT-OSS content-only roundtrip."""
messages = [{"role": "user", "content": "What is 2+2?"}]
gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
assistant_msg = {"role": "assistant", "content": "Four."}
output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
resp = await harmony_client.post(
"/v1/chat/completions/derender",
json={
"model": HARMONY_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": HARMONY_MODEL,
"messages": messages,
},
},
)
assert resp.status_code == 200, resp.text
content = resp.json()["choices"][0]["message"]["content"]
assert content is not None and len(content) > 0
assert "Four" in content
@pytest.mark.asyncio
async def test_e2e_harmony_reasoning(harmony_client, harmony_tokenizer):
"""GPT-OSS reasoning: analysis channel extracted."""
messages = [{"role": "user", "content": "Add 2 and 3."}]
gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
reasoning_text = "The user wants 2 plus 3."
answer_text = "The answer is 5."
assistant_msg = {
"role": "assistant",
"thinking": reasoning_text,
"content": answer_text,
}
output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
decoded = harmony_tokenizer.decode(output_ids)
if reasoning_text not in decoded:
pytest.skip("Harmony template did not render thinking")
resp = await harmony_client.post(
"/v1/chat/completions/derender",
json={
"model": HARMONY_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": HARMONY_MODEL,
"messages": messages,
"include_reasoning": True,
},
},
)
assert resp.status_code == 200, resp.text
msg = resp.json()["choices"][0]["message"]
assert msg["reasoning"] is not None
assert reasoning_text in msg["reasoning"]
assert answer_text in (msg["content"] or "")
@@ -0,0 +1,65 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests that validation_exception_handler populates the `param` field
in its error response using the Pydantic error's `loc`, even when no
custom VLLMValidationError context is present.
Previously, `param` was only populated for errors carrying a custom
VLLMValidationError in their Pydantic `ctx`. Plain validation failures
(missing fields, wrong types) left `param` as None, even though the
field name was readily available from `error['loc']`.
"""
import json
from types import SimpleNamespace
import pytest
from fastapi.exceptions import RequestValidationError
from vllm.entrypoints.serve.utils.server_utils import validation_exception_handler
def _fake_request(log_error_stack: bool = False) -> SimpleNamespace:
"""Minimal stand-in for a FastAPI Request - just enough for the
handler to read req.app.state.args.log_error_stack."""
return SimpleNamespace(
app=SimpleNamespace(
state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack))
),
state=SimpleNamespace(), # no request_metadata -> hasattr(...) is False
)
class TestValidationErrorParamFallback:
"""Ensure `param` falls back to the Pydantic error's `loc` when no
custom VLLMValidationError context is present."""
@pytest.mark.parametrize(
("error_type", "msg"),
[
("missing", "Field required"),
("list_type", "Input should be a valid list"),
],
ids=["missing-field", "wrong-type"],
)
@pytest.mark.asyncio
async def test_param_falls_back_to_loc(self, error_type: str, msg: str):
errors = [{"type": error_type, "loc": ("body", "messages"), "msg": msg}]
exc = RequestValidationError(errors)
response = await validation_exception_handler(_fake_request(), exc)
body = json.loads(response.body)
assert body["error"]["param"] == "body.messages"
@pytest.mark.asyncio
async def test_param_fallback_does_not_crash_on_non_dict_error(self):
"""Schemathesis fuzzing found that errors[0] isn't always a dict.
The fallback must not crash in that case - it should just leave
param as None instead of raising."""
exc = RequestValidationError(["some unexpected non-dict error"])
response = await validation_exception_handler(_fake_request(), exc)
body = json.loads(response.body)
assert body["error"]["param"] is None
@@ -78,7 +78,16 @@ def test_gsm8k_correctness(config_filename):
"Skipping DeepSeek-V3.2 and DeepSeek-R1 on ROCm platforms "
"due to agent pool disk space issues and pod evictions."
)
if current_platform.is_rocm() and (
"Qwen3.5-35B-A3B-MXFP4-AITER-TP2" in config_filename.name
):
from vllm.platforms.rocm import on_gfx950
if not on_gfx950():
pytest.skip(
"Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. "
"The quantization scheme is not supported on non-GFX950 platforms."
)
# Parse server arguments from config (use shlex to handle quoted strings)
server_args_str = eval_config.get("server_args", "")
server_args = shlex.split(server_args_str) if server_args_str else []
@@ -345,6 +345,63 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
)
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU",
)
def test_per_token_group_quant_fp8_packed_large_mn():
"""Regression test for https://github.com/vllm-project/vllm/issues/45099.
Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and
2^16 - 1, respectively.
Prior code introduced a bug where it incorrectly assumed grid.x and y both have
2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is
launched with a large mn that exceeds grid.y limit (2^16 - 1).
This issue doesn't surface often because each forward pass only processes a
bounded token batch, not the full context.
Quantizing tensors with more rows than that will fail at launch with
"CUDA error: invalid argument".
This is a differential test that compares fp8 output against Triton output
reference when token size sits just above the gridDim.y 2^16 - 1 limit.
"""
device = "cuda"
group_size = 128
# hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
# row, so any mn > 65535 overflowed grid.y before the fix.
num_tokens, hidden_dim = 65537, 2048
torch.manual_seed(42)
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
x,
group_size=group_size,
use_ue8m0=True,
)
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
x, group_size, use_ue8m0=True
)
assert torch.equal(out_q, ref_q), "Quantized output mismatch"
# Vectorized packed-scale check; the per-element loop used by the smaller
# tests is too slow at this size. groups_per_row is a multiple of 4 here,
# so there is no K padding and the packed view lines up.
mn = num_tokens
groups_per_row = hidden_dim // group_size
k_num_packed = (groups_per_row + 3) // 4
assert groups_per_row % 4 == 0
ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF
exp = ref_exponents.view(mn, k_num_packed, 4)
expected = (
exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24)
)
assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@@ -60,8 +60,10 @@ def test_rocm_compressed_tensors_w8a8(
vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
):
dtype = "bfloat16"
with vllm_runner(model_path, dtype=dtype) as vllm_model:
# Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179
with vllm_runner(
model_path, dtype=dtype, attention_backend="TRITON_ATTN"
) as vllm_model:
vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)
+124
View File
@@ -482,3 +482,127 @@ def test_kernels_hidden_size(
seq_length=128,
add_inputs=True,
)
@pytest.mark.parametrize("device", DEVICES)
def test_add_lora_fused_moe_early_exit(device):
"""
Ensures add_lora_fused_moe does not invoke the LoRA kernel or
modify the output tensor when no_lora_flag_cpu is True
"""
from types import SimpleNamespace
from vllm.lora.punica_wrapper.punica_gpu import PunicaWrapperGPU
torch.set_default_device(device)
torch.accelerator.set_device_index(device)
max_loras, num_tokens = 4, 16
num_experts, top_k, max_lora_rank = 8, 2, 16
K, N = 256, 128
# build PunicaWrapperGPU with minimal lora_config mock
lora_config = SimpleNamespace(
max_loras=max_loras,
specialize_active_lora=False,
)
wrapper = PunicaWrapperGPU(
max_num_batched_tokens=num_tokens,
max_batches=num_tokens,
device=device,
lora_config=lora_config,
)
# simulate a prior LoRA batch so the internal mapping is
# populated with stale LoRA IDs
lora_mapping = torch.zeros(
num_tokens,
dtype=torch.int32,
device=device,
)
lora_mapping[:8] = 1
lora_mapping[8:] = 2
wrapper.token_mapping_meta.prepare_tensors(lora_mapping)
# simulate a base-model batch (all -1)
base_mapping = torch.full(
(num_tokens,),
-1,
dtype=torch.int32,
device=device,
)
wrapper.token_mapping_meta.prepare_tensors(base_mapping)
assert wrapper.token_mapping_meta.no_lora_flag_cpu[0].item() is True
# dummy tensors for add_lora_fused_moe
y = torch.rand(num_tokens, top_k, N, dtype=torch.bfloat16, device=device)
y_snapshot = y.clone()
x = torch.rand(num_tokens, K, dtype=torch.bfloat16, device=device)
lora_a_stacked = (
torch.rand(
max_loras,
num_experts,
max_lora_rank,
K,
dtype=torch.bfloat16,
device=device,
),
)
lora_b_stacked = (
torch.rand(
max_loras,
num_experts,
N,
max_lora_rank,
dtype=torch.bfloat16,
device=device,
),
)
topk_weights = torch.ones(
num_tokens,
top_k,
dtype=torch.float32,
device=device,
)
adapter_enabled = torch.ones(
max_loras + 1,
dtype=torch.int32,
device=device,
)
shrink_config = expand_config = {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"NUM_WARPS": 4,
"NUM_STAGES": 3,
"SPLIT_K": 1,
}
# call add_lora_fused_moe - the early exit should prevent any
# modification to the output
wrapper.add_lora_fused_moe(
y=y,
x=x,
lora_a_stacked=lora_a_stacked,
lora_b_stacked=lora_b_stacked,
topk_weights=topk_weights,
sorted_token_ids=None,
expert_ids=torch.zeros(
num_tokens * top_k,
dtype=torch.int32,
device=device,
),
num_tokens_post_padded=None,
max_lora_rank=max_lora_rank,
top_k_num=top_k,
shrink_config=shrink_config,
expand_config=expand_config,
adapter_enabled=adapter_enabled,
)
assert torch.equal(y, y_snapshot), (
"add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True"
)
@@ -130,8 +130,12 @@ def test_models(
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
if model == "TitanML/tiny-mixtral":
# Untrained model: near-uniform logits make argmax sensitive to
# AITER's bfloat16 rounding error in plain rms_norm.
# AITER's bfloat16 rounding error. Route the plain rms_norm and the
# fused MoE (whose near-uniform router logits flip expert selection
# under ~1 ULP drift) through the native kernels for this model.
# See ROCm/aiter#3806 for the tracking issue and minimal repro.
monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
elif use_rocm_aiter and model not in AITER_MODEL_LIST:
# Skip model that are not using AITER tests.
# When more AITER kernels are added, this list will not be
@@ -25,7 +25,7 @@ TEST_IMAGE_NAMES = [
]
MAX_MODEL_LEN = 8192
REQUESTS_PER_ROUND = 4
WARMUP_ROUNDS = 1
WARMUP_ROUNDS = 2
MEASURED_ROUNDS = 16
GPU_GROWTH_THRESHOLD_MIB = 0
CPU_PEAK_GROWTH_THRESHOLD_MIB = 0
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
import transformers.utils
from PIL import Image
@@ -52,6 +53,7 @@ def _get_cherry_blossom_image() -> Image.Image:
)
@torch.inference_mode()
def _run_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
@@ -92,3 +92,49 @@ def test_processor_num_frames_timestamp(
assert len(video_phs) == 1, (
f"Expected exactly 1 video placeholder, got {len(video_phs)}"
)
@pytest.mark.parametrize("model_id", [MODEL_ID])
@pytest.mark.parametrize("num_videos", [2, 4])
def test_processor_multi_video(
model_id: str,
num_videos: int,
) -> None:
"""Verify that multi-video processing produces correct placeholders.
This exercises the token-level replacement path in
``_call_hf_processor`` which avoids the quadratic text-level
prompt expansion.
"""
ctx = build_model_context(
model_id,
limit_mm_per_prompt={"image": 0, "video": num_videos},
)
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
prompt = "<|vision_start|><|video_pad|><|vision_end|>" * num_videos
mm_data = {"video": [_build_video_mm_data(num_frames=8)["video"][0]] * num_videos}
processed = processor(
prompt,
mm_items=processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs={"num_frames": 8},
)
token_ids = processed["prompt_token_ids"]
assert len(token_ids) > 0
video_phs = processed["mm_placeholders"].get("video", [])
assert len(video_phs) == num_videos, (
f"Expected {num_videos} video placeholders, got {len(video_phs)}"
)
# All placeholders should have the same length (same video params)
# and must not overlap.
lengths = {ph.length for ph in video_phs}
assert len(lengths) == 1, f"Placeholder lengths differ: {lengths}"
for i in range(1, len(video_phs)):
prev_end = video_phs[i - 1].offset + video_phs[i - 1].length
assert video_phs[i].offset >= prev_end, (
f"Placeholder {i} overlaps with placeholder {i - 1}"
)
+3
View File
@@ -96,6 +96,9 @@ class MockTokenizer:
return "".join(parts)
CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
def make_mock_tokenizer(sample: Sample) -> MockTokenizer:
"""Build a mock tokenizer from a sample's vocab and token data."""
return MockTokenizer(
@@ -19,6 +19,7 @@ import pytest
from pydantic import TypeAdapter
from tests.parser.engine.replay_harness import (
CHUNK_SIZES,
MockTokenizer,
assert_parse_output,
collect_output,
@@ -113,8 +114,6 @@ _PAIRINGS = _discover_pairings()
_ALL_SAMPLES = [(p.parser_cls, s) for p in _PAIRINGS for s in p.samples]
CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
@pytest.mark.parametrize(
@@ -0,0 +1,181 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression test for U+FFFD leak at reasoning→content transition.
When byte-fallback tokens span the reasoning/content boundary,
decoding isolated content-side token IDs via tokenizer.decode()
produces U+FFFD (Unicode replacement character). The fix flushes
the reasoning parser's engine lexer instead.
Reproduces the bug at various chunk sizes and validates that the
fix prevents U+FFFD from leaking into streamed content.
"""
from __future__ import annotations
import pytest
from tests.parser.engine.replay_harness import (
CHUNK_SIZES,
MockTokenizer,
collect_output,
replay_streaming,
)
from vllm.parser.abstract_parser import DelegatingParser
from vllm.parser.engine.registered_adapters import (
Glm47MoeParserReasoningAdapter,
Glm47MoeParserToolAdapter,
Qwen3ParserReasoningAdapter,
Qwen3ParserToolAdapter,
)
class ByteFallbackMockTokenizer(MockTokenizer):
"""MockTokenizer that returns U+FFFD for specified token IDs.
Simulates byte-fallback tokenizer behavior where isolated
partial-byte tokens decode to the Unicode replacement character.
"""
def __init__(
self,
vocab: dict[str, int],
tokens: list[tuple[int, str]],
ufffd_token_ids: set[int],
) -> None:
super().__init__(vocab, tokens)
self._ufffd_token_ids = frozenset(ufffd_token_ids)
def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
parts: list[str] = []
for tid in ids:
if skip_special_tokens and tid in self._special_ids:
continue
if tid in self._ufffd_token_ids:
parts.append("")
else:
text = self._token_decode_map.get(tid, f"?{tid}?")
parts.append(text)
return "".join(parts)
# ── Model-specific DelegatingParser subclasses ───────────────────────
class _Glm47Delegating(DelegatingParser):
reasoning_parser_cls = Glm47MoeParserReasoningAdapter
tool_parser_cls = Glm47MoeParserToolAdapter
class _Qwen3Delegating(DelegatingParser):
reasoning_parser_cls = Qwen3ParserReasoningAdapter
tool_parser_cls = Qwen3ParserToolAdapter
# ── Shared test data ─────────────────────────────────────────────────
_SHARED_TOKENS: list[tuple[int, str]] = [
(100, "Let me"),
(101, " think"),
(102, " about"),
(103, " Samsung."),
(51, "</think>"),
(200, "삼성"),
(201, "전자의"),
(202, " 주가를"),
(203, " 분석합니다."),
]
_SHARED_UFFFD_IDS: set[int] = {200}
EXPECTED_REASONING = "Let me think about Samsung."
EXPECTED_CONTENT = "삼성전자의 주가를 분석합니다."
_MODEL_CONFIGS = [
pytest.param(
{
"<think>": 50,
"</think>": 51,
"<tool_call>": 60,
"</tool_call>": 61,
"<arg_key>": 62,
"</arg_key>": 63,
"<arg_value>": 64,
"</arg_value>": 65,
},
_Glm47Delegating,
id="glm47",
),
pytest.param(
{
"<think>": 50,
"</think>": 51,
"<tool_call>": 60,
"</tool_call>": 61,
},
_Qwen3Delegating,
id="qwen3",
),
]
# ── Tests ────────────────────────────────────────────────────────────
class TestUfffdReasoningTransition:
"""U+FFFD must not appear at the reasoning→content transition."""
@pytest.mark.parametrize("vocab,delegating_cls", _MODEL_CONFIGS)
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
def test_no_ufffd(self, chunk_size, vocab, delegating_cls):
tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS)
parser = delegating_cls(tokenizer)
deltas = replay_streaming(
parser,
_SHARED_TOKENS,
chunk_size=chunk_size,
finished_on_last=True,
)
output = collect_output(deltas)
assert "" not in output.content, (
f"U+FFFD leaked into content: {output.content!r}"
)
assert output.content == EXPECTED_CONTENT
assert output.reasoning == EXPECTED_REASONING
def test_byte_fallback_tokenizer_produces_ufffd(self):
"""Validate the fixture: decode() returns U+FFFD for isolated
byte-fallback token IDs, proving the old code path would leak."""
vocab = dict(_MODEL_CONFIGS[0].values[0])
tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS)
assert tokenizer.decode([200]) == ""
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
def test_multiple_ufffd_tokens_at_boundary(self, chunk_size):
"""Multiple consecutive byte-fallback tokens at the boundary."""
tokens: list[tuple[int, str]] = [
(100, "Reasoning."),
(51, "</think>"),
(200, ""),
(201, ""),
(202, "전자"),
]
ufffd_ids: set[int] = {200, 201}
vocab = dict(_MODEL_CONFIGS[0].values[0])
tokenizer = ByteFallbackMockTokenizer(vocab, tokens, ufffd_ids)
parser = _Glm47Delegating(tokenizer)
deltas = replay_streaming(
parser,
tokens,
chunk_size=chunk_size,
finished_on_last=True,
)
output = collect_output(deltas)
assert "" not in output.content, (
f"U+FFFD leaked into content: {output.content!r}"
)
assert output.content == "삼성전자"
assert output.reasoning == "Reasoning."
+17 -2
View File
@@ -143,6 +143,12 @@ SCENARIOS: list[Scenario] = [
tool_calls=[_READ_TOOL],
after_tool_response=True,
),
Scenario(
id="empty-tool-block",
description="Empty tool block followed by content (edge case recovery)",
content="Content after empty tools.",
tool_calls=[],
),
]
@@ -344,8 +350,11 @@ def _qwen3_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
if scenario.content is not None or scenario.tool_calls is not None:
segs.append(("</think>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<tool_call>", True))
segs.append(("</tool_call>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -437,8 +446,11 @@ def _minimax_m2_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
if scenario.content is not None or scenario.tool_calls is not None:
segs.append(("</think>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<minimax:tool_call>", True))
segs.append(("</minimax:tool_call>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -534,6 +546,9 @@ def _gemma4_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs.append((_GEMMA4_THOUGHT_PREFIX, False))
segs.append((scenario.reasoning, False))
segs.append(("<channel|>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<|tool_call>", True))
segs.append(("<tool_call|>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -83,6 +83,20 @@ class MiniMaxM3Tokenizer:
return "".join(tokens)
class SplitMiniMaxM3Tokenizer(MiniMaxM3Tokenizer):
"""Tokenizer that exposes marker vocab entries but encodes them as text."""
def tokenize(self, text: str) -> list[str]:
return list(text)
class RuntimeSplitMiniMaxM3Tokenizer(MiniMaxM3Tokenizer):
"""Tokenizer whose runtime output splits markers despite atomic encodes."""
def encode_runtime(self, text: str) -> list[int]:
return [self._add_token(token) for token in list(text)]
def make_parser(
chat_template_kwargs: dict[str, str] | None = None,
) -> tuple[MiniMaxM3ReasoningParser, MiniMaxM3Tokenizer]:
@@ -105,7 +119,8 @@ def run_streaming(
reasoning_end_states: list[bool] = []
for chunk in chunks:
delta_token_ids = tokenizer.encode(chunk, add_special_tokens=False)
encode_runtime = getattr(tokenizer, "encode_runtime", tokenizer.encode)
delta_token_ids = encode_runtime(chunk)
current_text = previous_text + chunk
current_token_ids = previous_token_ids + delta_token_ids
delta = parser.extract_reasoning_streaming(
@@ -174,14 +189,14 @@ def test_nonstreaming_drops_leading_end_tag():
assert content == "answer"
def test_nonstreaming_non_leading_end_tag_is_content():
def test_nonstreaming_end_tag_in_content_state_is_dropped():
parser, _ = make_parser()
request = ChatCompletionRequest(messages=[], model="test-model")
reasoning, content = parser.extract_reasoning("XXX</mm:think>YYY", request)
assert reasoning is None
assert content == "XXX</mm:think>YYY"
assert content == "XXXYYY"
def test_nonstreaming_enabled_mode_starts_in_reasoning():
@@ -246,7 +261,7 @@ def test_streaming_drops_leading_end_tag():
assert end_states == [True, True]
def test_streaming_non_leading_end_tag_is_content():
def test_streaming_end_tag_in_content_state_is_dropped():
parser, tokenizer = make_parser()
reasoning, content, end_states = run_streaming(
@@ -256,7 +271,7 @@ def test_streaming_non_leading_end_tag_is_content():
)
assert reasoning is None
assert content == "XXX</mm:think>YYY"
assert content == "XXXYYY"
assert end_states == [True]
@@ -288,6 +303,110 @@ def test_streaming_plain_content_ends_reasoning_phase():
assert end_states == [True, True]
def test_streaming_split_marker_tokens_are_not_returned():
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
parser = MiniMaxM3ReasoningParser(tokenizer)
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["<mm:think>", "Reasoning", " content", "</mm:think>", "content"],
)
assert reasoning == "Reasoning content"
assert content == "content"
assert end_states == [False, False, False, True, True]
def test_streaming_split_marker_text_drives_end_state():
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
parser = MiniMaxM3ReasoningParser(tokenizer)
previous_text = ""
previous_token_ids: list[int] = []
for chunk in ["<mm:think>", "Reasoning", " content", "</mm:think>"]:
delta_token_ids = tokenizer.encode_runtime(chunk)
current_text = previous_text + chunk
current_token_ids = previous_token_ids + delta_token_ids
parser.extract_reasoning_streaming(
previous_text=previous_text,
current_text=current_text,
delta_text=chunk,
previous_token_ids=previous_token_ids,
current_token_ids=current_token_ids,
delta_token_ids=delta_token_ids,
)
previous_text = current_text
previous_token_ids = current_token_ids
assert parser.is_reasoning_end_streaming(previous_token_ids, []) is True
def test_streaming_split_marker_tokens_enabled_mode():
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
parser = MiniMaxM3ReasoningParser(
tokenizer, chat_template_kwargs={"thinking_mode": "enabled"}
)
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["Reasoning", " content", "</mm:think>", "content"],
)
assert reasoning == "Reasoning content"
assert content == "content"
assert end_states == [False, False, True, True]
def test_streaming_split_marker_text_across_deltas():
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
parser = MiniMaxM3ReasoningParser(tokenizer)
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["<mm:", "think>", "Reasoning", " content", "</mm:", "think>", "content"],
)
assert reasoning == "Reasoning content"
assert content == "content"
assert end_states == [False, False, False, False, False, True, True]
def test_streaming_split_leading_end_marker_text_across_deltas():
tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
parser = MiniMaxM3ReasoningParser(tokenizer)
reasoning, content, end_states = run_streaming(
parser,
tokenizer,
["</mm:", "think>", "content"],
)
assert reasoning is None
assert content == "content"
assert end_states == [False, True, True]
def test_token_id_helpers_with_split_marker_tokens():
tokenizer = SplitMiniMaxM3Tokenizer()
parser = MiniMaxM3ReasoningParser(tokenizer)
output_ids = tokenizer.encode(
"<mm:think>abc</mm:think>def", add_special_tokens=False
)
open_reasoning_ids = tokenizer.encode("<mm:think>abc", add_special_tokens=False)
content_ids = tokenizer.encode("plain", add_special_tokens=False)
assert parser.is_reasoning_end(output_ids)
assert not parser.is_reasoning_end(open_reasoning_ids)
assert not parser.is_reasoning_end(content_ids)
assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
assert parser.extract_content_ids(open_reasoning_ids) == []
assert parser.extract_content_ids(content_ids) == content_ids
assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
def test_token_id_helpers():
parser, tokenizer = make_parser()
output_ids = tokenizer.encode(
@@ -18,6 +18,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
FP8ScaledMMLinearLayerConfig,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
get_fp8_min_max,
kFp8DynamicTokenSym,
kFp8StaticChannelSym,
kFp8StaticTensorSym,
@@ -309,7 +310,7 @@ def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel):
_check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens)
fp8_dtype = current_platform.fp8_dtype()
fp8_max = torch.finfo(fp8_dtype).max
fp8_max = get_fp8_min_max()[1]
device = torch.device("cuda")
# Build a bf16 weight and quantize per output channel (one scale per row).
+228
View File
@@ -0,0 +1,228 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for contiguous KV cache packing."""
from unittest.mock import MagicMock
import pytest
import torch
from vllm import envs
from vllm.v1.core.kv_cache_utils import (
_get_kv_cache_config_deepseek_v4,
get_kv_cache_config_from_groups,
)
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheGroupSpec,
KVCacheTensor,
MLAAttentionSpec,
SlidingWindowSpec,
UniformTypeKVCacheSpecs,
)
def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
return MLAAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=512,
dtype=torch.uint8,
page_size_padded=page_size,
cache_dtype_str="fp8_ds_mla",
model_version="deepseek_v4",
alignment=576,
)
def _make_full_spec() -> FullAttentionSpec:
return FullAttentionSpec(
block_size=16,
num_kv_heads=2,
head_size=64,
dtype=torch.float16,
)
def _make_sw_spec() -> SlidingWindowSpec:
return SlidingWindowSpec(
block_size=16,
num_kv_heads=2,
head_size=64,
dtype=torch.float16,
sliding_window=128,
)
def _make_groups(n_c4, n_c128, n_swa):
PS_C4_MLA = 37440
PS_C4_IDX = 8640
PS_C128 = 1728
PS_SWA = 37440
mla_specs = {}
for i in range(n_c4):
mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA)
mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX)
for i in range(n_c128):
mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128)
mla_group = KVCacheGroupSpec(
layer_names=list(mla_specs.keys()),
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs),
)
swa_specs = {}
for i in range(n_swa):
swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA)
swa_group = KVCacheGroupSpec(
layer_names=list(swa_specs.keys()),
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs),
)
return [mla_group, swa_group]
def _mock_vllm_config():
config = MagicMock()
config.cache_config.num_gpu_blocks_override = None
return config
def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
groups = _make_groups(n_c4, n_c128, n_swa)
return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem)
def _page_sizes_by_layer(
groups: list[KVCacheGroupSpec],
) -> dict[str, int]:
page_sizes = {}
for group in groups:
specs = group.kv_cache_spec.kv_cache_specs
for layer_name in group.layer_names:
page_sizes[layer_name] = specs[layer_name].page_size_bytes
return page_sizes
class TestInterleavedPacking:
def test_all_tensors_have_block_stride(self):
_, tensors = _run()
for t in tensors:
assert t.block_stride > 0
def test_all_tensors_share_same_size(self):
_, tensors = _run()
sizes = set(t.size for t in tensors)
assert len(sizes) == 1
assert sizes.pop() > 0
def test_offsets_within_one_block(self):
_, tensors = _run()
for t in tensors:
assert t.offset < t.block_stride
def test_all_layers_accounted_for(self):
n_c4, n_c128, n_swa = 5, 4, 7
_, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa)
all_names = set()
for t in tensors:
all_names.update(t.shared_by)
expected = n_c4 * 2 + n_c128 + n_swa
assert len(all_names) == expected
def test_strided_views_are_independent(self):
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
page_sizes = _page_sizes_by_layer(groups)
num_blocks, tensors = _get_kv_cache_config_deepseek_v4(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
views = []
for t in tensors:
page_size = page_sizes[t.shared_by[0]]
v = torch.as_strided(
backing,
size=(num_blocks, page_size),
stride=(t.block_stride, 1),
storage_offset=t.offset,
)
views.append(v)
for i, v in enumerate(views):
v.fill_(i + 1)
for i, v in enumerate(views):
assert (v == i + 1).all(), f"View {i} was corrupted"
def test_hma_attention_groups_keep_default_backing(self, monkeypatch):
monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False)
full = _make_full_spec()
sw = _make_sw_spec()
page_size = full.page_size_bytes
groups = [
KVCacheGroupSpec(["full.0", "full.1"], full),
KVCacheGroupSpec(["sw.0", "sw.2"], sw),
KVCacheGroupSpec(["sw.1", "sw.3"], sw),
]
config = get_kv_cache_config_from_groups(
_mock_vllm_config(), groups, available_memory=page_size * 2 * 32
)
assert config.num_blocks == 32
assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32
assert config.kv_cache_tensors == [
KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),
KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
]
def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch):
monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False)
full = _make_full_spec()
sw = _make_sw_spec()
page_size = full.page_size_bytes
groups = [
KVCacheGroupSpec(["full.0", "full.1"], full),
KVCacheGroupSpec(["sw.0", "sw.2"], sw),
KVCacheGroupSpec(["sw.1", "sw.3"], sw),
]
config = get_kv_cache_config_from_groups(
_mock_vllm_config(), groups, available_memory=page_size * 2 * 32
)
assert config.num_blocks == 32
assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}
assert config.kv_cache_tensors == [
KVCacheTensor(
size=page_size * 2 * 32,
shared_by=["full.0", "sw.0", "sw.1"],
offset=0,
block_stride=page_size * 2,
),
KVCacheTensor(
size=page_size * 2 * 32,
shared_by=["full.1", "sw.2", "sw.3"],
offset=page_size,
block_stride=page_size * 2,
),
]
def test_single_group_attention_keeps_unpacked_layout(self):
spec = _make_full_spec()
groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
config = get_kv_cache_config_from_groups(
_mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
)
assert sum(t.size for t in config.kv_cache_tensors) == (
spec.page_size_bytes * 2 * 32
)
assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]
if __name__ == "__main__":
pytest.main([__file__, "-v"])
+37
View File
@@ -144,6 +144,43 @@ def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders():
assert req.request_id in output.num_scheduled_tokens
def test_cached_request_data_resumed_all_token_ids_mrv1_only():
"""all_token_ids carries a resumed request's token ids to the connector
for the V1 model runner, but is skipped entirely for the V2 model runner.
"""
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
scheduler = create_scheduler()
(req,) = create_requests(num_requests=1, num_tokens=8)
req.append_output_token_ids([101, 102, 103])
# A resumed request was not scheduled in the previous step.
assert req.request_id not in scheduler.prev_step_scheduled_req_ids
empty_blocks = KVCacheBlocks(blocks=((),))
def make_cached():
return scheduler._make_cached_request_data(
running_reqs=[],
resumed_reqs=[req],
num_scheduled_tokens={req.request_id: 1},
spec_decode_tokens={},
req_to_new_blocks={req.request_id: empty_blocks},
)
# V1 model runner: the full token id list is propagated.
assert not scheduler.use_v2_model_runner
cached = make_cached()
assert req.request_id in cached.resumed_req_ids
assert cached.all_token_ids[req.request_id] == list(req.all_token_ids)
# V2 model runner: all_token_ids is skipped entirely.
scheduler.use_v2_model_runner = True
cached = make_cached()
assert req.request_id in cached.resumed_req_ids
assert cached.all_token_ids == {}
def test_schedule_partial_requests():
"""Test scheduling behavior with partial requests.
@@ -4,9 +4,16 @@
import pytest
from vllm import LLM, SamplingParams
from vllm.platforms import current_platform
from ....utils import create_new_process_for_each_test
if current_platform.is_rocm():
pytest.skip(
"Cascade attention backends FLASH_ATTN and FLASHINFER are notsupported on ROCm",
allow_module_level=True,
)
@create_new_process_for_each_test()
@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"])
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
import time
from unittest.mock import MagicMock, patch
from vllm.config import set_current_vllm_config
@@ -406,7 +408,9 @@ def test_lookup_key_client_lookup_prepends_typed_tag():
fake_socket = mock_make_socket.return_value
fake_socket.recv.return_value = (5).to_bytes(4, "big")
assert client.lookup(token_len=128, block_hashes=[]) == 5
# Blocking lookup (non_block defaults to False) runs on the executor and
# returns the resolved hit length.
assert client.lookup("req0", token_len=128, block_hashes=[]) == 5
sent_frames = fake_socket.send_multipart.call_args[0][0]
assert sent_frames[0] == protocol.LOOKUP_MSG
@@ -435,6 +439,127 @@ def test_lookup_key_client_reset_uses_typed_protocol():
assert client.reset() is False
def _poll_lookup(client, req_id, token_len=128, block_hashes=(), timeout=5.0):
"""Drive non-blocking lookup until the executor completes it."""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
result = client.lookup(req_id, token_len, list(block_hashes), non_block=True)
if result is not None:
return result
time.sleep(0.005)
return None
def _gated_recv(gate: threading.Event, value: int):
"""Mock recv side-effect that blocks until ``gate`` is set, so the
executor's lookup can be held pending deterministically."""
def recv():
gate.wait()
return value.to_bytes(4, "big")
return recv
def test_lookup_key_client_non_block_lookup_async():
"""Non-blocking lookup defers to the executor: None first, hit once the
Future resolves."""
vllm_config = _make_vllm_config()
with patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.make_zmq_socket"
) as mock_make_socket:
client = worker.LookupKeyClient(vllm_config)
fake_socket = mock_make_socket.return_value
# Hold the executor's lookup pending until we release the gate.
gate = threading.Event()
fake_socket.recv.side_effect = _gated_recv(gate, 7)
# First query submits the lookup and returns None while it is in flight.
assert client.lookup("req1", 128, [], non_block=True) is None
# Release the executor; a later poll returns the hit length.
gate.set()
assert _poll_lookup(client, "req1") == 7
# Future is consumed (popped) on read.
assert "req1" not in client.futures
def test_lookup_key_client_discard_clears_state():
"""discard() drops a completed lookup Future so it is not served stale."""
vllm_config = _make_vllm_config()
with patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.make_zmq_socket"
) as mock_make_socket:
client = worker.LookupKeyClient(vllm_config)
fake_socket = mock_make_socket.return_value
gate = threading.Event()
fake_socket.recv.side_effect = _gated_recv(gate, 9)
# Submit while gated so the call returns None and the Future stays in
# `futures` (unconsumed) once it resolves.
assert client.lookup("req2", 128, [], non_block=True) is None
gate.set()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
if client.futures["req2"].done():
break
time.sleep(0.005)
# discard() drops the completed result before any lookup consumes it.
client.discard("req2")
assert "req2" not in client.futures
# A fresh query re-submits rather than returning a stale value: hold the
# gate so the resubmitted lookup stays in flight.
gate.clear()
assert client.lookup("req2", 128, [], non_block=True) is None
gate.set() # release the executor so the worker thread can drain
def test_get_num_new_matched_tokens_async_defers_then_reports():
"""Async lookup returns (None, False) until ready, then the hit count."""
vllm_config = create_vllm_config(
kv_connector="MooncakeStoreConnector",
kv_role="kv_both",
kv_connector_extra_config={"lookup_async": True},
)
kv_cache_config = _make_kv_cache_config()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"scheduler.LookupKeyClient"
) as mock_client_cls,
):
sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)
assert sched.lookup_async is True
mock_client = mock_client_cls.return_value
block_size = sched._block_size
request = MagicMock()
request.request_id = "r1"
request.num_tokens = 4 * block_size
request.block_hashes = []
# Lookup not ready -> defer.
mock_client.lookup.return_value = None
assert sched.get_num_new_matched_tokens(request, 0) == (None, False)
assert "r1" not in sched.load_specs
# Lookup ready with a hit -> report need_to_allocate + async-load flag.
hit = 3 * block_size
mock_client.lookup.return_value = hit
need, load_async = sched.get_num_new_matched_tokens(request, 0)
assert need == hit
assert load_async == sched.load_async
assert sched.load_specs["r1"].kvpool_cached_tokens == hit
def test_protocol_tags_are_distinct_and_non_empty():
"""Protocol tags must be unique and non-empty to avoid collision."""
tags = {protocol.LOOKUP_MSG, protocol.RESET_MSG}
@@ -7,7 +7,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp
ExternalCachedBlockPool,
MooncakeStoreCoordinator,
)
from vllm.v1.core.kv_cache_utils import BlockHash, BlockHashListWithBlockSize
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
chunk_hashes_for_block_size,
)
from vllm.v1.core.kv_cache_utils import BlockHash
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheGroupSpec,
@@ -182,7 +185,7 @@ def test_coordinator_group_block_size_double_hash():
]
coord = _make_coord(groups, hash_block_size=16)
hs = _hashes(4)
big_hashes = list(BlockHashListWithBlockSize(hs, 16, 32))
big_hashes = list(chunk_hashes_for_block_size(hs, 16, 32))
exists = {(0, bytes(h)) for h in hs}
exists |= {(1, bytes(bh)) for bh in big_hashes}
cmap = ExternalCachedBlockPool(exists)
@@ -323,8 +323,8 @@ def test_recv_skips_swa_blocks_before_window():
def test_chunked_token_database_hash_block_size_smaller_than_block_size():
"""DSv4-style: hash_block_size=4, group block_size=16 — process_tokens
must merge every 4 fine hashes into one chunk hash via
BlockHashListWithBlockSize."""
keys each 16-token chunk by its last fine hash, keeping the Mooncake key
at one digest instead of concatenating all 4 fine hashes."""
md = KeyMetadata("m", 0, 0, 0, 0, group_id=3)
db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4)
db.set_kv_caches_base_addr([0])
@@ -335,8 +335,7 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size():
assert len(out) == 2
assert out[0][0] == 0 and out[0][1] == 16
assert out[1][0] == 16 and out[1][1] == 32
# Each chunk's hash is the concatenation of 4 fine hashes.
expected0 = b"".join(fine_hashes[0:4]).hex()
expected1 = b"".join(fine_hashes[4:8]).hex()
assert out[0][2].chunk_hash == expected0
assert out[1][2].chunk_hash == expected1
# Each chunk's hash is its last (4th) fine hash, which already chains the
# prior three.
assert out[0][2].chunk_hash == fine_hashes[3].hex()
assert out[1][2].chunk_hash == fine_hashes[7].hex()
@@ -16,6 +16,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor
def _make_bare_scheduler() -> MooncakeStoreScheduler:
scheduler = object.__new__(MooncakeStoreScheduler)
scheduler.kv_role = "kv_both"
scheduler.lookup_async = False
scheduler._block_size = 16
scheduler.load_specs = {}
scheduler._preempted_req_ids = set()
@@ -405,7 +406,13 @@ class _StubLookupClient:
def __init__(self, hit_tokens: int) -> None:
self._hit_tokens = hit_tokens
def lookup(self, token_len: int, block_hashes: list[bytes]) -> int:
def lookup(
self,
req_id: str,
token_len: int,
block_hashes: list[bytes],
non_block: bool = False,
) -> int:
return self._hit_tokens
@@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
worker as mooncake_store_worker,
)
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
BlobBlockHashes,
ChunkedTokenDatabase,
KeyMetadata,
LoadSpec,
@@ -32,6 +33,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
MooncakeStoreConnectorStats,
)
from vllm.v1.core.kv_cache_utils import BlockHash
def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:
@@ -175,14 +177,17 @@ class _FakeModelConfig:
def _make_vllm_config(
*, extra_config: dict[str, object] | None = None
*,
extra_config: dict[str, object] | None = None,
rank: int = 0,
decode_context_parallel_size: int = 1,
) -> SimpleNamespace:
return SimpleNamespace(
model_config=_FakeModelConfig(),
parallel_config=SimpleNamespace(
pipeline_parallel_size=1,
rank=0,
decode_context_parallel_size=1,
rank=rank,
decode_context_parallel_size=decode_context_parallel_size,
prefill_context_parallel_size=1,
),
kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
@@ -231,13 +236,23 @@ def _install_fake_mooncake(monkeypatch, store_instance: MagicMock):
return FakeReplicateConfig
def _patch_worker_runtime(monkeypatch, *, local_ip: str = "10.0.0.7") -> None:
def _patch_worker_runtime(
monkeypatch,
*,
local_ip: str = "10.0.0.7",
tp_rank: int = 0,
tp_size: int = 1,
dcp_size: int = 1,
) -> None:
single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0)
# DCP groups are contiguous splits of the TP group (see
# parallel_state.py), so dcp_rank == tp_rank % dcp_size.
dcp_group = SimpleNamespace(world_size=dcp_size, rank_in_group=tp_rank % dcp_size)
monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0)
monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: 1)
monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: tp_rank)
monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: tp_size)
monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
monkeypatch.setattr(worker, "get_dcp_group", lambda: single_rank_group)
monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group)
monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
@@ -884,6 +899,66 @@ def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
@pytest.mark.parametrize("dcp_size", [1, 4])
def test_worker_put_striding_covers_every_rank_get_namespace(
tmp_path, monkeypatch, dcp_size
):
"""Every key a rank GETs must have been PUT by some rank.
When num_kv_head < tp_size, ranks holding the same KV heads stripe
their PUTs across one shared key namespace. That dedup is only valid
when those ranks really share a namespace: with DCP > 1 each rank GETs
every key from its own ``@dcpN`` namespace, so striding must be
disabled.
"""
tp_size = 4
store = MagicMock()
store.setup.return_value = 0
_install_fake_mooncake(monkeypatch, store)
monkeypatch.setenv(
"MOONCAKE_CONFIG_PATH",
_write_mooncake_config(
tmp_path,
{
"metadata_server": "http://metadata/endpoint",
"protocol": "tcp",
"device_name": "",
"master_server_address": "10.0.0.7:50051",
},
),
)
# _FakeModelConfig has num_kv_head=1 < tp_size, which enables striding.
block_hashes = [f"hash-{i}".encode() for i in range(4)]
put_keys: set[str] = set()
get_keys_per_rank: dict[int, set[str]] = {}
for tp_rank in range(tp_size):
_patch_worker_runtime(
monkeypatch, tp_rank=tp_rank, tp_size=tp_size, dcp_size=dcp_size
)
w = worker.MooncakeStoreWorker(
_make_vllm_config(rank=tp_rank, decode_context_parallel_size=dcp_size),
_make_kv_cache_config(),
)
db = w.token_dbs[0]
token_len = len(block_hashes) * db.block_size
keys = [
key.to_string() for _, _, key in db.process_tokens(token_len, block_hashes)
]
assert len(keys) == len(block_hashes)
# PUT side: mirrors KVCacheStoreSendingThread's striding slice.
put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step])
# GET side: KVCacheStoreRecvingThread fetches every key.
get_keys_per_rank[tp_rank] = set(keys)
for tp_rank, rank_keys in get_keys_per_rank.items():
missing = rank_keys - put_keys
assert not missing, (
f"tp_rank={tp_rank} would GET {len(missing)}/{len(rank_keys)} keys "
f"that no rank PUT (Mooncake OBJECT_NOT_FOUND): {sorted(missing)}"
)
# ---------------------------------------------------------------------------
# Helpers for register_kv_caches tests
# ---------------------------------------------------------------------------
@@ -1106,9 +1181,9 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():
assert full_event.group_idx == 0
assert full_event.block_size == 32
assert full_event.token_ids == list(range(32))
assert full_event.block_hashes == [
maybe_convert_block_hash(BlockHash(b"".join(hs)))
]
# block_size=32 over hash_block_size=8 (scale 4): the chunk is keyed by its
# last sub-hash, not the concatenation of all four.
assert full_event.block_hashes == [maybe_convert_block_hash(BlockHash(hs[3]))]
assert swa_event.group_idx == 1
assert swa_event.block_size == 8
@@ -1676,3 +1751,33 @@ def test_store_worker_close_swallows_store_errors():
worker.close()
assert worker.store is None
def test_blob_block_hashes_wire_roundtrip():
"""The lookup wire format sends a ``hash_len`` frame plus the raw hashes
concatenated back-to-back; the server rebuilds them through a zero-copy
``BlobBlockHashes`` view over the frame buffer."""
hashes = [BlockHash(bytes([i]) * 16) for i in range(5)]
hash_len = len(hashes[0])
# Client side (LookupKeyClient._lookup): flat payload frame.
blob = b"".join(hashes)
# Server side (LookupKeyServer): view over the frame buffer (a memoryview),
# never materializing the full hash list upfront.
view = BlobBlockHashes(memoryview(blob), hash_len)
assert len(view) == 5
assert list(view) == hashes # default Sequence iter terminates via IndexError
assert [bytes(h) for h in view] == hashes
assert bytes(view[-1]) == hashes[-1]
assert [bytes(h) for h in view[1:3]] == hashes[1:3]
with pytest.raises(IndexError):
_ = view[5]
def test_blob_block_hashes_empty():
"""Empty lookups send hash_len=0 and an empty payload."""
view = BlobBlockHashes(memoryview(b""), 0)
assert len(view) == 0
assert list(view) == []
+41 -5
View File
@@ -14,12 +14,13 @@ from vllm.v1.kv_offload.base import (
ReqContext,
make_offload_key,
)
from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
from vllm.v1.kv_offload.cpu.common import (
CPULoadStoreSpec,
CPUOffloadingMetrics,
)
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
def make_req_context(
req_id: str = "", kv_transfer_params: dict | None = None
@@ -181,10 +182,45 @@ def test_filter_reused_manager_reports_stores_skipped_counter():
)
stats = manager.get_stats()
assert stats is not None
assert stats.reduce()[STORES_SKIPPED] == 3
assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 3
stats = manager.get_stats()
assert stats is not None
assert stats.reduce()[STORES_SKIPPED] == 0
assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 0
def test_cpu_manager_reports_cache_usage_gauge():
def check_usage_stats(manager: CPUOffloadingManager, value: float):
stats = manager.get_stats()
assert stats is not None
assert stats.reduce()[
CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC
] == pytest.approx(value)
# Zero-capacity manager always reports 0.0
manager = make_cpu_manager(num_blocks=0)
check_usage_stats(manager, 0.0)
# Empty manager (4 blocks, none allocated): usage = 0.0
manager = make_cpu_manager(num_blocks=4)
check_usage_stats(manager, 0.0)
# After allocating 2 of 4 blocks: usage = 0.5
manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
check_usage_stats(manager, 0.5)
# After filling all 4 blocks: usage = 1.0
manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
check_usage_stats(manager, 1.0)
# After completing store, the blocks becomes evictable as it is not actively used
# and usage drops.
manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
check_usage_stats(manager, 0.5)
# After completing store, the blocks becomes evictable as it is not actively used
# and usage drops.
manager.complete_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
check_usage_stats(manager, 0.0)
def test_cpu_manager():
@@ -145,7 +145,6 @@ def _generate_fake_sampling_metadata(
vllm_config.scheduler_config.max_num_seqs,
num_spec,
device,
PIN_MEMORY_AVAILABLE,
)
fake_sampling_metadata = SamplingMetadata(
temperature=torch.full((batch_size,), 0.0),
@@ -880,7 +879,6 @@ def test_maybe_create_thinking_budget_holder_without_reasoning():
cfg.scheduler_config.max_num_seqs,
0,
torch.device("cpu"),
False,
)
is None
)
@@ -6,6 +6,7 @@ from __future__ import annotations
from dataclasses import dataclass
import pytest
import torch
from vllm import SamplingParams
@@ -1528,3 +1529,232 @@ def test_reset_pending_loads() -> None:
# All GPU blocks free
num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()
assert num_used == 1, f"Expected only null block in use, got {num_used}"
def _make_cp_vllm_config(
dcp_world_size: int = 1,
pcp_world_size: int = 1,
) -> VllmConfig:
"""VllmConfig with context-parallel sizes set for scheduler-only tests."""
cfg = _make_vllm_config()
cfg.parallel_config.decode_context_parallel_size = dcp_world_size
cfg.parallel_config.prefill_context_parallel_size = pcp_world_size
return cfg
def _make_cp_scheduler(
*,
dcp_world_size: int = 1,
pcp_world_size: int = 1,
num_cpu_blocks: int = 8,
num_gpu_blocks: int = 16,
lazy: bool = False,
) -> SchedulerFixture:
"""Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size."""
cp_world_size = dcp_world_size * pcp_world_size
virtual_block_size = BLOCK_SIZE * cp_world_size
kv_cache_config = _make_kv_cache_config(num_gpu_blocks)
vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size)
cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks
sched = SimpleCPUOffloadScheduler(
vllm_config=vllm_config,
kv_cache_config=kv_cache_config,
cpu_capacity_bytes=cpu_capacity_bytes,
scheduler_block_size=virtual_block_size,
hash_block_size=virtual_block_size,
lazy_offload=lazy,
)
gpu_block_pool = BlockPool(
num_gpu_blocks=num_gpu_blocks,
enable_caching=True,
hash_block_size=virtual_block_size,
)
sched.bind_gpu_block_pool(gpu_block_pool)
return SchedulerFixture(
scheduler=sched,
gpu_block_pool=gpu_block_pool,
vllm_config=vllm_config,
kv_cache_config=kv_cache_config,
)
def _make_cp_request(
num_blocks: int,
virtual_block_size: int,
request_id: str | None = None,
) -> Request:
"""Create a request whose block hashes are computed at the virtual
(CP-scaled) block size, matching what the real scheduler does.
"""
global _req_counter
_req_counter += 1
if request_id is None:
request_id = f"req-cp-{_req_counter}"
num_tokens = num_blocks * virtual_block_size + 1
start = _req_counter * 10000
prompt_token_ids = list(range(start, start + num_tokens))
sampling_params = SamplingParams(max_tokens=1)
return Request(
request_id=request_id,
prompt_token_ids=prompt_token_ids,
sampling_params=sampling_params,
pooling_params=None,
mm_features=None,
block_hasher=get_request_block_hasher(virtual_block_size, sha256),
)
def _allocate_cp_gpu_blocks(
gpu_block_pool: BlockPool,
request: Request,
num_blocks: int,
virtual_block_size: int,
group_id: int = 0,
) -> list:
"""Allocate GPU blocks and cache them using the CP-scaled block size."""
blocks = gpu_block_pool.get_new_blocks(num_blocks)
num_full = min(num_blocks, len(request.block_hashes))
if num_full > 0:
gpu_block_pool.cache_full_blocks(
request=request,
blocks=blocks,
num_cached_blocks=0,
num_full_blocks=num_full,
block_size=virtual_block_size,
kv_cache_group_id=group_id,
)
return blocks
# ---------------------------------------------------------------------------
# Test 15: CP block size scaling is correct
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"dcp_world_size, pcp_world_size",
[
(2, 1), # DCP only
(1, 2), # PCP only
(2, 2), # DCP + PCP
],
)
def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None:
"""Verify that the scheduler's block_size and cp_world_size are correctly
scaled when context parallelism is enabled."""
fix = _make_cp_scheduler(
dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
)
sched = fix.scheduler
expected_cp = dcp_world_size * pcp_world_size
assert sched.cp_world_size == expected_cp
assert sched.block_size == BLOCK_SIZE * expected_cp
# ---------------------------------------------------------------------------
# Test 16: CP eager store-and-load roundtrip
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"dcp_world_size, pcp_world_size",
[
(2, 1),
(1, 2),
],
)
def test_cp_eager_store_and_load_roundtrip(
dcp_world_size: int, pcp_world_size: int
) -> None:
"""With CP enabled, store blocks to CPU and reload them for a new request
with matching tokens. Verifies that hash matching and transfer-pair
construction work with the virtual block size."""
fix = _make_cp_scheduler(
dcp_world_size=dcp_world_size,
pcp_world_size=pcp_world_size,
num_cpu_blocks=8,
num_gpu_blocks=16,
lazy=False,
)
sched = fix.scheduler
cp = dcp_world_size * pcp_world_size
vbs = BLOCK_SIZE * cp
num_blocks = 2
req = _make_cp_request(num_blocks, vbs)
# Allocate GPU blocks and register hashes
gpu_blocks = _allocate_cp_gpu_blocks(fix.gpu_block_pool, req, num_blocks, vbs)
kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,))
req.num_computed_tokens = num_blocks * vbs
sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
block_ids = kv_blocks.get_block_ids()
sched_out = make_scheduler_output(
{req.request_id: num_blocks * vbs},
new_reqs={req.request_id: block_ids},
)
meta = sched.build_connector_meta(sched_out)
assert meta.store_event >= 0, "Expected a store event"
assert len(meta.store_gpu_blocks) == num_blocks
assert len(meta.store_cpu_blocks) == num_blocks
simulate_store_completion(sched, meta.store_event)
# New request with same tokens — should get a full CPU cache hit.
req2 = Request(
request_id="req-cp-load",
prompt_token_ids=req.prompt_token_ids,
sampling_params=req.sampling_params,
pooling_params=None,
mm_features=None,
block_hasher=req._block_hasher,
)
hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
assert hit_tokens == num_blocks * vbs
assert is_async is True
# Allocate fresh GPU blocks for the load.
gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
sched_out2 = make_scheduler_output(
{req2.request_id: 1},
new_reqs={req2.request_id: kv_blocks2.get_block_ids()},
)
meta2 = sched.build_connector_meta(sched_out2)
assert meta2.load_event >= 0, "Expected a load event"
assert len(meta2.load_gpu_blocks) == num_blocks
assert len(meta2.load_cpu_blocks) == num_blocks
# ---------------------------------------------------------------------------
# Test 17: CP lazy target blocks are scaled correctly
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("cp_world_size", [1, 2, 4])
def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None:
"""_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1
because each virtual block covers more tokens."""
kv_cache_config = _make_kv_cache_config(num_blocks=16)
max_batched = 64
target_base = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
kv_cache_config, max_batched, cp_world_size=1
)
target_cp = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
kv_cache_config, max_batched, cp_world_size=cp_world_size
)
if cp_world_size == 1:
assert target_cp == target_base
else:
assert target_cp < target_base, (
f"cp_world_size={cp_world_size}: target_cp={target_cp} should be "
f"less than target_base={target_base}"
)
@@ -35,7 +35,6 @@ def mock_model_runner_with_input_batch():
max_model_len=1024,
max_num_batched_tokens=1024,
device="cpu",
pin_memory=False,
vocab_size=32000,
block_sizes=[16],
kernel_block_sizes=[16],
+227
View File
@@ -0,0 +1,227 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the Model Runner V2 Gumbel-max sampling kernel.
Accuracy: define a target categorical distribution as a non-negative int64
count tensor summing to N, turn it into logits (= log(count)), sample many
times with `gumbel_sample`, and check the empirical distribution matches.
The count tensor is deliberately heavy-tailed (one dominant token, the rest
~18 logits below). That tail is the sensitive part: the fp32 Gumbel noise must
reach ~18 to ever sample it. A flat distribution would keep every token within
a few logits of the top and would not exercise the noise tail at all.
"""
import math
import pytest
import torch
pytest.importorskip("triton")
if not torch.cuda.is_available():
pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True)
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
DEVICE = "cuda"
VOCAB_SIZE = 200_000
NUM_SAMPLES = 500_000
# Dominant token is exp(HEAD_LOG_GAP)x larger than the unit-count tail, so the
# tail sits ~HEAD_LOG_GAP logits below the top.
HEAD_LOG_GAP = 18.0
# 10-sigma band: a correct sampler effectively never trips it.
Z_TOLERANCE = 10.0
def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor:
"""Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N."""
gen = torch.Generator(device=DEVICE).manual_seed(seed)
counts = torch.randint(
1, 4, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE
)
counts[0] = round(math.exp(HEAD_LOG_GAP)) # dominant token
return counts
def _counts_to_logits(counts: torch.Tensor) -> torch.Tensor:
# softmax(log(count)) == count / sum(count); count 0 -> logit -inf -> prob 0.
return counts.double().log().to(torch.float32)
def _sample(
logits_1d: torch.Tensor,
num_samples: int,
*,
use_fp64: bool = False,
temperature: float = 1.0,
) -> torch.Tensor:
"""Sample `num_samples` tokens from one logit vector.
Fixed seed with a distinct `pos` per sample gives independent draws; the
logits are broadcast with a 0-stride view to avoid materializing
[num_samples, vocab_size].
"""
vocab_size = logits_1d.shape[0]
logits = logits_1d.unsqueeze(0).expand(num_samples, vocab_size)
idx_mapping = torch.zeros(num_samples, dtype=torch.int32, device=DEVICE)
temp = torch.tensor([temperature], dtype=torch.float32, device=DEVICE)
seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_samples, dtype=torch.int64, device=DEVICE)
return gumbel_sample(
logits,
idx_mapping,
temp,
seed,
pos,
apply_temperature=True,
use_fp64=use_fp64,
)
def _z_score(observed: int, expected: float, num_trials: int) -> float:
p = expected / num_trials
return (observed - expected) / math.sqrt(num_trials * p * (1 - p))
def _sample_histogram(
logits_1d: torch.Tensor, num_samples: int, *, chunk: int = 1_000_000
) -> torch.Tensor:
"""Histogram of `num_samples` draws, accumulated in chunks.
Chunking keeps the kernel's per-sample scratch ([chunk, num_blocks]) bounded
so a large sample count does not blow up memory.
"""
vocab_size = logits_1d.shape[0]
hist = torch.zeros(vocab_size, dtype=torch.float64, device=DEVICE)
for start in range(0, num_samples, chunk):
size = min(chunk, num_samples - start)
logits = logits_1d.unsqueeze(0).expand(size, vocab_size)
idx_mapping = torch.zeros(size, dtype=torch.int32, device=DEVICE)
temp = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
pos = torch.arange(start, start + size, dtype=torch.int64, device=DEVICE)
out = gumbel_sample(
logits, idx_mapping, temp, seed, pos, apply_temperature=True
)
hist += torch.bincount(out, minlength=vocab_size).double()
return hist
# ----------------------------- Accuracy ------------------------------------
@pytest.mark.parametrize("use_fp64", [False, True])
def test_sampling_matches_target_distribution(use_fp64: bool):
counts = _make_heavy_tailed_counts()
total = counts.sum().item()
logits = _counts_to_logits(counts)
sampled = _sample(logits, NUM_SAMPLES, use_fp64=use_fp64)
assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
# The dominant token (index 0) and the aggregate tail are the two
# statistically resolvable bins (individual tail tokens are far below the
# ~5/N detectability floor). The tail mass is small but well above noise,
# and it lives beyond the fp32 Gumbel cap -- the regime sensitive to noise
# precision -- so matching it is the meaningful check.
tail_prob = (total - counts[0].item()) / total
tail_count = (sampled != 0).sum().item()
z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES)
assert abs(z) < Z_TOLERANCE, (
f"sampled tail mass {tail_count / NUM_SAMPLES:.3e} != target "
f"{tail_prob:.3e} (z={z:.2f})"
)
def test_full_vocab_distribution_fidelity():
"""The sampled distribution matches the target across the WHOLE vocab.
A near-flat count tensor makes every one of the 200K bins individually
measurable. With ~20 samples/bin, a goodness-of-fit over all bins checks
that no part of the vocab is over- or under-represented (the heavy-tailed
test above only resolves head vs aggregate tail). Empirically the fp32
sampler is as faithful here as torch.multinomial; the residual error is the
multinomial sampling-noise floor, not the kernel.
"""
gen = torch.Generator(device=DEVICE).manual_seed(2024)
counts = torch.randint(
500, 1500, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE
)
total = counts.sum().item()
logits = _counts_to_logits(counts)
num_samples = 4_000_000
hist = _sample_histogram(logits, num_samples)
# Diversity: essentially every token must be reachable (no starved region).
coverage = (hist > 0).sum().item() / VOCAB_SIZE
assert coverage > 0.99, f"only {coverage:.4f} of the vocab was ever sampled"
# Goodness-of-fit across all bins (each has expected count >= ~10).
expected = (counts.double() / total) * num_samples
chi2 = (((hist - expected) ** 2) / expected).sum().item()
df = VOCAB_SIZE - 1
assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.0f}, df={df}"
# ----------------------------- Edge cases ----------------------------------
def test_greedy_temperature_zero_returns_argmax():
"""temperature == 0 skips Gumbel noise and returns the exact argmax."""
torch.manual_seed(0)
num_reqs = 128
logits = torch.randn(num_reqs, VOCAB_SIZE, device=DEVICE, dtype=torch.float32)
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
temp = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE)
seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
sampled = gumbel_sample(
logits, idx_mapping, temp, seed, pos, apply_temperature=True
)
assert torch.equal(sampled, logits.argmax(dim=-1))
def test_zero_count_tokens_are_never_sampled():
"""Count 0 -> -inf logit -> probability 0; must never be selected."""
counts = _make_heavy_tailed_counts(seed=7)
zeroed = torch.arange(1, VOCAB_SIZE, 2, device=DEVICE) # odd indices (not head)
counts[zeroed] = 0
logits = _counts_to_logits(counts)
sampled = _sample(logits, NUM_SAMPLES)
assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
assert not torch.isin(sampled, zeroed).any(), "sampled a zero-probability token"
def test_single_nonzero_token_is_always_sampled():
"""A lone finite logit must win every draw, regardless of its index."""
counts = torch.zeros(VOCAB_SIZE, dtype=torch.int64, device=DEVICE)
counts[123_456] = 1000
logits = _counts_to_logits(counts)
sampled = _sample(logits, 10_000)
assert (sampled == 123_456).all()
@pytest.mark.parametrize("vocab_size", [1, 999, 1024, 4097])
def test_vocab_size_not_multiple_of_block(vocab_size: int):
"""Per-block tail masking for non-block-aligned vocab; all bins measurable."""
gen = torch.Generator(device=DEVICE).manual_seed(vocab_size)
counts = torch.randint(
20, 200, (vocab_size,), generator=gen, dtype=torch.int64, device=DEVICE
)
total = counts.sum().item()
logits = _counts_to_logits(counts)
num_samples = max(40 * vocab_size, 50_000)
sampled = _sample(logits, num_samples)
assert sampled.min() >= 0 and sampled.max() < vocab_size
observed = torch.bincount(sampled, minlength=vocab_size).double()
expected = (counts.double() / total) * num_samples
chi2 = (((observed - expected) ** 2) / expected).sum().item()
df = vocab_size - 1
if df >= 1:
assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.1f}, df={df}"
-6
View File
@@ -10,7 +10,6 @@ import torch
from vllm.platforms import current_platform
from vllm.sampling_params import SamplingParams
from vllm.utils.platform_utils import is_pin_memory_available
from vllm.utils.torch_utils import make_tensor_with_pad
from vllm.v1.pool.metadata import PoolingMetadata
from vllm.v1.sample.logits_processor import LogitsProcessors
@@ -236,7 +235,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
max_model_len=1024,
max_num_batched_tokens=1024,
device=torch.device(device),
pin_memory=is_pin_memory_available(),
vocab_size=1024,
block_sizes=[1],
kernel_block_sizes=[1],
@@ -331,7 +329,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
max_model_len=1024,
max_num_batched_tokens=1024,
device=torch.device(device),
pin_memory=is_pin_memory_available(),
vocab_size=1024,
block_sizes=[1],
kernel_block_sizes=[1],
@@ -341,7 +338,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
max_model_len=1024,
max_num_batched_tokens=1024,
device=torch.device(device),
pin_memory=is_pin_memory_available(),
vocab_size=1024,
block_sizes=[1],
kernel_block_sizes=[1],
@@ -410,7 +406,6 @@ def test_pooling_prompt_lens_not_aliased(device: str):
max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
max_num_batched_tokens=batch_size * (MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS),
device=torch.device(device),
pin_memory=is_pin_memory_available(),
vocab_size=VOCAB_SIZE,
block_sizes=[16],
kernel_block_sizes=[16],
@@ -459,7 +454,6 @@ def test_pooling_metadata_token_id_buffers(
max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
device=torch.device("cpu"),
pin_memory=False,
vocab_size=VOCAB_SIZE,
block_sizes=[16],
kernel_block_sizes=[16],
-4
View File
@@ -85,7 +85,6 @@ def initialize_kv_cache(runner: GPUModelRunner):
max_model_len=runner.max_model_len,
max_num_batched_tokens=runner.max_num_tokens,
device=runner.device,
pin_memory=runner.pin_memory,
vocab_size=runner.model_config.get_vocab_size(),
block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
kernel_block_sizes=[
@@ -1405,7 +1404,6 @@ def test_input_batch_with_kernel_block_sizes():
max_model_len = 512
max_num_batched_tokens = 512
device = torch.device(DEVICE_TYPE)
pin_memory = False
vocab_size = 50272
# Test with different kernel block sizes
@@ -1417,7 +1415,6 @@ def test_input_batch_with_kernel_block_sizes():
max_model_len=max_model_len,
max_num_batched_tokens=max_num_batched_tokens,
device=device,
pin_memory=pin_memory,
vocab_size=vocab_size,
block_sizes=block_sizes,
kernel_block_sizes=kernel_block_sizes,
@@ -1478,7 +1475,6 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
max_model_len=runner.max_model_len,
max_num_batched_tokens=runner.max_num_tokens,
device=runner.device,
pin_memory=runner.pin_memory,
vocab_size=runner.model_config.get_vocab_size(),
block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
kernel_block_sizes=[16],
+1 -1
View File
@@ -991,7 +991,7 @@ class VllmBackend:
},
payload_fn=lambda: json.dumps(
{
"model": self.vllm_config.model_config.model,
"model": getattr(self.vllm_config.model_config, "model", "unknown"),
"prefix": self.prefix,
"mode": str(cc.mode),
"backend": cc.backend,
+9
View File
@@ -302,6 +302,14 @@ class ParallelConfig:
Each entry must use `numactl --physcpubind` CPU-list syntax, for example
`"0-3"` or `"0,2,4-7"`.
"""
assigned_physical_gpu_ids: list[int] | None = None
"""Mapping from vLLM-local logical GPU IDs to physical GPU IDs.
For example, ``[2, 3]`` means logical GPU 0 maps to physical GPU 2,
and logical GPU 1 maps to physical GPU 3. Physical IDs are used only
at platform/topology boundaries such as NVML, NIC affinity, P2P
checks, and final CUDA device selection when needed. When None,
logical IDs map to visible device IDs in order."""
distributed_timeout_seconds: int | None = None
"""Timeout in seconds for distributed operations (e.g., init_process_group).
@@ -772,6 +780,7 @@ class ParallelConfig:
"numa_bind",
"numa_bind_nodes",
"numa_bind_cpus",
"assigned_physical_gpu_ids",
}
from vllm.config.utils import get_hash_factors, hash_factors
+2 -2
View File
@@ -18,8 +18,8 @@ import torch
from vllm.device_allocator import AllocationData, HandleType
from vllm.logger import init_logger
from vllm.utils.platform_utils import is_pin_memory_available
from vllm.utils.system_utils import find_loaded_library
from vllm.utils.torch_utils import PIN_MEMORY
logger = init_logger(__name__)
@@ -196,7 +196,7 @@ class CuMemAllocator:
size_in_bytes,
dtype=torch.uint8,
device="cpu",
pin_memory=is_pin_memory_available(),
pin_memory=PIN_MEMORY,
)
cpu_ptr = cpu_backup_tensor.data_ptr()
libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes)
+2 -2
View File
@@ -11,7 +11,7 @@ import torch
from vllm.device_allocator import AllocationData, HandleType
from vllm.logger import init_logger
from vllm.utils.platform_utils import is_pin_memory_available
from vllm.utils.torch_utils import PIN_MEMORY
logger = init_logger(__name__)
@@ -188,7 +188,7 @@ class XpuMemAllocator:
size_in_bytes,
dtype=torch.uint8,
device="cpu",
pin_memory=is_pin_memory_available(),
pin_memory=PIN_MEMORY,
)
cpu_ptr = cpu_backup_tensor.data_ptr()
_xpu_memcpy_sync(
@@ -704,7 +704,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
self.num_experts = num_experts
self.cleanup()
gpus_per_node = torch.accelerator.device_count()
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
gpus_per_node = (
len(assigned_physical_gpu_ids)
if assigned_physical_gpu_ids is not None
else torch.accelerator.device_count()
)
logger.debug(
"Making One-sided NVLink mapping: rank=%d, world size=%d",
self.rank,
@@ -320,13 +320,21 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
is_distributed = dist.is_initialized()
num_dev = current_platform.device_count()
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices is None:
cuda_visible_devices = ",".join(str(i) for i in range(num_dev))
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
if assigned_physical_gpu_ids is not None:
# Key by the ordered list: the cache stores directed local-index
# pairs, so permutations of the same set are distinct mappings.
cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids)
num_dev = len(assigned_physical_gpu_ids)
else:
num_dev = current_platform.device_count()
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev))
path = os.path.join(
envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json"
envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json"
)
os.makedirs(os.path.dirname(path), exist_ok=True)
from vllm.distributed.parallel_state import get_world_group
@@ -338,7 +346,15 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
# enter this block to calculate the cache
logger.info("generating GPU P2P access cache in %s", path)
cache: dict[str, bool] = {}
ids = list(range(num_dev))
# The probe subprocesses inherit this process's device-control env
# var, so they must be given visible ordinals, not physical IDs.
if assigned_physical_gpu_ids is not None:
ids = [
current_platform.logical_device_id_to_visible_device_id(local)
for local in range(num_dev)
]
else:
ids = list(range(num_dev))
# batch of all pairs of GPUs
batch_src, batch_tgt = zip(*list(product(ids, ids)))
# NOTE: we use `subprocess` rather than `multiprocessing` here
@@ -368,8 +384,11 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
) from e
with open(output_file.name, "rb") as f:
result = pickle.load(f)
# Cache entries must be keyed by local indices (0..N-1) because
# gpu_p2p_access_check() is called with local ranks.
id_to_local = {device_id: local for local, device_id in enumerate(ids)}
for _i, _j, r in zip(batch_src, batch_tgt, result):
cache[f"{_i}->{_j}"] = r
cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r
with open(path, "w") as f:
json.dump(cache, f, indent=4)
if is_distributed:
@@ -34,7 +34,12 @@ def _can_p2p(rank: int, world_size: int) -> bool:
continue
if envs.VLLM_SKIP_P2P_CHECK:
logger.debug("Skipping P2P check and trusting the driver's P2P report.")
return torch.cuda.can_device_access_peer(rank, i)
# can_device_access_peer takes visible device ordinals, while
# rank and i are logical local IDs.
return torch.cuda.can_device_access_peer(
current_platform.logical_device_id_to_visible_device_id(rank),
current_platform.logical_device_id_to_visible_device_id(i),
)
if not gpu_p2p_access_check(rank, i):
return False
return True
@@ -126,13 +131,10 @@ class CustomAllreduce:
CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],
max_size,
)
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices:
device_ids = list(map(int, cuda_visible_devices.split(",")))
else:
device_ids = list(range(current_platform.device_count()))
physical_device_id = device_ids[device.index]
# device.index is a visible ordinal, not a logical local ID.
physical_device_id = current_platform.visible_device_id_to_physical_device_id(
device.index
)
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
gather_list = [
torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
@@ -129,12 +129,10 @@ class QuickAllReduce:
assert isinstance(device, torch.device)
self.device = device
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices:
device_ids = list(map(int, cuda_visible_devices.split(",")))
else:
device_ids = list(range(current_platform.device_count()))
physical_device_id = device_ids[device.index]
# device.index is a visible ordinal, not a logical local ID.
physical_device_id = current_platform.visible_device_id_to_physical_device_id(
device.index
)
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
gather_list = [
torch.tensor([0], dtype=torch.int, device="cpu")
@@ -840,7 +840,13 @@ class MessageQueue:
The MessageQueue instance for the calling process,
and a list of handles (only non-empty for the reader process).
"""
local_size = current_platform.device_count()
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
if assigned_physical_gpu_ids is not None:
local_size = len(assigned_physical_gpu_ids)
else:
local_size = current_platform.device_count()
rank = dist.get_rank()
same_node = rank // local_size == reader_rank // local_size
buffer_io = MessageQueue(
@@ -482,10 +482,11 @@ def _init_lmcache_engine(
)
# Change current device.
num_gpus = torch.accelerator.device_count()
local_rank = parallel_config.rank % num_gpus
torch.accelerator.set_device_index(local_rank)
device = torch.device(f"cuda:{local_rank}")
from vllm.distributed.parallel_state import get_world_group
device_index = get_world_group().device_index
torch.accelerator.set_device_index(device_index)
device = torch.device(f"cuda:{device_index}")
metadata = LMCacheEngineMetadata(
model_config.model,
parallel_config.world_size,
@@ -176,7 +176,7 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
self,
request: Request,
num_computed_tokens: int,
) -> tuple[int, bool]:
) -> tuple[int | None, bool]:
assert self.connector_scheduler is not None
return self.connector_scheduler.get_num_new_matched_tokens(
request, num_computed_tokens
@@ -2,13 +2,15 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""External-store cache-hit coordinator for MooncakeStoreConnector."""
from collections.abc import Sequence
from typing import cast
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
chunk_hashes_for_block_size,
)
from vllm.v1.core.block_pool import BlockPool
from vllm.v1.core.kv_cache_utils import (
BlockHash,
BlockHashList,
BlockHashListWithBlockSize,
KVCacheBlock,
)
from vllm.v1.core.single_type_kv_cache_manager import (
@@ -120,7 +122,7 @@ class MooncakeStoreCoordinator:
def find_longest_cache_hit(
self,
block_hashes: list[BlockHash],
block_hashes: Sequence[BlockHash],
max_length: int,
cached_block_pool: ExternalCachedBlockPool,
*,
@@ -147,7 +149,7 @@ class MooncakeStoreCoordinator:
def load_mask(
self,
block_hashes: list[BlockHash],
block_hashes: Sequence[BlockHash],
token_len: int,
) -> tuple[list[bool], ...]:
"""Per-group load masks: ``mask[g][i]`` is True iff group ``g``'s
@@ -236,17 +238,15 @@ class MooncakeStoreCoordinator:
return tuple(masks)
def block_hashes_for_spec(
self, block_hashes: list[BlockHash], spec: KVCacheSpec
) -> BlockHashList:
if spec.block_size == self.hash_block_size:
return block_hashes
return BlockHashListWithBlockSize(
self, block_hashes: Sequence[BlockHash], spec: KVCacheSpec
) -> Sequence[BlockHash]:
return chunk_hashes_for_block_size(
block_hashes, self.hash_block_size, spec.block_size
)
def _find_hit_blocks(
self,
block_hashes: list[BlockHash],
block_hashes: Sequence[BlockHash],
max_length: int,
cached_block_pool: ExternalCachedBlockPool,
*,
@@ -264,7 +264,7 @@ class MooncakeStoreCoordinator:
spec, group_ids, manager_cls = self.attention_groups[0]
hashes = self.block_hashes_for_spec(block_hashes, spec)
hit_blocks = manager_cls.find_longest_cache_hit(
block_hashes=hashes,
block_hashes=hashes, # type: ignore[arg-type]
max_length=max_length,
kv_cache_group_ids=group_ids,
block_pool=cast(BlockPool, cached_block_pool),
@@ -304,7 +304,7 @@ class MooncakeStoreCoordinator:
_max_length = min(curr_hit_length + spec.block_size, max_length)
hashes = self.block_hashes_for_spec(block_hashes, spec)
hit_blocks = manager_cls.find_longest_cache_hit(
block_hashes=hashes,
block_hashes=hashes, # type: ignore[arg-type]
max_length=_max_length,
kv_cache_group_ids=group_ids,
block_pool=cast(BlockPool, cached_block_pool),
@@ -5,8 +5,9 @@
# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
"""Data classes for MooncakeStoreConnector."""
from collections.abc import Iterable
from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from typing import cast
import torch
@@ -23,6 +24,77 @@ from vllm.v1.core.kv_cache_utils import (
logger = init_logger(__name__)
class BlobBlockHashes(Sequence[BlockHash]):
"""Lazy view over a flat buffer of fixed-size block hashes to avoid the overhead
of materializing all hashes upfront.
"""
def __init__(self, blob: memoryview, hash_len: int):
self._blob = blob
self._hash_len = hash_len
self._n = len(blob) // hash_len if hash_len else 0
def __len__(self) -> int:
return self._n
def __getitem__(self, idx):
if isinstance(idx, slice):
return [self[i] for i in range(*idx.indices(self._n))]
if idx < 0:
idx += self._n
if not 0 <= idx < self._n:
raise IndexError(idx)
off = idx * self._hash_len
return BlockHash(self._blob[off : off + self._hash_len])
class _CompactChunkHashList(BlockHashListWithBlockSize):
"""View that keys each ``block_size`` chunk by the last constituent
``hash_block_size`` hash instead of concatenating all of them.
The engine chains block hashes (each hash folds in the previous one), so the
final sub-block hash of a chunk already uniquely identifies the whole chunk
and its prefix. Using it keeps a Mooncake key at a single hash digest
regardless of the ``block_size`` / ``hash_block_size`` ratio, instead of
growing the key linearly with it (e.g. 64x for ``block_size=256``,
``hash_block_size=4``).
"""
def __init__(
self,
block_hashes: Sequence[BlockHash],
hash_block_size: int,
target_block_size: int,
):
# Accept any indexable sequence (e.g. the lazy ``BlobBlockHashes``), not
# just ``list``; the base only indexes/sizes it.
assert target_block_size % hash_block_size == 0
self.block_hashes = block_hashes # type: ignore[assignment]
self.scale_factor = target_block_size // hash_block_size
def _get_value_at(self, idx: int) -> BlockHash:
return self.block_hashes[idx * self.scale_factor + self.scale_factor - 1]
def chunk_hashes_for_block_size(
block_hashes: Sequence[BlockHash],
hash_block_size: int,
block_size: int,
) -> Sequence[BlockHash]:
"""Map ``hash_block_size``-granular block hashes to one compact hash per
``block_size`` chunk (the chunk's last sub-hash). Returns ``block_hashes``
unchanged when the two sizes are equal.
"""
if block_size == hash_block_size:
return block_hashes
# Structurally a Sequence[BlockHash] (indexable + sized); the base class
# just isn't declared as one.
return cast(
"Sequence[BlockHash]",
_CompactChunkHashList(block_hashes, hash_block_size, block_size),
)
@dataclass
class KeyMetadata:
"""Metadata for constructing pool keys."""
@@ -138,18 +210,15 @@ class ChunkedTokenDatabase:
Args:
token_len: Total number of tokens.
block_hashes: Block hashes computed at ``hash_block_size`` granularity.
When ``block_size > hash_block_size`` consecutive hashes are merged
up to the group's ``block_size`` via ``BlockHashListWithBlockSize``.
When ``block_size > hash_block_size`` each group's ``block_size`` chunk
is keyed by its last sub-hash via ``chunk_hashes_for_block_size``.
mask_num: Number of tokens to skip from the beginning.
"""
if not block_hashes:
return
if self.block_size == self.hash_block_size:
chunk_hashes: Iterable[BlockHash] = block_hashes
else:
chunk_hashes = BlockHashListWithBlockSize(
block_hashes, self.hash_block_size, self.block_size
)
chunk_hashes: Iterable[BlockHash] = chunk_hashes_for_block_size(
block_hashes, self.hash_block_size, self.block_size
)
for chunk_id, h in enumerate(chunk_hashes):
start_idx = chunk_id * self.block_size
if start_idx >= token_len:
@@ -11,7 +11,10 @@ Wire format (REQ/REP over IPC):
msg_type == LOOKUP_MSG:
frame 1: token_len (u32 big-endian, 4 bytes)
frame 2..n: msgpack-encoded list[str] of block-hash hex digests
frame 2: hash_len (u16 big-endian, 2 bytes) byte length of each
fixed-size block hash (0 when there are no hashes)
frame 3: raw block hashes concatenated back-to-back (each hash_len
bytes); the server splits on hash_len
Response: [hit_count: u32 big-endian, 4 bytes]
msg_type == RESET_MSG:

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