Compare commits

...
Author SHA1 Message Date
khluuandClaude Opus 4.6 a8b0f65f44 Remove broken allocator pre-warming (initializes CUDA before fork)
The pre-warming called torch.cuda.mem_get_info() in the parent
process before the EngineCore subprocess forks, which initialized
CUDA. This caused "Cannot re-initialize CUDA in forked subprocess"
in the child process, breaking tests that previously passed.

The pre-warming also wouldn't have helped — PyTorch's allocator
pool is per-process, so the child's pool starts empty regardless.

Signed-off-by: khluu <khluu000@gmail.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-13 02:26:43 -07:00
khluuandClaude Opus 4.6 8a93cda9a9 [CI] Migrate all remaining gpu_1_queue jobs to h200_18gb MIG
Move all test steps that were falling through to gpu_1_queue (L4)
to explicitly use device: h200_18gb (H200 18GB MIG partition).

This covers 52 steps across 15 test area files. Requires ci-infra
PR #355 (PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False) to be
merged first — without it, PyTorch's CUDACachingAllocator hits an
NVML assertion on MIG partitions when allocating large tensors.

Signed-off-by: khluu <khluu000@gmail.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-12 18:50:08 -07:00
khluuandClaude Opus 4.6 0ab46391b7 [Core] Add MIG detection and workaround for PyTorch NVML assertion
On MIG (Multi-Instance GPU) partitions, NVML access is restricted.
PyTorch's CUDACachingAllocator uses NVML internally when growing
the memory pool, causing:

  NVML_SUCCESS == r INTERNAL ASSERT FAILED at
  "/pytorch/c10/cuda/CUDACachingAllocator.cpp":1165

This change:
1. Detects MIG via NVIDIA_VISIBLE_DEVICES="MIG-..." prefix
2. On MIG, sets PYTORCH_NVML_BASED_CUDA_CHECK=0 (use CUDA runtime
   instead of NVML for device checks)
3. On MIG, sets PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
   (avoids one NVML code path in the allocator)
4. Pre-warms the PyTorch allocator pool by allocating 90% of GPU
   memory upfront, so later KV cache allocations reuse the pool
   without triggering NVML-based pool growth

Signed-off-by: khluu <khluu000@gmail.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-12 15:57:19 -07:00
Kevin H. LuuandGitHub 1ff9d33535 [CI] Migrate remaining B200 jobs to b200-k8s with test fixes (#42387)
Signed-off-by: khluu <khluu000@gmail.com>
2026-05-12 02:00:37 -07:00
7f65f84428 [Bugfix] Fix empty channel/recipient in harmony for /v1/responses (#35540)
Signed-off-by: kg6-sleipnir <christopherhazen42@gmail.com>
Signed-off-by: chazen <45186108+kg6-sleipnir@users.noreply.github.com>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
2026-05-12 08:45:51 +00:00
amitz-nvandGitHub ef34592a1a [Bugfix] Fix double reduce in flashinfer_nvlink_two_sided and flashinfer_nvlink_one_sided backends (#41382)
Signed-off-by: amitz-nv <203509407+amitz-nv@users.noreply.github.com>
2026-05-12 07:47:47 +00:00
Kevin H. LuuandGitHub f69644caf8 [CI] Migrate more B200 jobs to b200-k8s queue (#42356)
Signed-off-by: khluu <khluu000@gmail.com>
2026-05-12 00:38:31 -07:00
wang.yuqiGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
d37e25ffbe [Frontend] Consolidate Speech to Text entrypoints. (#42370)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
Signed-off-by: wang.yuqi <noooop@126.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-05-12 07:06:57 +00:00
8517cdaf90 [XPU] update dp rank w/o env-var isolation (#39856)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-05-12 14:49:54 +08:00
Lucas KabelaandGitHub 4e498b5e5c [Bugfix][Performance Improvement] Improve penalties triton kernel performance (#40657)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
2026-05-12 05:47:20 +00:00
Yasmin MoslemGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
28ee78af54 Implement custom dataset class for ASR benchmarking (#41576)
Signed-off-by: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-05-12 12:17:58 +08:00
ZiTian ZhaoandGitHub 630492da30 [Fix] Gemma4 Mixed-Resolution Image Co-Batching Crash (#42217)
Signed-off-by: zitian.zhao <zitian.zhao@tencentmusic.com>
2026-05-12 03:13:03 +00:00
ChaunceyandGitHub 920bf3ec84 [Bugifx] [Qwen3CoderTool] Restore supports_required_and_named for required tool_choice (#42292)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-05-12 02:09:56 +00:00
pschlan-amdandGitHub 39dff5ff39 Add VLLM_USE_SPINLOOP_EXT to use more efficient busy polling (#36517)
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
2026-05-11 16:11:49 -07:00
Wentao YeGitHubNick Hillmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
d7af6b34d8 [Model Runner V2] Bug fix: logprob dtype int64/int32 issue (#41761)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-05-11 21:55:43 +00:00
Nick HillGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
bbee532988 [Perf][1/n] Eliminate various GPU<->CPU syncs (#41429)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-05-11 20:36:03 +00:00
53181384e0 [Bugfix] Fix DSV4 swiglu_limit on marlin backend (#42287)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-11 13:03:56 -07:00
wang.yuqiandGitHub a0dc7a0f36 [CI] Consolidate Speech to Text tests (#42274)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-05-11 19:50:17 +00:00
56e5810ff1 [BugFix] Prevent orphaned process on NCCL destroy (#39846)
Signed-off-by: Jeffrey Wang <jeffreywang@anyscale.com>
Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
2026-05-11 15:25:26 -04:00
Flora FengandGitHub 639cbfd274 [CI] Add tests/parser to CI coverage (#41877)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-05-11 19:08:54 +00:00
a721315488 [ROCm][Perf] Fix RMSNorm+Quant fusion for gfx950 (non-fnuz) (#41825)
Signed-off-by: Frida Andersson <fanderss@amd.com>
Signed-off-by: Chuan Li <chuali@amd.com>
Co-authored-by: Markus Hartikainen <markus.hartikainen@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Chuan Li <chuali@amd.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Co-authored-by: Frida Andersson <frida-andersson@users.noreply.github.com>
Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
2026-05-11 15:00:51 -04:00
VensenGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
6fdb49392e [Bugfix] Fix int32 overflow in DeepGEMM SiLU/mul FP8 Triton kernel (#42201)
Signed-off-by: vensen <vensenmu@gmail.com>
Signed-off-by: Vensen <vensenmu@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-05-11 14:52:31 -04:00
cf0d279142 [Docs] Add Apple Silicon documentation for vLLM-Metal GPU support (#41987)
Signed-off-by: alexagriffith <agriffith96@gmail.com>
Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
2026-05-11 11:34:25 -07:00
5497ffbf7c Add documentation about vLLM FIPS compliance (#42190)
Signed-off-by: Vinay Damodaran <vrdn@hey.com>
Signed-off-by: Vinay R Damodaran <vrdn@hey.com>
Co-authored-by: Russell Bryant <russell.bryant@gmail.com>
2026-05-11 18:17:02 +00:00
Nick HillandGitHub 9af6a5ed75 [Model Runner V2] Fix seq_lens_cpu_upper_bound (#42202)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-05-11 10:37:50 -07:00
Hexiang WangandGitHub 7863fff6e5 [ROCm][DSv4] implement flash sparse mla with triton kernels (#41812)
Signed-off-by: whx-sjtu <xiaowang990929@gmail.com>
2026-05-11 09:27:11 -07:00
Wentao YeandGitHub 0d453e2336 [Perf] Batch invariance with Cutlass fp8 support, 28.9% E2E latency improvement (#40408)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
2026-05-11 12:20:58 -04:00
Wentao YeandGitHub 3f9c0c25b3 [Bug] Fix kimi dtype issue with mm_projector_forward (#42081)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-05-11 11:45:24 -04:00
Vadim GimpelsonandGitHub a2e776d716 [Bugfix] Accept canonicalized modelopt_* quant_method in _extract_modelopt_quant_algo (#42181)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
2026-05-11 11:10:57 -04:00
Martin HickeyGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
4955990f1b [kv_offload] Move FilterReusedOffloadingManager logic to CPUOffloadingManager (#41727)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-05-11 18:09:29 +03:00
Wentao YeandGitHub 4b64fc2cbf [Refactor] Cleanup batch invariant dead code (#41993)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-05-11 10:48:39 -04:00
pschlan-amdandGitHub 5f1b313900 [ROCm] Clean up a bit the AITER FA backend (#41942)
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
2026-05-11 22:45:18 +08:00
724ed2fc35 [DSv4] Improved dequant gather K cache kernel (#42236)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-11 10:41:12 -04:00
a51376b3f0 [Performance][DSR1]: Fused RoPE+KVCache+q_concat for MLA (#40392)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
Co-authored-by: ElizaWszola <ewszola@redhat.com>
2026-05-11 14:10:50 +00:00
Martin HickeyandGitHub 8415bf2cdb [kv_offload] Set offloading connector to prefer HND layout (#41928)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>
2026-05-11 15:05:41 +03:00
Noa NeriaandGitHub ac062147fa Avoid silent weights corruption when loading Nemotron Nano VL with reusable-buffer loaders like runai distributed streaming (#42244)
Signed-off-by: Noa Neria <nneria@nvidia.com>
2026-05-11 12:03:14 +00:00
ChaunceyandGitHub 617239b70c [Frontend]Responses API supports chat_template_kwargs (#42272)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-05-11 11:59:39 +00:00
Kyungmin LeeandGitHub 27ae676364 Fix EXAONE-4.5 to align with Transformers update (#42246)
Signed-off-by: lkm2835 <lkm2835@gmail.com>
2026-05-11 10:25:31 +00:00
haosdentandGitHub 17ed5e61f5 [CI] Make Python-only Installation optional (#42293)
Signed-off-by: haosdent <haosdent@gmail.com>
2026-05-11 09:47:16 +00:00
Nicolò LucchesiandGitHub 5672d100ed [KV Connector][NIXL][Bugfix] Fix NIXL handshake failures not honoring kv_load_failure_policy (#40364)
When NIXL handshake fails (e.g., due to compatibility hash mismatch
between prefill and decode instances), requests fail with "engine dead"
error instead of gracefully falling back to local recomputation as configured
by kv_load_failure_policy='recompute'.

Signed-off-by: NickLucche <nlucches@redhat.com>
2026-05-11 09:37:21 +00:00
Nicolò LucchesiandGitHub 770e9bd6b3 [Nixl][PD] Lease renewal TTL KV blocks on P (#41383)
Signed-off-by: NickLucche <nlucches@redhat.com>
2026-05-11 09:27:30 +00:00
Cyrus LeungandGitHub 9efdddca28 [Model] Fix missing maybe_prefix (#42280)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
2026-05-11 09:04:06 +00:00
QiuandGitHub b1b59720b2 bugfix(flashinfer,dcp): remove kv_cache_layout for BatchDCPPrefillWrapper._new_tokens. (#38895)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>
2026-05-11 08:11:49 +00:00
shaharmor98GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
f9f770ca0b fix nixl side-channel host selection (#41806)
Signed-off-by: Shahar Mor <smor@nvidia.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-05-11 07:40:37 +00:00
Haoqing WangandGitHub 5cba6839e6 Document MolmoWeb hf_overrides (#42163)
Signed-off-by: Haoqi Wang <78337154+hqhq1025@users.noreply.github.com>
2026-05-10 23:58:22 -07:00
Jee Jee LiandGitHub 05d610e5cd [CI/Build] Reduce LoRA model tests. (#42266)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-05-11 14:49:08 +08:00
209 changed files with 7899 additions and 1705 deletions
+34 -6
View File
@@ -460,7 +460,7 @@ steps:
- tests/lora
- vllm/platforms/rocm.py
commands:
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
#------------------------------------------------------ mi250 · model_executor -------------------------------------------------------#
@@ -929,6 +929,7 @@ steps:
- tests/tokenizers_
- tests/reasoning
- tests/tool_parsers
- tests/parser
- tests/transformers_utils
- tests/config
commands:
@@ -942,6 +943,7 @@ steps:
- pytest -v -s tokenizers_
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
- pytest -v -s tool_parsers
- pytest -v -s parser
- pytest -v -s transformers_utils
- pytest -v -s config
@@ -1320,7 +1322,6 @@ steps:
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai/speech_to_text/
- pytest -v -s entrypoints/test_chat_utils.py
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
@@ -1336,7 +1337,21 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- label: Entrypoints Integration (Speech to Text) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
fast_check: true
torch_nightly: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/speech_to_text
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/speech_to_text
- label: Entrypoints Integration (LLM) # TBD
timeout_in_minutes: 180
@@ -1760,7 +1775,7 @@ steps:
- export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
- pytest -v -s -x lora/test_chatglm3_tp.py
- pytest -v -s -x lora/test_llama_tp.py
- pytest -v -s -x lora/test_llm_with_multi_loras.py
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
- pytest -v -s -x lora/test_olmoe_tp.py
- pytest -v -s -x lora/test_gptoss_tp.py
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -2766,7 +2781,6 @@ steps:
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai/speech_to_text/
- pytest -v -s entrypoints/test_chat_utils.py
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
@@ -2782,7 +2796,21 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- label: Entrypoints Integration (Speech to Text) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
agent_pool: mi355_1
fast_check: true
torch_nightly: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/speech_to_text
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/speech_to_text
- label: Entrypoints Integration (Pooling) # TBD
timeout_in_minutes: 180
+1
View File
@@ -20,6 +20,7 @@ steps:
- label: Cudagraph
key: cudagraph
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
- tests/v1/cudagraph
- vllm/v1/cudagraph_dispatcher.py
+2
View File
@@ -30,6 +30,7 @@ steps:
- label: Engine (1 GPU)
key: engine-1-gpu
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/v1/engine/
- tests/v1/engine/
@@ -50,6 +51,7 @@ steps:
- label: e2e Core (1 GPU)
key: e2e-core-1-gpu
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
+22 -5
View File
@@ -6,17 +6,19 @@ steps:
key: entrypoints-unit-tests
timeout_in_minutes: 10
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/entrypoints
- tests/entrypoints/
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
- label: Entrypoints Integration (LLM)
key: entrypoints-integration-llm
timeout_in_minutes: 40
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/entrypoints/llm
@@ -36,6 +38,7 @@ steps:
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
@@ -44,18 +47,17 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
- label: Entrypoints Integration (API Server openai - Part 2)
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai/speech_to_text/
- pytest -v -s entrypoints/test_chat_utils.py
- label: Entrypoints Integration (API Server openai - Part 3)
@@ -69,12 +71,13 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- label: Entrypoints Integration (API Server 2)
key: entrypoints-integration-api-server-2
timeout_in_minutes: 130
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/entrypoints/rpc
@@ -86,10 +89,23 @@ steps:
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- label: Entrypoints Integration (Speech to Text)
key: entrypoints-integration-speech_to_text
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/entrypoints/speech_to_text
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/speech_to_text
- label: Entrypoints Integration (Pooling)
key: entrypoints-integration-pooling
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/entrypoints/pooling
@@ -101,6 +117,7 @@ steps:
key: entrypoints-integration-responses-api
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/entrypoints/openai/responses
@@ -115,5 +132,5 @@ steps:
- csrc/
- vllm/entrypoints/openai/
- vllm/model_executor/models/whisper.py
commands: # LMEval+Transcription WER check
commands: # LMEval
- pytest -s entrypoints/openai/correctness/
+7 -1
View File
@@ -17,6 +17,7 @@ steps:
- label: Kernels Core Operation Test
key: kernels-core-operation-test
timeout_in_minutes: 75
device: h200_18gb
source_file_dependencies:
- csrc/
- tests/kernels/core
@@ -41,6 +42,7 @@ steps:
- label: Kernels Attention Test %N
key: kernels-attention-test
timeout_in_minutes: 35
device: h200_18gb
source_file_dependencies:
- csrc/attention/
- vllm/v1/attention
@@ -55,6 +57,7 @@ steps:
- label: Kernels Quantization Test %N
key: kernels-quantization-test
timeout_in_minutes: 90
device: h200_18gb
source_file_dependencies:
- csrc/quantization/
- vllm/model_executor/layers/quantization
@@ -66,6 +69,7 @@ steps:
- label: Kernels MoE Test %N
key: kernels-moe-test
timeout_in_minutes: 25
device: h200_18gb
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
- csrc/moe/
@@ -82,6 +86,7 @@ steps:
- label: Kernels Mamba Test
key: kernels-mamba-test
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- csrc/mamba/
- tests/kernels/mamba
@@ -91,6 +96,7 @@ steps:
- label: Kernels KDA Test
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
- vllm/model_executor/layers/fla/ops/kda.py
- vllm/model_executor/layers/fla/ops/chunk_delta_h.py
@@ -242,7 +248,7 @@ steps:
- label: Kernels FusedMoE Layer Test (2 B200s)
key: kernels-fusedmoe-layer-test-2-b200s
timeout_in_minutes: 90
device: b200
device: b200-k8s
num_devices: 2
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
+18 -3
View File
@@ -5,6 +5,7 @@ steps:
- label: LM Eval Small Models
key: lm-eval-small-models
timeout_in_minutes: 75
device: h200_18gb
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
@@ -40,7 +41,7 @@ steps:
- label: LM Eval Small Models (B200)
key: lm-eval-small-models-b200
timeout_in_minutes: 120
device: b200
device: b200-k8s
optional: true
source_file_dependencies:
- csrc/
@@ -48,6 +49,18 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
- label: LM Eval Large Models (B200, EP)
key: lm-eval-large-models-b200-ep
timeout_in_minutes: 120
device: b200-k8s
optional: true
num_devices: 2
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep.txt
- label: LM Eval Qwen3.5 Models (B200)
key: lm-eval-qwen3-5-models-b200
timeout_in_minutes: 120
@@ -92,7 +105,7 @@ steps:
- label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
key: moe-refactor-integration-test-b200-dp-temporary
device: b200
device: b200-k8s
optional: true
num_devices: 2
commands:
@@ -102,6 +115,7 @@ steps:
- label: LM Eval TurboQuant KV Cache
key: lm-eval-turboquant-kv-cache
timeout_in_minutes: 75
device: h200_18gb
source_file_dependencies:
- vllm/model_executor/layers/quantization/turboquant/
- vllm/v1/attention/backends/turboquant_attn.py
@@ -127,7 +141,7 @@ steps:
- label: GPQA Eval (GPT-OSS) (B200)
key: gpqa-eval-gpt-oss-b200
timeout_in_minutes: 120
device: b200
device: b200-k8s
optional: true
num_devices: 2
source_file_dependencies:
@@ -140,6 +154,7 @@ steps:
- label: MRCR Eval Small Models
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- tests/evals/mrcr/
commands:
+3 -2
View File
@@ -5,11 +5,12 @@ steps:
- label: LoRA %N
key: lora
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/lora
- tests/lora
commands:
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
parallelism: 4
@@ -31,7 +32,7 @@ steps:
# requires multi-GPU testing for validation.
- pytest -v -s -x lora/test_chatglm3_tp.py
- pytest -v -s -x lora/test_llama_tp.py
- pytest -v -s -x lora/test_llm_with_multi_loras.py
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
- pytest -v -s -x lora/test_olmoe_tp.py
- pytest -v -s -x lora/test_gptoss_tp.py
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
+8
View File
@@ -5,6 +5,7 @@ steps:
- label: V1 Spec Decode
key: v1-spec-decode
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/config/
- vllm/distributed/
@@ -58,6 +59,7 @@ steps:
- label: V1 Core + KV + Metrics
key: v1-core-kv-metrics
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/config/
- vllm/distributed/
@@ -153,6 +155,7 @@ steps:
key: examples
timeout_in_minutes: 45
working_dir: "/vllm-workspace/examples"
device: h200_18gb
source_file_dependencies:
- vllm/entrypoints
- vllm/multimodal
@@ -210,7 +213,9 @@ steps:
- label: Python-only Installation
key: python-only-installation
depends_on: ~
optional: true
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
- tests/standalone_tests/python_only_compile.sh
- setup.py
@@ -220,6 +225,7 @@ steps:
- label: Async Engine, Inputs, Utils, Worker
key: async-engine-inputs-utils-worker
timeout_in_minutes: 50
device: h200_18gb
source_file_dependencies:
- vllm/assets/
- vllm/config/
@@ -282,6 +288,7 @@ steps:
- tests/tokenizers_
- tests/reasoning
- tests/tool_parsers
- tests/parser
- tests/transformers_utils
- tests/config
device: cpu-small
@@ -296,6 +303,7 @@ steps:
- pytest -v -s tokenizers_
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
- pytest -v -s tool_parsers
- pytest -v -s parser
- pytest -v -s transformers_utils
- pytest -v -s config
@@ -5,6 +5,7 @@ steps:
- label: Model Executor
key: model-executor
timeout_in_minutes: 35
device: h200_18gb
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
@@ -5,6 +5,7 @@ steps:
- label: Model Runner V2 Core Tests
key: model-runner-v2-core-tests
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
@@ -29,6 +30,7 @@ steps:
key: model-runner-v2-examples
timeout_in_minutes: 45
working_dir: "/vllm-workspace/examples"
device: h200_18gb
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/core/sched/
@@ -102,6 +104,7 @@ steps:
key: model-runner-v2-spec-decode
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
device: h200_18gb
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
+5
View File
@@ -6,6 +6,7 @@ steps:
key: basic-models-tests-initialization
timeout_in_minutes: 45
torch_nightly: true
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/test_initialization.py
@@ -19,6 +20,7 @@ steps:
- label: Basic Models Tests (Extra Initialization) %N
key: basic-models-tests-extra-initialization
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/model_executor/models/
- tests/models/test_initialization.py
@@ -35,6 +37,7 @@ steps:
- label: Basic Models Tests (Other)
key: basic-models-tests-other
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/test_terratorch.py
@@ -61,6 +64,7 @@ steps:
working_dir: "/vllm-workspace/"
optional: true
soft_fail: true
device: h200_18gb
commands:
- pip install --upgrade git+https://github.com/huggingface/transformers
- pytest -v -s tests/models/test_initialization.py
@@ -77,6 +81,7 @@ steps:
working_dir: "/vllm-workspace/"
optional: true
soft_fail: true
device: h200_18gb
commands:
- pip install transformers==4.57.5
- pytest -v -s tests/models/test_initialization.py
@@ -5,6 +5,7 @@ steps:
- label: Language Models Tests (Standard)
key: language-models-tests-standard
timeout_in_minutes: 25
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/language
@@ -18,6 +19,7 @@ steps:
- label: Language Models Tests (Extra Standard) %N
key: language-models-tests-extra-standard
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/model_executor/models/
- tests/models/language/pooling/test_embedding.py
@@ -35,6 +37,7 @@ steps:
- label: Language Models Tests (Hybrid) %N
key: language-models-tests-hybrid
timeout_in_minutes: 75
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/language/generation
@@ -60,6 +63,7 @@ steps:
- label: Language Models Test (Extended Generation) # 80min
key: language-models-test-extended-generation
timeout_in_minutes: 110
device: h200_18gb
optional: true
source_file_dependencies:
- vllm/
@@ -85,6 +89,7 @@ steps:
- label: Language Models Test (Extended Pooling) # 36min
key: language-models-test-extended-pooling
timeout_in_minutes: 50
device: h200_18gb
optional: true
source_file_dependencies:
- vllm/
@@ -40,6 +40,7 @@ steps:
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
key: multi-modal-models-standard-3-llava-qwen2-vl
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -56,6 +57,7 @@ steps:
- label: "Multi-Modal Models (Standard) 4: other + whisper"
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -94,6 +96,7 @@ steps:
- label: Multi-Modal Accuracy Eval (Small Models) # 50min
key: multi-modal-accuracy-eval-small-models
timeout_in_minutes: 70
device: h200_18gb
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
- vllm/multimodal/
@@ -105,6 +108,7 @@ steps:
- label: Multi-Modal Models (Extended Generation 1)
key: multi-modal-models-extended-generation-1
optional: true
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal/generation
@@ -122,6 +126,7 @@ steps:
- label: Multi-Modal Models (Extended Generation 2)
key: multi-modal-models-extended-generation-2
optional: true
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal/generation
@@ -132,6 +137,7 @@ steps:
- label: Multi-Modal Models (Extended Generation 3)
key: multi-modal-models-extended-generation-3
optional: true
device: h200_18gb
source_file_dependencies:
- vllm/
- tests/models/multimodal/generation
+3
View File
@@ -5,6 +5,7 @@ steps:
- label: PyTorch Compilation Unit Tests
key: pytorch-compilation-unit-tests
timeout_in_minutes: 10
device: h200_18gb
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
@@ -78,6 +79,7 @@ steps:
- label: PyTorch Compilation Passes Unit Tests
key: pytorch-compilation-passes-unit-tests
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
@@ -110,6 +112,7 @@ steps:
- label: PyTorch Fullgraph Smoke Test
key: pytorch-fullgraph-smoke-test
timeout_in_minutes: 35
device: h200_18gb
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
+2
View File
@@ -5,6 +5,7 @@ steps:
- label: Quantization
key: quantization
timeout_in_minutes: 90
device: h200_18gb
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
@@ -42,6 +43,7 @@ steps:
- label: Quantized Models Test
key: quantized-models-test
timeout_in_minutes: 60
device: h200_18gb
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/models/quantization
+1
View File
@@ -5,6 +5,7 @@ steps:
- label: Samplers Test
key: samplers-test
timeout_in_minutes: 75
device: h200_18gb
source_file_dependencies:
- vllm/model_executor/layers
- vllm/sampling_metadata.py
+3 -3
View File
@@ -16,7 +16,7 @@ steps:
- label: Spec Decode Eagle Nightly B200
key: spec-decode-eagle-nightly-b200
timeout_in_minutes: 30
device: b200
device: b200-k8s
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
@@ -40,7 +40,7 @@ steps:
- label: Spec Decode Speculators + MTP Nightly B200
key: spec-decode-speculators-mtp-nightly-b200
timeout_in_minutes: 30
device: b200
device: b200-k8s
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
@@ -100,7 +100,7 @@ steps:
- label: Spec Decode MTP hybrid (B200)
timeout_in_minutes: 30
device: b200
device: b200-k8s
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
+2 -2
View File
@@ -31,8 +31,8 @@
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
/vllm/entrypoints/mcp @heheda12345
/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
/vllm/entrypoints/openai/realtime @njhill
/vllm/entrypoints/openai/speech_to_text @NickLucche
/vllm/entrypoints/speech_to_text/realtime @njhill
/vllm/entrypoints/speech_to_text @NickLucche
/vllm/entrypoints/pooling @noooop
/vllm/entrypoints/sagemaker @DarkLight1337
/vllm/entrypoints/serve @njhill
+18
View File
@@ -109,6 +109,24 @@ else()
set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
endif()
#
# spinloop extension (pure CXX; must stay above the non-CUDA device branch so
# CPU builds define the target before the early return)
#
set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp")
set(SPINLOOP_COMPILE_FLAGS "")
if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64")
list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx")
endif()
define_extension_target(
spinloop
DESTINATION vllm
LANGUAGE CXX
SOURCES ${VLLM_SPINLOOP_EXT_SRC}
COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS}
USE_SABI 3.11
WITH_SOABI)
#
# Forward the non-CUDA device extensions to external CMake scripts.
#
+75 -60
View File
@@ -21,28 +21,33 @@ namespace vllm {
// NOTE Be EXTRA careful with raw_kv_scalar_t, for __half and __nv_bfloat16 it's
// using u16 as the backing type.
template <typename qk_t, bool IS_NEOX, typename raw_kv_scalar_t,
typename cache_t, Fp8KVCacheDataType kv_dt>
template <typename qk_t, typename cos_sin_t, bool IS_NEOX,
typename raw_kv_scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
__global__ void concat_and_cache_mla_rope_fused_kernel(
const int64_t* __restrict__ positions, // [num_tokens]
qk_t* __restrict__ q_pe, // [num_tokens, num_q_heads, rot_dim]
qk_t* __restrict__ k_pe, // [num_tokens, rot_dim]
const qk_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
const qk_t* __restrict__ rope_cos_sin_cache, // [max_position, 2,
// rot_dim // 2]
const cos_sin_t* __restrict__ rope_cos_sin_cache, // [max_position, 2,
// rot_dim // 2]
const int rot_dim, const int64_t q_pe_stride_token,
const int64_t q_pe_stride_head, const int64_t k_pe_stride,
const int64_t kv_c_stride, const int num_q_heads,
cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank +
// rot_dim)]
const int64_t* __restrict__ kv_cache_slot_mapping, // [num_tokens]
const int64_t* __restrict__ slot_mapping, // [num_tokens]
const int block_stride, const int entry_stride, const int kv_lora_rank,
const int block_size, const float* kv_cache_quant_scale) {
// Each thread block is responsible for one token.
const int64_t token_idx = blockIdx.x;
const int64_t slot_idx = slot_mapping[token_idx];
// NOTE: slot_idx can be -1 if the token is padded
if (slot_idx < 0) {
return;
}
const int64_t pos = positions[token_idx];
const qk_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
const cos_sin_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
const int embed_dim = rot_dim / 2;
@@ -54,8 +59,8 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
// NOTE: Would be nice to have interleaved sin/cos so we could just load
// both at the same time.
qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim));
qk_t* q_pe_head_ptr =
q_pe + token_idx * q_pe_stride_token + head_idx * q_pe_stride_head;
@@ -81,21 +86,15 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
q_pe_head_ptr[pair_idx_y] = y_dst;
}
const int64_t slot_idx = kv_cache_slot_mapping[token_idx];
const int64_t block_idx = slot_idx / block_size;
const int64_t entry_idx = slot_idx % block_size;
// NOTE: slot_idx can be -1 if the token is padded
if (slot_idx < 0) {
return;
}
// K with 1 HEAD
for (int i = threadIdx.x; i < embed_dim; i += blockDim.x) {
int pair_idx = i;
qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim));
qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride;
@@ -165,36 +164,43 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
} // namespace vllm
#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
do { \
VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
using qk_t = scalar_t; \
if (rope_is_neox) { \
vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, true, RAW_KV_T, \
CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} else { \
vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, false, RAW_KV_T, \
CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} \
}); \
#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
do { \
VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
using qk_t = scalar_t; \
VLLM_DISPATCH_FLOATING_TYPES( \
rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \
[&] { \
using cos_sin_t = scalar_t; \
if (rope_is_neox) { \
vllm::concat_and_cache_mla_rope_fused_kernel< \
qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
kv_c_stride, num_q_heads, \
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} else { \
vllm::concat_and_cache_mla_rope_fused_kernel< \
qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
<<<grid, block, 0, stream>>>( \
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
kv_c_stride, num_q_heads, \
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
slot_mapping.data_ptr<int64_t>(), block_stride, \
entry_stride, kv_lora_rank, block_size, \
kv_cache_quant_scale.data_ptr<float>()); \
} \
}); \
}); \
} while (false)
// Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
@@ -208,43 +214,52 @@ void concat_and_cache_mla_rope_fused(
torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
bool rope_is_neox,
torch::Tensor&
kv_cache_slot_mapping, // [num_tokens] or [num_actual_tokens]
torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
torch::Tensor&
kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
const int64_t num_tokens = q_pe.size(0);
// NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from
// slot_mapping.size(0) because of padding for CUDA graphs.
// In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
// both include padding.
// In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
// since key includes padding for CUDA graphs, while slot_mapping does not.
// In this case, slot_mapping.size(0) represents the actual number of tokens
// before padding.
// For compatibility with both cases, we use slot_mapping.size(0) as the
// number of tokens.
int num_tokens = slot_mapping.size(0);
int num_padded_tokens = q_pe.size(0);
TORCH_CHECK_GE(num_padded_tokens, num_tokens);
const int num_q_heads = q_pe.size(1);
const int rot_dim = q_pe.size(2);
const int kv_lora_rank = kv_c.size(1);
TORCH_CHECK(positions.size(0) >=
num_tokens); // CUDA Graphs might pad this for us
TORCH_CHECK_EQ(positions.size(0), num_padded_tokens);
TORCH_CHECK_EQ(positions.dim(), 1);
TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
TORCH_CHECK_EQ(q_pe.size(0), num_tokens);
TORCH_CHECK_EQ(q_pe.dim(), 3);
TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens);
TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
TORCH_CHECK_EQ(q_pe.dim(), 3);
TORCH_CHECK_EQ(k_pe.size(0), num_tokens);
TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
TORCH_CHECK_EQ(k_pe.dim(), 2);
TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens);
TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
TORCH_CHECK_EQ(kv_c.size(0), num_tokens);
TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
TORCH_CHECK_EQ(kv_c.dim(), 2);
TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens);
TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
TORCH_CHECK_EQ(rope_cos_sin_cache.scalar_type(), q_pe.scalar_type());
TORCH_CHECK_EQ(kv_cache_slot_mapping.size(0), num_tokens);
TORCH_CHECK_EQ(kv_cache_slot_mapping.scalar_type(), c10::ScalarType::Long);
TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens);
TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long);
TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
TORCH_CHECK_EQ(kv_cache.dim(), 3);
-1
View File
@@ -1,7 +1,6 @@
#pragma once
#include <cstdlib>
#include <string>
#include <cctype>
namespace vllm {
@@ -1,5 +1,6 @@
#include "scaled_mm_kernels.hpp"
#include "scaled_mm_sm100_fp8_dispatch.cuh"
#include "core/batch_invariant.hpp"
namespace vllm {
@@ -13,9 +14,17 @@ void cutlass_scaled_mm_sm100_fp8(
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
"currently bias dtype must match output dtype ",
out.scalar_type());
if (vllm_is_batch_invariant()) {
return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue<true>(
out, a, b, a_scales, b_scales, *bias);
}
return cutlass_scaled_mm_sm100_fp8_epilogue<true>(out, a, b, a_scales,
b_scales, *bias);
} else {
if (vllm_is_batch_invariant()) {
return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue<false>(
out, a, b, a_scales, b_scales);
}
return cutlass_scaled_mm_sm100_fp8_epilogue<false>(out, a, b, a_scales,
b_scales);
}
@@ -294,6 +294,34 @@ inline void cutlass_gemm_sm100_fp8_dispatch(
}
}
template <typename InType, typename OutType, bool EnableBias,
typename... EpilogueArgs>
inline void cutlass_gemm_sm100_fp8_batch_invariant_dispatch(
torch::stable::Tensor& out, torch::stable::Tensor const& a,
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
STD_TORCH_CHECK(a.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
STD_TORCH_CHECK(b.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
using Cutlass3xGemmM64SwapAB =
typename sm100_fp8_config_M64_swap_ab<InType, OutType,
EnableBias>::Cutlass3xGemm;
using Cutlass3xGemmM64 =
typename sm100_fp8_config_M64<InType, OutType, EnableBias>::Cutlass3xGemm;
// keep the CUTLASS config independent of M for batch invariance
uint32_t const k = a.size(1);
if (k < 4096) {
return cutlass_gemm_caller_sm100_fp8<Cutlass3xGemmM64>(
out, a, b, a_scales, b_scales, std::forward<EpilogueArgs>(args)...);
}
return cutlass_gemm_caller_sm100_fp8<Cutlass3xGemmM64SwapAB>(
out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
}
template <bool EnableBias, typename... EpilogueArgs>
void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
torch::stable::Tensor const& a,
@@ -320,4 +348,28 @@ void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
}
}
template <bool EnableBias, typename... EpilogueArgs>
void cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue(
torch::stable::Tensor& out, torch::stable::Tensor const& a,
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) {
STD_TORCH_CHECK(a.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
STD_TORCH_CHECK(b.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
return cutlass_gemm_sm100_fp8_batch_invariant_dispatch<
cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>(
out, a, b, a_scales, b_scales,
std::forward<EpilogueArgs>(epilogue_args)...);
} else {
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
return cutlass_gemm_sm100_fp8_batch_invariant_dispatch<
cutlass::float_e4m3_t, cutlass::half_t, EnableBias>(
out, a, b, a_scales, b_scales,
std::forward<EpilogueArgs>(epilogue_args)...);
}
}
} // namespace vllm
@@ -1,5 +1,6 @@
#include "scaled_mm_kernels.hpp"
#include "scaled_mm_sm120_fp8_dispatch.cuh"
#include "core/batch_invariant.hpp"
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
namespace vllm {
@@ -14,9 +15,17 @@ void cutlass_scaled_mm_sm120_fp8(
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
"currently bias dtype must match output dtype ",
out.scalar_type());
if (vllm_is_batch_invariant()) {
return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
c3x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
}
return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
out, a, b, a_scales, b_scales, *bias);
} else {
if (vllm_is_batch_invariant()) {
return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
c3x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
}
return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
out, a, b, a_scales, b_scales);
}
@@ -179,6 +179,26 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::stable::Tensor& out,
out, a, b, std::forward<EpilogueArgs>(args)...);
}
template <typename InType, typename OutType,
template <typename, typename, typename> typename Epilogue,
typename... EpilogueArgs>
inline void cutlass_gemm_sm120_fp8_batch_invariant_dispatch(
torch::stable::Tensor& out, torch::stable::Tensor const& a,
torch::stable::Tensor const& b, EpilogueArgs&&... args) {
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
STD_TORCH_CHECK(a.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
STD_TORCH_CHECK(b.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
using Cutlass3xGemmM64 =
typename sm120_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
// keep the CUTLASS config independent of M for batch invariance
return cutlass_gemm_caller<Cutlass3xGemmM64>(
out, a, b, std::forward<EpilogueArgs>(args)...);
}
template <template <typename, typename, typename> typename Epilogue,
typename... EpilogueArgs>
void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
@@ -202,4 +222,26 @@ void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
}
}
template <template <typename, typename, typename> typename Epilogue,
typename... EpilogueArgs>
void cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue(
torch::stable::Tensor& out, torch::stable::Tensor const& a,
torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
STD_TORCH_CHECK(a.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
STD_TORCH_CHECK(b.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
return cutlass_gemm_sm120_fp8_batch_invariant_dispatch<
cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
} else {
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
return cutlass_gemm_sm120_fp8_batch_invariant_dispatch<
cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
}
}
} // namespace vllm
@@ -1,5 +1,6 @@
#include "scaled_mm_kernels.hpp"
#include "scaled_mm_sm90_fp8_dispatch.cuh"
#include "core/batch_invariant.hpp"
namespace vllm {
@@ -13,9 +14,17 @@ void cutlass_scaled_mm_sm90_fp8(
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
"currently bias dtype must match output dtype ",
out.scalar_type());
if (vllm_is_batch_invariant()) {
return cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue<true>(
out, a, b, a_scales, b_scales, *bias);
}
return cutlass_scaled_mm_sm90_fp8_epilogue<true>(out, a, b, a_scales,
b_scales, *bias);
} else {
if (vllm_is_batch_invariant()) {
return cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue<false>(
out, a, b, a_scales, b_scales);
}
return cutlass_scaled_mm_sm90_fp8_epilogue<false>(out, a, b, a_scales,
b_scales);
}
@@ -349,6 +349,35 @@ inline void cutlass_gemm_sm90_fp8_dispatch(
}
}
template <typename InType, typename OutType, bool EnableBias,
typename... EpilogueArgs>
inline void cutlass_gemm_sm90_fp8_batch_invariant_dispatch(
torch::stable::Tensor& out, torch::stable::Tensor const& a,
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
STD_TORCH_CHECK(a.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
STD_TORCH_CHECK(b.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
using Cutlass3xGemmM64_N1280 =
typename sm90_fp8_config_M64_N1280<InType, OutType,
EnableBias>::Cutlass3xGemm;
using Cutlass3xGemmM64_N8192 =
typename sm90_fp8_config_M64_N8192<InType, OutType,
EnableBias>::Cutlass3xGemm;
// keep the CUTLASS config independent of M for batch invariance
uint32_t const n = b.size(1);
if (n <= 1280) {
return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM64_N1280>(
out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
}
return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM64_N8192>(
out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
}
template <bool EnableBias, typename... EpilogueArgs>
void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
torch::stable::Tensor const& a,
@@ -375,4 +404,28 @@ void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
}
}
template <bool EnableBias, typename... EpilogueArgs>
void cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue(
torch::stable::Tensor& out, torch::stable::Tensor const& a,
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) {
STD_TORCH_CHECK(a.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
STD_TORCH_CHECK(b.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
return cutlass_gemm_sm90_fp8_batch_invariant_dispatch<
cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>(
out, a, b, a_scales, b_scales,
std::forward<EpilogueArgs>(epilogue_args)...);
} else {
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
return cutlass_gemm_sm90_fp8_batch_invariant_dispatch<
cutlass::float_e4m3_t, cutlass::half_t, EnableBias>(
out, a, b, a_scales, b_scales,
std::forward<EpilogueArgs>(epilogue_args)...);
}
}
} // namespace vllm
@@ -9,6 +9,7 @@
#include "scaled_mm_c2x_sm89_fp8_dispatch.cuh"
#include "scaled_mm_c2x_sm89_int8_dispatch.cuh"
#include "core/batch_invariant.hpp"
#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp"
using namespace vllm;
@@ -191,9 +192,17 @@ void cutlass_scaled_mm_sm89(torch::stable::Tensor& out,
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
"currently bias dtype must match output dtype ",
out.scalar_type());
if (vllm_is_batch_invariant()) {
return cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue<
c2x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
}
return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBias>(
out, a, b, a_scales, b_scales, *bias);
} else {
if (vllm_is_batch_invariant()) {
return cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue<
c2x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
}
return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogue>(
out, a, b, a_scales, b_scales);
}
@@ -381,4 +381,43 @@ inline void cutlass_gemm_sm89_fp8_dispatch(torch::stable::Tensor& out,
}
}
template <typename InType, typename OutType,
template <typename, typename> typename Epilogue,
typename... EpilogueArgs>
inline void cutlass_gemm_sm89_fp8_batch_invariant_dispatch(
torch::stable::Tensor& out, torch::stable::Tensor const& a,
torch::stable::Tensor const& b, EpilogueArgs&&... args) {
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
STD_TORCH_CHECK(a.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
STD_TORCH_CHECK(b.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
// keep the CUTLASS config independent of M for batch invariance
return sm89_fp8_config_M64::dispatch<InType, OutType, Epilogue>(
out, a, b, std::forward<EpilogueArgs>(args)...);
}
template <template <typename, typename> typename Epilogue,
typename... EpilogueArgs>
void cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue(
torch::stable::Tensor& out, torch::stable::Tensor const& a,
torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
STD_TORCH_CHECK(a.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
STD_TORCH_CHECK(b.scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn);
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
} else {
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
}
}
} // namespace vllm
+204
View File
@@ -0,0 +1,204 @@
#include <Python.h>
extern "C" {
#include <stdbool.h>
#include <time.h>
#if defined(__i386__) || defined(__x86_64__)
#include <cpuid.h>
#include <mwaitxintrin.h>
#endif
#if defined(CLOCK_MONOTONIC_RAW)
#define TIMEOUT_CLOCK CLOCK_MONOTONIC_RAW
#else
#define TIMEOUT_CLOCK CLOCK_MONOTONIC
#endif
#define CPU_SUPPORT_NONE 0
#define CPU_SUPPORT_MONITORX 1
#define MWAITX_DEFAULT_TIMEOUT_CYCLES 1000000
typedef struct {
unsigned int cpu_support;
unsigned int max_monitor_line_size;
} spinloop_state_t;
static void determine_cpu_support(spinloop_state_t* state) {
state->cpu_support = CPU_SUPPORT_NONE;
state->max_monitor_line_size = 0;
#if defined(__i386__) || defined(__x86_64__)
unsigned int eax, ebx, ecx, edx;
if (__get_cpuid(0, &eax, &ebx, &ecx, &edx) == 1) {
// AMD CPU (possible monitorx/mwaitx support)
if (ebx == 0x68747541 && edx == 0x69746e65 && ecx == 0x444d4163) {
if (__get_cpuid(0x80000000, &eax, &ebx, &ecx, &edx) == 1 &&
eax >= 0x80000001 &&
__get_cpuid(0x80000001, &eax, &ebx, &ecx, &edx) == 1) {
if ((ecx & (1 << 29)) != 0) {
state->cpu_support = CPU_SUPPORT_MONITORX;
}
}
}
}
if (state->cpu_support == CPU_SUPPORT_MONITORX) {
if (__get_cpuid(5, &eax, &ebx, &ecx, &edx) == 1) {
state->max_monitor_line_size = ebx & 0xff;
}
}
#endif
}
static PyObject* method_spinloop(PyObject* self, PyObject* args,
PyObject* kwargs) {
Py_buffer buffer;
PyObject* callback;
double timeout = 0.;
spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(self);
if (state == NULL) {
PyErr_SetString(PyExc_TypeError, "Failed to retrieve module state!");
return NULL;
}
static const char* keywords[] = {"buffer", "callback", "timeout", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*O|d", (char**)keywords,
&buffer, &callback, &timeout)) {
return NULL;
}
if (!PyCallable_Check(callback)) {
PyErr_SetString(PyExc_TypeError, "callback parameter must be callable!");
PyBuffer_Release(&buffer);
return NULL;
}
struct timespec t_start;
if (clock_gettime(TIMEOUT_CLOCK, &t_start) != 0) {
PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
PyBuffer_Release(&buffer);
return NULL;
}
bool result = false;
bool error = false;
bool have_timeout = (timeout > 1e-9);
unsigned int iteration = 0;
const bool buffer_qualifies = (buffer.len <= state->max_monitor_line_size);
while (true) {
PyObject* res = PyObject_CallNoArgs(callback);
if (res == NULL) {
error = true;
break;
}
int ok = (res == Py_True);
Py_DECREF(res);
if (ok) {
result = true;
break;
}
// Check timeout at most every 16 iterations to avoid clock_gettime and
// comparison cost
if (have_timeout && (iteration & 15u) == 0) {
struct timespec t_now;
if (clock_gettime(TIMEOUT_CLOCK, &t_now) != 0) {
PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
error = true;
break;
}
const double elapsed = (double)(t_now.tv_sec - t_start.tv_sec) +
(t_now.tv_nsec - t_start.tv_nsec) * 1e-9;
if (elapsed >= timeout) {
result = false;
break;
}
}
++iteration;
#if defined(__i386__) || defined(__x86_64__)
// monitorx + mwaitx with qualified buffer
if (buffer_qualifies && state->cpu_support == CPU_SUPPORT_MONITORX) {
_mm_monitorx(buffer.buf, 0, 0);
// Check once more in case the buffer has been modified while we were
// arming the monitor hardware
res = PyObject_CallNoArgs(callback);
if (res == NULL) {
error = true;
break;
}
ok = (res == Py_True);
Py_DECREF(res);
if (ok) {
result = true;
break;
}
// Run mwaitx with enabled timeout (bit 1). The actual timeout value
// is not very important, we just want to ensure we don't lock up
// here for too long.
Py_BEGIN_ALLOW_THREADS _mm_mwaitx((1 << 1), 0,
MWAITX_DEFAULT_TIMEOUT_CYCLES);
Py_END_ALLOW_THREADS
}
// Fallback: Busy poll
else {
#endif
// Give other threads a chance to be scheduled
Py_BEGIN_ALLOW_THREADS
#if defined(__i386__) || defined(__x86_64__)
__builtin_ia32_pause();
#elif defined(__aarch64__)
__asm__ volatile("yield" :: : "memory");
#endif
Py_END_ALLOW_THREADS
#if defined(__i386__) || defined(__x86_64__)
}
#endif
}
PyBuffer_Release(&buffer);
if (error) {
return NULL;
}
if (result) {
Py_RETURN_TRUE;
}
Py_RETURN_FALSE;
}
static PyMethodDef spinloop_methods[] = {
{"spinloop", (PyCFunction)method_spinloop, METH_VARARGS | METH_KEYWORDS,
"Wait for store with callback"},
{NULL, NULL, 0, NULL}};
static struct PyModuleDef spinloop_module = {
PyModuleDef_HEAD_INIT, "spinloop",
"Hardware-optimized spinloops for Python", sizeof(spinloop_state_t),
spinloop_methods};
PyMODINIT_FUNC PyInit_spinloop(void) {
PyObject* m = PyModule_Create(&spinloop_module);
if (m != NULL) {
spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(m);
if (state != NULL) {
determine_cpu_support(state);
}
}
return m;
}
} // extern "C"
+60 -6
View File
@@ -39,7 +39,8 @@ th {
| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
| Custom | ✅ | ✅ | Local file: `data.jsonl` |
| Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
| Custom Audio | ✅ | ✅ | Local file: `audio_data.jsonl` |
| Custom Image | ✅ | ✅ | Local file: `image_data.jsonl` |
Legend:
@@ -142,7 +143,7 @@ Example output: ![Dataset Statistics](../assets/contributing/vllm_bench_serve_da
#### Custom Dataset
If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. At inference time, use the option `--dataset-name custom`. Your data needs to be in the `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
```json
{"prompt": "What is the capital of India?"}
@@ -173,9 +174,62 @@ vllm bench serve --port 9001 --save-result --save-detailed \
You can skip applying chat template if your data already has it by using `--custom-skip-chat-template`.
#### Custom multimodal dataset
#### Custom Audio Dataset
If the multimodal dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomMMDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" and "image_files" field per entry, e.g., `mm_data.jsonl`:
If the audio dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomAudioDataset`. At inference time, use the option `--dataset-name custom_audio`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "audio" fields per entry, e.g., `audio_data.jsonl`:
```json
{"prompt": "What does this audio say?", "audio": "/path/to/audio_1.wav"}
{"prompt": "Transcribe the audio.", "audio": "/path/to/audio_2.wav"}
```
- **Supported models:** The `CustomAudioDataset` class supports two types of audio models: ASR models (e.g. Whisper) which do not require a "prompt" field; and multimodal audio-text chat models (e.g. Qwen2-Audio). Since these model types require different arguments at inference, we are giving two examples.
- **Example 1: Whisper**
Whisper is a dedicated ASR encoder-decoder model, so it uses `--backend openai-audio` and `--endpoint /v1/audio/transcriptions`.
```bash
# start server
vllm serve openai/whisper-tiny
```
```bash
vllm bench serve \
--model openai/whisper-tiny \
--backend openai-audio \
--endpoint /v1/audio/transcriptions \
--dataset-name custom_audio \
--dataset-path audio_data.jsonl \
--no-oversample \
--custom-output-len 256 \
--save-result \
--save-detailed \
--result-filename whisper_bench.json
```
- **Example 2: Qwen2-Audio**
Qwen2-Audio is a multimodal chat model that can do ASR and speech analysis, so it uses `--backend openai-chat`, and `--endpoint /v1/chat/completions`. It also requires `--enable-multimodal-chat` to enable multimodal chat transformation.
```bash
vllm bench serve \
--model Qwen/Qwen2-Audio-7B-Instruct \
--backend openai-chat \
--endpoint /v1/chat/completions \
--dataset-name custom_audio \
--dataset-path audio_data.jsonl \
--no-oversample \
--custom-output-len 256 \
--enable-multimodal-chat \
--save-result \
--save-detailed \
--result-filename qwen_bench.json
```
#### Custom Image Dataset
If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
```json
{"prompt": "How many animals are present in the given image?", "image_files": ["/path/to/image/folder/horsepony.jpg"]}
@@ -193,8 +247,8 @@ vllm bench serve--save-result --save-detailed \
--backend openai-chat \
--model Qwen/Qwen2-VL-7B-Instruct \
--endpoint /v1/chat/completions \
--dataset-name custom_mm \
--dataset-path <path-to-your-mm-data-jsonl> \
--dataset-name custom_image \
--dataset-path <path-to-your-image-data-jsonl> \
--allowed-local-media-path /path/to/image/folder
```
+8 -2
View File
@@ -126,9 +126,15 @@ python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
- Set when prefiller and decoder are on different machines
- Connection info is passed via KVTransferParams from prefiller to decoder for handshake
- `VLLM_NIXL_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefillers KV cache for a particular request. (Optional)
- `kv_lease_duration` (via `kv_connector_extra_config`): Lease duration (in seconds) for the prefiller's KV cache blocks. (Optional)
- Default: 30
- When a prefill request finishes, its KV blocks are held for this duration waiting for the decoder to read them. While the request is queued on the decoder, periodic heartbeats automatically extend the lease. If neither a heartbeat nor a read notification arrives before the lease expires, the blocks are freed. The heartbeat interval and extension amount are derived automatically from this value.
- Example: `--kv-transfer-config '{"kv_connector_extra_config": {"kv_lease_duration": 60}}'`
- `decoder_kv_blocks_ttl` (via `kv_connector_extra_config`): TTL (in seconds) for KV blocks cached on the decoder in bidirectional transfer mode. (Optional)
- Default: 480
- If a request is aborted and the decoder has not yet read the KV-cache blocks through the nixl channel, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely.
- In bidirectional mode, the decoder caches KV blocks for multi-turn conversations. This TTL controls how long those blocks are held before being released. Unlike the prefiller lease, this TTL is not renewed via heartbeats.
- Example: `--kv-transfer-config '{"kv_connector_extra_config": {"decoder_kv_blocks_ttl": 600}}'`
## Multi-Instance Setup
+4 -3
View File
@@ -3,9 +3,10 @@
vLLM supports the following hardware platforms:
- [GPU](gpu.md)
- [NVIDIA CUDA](gpu.md#nvidia-cuda)
- [AMD ROCm](gpu.md#amd-rocm)
- [Intel XPU](gpu.md#intel-xpu)
- [NVIDIA CUDA](gpu.md)
- [AMD ROCm](gpu.md)
- [Intel XPU](gpu.md)
- [Apple Silicon](gpu.md) (via [vLLM-Metal](https://github.com/vllm-project/vllm-metal))
- [CPU](cpu.md)
- [Intel/AMD x86](cpu.md#intelamd-x86)
- [ARM AArch64](cpu.md#arm-aarch64)
@@ -0,0 +1,125 @@
<!-- markdownlint-disable MD041 -->
--8<-- [start:installation]
For GPU-accelerated inference on Apple Silicon, use [vLLM-Metal](https://github.com/vllm-project/vllm-metal), a community-maintained hardware plugin that uses MLX as the compute backend and provides native GPU acceleration via Apple's Metal framework.
vLLM-Metal works with MLX-optimized models from the [mlx-community](https://huggingface.co/mlx-community) organization on Hugging Face, which provides quantized versions of popular models optimized for Apple Silicon.
!!! tip
For installation and usage instructions, see the [Set up using vLLM-Metal](#set-up-using-vllm-metal) section below.
--8<-- [end:installation]
--8<-- [start:requirements]
- OS: macOS Sonoma or later
- Hardware: Apple Silicon
- Metal support enabled
!!! note
See the [Set up using vLLM-Metal](#set-up-using-vllm-metal) section below for installation instructions.
--8<-- [end:requirements]
--8<-- [start:set-up-using-python]
## Set up using vLLM-Metal
vLLM-Metal is distributed as a separate package that provides native GPU acceleration on Apple Silicon.
To install vLLM-Metal, follow the installation instructions in the [vLLM-Metal documentation](https://github.com/vllm-project/vllm-metal#installation).
The installation will:
1. Set up the appropriate Python environment
2. Install MLX and required dependencies
3. Install the vLLM-Metal package
After installation, you can start using vLLM with Metal GPU acceleration.
!!! tip
When using vLLM-Metal, use models from the [mlx-community](https://huggingface.co/mlx-community) on Hugging Face for best performance. These models are optimized for MLX and often include quantized versions (4-bit, 8-bit) that run efficiently on Apple Silicon.
Example model: `mlx-community/Qwen2.5-0.5B-Instruct-4bit`
### Using vLLM-Metal
After installation, vLLM-Metal provides an easy-to-use CLI for running an OpenAI-compatible API server:
```bash
# Activate the vLLM-Metal environment
source ~/.venv-vllm-metal/bin/activate
# Start the API server (specify your mlx-community model or it will use default)
vllm serve
```
Once the server is running, you have multiple options to interact with it:
#### Option 1: Interactive chat
Open a new terminal and start an interactive chat session:
```bash
source ~/.venv-vllm-metal/bin/activate
vllm chat
```
#### Option 2: API requests with curl
```bash
curl http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": "Hello!"}],
"max_tokens": 50
}'
```
#### Option 3: Python with OpenAI SDK
```python
from openai import OpenAI
client = OpenAI(
base_url="http://localhost:8000/v1",
api_key="dummy" # No auth required for local server
)
response = client.chat.completions.create(
model="mlx-community/Qwen2.5-0.5B-Instruct-4bit",
messages=[{"role": "user", "content": "Hello!"}]
)
print(response.choices[0].message.content)
```
For more details on the `vllm` CLI commands, see the [OpenAI-compatible server documentation](../../serving/openai_compatible_server.md).
--8<-- [end:set-up-using-python]
--8<-- [start:pre-built-wheels]
vLLM-Metal is installed via the vLLM-Metal package. See the [Set up using vLLM-Metal](#set-up-using-vllm-metal) section above.
--8<-- [end:pre-built-wheels]
--8<-- [start:build-wheel-from-source]
For build instructions from source, refer to the [vLLM-Metal documentation](https://github.com/vllm-project/vllm-metal#installation).
--8<-- [end:build-wheel-from-source]
--8<-- [start:pre-built-images]
--8<-- [end:pre-built-images]
--8<-- [start:build-image-from-source]
--8<-- [end:build-image-from-source]
--8<-- [start:supported-features]
vLLM-Metal provides:
- Native GPU acceleration using Metal
- MLX-based compute backend optimized for Apple Silicon
- OpenAI-compatible API server
- Support for popular model architectures
For specific feature support and limitations, refer to the [vLLM-Metal documentation](https://github.com/vllm-project/vllm-metal).
--8<-- [end:supported-features]
+32
View File
@@ -18,6 +18,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:installation"
=== "Apple Silicon"
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:installation"
## Requirements
- OS: Linux
@@ -38,6 +42,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:requirements"
=== "Apple Silicon"
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:requirements"
## Set up using Python
### Create a new Python environment
@@ -56,6 +64,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:set-up-using-python"
=== "Apple Silicon"
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:set-up-using-python"
### Pre-built wheels {#pre-built-wheels}
=== "NVIDIA CUDA"
@@ -70,6 +82,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:pre-built-wheels"
=== "Apple Silicon"
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:pre-built-wheels"
### Build wheel from source
=== "NVIDIA CUDA"
@@ -84,6 +100,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:build-wheel-from-source"
=== "Apple Silicon"
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:build-wheel-from-source"
## Set up using Docker
### Pre-built images
@@ -102,6 +122,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:pre-built-images"
=== "Apple Silicon"
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:pre-built-images"
--8<-- [end:pre-built-images]
### Build image from source
@@ -120,6 +144,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:build-image-from-source"
=== "Apple Silicon"
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:build-image-from-source"
--8<-- [end:build-image-from-source]
## Supported features
@@ -135,3 +163,7 @@ vLLM is a Python library that supports the following GPU variants. Select your G
=== "Intel XPU"
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:supported-features"
=== "Apple Silicon"
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:supported-features"
+15
View File
@@ -10,6 +10,9 @@ This guide will help you quickly get started with vLLM to perform:
- OS: Linux
- Python: 3.10 -- 3.13
!!! note
vLLM also works on macOS with [vLLM-Metal](https://github.com/vllm-project/vllm-metal) for Apple Silicon GPU acceleration. See the [GPU installation guide](installation/gpu.md) and select the "Apple Silicon" tab.
## Installation
=== "NVIDIA CUDA"
@@ -73,6 +76,18 @@ This guide will help you quickly get started with vLLM to perform:
!!! note
For more detailed instructions, including Docker, installing from source, and troubleshooting, please refer to the [vLLM on TPU documentation](https://docs.vllm.ai/projects/tpu/en/latest/).
=== "Apple Silicon (Mac)"
If you are using Apple Silicon Macs, you can use vLLM-Metal for GPU-accelerated inference via Apple's Metal framework.
Follow the installation instructions in the [vLLM-Metal documentation](https://github.com/vllm-project/vllm-metal#installation).
!!! note
vLLM-Metal uses MLX instead of PyTorch as the compute backend and requires MLX-optimized models from the [mlx-community](https://huggingface.co/mlx-community) on Hugging Face.
!!! tip
For more detailed instructions, please refer to the [GPU installation guide](installation/gpu.md) and select the "Apple Silicon" tab.
!!! note
For more detail and non-CUDA platforms, please refer to the [installation guide](installation/README.md) for specific instructions on how to install vLLM.
+6 -1
View File
@@ -598,7 +598,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B` | ✅︎ | ✅︎ |
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
@@ -666,6 +666,11 @@ Some models are supported only via the [Transformers modeling backend](#transfor
!!! note
For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
!!! note
To use `allenai/MolmoWeb-4B` or `allenai/MolmoWeb-8B`, serve the checkpoint
with the Molmo2 architecture and disable multimodal-prefix attention:
`--hf-overrides '{"architectures": ["Molmo2ForConditionalGeneration"], "is_mm_prefix_lm": false}'`.
!!! note
`Moondream3ForCausalLM` uses task-specific prompt templates for `query`
and `caption`. The native `detect` and `point` skills require custom
+37
View File
@@ -333,6 +333,43 @@ Most cache paths default to subdirectories under a single root. Changing `VLLM_C
- **Do not copy cache contents from untrusted sources.** If you distribute cache artifacts between environments, ensure they originate from a trusted build pipeline.
- **Container deployments:** If mounting cache directories into containers, ensure the volume source is trusted.
## FIPS Compatibility
FIPS compliance depends on many factors, so a vLLM deployment is not automatically FIPS compliant. Recent changes have improved vLLM's *tolerance* of FIPS-enabled hosts — that is, avoiding crashes when non-approved algorithms are blocked — but tolerance is not the same as compliance. Whether a deployment satisfies FIPS requirements depends on the host operating system, the OpenSSL provider backing Python's `hashlib` and `ssl` modules, and which optional dependencies are installed.
### FIPS-relevant configuration
Operators running vLLM on FIPS-enabled hosts should select FIPS-approved algorithms via the following knobs:
- **Multimodal input hashing** — `VLLM_MM_HASHER_ALGORITHM` defaults to `blake3`, which is not FIPS-approved. Set it to `sha256` or `sha512` in FIPS-enabled environments.
- **Prefix-cache hashing** — set `--prefix-caching-hash-algo` (config field `prefix_caching_hash_algo`) to `sha256` or `sha256_cbor`. The `xxhash` and `xxhash_cbor` options are not FIPS-approved.
- **TLS ciphers** — use `--ssl-ciphers` to restrict the API server's TLS handshake to FIPS-approved cipher suites that match your environment's policy.
### Automatic fallback for non-security MD5 use
vLLM uses MD5 in a few places to derive non-security cache keys (for example, configuration hashes). These call sites pass `usedforsecurity=False` and additionally fall back to SHA-256 when the underlying OpenSSL provider refuses MD5 outright (see `safe_hash()` in `vllm/utils/hashing.py`). No user action is required; this behavior is documented so that auditors and security reviewers can identify the MD5 references and understand their purpose.
### Dependencies that provide non-FIPS hash implementations
Some dependencies expose hash implementations that are not FIPS-approved. vLLM only invokes them when the corresponding algorithm is selected, but operators with strict cryptographic controls may want to ensure the code paths are not exercised — and, where policy requires, that the packages themselves are absent:
- `blake3` — currently listed in `requirements/common.txt`, so a standard install pulls it in. It is imported lazily and only used when `VLLM_MM_HASHER_ALGORITHM=blake3` (the default). Setting `VLLM_MM_HASHER_ALGORITHM` to `sha256` or `sha512` is sufficient to keep the non-FIPS code path dormant. If your policy additionally forbids the package being present, uninstall it after `pip install` (`pip uninstall blake3`); vLLM will continue to function as long as `VLLM_MM_HASHER_ALGORITHM` is set to a non-blake3 value.
- `xxhash` — a true optional dependency (not in `requirements/common.txt`). It is only imported when an `xxhash`-based prefix-cache algorithm is selected. Leave it uninstalled and select a `sha256`-based prefix-cache algorithm.
### Beyond hashing: other FIPS considerations
Hashing is the area where vLLM has explicit FIPS-aware code, but a FIPS-compliant deployment depends on several factors that sit outside vLLM itself. Operators should evaluate the following with their platform and security teams:
- **Host crypto provider.** Python's `hashlib` and `ssl` modules are FIPS-aware only when Python is linked against a FIPS-validated OpenSSL (or equivalent) provider supplied by the host OS. vLLM inherits whatever provider the host configures — it does not bundle one.
- **API server TLS.** TLS termination for the OpenAI-compatible API server uses the host's OpenSSL via Python's `ssl` module. Restrict the cipher suite with `--ssl-ciphers` to match your environment's FIPS policy, and ensure server certificates are issued with FIPS-approved algorithms and key sizes.
- **Outbound HTTPS.** Model and asset downloads (for example, via `huggingface_hub`) use the same host TLS stack. The same provider/cipher considerations apply.
- **Inter-node communication is unencrypted by default.** As described in [Inter-Node Communication](#inter-node-communication), PyTorch Distributed, KV-cache transfer, and data-parallel channels do not encrypt traffic. FIPS environments that require FIPS-approved cryptography for data in transit must provide that protection externally — for example, via an mTLS sidecar or IPsec terminated by a FIPS-validated module — since vLLM's internal channels cannot satisfy the requirement on their own. Network isolation alone is not cryptography and does not meet a "FIPS-approved cryptography for data in transit" requirement, though it remains a useful defense-in-depth measure.
- **Dependencies that bundle their own OpenSSL.** Some Python wheels statically link OpenSSL builds that fail the kernel FIPS self-test on FIPS-enabled hosts (`FATAL FIPS SELFTEST FAILURE`). `opencv-python-headless` is a known example; other manylinux wheels may behave similarly. Audit your installed wheels for bundled crypto libraries when troubleshooting FIPS startup failures.
- **Accelerator and ML libraries.** PyTorch, CUDA, cuDNN, NCCL, and similar components have their own crypto and FIPS posture independent of vLLM. NVIDIA publishes FIPS-validated builds for some libraries; vLLM does not pin to those builds, so selecting and validating them is the operator's responsibility.
- **What is *not* a FIPS concern in vLLM.** Random number generation used for token sampling (Python/NumPy/PyTorch RNGs) is not a cryptographic use and is out of scope for FIPS. Pickled cache artifacts are a separate security concern covered under [Cache Directory Security](#cache-directory-security).
In short: the configuration knobs above let vLLM avoid non-approved algorithms, and the automatic fallbacks let it run without crashing on FIPS-enabled hosts. End-to-end FIPS compliance, however, is a property of the full deployment — host OS, crypto provider, transitive dependencies, and network architecture — not of vLLM alone.
## Reporting Security Vulnerabilities
If you believe you have found a security vulnerability in vLLM, please report it following the project's security policy. For more information on how to report security issues and the project's security policy, please see the [vLLM Security Policy](https://github.com/vllm-project/vllm/blob/main/SECURITY.md).
+2 -1
View File
@@ -123,7 +123,8 @@ python = "./.venv"
extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*",
"benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
"examples/pooling/token_embed/*", "tests/models/language/pooling/*",
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py",
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
"tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py",
"docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html",
"tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*"]
ignore-hidden = false
+3
View File
@@ -686,6 +686,7 @@ class precompiled_wheel_utils:
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
"vllm/cumem_allocator.abi3.so",
"vllm/spinloop.abi3.so",
# ROCm-specific libraries
"vllm/_rocm_C.abi3.so",
]
@@ -993,6 +994,8 @@ if _is_cuda() or _is_hip():
# copying the relevant .py files from the source repository.
ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))
ext_modules.append(CMakeExtension(name="vllm.spinloop"))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
@@ -0,0 +1,166 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for the DoubleQuant fan-out variants registered by
``RocmAiterRMSNormQuantFusionPass``.
Both variants target a 1-to-2 fan-out where one ``rms_norm`` output feeds
two distinct ``rocm_aiter_group_fp8_quant`` consumers and rewrite it into
two independent fused ``rms_norm + group_fp8_quant`` ops:
* ``DoubleAiterRMSFp8GroupQuantPattern`` matches the un-viewed shape
(e.g. Kimi-K2.5 / DSR1).
* ``DoubleAiterRMSFp8GroupQuantViewPattern`` (this PR) is the view-tolerant
sibling that additionally matches the
``rms_norm -> view -> group_fp8_quant`` shape that DSv3.2's MLA indexer
q_c norm exposes through ``Fp8BlockScaledMMLinearKernel.apply_weights``'s
2D-flatten boilerplate.
"""
import pytest
import torch
import vllm.config
from tests.compile.backend import TestBackend
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
from vllm.config import (
CompilationConfig,
CompilationMode,
ModelConfig,
PassConfig,
VllmConfig,
)
EPS = 1e-5
HIDDEN_SIZE = 256
GROUP_SIZE = 128
class _NoViewDoubleQuantModel(torch.nn.Module):
"""``rms_norm -> 2x group_fp8_quant`` fan-out (Kimi-K2.5 / DSR1 shape)."""
def __init__(self) -> None:
super().__init__()
self.weight = torch.nn.Parameter(torch.ones(HIDDEN_SIZE, dtype=torch.bfloat16))
def forward(
self, x: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
# avoid graph input being a direct arg to a matched pattern node
x = torch.relu(x)
rms = torch.ops.vllm_ir.rms_norm(x, self.weight, EPS)
q1, s1 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(rms, GROUP_SIZE)
q2, s2 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(rms, GROUP_SIZE)
return q1, s1, q2, s2
class _ViewDoubleQuantModel(torch.nn.Module):
"""``rms_norm -> view -> 2x group_fp8_quant`` fan-out (DSv3.2 shape).
Reproduces the FX-graph shape produced by ``Fp8BlockScaledMMLinearKernel``'s
2D-flatten before the FP8 group quant op.
"""
def __init__(self) -> None:
super().__init__()
self.weight = torch.nn.Parameter(torch.ones(HIDDEN_SIZE, dtype=torch.bfloat16))
def forward(
self, x: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
x = torch.relu(x)
rms = torch.ops.vllm_ir.rms_norm(x, self.weight, EPS)
view = rms.view(-1, rms.shape[-1])
q1, s1 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(view, GROUP_SIZE)
q2, s2 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(view, GROUP_SIZE)
return q1, s1, q2, s2
@pytest.mark.parametrize(
"model_cls",
[_NoViewDoubleQuantModel, _ViewDoubleQuantModel],
ids=["no_view", "with_view"],
)
@pytest.mark.skipif(
not is_aiter_found_and_supported(),
reason="Only test on ROCm with AITER installed and supported",
)
def test_double_aiter_rms_fp8_group_quant_fusion(
model_cls: type[torch.nn.Module],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
Both fan-out shapes (with and without an intermediate view) must fuse
into ``rocm_aiter_rmsnorm_fp8_group_quant``: the no-view shape via
``DoubleAiterRMSFp8GroupQuantPattern`` and the viewed shape via the
new ``DoubleAiterRMSFp8GroupQuantViewPattern`` sibling.
A failure on the ``with_view`` parametrization is a regression on the
DSv3.2 q_c norm path that this PR's view-tolerant pattern is intended
to cover.
"""
torch._dynamo.reset()
vllm_config = VllmConfig(
model_config=ModelConfig(dtype=torch.bfloat16),
compilation_config=CompilationConfig(
mode=CompilationMode.VLLM_COMPILE,
custom_ops=["+rms_norm", "+quant_fp8"],
pass_config=PassConfig(
fuse_norm_quant=True,
eliminate_noops=True,
),
),
)
with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
RocmAiterRMSNormQuantFusionPass,
)
torch.set_default_device("cuda")
torch.set_default_dtype(torch.bfloat16)
torch.manual_seed(0)
m.setenv("VLLM_ROCM_USE_AITER", "1")
rocm_aiter_ops.refresh_env_variables()
fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config)
passes = [
NoOpEliminationPass(vllm_config),
fusion_pass,
PostCleanupPass(vllm_config),
]
backend = TestBackend(*passes)
model = model_cls()
x = torch.randn(8, HIDDEN_SIZE)
torch._dynamo.mark_dynamic(x, 0)
outputs_unfused = model(x)
model_fused = torch.compile(model, backend=backend)
outputs_fused = model_fused(x)
# Both consumers must be rewritten into the fused op (one
# ``register_replacement`` rewrite covers the whole 1-to-2 fan-out).
assert fusion_pass.matched_count == 1, (
f"Expected the {model_cls.__name__} fan-out to fuse via the "
f"DoubleQuant pattern (matched_count == 1), got "
f"{fusion_pass.matched_count}"
)
fused_op = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
backend.check_after_ops([fused_op])
# Numerical parity sanity-check: the fused pair must match the
# unfused pair on FP8 outputs (exact byte-equality is the goal,
# but allow a tiny tolerance for any residual numeric noise).
for fused_t, unfused_t in zip(outputs_fused, outputs_unfused):
torch.testing.assert_close(
fused_t.to(torch.float32),
unfused_t.to(torch.float32),
atol=1e-2,
rtol=1e-2,
)
@@ -0,0 +1,413 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import vllm.config
from tests.compile.backend import TestBackend
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
from vllm.compilation.passes.fusion.mla_rope_kvcache_cat_fusion import (
MLARoPEKVCacheCatFusionPass,
)
from vllm.compilation.passes.utility.fix_functionalization import (
FixFunctionalizationPass,
)
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
from vllm.config import (
CacheConfig,
CompilationConfig,
CompilationMode,
ModelConfig,
PassConfig,
VllmConfig,
)
from vllm.forward_context import get_forward_context, set_forward_context
from vllm.model_executor.layers.attention import MLAAttention
from vllm.model_executor.layers.linear import ColumnParallelLinear
from vllm.model_executor.layers.rotary_embedding import (
DeepseekScalingRotaryEmbedding,
RotaryEmbedding,
)
from vllm.platforms import current_platform
from vllm.utils.torch_utils import _encode_layer_name
from vllm.v1.attention.backend import (
AttentionBackend,
CommonAttentionMetadata,
)
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
from vllm.v1.attention.backends.registry import AttentionBackendEnum
INDEX_SELECT_OP = torch.ops.aten.index.Tensor
VLLM_UNIFIED_MLA_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_mla_kv_cache_update
FP8_DTYPE = current_platform.fp8_dtype()
class MLARoPEKVCacheCatTestModel(torch.nn.Module):
def __init__(
self,
vllm_config: VllmConfig,
attn_backend: AttentionBackendEnum,
use_deepseek_scaling_rope: bool,
num_heads: int,
qk_nope_head_dim: int,
qk_rope_head_dim: int,
v_head_dim: int,
q_lora_rank: int,
kv_lora_rank: int,
is_neox: bool,
dtype: torch.dtype,
device: torch.device,
prefix: str = "model.layers.0.self_attn.attn",
):
super().__init__()
self.num_heads = num_heads
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
self.v_head_dim = v_head_dim
self.q_lora_rank = q_lora_rank
self.kv_lora_rank = kv_lora_rank
self.dtype = dtype
self.device = device
self.layer_name = prefix
self.num_kv_heads = 1
self.head_size = kv_lora_rank + qk_rope_head_dim
self.block_size = vllm_config.cache_config.block_size
self.scale = self.qk_head_dim**-0.5
if use_deepseek_scaling_rope:
self.rotary_emb = DeepseekScalingRotaryEmbedding(
head_size=qk_rope_head_dim,
rotary_dim=qk_rope_head_dim,
max_position_embeddings=4096,
base=10000,
is_neox_style=is_neox,
scaling_factor=1.0,
dtype=dtype,
)
else:
self.rotary_emb = RotaryEmbedding(
head_size=qk_rope_head_dim,
rotary_dim=qk_rope_head_dim,
max_position_embeddings=4096,
base=10000,
is_neox_style=is_neox,
dtype=dtype,
)
# Initialize intermediate mm layers for unit test
self.q_b_proj = ColumnParallelLinear(
self.q_lora_rank,
self.num_heads * self.qk_head_dim,
bias=False,
prefix=f"{prefix}.q_b_proj",
).to(device)
self.kv_b_proj = ColumnParallelLinear(
self.kv_lora_rank,
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
bias=False,
prefix=f"{prefix}.kv_b_proj",
).to(device)
# ColumnParallelLinear default init in bf16 with seed 0 produces
# near-zero weights (7/4.7M nonzero), making the GEMM output almost
# entirely zero and masking correctness bugs. Reinitialize to get
# dense outputs.
with torch.no_grad():
torch.nn.init.normal_(self.q_b_proj.weight, std=0.02)
torch.nn.init.normal_(self.kv_b_proj.weight, std=0.02)
# Register layer metadata for the fusion pass via MLAAttention
self.mla_attn = MLAAttention(
num_heads=self.num_heads,
scale=self.scale,
qk_nope_head_dim=self.qk_nope_head_dim,
qk_rope_head_dim=self.qk_rope_head_dim,
v_head_dim=self.v_head_dim,
q_lora_rank=self.q_lora_rank,
kv_lora_rank=self.kv_lora_rank,
kv_b_proj=self.kv_b_proj,
cache_config=vllm_config.cache_config,
quant_config=vllm_config.quant_config,
prefix=prefix,
attn_backend=attn_backend.get_class(),
)
self.attn_backend: type[AttentionBackend] = self.mla_attn.get_attn_backend()
self.mla_attn._k_scale = self.mla_attn._k_scale.to(device)
self.mla_attn._v_scale = self.mla_attn._v_scale.to(device)
# Keep both the string dtype (for ops) and torch dtype (for tensors)
self.kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
self.kv_cache_dtype = (
FP8_DTYPE if self.kv_cache_dtype_str.startswith("fp8") else self.dtype
)
# Initialize attn MetadataBuilder
self.builder = self.attn_backend.get_builder_cls()(
kv_cache_spec=self.mla_attn.get_kv_cache_spec(vllm_config),
layer_names=[self.mla_attn.layer_name],
vllm_config=vllm_config,
device=device,
)
def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
"""Initialize attention metadata."""
# Create common attn metadata
batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
common_attn_metadata = create_common_attn_metadata(
batch_spec, self.block_size, self.device, arange_block_indices=True
)
max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
num_blocks = batch_size * max_blocks
# Fetch the attention backend and kv cache shape and stride order
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
num_blocks, self.block_size, self.num_kv_heads, self.head_size
)
try:
kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
except (AttributeError, NotImplementedError):
kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
inv_order = [
kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
]
raw_tensor = torch.zeros(
num_blocks * self.block_size * self.num_kv_heads * self.head_size,
dtype=self.kv_cache_dtype,
device=self.device,
)
raw_tensor = raw_tensor.view(kv_cache_shape)
kv_cache = raw_tensor.permute(*inv_order)
self.mla_attn.kv_cache = kv_cache
# Build attn metadata
attn_metadata = self.builder.build(
common_prefix_len=0, common_attn_metadata=common_attn_metadata
)
return attn_metadata
def forward(
self, qkv_lora: torch.Tensor, positions: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
qkv_lora = qkv_lora.clone()
q_c, kv_lora = qkv_lora.split(
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
dim=-1,
)
q = self.q_b_proj(q_c)[0]
kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
q = q.view(-1, self.num_heads, self.qk_head_dim)
k_pe = k_pe.unsqueeze(1)
q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
positions, q[..., self.qk_nope_head_dim :], k_pe
)
dummy = torch.ops.vllm.unified_mla_kv_cache_update(
kv_c,
k_pe,
_encode_layer_name(self.layer_name),
self.kv_cache_dtype_str,
self.mla_attn._k_scale,
)
return q, kv_c, k_pe, dummy
def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
ops = [
INDEX_SELECT_OP,
torch.ops.vllm.unified_mla_kv_cache_update.default,
]
return ops
def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
return [torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default]
MLA_BACKENDS = [AttentionBackendEnum.TRITON_MLA]
if flash_attn_supports_mla():
MLA_BACKENDS += [AttentionBackendEnum.FLASH_ATTN_MLA]
if is_aiter_found_and_supported():
MLA_BACKENDS += [AttentionBackendEnum.ROCM_AITER_MLA]
@pytest.mark.parametrize("attn_backend", MLA_BACKENDS)
@pytest.mark.parametrize("use_deepseek_scaling_rope", [True])
@pytest.mark.parametrize("num_heads", [16])
@pytest.mark.parametrize("qk_nope_head_dim", [128])
@pytest.mark.parametrize("qk_rope_head_dim", [64])
@pytest.mark.parametrize("v_head_dim", [128])
@pytest.mark.parametrize("q_lora_rank", [1536])
@pytest.mark.parametrize("kv_lora_rank", [512])
@pytest.mark.parametrize("block_size", [16])
@pytest.mark.parametrize("is_neox", [True, False])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="MLA RoPE+KVCache+Cat fusion is only supported on CUDA and ROCm.",
)
def test_mla_rope_kvcache_cat_fusion(
attn_backend: AttentionBackendEnum,
use_deepseek_scaling_rope: bool,
num_heads: int,
qk_nope_head_dim: int,
qk_rope_head_dim: int,
v_head_dim: int,
q_lora_rank: int,
kv_lora_rank: int,
block_size: int,
is_neox: bool,
dtype: torch.dtype,
kv_cache_dtype: str,
monkeypatch: pytest.MonkeyPatch,
):
torch.set_default_device("cuda")
torch.set_default_dtype(dtype)
torch.manual_seed(0)
vllm_config = VllmConfig(
model_config=ModelConfig(
model="deepseek-ai/DeepSeek-V2-Lite",
dtype=dtype,
),
cache_config=CacheConfig(
block_size=block_size,
cache_dtype=kv_cache_dtype,
),
compilation_config=CompilationConfig(
mode=CompilationMode.VLLM_COMPILE,
pass_config=PassConfig(
fuse_rope_kvcache_cat_mla=True,
eliminate_noops=True,
),
),
)
with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
if not torch.distributed.is_initialized():
from vllm.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
)
from vllm.utils.system_utils import update_environment_variables
update_environment_variables(
{
"RANK": "0",
"LOCAL_RANK": "0",
"WORLD_SIZE": "1",
"MASTER_ADDR": "localhost",
"MASTER_PORT": "54321",
}
)
init_distributed_environment()
initialize_model_parallel()
if attn_backend == AttentionBackendEnum.ROCM_AITER_MLA:
m.setenv("VLLM_ROCM_USE_AITER", "1")
rocm_aiter_ops.refresh_env_variables()
model = MLARoPEKVCacheCatTestModel(
vllm_config=vllm_config,
attn_backend=attn_backend,
use_deepseek_scaling_rope=use_deepseek_scaling_rope,
num_heads=num_heads,
qk_nope_head_dim=qk_nope_head_dim,
qk_rope_head_dim=qk_rope_head_dim,
v_head_dim=v_head_dim,
q_lora_rank=q_lora_rank,
kv_lora_rank=kv_lora_rank,
is_neox=is_neox,
dtype=dtype,
device=torch.get_default_device(),
)
fusion_pass = MLARoPEKVCacheCatFusionPass(vllm_config)
# note: FixFunctionalizationPass is required to correctly lower
# the fused op to its inplace version with auto-functionalization v1.
# Without it, decompose_auto_functionalized calls clone_preserve_strides
# on the non-contiguous q_pe slice directly, and inductor's lowering
# of the resulting as_strided chain incorrectly drops the storage offset.
# auto-functionalization v2 avoids this: it clones the contiguous base
# tensor (_all_bases) and reconstructs the slice as a view, so the
# offset is never passed through as_strided lowering.
passes = [
NoOpEliminationPass(vllm_config),
fusion_pass,
PostCleanupPass(vllm_config),
FixFunctionalizationPass(vllm_config),
]
backend = TestBackend(*passes)
T = 5
qkv_lora = torch.randn(
T,
q_lora_rank + kv_lora_rank + qk_rope_head_dim,
dtype=dtype,
)
pos = torch.arange(T, dtype=torch.long)
qkv_unfused = qkv_lora.clone()
pos_unfused = pos.clone()
# Run unfused version
with set_forward_context(None, vllm_config):
forward_context = get_forward_context()
attn_metadata = model.build_attn_metadata(T)
forward_context.slot_mapping = {
model.layer_name: attn_metadata.slot_mapping
}
q_unfused, kv_c_unfused, k_pe_unfused, dummy = model(
qkv_unfused, pos_unfused
)
attn_layer = forward_context.no_compile_layers[model.layer_name]
kv_cache_unfused = attn_layer.kv_cache.clone()
del dummy
# Run fused version (compiled)
torch._dynamo.mark_dynamic(qkv_lora, 0)
torch._dynamo.mark_dynamic(pos, 0)
with set_forward_context(None, vllm_config):
model_fused = torch.compile(model, backend=backend)
forward_context = get_forward_context()
attn_metadata = model.build_attn_metadata(T)
forward_context.slot_mapping = {
model.layer_name: attn_metadata.slot_mapping
}
q_fused, kv_c_fused, k_pe_fused, dummy = model_fused(qkv_lora, pos)
attn_layer = forward_context.no_compile_layers[model.layer_name]
kv_cache_fused = attn_layer.kv_cache
del dummy
assert fusion_pass.matched_count == 1
backend.check_before_ops(model.ops_in_model_before())
backend.check_after_ops(model.ops_in_model_after())
if dtype == torch.float16:
ATOL, RTOL = (2e-3, 2e-3)
else:
ATOL, RTOL = (1e-2, 1e-2)
torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
torch.testing.assert_close(kv_c_unfused, kv_c_fused, atol=ATOL, rtol=RTOL)
torch.testing.assert_close(k_pe_unfused, k_pe_fused, atol=ATOL, rtol=RTOL)
# Cannot compare fp8_* directly here, cast to model dtype instead
torch.testing.assert_close(
kv_cache_unfused.view(dtype),
kv_cache_fused.view(dtype),
atol=ATOL,
rtol=RTOL,
)
@@ -34,7 +34,6 @@ from vllm.v1.attention.backend import (
CommonAttentionMetadata,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.kv_cache_interface import AttentionSpec
INDEX_SELECT_OP = torch.ops.aten.index.Tensor
VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
@@ -102,13 +101,8 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
)
# Initialize attn MetadataBuilder
self.builder = self.attn.attn_backend.get_builder_cls()(
kv_cache_spec=AttentionSpec(
block_size=self.block_size,
num_kv_heads=self.num_kv_heads,
head_size=head_size,
dtype=self.kv_cache_dtype,
),
self.builder = self.attn_backend.get_builder_cls()(
kv_cache_spec=self.attn.get_kv_cache_spec(vllm_config),
layer_names=[self.attn.layer_name],
vllm_config=vllm_config,
device=device,
@@ -126,12 +120,11 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
num_blocks = batch_size * max_blocks
# Fetch the attention backend and kv cache shape and stride order
attn_backend = self.attn.attn_backend
kv_cache_shape = attn_backend.get_kv_cache_shape(
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
num_blocks, self.block_size, self.num_kv_heads, self.head_size
)
try:
kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
except (AttributeError, NotImplementedError):
kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
@@ -67,58 +67,3 @@ async def test_chat_with_enable_force_include_usage(
chunk.usage.prompt_tokens + chunk.usage.completion_tokens
)
last_completion_tokens = chunk.usage.completion_tokens
@pytest.fixture(scope="module")
def transcription_server_with_force_include_usage():
args = [
# use half precision for speed and memory savings in CI environment
"--dtype",
"bfloat16",
"--max-num-seqs",
"4",
"--enforce-eager",
"--enable-force-include-usage",
"--gpu-memory-utilization",
"0.2",
]
with RemoteOpenAIServer("openai/whisper-large-v3-turbo", args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def transcription_client_with_force_include_usage(
transcription_server_with_force_include_usage,
):
async with (
transcription_server_with_force_include_usage.get_async_client() as async_client
):
yield async_client
@pytest.mark.asyncio
async def test_transcription_with_enable_force_include_usage(
transcription_client_with_force_include_usage, winning_call
):
res = (
await transcription_client_with_force_include_usage.audio.transcriptions.create(
model="openai/whisper-large-v3-turbo",
file=winning_call,
language="en",
temperature=0.0,
stream=True,
timeout=30,
)
)
async for chunk in res:
if not len(chunk.choices):
# final usage sent
usage = chunk.usage
assert isinstance(usage, dict)
assert usage["prompt_tokens"] > 0
assert usage["completion_tokens"] > 0
assert usage["total_tokens"] > 0
else:
assert not hasattr(chunk, "usage")
@@ -0,0 +1,492 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Cross-API render parity tests.
Verifies that the chat completion input path (parse_chat_input_to_harmony_message)
and the responses API input path (response_input_to_harmony) produce identical
Harmony messages and identical rendered token sequences when given equivalent
conversation representations.
The chat completion API encodes reasoning and tool calls as fields on a single
assistant message dict; the responses API encodes them as separate typed items
in request.input. Both paths must converge on the same Harmony message list and
therefore the same rendered prompt.
Each test:
1. Builds Harmony messages from each path for a single message or sequence.
2. Asserts message-level properties (role, channel, recipient, content)
using verify_harmony_messages.
3. Asserts that render_for_completion produces identical token sequences.
"""
from openai.types.responses import ResponseFunctionToolCall
from tests.entrypoints.openai.utils import verify_harmony_messages
from vllm.entrypoints.openai.parser.harmony_utils import (
get_system_message,
parse_chat_input_to_harmony_message,
render_for_completion,
)
from vllm.entrypoints.openai.responses.harmony import response_input_to_harmony
# Use a fixed date so the system message is deterministic across both paths.
_DATE = "2025-01-01"
def _system():
return get_system_message(start_date=_DATE)
class TestResponseInputToHarmonyRenderParity:
"""Each test drives the same conversation through both APIs and asserts
identical Harmony messages and rendered token sequences."""
# -----------------------------------------------------------------------
# Single-message cases
# -----------------------------------------------------------------------
def test_user_message(self):
chat_msgs = parse_chat_input_to_harmony_message(
{"role": "user", "content": "What's the weather in Paris?"}
)
resp_msgs = [
response_input_to_harmony(
{
"type": "message",
"role": "user",
"content": "What's the weather in Paris?",
},
prev_responses=[],
)
]
expected = [{"role": "user", "content": "What's the weather in Paris?"}]
verify_harmony_messages(chat_msgs, expected)
verify_harmony_messages(resp_msgs, expected)
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
[_system()] + resp_msgs
)
def test_assistant_final_message(self):
chat_msgs = parse_chat_input_to_harmony_message(
{"role": "assistant", "content": "It is 18°C in Paris."}
)
resp_msgs = [
response_input_to_harmony(
{
"type": "message",
"role": "assistant",
"content": "It is 18°C in Paris.",
},
prev_responses=[],
)
]
expected = [
{"role": "assistant", "channel": "final", "content": "It is 18°C in Paris."}
]
verify_harmony_messages(chat_msgs, expected)
verify_harmony_messages(resp_msgs, expected)
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
[_system()] + resp_msgs
)
def test_reasoning_item(self):
# Chat path: assistant message with only a reasoning field and no content.
chat_msgs = parse_chat_input_to_harmony_message(
{
"role": "assistant",
"reasoning": "I should call get_weather.",
"content": "",
}
)
resp_msgs = [
response_input_to_harmony(
{
"type": "reasoning",
"content": [
{"type": "reasoning_text", "text": "I should call get_weather."}
],
},
prev_responses=[],
)
]
expected = [
{
"role": "assistant",
"channel": "analysis",
"content": "I should call get_weather.",
}
]
verify_harmony_messages(chat_msgs, expected)
verify_harmony_messages(resp_msgs, expected)
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
[_system()] + resp_msgs
)
def test_function_call(self):
chat_msgs = parse_chat_input_to_harmony_message(
{
"role": "assistant",
"tool_calls": [
{
"id": "call_1",
"function": {
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
}
],
}
)
resp_msgs = [
response_input_to_harmony(
{
"type": "function_call",
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
prev_responses=[],
)
]
expected = [
{
"role": "assistant",
"channel": "commentary",
"recipient": "functions.get_weather",
"content": '{"location": "Paris"}',
"content_type": "json",
}
]
verify_harmony_messages(chat_msgs, expected)
verify_harmony_messages(resp_msgs, expected)
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
[_system()] + resp_msgs
)
def test_tool_output(self):
prev_call = ResponseFunctionToolCall(
id="fc_1",
call_id="call_1",
name="get_weather",
arguments='{"location": "Paris"}',
type="function_call",
)
chat_msgs = parse_chat_input_to_harmony_message(
{"role": "tool", "tool_call_id": "call_1", "content": "18°C, clear skies."},
tool_id_names={"call_1": "get_weather"},
)
resp_msgs = [
response_input_to_harmony(
{
"type": "function_call_output",
"call_id": "call_1",
"output": "18°C, clear skies.",
},
prev_responses=[prev_call],
)
]
expected = [
{
"role": "tool",
"author_name": "functions.get_weather",
"channel": "commentary",
"recipient": "assistant",
"content": "18°C, clear skies.",
}
]
verify_harmony_messages(chat_msgs, expected)
verify_harmony_messages(resp_msgs, expected)
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
[_system()] + resp_msgs
)
# -----------------------------------------------------------------------
# Combined and multi-turn cases
# -----------------------------------------------------------------------
def test_reasoning_combined_with_function_call(self):
"""Chat API packs reasoning + tool_calls into one dict; responses API
represents them as two separate items. Both must produce the same two
Harmony messages in the same order: analysis then commentary."""
chat_msgs = parse_chat_input_to_harmony_message(
{
"role": "assistant",
"reasoning": "I should get the weather for Paris.",
"tool_calls": [
{
"id": "call_1",
"function": {
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
}
],
}
)
resp_msgs = [
response_input_to_harmony(
{
"type": "reasoning",
"content": [
{
"type": "reasoning_text",
"text": "I should get the weather for Paris.",
}
],
},
prev_responses=[],
),
response_input_to_harmony(
{
"type": "function_call",
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
prev_responses=[],
),
]
expected = [
{
"role": "assistant",
"channel": "analysis",
"content": "I should get the weather for Paris.",
},
{
"role": "assistant",
"channel": "commentary",
"recipient": "functions.get_weather",
"content": '{"location": "Paris"}',
"content_type": "json",
},
]
verify_harmony_messages(chat_msgs, expected)
verify_harmony_messages(resp_msgs, expected)
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
[_system()] + resp_msgs
)
def test_full_multi_turn_tool_call_conversation(self):
"""Full conversation: user -> reasoning + tool_call -> tool_output -> final.
Both APIs must render the complete conversation to identical token sequences.
This exercises the entire input pipeline including all message types and
the Rust harmony encoder.
"""
prev_call = ResponseFunctionToolCall(
id="fc_1",
call_id="call_1",
name="get_weather",
arguments='{"location": "Paris"}',
type="function_call",
)
# --- Chat completion API path ---
tool_id_names = {"call_1": "get_weather"}
chat_msgs = []
chat_msgs += parse_chat_input_to_harmony_message(
{"role": "user", "content": "What's the weather in Paris?"}
)
chat_msgs += parse_chat_input_to_harmony_message(
{
"role": "assistant",
"reasoning": "I should call get_weather for Paris.",
"tool_calls": [
{
"id": "call_1",
"function": {
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
}
],
}
)
chat_msgs += parse_chat_input_to_harmony_message(
{"role": "tool", "tool_call_id": "call_1", "content": "18°C, clear skies."},
tool_id_names=tool_id_names,
)
chat_msgs += parse_chat_input_to_harmony_message(
{
"role": "assistant",
"content": "It is currently 18°C in Paris with clear skies.",
}
)
# --- Responses API path ---
resp_input = [
{
"type": "message",
"role": "user",
"content": "What's the weather in Paris?",
},
{
"type": "reasoning",
"content": [
{
"type": "reasoning_text",
"text": "I should call get_weather for Paris.",
}
],
},
{
"type": "function_call",
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
{
"type": "function_call_output",
"call_id": "call_1",
"output": "18°C, clear skies.",
},
{
"type": "message",
"role": "assistant",
"content": "It is currently 18°C in Paris with clear skies.",
},
]
resp_msgs = [
response_input_to_harmony(item, prev_responses=[prev_call])
for item in resp_input
]
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
[_system()] + resp_msgs
)
def test_multi_turn_two_tool_calls_with_reasoning_between(self):
"""Validates parity for a chain of two tool calls, each with its own
reasoning trace. Reasoning traces in between commentary-channel tool
calls must survive as analysis-channel messages in both paths.
"""
prev_call_1 = ResponseFunctionToolCall(
id="fc_1",
call_id="call_1",
name="get_weather",
arguments='{"location": "Paris"}',
type="function_call",
)
prev_call_2 = ResponseFunctionToolCall(
id="fc_2",
call_id="call_2",
name="get_forecast",
arguments='{"location": "Paris", "days": 7}',
type="function_call",
)
# --- Chat completion API path ---
tool_id_names = {"call_1": "get_weather", "call_2": "get_forecast"}
chat_msgs = []
chat_msgs += parse_chat_input_to_harmony_message(
{"role": "user", "content": "What's the weather and forecast for Paris?"}
)
# First reasoning + tool call
chat_msgs += parse_chat_input_to_harmony_message(
{
"role": "assistant",
"reasoning": "I need current weather first.",
"tool_calls": [
{
"id": "call_1",
"function": {
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
}
],
}
)
chat_msgs += parse_chat_input_to_harmony_message(
{"role": "tool", "tool_call_id": "call_1", "content": "18°C, clear skies."},
tool_id_names=tool_id_names,
)
# Second reasoning + tool call
chat_msgs += parse_chat_input_to_harmony_message(
{
"role": "assistant",
"reasoning": "Now I need the weekly forecast.",
"tool_calls": [
{
"id": "call_2",
"function": {
"name": "get_forecast",
"arguments": '{"location": "Paris", "days": 7}',
},
}
],
}
)
chat_msgs += parse_chat_input_to_harmony_message(
{
"role": "tool",
"tool_call_id": "call_2",
"content": "Mon 17°C, Tue 19°C, Wed 16°C",
},
tool_id_names=tool_id_names,
)
# --- Responses API path ---
prev_responses = [prev_call_1, prev_call_2]
resp_input = [
{
"type": "message",
"role": "user",
"content": "What's the weather and forecast for Paris?",
},
# First reasoning + tool call
{
"type": "reasoning",
"content": [
{"type": "reasoning_text", "text": "I need current weather first."}
],
},
{
"type": "function_call",
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
{
"type": "function_call_output",
"call_id": "call_1",
"output": "18°C, clear skies.",
},
# Second reasoning + tool call
{
"type": "reasoning",
"content": [
{
"type": "reasoning_text",
"text": "Now I need the weekly forecast.",
}
],
},
{
"type": "function_call",
"name": "get_forecast",
"arguments": '{"location": "Paris", "days": 7}',
},
{
"type": "function_call_output",
"call_id": "call_2",
"output": "Mon 17°C, Tue 19°C, Wed 16°C",
},
]
resp_msgs = [
response_input_to_harmony(item, prev_responses=prev_responses)
for item in resp_input
]
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
[_system()] + resp_msgs
)
@@ -325,8 +325,12 @@ async def test_function_calling_with_streaming_expected_arguments(
"tool_choice",
["auto", "required", {"type": "function", "name": "get_current_weather"}],
)
@pytest.mark.parametrize(
"enable_thinking",
[True, False],
)
async def test_function_calling_with_streaming_types(
client: openai.AsyncOpenAI, model_name: str, tool_choice
client: openai.AsyncOpenAI, model_name: str, tool_choice, enable_thinking: bool
):
# this links the "done" type with the "start" type
# so every "done" type should have a corresponding "start" type
@@ -436,6 +440,7 @@ async def test_function_calling_with_streaming_types(
input=input_list,
tools=tools,
tool_choice=tool_choice,
extra_body={"chat_template_kwargs": {"enable_thinking": enable_thinking}},
stream=True,
)
@@ -0,0 +1,274 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for response_input_to_harmony.
Covers every type branch in the function and verifies that each produced
Harmony Message has the correct role, channel, recipient, content_type,
author name, and text content.
"""
import pytest
from openai.types.responses import ResponseFunctionToolCall, ResponseReasoningItem
from openai.types.responses.response_reasoning_item import (
Content as ReasoningTextContent,
)
from openai_harmony import Role
from vllm.entrypoints.openai.responses.harmony import response_input_to_harmony
# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------
_PREV_CALL = ResponseFunctionToolCall(
id="fc_test",
call_id="call_test",
name="get_weather",
arguments='{"location": "Paris"}',
type="function_call",
)
_REASONING_ITEM = ResponseReasoningItem(
id="rs_test",
type="reasoning",
content=[ReasoningTextContent(type="reasoning_text", text="Thinking hard.")],
summary=[],
status=None,
)
class TestResponseInputToHarmonyMessage:
"""Unit tests for every message type handled by response_input_to_harmony."""
# -----------------------------------------------------------------------
# type="message" (or no type key)
# -----------------------------------------------------------------------
def test_user_message_string_content(self):
msg = response_input_to_harmony(
{"type": "message", "role": "user", "content": "Hello"},
prev_responses=[],
)
assert msg.author.role == Role.USER
assert msg.content[0].text == "Hello"
assert msg.channel is None
def test_no_type_key_defaults_to_message_branch(self):
"""Omitting 'type' should fall through to the message branch."""
msg = response_input_to_harmony(
{"role": "user", "content": "Hello"},
prev_responses=[],
)
assert msg.author.role == Role.USER
assert msg.content[0].text == "Hello"
def test_system_message(self):
msg = response_input_to_harmony(
{"type": "message", "role": "system", "content": "Be helpful."},
prev_responses=[],
)
assert msg.author.role == Role.SYSTEM
assert msg.content[0].text == "Be helpful."
assert msg.channel is None
def test_assistant_message_gets_final_channel(self):
msg = response_input_to_harmony(
{"type": "message", "role": "assistant", "content": "The answer is 42."},
prev_responses=[],
)
assert msg.author.role == Role.ASSISTANT
assert msg.channel == "final"
assert msg.content[0].text == "The answer is 42."
def test_developer_message_gets_instructions_prefix(self):
msg = response_input_to_harmony(
{"type": "message", "role": "developer", "content": "Be concise."},
prev_responses=[],
)
assert msg.author.role == Role.DEVELOPER
assert msg.content[0].text == "Instructions:\nBe concise."
assert msg.channel is None
def test_message_with_array_content(self):
msg = response_input_to_harmony(
{
"type": "message",
"role": "user",
"content": [
{"type": "text", "text": "Part one. "},
{"type": "text", "text": "Part two."},
],
},
prev_responses=[],
)
assert msg.author.role == Role.USER
assert len(msg.content) == 2
assert msg.content[0].text == "Part one. "
assert msg.content[1].text == "Part two."
def test_developer_message_array_content_gets_prefix_on_each_part(self):
msg = response_input_to_harmony(
{
"type": "message",
"role": "developer",
"content": [
{"type": "text", "text": "Rule 1."},
{"type": "text", "text": "Rule 2."},
],
},
prev_responses=[],
)
assert msg.content[0].text == "Instructions:\nRule 1."
assert msg.content[1].text == "Instructions:\nRule 2."
# -----------------------------------------------------------------------
# type="reasoning"
# -----------------------------------------------------------------------
def test_reasoning_gets_analysis_channel(self):
msg = response_input_to_harmony(
{
"type": "reasoning",
"content": [
{"type": "reasoning_text", "text": "I should call get_weather."}
],
},
prev_responses=[],
)
assert msg.author.role == Role.ASSISTANT
assert msg.channel == "analysis"
assert msg.content[0].text == "I should call get_weather."
def test_reasoning_pydantic_model_input(self):
"""A Pydantic ResponseReasoningItem should be model_dump()'d before parsing."""
msg = response_input_to_harmony(_REASONING_ITEM, prev_responses=[])
assert msg.author.role == Role.ASSISTANT
assert msg.channel == "analysis"
assert msg.content[0].text == "Thinking hard."
# -----------------------------------------------------------------------
# type="function_call"
# -----------------------------------------------------------------------
def test_function_call_channel_recipient_and_content_type(self):
msg = response_input_to_harmony(
{
"type": "function_call",
"name": "get_weather",
"arguments": '{"location": "Paris"}',
},
prev_responses=[],
)
assert msg.author.role == Role.ASSISTANT
assert msg.channel == "commentary"
assert msg.recipient == "functions.get_weather"
assert msg.content_type == "json"
assert msg.content[0].text == '{"location": "Paris"}'
def test_function_call_empty_arguments(self):
msg = response_input_to_harmony(
{"type": "function_call", "name": "ping", "arguments": ""},
prev_responses=[],
)
assert msg.recipient == "functions.ping"
assert msg.content[0].text == ""
# -----------------------------------------------------------------------
# type="function_call_output"
# -----------------------------------------------------------------------
def test_function_call_output_channel_recipient_and_author_name(self):
msg = response_input_to_harmony(
{"type": "function_call_output", "call_id": "call_test", "output": "18°C"},
prev_responses=[_PREV_CALL],
)
assert msg.author.role == Role.TOOL
assert msg.author.name == "functions.get_weather"
assert msg.channel == "commentary"
assert msg.recipient == "assistant"
assert msg.content[0].text == "18°C"
def test_function_call_output_uses_most_recent_matching_call(self):
"""When multiple prev_responses share a call_id, the last one wins
because the search is reversed."""
earlier = ResponseFunctionToolCall(
id="fc_old",
call_id="call_test",
name="old_func",
arguments="{}",
type="function_call",
)
later = ResponseFunctionToolCall(
id="fc_new",
call_id="call_test",
name="get_weather",
arguments="{}",
type="function_call",
)
msg = response_input_to_harmony(
{
"type": "function_call_output",
"call_id": "call_test",
"output": "result",
},
prev_responses=[earlier, later],
)
assert msg.author.name == "functions.get_weather"
def test_function_call_output_skips_non_function_call_items_in_prev_responses(
self,
):
"""ResponseReasoningItem entries in prev_responses should be ignored."""
msg = response_input_to_harmony(
{
"type": "function_call_output",
"call_id": "call_test",
"output": "18°C",
},
prev_responses=[_REASONING_ITEM, _PREV_CALL],
)
assert msg.author.name == "functions.get_weather"
def test_function_call_output_raises_if_no_matching_call(self):
with pytest.raises(ValueError, match="No call message found for"):
response_input_to_harmony(
{
"type": "function_call_output",
"call_id": "no_such_id",
"output": "x",
},
prev_responses=[_PREV_CALL],
)
def test_function_call_output_raises_on_empty_prev_responses(self):
with pytest.raises(ValueError, match="No call message found for"):
response_input_to_harmony(
{"type": "function_call_output", "call_id": "call_test", "output": "x"},
prev_responses=[],
)
# -----------------------------------------------------------------------
# Error cases
# -----------------------------------------------------------------------
def test_unknown_type_raises_value_error(self):
with pytest.raises(ValueError, match="Unknown input type"):
response_input_to_harmony(
{"type": "image_url", "url": "https://example.com/img.png"},
prev_responses=[],
)
@@ -10,7 +10,7 @@ import pybase64 as base64
import pytest
import websockets
from tests.entrypoints.openai.conftest import add_attention_backend
from tests.entrypoints.speech_to_text.conftest import add_attention_backend
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
from vllm.assets.audio import AudioAsset
from vllm.multimodal.media.audio import load_audio
@@ -163,6 +163,11 @@ async def test_multi_chunk_streaming(
" A little piece of practical poetry. Mary had a little lamb,"
" it sleeps with quite a flow, and everywhere that Mary went,"
" the lamb was sure to go."
) or full_text == (
" First words I spoke in the original phonograph."
" A little piece of practical poetry. Mary had a little lamb,"
" it squeaked with quite a flow, and everywhere that Mary went,"
" the lamb was sure to go."
)
@@ -7,8 +7,8 @@ from unittest.mock import AsyncMock, Mock
import pytest
from vllm.entrypoints.openai.speech_to_text.protocol import TranscriptionResponse
from vllm.entrypoints.openai.speech_to_text.speech_to_text import OpenAISpeechToText
from vllm.entrypoints.speech_to_text.base.serving import OpenAISpeechToText
from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionResponse
async def _never_finishes():
@@ -0,0 +1,61 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import pytest_asyncio
from tests.utils import RemoteOpenAIServer
@pytest.fixture(scope="module")
def transcription_server_with_force_include_usage():
args = [
# use half precision for speed and memory savings in CI environment
"--dtype",
"bfloat16",
"--max-num-seqs",
"4",
"--enforce-eager",
"--enable-force-include-usage",
"--gpu-memory-utilization",
"0.2",
]
with RemoteOpenAIServer("openai/whisper-large-v3-turbo", args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def transcription_client_with_force_include_usage(
transcription_server_with_force_include_usage,
):
async with (
transcription_server_with_force_include_usage.get_async_client() as async_client
):
yield async_client
@pytest.mark.asyncio
async def test_transcription_with_enable_force_include_usage(
transcription_client_with_force_include_usage, winning_call
):
res = (
await transcription_client_with_force_include_usage.audio.transcriptions.create(
model="openai/whisper-large-v3-turbo",
file=winning_call,
language="en",
temperature=0.0,
stream=True,
timeout=30,
)
)
async for chunk in res:
if not len(chunk.choices):
# final usage sent
usage = chunk.usage
assert isinstance(usage, dict)
assert usage["prompt_tokens"] > 0
assert usage["completion_tokens"] > 0
assert usage["total_tokens"] > 0
else:
assert not hasattr(chunk, "usage")
@@ -24,12 +24,14 @@ from vllm.entrypoints.openai.engine.protocol import (
RequestResponseMetadata,
)
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.openai.speech_to_text.protocol import TranscriptionRequest
from vllm.entrypoints.openai.speech_to_text.serving import OpenAIServingTranscription
from vllm.entrypoints.openai.speech_to_text.speech_to_text import (
from vllm.entrypoints.speech_to_text.base.serving import (
OpenAISpeechToText,
asr_inter_chunk_separator,
)
from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionRequest
from vllm.entrypoints.speech_to_text.transcription.serving import (
OpenAIServingTranscription,
)
from vllm.model_executor.models.interfaces import SupportsTranscription
from vllm.outputs import CompletionOutput, RequestOutput
@@ -6,7 +6,7 @@ import json
import pytest
from tests.entrypoints.openai.conftest import add_attention_backend
from tests.entrypoints.speech_to_text.conftest import add_attention_backend
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
MISTRAL_FORMAT_ARGS = [
@@ -13,7 +13,7 @@ import pytest
import pytest_asyncio
import soundfile as sf
from tests.entrypoints.openai.conftest import add_attention_backend
from tests.entrypoints.speech_to_text.conftest import add_attention_backend
from tests.utils import RemoteOpenAIServer
from vllm.logger import init_logger
from vllm.multimodal.media.audio import load_audio
@@ -0,0 +1,3 @@
Qwen3-Next-80B-A3B-NVFP4-EP2.yaml
Qwen3-Next-FP8-EP2.yaml
Nemotron-3-Super-120B-A12B-NVFP4.yaml
@@ -3,6 +3,3 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
Qwen1.5-MoE-W4A16-CT.yaml
DeepSeek-V2-Lite-Instruct-FP8.yaml
Qwen3-30B-A3B-NVFP4.yaml
Qwen3-Next-80B-A3B-NVFP4-EP2.yaml
Qwen3-Next-FP8-EP2.yaml
Nemotron-3-Super-120B-A12B-NVFP4.yaml
@@ -0,0 +1,377 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm.platforms import current_platform
pytestmark = pytest.mark.skipif(
not current_platform.is_rocm(), reason="Only used by ROCm"
)
NOPE_HEAD_DIM = 448
ROPE_HEAD_DIM = 64
HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
def _ref_global_topk_ragged(
topk_indices: torch.Tensor,
token_to_req_indices: torch.Tensor,
block_table: torch.Tensor,
block_size: int,
is_valid_token: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
topk = topk_indices.reshape(topk_indices.shape[0], -1)
valid = (topk >= 0) & is_valid_token[:, None]
lens = valid.sum(dim=1, dtype=torch.int32)
indptr = torch.zeros(lens.shape[0] + 1, dtype=torch.int32, device=topk.device)
torch.cumsum(lens, dim=0, out=indptr[1:])
safe_topk = torch.clamp(topk, min=0)
block_indices = safe_topk // block_size
block_offsets = safe_topk % block_size
req_indices = token_to_req_indices[:, None].expand_as(topk)
slot_ids = block_table[req_indices, block_indices] * block_size + block_offsets
offsets = torch.arange(topk.shape[1], dtype=torch.int32, device=topk.device)
positions = indptr[:-1, None] + offsets[None, :]
return slot_ids[valid], positions[valid].to(torch.long), indptr, lens
def _ref_sparse_prefill_ragged(
q: torch.Tensor,
kv: torch.Tensor,
rows: list[list[int]],
scale: float,
attn_sink: torch.Tensor | None,
) -> torch.Tensor:
q_f32 = q.float()
kv_f32 = kv.float()
out = torch.empty_like(q_f32)
for query_idx in range(q.shape[0]):
row_indices = rows[query_idx]
for head_idx in range(q.shape[1]):
if row_indices:
selected_kv = kv_f32[row_indices]
scores = torch.mv(selected_kv, q_f32[query_idx, head_idx]) * scale
if attn_sink is not None:
scores_with_sink = torch.cat(
[scores, attn_sink[head_idx].float().reshape(1)]
)
probs = torch.softmax(scores_with_sink, dim=0)[:-1]
else:
probs = torch.softmax(scores, dim=0)
out[query_idx, head_idx] = torch.sum(
probs[:, None] * selected_kv, dim=0
)
else:
out[query_idx, head_idx] = 0
return out.to(torch.bfloat16)
def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
assert kv.shape[-1] == HEAD_DIM
num_tokens = kv.shape[0]
num_blocks = (num_tokens + block_size - 1) // block_size
cache = torch.zeros(
(num_blocks, block_size, 584),
dtype=torch.uint8,
device=kv.device,
)
cache_flat = cache.view(torch.uint8).flatten()
kv_nope_fp8 = (
kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
)
kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
for slot in range(num_tokens):
block_idx = slot // block_size
pos = slot % block_size
block_base = block_idx * cache.stride(0)
token_base = block_base + pos * 576
scale_base = block_base + block_size * 576 + pos * 8
cache_flat[token_base : token_base + NOPE_HEAD_DIM].copy_(kv_nope_fp8[slot])
cache_flat[
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
].copy_(kv_rope_u8[slot])
cache_flat[scale_base : scale_base + 7].fill_(127)
return cache
def _read_fp8_ds_mla_cache(
cache: torch.Tensor, slot: int, block_size: int
) -> torch.Tensor:
cache_flat = cache.view(torch.uint8).flatten()
block_idx = slot // block_size
pos = slot % block_size
block_base = block_idx * cache.stride(0)
token_base = block_base + pos * 576
nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
rope_u8 = cache_flat[
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
]
rope = rope_u8.view(torch.bfloat16).to(torch.float32)
return torch.cat([nope, rope])
def _ref_sparse_decode_ragged(
q: torch.Tensor,
main_cache: torch.Tensor,
main_rows: list[list[int]],
scale: float,
attn_sink: torch.Tensor | None,
block_size: int,
extra_cache: torch.Tensor | None = None,
extra_rows: list[list[int]] | None = None,
) -> torch.Tensor:
q_f32 = q.float()
out = torch.empty_like(q_f32)
for query_idx in range(q.shape[0]):
row_kv = [
_read_fp8_ds_mla_cache(main_cache, int(slot), block_size)
for slot in main_rows[query_idx]
]
if extra_cache is not None and extra_rows is not None:
row_kv.extend(
_read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
for slot in extra_rows[query_idx]
)
kv = torch.stack(row_kv).to(q.device)
for head_idx in range(q.shape[1]):
scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale
if attn_sink is not None:
scores_with_sink = torch.cat(
[scores, attn_sink[head_idx].float().reshape(1)]
)
probs = torch.softmax(scores_with_sink, dim=0)[:-1]
else:
probs = torch.softmax(scores, dim=0)
out[query_idx, head_idx] = torch.sum(probs[:, None] * kv, dim=0)
return out.to(torch.bfloat16)
def _ref_combine_topk_swa_ragged(
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
expected_ragged = torch.tensor(
[
100,
101,
7,
8,
9,
110,
111,
8,
9,
10,
120,
121,
122,
9,
10,
11,
150,
27,
28,
29,
160,
161,
28,
29,
30,
],
dtype=torch.int32,
device=device,
)
expected_lens = torch.tensor([5, 5, 6, 4, 5], dtype=torch.int32, device=device)
expected_indptr = torch.zeros(6, dtype=torch.int32, device=device)
torch.cumsum(expected_lens, dim=0, out=expected_indptr[1:])
return expected_ragged, expected_indptr, expected_lens
@torch.inference_mode()
def test_compute_global_topk_ragged_indices_and_indptr() -> None:
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
compute_global_topk_ragged_indices_and_indptr,
)
device = torch.device("cuda")
block_size = 4
topk_indices = torch.tensor(
[
[0, 3, 4, -1],
[5, 8, -1, -1],
[2, 7, 9, -1],
],
dtype=torch.int32,
device=device,
)
token_to_req_indices = torch.tensor([0, 1, 1], dtype=torch.int32, device=device)
block_table = torch.tensor(
[
[10, 11, 12],
[20, 21, 22],
],
dtype=torch.int32,
device=device,
)
is_valid_token = torch.tensor([True, False, True], dtype=torch.bool, device=device)
actual_ragged, actual_indptr, actual_lens = (
compute_global_topk_ragged_indices_and_indptr(
topk_indices,
token_to_req_indices,
block_table,
block_size,
is_valid_token,
)
)
expected_values, expected_positions, expected_indptr, expected_lens = (
_ref_global_topk_ragged(
topk_indices,
token_to_req_indices,
block_table,
block_size,
is_valid_token,
)
)
torch.testing.assert_close(actual_ragged[expected_positions], expected_values)
torch.testing.assert_close(actual_indptr, expected_indptr)
torch.testing.assert_close(actual_lens, expected_lens)
@torch.inference_mode()
def test_sparse_attn_prefill_ragged_kernel() -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_prefill_ragged_triton,
)
device = torch.device("cuda")
torch.manual_seed(0)
q = torch.randn(3, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
indices = torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device)
indptr = torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device)
attn_sink = torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device)
scale = HEAD_DIM**-0.5
actual = _rocm_sparse_attn_prefill_ragged_triton(
q=q,
kv=kv,
indices=indices,
indptr=indptr,
scale=scale,
attn_sink=attn_sink,
nope_head_dim=NOPE_HEAD_DIM,
rope_head_dim=ROPE_HEAD_DIM,
)
expected = _ref_sparse_prefill_ragged(
q, kv, [[0, 2], [1, 3, 4], []], scale, attn_sink
)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
@torch.inference_mode()
def test_sparse_attn_decode_ragged_kernel() -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_decode_ragged_triton,
)
device = torch.device("cuda")
torch.manual_seed(1)
block_size = 4
q = torch.randn(2, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
extra_indptr = torch.tensor([0, 1, 3], dtype=torch.int32, device=device)
attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
scale = HEAD_DIM**-0.5
actual = _rocm_sparse_attn_decode_ragged_triton(
q=q,
main_cache=main_cache,
main_indices=main_indices,
main_indptr=main_indptr,
scale=scale,
attn_sink=attn_sink,
nope_head_dim=NOPE_HEAD_DIM,
rope_head_dim=ROPE_HEAD_DIM,
extra_cache=extra_cache,
extra_indices=extra_indices,
extra_indptr=extra_indptr,
)
expected = _ref_sparse_decode_ragged(
q=q,
main_cache=main_cache,
main_rows=[[0, 2], [4, 1]],
scale=scale,
attn_sink=attn_sink,
block_size=block_size,
extra_cache=extra_cache,
extra_rows=[[1], [3, 0]],
)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
@torch.inference_mode()
def test_combine_topk_swa_indices_ragged() -> None:
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
combine_topk_swa_indices_ragged,
)
device = torch.device("cuda")
topk_indices = torch.tensor(
[
[100, 101, 102, 103],
[110, 111, 112, 113],
[120, 121, 122, 123],
[130, 131, 132, 133],
[140, 141, 142, 143],
],
dtype=torch.int32,
device=device,
)
query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
seq_lens = torch.tensor([6, 4], dtype=torch.int32, device=device)
gather_lens = torch.tensor([4, 3], dtype=torch.int32, device=device)
window_size = 3
compress_ratio = 2
topk = 4
M = 20
N = 8
actual_ragged, actual_indptr, actual_lens = combine_topk_swa_indices_ragged(
topk_indices,
query_start_loc,
seq_lens,
gather_lens,
window_size,
compress_ratio,
topk,
M,
N,
)
expected_ragged, expected_indptr, expected_lens = _ref_combine_topk_swa_ragged(
device
)
torch.testing.assert_close(
actual_ragged[: expected_ragged.numel()], expected_ragged
)
torch.testing.assert_close(actual_indptr, expected_indptr)
torch.testing.assert_close(actual_lens, expected_lens)
+141 -7
View File
@@ -3,11 +3,12 @@
"""
Round-trip tests for compressor FP8 quant + KV cache insert gather + dequant.
Four test functions cover five paths:
These tests cover:
A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
B) Indexer: head_dim=128 (all FP8), quant_block=128
C) DeepseekV4 Attention magnitude range: correctness across small/large values
D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
B) Fused dequant+gather K cache
C) Indexer: head_dim=128 (all FP8), quant_block=128
D) DeepseekV4 Attention magnitude range: correctness across small/large values
E) Indexer fused Triton kernel: compress+norm+rope+quant+insert
"""
import math
@@ -134,7 +135,140 @@ def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size
)
# ── Test B: Indexer path ────────────────────────────────────────────────────
# ── Test B: Fused dequant+gather K cache ────────────────────────────────────
def _dequantize_and_gather_k_cache_reference(
out: torch.Tensor,
k_cache: torch.Tensor,
seq_lens: torch.Tensor,
gather_lens: torch.Tensor | None,
block_table: torch.Tensor,
block_size: int,
offset: int,
) -> None:
fp8_dim = 448
bf16_dim = 64
scale_dim = 8
quant_block = 64
token_data_size = fp8_dim + bf16_dim * 2
for req_id in range(seq_lens.shape[0]):
seq_len = seq_lens[req_id].item()
gather_len = gather_lens[req_id].item() if gather_lens is not None else seq_len
start_pos = seq_len - gather_len
for i in range(gather_len):
pos = start_pos + i
pos_in_block = pos % block_size
block_idx = block_table[req_id, pos // block_size].item()
cache_block = k_cache[block_idx].view(-1)
token_data_start = pos_in_block * token_data_size
fp8_bytes = cache_block[token_data_start : token_data_start + fp8_dim]
fp8_vals = fp8_bytes.view(torch.float8_e4m3fn).float()
scale_start = block_size * token_data_size + pos_in_block * scale_dim
encoded_scales = cache_block[scale_start : scale_start + scale_dim]
scales = torch.exp2(encoded_scales[:7].float() - 127.0)
dequant = fp8_vals * scales.repeat_interleave(quant_block)
bf16_start = token_data_start + fp8_dim
bf16_bytes = cache_block[bf16_start : bf16_start + bf16_dim * 2]
bf16_tail = bf16_bytes.view(torch.bfloat16)
out[req_id, offset + i, :fp8_dim] = dequant
out[req_id, offset + i, fp8_dim:] = bf16_tail
@pytest.mark.parametrize(
("seq_lens_host", "gather_lens_host", "offset"),
[
([9, 23, 7], None, 0),
([19, 8, 257], [6, 8, 129], 5),
],
)
def test_dequantize_and_gather_k_cache(
seq_lens_host: list[int],
gather_lens_host: list[int] | None,
offset: int,
):
block_size = 64
head_dim = 512
nope_dim = 448
scale_dim = 8
head_bytes = nope_dim + (head_dim - nope_dim) * 2 + scale_dim
device = "cuda"
num_reqs = len(seq_lens_host)
num_tokens = sum(seq_lens_host)
max_gather_len = max(gather_lens_host or seq_lens_host)
max_blocks_per_seq = math.ceil(max(seq_lens_host) / block_size)
num_blocks = sum(math.ceil(seq_len / block_size) for seq_len in seq_lens_host)
compressed_kv = torch.randn(
num_tokens, head_dim, dtype=torch.bfloat16, device=device
)
# Randomize physical pages so the test covers block-table translation.
# Keep padded block-table entries invalid to catch accidental reads.
physical_blocks = torch.randperm(num_blocks, device=device)
block_table = torch.full(
(num_reqs, max_blocks_per_seq), int(-1e6), dtype=torch.int32, device=device
)
start = 0
for req_id, seq_len in enumerate(seq_lens_host):
num_req_blocks = math.ceil(seq_len / block_size)
req_blocks = physical_blocks[start : start + num_req_blocks]
block_table[req_id, :num_req_blocks] = req_blocks
start += num_req_blocks
# Build slot_mapping for quantize_and_insert_k_cache.
slot_mapping = torch.empty(num_tokens, dtype=torch.int64, device=device)
start = 0
for req_id, seq_len in enumerate(seq_lens_host):
logical_pos = torch.arange(seq_len, dtype=torch.int64, device=device)
block_idx = block_table[req_id, logical_pos // block_size].to(torch.int64)
token_slots = block_idx * block_size + logical_pos % block_size
slot_mapping[start : start + seq_len] = token_slots
start += seq_len
# Insert compressed K into the paged cache layout used by the gather op.
k_cache = torch.empty(
num_blocks, block_size, head_bytes, dtype=torch.uint8, device=device
)
k_cache_2d = k_cache.view(num_blocks, -1)
quantize_and_insert_k_cache(compressed_kv, k_cache_2d, slot_mapping, block_size)
out_shape = (num_reqs, offset + max_gather_len + 3, head_dim)
ref_out = torch.empty(out_shape, dtype=torch.bfloat16, device=device)
actual_out = torch.empty_like(ref_out)
seq_lens = torch.tensor(seq_lens_host, dtype=torch.int32, device=device)
gather_lens = (
torch.tensor(gather_lens_host, dtype=torch.int32, device=device)
if gather_lens_host is not None
else None
)
# Compare production gather against a PyTorch reference for valid output rows.
_dequantize_and_gather_k_cache_reference(
ref_out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
)
dequantize_and_gather_k_cache(
actual_out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
)
torch.accelerator.synchronize()
# only check non-padded content
for req_id, seq_len in enumerate(seq_lens_host):
gather_len = (
gather_lens_host[req_id] if gather_lens_host is not None else seq_len
)
actual = actual_out[req_id, offset : offset + gather_len]
expected = ref_out[req_id, offset : offset + gather_len]
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
# ── Test C: Indexer path ────────────────────────────────────────────────────
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
@@ -254,7 +388,7 @@ def test_indexer_gather_accepts_upper_bound_output():
assert torch.all(dst_scale[valid_tokens:] == sentinel)
# ── Test C: DeepseekV4 attention with values at different magnitudes ───────────
# ── Test D: DeepseekV4 attention with values at different magnitudes ───────────
def test_deepseek_v4_quant_magnitude_range():
@@ -316,7 +450,7 @@ def test_deepseek_v4_quant_magnitude_range():
)
# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
# ── Test E: Indexer fused K-cache insert (Triton kernels) ────────────────────
#
# Both kernels share the same Triton signature; use_fp4 selects between them.
# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
+9
View File
@@ -1,9 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import vllm
import vllm.config
from vllm.lora.request import LoRARequest
from vllm.platforms import current_platform
from ..utils import create_new_process_for_each_test, multi_gpu_test
@@ -50,6 +53,9 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
return generated_texts
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
@create_new_process_for_each_test()
def test_chatglm3_lora(chatglm3_lora_files):
llm = vllm.LLM(
@@ -70,6 +76,9 @@ def test_chatglm3_lora(chatglm3_lora_files):
assert output2[i] == EXPECTED_LORA_OUTPUT[i]
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
@multi_gpu_test(num_gpus=4)
def test_chatglm3_lora_tp4(chatglm3_lora_files):
llm = vllm.LLM(
+10
View File
@@ -11,6 +11,7 @@ import pytest
from huggingface_hub import snapshot_download
from vllm.lora.request import LoRARequest
from vllm.platforms import current_platform
from ..conftest import AudioTestAssets, VllmRunner
from ..utils import create_new_process_for_each_test
@@ -76,6 +77,9 @@ def test_active_default_mm_lora(
)
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
@create_new_process_for_each_test()
def test_inactive_default_mm_lora(
vllm_runner: type[VllmRunner],
@@ -92,6 +96,9 @@ def test_inactive_default_mm_lora(
)
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
@create_new_process_for_each_test()
def test_default_mm_lora_succeeds_with_redundant_lora_request(
vllm_runner: type[VllmRunner],
@@ -107,6 +114,9 @@ def test_default_mm_lora_succeeds_with_redundant_lora_request(
)
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
@create_new_process_for_each_test()
def test_default_mm_lora_fails_with_overridden_lora_request(
vllm_runner: type[VllmRunner],
+4
View File
@@ -10,6 +10,7 @@ import vllm.config
from vllm import LLM
from vllm.lora.request import LoRARequest
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
from vllm.platforms import current_platform
from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test
@@ -139,6 +140,9 @@ def test_llama_lora(llama32_lora_files, cudagraph_specialize_lora: bool):
generate_and_test(llm, llama32_lora_files)
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
@multi_gpu_test(num_gpus=4)
def test_llama_lora_tp4(llama32_lora_files):
llm = vllm.LLM(
+3
View File
@@ -68,6 +68,9 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
return generated_texts
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
def test_minicpmv_lora(minicpmv_lora_files):
llm = vllm.LLM(
MODEL_PATH,
+7
View File
@@ -11,6 +11,7 @@ from safetensors.torch import load_file, save_file
import vllm
from vllm.lora.request import LoRARequest
from vllm.platforms import current_platform
from ..utils import multi_gpu_test
@@ -110,6 +111,9 @@ def generate_and_test(
)
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
def test_olmoe_lora(olmoe_lora_files, maybe_enable_lora_dual_stream):
# We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
# Otherwise, the lora-test will fail due to CUDA OOM.
@@ -178,6 +182,9 @@ def test_olmoe_lora_mixed_random(
assert outputs[0].outputs[0].text.strip().startswith(EXPECTED_LORA_OUTPUT[0])
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
@pytest.mark.parametrize("fully_sharded_loras", [False, True])
@multi_gpu_test(num_gpus=2)
def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
@@ -1,12 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from transformers import AutoTokenizer
import vllm
import vllm.config
from vllm.assets.image import ImageAsset
from vllm.lora.request import LoRARequest
from vllm.platforms import current_platform
from ..utils import create_new_process_for_each_test, multi_gpu_test
@@ -311,6 +313,9 @@ def _assert_qwen35_text_vl_and_mixed_lora(
)
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
@create_new_process_for_each_test()
def test_qwen35_text_lora(
qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream
+8
View File
@@ -2,12 +2,14 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
import pytest
from packaging.version import Version
from transformers import __version__ as TRANSFORMERS_VERSION
import vllm
from vllm.assets.image import ImageAsset
from vllm.lora.request import LoRARequest
from vllm.platforms import current_platform
from vllm.sampling_params import BeamSearchParams
@@ -206,6 +208,9 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
)
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
def test_qwen25vl_lora(qwen25vl_lora_files):
"""Test Qwen 2.5 VL model with LoRA"""
config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files)
@@ -216,6 +221,9 @@ def test_qwen25vl_lora(qwen25vl_lora_files):
tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
@pytest.mark.skipif(
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
)
def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files):
config = TestConfig(
model_path=QWEN25VL_MODEL_PATH,
-27
View File
@@ -124,30 +124,3 @@ def test_whisper_multi_lora(whisper_lora_files):
f"Expected same outputs for same adapter with different IDs. "
f"Got: {outputs_lora1} vs {outputs_lora2}"
)
@create_new_process_for_each_test()
def test_whisper_with_and_without_lora(whisper_lora_files):
"""Test that Whisper produces different outputs with and without LoRA.
This test verifies that the LoRA adapter actually affects the model output.
"""
llm = create_whisper_llm(enable_lora=True)
# Run with LoRA
outputs_with_lora = run_whisper_inference(
llm, lora_path=whisper_lora_files, lora_id=1
)
# Run without LoRA (base model only)
outputs_without_lora = run_whisper_inference(llm, lora_path=None)
# Both should produce valid outputs
assert len(outputs_with_lora[0]) > 0
assert len(outputs_without_lora[0]) > 0
print(f"Output with LoRA: {outputs_with_lora[0]}")
print(f"Output without LoRA: {outputs_without_lora[0]}")
# Note: Outputs may or may not differ depending on the adapter
# The main verification is that both configurations work
@@ -4,9 +4,12 @@
from collections.abc import Mapping
import pytest
import torch
from PIL import Image as PILImage
from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import MultiModalFieldConfig
from ....conftest import ImageTestAssets
from ...utils import build_model_context
@@ -15,6 +18,36 @@ from ...utils import build_model_context
GEMMA4_MODEL_ID = "google/gemma-4-E2B-it"
def test_gemma4_image_schema_accepts_variable_patch_counts():
Gemma4ImagePixelInputs(
pixel_values=[
torch.randn(10080, 768),
torch.randn(2520, 768),
],
pixel_position_ids=[
torch.zeros(10080, 2, dtype=torch.long),
torch.zeros(2520, 2, dtype=torch.long),
],
)
def test_gemma4_image_batching_keeps_variable_patch_counts_unstacked():
field = MultiModalFieldConfig.batched("image").field
elems = field.build_elems(
"image",
"pixel_values",
[torch.randn(10080, 768), torch.randn(2520, 768)],
)
reduced = field.reduce_data(list(elems))
assert isinstance(reduced, list)
assert [tensor.shape for tensor in reduced] == [
torch.Size([10080, 768]),
torch.Size([2520, 768]),
]
@pytest.mark.parametrize(
"image_width,image_height,max_soft_tokens",
[
@@ -53,6 +53,19 @@ class _VisionModel:
self.loaded_weights = list(weights)
class _FakeTensor:
"""Sentinel stand-in for torch.Tensor in load_weights tests. Supports the
.detach().clone() chain used by load_weights for buffered mm weights;
both methods return self so identity (and the existing equality
assertions) are preserved through cloning."""
def detach(self):
return self
def clone(self):
return self
def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode():
model = object.__new__(NemotronH_Nano_VL_V2)
language_model = _LanguageModel()
@@ -86,7 +99,7 @@ def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder():
object.__setattr__(model, "sound_encoder", None)
language_weight = object()
vision_weight = object()
vision_weight = _FakeTensor()
model.load_weights(
[
("language_model.layers.0.weight", language_weight),
+8
View File
@@ -28,6 +28,14 @@ def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch):
assert not hasattr(envs.__getattr__, "cache_info")
def test_nixl_side_channel_host_is_not_compile_factor(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setenv("VLLM_NIXL_SIDE_CHANNEL_HOST", "10.0.0.15")
assert "VLLM_NIXL_SIDE_CHANNEL_HOST" not in envs.compile_factors()
def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("VLLM_HOST_IP", "1.1.1.1")
monkeypatch.setenv("VLLM_PORT", "1234")
+14
View File
@@ -6,6 +6,7 @@ import os
from unittest.mock import patch
from vllm.ray.ray_env import get_env_vars_to_copy
from vllm.v1.executor.ray_utils import WORKER_SPECIFIC_ENV_VARS
# ---------------------------------------------------------------------------
# Default prefix matching
@@ -106,6 +107,19 @@ class TestExclusion:
result = get_env_vars_to_copy(exclude_vars={"CUDA_VISIBLE_DEVICES"})
assert "CUDA_VISIBLE_DEVICES" not in result
@patch.dict(
os.environ,
{
"VLLM_HOST_IP": "10.0.0.1",
"VLLM_NIXL_SIDE_CHANNEL_HOST": "10.0.0.1",
},
clear=False,
)
def test_worker_specific_host_vars_are_excluded(self):
result = get_env_vars_to_copy(exclude_vars=WORKER_SPECIFIC_ENV_VARS)
assert "VLLM_HOST_IP" not in result
assert "VLLM_NIXL_SIDE_CHANNEL_HOST" not in result
@patch.dict(os.environ, {"LMCACHE_LOCAL_CPU": "True"}, clear=False)
@patch(
"vllm.ray.ray_env.RAY_NON_CARRY_OVER_ENV_VARS",
@@ -0,0 +1,68 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import vllm.envs as envs
from tests.utils import TestFP8Layer, requires_fp8
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
CutlassFP8ScaledMMLinearKernel,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
kFp8DynamicTokenSym,
kFp8StaticTensorSym,
)
from vllm.platforms import current_platform
pytest.importorskip("torch.cuda")
@pytest.fixture(autouse=True)
def setup_cuda():
if not current_platform.is_cuda():
pytest.skip("CUTLASS FP8 kernels require CUDA.")
torch.set_default_device("cuda")
@requires_fp8
@pytest.mark.parametrize("weight_shape", [(1024, 2048), (4608, 4096)])
@pytest.mark.parametrize("batch_size", [1, 16, 17, 32, 64, 65, 256, 257])
@torch.inference_mode()
def test_cutlass_fp8_batch_invariant_fixed_config(
weight_shape: tuple[int, int],
batch_size: int,
default_vllm_config,
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
torch.manual_seed(0)
layer = TestFP8Layer(
weight_shape=weight_shape,
activation_quant_key=kFp8DynamicTokenSym,
weight_quant_key=kFp8StaticTensorSym,
input_dtype=torch.bfloat16,
out_dtype=torch.bfloat16,
device=torch.device("cuda"),
force_kernel=CutlassFP8ScaledMMLinearKernel,
)
assert isinstance(layer.kernel, CutlassFP8ScaledMMLinearKernel)
in_features = weight_shape[1]
needle = torch.randn((1, in_features), device="cuda", dtype=torch.bfloat16)
baseline = layer(needle)[0]
filler = torch.randn(
(max(batch_size - 1, 0), in_features), device="cuda", dtype=torch.bfloat16
)
front_batch = torch.cat([needle, filler], dim=0)
back_batch = torch.cat([filler, needle], dim=0)
front_output = layer(front_batch)[0]
back_output = layer(back_batch)[-1]
torch.testing.assert_close(front_output, baseline, rtol=0, atol=0)
torch.testing.assert_close(back_output, baseline, rtol=0, atol=0)
+13 -1
View File
@@ -488,6 +488,10 @@ def _run_eagle_correctness(
@single_gpu_only
@pytest.mark.skipif(
current_platform.is_device_capability_family(100),
reason="DeepSeek head_dim=192 not supported on SM100/SM110 (Blackwell)",
)
@pytest.mark.parametrize(
[
"model_setup",
@@ -718,7 +722,15 @@ def test_eagle_correctness_heavy(
["model_setup", "mm_enabled", "expected_accuracy_threshold"],
[
(("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70%
(("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1), False, 0.0), # dummy model
pytest.param(
("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1),
False,
0.0,
marks=pytest.mark.skipif(
current_platform.is_device_capability_family(100),
reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell",
),
), # dummy model
(
("mtp", "Qwen/Qwen3.5-0.8B-Base", 1),
False,
@@ -0,0 +1,136 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
import uuid
from pathlib import Path
from types import SimpleNamespace
from typing import Any
import pytest
import ray
from vllm.v1.engine.core import EngineCoreActorMixin
from vllm.v1.engine.utils import CoreEngineActorManager, EngineZmqAddresses
class _StubEngineCoreActor(EngineCoreActorMixin):
def __init__(
self,
vllm_config: Any,
local_client: bool,
addresses: EngineZmqAddresses,
executor_class: type[Any],
log_stats: bool,
dp_rank: int = 0,
local_dp_rank: int = 0,
):
# Exercise the production Ray actor mixin without loading a model.
EngineCoreActorMixin.__init__(
self, vllm_config, addresses, dp_rank, local_dp_rank
)
def _set_visible_devices(self, vllm_config: Any, local_dp_rank: int) -> None:
pass
def wait_for_init(self) -> None:
pass
def run(self) -> None:
pass
def get_nixl_side_channel_host(self) -> str | None:
return os.environ.get("VLLM_NIXL_SIDE_CHANNEL_HOST")
class _DummyExecutor:
pass
def _make_vllm_config() -> SimpleNamespace:
return SimpleNamespace(
parallel_config=SimpleNamespace(
data_parallel_size=1,
data_parallel_size_local=1,
enable_elastic_ep=False,
world_size=1,
),
model_config=SimpleNamespace(is_moe=False),
kv_transfer_config=None,
)
def _make_addresses() -> EngineZmqAddresses:
return EngineZmqAddresses(
inputs=["tcp://127.0.0.1:12345"],
outputs=["tcp://127.0.0.1:12346"],
)
def _make_cpu_placement_group():
pg = ray.util.placement_group(
[{"CPU": 0.001}, {"CPU": 1.0}],
strategy="PACK",
)
ray.get(pg.ready())
return pg
@pytest.fixture
def ray_context():
started_ray = False
if not ray.is_initialized():
project_root = str(Path(__file__).resolve().parents[3])
ray.init(
num_cpus=2,
runtime_env={"env_vars": {"PYTHONPATH": project_root}},
log_to_driver=False,
)
started_ray = True
yield
if started_ray:
ray.shutdown()
@pytest.mark.usefixtures("ray_context")
def test_driver_nixl_side_channel_host_does_not_leak_to_engine_core_actor(
monkeypatch: pytest.MonkeyPatch,
) -> None:
driver_marker = f"driver-only-nixl-host-{uuid.uuid4()}"
created_placement_groups: list[Any] = []
manager: CoreEngineActorManager | None = None
def create_dp_placement_groups(vllm_config: Any):
pg = _make_cpu_placement_group()
created_placement_groups.append(pg)
return [pg], [0]
monkeypatch.setenv("VLLM_NIXL_SIDE_CHANNEL_HOST", driver_marker)
monkeypatch.setattr("vllm.v1.engine.core.EngineCoreActor", _StubEngineCoreActor)
monkeypatch.setattr(
CoreEngineActorManager,
"create_dp_placement_groups",
staticmethod(create_dp_placement_groups),
)
try:
manager = CoreEngineActorManager(
vllm_config=_make_vllm_config(),
addresses=_make_addresses(),
executor_class=_DummyExecutor,
log_stats=False,
)
actor = manager.local_engine_actors[0]
actor_host = ray.get(actor.get_nixl_side_channel_host.remote())
node_host = ray.util.get_node_ip_address()
assert actor_host != driver_marker
assert actor_host == node_host
finally:
if manager is not None:
manager.shutdown()
else:
for pg in created_placement_groups:
ray.util.remove_placement_group(pg)
@@ -0,0 +1,66 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for LogprobsProcessor.
These tests exercise the truncation invariant that the MRV2 sampler relies
on: when the sampler returns a row wider than a request's own
`num_logprobs + 1` (because another request in the batch needed a wider
row), the trailing positions are populated with sentinel values
(`token_id=0`, `logprob=-inf`). LogprobsProcessor must read only the first
`num_logprobs + 1` entries so those sentinels never reach the user.
"""
import numpy as np
from vllm.logprobs import create_sample_logprobs
from vllm.v1.engine.logprobs import LogprobsProcessor
from vllm.v1.outputs import LogprobsLists
def _make_processor(num_logprobs: int) -> LogprobsProcessor:
return LogprobsProcessor(
tokenizer=None,
logprobs=create_sample_logprobs(flat_logprobs=False),
prompt_logprobs=None,
cumulative_logprob=0.0,
num_logprobs=num_logprobs,
num_prompt_logprobs=None,
)
def test_drops_trailing_sentinel_columns():
"""A request that asked for 3 custom token logprobs but ended up in a
batch padded to width 5 must not surface the trailing -inf entries."""
processor = _make_processor(num_logprobs=3)
sampled = 42
# Layout: [sampled, custom_1, custom_2, custom_3, SENTINEL, SENTINEL]
# Use float32-exact values so cumulative_logprob compares cleanly.
token_ids = np.array([[sampled, 100, 200, 300, 0, 0]], dtype=np.int32)
logprobs = np.array([[-0.5, -1.0, -2.0, -3.0, -np.inf, -np.inf]], dtype=np.float32)
ranks = np.array([1], dtype=np.int32)
processor._update_sample_logprobs(LogprobsLists(token_ids, logprobs, ranks))
assert len(processor.logprobs) == 1
pos = processor.logprobs[0]
# Exactly sampled + 3 requested tokens; trailing sentinels dropped.
assert set(pos.keys()) == {sampled, 100, 200, 300}
assert 0 not in pos
assert all(np.isfinite(lp.logprob) for lp in pos.values())
# cumulative_logprob comes from the sampled token's logprob only.
assert processor.cumulative_logprob == -0.5
def test_accepts_exactly_sized_row():
"""When the row is exactly num_logprobs+1, no truncation needed."""
processor = _make_processor(num_logprobs=2)
token_ids = np.array([[7, 11, 13]], dtype=np.int32)
logprobs = np.array([[-0.5, -1.5, -2.5]], dtype=np.float32)
ranks = np.array([1], dtype=np.int32)
processor._update_sample_logprobs(LogprobsLists(token_ids, logprobs, ranks))
pos = processor.logprobs[0]
assert set(pos.keys()) == {7, 11, 13}
@@ -262,9 +262,11 @@ def test_multi_example_connector_consistency():
events = get_connector_events()
# First event is set_xfer_handshake_metadata from initialization, then
# on_new_request when the request is enqueued, then
# get_num_new_matched_tokens and update_state_after_alloc from generate().
assert events["storage1-SCHEDULER"][:4] == [
assert events["storage1-SCHEDULER"][:5] == [
"set_xfer_handshake_metadata",
"on_new_request",
"get_num_new_matched_tokens 0",
"update_state_after_alloc num_blocks=[0] 0",
"build_connector_meta",
@@ -281,8 +283,9 @@ def test_multi_example_connector_consistency():
"wait_for_layer_load",
"save_kv_layer",
]
assert events["storage2-SCHEDULER"][:4] == [
assert events["storage2-SCHEDULER"][:5] == [
"set_xfer_handshake_metadata",
"on_new_request",
"get_num_new_matched_tokens 0",
"update_state_after_alloc num_blocks=[0] 0",
"build_connector_meta",
@@ -310,12 +313,14 @@ def test_multi_example_connector_consistency():
# connector so update_state_after_alloc will be with allocated blocks
# on that one but with zero blocks for others (first nonzero match is
# chosen).
assert events["storage1-SCHEDULER"][:3] == [
assert events["storage1-SCHEDULER"][:4] == [
"on_new_request",
"get_num_new_matched_tokens 0",
"update_state_after_alloc num_blocks=[7] 96",
"build_connector_meta",
]
assert events["storage2-SCHEDULER"][:3] == [
assert events["storage2-SCHEDULER"][:4] == [
"on_new_request",
"get_num_new_matched_tokens 0",
"update_state_after_alloc num_blocks=[0] 0",
"build_connector_meta",
@@ -336,12 +341,14 @@ def test_multi_example_connector_consistency():
# return 0 from the first connector, but the second connector should have
# a hit, so update_state_after_alloc will only be called with allocated
# blocks for the second connector.
assert events["storage1-SCHEDULER"][:3] == [
assert events["storage1-SCHEDULER"][:4] == [
"on_new_request",
"get_num_new_matched_tokens 0",
"update_state_after_alloc num_blocks=[0] 0",
"build_connector_meta",
]
assert events["storage2-SCHEDULER"][:3] == [
assert events["storage2-SCHEDULER"][:4] == [
"on_new_request",
"get_num_new_matched_tokens 0",
"update_state_after_alloc num_blocks=[7] 96",
"build_connector_meta",
@@ -1374,9 +1374,11 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
| {eventually free blocks}
"""
model_name = "Qwen/Qwen3-0.6B"
timeout = 6
kv_transfer_config = KVTransferConfig(
kv_connector="NixlConnector",
kv_role="kv_both",
kv_connector_extra_config={"kv_lease_duration": timeout},
)
llm_kwargs = {
"model": model_name,
@@ -1386,9 +1388,7 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
"distributed_executor_backend": distributed_executor_backend,
}
timeout = 6
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
monkeypatch.setenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", str(timeout))
def run_test_and_cleanup():
llm = LLM(**llm_kwargs)
@@ -1403,8 +1403,6 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
runtime_env = {
"working_dir": working_dir, # ship fake nixl package
"env_vars": {
"VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout),
# TODO: for ray to carry over, remove once we set
"NIXL_TELEMETRY_ENABLE": "1",
},
}
@@ -2269,6 +2267,127 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
assert request_id in done_recving
@patch(
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
FailingNixlWrapper,
)
@pytest.mark.parametrize(
"failure_mode",
[
"handshake",
"transfer_setup",
"transfer_failed",
"transfer_exception",
],
)
def test_failed_request_skips_kv_postprocessing(
default_vllm_config, dist_init, failure_mode
):
"""Test that failed requests skip KV sync and post-processing in
get_finished().
This is the core safety behavior: when a KV transfer fails at any stage,
the request must still appear in done_recving (so the scheduler can apply
kv_load_failure_policy), but sync_recved_kv_to_device and post-processing
must NOT be called since no valid KV data was received.
Covers all failure paths that involve an actual (attempted) KV transfer:
- handshake: add_remote_agent raises during async handshake
- transfer_setup: make_prepped_xfer raises before handle is in transfers
- transfer_failed: check_xfer_state returns bad state ("ERR") in
_pop_done_transfers this is the path that previously had the bug
where post-processing was NOT skipped
- transfer_exception: check_xfer_state raises in _pop_done_transfers
Note: notification_failed (send_notif raises on the full-cache-hit path)
is intentionally excluded. That path is a best-effort DP courtesy
notification; the blocks are already in D's cache, so no KV transfer
was attempted and done_recving is correctly empty.
"""
# Map each failure mode to the FailingNixlWrapper attribute to set.
_WRAPPER_CONFIG: dict[str, str] = {
"handshake": "fail_handshake",
"transfer_setup": "fail_transfer_setup",
"transfer_failed": "fail_transfer_state",
"transfer_exception": "fail_transfer_exception",
}
# Use enable_permute_local_kv=True so that
# post_process_device_kv_on_receive would be called on the success path,
# making the assertion meaningful (not trivially true).
vllm_config = create_vllm_config(enable_permute_local_kv=True)
connector = NixlConnector(
vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
)
connector.connector_worker = FakeNixlConnectorWorker(
vllm_config,
connector.engine_id,
hand_shake_latency=0.1 if failure_mode == "handshake" else 0,
)
worker = connector.connector_worker
setattr(worker.nixl_wrapper, _WRAPPER_CONFIG[failure_mode], True)
request_id = f"test_{failure_mode}_skip_postprocess"
metadata = NixlConnectorMetadata()
metadata.add_new_req_to_recv(
request_id=request_id,
local_block_ids=([1, 2, 3],),
kv_transfer_params={
"remote_block_ids": ([4, 5, 6],),
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
"remote_request_id": f"prefill-{request_id}",
"remote_host": "localhost",
"remote_port": 1234,
"remote_tp_size": 1,
},
)
connector.bind_connector_metadata(metadata)
dummy_ctx = ForwardContext(
no_compile_layers={},
attn_metadata={},
slot_mapping={},
)
connector.start_load_kv(dummy_ctx)
if failure_mode == "handshake":
# Wait for async handshake to fail.
time.sleep(0.3)
else:
# All other modes: let the handshake complete, then process the
# ready_requests queue. For transfer_failed / transfer_exception the
# handle ends up in _recving_transfers; the failure surfaces in
# get_finished() via _pop_done_transfers below.
connector.bind_connector_metadata(NixlConnectorMetadata())
time.sleep(0.1)
connector.start_load_kv(dummy_ctx)
# Spy on sync_recved_kv_to_device and post_process_device_kv_on_receive
# to verify they are NOT called for the failed request.
with (
patch.object(worker, "sync_recved_kv_to_device") as mock_sync,
patch.object(worker, "post_process_device_kv_on_receive") as mock_postprocess,
):
_, done_recving = connector.get_finished(finished_req_ids=set())
# The failed request must appear in done_recving so the scheduler
# can handle it (e.g., trigger recompute via kv_load_failure_policy).
assert request_id in done_recving
# Critical: KV sync and post-processing must NOT have been called
# since no valid KV data was received for the failed request.
mock_sync.assert_not_called()
mock_postprocess.assert_not_called()
# Metadata for the request should have been cleaned up.
assert request_id not in worker._recving_metadata
# Blocks should have been marked as invalid.
invalid_blocks = connector.get_block_ids_with_load_errors()
assert invalid_blocks == {1, 2, 3}
@pytest.mark.parametrize(
"mismatch_type,config_overrides,version_override,should_fail,enforce_handshake_compat",
[
@@ -0,0 +1,165 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the scheduler-driven heartbeat / lease-renewal system."""
import time
from unittest.mock import MagicMock
import pytest
from vllm.v1.outputs import KVConnectorOutput
from .utils import create_request, make_nixl_scheduler
_ENGINE_A = "my-engine-id"
def _sched(kv_lease_duration: int = 30):
return make_nixl_scheduler(heartbeat=True, kv_lease_duration=kv_lease_duration)
def _req(request_id: int = 1):
return create_request(request_id=request_id, do_remote_prefill=True)
def _worker_stub():
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
)
w = object.__new__(NixlConnectorWorker)
w._reqs_to_send = {}
w._lease_extension = 20
return w
# ===================================================================
# Scheduler: on_new_request
# ===================================================================
def test_on_new_request_tracks_and_groups():
"""Add two reqs to same engine, one to another; verify grouping."""
s = _sched()
s.on_new_request(_req(1))
s.on_new_request(_req(2))
assert s._heartbeat_by_engine[_ENGINE_A].req_ids == {"prefill-1", "prefill-2"}
info = s._heartbeat_by_engine[_ENGINE_A]
assert (info.host, info.port, info.tp_size) == ("my-host", 1234, 1)
assert s._heartbeat_req_engine["id-1"] == (_ENGINE_A, "prefill-1")
# Different engine.
r3 = _req(3)
r3.kv_transfer_params["remote_engine_id"] = "engine-b"
s.on_new_request(r3)
assert len(s._heartbeat_by_engine) == 2
@pytest.mark.parametrize(
"make_req",
[
lambda: create_request(request_id=2, do_remote_decode=True),
lambda: create_request(request_id=3), # no kv_transfer_params
],
ids=["decode", "plain"],
)
def test_on_new_request_ignores_non_prefill(make_req):
s = _sched()
s.on_new_request(make_req())
assert len(s._heartbeat_by_engine) == 0
# ===================================================================
# Scheduler: _stop_heartbeat
# ===================================================================
def test_stop_heartbeat_partial_and_full():
"""Stop one of two reqs on same engine, then stop the other."""
s = _sched()
s.on_new_request(_req(1))
s.on_new_request(_req(2))
s._stop_heartbeat("id-1")
assert s._heartbeat_by_engine[_ENGINE_A].req_ids == {"prefill-2"}
assert "id-1" not in s._heartbeat_req_engine
s._stop_heartbeat("id-2")
assert len(s._heartbeat_by_engine) == 0
assert len(s._heartbeat_req_engine) == 0
# ===================================================================
# Scheduler: build_connector_meta throttling
# ===================================================================
def test_build_connector_meta_heartbeat_throttling():
# kv_lease_duration=30 => _heartbeat_interval = 30 // 6 = 5
s = _sched(kv_lease_duration=30)
s.on_new_request(_req(1))
# Ensure the first call triggers by placing last_heartbeat far in the past.
s._last_heartbeat_time = time.perf_counter() - 10
meta1 = s.build_connector_meta(MagicMock())
assert _ENGINE_A in meta1.heartbeat_by_engine
# Immediate second call is throttled (< 5s since last).
meta2 = s.build_connector_meta(MagicMock())
assert len(meta2.heartbeat_by_engine) == 0
# ===================================================================
# Scheduler: cleanup paths (update_connector_output / request_finished)
# ===================================================================
def test_update_connector_output_stops_heartbeat():
s = _sched()
s.on_new_request(_req(1))
s.update_connector_output(
KVConnectorOutput(
finished_sending=None,
finished_recving={"id-1"},
invalid_block_ids=set(),
)
)
assert len(s._heartbeat_by_engine) == 0
assert len(s._heartbeat_req_engine) == 0
def test_request_finished_stops_heartbeat():
s = _sched()
r = _req(1)
s.on_new_request(r)
# Simulate update_state_after_alloc having consumed do_remote_prefill.
r.kv_transfer_params["do_remote_prefill"] = False
s.request_finished(r, block_ids=())
assert len(s._heartbeat_by_engine) == 0
assert len(s._heartbeat_req_engine) == 0
# ===================================================================
# Worker: _handle_heartbeat
# ===================================================================
def test_handle_heartbeat():
w = _worker_stub()
far_future = time.perf_counter() + 99999
w._reqs_to_send = {"req-a": 100.0, "req-b": far_future}
before = time.perf_counter()
w._handle_heartbeat("req-a,req-b,req-unknown")
# req-a: pushed forward to ~now+20.
assert w._reqs_to_send["req-a"] >= before + 20
# req-b: already far out, max() keeps it.
assert w._reqs_to_send["req-b"] >= far_future
# req-unknown: not added.
assert "req-unknown" not in w._reqs_to_send
+28 -2
View File
@@ -223,6 +223,7 @@ def create_request(
remote_block_ids=list(range(num_remote_blocks)),
remote_host="my-host",
remote_port=1234,
tp_size=1,
)
max_tokens = 1 if do_remote_decode else max_tokens
@@ -482,10 +483,16 @@ def make_kv_cache_config(
)
def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False):
def make_nixl_scheduler(
has_mamba: bool = False,
is_hma_required: bool = False,
heartbeat: bool = False,
kv_lease_duration: int = 30,
):
"""Create a NixlConnectorScheduler via __new__ (skipping __init__).
Only sets the two flags needed by the N-1 prefill logic.
Only sets the flags needed by the tests. When *heartbeat=True* the
scheduler-side heartbeat bookkeeping fields are also initialised.
"""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
NixlConnectorScheduler,
@@ -494,4 +501,23 @@ def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False):
sched = object.__new__(NixlConnectorScheduler)
sched._has_mamba = has_mamba
sched._is_hma_required = is_hma_required
if heartbeat:
sched._heartbeat_by_engine = {}
sched._heartbeat_req_engine = {}
sched._last_heartbeat_time = 0.0
sched._kv_lease_duration = kv_lease_duration
sched._heartbeat_interval = kv_lease_duration // 6
# Fields touched by build_connector_meta / request_finished:
sched._reqs_need_recv = {}
sched._reqs_need_send = {}
sched._reqs_in_batch = set()
sched._reqs_not_processed = set()
sched._reqs_need_save = {}
sched.use_host_buffer = False
sched.engine_id = "test-engine"
sched.side_channel_host = "localhost"
sched.side_channel_port = 5555
sched.blocks_per_sw = []
sched.is_bidirectional_kv_xfer_enabled = False
return sched
+7 -8
View File
@@ -17,7 +17,6 @@ from vllm.v1.kv_offload.base import (
from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext:
@@ -565,14 +564,14 @@ class TestARCPolicy:
def test_filter_reused_manager():
"""
Tests FilterReusedOffloadingManager with a CPUOffloadingManager.
Tests CPUOffloadingManager reuse filtering (store_threshold=2).
"""
lru_manager = CPUOffloadingManager(
num_blocks=4, cache_policy="lru", enable_events=True
)
manager = FilterReusedOffloadingManager(
backing=lru_manager, store_threshold=2, max_tracker_size=3
manager = CPUOffloadingManager(
num_blocks=4,
cache_policy="lru",
enable_events=True,
store_threshold=2,
max_tracker_size=3,
)
# Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
+6 -2
View File
@@ -9,8 +9,9 @@ from vllm.benchmarks.datasets.datasets import (
BlazeditDataset,
BurstGPTDataset,
ConversationDataset,
CustomAudioDataset,
CustomDataset,
CustomMMDataset,
CustomImageDataset,
HuggingFaceDataset,
InstructCoderDataset,
MLPerfDataset,
@@ -36,6 +37,7 @@ from vllm.benchmarks.datasets.datasets import (
is_valid_sequence,
lora_path_on_disk,
lora_tokenizer_cache,
process_audio,
process_image,
process_video,
zeta_prompt,
@@ -51,7 +53,8 @@ __all__ = [
"BurstGPTDataset",
"ConversationDataset",
"CustomDataset",
"CustomMMDataset",
"CustomAudioDataset",
"CustomImageDataset",
"HuggingFaceDataset",
"InstructCoderDataset",
"MLPerfDataset",
@@ -77,6 +80,7 @@ __all__ = [
"is_valid_sequence",
"lora_path_on_disk",
"lora_tokenizer_cache",
"process_audio",
"process_image",
"process_video",
"RangeRatio",
+150 -4
View File
@@ -60,6 +60,11 @@ try:
except ImportError:
pd = PlaceholderModule("pandas")
try:
import soundfile as sf
except ImportError:
sf = PlaceholderModule("soundfile")
logger = logging.getLogger(__name__)
@@ -441,6 +446,27 @@ def process_video(video: Any) -> Mapping[str, Any]:
)
def process_audio(audio: Any) -> tuple:
"""
Process a single audio input and return a (array, sample_rate) tuple.
Supports:
1. String: treated as a file path, loaded with soundfile.
2. Dict with 'array' and 'sampling_rate' keys: HuggingFace audio format.
3. Tuple (array, sr): passed through directly.
"""
if isinstance(audio, str):
return sf.read(audio)
if isinstance(audio, dict) and "array" in audio and "sampling_rate" in audio:
return audio["array"], audio["sampling_rate"]
if isinstance(audio, tuple) and len(audio) == 2:
return audio
raise ValueError(
f"Invalid audio input {audio}. Must be a file path string, "
"a dict with 'array' and 'sampling_rate', or a (array, sr) tuple."
)
def gen_prompt_decode_to_target_len(
tokenizer: TokenizerLike,
token_sequence: list[int],
@@ -1399,6 +1425,8 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
"random-rerank",
"hf",
"custom",
"custom_audio",
"custom_image",
"custom_mm",
"prefix_repetition",
"spec_bench",
@@ -1816,8 +1844,28 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
no_oversample=args.no_oversample,
)
elif args.dataset_name == "custom_mm":
dataset = CustomMMDataset(
elif args.dataset_name in ("custom_image", "custom_mm"):
if args.dataset_name == "custom_mm":
logger.warning(
"Dataset name 'custom_mm' is deprecated and will be removed in v0.24. "
"Use '--dataset-name custom_image' instead."
)
dataset = CustomImageDataset(
dataset_path=args.dataset_path,
disable_shuffle=args.disable_shuffle,
random_seed=args.seed,
)
input_requests = dataset.sample(
num_requests=args.num_prompts,
tokenizer=tokenizer,
output_len=args.custom_output_len,
enable_multimodal_chat=args.enable_multimodal_chat,
request_id_prefix=args.request_id_prefix,
no_oversample=args.no_oversample,
)
elif args.dataset_name == "custom_audio":
dataset = CustomAudioDataset(
dataset_path=args.dataset_path,
disable_shuffle=args.disable_shuffle,
random_seed=args.seed,
@@ -2249,9 +2297,9 @@ class CustomDataset(BenchmarkDataset):
return sampled_requests
class CustomMMDataset(CustomDataset):
class CustomImageDataset(CustomDataset):
"""
Implements the Custom MultiModal dataset. Loads data from a JSONL file and generates
Implements the Custom image dataset. Loads data from a JSONL file and generates
sample requests based on conversation turns. E.g.,
```
{
@@ -2328,6 +2376,104 @@ class CustomMMDataset(CustomDataset):
return sampled_requests
class CustomAudioDataset(CustomDataset):
"""
Custom dataset for audio benchmarking. Loads data from a JSONL file. E.g.,
{"prompt": "Transcribe the audio.", "audio": "/path/to/audio.wav"}
Supports both:
- Dedicated ASR models (e.g. Whisper) via openai-audio & /v1/audio/transcriptions
- Chat-based audio models (e.g. Qwen2-Audio) via openai-chat & /v1/chat/completions
"""
IS_MULTIMODAL = True
def sample(
self,
tokenizer: TokenizerLike,
num_requests: int,
output_len: int | None = None,
request_id_prefix: str = "",
no_oversample: bool = False,
skip_chat_template: bool = False,
enable_multimodal_chat: bool = False,
**kwargs,
) -> list[SampleRequest]:
self.num_available_samples = len(self.data)
if num_requests <= 0:
num_requests = self.num_available_samples
sampled_requests = []
for i, item in enumerate(self.data):
if len(sampled_requests) >= num_requests:
break
prompt = item.get("prompt", "")
if tokenizer is None:
prompt_len = 1
new_output_len = output_len if output_len not in (None, -1) else 256
mm_content = None
else:
use_chat_template = (
not skip_chat_template
and hasattr(tokenizer, "chat_template")
and tokenizer.chat_template is not None
)
if enable_multimodal_chat:
# Chat-based audio models (e.g., Qwen2-Audio):
# encode audio as base64; serve.py assembles the chat message
# as: {"role": "user", "content": [
# {"type": "text", "text": prompt},
# {"type": "input_audio", "input_audio": {...}}
# ]}
y, sr = process_audio(item["audio"])
buf = io.BytesIO()
sf.write(buf, y, sr, format="WAV")
audio_base64 = base64.b64encode(buf.getvalue()).decode("utf-8")
mm_content = {
"type": "input_audio",
"input_audio": {
"data": audio_base64,
"format": "wav",
},
}
# prompt stays as plain string; serve.py handles wrapping
else:
# Whisper-style models: load audio array locally
y, sr = process_audio(item["audio"])
mm_content = {"audio": (y, sr)}
if use_chat_template:
# ASR models with a chat template but not multimodal chat
prompt = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}],
add_generation_prompt=True,
tokenize=False,
)
# else: plain prompt for Whisper-style models
prompt_len = (
len(tokenizer(prompt).input_ids) if isinstance(prompt, str) else 1
)
new_output_len = output_len
if output_len is None or output_len == -1:
if "output_tokens" not in item:
raise ValueError(
"If no output length is provided the "
"custom dataset must contain an 'output_tokens' field."
)
new_output_len = int(item["output_tokens"])
sampled_requests.append(
SampleRequest(
prompt=prompt,
prompt_len=prompt_len,
expected_output_len=new_output_len,
multi_modal_data=mm_content,
request_id=request_id_prefix + str(i),
)
)
self.maybe_oversample_requests(
sampled_requests, num_requests, request_id_prefix, no_oversample
)
return sampled_requests
# -----------------------------------------------------------------------------
# Spec Bench Dataset Implementation
# -----------------------------------------------------------------------------
@@ -23,6 +23,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kNvfp4Dynamic,
)
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
DeepseekScalingRotaryEmbedding,
)
from vllm.platforms import current_platform
RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
@@ -158,6 +161,87 @@ class MatcherRotaryEmbedding(MatcherCustomOp):
return result
class MatcherDeepseekScalingRotaryEmbedding(MatcherCustomOp):
def __init__(
self,
is_neox: bool,
head_size: int,
num_heads: int,
num_kv_heads: int,
use_flashinfer: bool = False,
enabled: bool | None = None,
) -> None:
if enabled is None:
enabled = DeepseekScalingRotaryEmbedding.enabled()
super().__init__(enabled)
self.is_neox = is_neox
self.head_size = head_size
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.q_size = self.num_heads * self.head_size
self.kv_size = self.num_kv_heads * self.head_size
self.rotary_dim = head_size
self.use_flashinfer = use_flashinfer
def inputs(self) -> list[torch.Tensor]:
positions = self.empty_int64(5)
query = self.empty(5, self.num_heads, self.head_size)
key = self.empty(5, self.num_kv_heads, self.head_size)
cos_sin_cache = self.empty(4096, self.rotary_dim)
return [positions, query, key, cos_sin_cache]
def forward_custom(
self,
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor | None,
cos_sin_cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor | None]:
if self.use_flashinfer:
torch.ops.vllm.flashinfer_rotary_embedding(
positions,
query,
key,
self.head_size,
cos_sin_cache,
self.is_neox,
)
return query, key
result: tuple[torch.Tensor, torch.Tensor | None] = (
DeepseekScalingRotaryEmbedding.forward_static(
positions,
query,
key,
self.head_size,
self.rotary_dim,
cos_sin_cache,
self.is_neox,
)
)
return result
def forward_native(
self,
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor | None,
cos_sin_cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor | None]:
result: tuple[torch.Tensor, torch.Tensor | None] = (
DeepseekScalingRotaryEmbedding.forward_static(
positions,
query,
key,
self.head_size,
self.rotary_dim,
cos_sin_cache,
self.is_neox,
)
)
return result
class MatcherQuantFP8(MatcherCustomOp):
def __init__(
self,
@@ -189,12 +273,7 @@ class MatcherQuantFP8(MatcherCustomOp):
"ROCm aiter fusion pass currently supports "
"quantization operation with group_size 128"
)
if current_platform.is_fp8_fnuz():
self.QUANT_OP = rocm_aiter_ops.get_group_quant_op()
else:
self.QUANT_OP = (
torch.ops.vllm.triton_per_token_group_quant_fp8.default
)
self.QUANT_OP = rocm_aiter_ops.get_group_quant_op()
else:
assert quant_key in QUANT_OPS, (
@@ -0,0 +1,271 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from torch._higher_order_ops.auto_functionalize import auto_functionalized
import vllm._custom_ops as ops
from vllm.config import VllmConfig, get_layers_from_vllm_config
from vllm.logger import init_logger
from vllm.model_executor.layers.attention import MLAAttention
from vllm.model_executor.layers.attention.attention import get_attention_context
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
from vllm.utils.torch_utils import (
_USE_LAYERNAME,
LayerNameType,
_encode_layer_name,
_resolve_layer_name,
direct_register_custom_op,
)
from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
from .matcher_utils import MatcherDeepseekScalingRotaryEmbedding, MatcherRotaryEmbedding
logger = init_logger(__name__)
def fused_rope_unified_mla_kv_cache_update_impl(
positions: torch.Tensor,
q_pe: torch.Tensor,
k_pe: torch.Tensor,
kv_c: torch.Tensor,
cos_sin_cache: torch.Tensor,
is_neox: bool,
kv_cache_dtype: str,
kv_cache_scale: torch.Tensor,
layer_name: LayerNameType,
) -> torch.Tensor:
layer_name = _resolve_layer_name(layer_name)
attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
if layer_slot_mapping is not None:
ops.concat_and_cache_mla_rope_fused(
positions,
q_pe,
k_pe,
kv_c,
cos_sin_cache,
is_neox,
layer_slot_mapping,
kv_cache,
kv_cache_dtype,
kv_cache_scale,
)
return torch.empty(0, device=kv_c.device, dtype=kv_c.dtype)
def fused_rope_unified_mla_kv_cache_update_fake(
positions: torch.Tensor,
q_pe: torch.Tensor,
k_pe: torch.Tensor,
kv_c: torch.Tensor,
cos_sin_cache: torch.Tensor,
is_neox: bool,
kv_cache_dtype: str,
kv_cache_scale: torch.Tensor,
layer_name: LayerNameType,
) -> torch.Tensor:
return torch.empty(0, dtype=kv_c.dtype, device=kv_c.device)
direct_register_custom_op(
op_name="fused_rope_unified_mla_kv_cache_update",
op_func=fused_rope_unified_mla_kv_cache_update_impl,
fake_impl=fused_rope_unified_mla_kv_cache_update_fake,
mutates_args=["q_pe", "k_pe"],
)
class MLARoPEKVCacheCatPattern(VllmPatternReplacement):
FUSED_OP = torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default
def __init__(
self,
layer: MLAAttention,
is_neox: bool,
use_flashinfer: bool = False,
use_deepseek_scaling: bool = False,
) -> None:
self.layer_name = layer.layer_name
self.kv_cache_dtype = layer.kv_cache_dtype
self.num_heads = layer.num_heads
self.num_kv_heads = layer.num_kv_heads
self.kv_lora_rank = layer.kv_lora_rank
self.qk_rope_head_dim = layer.qk_rope_head_dim
self.is_neox = is_neox
self.use_flashinfer = use_flashinfer
self._ln = _encode_layer_name(self.layer_name)
if use_deepseek_scaling:
self.rope_matcher = MatcherDeepseekScalingRotaryEmbedding(
is_neox=self.is_neox,
head_size=self.qk_rope_head_dim,
num_heads=self.num_heads,
num_kv_heads=self.num_kv_heads,
use_flashinfer=self.use_flashinfer,
)
else:
self.rope_matcher = MatcherRotaryEmbedding( # type: ignore
is_neox=self.is_neox,
head_size=self.qk_rope_head_dim,
num_heads=self.num_heads,
num_kv_heads=self.num_kv_heads,
use_flashinfer=self.use_flashinfer,
)
def get_inputs(self) -> list[torch.Tensor]:
T = 5
L = 4096
q_pe = self.empty_bf16(T, self.num_heads, self.qk_rope_head_dim)
k_pe = self.empty_bf16(T, self.qk_rope_head_dim)
kv_c_normed = self.empty_bf16(T, self.kv_lora_rank)
cos_sin_cache = self.empty_bf16(L, self.qk_rope_head_dim)
positions = self.empty(T, dtype=torch.int64)
k_scale = self.empty(0, dtype=torch.float32)
inputs = [
q_pe,
k_pe,
kv_c_normed,
positions,
cos_sin_cache,
k_scale,
]
if _USE_LAYERNAME:
inputs.append(self._ln)
return inputs
@property
def pattern(self):
_ln = self._ln
if _USE_LAYERNAME:
def _pattern_with_ln(
q_pe: torch.Tensor,
k_pe: torch.Tensor,
kv_c_normed: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
k_scale: torch.Tensor,
layer_name: LayerNameType,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
k_pe_unsqueezed = k_pe.unsqueeze(1)
q_pe, k_pe = self.rope_matcher(
positions, q_pe, k_pe_unsqueezed, cos_sin_cache
)
dummy = torch.ops.vllm.unified_mla_kv_cache_update(
kv_c_normed, k_pe, layer_name, self.kv_cache_dtype, k_scale
)
return dummy, q_pe, k_pe
return _pattern_with_ln
def _pattern(
q_pe: torch.Tensor,
k_pe: torch.Tensor,
kv_c_normed: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
k_scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
k_pe_unsqueezed = k_pe.unsqueeze(1)
q_pe, k_pe = self.rope_matcher(
positions, q_pe, k_pe_unsqueezed, cos_sin_cache
)
dummy = torch.ops.vllm.unified_mla_kv_cache_update(
kv_c_normed, k_pe, _ln, self.kv_cache_dtype, k_scale
)
return dummy, q_pe, k_pe
return _pattern
@property
def replacement(self):
_ln = self._ln
if _USE_LAYERNAME:
def _replacement_with_ln(
q_pe: torch.Tensor,
k_pe: torch.Tensor,
kv_c_normed: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
k_scale: torch.Tensor,
layer_name: LayerNameType,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
at = auto_functionalized(
self.FUSED_OP,
positions=positions,
q_pe=q_pe,
k_pe=k_pe,
kv_c=kv_c_normed,
cos_sin_cache=cos_sin_cache,
is_neox=self.is_neox,
kv_cache_dtype=self.kv_cache_dtype,
kv_cache_scale=k_scale,
layer_name=layer_name,
)
dummy, q_pe, k_pe_squeezed = at
k_pe = k_pe_squeezed.unsqueeze(1)
return dummy, q_pe, k_pe
return _replacement_with_ln
def _replacement(
q_pe: torch.Tensor,
k_pe: torch.Tensor,
kv_c_normed: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
k_scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
at = auto_functionalized(
self.FUSED_OP,
positions=positions,
q_pe=q_pe,
k_pe=k_pe,
kv_c=kv_c_normed,
cos_sin_cache=cos_sin_cache,
is_neox=self.is_neox,
kv_cache_dtype=self.kv_cache_dtype,
kv_cache_scale=k_scale,
layer_name=_ln,
)
dummy, q_pe, k_pe_squeezed = at
k_pe = k_pe_squeezed.unsqueeze(1)
return dummy, q_pe, k_pe
return _replacement
class MLARoPEKVCacheCatFusionPass(VllmFusionPatternMatcherPass):
def __init__(self, config: VllmConfig) -> None:
super().__init__(config, "mla_rope_kv_cache_fusion_pass")
attn_layers = get_layers_from_vllm_config(config, MLAAttention)
for _, layer in attn_layers.items():
for is_neox in [False, True]:
for use_deepseek_scaling in [False, True]:
if RotaryEmbedding.enabled():
for use_flashinfer in [False, True]:
self.register(
MLARoPEKVCacheCatPattern(
layer,
is_neox,
use_flashinfer,
use_deepseek_scaling,
)
)
else:
self.register(
MLARoPEKVCacheCatPattern(
layer,
is_neox,
use_deepseek_scaling=use_deepseek_scaling,
)
)
if _USE_LAYERNAME:
break
self.dump_patterns(config, self.pm_pass)
@@ -6,6 +6,7 @@ from typing import Any
import torch
import torch._inductor.pattern_matcher as pm
from torch import fx
from torch._inductor.fx_passes.post_grad import view_to_reshape
from torch._inductor.pattern_matcher import PatternMatcherPass
import vllm.ir.ops
@@ -293,6 +294,161 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
class DoubleAiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
"""
Pattern matching ``rms_norm`` whose output feeds *two* distinct
``rocm_aiter_group_fp8_quant`` consumers, replacing it with two
independent fused ``rms_norm_group_fp8_quant`` ops.
Repeating the rms_norm in the replacement is preferable to leaving
the fused 16-bit rms output materialized for two unfused quant
consumers, and matches what the previous manual graph surgery
achieved by cloning the rms_norm node.
"""
FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
def __init__(
self,
epsilon: float,
quant_dtype: torch.dtype,
group_shape: GroupShape,
match_aiter_quant: bool = True,
symmetric: bool = True,
) -> None:
scale = ScaleDesc(torch.float32, False, group_shape)
key = FusedRMSQuantKey(
fused_add=False,
quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
)
super().__init__(epsilon, key, match_aiter_quant)
def register(self, pm_pass: PatternMatcherPass) -> None:
def pattern(
input: torch.Tensor,
weight: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon)
result1, scale1 = self.quant_matcher(result_rms)
result2, scale2 = self.quant_matcher(result_rms)
return result1, scale1, result2, scale2
def replacement(
input: torch.Tensor,
weight: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
at1 = self.FUSED_OP(
x=input,
weight=weight,
variance_epsilon=self.epsilon,
group_size=128,
)
at2 = self.FUSED_OP(
x=input,
weight=weight,
variance_epsilon=self.epsilon,
group_size=128,
)
return at1[0], at1[1], at2[0], at2[1]
pm.register_replacement(
pattern,
replacement,
# input, weight
[self.empty(5, 16), self.empty(16)],
pm.fwd_only,
pm_pass,
)
class DoubleAiterRMSFp8GroupQuantViewPattern(AiterRMSNormQuantPattern):
"""
View-tolerant variant of ``DoubleAiterRMSFp8GroupQuantPattern``.
Matches the same 1-to-2 fan-out, but with a ``view``/``reshape`` between
the ``rms_norm`` output and the two ``rocm_aiter_group_fp8_quant``
consumers::
rms_norm -> view -> rocm_aiter_group_fp8_quant
\\-> view -> rocm_aiter_group_fp8_quant
This shape arises in DeepSeek-V3.2's MLA indexer q_c norm, where the
FP8 linear path's 2D-flatten boilerplate
(``Fp8BlockScaledMMLinearKernel.apply_weights``) inserts a view between
the rms_norm output and each FP8 group quant op. The non-view sibling
pattern silently no-ops on this graph because the pattern matcher
requires the in-graph and in-pattern node shapes to align.
The trace_fn runs Inductor's ``view_to_reshape`` post-grad pass to
normalize ``view`` to ``reshape`` in both the pattern and the input
graph, widening the match without touching the no-view sibling.
"""
FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
def __init__(
self,
epsilon: float,
quant_dtype: torch.dtype,
group_shape: GroupShape,
match_aiter_quant: bool = True,
symmetric: bool = True,
) -> None:
scale = ScaleDesc(torch.float32, False, group_shape)
key = FusedRMSQuantKey(
fused_add=False,
quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
)
super().__init__(epsilon, key, match_aiter_quant)
def register(self, pm_pass: PatternMatcherPass) -> None:
def pattern(
input: torch.Tensor,
weight: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon)
view_rms = result_rms.view(-1, result_rms.shape[-1])
result1, scale1 = self.quant_matcher(view_rms)
result2, scale2 = self.quant_matcher(view_rms)
return result1, scale1, result2, scale2
def replacement(
input: torch.Tensor,
weight: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
at1 = self.FUSED_OP(
x=input,
weight=weight,
variance_epsilon=self.epsilon,
group_size=128,
)
at2 = self.FUSED_OP(
x=input,
weight=weight,
variance_epsilon=self.epsilon,
group_size=128,
)
return at1[0], at1[1], at2[0], at2[1]
def trace_with_view_to_reshape(*args: Any, **kwargs: Any) -> fx.GraphModule:
gm = pm.fwd_only(*args, **kwargs)
view_to_reshape(gm)
return gm
pm.register_replacement(
pattern,
replacement,
# input, weight
[self.empty(5, 16), self.empty(16)],
trace_with_view_to_reshape,
pm_pass,
)
class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
"""
This pass fuses aiter rms_norm & vllm/aiter quant custom ops
@@ -309,8 +465,24 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
)
# Make sure fused add patterns are before simple rms norm,
# as the latter is a subset of the former in torch ops
# as the latter is a subset of the former in torch ops.
# The DoubleQuant patterns handle 1 rms_norm -> 2 group_fp8_quant
# fan-out (e.g. DSv3.2) and must be registered before the single
# group-quant pattern so they match first. The view-tolerant variant
# additionally covers the rms_norm -> view -> 2x quant shape that
# appears when the FP8 linear path inserts a 2D-flatten boilerplate
# (DSv3.2 MLA indexer q_c norm).
for epsilon in [1e-5, 1e-6]:
# Fuse aiter rms_norm + 2x aiter group fp8 quant
DoubleAiterRMSFp8GroupQuantPattern(
epsilon, FP8_DTYPE, GroupShape(1, 128)
).register(self.patterns)
# View-tolerant sibling for DSv3.2 q_c norm fan-out
DoubleAiterRMSFp8GroupQuantViewPattern(
epsilon, FP8_DTYPE, GroupShape(1, 128)
).register(self.patterns)
# Fuse aiter rms_norm + aiter dynamic group fp8 quant
AiterRMSFp8GroupQuantPattern(
epsilon, FP8_DTYPE, GroupShape(1, 128)
@@ -360,6 +532,8 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
AiterFusedAddRMSNormDynamicQuantPattern,
AiterRMSFp8GroupQuantPattern,
AiterFusedAddRMSFp8GroupQuantPattern,
DoubleAiterRMSFp8GroupQuantPattern,
DoubleAiterRMSFp8GroupQuantViewPattern,
]
return self.hash_source(self, *fusion_patterns)
+4
View File
@@ -33,6 +33,7 @@ if current_platform.is_cuda_alike():
from .fusion.act_quant_fusion import ActivationQuantFusionPass
from .fusion.attn_quant_fusion import AttnQuantFusionPass
from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass
from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass
from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
@@ -174,6 +175,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
self.passes += [ScatterSplitReplacementPass(config)]
self.passes += [RopeKVCacheFusionPass(config)]
if self.pass_config.fuse_rope_kvcache_cat_mla:
self.passes += [MLARoPEKVCacheCatFusionPass(config)]
if self.pass_config.fuse_attn_quant:
self.passes += [AttnQuantFusionPass(config)]
self.passes += [MLAAttnQuantFusionPass(config)]
@@ -181,6 +181,45 @@ class FixFunctionalizationPass(VllmInductorPass):
2: "key",
}
self.defunctionalize(graph, node, mutated_args=mutated_args)
elif (
hasattr(torch.ops.vllm, "fused_rope_unified_mla_kv_cache_update")
and at_target
== torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default
):
# AOTAutograd functionalizes `q[..., nope_dim:] = rope_result` into
# a sequence of aten ops on q: view+slice+copy+slice_scatter.
# Since the fused MLA RoPE op mutates q_pe in-place, we can remove
# the redundant copy and slice_scatter ops during defunctionalization.
getitem_nodes = self.getitem_users(node)
q_pe_out = getitem_nodes[1]
for user in list(q_pe_out.users):
if is_func(user, torch.ops.aten.copy.default):
copy_temp = user
slice_temp = copy_temp.args[0]
for user in list(copy_temp.users):
if is_func(user, torch.ops.aten.slice_scatter.default):
slice_scatter_temp = user
view_temp = slice_scatter_temp.args[0]
view_orig = slice_temp.args[0]
slice_scatter_temp.replace_all_uses_with(view_orig)
self._remove(slice_scatter_temp)
self._remove(copy_temp)
self._remove(slice_temp)
self._remove(view_temp)
self._remove(q_pe_out)
# defunctionalize k_pe manually; self.replace_users_with_mutated_args
# does not support only replacing specific kwargs
k_pe_in = node.kwargs["k_pe"]
k_pe_out = getitem_nodes[2]
k_pe_out.replace_all_uses_with(k_pe_in)
self._remove(k_pe_out)
self.insert_defunctionalized(graph, node)
self._remove(node)
# only used for test_functionalization::TestFunctionWithMutatedArgsAndReturn
elif (
hasattr(torch.ops.vllm, "function_with_mutated_args_and_return")
+10 -1
View File
@@ -136,8 +136,10 @@ class PassConfig:
"""Enable flashinfer allreduce fusion."""
fuse_minimax_qk_norm: bool = None # type: ignore[assignment]
"""Enable fused allreduce+RMSNorm for MiniMax QK norm."""
enable_qk_norm_rope_fusion: bool = False
enable_qk_norm_rope_fusion: bool = None # type: ignore[assignment]
"""Enable fused Q/K RMSNorm + RoPE pass."""
fuse_rope_kvcache_cat_mla: bool = None # type: ignore[assignment]
"""Enable fused MLA KV cache update with RoPE."""
# ROCm/AITER specific fusions
fuse_act_padding: bool = None # type: ignore[assignment]
@@ -228,6 +230,7 @@ class PassConfig:
"fuse_act_padding",
"fuse_mla_dual_rms_norm",
"fuse_rope_kvcache",
"fuse_rope_kvcache_cat_mla",
mode="wrap",
)
@classmethod
@@ -285,6 +288,12 @@ class PassConfig:
"The fusion will be disabled."
)
self.fuse_rope_kvcache = False
if self.fuse_rope_kvcache_cat_mla and not current_platform.is_cuda_alike():
logger.warning_once(
"MLA KV cache update with RoPE fusion enabled but the "
"current platform is not CUDA or ROCm. The fusion will be disabled."
)
self.fuse_rope_kvcache_cat_mla = False
def log_enabled_passes(self) -> None:
"""
+13
View File
@@ -155,6 +155,15 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool:
)
def enable_rope_kvcache_mla_fusion(cfg: "VllmConfig") -> bool:
"""Enable if use_inductor_graph_partition is enabled."""
return (
cfg.compilation_config.use_inductor_graph_partition
or not cfg.compilation_config.splitting_ops_contain_kv_cache_update()
)
def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
"""Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss."""
@@ -184,6 +193,7 @@ OPTIMIZATION_LEVEL_00 = {
"fuse_act_padding": False,
"fuse_mla_dual_rms_norm": False,
"fuse_rope_kvcache": False,
"fuse_rope_kvcache_cat_mla": False,
},
"cudagraph_mode": CUDAGraphMode.NONE,
"use_inductor_graph_partition": False,
@@ -204,6 +214,7 @@ OPTIMIZATION_LEVEL_01 = {
"fuse_act_padding": enable_norm_pad_fusion,
"fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
"fuse_rope_kvcache": False,
"fuse_rope_kvcache_cat_mla": False,
},
"cudagraph_mode": CUDAGraphMode.PIECEWISE,
"use_inductor_graph_partition": False,
@@ -226,6 +237,7 @@ OPTIMIZATION_LEVEL_02 = {
"fuse_act_padding": enable_norm_pad_fusion,
"fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
"fuse_rope_kvcache": enable_rope_kvcache_fusion,
"fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion,
},
"cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
"use_inductor_graph_partition": False,
@@ -248,6 +260,7 @@ OPTIMIZATION_LEVEL_03 = {
"fuse_act_padding": enable_norm_pad_fusion,
"fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
"fuse_rope_kvcache": enable_rope_kvcache_fusion,
"fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion,
},
"cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
"use_inductor_graph_partition": False,
@@ -3,6 +3,8 @@
# ===================== import region =====================
import threading
import torch
import torch.distributed as dist
from torch.distributed import ProcessGroup, ReduceOp
@@ -145,8 +147,19 @@ class PyNcclCommunicator:
def destroy(self):
if self.available and not self.disabled:
with torch.accelerator.device_index(self.device.index):
self.nccl.ncclCommDestroy(self.comm)
# ncclCommAbort can block until all CUDA graphs that
# captured NCCL ops on this comm are destroyed — and
# those graphs are released later in this same main-
# thread teardown, so a direct call here self-deadlocks.
# Run it in a daemon thread with a timeout: the main
# thread proceeds, the graphs drop, and the abort returns.
def _abort():
with torch.accelerator.device_index(self.device.index):
self.nccl.ncclCommAbort(self.comm)
abort_thread = threading.Thread(target=_abort, daemon=True)
abort_thread.start()
abort_thread.join(timeout=5.0)
self.available = False
self.disabled = True
@@ -290,6 +290,12 @@ class NCCLLibrary:
# it is better not to call it at all.
# ncclResult_t ncclCommDestroy(ncclComm_t comm);
Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
# ncclCommAbort frees resources associated with the communicator
# without requiring a collective synchronization. Unlike
# ncclCommDestroy, it is safe to call during an uncoordinated
# shutdown when peer ranks may already be gone.
# ncclResult_t ncclCommAbort(ncclComm_t comm);
Function("ncclCommAbort", ncclResult_t, [ncclComm_t]),
# ncclResult_t ncclGroupStart();
Function("ncclGroupStart", ncclResult_t, []),
# ncclResult_t ncclGroupEnd();
@@ -548,6 +554,9 @@ class NCCLLibrary:
def ncclCommDestroy(self, comm: ncclComm_t) -> None:
self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
def ncclCommAbort(self, comm: ncclComm_t) -> None:
self.NCCL_CHECK(self._funcs["ncclCommAbort"](comm))
def ncclGroupStart(self) -> None:
self.NCCL_CHECK(self._funcs["ncclGroupStart"]())
@@ -38,6 +38,11 @@ from vllm.utils.network_utils import (
is_valid_ipv6_address,
)
if envs.VLLM_USE_SPINLOOP_EXT:
from vllm.spinloop import spinloop
SPINLOOP_TIMEOUT_SECONDS = 0.1
if TYPE_CHECKING:
from _typeshed import SizedBuffer
@@ -540,13 +545,17 @@ class MessageQueue:
n_warning = 1
while True:
with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
# Memory fence ensures we see the latest read flags from readers.
# Without this, we may read stale flags from our CPU cache and
# spin indefinitely even though readers have completed.
memory_fence()
read_count = sum(metadata_buffer[1:])
written_flag = metadata_buffer[0]
if written_flag and read_count != self.buffer.n_reader:
def check():
memory_fence()
read_count = sum(metadata_buffer[1:])
written_flag = metadata_buffer[0]
return not (written_flag and read_count != self.buffer.n_reader)
if envs.VLLM_USE_SPINLOOP_EXT and not check():
spinloop(metadata_buffer, check, timeout=SPINLOOP_TIMEOUT_SECONDS)
if not check():
# this block is written and not read by all readers
# for writers, `self.current_idx` is the next block to write
# if this block is not ready to write,
@@ -657,13 +666,21 @@ class MessageQueue:
)
with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
while True:
# Memory fence ensures we see the latest writes from the writer.
# Without this, we may read stale flags from our CPU cache
# and spin indefinitely even though writer has updated them.
memory_fence()
read_flag = metadata_buffer[self.local_reader_rank + 1]
written_flag = metadata_buffer[0]
if not written_flag or read_flag:
def check():
memory_fence()
read_flag = metadata_buffer[self.local_reader_rank + 1]
written_flag = metadata_buffer[0]
return not (not written_flag or read_flag)
if envs.VLLM_USE_SPINLOOP_EXT and not check():
spinloop(
metadata_buffer[0 : self.local_reader_rank + 1],
check,
timeout=SPINLOOP_TIMEOUT_SECONDS,
)
if not check():
# this block is either
# (1) not written
# (2) already read by this reader
@@ -510,6 +510,14 @@ class KVConnectorBase_V1(ABC):
"""
pass
def on_new_request(self, request: "Request") -> None:
"""Called by the scheduler when a new request is added.
Connectors can override this to inspect the request and perform
bookkeeping. The default implementation is a no-op.
"""
return
def update_connector_output(self, connector_output: KVConnectorOutput):
"""
Update KVConnector state from worker-side connectors output.

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