Merge branch 'main' into woosuk/triton-fix

This commit is contained in:
Woosuk Kwon
2026-06-22 00:29:09 +00:00
274 changed files with 8665 additions and 3855 deletions
+1 -1
View File
@@ -67,7 +67,7 @@ steps:
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
pytest -v -s v1/structured_output &&
pytest -v -s v1/test_serial_utils.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
- label: "XPU server test"
depends_on:
+26 -39
View File
@@ -415,22 +415,6 @@ steps:
commands:
- pytest -v -s kernels/mamba
#----------------------------------------------------------- mi250 · lora ------------------------------------------------------------#
- label: LoRA %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
parallelism: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/lora
- 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_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 · models / basic -------------------------------------------------------#
- label: Basic Models Test (Other CPU) # TBD
@@ -663,7 +647,7 @@ steps:
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- label: e2e Core (1 GPU) # TBD
timeout_in_minutes: 180
timeout_in_minutes: 35
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
optional: true
@@ -1610,9 +1594,10 @@ steps:
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 180
timeout_in_minutes: 55
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
parallelism: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -1643,10 +1628,11 @@ steps:
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 180
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
parallelism: 4
optional: true
parallelism: 5
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
@@ -1699,6 +1685,20 @@ steps:
#----------------------------------------------------------- mi300 · lora ------------------------------------------------------------#
- label: LoRA %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
parallelism: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/lora
- 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_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
- label: LoRA TP (Distributed) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2077,19 +2077,6 @@ steps:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
- label: e2e Core (1 GPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
- label: e2e Scheduling (1 GPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2135,9 +2122,10 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: Spec Decode Eagle # TBD
timeout_in_minutes: 180
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
@@ -3055,7 +3043,7 @@ steps:
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
- label: Kernels (B200-MI355) # TBD
timeout_in_minutes: 180
timeout_in_minutes: 15
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
working_dir: "/vllm-workspace/"
@@ -3079,11 +3067,10 @@ steps:
- pytest -v -s tests/kernels/attention/test_attention_selector.py
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 180
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
parallelism: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/attention/
@@ -3097,10 +3084,10 @@ steps:
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 180
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
parallelism: 4
parallelism: 5
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
+10
View File
@@ -74,6 +74,16 @@ steps:
- tests/v1/e2e/general/
commands:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
mirror:
amd:
device: mi250_1
timeout_in_minutes: 35
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
- vllm/platforms/rocm.py
- label: V1 e2e (2 GPUs)
key: v1-e2e-2-gpus
+31
View File
@@ -74,6 +74,20 @@ steps:
commands:
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 2
mirror:
amd:
device: mi325_1
timeout_in_minutes: 55
depends_on:
- image-build-amd
source_file_dependencies:
- csrc/attention/
- vllm/v1/attention
- vllm/model_executor/layers/attention
- tests/kernels/attention
- vllm/_aiter_ops.py
- vllm/envs.py
- vllm/platforms/rocm.py
- label: Kernels Attention DiffKV Test (H100)
key: kernels-attention-diffkv-test-h100
@@ -104,6 +118,7 @@ steps:
source_file_dependencies:
- csrc/quantization/
- vllm/model_executor/layers/quantization
- vllm/config/
- tests/kernels/quantization
- tests/kernels/quantization/test_rocm_skinny_gemms.py
- vllm/_aiter_ops.py
@@ -127,6 +142,22 @@ steps:
- pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
- pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 5
mirror:
amd:
device: mi325_1
timeout_in_minutes: 50
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
- csrc/moe/
- tests/kernels/moe
- vllm/model_executor/layers/fused_moe/
- vllm/distributed/device_communicators/
- vllm/envs.py
- vllm/config
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
depends_on:
- image-build-amd
- label: Kernels Mamba Test
key: kernels-mamba-test
+10
View File
@@ -101,6 +101,16 @@ steps:
num_devices: 8
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
mirror:
amd:
device: mi300_8
timeout_in_minutes: 180
depends_on:
- image-build-amd
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
key: moe-refactor-integration-test-h100-temporary
+11
View File
@@ -12,6 +12,17 @@ steps:
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_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
mirror:
amd:
device: mi325_1
working_dir: "/vllm-workspace/tests"
timeout_in_minutes: 60
source_file_dependencies:
- vllm/lora
- tests/lora
- vllm/platforms/rocm.py
depends_on:
- image-build-amd
- label: LoRA TP (Distributed)
+6
View File
@@ -21,6 +21,12 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# TODO: create another `optional` test group for slow tests
- pytest -v -s -m 'not slow_test' v1/spec_decode
mirror:
amd:
device: mi300_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
- label: V1 Sample + Logits
key: v1-sample-logits
@@ -68,7 +68,6 @@ steps:
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
mirror:
amd:
soft_fail: true
device: mi325_1
depends_on:
- image-build-amd
+6
View File
@@ -107,6 +107,12 @@ steps:
- tests/compile/passes
commands:
- pytest -s -v compile/passes --ignore compile/passes/distributed
mirror:
amd:
device: mi300_1
timeout_in_minutes: 180
depends_on:
- image-build-amd
- label: PyTorch Fullgraph Smoke Test
key: pytorch-fullgraph-smoke-test
+14
View File
@@ -12,6 +12,20 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
mirror:
amd:
device: mi325_1
timeout_in_minutes: 45
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- label: Spec Decode Eagle Nightly B200
key: spec-decode-eagle-nightly-b200
@@ -0,0 +1,35 @@
---
name: ci-fails-buildkite
description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs.
---
# Diagnosing vLLM Buildkite CI Failures
Buildkite logs are public; no login needed.
`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch.
## Fetching logs
```bash
# All failed jobs in a PR's latest build (current branch's PR if omitted):
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
# All failed jobs in a build (--soft also includes soft-failed jobs;
# --all fetches every finished job):
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
# work; pass "-" as a second argument to stream to stdout:
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
```
To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`:
```bash
./ci-clean-log.sh ci.log
```
## Reference
See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage.
+2 -3
View File
@@ -2,15 +2,14 @@
# for more info about CODEOWNERS file
# This lists cover the "core" components of vLLM that require careful review
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
/vllm/lora @jeejeelee
/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
/vllm/model_executor/layers/mamba @tdoublep @tomeras91
/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
/vllm/model_executor/model_loader @22quinn
/vllm/model_executor/layers/batch_invariant.py @yewentao256
/vllm/ir @ProExpertProg
+3 -1
View File
@@ -199,7 +199,9 @@ cython_debug/
.vscode/
# Claude
.claude/
.claude/*
!.claude/skills/
!.claude/skills/**
# Codex
.codex/
-11
View File
@@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository:
- Keep comments and docstrings minimal and concise.
- Assume the reader is familiar with vLLM.
### Diagnosing CI failures
Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
```bash
# All failed-job logs for a PR's latest build (current branch's PR if omitted):
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
# Any Buildkite build or job URL also works:
.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
```
### Commit messages
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
+56 -72
View File
@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
#
# _C extension
# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
#
set(VLLM_EXT_SRC
"csrc/quantization/activation_kernels.cu"
"csrc/torch_bindings.cpp")
if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
set(CUTLASS_REVISION "v4.4.2")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
endif()
if(VLLM_CUTLASS_SRC_DIR)
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
endif()
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
# Please keep this in sync with CUTLASS_REVISION line above.
GIT_TAG ${CUTLASS_REVISION}
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
GIT_SHALLOW TRUE
)
endif()
FetchContent_MakeAvailable(cutlass)
set_gencode_flags_for_srcs(
SRCS "${VLLM_EXT_SRC}"
CUDA_ARCHS "${CUDA_ARCHS}")
# if CUDA endif
endif()
if (VLLM_GPU_LANG STREQUAL "HIP")
# Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
# TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
list(APPEND VLLM_EXT_SRC
if(VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_EXT_SRC
"csrc/torch_bindings.cpp"
"csrc/custom_quickreduce.cu"
"csrc/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
)
# if ROCM endif
endif()
"csrc/libtorch_stable/cuda_utils_kernels.cu")
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
# driver API. This causes problems when linking with earlier versions of CUDA.
# Setting this variable sidesteps the issue by calling the driver directly.
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
# driver API. This causes problems when linking with earlier versions of CUDA.
# Setting this variable sidesteps the issue by calling the driver directly.
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
endif() # _C HIP endif
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
#
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/torch_bindings.cpp"
"csrc/libtorch_stable/activation_kernels.cu"
"csrc/libtorch_stable/quantization/activation_kernels.cu"
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
@@ -429,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
set(CUTLASS_REVISION "v4.4.2")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
endif()
if(VLLM_CUTLASS_SRC_DIR)
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
endif()
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
# Please keep this in sync with CUTLASS_REVISION line above.
GIT_TAG ${CUTLASS_REVISION}
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
GIT_SHALLOW TRUE
)
endif()
FetchContent_MakeAvailable(cutlass)
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${FP4_SM120_SRCS}"
CUDA_ARCHS "${FP4_SM120_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
@@ -962,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${FP4_SM100_SRCS}"
CUDA_ARCHS "${FP4_SM100_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
+29 -5
View File
@@ -60,6 +60,7 @@ endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
set(QUTLASS_SOURCES
csrc/qutlass_registration.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
else()
get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
endif()
endif()
elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
list(APPEND QUTLASS_INCLUDES
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
else()
message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
CUDA_ARCHS "${QUTLASS_ARCHS}"
)
target_sources(_C PRIVATE ${QUTLASS_SOURCES})
target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
target_compile_definitions(_C PRIVATE
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
# Keep it as its own extension (registers torch.ops._qutlass_C).
define_extension_target(
_qutlass_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${QUTLASS_SOURCES}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
USE_SABI 3
WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE
QUTLASS_DISABLE_PYBIND=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
)
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +134,5 @@ else()
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
"CUDA_ARCHS='${CUDA_ARCHS}'.")
endif()
add_custom_target(_qutlass_C)
endif()
@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
using TileShapeD = typename MlaSm100Type::TileShapeD;
arguments.problem_shape =
cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
// Assumes device 0 when getting sm_count.
arguments.hw_info.sm_count =
sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
if (sm_count <= 0) {
int current_device = 0;
cudaGetDevice(&current_device);
arguments.hw_info.sm_count =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
} else {
arguments.hw_info.sm_count = sm_count;
}
arguments.split_kv = static_cast<int>(num_kv_splits);
MlaSm100Type::Fmha::set_split_kv(arguments);
+7 -6
View File
@@ -81,11 +81,11 @@ __global__ void rms_norm_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
scalar_t normalized = static_cast<scalar_t>(x * s_variance);
if constexpr (HasWeight) {
dst.val[j] = normalized * src2.val[j];
float w = static_cast<float>(src2.val[j]);
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
} else {
dst.val[j] = normalized;
dst.val[j] = static_cast<scalar_t>(x * s_variance);
}
}
v_out[i] = dst;
@@ -151,7 +151,8 @@ fused_add_rms_norm_kernel(
#pragma unroll
for (int j = 0; j < width; ++j) {
float x = Converter::convert(res.data[j]);
out.data[j] = Converter::convert(x * s_variance) * w.data[j];
float wf = Converter::convert(w.data[j]);
out.data[j] = Converter::convert(x * s_variance * wf);
}
} else {
#pragma unroll
@@ -198,8 +199,8 @@ fused_add_rms_norm_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
if constexpr (HasWeight) {
input[blockIdx.x * input_stride + idx] =
(scalar_t)(x * s_variance) * weight[idx];
float w = (float)weight[idx];
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
} else {
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance);
}
@@ -66,8 +66,13 @@ __global__ void rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
// Multiply in weight's native dtype to match rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
float w = static_cast<float>(src2.val[j]);
// Round normalized result through scalar_t to match the precision of the
// unfused composite (rms_norm writes scalar_t, then
// static_scaled_fp8_quant re-loads it as float before FP8 conversion).
// Without this round, the fused path is strictly more accurate and
// disagrees with the composite at exact E4M3 quantization tie boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
scale_inv);
@@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int i = 0; i < width; ++i) {
float x = Converter::convert(res.data[i]);
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
float wf = Converter::convert(w.data[i]);
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries. We use the
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
HipT out_norm_h = Converter::convert(x * s_variance * wf);
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
Converter::convert(out_norm_h), scale_inv);
}
@@ -183,8 +192,10 @@ fused_add_rms_norm_static_fp8_quant_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
float w = (float)weight[idx];
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
static_cast<float>(out_norm), scale_inv);
}
@@ -9,7 +9,7 @@
#include <torch/headeronly/core/ScalarType.h>
#include "../../cuda_compat.h"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/dispatch_utils.h"
#include "libtorch_stable/torch_utils.h"
+28
View File
@@ -2,9 +2,25 @@
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/util/Exception.h>
#include <optional>
#include <string>
#include <vector>
#include <torch/csrc/stable/ops.h>
inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
// Ensure tensor is on CUDA
STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
// Get the raw data pointer
void* data_ptr = tensor.mutable_data_ptr();
/// Create a new tensor from the raw data pointer
return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
tensor.device(), tensor.scalar_type());
}
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
torch::stable::Tensor& output_q,
@@ -371,6 +387,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void silu_and_mul_clamp(torch::stable::Tensor& out,
torch::stable::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::stable::Tensor& out,
torch::stable::Tensor& input,
torch::stable::Tensor& scale);
void persistent_masked_m_silu_mul_quant(
const torch::stable::Tensor& input, // (E, T, 2*H)
const torch::stable::Tensor& tokens_per_expert, // (E)
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool use_ue8m0);
void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_tanh_and_mul(torch::stable::Tensor& out,
@@ -1,16 +1,12 @@
#include <ATen/cuda/CUDAContext.h>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include "libtorch_stable/torch_utils.h"
#include <cmath>
#include "core/math.hpp"
#include "../cuda_compat.h"
#include "dispatch_utils.h"
#include "libtorch_stable/core/math.hpp"
#include "cuda_compat.h"
#include "libtorch_stable/dispatch_utils.h"
#include "quantization/w8a8/fp8/common.cuh"
#include <c10/util/Float8_e4m3fn.h>
#ifndef USE_ROCM
#include <cuda_bf16.h>
#include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
#endif
#endif
#include "core/registration.h"
namespace vllm {
template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
} // namespace vllm
// Launch activation, gating, and quantize kernel.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
dim3 block(std::min(d, 512)); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
VLLM_DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "act_and_mul_kernel", [&] { \
VLLM_DISPATCH_FP8_TYPES( \
out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
fp8_t> \
<<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(), \
input.data_ptr<scalar_t>(), \
scale.data_ptr<float>(), d); \
}); \
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
dim3 block(std::min(d, 512)); \
const torch::stable::accelerator::DeviceGuard device_guard( \
input.get_device_index()); \
const cudaStream_t stream = \
get_current_cuda_stream(input.get_device_index()); \
VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "act_and_mul_kernel", [&] { \
VLLM_STABLE_DISPATCH_FP8_TYPES( \
out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
fp8_t> \
<<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<fp8_t>(), \
input.const_data_ptr<scalar_t>(), \
scale.const_data_ptr<float>(), d); \
}); \
});
void silu_and_mul_quant(torch::Tensor& out, // [..., d]
torch::Tensor& input, // [..., 2 * d]
torch::Tensor& scale) {
TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
out.dtype() == torch::kFloat8_e4m3fnuz);
TORCH_CHECK(input.dtype() == torch::kFloat16 ||
input.dtype() == torch::kBFloat16);
TORCH_CHECK(input.size(-1) % 2 == 0);
void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input, // [..., 2 * d]
torch::stable::Tensor& scale) {
STD_TORCH_CHECK(
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
STD_TORCH_CHECK(
input.scalar_type() == torch::headeronly::ScalarType::Half ||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"Input must be FP16 or BF16");
STD_TORCH_CHECK(input.size(-1) % 2 == 0);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
}
void persistent_masked_m_silu_mul_quant(
const at::Tensor& input, // (E, T, 2*H)
const at::Tensor& tokens_per_expert, // (E)
at::Tensor& y_q, // (E, T, H) [OUT]
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
const torch::stable::Tensor& input, // (E, T, 2*H)
const torch::stable::Tensor& tokens_per_expert, // (E)
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool cast_scale_ue8m0) {
#ifndef USE_ROCM
@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
// fixed GROUP_SIZE of 128.
static constexpr int GROUP_SIZE = 128;
TORCH_CHECK(input.dtype() == torch::kBFloat16);
TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
y_q.dtype() == torch::kFloat8_e4m3fnuz);
TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
STD_TORCH_CHECK(input.scalar_type() ==
torch::headeronly::ScalarType::BFloat16);
STD_TORCH_CHECK(
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
bool const is_packed_ue8m0 =
(y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
(y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
cast_scale_ue8m0);
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
is_packed_ue8m0);
using Idx_t = int64_t;
@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(
int const NUM_GROUPS = H / GROUP_SIZE;
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
// TODO: Get this from cuda_arch ?
static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
static constexpr int max_shared_mem_bytes = \
GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \
dim3 grid(sms), block(THREAD_COUNT); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
VLLM_DISPATCH_FP8_TYPES( \
const torch::stable::accelerator::DeviceGuard device_guard( \
input.get_device_index()); \
VLLM_STABLE_DISPATCH_FP8_TYPES( \
y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \
vllm::silu_mul_fp8_quant_deep_gemm_kernel< \
BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \
<<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>( \
reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
(fp8_t*)y_q.data_ptr(), \
reinterpret_cast<scale_t*>(y_s.data_ptr()), \
reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
reinterpret_cast<const __nv_bfloat16*>( \
input.const_data_ptr()), \
y_q.mutable_data_ptr<fp8_t>(), \
reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()), \
reinterpret_cast<const int32_t*>( \
tokens_per_expert.const_data_ptr()), \
E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \
STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \
});
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
Idx_t stride_ys_g = y_s.stride(2);
Idx_t stride_ys_p = 0;
if (!cast_scale_ue8m0) {
TORCH_CHECK(!is_packed_ue8m0);
STD_TORCH_CHECK(!is_packed_ue8m0);
LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
false);
return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
return;
}
TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
TORCH_CHECK(y_s.dtype() == torch::kInt32);
STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);
// Int32 packed ue8m0 scales tensor.
// Let E, T, G be the number to experts, number of tokens and number of groups
@@ -31,7 +31,7 @@
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "core/batch_invariant.hpp"
using namespace cute;
@@ -31,7 +31,7 @@
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "core/batch_invariant.hpp"
using namespace cute;
@@ -24,13 +24,21 @@ __device__ inline void vectorize_with_alignment(
ScaOp&& scalar_op) { // InT -> OutT
static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
"VEC_SIZE must be a positive power-of-two");
constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 64 B
constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 16 B
constexpr int OUT_WIDTH = VEC_SIZE * sizeof(OutT); // eg: 16 B
uintptr_t addr = reinterpret_cast<uintptr_t>(in);
uintptr_t out_addr = reinterpret_cast<uintptr_t>(out);
// fast path when the whole region is already aligned
// Note: currently the output is guaranteed to be same as the input, so we
// don't check it here, comments here just for future reference.
bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0);
// fast path when input and output are both fully aligned. The vector
// load/store below go through vec_n_t<T, VEC_SIZE>, declared
// __align__(VEC_SIZE * sizeof(T)), so each side must be aligned to its
// own vector width. out is NOT generally co-aligned with in: e.g.
// reshape_and_cache_flash writes KV-cache rows whose byte offset is a
// multiple of head_size, which for head sizes that are not a multiple
// of VEC_SIZE puts some rows off the vector-width boundary.
bool can_vec = ((addr & (WIDTH - 1)) == 0) &&
((out_addr & (OUT_WIDTH - 1)) == 0) &&
((len & (VEC_SIZE - 1)) == 0);
if (can_vec) {
int num_vec = len / VEC_SIZE;
@@ -55,6 +63,16 @@ __device__ inline void vectorize_with_alignment(
prefix_elems /= sizeof(InT);
prefix_elems = min(prefix_elems, len); // 0 ≤ prefix < 16
// the prefix below aligns in; if that does not also align out (their
// addresses differ modulo the vector width), vectorizing is impossible
// and the whole copy must stay scalar.
if (((out_addr + prefix_elems * sizeof(OutT)) & (OUT_WIDTH - 1)) != 0) {
for (int i = tid; i < len; i += stride) {
scalar_op(out[i], in[i]);
}
return;
}
// 1. prefill the when it is unsafe to vectorize
for (int i = tid; i < prefix_elems; i += stride) {
scalar_op(out[i], in[i]);
@@ -19,7 +19,7 @@
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -14,7 +14,7 @@
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -22,7 +22,7 @@
#include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
const int sf_k_local = local_group_id % kGroupsPerBlockX;
const int row_local = local_group_id / kGroupsPerBlockX;
const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
// Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;");
@@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
" is not a multiple of 4.");
const int kx = GetGroupsPerBlockX(padded_groups_per_row);
const int ry = 16 / kx;
const int64_t blocks_x = padded_groups_per_row / kx;
const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
const int64_t sf_k_blocks = padded_groups_per_row / kx;
const int num_threads = (kx * ry) * THREADS_PER_GROUP;
// CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
blocks_y <= static_cast<int64_t>(INT32_MAX),
// CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
constexpr int64_t kMaxGridDimYZ = 65535;
STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
sf_k_blocks <= kMaxGridDimYZ,
"per_token_group_quant_8bit_packed grid too large: (",
blocks_x, ", ", blocks_y, ").");
row_blocks, ", ", sf_k_blocks, ").");
auto dst_type = output_q.scalar_type();
@@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
do { \
cudaLaunchConfig_t config = {}; \
config.gridDim = dim3(static_cast<unsigned int>(blocks_x), \
static_cast<unsigned int>(blocks_y)); \
config.gridDim = dim3(static_cast<unsigned int>(row_blocks), \
static_cast<unsigned int>(sf_k_blocks)); \
config.blockDim = dim3(num_threads); \
config.dynamicSmemBytes = 0; \
config.stream = stream; \
@@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
#else
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
do { \
dim3 grid(static_cast<unsigned int>(blocks_x), \
static_cast<unsigned int>(blocks_y)); \
dim3 grid(static_cast<unsigned int>(row_blocks), \
static_cast<unsigned int>(sf_k_blocks)); \
dim3 block(num_threads); \
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
RY> \
+27
View File
@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
// TODO: Remove this once ROCm upgrade to torch 2.11.
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
// But since they are set back to zero once the kernel is finished we can
// hand wave and say that they have no net effect.
//
// The reason to mark 'workspace' as immutable is so that they don't interfere
// with using ScalarType arguments in the ops. If they are marked as mutable,
// pytorch throws an assert in
// 'torch._higher_order_ops._register_effectful_op' that prevents these
// kernels from being torch.compile'd.
// See the following document for more info on custom types and ops that use
// custom types:
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
ops.def(
"machete_supported_schedules("
@@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor workspace, int k, int max_seq_len) -> ()");
// Activation ops
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s, bool use_ue8m0) -> ()");
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
// Activation function used in SwiGLU.
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
@@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
"float alpha=1.0, float beta=0.0) -> ()");
// SwiGLU activation with FP8 quantization.
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
// Activation function used in GeGLU with `none` approximation.
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
// Activation kernels (shared CUDA/ROCm)
ops.impl("persistent_masked_m_silu_mul_quant",
TORCH_BOX(&persistent_masked_m_silu_mul_quant));
ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
-32
View File
@@ -9,28 +9,6 @@
#include <vector>
torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
// Ensure tensor is on CUDA
if (!tensor.is_cuda()) {
throw std::runtime_error("Tensor must be on CUDA device");
}
// Get the raw data pointer
void* data_ptr = tensor.data_ptr();
// Get tensor sizes and strides
std::vector<int64_t> sizes = tensor.sizes().vec();
std::vector<int64_t> strides = tensor.strides().vec();
// Get tensor options (dtype, device)
auto options = tensor.options();
// Create a new tensor from the raw data pointer
auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
return new_tensor;
}
// rms_norm and fused_add_rms_norm declarations also exist in
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
// because the CPU build still uses these torch::Tensor declarations.
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
torch::Tensor& scale);
void persistent_masked_m_silu_mul_quant(
const at::Tensor& input, // (E, T, 2*H)
const at::Tensor& counts, // (E)
at::Tensor& y_q, // (E, T, H) [OUT]
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool use_ue8m0);
void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
+5
View File
@@ -0,0 +1,5 @@
#include "core/registration.h"
// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
// This stub lets Python import vllm._qutlass_C to trigger op registration.
REGISTER_EXTENSION(_qutlass_C)
-40
View File
@@ -20,17 +20,6 @@
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// vLLM custom ops
//
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s,"
"bool use_ue8m0) -> ()");
ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
&persistent_masked_m_silu_mul_quant);
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
#ifdef USE_ROCM
// TODO: Remove this once we upgrade to torch 2.11.
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
&get_cuda_view_from_cpu_tensor);
#endif
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
// kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
// full-cache bf16/fp8 variants).
// Quantization ops
#ifndef USE_ROCM
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
// But since they are set back to zero once the kernel is finished we can
// hand wave and say that they have no net effect.
//
// The reason to mark 'workspace' as immutable is so that they don't interfere
// with using ScalarType arguments in the ops. If they are marked as mutable,
// pytorch throws an assert in
// 'torch._higher_order_ops._register_effectful_op' that prevents these
// kernels from being torch.compile'd.
// See the following document for more info on custom types and ops that use
// custom types:
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
#endif
}
+7 -7
View File
@@ -75,13 +75,13 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
# Install UMD
RUN mkdir neo && \
cd neo && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-core-2_2.34.4+21428_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-opencl-2_2.34.4+21428_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-ocloc_26.18.38308.1-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-opencl-icd_26.18.38308.1-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libigdgmm12_22.10.0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libze-intel-gpu1_26.18.38308.1-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero_1.28.2+u24.04_amd64.deb && \
dpkg -i *.deb && \
cd .. && \
rm -rf neo
+1 -1
View File
@@ -133,7 +133,7 @@ The model should inherit protocol `IsAttentionFree` and also implement class met
For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).
For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
@@ -4,7 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
* **Upstream vLLM compatibility** It wraps around upstream vLLM without modifying its code.
* **Ease of use** Simplified deployment via Helm charts and observability through Grafana dashboards.
* **High performance** Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache), among others.
* **High performance** Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache) (wired up in vLLM via `--kv-offloading-backend lmcache`; see the [LMCache examples](https://github.com/vllm-project/vllm/tree/main/examples/disaggregated/lmcache) and [docs.lmcache.ai](https://docs.lmcache.ai)), among others.
If you are new to Kubernetes, don't worry: in the vLLM production stack [repo](https://github.com/vllm-project/production-stack), we provide a step-by-step [guide](https://github.com/vllm-project/production-stack/blob/main/tutorials/00-install-kubernetes-env.md) and a [short video](https://www.youtube.com/watch?v=EsTJbQtzj0g) to set up everything and get started in **4 minutes**!
+1 -1
View File
@@ -20,7 +20,7 @@ Two main reasons:
Now supports 9 types of connectors:
- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission. LMCache also offers a multi-process (MP) mode via `LMCacheMPConnector`, where a standalone `lmcache server` holds the KV cache shared by one or more vLLM instances; see the [LMCache examples](../../examples/disaggregated/lmcache/README.md) and the [LMCache docs](https://docs.lmcache.ai) for setup.
- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as:
```bash
@@ -203,6 +203,7 @@ the vLLM JSON config.
### kv_connector_extra_config
- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
- `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`.
- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
- `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
@@ -27,6 +27,7 @@ Currently, there are no pre-built XPU wheels.
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
```bash
git clone https://github.com/vllm-project/vllm.git
+34 -14
View File
@@ -1,10 +1,38 @@
# LMCache Examples
This folder demonstrates how to use LMCache for disaggregated prefilling, CPU offloading and KV cache sharing.
This folder demonstrates how to use LMCache with vLLM v1 for KV cache
offloading, disaggregated prefilling, and KV cache sharing.
## 1. Disaggregated Prefill in vLLM v1
## Integration modes
This example demonstrates how to run LMCache with disaggregated prefill using NIXL on a single node.
LMCache integrates with vLLM v1 in two ways:
- **In-process mode** (`LMCacheConnectorV1`): LMCache runs inside the vLLM
process and is configured through environment variables or a YAML config
file (`LMCACHE_CONFIG_FILE`). This is the simplest way to add single-node
CPU/disk offloading.
- **Multi-process (MP) mode** (`LMCacheMPConnector`): LMCache runs as a
standalone server (`lmcache server`) that owns the KV cache storage; one or
more vLLM instances connect to it. This is the recommended mode for
distributed KV storage and for sharing KV cache across instances. See the
[LMCache docs](https://docs.lmcache.ai) for the full MP setup.
## 1. CPU offload (in-process)
- `python cpu_offload_lmcache.py` - CPU offloading with `LMCacheConnectorV1`
for vLLM v1.
## 2. CPU offload (multi-process)
- `bash cpu_offload_lmcache_mp.sh` - CPU offloading with `LMCacheMPConnector`,
using a standalone `lmcache server`. vLLM provides a built-in shortcut for
this setup via `--kv-offloading-backend lmcache` and
`--kv-offloading-size <GiB>`.
## 3. Disaggregated Prefill in vLLM v1
This example demonstrates how to run LMCache with disaggregated prefill using
NIXL on a single node.
### Prerequisites
@@ -46,15 +74,7 @@ The main script generates several log files:
- `decoder.log` - Logs from the decode server
- `proxy.log` - Logs from the proxy server
## 2. CPU Offload Examples
## 4. KV Cache Sharing
- `python cpu_offload_lmcache.py -v v0` - CPU offloading implementation for vLLM v0
- `python cpu_offload_lmcache.py -v v1` - CPU offloading implementation for vLLM v1
## 3. KV Cache Sharing
The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV caches between vLLM v1 instances.
## 4. Disaggregated Prefill in vLLM v0
The `disaggregated_prefill_lmcache_v0.py` provides an example of how to run disaggregated prefill in vLLM v0.
The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV
caches between vLLM v1 instances through a centralized LMCache server.
@@ -1,20 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of cpu offloading
with LMCache in vLLM v1 or v0.
Usage:
Specify vLLM version
-v v0 : Use LMCacheConnector
model = mistralai/Mistral-7B-Instruct-v0.2
(Includes enable_chunked_prefill = True)
-v v1 : Use LMCacheConnectorV1 (default)
model = meta-llama/Meta-Llama-3.1-8B-Instruct
(Without enable_chunked_prefill)
This file demonstrates the example usage of CPU offloading
with LMCache in vLLM v1.
Note that `lmcache` is needed to run this example.
Requirements:
@@ -23,7 +11,6 @@ Learn more about LMCache environment setup, please refer to:
https://docs.lmcache.ai/getting_started/installation.html
"""
import argparse
import contextlib
import os
import time
@@ -39,8 +26,6 @@ from vllm.engine.arg_utils import EngineArgs
def setup_environment_variables():
# LMCache-related environment variables
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Enable local CPU backend in LMCache
@@ -50,9 +35,9 @@ def setup_environment_variables():
@contextlib.contextmanager
def build_llm_with_lmcache(lmcache_connector: str, model: str):
def build_llm_with_lmcache(model: str):
ktc = KVTransferConfig(
kv_connector=lmcache_connector,
kv_connector="LMCacheConnectorV1",
kv_role="kv_both",
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
@@ -92,23 +77,10 @@ def print_output(
print("-" * 50)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"-v",
"--version",
choices=["v0", "v1"],
default="v1",
help="Specify vLLM version (default: v1)",
)
return parser.parse_args()
def main():
lmcache_connector = "LMCacheConnectorV1"
model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
setup_environment_variables()
with build_llm_with_lmcache(lmcache_connector, model) as llm:
with build_llm_with_lmcache(model) as llm:
# This example script runs two requests with a shared prefix.
# Define the shared prompt and specific prompts
shared_prompt = "Hello, how are you?" * 1000
+43
View File
@@ -0,0 +1,43 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# CPU offloading with LMCache in multi-process (MP) mode.
#
# In MP mode, LMCache runs as a standalone server process (`lmcache server`)
# that owns the KV cache storage. One or more vLLM instances connect to it via
# the `LMCacheMPConnector`. This is the recommended way to run LMCache for
# distributed KV storage and for sharing KV cache across vLLM instances.
#
# vLLM ships a built-in shortcut for this setup: pass `--kv-offloading-backend
# lmcache` together with `--kv-offloading-size <GiB>` and vLLM wires up the
# `LMCacheMPConnector` for you (it defaults to the LMCache server at
# tcp://localhost:5555, matching the `lmcache server` default).
#
# Requires `lmcache` to be installed (`pip install lmcache`).
# Learn more: https://docs.lmcache.ai
set -euo pipefail
MODEL=${MODEL:-meta-llama/Meta-Llama-3.1-8B-Instruct}
# 1. Launch the standalone LMCache server (binds tcp://localhost:5555 by
# default). `--l1-size-gb` sets the CPU memory budget for the L1 cache.
echo "Starting LMCache server..."
lmcache server --host localhost --port 5555 --l1-size-gb 5 &
LMCACHE_SERVER_PID=$!
trap 'kill $LMCACHE_SERVER_PID 2>/dev/null || true' EXIT
# 2. Launch vLLM and offload KV cache to the LMCache server.
# The MP connector currently requires the non-hybrid KV cache manager.
echo "Starting vLLM server with LMCache MP offloading..."
vllm serve "$MODEL" \
--port 8000 \
--kv-offloading-size 5 \
--kv-offloading-backend lmcache \
--disable-hybrid-kv-cache-manager
# Equivalent explicit configuration (instead of the two flags above):
# --kv-transfer-config \
# '{"kv_connector":"LMCacheMPConnector","kv_role":"kv_both",
# "kv_connector_extra_config":{"lmcache.mp.host":"tcp://localhost",
# "lmcache.mp.port":5555}}'
@@ -1,144 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of disaggregated prefilling
with LMCache.
We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
and launch an additional LMCache server.
KV cache is transferred in the following manner:
vLLM prefill node -> LMCache server -> vLLM decode node.
Note that `pip install lmcache` is needed to run this example.
Learn more about LMCache in https://github.com/LMCache/LMCache.
"""
import os
import subprocess
import time
from multiprocessing import Event, Process
from lmcache.experimental.cache_engine import LMCacheEngineBuilder
from lmcache.integration.vllm.utils import ENGINE_NAME
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
# LMCache-related environment variables
# The port to start LMCache server
port = 8100
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Disable local CPU backend in LMCache
os.environ["LMCACHE_LOCAL_CPU"] = "False"
# Set local CPU memory buffer limit to 5.0 GB
os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0"
# Set the remote URL for LMCache server
os.environ["LMCACHE_REMOTE_URL"] = f"lm://localhost:{port}"
# Set the serializer/deserializer between vllm and LMCache server
# `naive` indicates using raw bytes of the tensor without any compression
os.environ["LMCACHE_REMOTE_SERDE"] = "naive"
prompts = [
"Hello, how are you?" * 1000,
]
def run_prefill(prefill_done, prompts):
# We use GPU 0 for prefill node.
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
ktc = KVTransferConfig(
kv_connector="LMCacheConnector",
kv_role="kv_producer",
kv_rank=0,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
# memory. Reduce the value if your GPU has less memory.
llm = LLM(
model="mistralai/Mistral-7B-Instruct-v0.2",
kv_transfer_config=ktc,
max_model_len=8000,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
# llm.generate(prompts, sampling_params)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
generated_text = output.outputs[0].text
print(f"Generated text: {generated_text!r}")
print("Prefill node is finished.")
prefill_done.set()
# Clean up lmcache backend
LMCacheEngineBuilder.destroy(ENGINE_NAME)
def run_decode(prefill_done, prompts, timeout=1):
# We use GPU 1 for decode node.
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
ktc = KVTransferConfig(
kv_connector="LMCacheConnector",
kv_role="kv_consumer",
kv_rank=1,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
# of memory. Reduce the value if your GPU has less memory.
llm = LLM(
model="mistralai/Mistral-7B-Instruct-v0.2",
kv_transfer_config=ktc,
max_model_len=8000,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
print("Waiting for prefill node to finish...")
prefill_done.wait()
time.sleep(timeout)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
generated_text = output.outputs[0].text
print(f"Generated text: {generated_text!r}")
# Clean up lmcache backend
LMCacheEngineBuilder.destroy(ENGINE_NAME)
def run_lmcache_server(port):
server_proc = subprocess.Popen(
["python", "-m", "lmcache.experimental.server", "localhost", str(port)]
)
return server_proc
def main():
prefill_done = Event()
prefill_process = Process(target=run_prefill, args=(prefill_done, prompts))
decode_process = Process(target=run_decode, args=(prefill_done, prompts))
lmcache_server_process = run_lmcache_server(port)
# Start prefill node
prefill_process.start()
# Start decode node
decode_process.start()
# Clean up the processes
decode_process.join()
prefill_process.terminate()
lmcache_server_process.terminate()
lmcache_server_process.wait()
if __name__ == "__main__":
main()
@@ -30,7 +30,6 @@ if [[ $1 == "prefiller" ]]; then
UCX_TLS=cuda_ipc,cuda_copy,tcp \
LMCACHE_CONFIG_FILE=$prefill_config_file \
LMCACHE_USE_EXPERIMENTAL=True \
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES=0 \
@@ -47,7 +46,6 @@ elif [[ $1 == "decoder" ]]; then
UCX_TLS=cuda_ipc,cuda_copy,tcp \
LMCACHE_CONFIG_FILE=$decode_config_file \
LMCACHE_USE_EXPERIMENTAL=True \
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES=1 \
@@ -26,8 +26,6 @@ from vllm.config import KVTransferConfig
# LMCache-related environment variables
# The port to start LMCache server
port = 8100
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Disable local CPU backend in LMCache
+3 -2
View File
@@ -11,13 +11,14 @@ transformers >= 5.5.3
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides.
starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1
aiohttp >= 3.13.3
openai >= 2.0.0 # For Responses API with reasoning content
pydantic >= 2.12.0
prometheus_client >= 0.18.0
pillow # Required for image processing
prometheus-fastapi-instrumentator >= 7.0.0
prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0
tiktoken >= 0.6.0 # Required for DBRX tokenizer
lm-format-enforcer == 0.11.3
llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
+1 -1
View File
@@ -40,7 +40,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test.
schemathesis>=4.0.0 # Required for openai schema test.
# quantization
bitsandbytes==0.49.2
buildkite-test-collector==0.1.9
+23 -48
View File
@@ -35,14 +35,11 @@ arctic-inference==0.1.1
# via -r requirements/test/cuda.in
argcomplete==3.5.1
# via datamodel-code-generator
arrow==1.3.0
# via isoduration
attrs==24.2.0
# via
# aiohttp
# hypothesis
# jsonschema
# pytest-subtests
# referencing
audioread==3.0.1
# via librosa
@@ -57,9 +54,7 @@ azure-identity==1.25.2
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test/cuda.in
# schemathesis
# via -r requirements/test/cuda.in
bitsandbytes==0.49.2
# via -r requirements/test/cuda.in
black==24.10.0
@@ -110,7 +105,6 @@ colorama==0.4.6
# via
# perceptron
# sacrebleu
# schemathesis
colorful==0.5.6
# via ray
colorlog==6.10.1
@@ -183,7 +177,7 @@ et-xmlfile==2.0.0
# via openpyxl
evaluate==0.4.3
# via lm-eval
fastapi==0.128.0
fastapi==0.136.3
# via
# -c requirements/common.txt
# gpt-oss
@@ -206,8 +200,6 @@ filelock==3.16.1
# virtualenv
fonttools==4.55.0
# via matplotlib
fqdn==1.5.1
# via jsonschema
frozendict==2.4.6
# via einx
frozenlist==1.5.0
@@ -269,7 +261,7 @@ h11==0.14.0
# uvicorn
h2==4.3.0
# via httpx
harfile==0.3.0
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
# via huggingface-hub
@@ -309,7 +301,7 @@ hypothesis==6.131.0
# hypothesis-graphql
# hypothesis-jsonschema
# schemathesis
hypothesis-graphql==0.11.1
hypothesis-graphql==0.13.0
# via schemathesis
hypothesis-jsonschema==0.23.1
# via schemathesis
@@ -318,7 +310,6 @@ idna==3.10
# anyio
# email-validator
# httpx
# jsonschema
# requests
# yarl
imagehash==4.3.2
@@ -335,8 +326,6 @@ instanttensor==0.1.5
# via -r requirements/test/cuda.in
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==5.13.2
# via datamodel-code-generator
jinja2==3.1.6
@@ -356,15 +345,14 @@ joblib==1.4.2
# librosa
# nltk
# scikit-learn
jsonpointer==3.0.0
# via jsonschema
jsonschema==4.23.0
# via
# -c requirements/common.txt
# hypothesis-jsonschema
# mistral-common
# ray
# schemathesis
jsonschema-rs==0.46.5
# via schemathesis
jsonschema-specifications==2024.10.1
# via jsonschema
junit-xml==1.9
@@ -715,18 +703,20 @@ pydantic-core==2.41.1
pydantic-extra-types==2.10.5
# via mistral-common
pygments==2.18.0
# via rich
# via
# pytest
# rich
pyjwt==2.11.0
# via msal
pyparsing==3.2.0
# via matplotlib
pyrate-limiter==3.7.0
pyrate-limiter==4.4.0
# via schemathesis
pystemmer==3.0.0
# via mteb
pytablewriter==1.2.0
# via lm-eval
pytest==8.3.5
pytest==9.1.0
# via
# -r requirements/test/cuda.in
# buildkite-test-collector
@@ -737,10 +727,9 @@ pytest==8.3.5
# pytest-mock
# pytest-rerunfailures
# pytest-shard
# pytest-subtests
# pytest-timeout
# schemathesis
pytest-asyncio==0.24.0
pytest-asyncio==1.4.0
# via -r requirements/test/cuda.in
pytest-cov==6.3.0
# via -r requirements/test/cuda.in
@@ -752,13 +741,10 @@ pytest-rerunfailures==14.0
# via -r requirements/test/cuda.in
pytest-shard==0.1.2
# via -r requirements/test/cuda.in
pytest-subtests==0.14.1
# via schemathesis
pytest-timeout==2.3.1
# via -r requirements/test/cuda.in
python-dateutil==2.9.0.post0
# via
# arrow
# botocore
# matplotlib
# pandas
@@ -829,15 +815,12 @@ requests==2.32.3
# tiktoken
responses==0.25.3
# via genai-perf
rfc3339-validator==0.1.4
# via jsonschema
rfc3987==1.3.8
# via jsonschema
rich==13.9.4
# via
# genai-perf
# mteb
# perceptron
# schemathesis
# typer
rouge-score==0.1.2
# via lm-eval
@@ -868,7 +851,7 @@ safetensors==0.7.0
# segmentation-models-pytorch
# timm
# transformers
schemathesis==3.39.15
schemathesis==4.21.6
# via -r requirements/test/cuda.in
scikit-image==0.25.2
# via albumentations
@@ -912,7 +895,6 @@ six==1.16.0
# junit-xml
# opencensus
# python-dateutil
# rfc3339-validator
# rouge-score
smart-open==7.1.0
# via ray
@@ -938,10 +920,10 @@ sqlalchemy==2.0.41
# optuna
sqlitedict==2.1.0
# via lm-eval
starlette==0.50.0
starlette==1.3.1
# via
# -c requirements/common.txt
# fastapi
# schemathesis
# starlette-testclient
starlette-testclient==0.4.1
# via schemathesis
@@ -966,6 +948,7 @@ tenacity==9.1.2
# gpt-oss
# lm-eval
# plotly
# schemathesis
tensorizer==2.10.1
# via -r requirements/test/cuda.in
termcolor==3.1.0
@@ -990,10 +973,6 @@ tokenizers==0.22.2
# -c requirements/common.txt
# -r requirements/test/cuda.in
# transformers
tomli==2.2.1
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch==2.11.0+cu130
# via
# -c requirements/cuda.txt
@@ -1066,8 +1045,6 @@ typer==0.15.2
# huggingface-hub
# perceptron
# transformers
types-python-dateutil==2.9.0.20241206
# via arrow
typing-extensions==4.15.0
# via
# -c requirements/common.txt
@@ -1092,6 +1069,8 @@ typing-extensions==4.15.0
# pydantic
# pydantic-core
# pydantic-extra-types
# pytest-asyncio
# schemathesis
# sentence-transformers
# sqlalchemy
# starlette
@@ -1099,11 +1078,11 @@ typing-extensions==4.15.0
# typer
# typing-inspection
typing-inspection==0.4.2
# via pydantic
# via
# fastapi
# pydantic
tzdata==2024.2
# via pandas
uri-template==1.3.0
# via jsonschema
urllib3==2.2.3
# via
# blobfile
@@ -1122,8 +1101,6 @@ vocos==0.1.0
# via -r requirements/test/cuda.in
wcwidth==0.2.13
# via ftfy
webcolors==24.11.1
# via jsonschema
werkzeug==3.1.3
# via schemathesis
word2number==1.1
@@ -1135,8 +1112,6 @@ xxhash==3.5.0
# datasets
# evaluate
yarl==1.17.1
# via
# aiohttp
# schemathesis
# via aiohttp
zipp==3.23.0
# via importlib-metadata
+1 -1
View File
@@ -31,7 +31,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test.
schemathesis>=4.0.0 # Required for openai schema test.
# quantization
bitsandbytes>=0.49.2
buildkite-test-collector==0.1.9
+1 -1
View File
@@ -39,7 +39,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test
schemathesis>=4.0.0 # Required for openai schema test
# quantization
bitsandbytes==0.49.2
buildkite-test-collector==0.1.9
+21 -47
View File
@@ -51,15 +51,12 @@ arctic-inference==0.1.1
# via -r requirements/test/rocm.in
argcomplete==3.6.3
# via datamodel-code-generator
arrow==1.4.0
# via isoduration
astor==0.8.1
# via depyf
attrs==26.1.0
# via
# aiohttp
# jsonschema
# pytest-subtests
# referencing
audioread==3.0.1
# via librosa
@@ -74,9 +71,7 @@ azure-identity==1.25.3
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test/rocm.in
# schemathesis
# via -r requirements/test/rocm.in
bitsandbytes==0.49.2
# via -r requirements/test/rocm.in
black==26.3.1
@@ -139,7 +134,6 @@ colorama==0.4.6
# via
# perceptron
# sacrebleu
# schemathesis
colorful==0.5.8
# via ray
colorlog==6.10.1
@@ -258,8 +252,6 @@ filelock==3.25.2
# virtualenv
fonttools==4.62.1
# via matplotlib
fqdn==1.5.1
# via jsonschema
frozendict==2.4.7
# via einx
frozenlist==1.8.0
@@ -328,7 +320,7 @@ h11==0.16.0
# uvicorn
h2==4.3.0
# via httpx
harfile==0.4.0
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
# via huggingface-hub
@@ -378,7 +370,7 @@ hypothesis==6.151.9
# hypothesis-graphql
# hypothesis-jsonschema
# schemathesis
hypothesis-graphql==0.12.0
hypothesis-graphql==0.13.0
# via schemathesis
hypothesis-jsonschema==0.23.1
# via schemathesis
@@ -387,7 +379,6 @@ idna==3.11
# anyio
# email-validator
# httpx
# jsonschema
# requests
# yarl
ijson==3.5.0
@@ -408,8 +399,6 @@ interegular==0.3.3
# via lm-format-enforcer
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==8.0.1
# via datamodel-code-generator
jinja2==3.1.6
@@ -435,8 +424,6 @@ joblib==1.5.3
# librosa
# nltk
# scikit-learn
jsonpointer==3.1.0
# via jsonschema
jsonschema==4.26.0
# via
# -c requirements/common.txt
@@ -445,7 +432,8 @@ jsonschema==4.26.0
# mcp
# mistral-common
# ray
# schemathesis
jsonschema-rs==0.46.5
# via schemathesis
jsonschema-specifications==2025.9.1
# via jsonschema
junit-xml==1.9
@@ -792,7 +780,7 @@ prometheus-client==0.24.1
# opentelemetry-exporter-prometheus
# prometheus-fastapi-instrumentator
# ray
prometheus-fastapi-instrumentator==7.1.0
prometheus-fastapi-instrumentator==8.0.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -876,20 +864,22 @@ pydantic-settings==2.13.1
# fastapi
# mcp
pygments==2.19.2
# via rich
# via
# pytest
# rich
pyjwt==2.12.1
# via
# mcp
# msal
pyparsing==3.3.2
# via matplotlib
pyrate-limiter==3.9.0
pyrate-limiter==4.4.0
# via schemathesis
pystemmer==3.0.0
# via mteb
pytablewriter==1.2.1
# via lm-eval
pytest==8.3.5
pytest==9.1.0
# via
# -r requirements/test/rocm.in
# buildkite-test-collector
@@ -900,10 +890,9 @@ pytest==8.3.5
# pytest-mock
# pytest-rerunfailures
# pytest-shard
# pytest-subtests
# pytest-timeout
# schemathesis
pytest-asyncio==0.24.0
pytest-asyncio==1.4.0
# via -r requirements/test/rocm.in
pytest-cov==6.3.0
# via -r requirements/test/rocm.in
@@ -915,13 +904,10 @@ pytest-rerunfailures==14.0
# via -r requirements/test/rocm.in
pytest-shard==0.1.2
# via -r requirements/test/rocm.in
pytest-subtests==0.14.2
# via schemathesis
pytest-timeout==2.3.1
# via -r requirements/test/rocm.in
python-dateutil==2.9.0.post0
# via
# arrow
# botocore
# matplotlib
# pandas
@@ -1016,16 +1002,13 @@ requests==2.32.5
# tiktoken
responses==0.26.0
# via genai-perf
rfc3339-validator==0.1.4
# via jsonschema
rfc3987==1.3.8
# via jsonschema
rich==14.3.3
# via
# genai-perf
# mteb
# perceptron
# rich-toolkit
# schemathesis
# typer
rich-toolkit==0.19.7
# via
@@ -1063,7 +1046,7 @@ safetensors==0.7.0
# segmentation-models-pytorch
# timm
# transformers
schemathesis==3.39.15
schemathesis==4.21.6
# via -r requirements/test/rocm.in
scikit-image==0.26.0
# via albumentations
@@ -1120,7 +1103,6 @@ six==1.17.0
# junit-xml
# opencensus
# python-dateutil
# rfc3339-validator
# rouge-score
smart-open==7.5.1
# via ray
@@ -1149,13 +1131,14 @@ sqlitedict==2.1.0
# via lm-eval
sse-starlette==3.3.4
# via mcp
starlette==0.52.1
starlette==1.3.1
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
# fastapi
# mcp
# model-hosting-container-standards
# prometheus-fastapi-instrumentator
# schemathesis
# sse-starlette
# starlette-testclient
starlette-testclient==0.4.1
@@ -1182,6 +1165,7 @@ tenacity==9.1.4
# via
# gpt-oss
# lm-eval
# schemathesis
tensorizer==2.10.1
# via
# -c requirements/rocm.txt
@@ -1215,10 +1199,6 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
# transformers
tomli==2.4.0
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch-c-dlpack-ext==0.1.5
# via tilelang
tqdm==4.67.3
@@ -1301,8 +1281,10 @@ typing-extensions==4.15.0
# pydantic
# pydantic-core
# pydantic-extra-types
# pytest-asyncio
# referencing
# rich-toolkit
# schemathesis
# sentence-transformers
# sqlalchemy
# starlette
@@ -1317,10 +1299,6 @@ typing-inspection==0.4.2
# mcp
# pydantic
# pydantic-settings
tzdata==2025.3
# via arrow
uri-template==1.3.0
# via jsonschema
urllib3==2.6.3
# via
# blobfile
@@ -1351,8 +1329,6 @@ watchfiles==1.1.1
# uvicorn
wcwidth==0.6.0
# via ftfy
webcolors==25.10.0
# via jsonschema
websockets==16.0
# via uvicorn
werkzeug==3.1.6
@@ -1370,9 +1346,7 @@ xxhash==3.6.0
# datasets
# evaluate
yarl==1.23.0
# via
# aiohttp
# schemathesis
# via aiohttp
z3-solver==4.15.4.0
# via tilelang
zipp==3.23.0
+2 -1
View File
@@ -593,8 +593,9 @@ soxr==0.5.0.post1
# mistral-common
sqlitedict==2.1.0
# via lm-eval
starlette==1.0.0
starlette==1.3.1
# via
# -c requirements/common.txt
# fastapi
# starlette-testclient
starlette-testclient==0.4.1
+1 -1
View File
@@ -17,4 +17,4 @@ torchaudio
torchvision
auto_round_lib>=0.13.3
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10/vllm_xpu_kernels-0.1.10-cp38-abi3-manylinux_2_28_x86_64.whl
+4 -1
View File
@@ -769,6 +769,7 @@ class precompiled_wheel_utils:
"vllm/_C.abi3.so",
"vllm/_C_stable_libtorch.abi3.so",
"vllm/_moe_C_stable_libtorch.abi3.so",
"vllm/_qutlass_C.abi3.so",
"vllm/_flashmla_C.abi3.so",
"vllm/_flashmla_extension_C.abi3.so",
"vllm/_sparse_flashmla_C.abi3.so",
@@ -1135,6 +1136,7 @@ if _is_cuda():
# DeepGEMM requires CUDA 12.3+ (SM90/SM100)
# Optional since it won't build on unsupported architectures
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
# fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
@@ -1149,7 +1151,8 @@ if _is_cpu():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _build_custom_ops():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
@@ -22,7 +22,7 @@ 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._aiter_ops import 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 (
@@ -83,9 +83,8 @@ class _ViewDoubleQuantModel(torch.nn.Module):
[_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",
@pytest.mark.skip(
reason="Skipping for now because pytorch compiler removes one the two quant ops"
)
def test_double_aiter_rms_fp8_group_quant_fusion(
model_cls: type[torch.nn.Module],
+2 -2
View File
@@ -175,7 +175,7 @@ MULTIMODAL_MODELS = {
"facebook/chameleon-7b": PPTestSettings.fast(),
"adept/fuyu-8b": PPTestSettings.fast(),
"zai-org/glm-4v-9b": PPTestSettings.fast(),
"OpenGVLab/InternVL2-1B": PPTestSettings.fast(),
"OpenGVLab/InternVL3-1B": PPTestSettings.fast(),
"llava-hf/llava-1.5-7b-hf": PPTestSettings.fast(),
"llava-hf/llava-v1.6-mistral-7b-hf": PPTestSettings.fast(),
"llava-hf/LLaVA-NeXT-Video-7B-hf": PPTestSettings.fast(),
@@ -203,7 +203,7 @@ TEST_MODELS = [
"intfloat/e5-mistral-7b-instruct",
"BAAI/bge-multilingual-gemma2",
# [MULTIMODAL GENERATION]
"OpenGVLab/InternVL2-1B",
"OpenGVLab/InternVL3-1B",
"microsoft/Phi-3.5-vision-instruct",
"fixie-ai/ultravox-v0_5-llama-3_2-1b",
# [LANGUAGE GENERATION - HYBRID ARCH]
+193
View File
@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")
assert args.model == "s3://bucket/model"
assert args.tokenizer == "s3://bucket/tokenizer"
class TestDeviceIds:
def test_device_ids_with_cvd_out_of_range(self, monkeypatch):
"""--device-ids index beyond the CVD set raises ValueError."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "4,5")
args = EngineArgs(model="m", device_ids=[0, 2])
with pytest.raises(ValueError, match="out of range"):
args._resolve_device_ids()
def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):
"""--device-ids are CVD-local indices resolved to physical ids."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "4,5")
args = EngineArgs(model="m", device_ids=[0, 1])
assert args._resolve_device_ids() == [4, 5]
def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):
"""--device-ids support UUID CVD values resolved by the platform."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
),
)
args = EngineArgs(model="m", device_ids=[0, 1])
assert args._resolve_device_ids() == [4, 5]
def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):
"""UUID --device-ids are resolved to physical IDs immediately."""
from vllm.platforms import current_platform
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),
)
args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])
assert args._resolve_device_ids() == [4]
def test_device_ids_reject_mixed_integer_and_uuid_args(self):
"""--device-ids must not mix CVD indices and UUIDs."""
args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])
with pytest.raises(ValueError, match="must not mix"):
args._resolve_device_ids()
def test_no_device_ids(self):
"""No --device-ids returns None."""
args = EngineArgs(model="m")
assert args._resolve_device_ids() is None
def test_cli_parsing(self):
"""--device-ids parses comma-separated string from CLI."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])
assert parsed.device_ids == [0, 2, 4]
def test_cli_parsing_uuid(self):
"""--device-ids parses comma-separated UUID strings from CLI."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(
["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]
)
assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]
def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):
"""assigned_physical_gpu_ids are already physical and not composed with CVD."""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
assert current_platform.device_id_to_physical_device_id(0) == 4
assert current_platform.device_id_to_physical_device_id(1) == 5
assert current_platform.logical_device_id_to_visible_device_id(0) == 0
assert current_platform.logical_device_id_to_visible_device_id(1) == 1
def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):
"""Physical IDs map back to visible ordinals when CVD uses UUIDs."""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])
monkeypatch.setenv(
current_platform.device_control_env_var,
"GPU-abcd1234,GPU-ef567890",
)
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
),
)
assert current_platform.logical_device_id_to_visible_device_id(0) == 1
def test_device_ids_reject_duplicates(self):
"""--device-ids must not contain duplicate entries."""
args = EngineArgs(model="m", device_ids=[2, 2])
with pytest.raises(ValueError, match="duplicates"):
args._resolve_device_ids()
def test_cli_parsing_strips_whitespace(self):
"""--device-ids tolerates whitespace around commas."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])
assert parsed.device_ids == [0, 2, 4]
def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):
"""visible_device_id_to_physical_device_id maps torch device ordinals,
independent of the logical-to-physical mapping.
Regression test: CustomAllreduce passes device.index (a visible
ordinal) and must not index into assigned_physical_gpu_ids, which
raised IndexError for non-identity --device-ids like [2, 3].
"""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])
monkeypatch.delenv(current_platform.device_control_env_var, raising=False)
# CVD unset: visible ordinal == physical ID, even beyond the
# assigned list's length.
assert current_platform.visible_device_id_to_physical_device_id(2) == 2
assert current_platform.visible_device_id_to_physical_device_id(3) == 3
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
assert current_platform.visible_device_id_to_physical_device_id(1) == 5
with pytest.raises(IndexError, match="out of range"):
current_platform.visible_device_id_to_physical_device_id(2)
class TestDpDeviceIdSharding:
def test_dp_supervisor_device_ids_stay_env_relative(self):
"""Regression test: the DP supervisor must pass env-relative indices,
not physical IDs, because each child re-resolves --device-ids
against its inherited device-control env var."""
import argparse
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
args = argparse.Namespace(
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None
)
assert _build_device_ids(args, local_rank=0) == [0, 1]
assert _build_device_ids(args, local_rank=1) == [2, 3]
def test_dp_supervisor_shards_user_device_ids(self):
"""User-provided --device-ids are sharded across DP children."""
import argparse
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
args = argparse.Namespace(
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]
)
assert _build_device_ids(args, local_rank=0) == [4, 5]
assert _build_device_ids(args, local_rank=1) == [6, 7]
with pytest.raises(ValueError, match="needs devices"):
_build_device_ids(args, local_rank=2)
def test_dp_rank_shards_user_assigned_gpu_ids(self):
"""get_physical_gpu_ids_for_local_dp_rank slices the user-provided
--device-ids list instead of recomputing from the env var."""
from vllm.platforms import current_platform
from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
evar = current_platform.device_control_env_var
assert get_physical_gpu_ids_for_local_dp_rank(
evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
) == [6, 7]
with pytest.raises(ValueError, match="needs devices"):
get_physical_gpu_ids_for_local_dp_rank(
evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
)
@@ -8,6 +8,8 @@ AnthropicServingMessages._convert_anthropic_to_openai_request().
Also covers extended-thinking edge cases such as ``redacted_thinking``
blocks echoed back by Anthropic clients, and streaming conversion in
``message_stream_converter``.
Also covers cache usage computation in ``_build_anthropic_usage``.
"""
import json
@@ -18,7 +20,11 @@ import pytest
from vllm.entrypoints.anthropic.protocol import (
AnthropicMessagesRequest,
)
from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
from vllm.entrypoints.anthropic.serving import (
AnthropicServingMessages,
_build_anthropic_usage,
_get_cached_tokens,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionResponseStreamChoice,
ChatCompletionStreamResponse,
@@ -27,6 +33,7 @@ from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
PromptTokenUsageInfo,
UsageInfo,
)
@@ -653,6 +660,108 @@ class TestThinkingBlockConversion:
assert asst.get("content") == "Hi!"
# ======================================================================
# Cache usage computation
# ======================================================================
class TestGetCachedTokens:
"""Tests for _get_cached_tokens helper."""
def test_none_usage(self):
assert _get_cached_tokens(None) is None
def test_no_prompt_tokens_details(self):
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
assert _get_cached_tokens(usage) is None
def test_cached_tokens_present(self):
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
)
assert _get_cached_tokens(usage) == 80
def test_cached_tokens_zero(self):
"""Zero cached tokens should return 0, not None."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
)
assert _get_cached_tokens(usage) == 0
def test_cached_tokens_none_in_details(self):
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
)
assert _get_cached_tokens(usage) is None
class TestBuildAnthropicUsage:
"""Tests for _build_anthropic_usage helper.
Anthropic defines: total_input = input_tokens + cache_read + cache_creation
vLLM's prompt_tokens is the total.
"""
def test_no_cache_info(self):
"""When cache info is unavailable, return raw prompt_tokens."""
result = _build_anthropic_usage(100, 10, None)
assert result.input_tokens == 100
assert result.output_tokens == 10
assert result.cache_read_input_tokens is None
assert result.cache_creation_input_tokens is None
def test_cache_hit(self):
"""When cache is hit, input_tokens excludes cached tokens."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 20 # 100 - 80
assert result.output_tokens == 10
assert result.cache_read_input_tokens == 80
assert result.cache_creation_input_tokens == 0
def test_zero_cached_tokens(self):
"""Zero cached tokens should still set cache_creation to 0."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 100 # 100 - 0
assert result.cache_read_input_tokens == 0
assert result.cache_creation_input_tokens == 0
def test_all_tokens_cached(self):
"""When all tokens are cached, input_tokens should be 0."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 0
assert result.cache_read_input_tokens == 100
assert result.cache_creation_input_tokens == 0
def test_no_prompt_tokens_details(self):
"""UsageInfo without prompt_tokens_details returns no cache info."""
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 100
assert result.cache_read_input_tokens is None
assert result.cache_creation_input_tokens is None
class TestInlineSystemMessageInMessagesArray:
"""Verify that ``role: system`` messages embedded inside the ``messages``
array are preserved in their original position.
@@ -1096,3 +1205,179 @@ class TestMessageStartIncludesTypeAndRole:
message = events[0][1]["message"]
assert message["type"] == "message"
assert message["role"] == "assistant"
class TestStreamingCacheUsageSemantics:
"""Locks in the documented streaming behavior of cache usage fields.
vLLM's OpenAI chat completion streaming only attaches
``prompt_tokens_details`` to the terminal usage chunk. The Anthropic layer
mirrors that contract: cache fields are omitted on ``message_start`` (key
absence signals "unknown") and populated on ``message_delta`` (the final
cumulative count). This is intentionally consistent with vLLM's OpenAI
behavior, even though Anthropic's upstream API populates cache fields on
``message_start``; closing that gap requires plumbing cache info into the
first chunk at the OpenAI layer, which is out of scope here.
"""
@pytest.mark.asyncio
async def test_streaming_cache_fields_absent_then_populated(self):
"""First chunk lacks prompt_tokens_details (vLLM contract);
message_start omits cache fields. The final chunk carries
prompt_tokens_details, so message_delta carries resolved values."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant", content="hi"),
usage=UsageInfo(prompt_tokens=100, total_tokens=100),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(
prompt_tokens=100,
completion_tokens=5,
total_tokens=105,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
output.append(event)
events = _parse_sse_events(output)
# message_start: cache fields unknown → omitted from JSON entirely.
start_usage = events[0][1]["message"]["usage"]
assert events[0][0] == "message_start"
assert start_usage["input_tokens"] == 100
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
# message_delta: authoritative usage with cache fields populated.
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert delta_usage["input_tokens"] == 20 # 100 - 80
assert delta_usage["cache_read_input_tokens"] == 80
assert delta_usage["cache_creation_input_tokens"] == 0
@pytest.mark.asyncio
async def test_streaming_no_cache_hit(self):
"""When the final chunk reports cached_tokens=0, message_delta carries
cache fields = 0 (cache miss); message_start still omits them."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=50, total_tokens=50),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(
prompt_tokens=50,
completion_tokens=5,
total_tokens=55,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
output.append(event)
events = _parse_sse_events(output)
start_usage = events[0][1]["message"]["usage"]
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert start_usage["input_tokens"] == 50
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
assert delta_usage["input_tokens"] == 50 # 50 - 0
assert delta_usage["cache_read_input_tokens"] == 0
assert delta_usage["cache_creation_input_tokens"] == 0
@pytest.mark.asyncio
async def test_streaming_no_prompt_tokens_details_at_all(self):
"""If --enable-prompt-tokens-details is off, no chunk carries cache
info; both message_start and message_delta omit cache fields."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=30, total_tokens=30),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32),
)
yield "data: [DONE]"
converter = _make_stream_converter()
output = []
async for event in converter.message_stream_converter(sse_input()):
output.append(event)
events = _parse_sse_events(output)
start_usage = events[0][1]["message"]["usage"]
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
assert "cache_read_input_tokens" not in delta_usage
assert "cache_creation_input_tokens" not in delta_usage
# ======================================================================
# Auto-detection of system-first template requirement
# ======================================================================
Q35_TEMPLATE = (
"{%- for message in messages %}"
"{%- if message.role == 'system' %}"
"{%- if not loop.first %}"
"{{- raise_exception('System message must be at the beginning.') }}"
"{%- endif %}"
"{%- endif %}"
"{%- endfor %}"
)
class TestDetectMergeInlineSystem:
"""Verify _detect_merge_inline_system auto-detection.
Tests three scenarios:
1. Template with system-first guard (e.g. Qwen) → merge needed
2. Template without restrictions → no merge, cache-friendly
3. No template provided → safe default: merge
"""
def test_qwen_template_requires_merge(self):
"""Template with loop.first guard rejects mid-conversation system."""
assert (
AnthropicServingMessages._detect_merge_inline_system(Q35_TEMPLATE) is True
)
def test_no_restriction_no_merge(self):
"""Template without restriction accepts mid-conversation system."""
assert (
AnthropicServingMessages._detect_merge_inline_system(
"{%- for message in messages %}"
"{{- message.role }}: {{ message.content }}\n"
"{%- endfor %}"
)
is False
)
def test_no_template_defaults_merge(self):
"""No chat_template → conservative default: merge."""
assert AnthropicServingMessages._detect_merge_inline_system(None) is True
@@ -364,7 +364,7 @@ class MockVLLMServer:
await self._serve_task
def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
def launch_mock_vllm(child_args: argparse.Namespace):
logger.info("Launching mock vLLM on port %s", child_args.port)
mock_vllm = MockVLLMServer(
port=child_args.port,
@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]
def launch_mock_vllm_with_drain(
child_args: argparse.Namespace, env_updates: dict[str, str]
child_args: argparse.Namespace,
):
logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
mock_vllm = MockVLLMServer(
+57 -44
View File
@@ -6,15 +6,22 @@ from typing import Final
import pytest
import schemathesis
from hypothesis import HealthCheck, settings
from schemathesis import GenerationConfig
from schemathesis.models import Case
from schemathesis import GenerationMode
from schemathesis.config import (
ChecksConfig,
CoveragePhaseConfig,
GenerationConfig,
PhasesConfig,
PositiveDataAcceptanceConfig,
ProjectConfig,
ProjectsConfig,
SchemathesisConfig,
)
from vllm.platforms import current_platform
from ...utils import RemoteOpenAIServer
schemathesis.experimental.OPEN_API_3_1.enable()
MODEL_NAME = "HuggingFaceTB/SmolVLM-256M-Instruct"
MAXIMUM_IMAGES = 2
_ROCM_TIMEOUT_MULTIPLIER = 3 if current_platform.is_rocm() else 1
@@ -44,21 +51,38 @@ def server():
@pytest.fixture(scope="module")
def get_schema(server):
# avoid generating null (\x00) bytes in strings during test case generation
return schemathesis.openapi.from_uri(
return schemathesis.openapi.from_url(
f"{server.url_root}/openapi.json",
generation_config=GenerationConfig(allow_x00=False),
config=SchemathesisConfig(
projects=ProjectsConfig(
default=ProjectConfig(
generation=GenerationConfig(
allow_x00=False,
modes=[GenerationMode.POSITIVE],
),
checks=ChecksConfig(
positive_data_acceptance=PositiveDataAcceptanceConfig(
enabled=False,
),
),
phases=PhasesConfig(
coverage=CoveragePhaseConfig(enabled=False),
),
),
),
),
)
schema = schemathesis.from_pytest_fixture("get_schema")
schema = schemathesis.pytest.from_fixture("get_schema")
@schemathesis.hook
def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
def before_generate_case(context: schemathesis.HookContext, strategy):
op = context.operation
assert op is not None
def no_invalid_types(case: schemathesis.models.Case):
def no_invalid_types(case: schemathesis.Case):
"""
Skips tool_calls with `"type": "custom"` which schemathesis incorrectly
generates instead of the valid `"type": "function"`.
@@ -68,39 +92,25 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
-d '{"messages": [{"role": "assistant", "tool_calls": [{"custom": {"input": "", "name": ""}, "id": "", "type": "custom"}]}]}' \
http://localhost:8000/v1/chat/completions
""" # noqa: E501
if hasattr(case, "body") and isinstance(case.body, dict):
if (
"messages" in case.body
and isinstance(case.body["messages"], list)
and len(case.body["messages"]) > 0
):
for message in case.body["messages"]:
if not isinstance(message, dict):
continue
if (
hasattr(case, "body")
and isinstance(case.body, dict)
and "messages" in case.body
and isinstance(case.body["messages"], list)
and len(case.body["messages"]) > 0
):
for message in case.body["messages"]:
if not isinstance(message, dict):
continue
tool_calls = message.get("tool_calls", [])
if isinstance(tool_calls, list):
for tool_call in tool_calls:
if isinstance(tool_call, dict):
if tool_call.get("type") != "function":
return False
if "custom" in tool_call:
return False
# Sometimes structured_outputs.grammar is generated to be empty
# Causing a server error in EBNF grammar parsing
# https://github.com/vllm-project/vllm/pull/22587#issuecomment-3195253421
structured_outputs = case.body.get("structured_outputs", {})
grammar = (
structured_outputs.get("grammar")
if isinstance(structured_outputs, dict)
else None
)
if grammar == "":
# Allow None (will be handled as no grammar)
# But skip empty strings
return False
tool_calls = message.get("tool_calls", [])
if isinstance(tool_calls, list):
for tool_call in tool_calls:
if isinstance(tool_call, dict):
if tool_call.get("type") != "function":
return False
if "custom" in tool_call:
return False
return True
@@ -108,7 +118,6 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
@schema.parametrize()
@schema.override(headers={"Content-Type": "application/json"})
@settings(
deadline=LONG_TIMEOUT_SECONDS * 1000,
max_examples=50,
@@ -122,7 +131,7 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
# generating large-but-valid request bodies before vLLM is called.
suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.data_too_large],
)
def test_openapi_stateless(case: Case):
def test_openapi_stateless(case: schemathesis.Case):
key = (
case.operation.method.upper(),
case.operation.path,
@@ -151,4 +160,8 @@ def test_openapi_stateless(case: Case):
}.get(key, DEFAULT_TIMEOUT_SECONDS)
# No need to verify SSL certificate for localhost
case.call_and_validate(verify=False, timeout=timeout)
case.call_and_validate(
verify=False,
timeout=timeout,
headers={"Content-Type": "application/json"},
)
@@ -25,7 +25,7 @@ def server():
"--runner",
"pooling",
"--max-model-len",
"5000",
"16384",
"--enforce-eager",
"--limit-mm-per-prompt",
json.dumps({"video": MAXIMUM_VIDEOS}),
@@ -143,4 +143,4 @@ def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str):
assert output.model == model_name
assert len(output.data) == 1
assert len(output.data[0].probs) == 2
assert output.usage.prompt_tokens == 4807
assert output.usage.prompt_tokens == 8993
@@ -8,6 +8,7 @@ import pytest
import pytest_asyncio
from tests.utils import RemoteLaunchRenderServer
from vllm.tokenizers import get_tokenizer
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
@@ -486,3 +487,438 @@ async def test_derender_completion_kv_transfer_params_passthrough(client):
)
assert response.status_code == 200
assert response.json()["kv_transfer_params"] == kv
# ---------------------------------------------------------------------------
# E2E: render -> derender roundtrip with parser (reasoning + tool calls)
# ---------------------------------------------------------------------------
PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
_E2E_TOOLS = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a city",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
},
},
}
]
@pytest.fixture(scope="module")
def parser_server():
args = [
"--enable-auto-tool-choice",
"--tool-call-parser",
"hermes",
"--reasoning-parser",
"deepseek_r1",
]
with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def parser_client(parser_server):
async with httpx.AsyncClient(
base_url=parser_server.url_for(""), timeout=60.0
) as http_client:
yield http_client
@pytest.fixture(scope="module")
def parser_tokenizer():
return get_tokenizer(PARSER_MODEL)
def _encode(tokenizer, text: str) -> list[int]:
return tokenizer.encode(text, add_special_tokens=False)
def _decoded(tokenizer, token_ids: list[int]) -> str:
return tokenizer.decode(token_ids, skip_special_tokens=True)
def _require_markers_survive(tokenizer, text: str, *markers: str) -> list[int]:
"""Encode text and skip the test if any marker is lost in roundtrip."""
ids = _encode(tokenizer, text)
decoded = tokenizer.decode(ids, skip_special_tokens=False)
for m in markers:
if m not in decoded:
pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip")
return ids
async def _e2e_render_chat(
client: httpx.AsyncClient,
model: str,
messages: list[dict],
) -> dict:
resp = await client.post(
"/v1/chat/completions/render",
json={"model": model, "messages": messages},
)
assert resp.status_code == 200, resp.text
return resp.json()
def _e2e_generate_response(
token_ids: list[int],
request_id: str = "chatcmpl-e2e-test",
) -> dict:
return {
"request_id": request_id,
"choices": [
{
"index": 0,
"token_ids": token_ids,
"finish_reason": "stop",
}
],
}
@pytest.mark.asyncio
async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer):
"""Plain text without reasoning markers roundtrips correctly."""
messages = [{"role": "user", "content": "What is 2+2?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
answer = "The answer is four."
output_ids = _encode(parser_tokenizer, answer)
expected = _decoded(parser_tokenizer, output_ids)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
},
)
assert resp.status_code == 200, resp.text
content = resp.json()["choices"][0]["message"]["content"]
assert content == expected
@pytest.mark.asyncio
async def test_e2e_token_identity(parser_client, parser_tokenizer):
"""encode(derender(token_ids)) == token_ids (RL invariant)."""
messages = [{"role": "user", "content": "Hi"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
answer = "Hello! How can I help?"
output_ids = _encode(parser_tokenizer, answer)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
},
)
assert resp.status_code == 200
content = resp.json()["choices"][0]["message"]["content"]
re_encoded = _encode(parser_tokenizer, content)
assert output_ids == re_encoded
@pytest.mark.asyncio
async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer):
"""CJK + emoji roundtrip without U+FFFD."""
messages = [{"role": "user", "content": "Reply in Chinese"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
answer = "你好世界 😀"
output_ids = _encode(parser_tokenizer, answer)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
},
)
assert resp.status_code == 200
content = resp.json()["choices"][0]["message"]["content"]
assert "" not in content
@pytest.mark.asyncio
async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer):
"""<think>...</think> splits into reasoning + content."""
messages = [{"role": "user", "content": "What is 2+3?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
reasoning_text = "The user wants 2 plus 3. That is 5."
answer_text = "The answer is 5."
output_text = f"<think>{reasoning_text}</think>{answer_text}"
output_ids = _require_markers_survive(parser_tokenizer, output_text, "</think>")
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": PARSER_MODEL,
"messages": messages,
"include_reasoning": True,
},
},
)
assert resp.status_code == 200, resp.text
msg = resp.json()["choices"][0]["message"]
assert msg["reasoning"] is not None
assert reasoning_text in msg["reasoning"]
assert answer_text in msg["content"]
assert "<think>" not in msg["content"]
@pytest.mark.asyncio
async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer):
"""<tool_call> extracted into tool_calls field."""
messages = [{"role": "user", "content": "Weather in Paris?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
output_text = (
"<think>Let me check the weather.</think>"
'<tool_call>\n{"name": "get_weather", '
'"arguments": {"city": "Paris"}}\n</tool_call>'
)
output_ids = _require_markers_survive(
parser_tokenizer,
output_text,
"</think>",
"<tool_call>",
"</tool_call>",
)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": PARSER_MODEL,
"messages": messages,
"tools": _E2E_TOOLS,
"tool_choice": "auto",
},
},
)
assert resp.status_code == 200, resp.text
choice = resp.json()["choices"][0]
assert choice["message"]["tool_calls"]
assert choice["message"]["tool_calls"][0]["function"]["name"] == "get_weather"
@pytest.mark.asyncio
async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenizer):
"""Reasoning + tool call in the same output."""
messages = [{"role": "user", "content": "Weather in Paris?"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
reasoning_text = "I should look up the weather."
tool_text = (
'<tool_call>\n{"name": "get_weather", '
'"arguments": {"city": "Paris"}}\n</tool_call>'
)
output_text = f"<think>{reasoning_text}</think>{tool_text}"
output_ids = _require_markers_survive(
parser_tokenizer, output_text, "</think>", "<tool_call>"
)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": PARSER_MODEL,
"messages": messages,
"tools": _E2E_TOOLS,
"tool_choice": "auto",
"include_reasoning": True,
},
},
)
assert resp.status_code == 200, resp.text
choice = resp.json()["choices"][0]
assert choice["message"]["reasoning"] is not None
assert reasoning_text in choice["message"]["reasoning"]
assert choice["message"]["tool_calls"]
@pytest.mark.asyncio
async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer):
"""Without chat_request, derender falls back to plain detokenization."""
messages = [{"role": "user", "content": "Hello"}]
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
answer = "Hi there!"
output_ids = _encode(parser_tokenizer, answer)
resp = await parser_client.post(
"/v1/chat/completions/derender",
json={
"model": PARSER_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
},
)
assert resp.status_code == 200
content = resp.json()["choices"][0]["message"]["content"]
assert "Hi" in content
# ---------------------------------------------------------------------------
# E2E: HarmonyParser + GPT-OSS
# ---------------------------------------------------------------------------
HARMONY_MODEL = "openai/gpt-oss-20b"
def _ensure_harmony_vocab():
"""Pre-cache the o200k_base BPE file needed by openai-harmony.
The Rust tiktoken-rs backend downloads from Azure Blob Storage, which
may be unreachable in some environments. When the cache is cold we
fetch the file ourselves and place it in ``/tmp/tiktoken-rs-cache/``
using the SHA-1(URL) filename that tiktoken-rs expects.
"""
import hashlib
import urllib.request
from pathlib import Path
url = "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
cache_dir = Path("/tmp/tiktoken-rs-cache")
cache_key = hashlib.sha1(url.encode()).hexdigest()
cache_file = cache_dir / cache_key
if not cache_file.exists():
cache_dir.mkdir(parents=True, exist_ok=True)
urllib.request.urlretrieve(url, cache_file)
@pytest.fixture(scope="module")
def harmony_server():
_ensure_harmony_vocab()
args = [
"--trust-remote-code",
"--enable-auto-tool-choice",
"--tool-call-parser",
"openai",
"--reasoning-parser",
"openai_gptoss",
]
with RemoteLaunchRenderServer(HARMONY_MODEL, args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def harmony_client(harmony_server):
async with httpx.AsyncClient(
base_url=harmony_server.url_for(""), timeout=60.0
) as http_client:
yield http_client
@pytest.fixture(scope="module")
def harmony_tokenizer():
return get_tokenizer(HARMONY_MODEL, trust_remote_code=True)
def _harmony_extract_assistant_ids(
tokenizer, assistant_msg: dict, user_content: str = "test"
) -> list[int]:
"""Extract assistant token IDs via apply_chat_template diff."""
prompt = [{"role": "user", "content": user_content}]
full = prompt + [assistant_msg]
text_prompt = tokenizer.apply_chat_template(
prompt, add_generation_prompt=True, tokenize=False
)
text_full = tokenizer.apply_chat_template(
full, add_generation_prompt=False, tokenize=False
)
prompt_ids = tokenizer.encode(text_prompt)
full_ids = tokenizer.encode(text_full)
assistant_ids = list(full_ids[len(prompt_ids) :])
if not assistant_ids:
pytest.skip("Could not extract assistant tokens for Harmony")
return assistant_ids
@pytest.mark.asyncio
async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer):
"""GPT-OSS content-only roundtrip."""
messages = [{"role": "user", "content": "What is 2+2?"}]
gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
assistant_msg = {"role": "assistant", "content": "Four."}
output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
resp = await harmony_client.post(
"/v1/chat/completions/derender",
json={
"model": HARMONY_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": HARMONY_MODEL,
"messages": messages,
},
},
)
assert resp.status_code == 200, resp.text
content = resp.json()["choices"][0]["message"]["content"]
assert content is not None and len(content) > 0
assert "Four" in content
@pytest.mark.asyncio
async def test_e2e_harmony_reasoning(harmony_client, harmony_tokenizer):
"""GPT-OSS reasoning: analysis channel extracted."""
messages = [{"role": "user", "content": "Add 2 and 3."}]
gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
reasoning_text = "The user wants 2 plus 3."
answer_text = "The answer is 5."
assistant_msg = {
"role": "assistant",
"thinking": reasoning_text,
"content": answer_text,
}
output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
decoded = harmony_tokenizer.decode(output_ids)
if reasoning_text not in decoded:
pytest.skip("Harmony template did not render thinking")
resp = await harmony_client.post(
"/v1/chat/completions/derender",
json={
"model": HARMONY_MODEL,
"generate_response": _e2e_generate_response(output_ids),
"prompt_tokens": len(gen_req["token_ids"]),
"chat_request": {
"model": HARMONY_MODEL,
"messages": messages,
"include_reasoning": True,
},
},
)
assert resp.status_code == 200, resp.text
msg = resp.json()["choices"][0]["message"]
assert msg["reasoning"] is not None
assert reasoning_text in msg["reasoning"]
assert answer_text in (msg["content"] or "")
@@ -0,0 +1,65 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests that validation_exception_handler populates the `param` field
in its error response using the Pydantic error's `loc`, even when no
custom VLLMValidationError context is present.
Previously, `param` was only populated for errors carrying a custom
VLLMValidationError in their Pydantic `ctx`. Plain validation failures
(missing fields, wrong types) left `param` as None, even though the
field name was readily available from `error['loc']`.
"""
import json
from types import SimpleNamespace
import pytest
from fastapi.exceptions import RequestValidationError
from vllm.entrypoints.serve.utils.server_utils import validation_exception_handler
def _fake_request(log_error_stack: bool = False) -> SimpleNamespace:
"""Minimal stand-in for a FastAPI Request - just enough for the
handler to read req.app.state.args.log_error_stack."""
return SimpleNamespace(
app=SimpleNamespace(
state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack))
),
state=SimpleNamespace(), # no request_metadata -> hasattr(...) is False
)
class TestValidationErrorParamFallback:
"""Ensure `param` falls back to the Pydantic error's `loc` when no
custom VLLMValidationError context is present."""
@pytest.mark.parametrize(
("error_type", "msg"),
[
("missing", "Field required"),
("list_type", "Input should be a valid list"),
],
ids=["missing-field", "wrong-type"],
)
@pytest.mark.asyncio
async def test_param_falls_back_to_loc(self, error_type: str, msg: str):
errors = [{"type": error_type, "loc": ("body", "messages"), "msg": msg}]
exc = RequestValidationError(errors)
response = await validation_exception_handler(_fake_request(), exc)
body = json.loads(response.body)
assert body["error"]["param"] == "body.messages"
@pytest.mark.asyncio
async def test_param_fallback_does_not_crash_on_non_dict_error(self):
"""Schemathesis fuzzing found that errors[0] isn't always a dict.
The fallback must not crash in that case - it should just leave
param as None instead of raising."""
exc = RequestValidationError(["some unexpected non-dict error"])
response = await validation_exception_handler(_fake_request(), exc)
body = json.loads(response.body)
assert body["error"]["param"] is None
@@ -78,7 +78,16 @@ def test_gsm8k_correctness(config_filename):
"Skipping DeepSeek-V3.2 and DeepSeek-R1 on ROCm platforms "
"due to agent pool disk space issues and pod evictions."
)
if current_platform.is_rocm() and (
"Qwen3.5-35B-A3B-MXFP4-AITER-TP2" in config_filename.name
):
from vllm.platforms.rocm import on_gfx950
if not on_gfx950():
pytest.skip(
"Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. "
"The quantization scheme is not supported on non-GFX950 platforms."
)
# Parse server arguments from config (use shlex to handle quoted strings)
server_args_str = eval_config.get("server_args", "")
server_args = shlex.split(server_args_str) if server_args_str else []
@@ -15,16 +15,14 @@ from vllm.config import (
from vllm.platforms import current_platform
from vllm.platforms.cpu import CpuPlatform
# CudaPlatform and RocmPlatform import their respective compiled C extensions
# at module level, raising ModuleNotFoundError on incompatible builds.
try:
if current_platform.is_cuda():
from vllm.platforms.cuda import CudaPlatform
except (ImportError, ModuleNotFoundError):
else:
CudaPlatform = None
try:
if current_platform.is_rocm():
from vllm.platforms.rocm import RocmPlatform
except (ImportError, ModuleNotFoundError):
else:
RocmPlatform = None
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -434,9 +432,15 @@ def test_per_head_quant_scales_backend_selection(
[
("FLASH_ATTN", True, True), # FlashAttn supports non-causal
("FLASH_ATTN", False, True), # FlashAttn also works with causal
("FLASHINFER", True, False), # FlashInfer does not support non-causal
("FLASHINFER", False, True), # FlashInfer works with causal
],
]
+ (
[
("FLASHINFER", True, False), # FlashInfer does not support non-causal
("FLASHINFER", False, True), # FlashInfer works with causal
]
if CudaPlatform is not None
else []
),
)
def test_non_causal_backend_selection(
backend_name: str, use_non_causal: bool, should_succeed: bool
@@ -459,11 +463,12 @@ def test_non_causal_backend_selection(
attention_config=attention_config, cache_config=cache_config
)
if CudaPlatform is None:
pytest.skip("CudaPlatform not available")
platform = CudaPlatform or RocmPlatform
if platform is None:
pytest.skip("CudaPlatform and RocmPlatform are not available")
with (
set_current_vllm_config(vllm_config),
patch("vllm.platforms.current_platform", CudaPlatform()),
patch("vllm.platforms.current_platform", platform()),
):
if should_succeed:
backend = get_attn_backend(
+37
View File
@@ -428,6 +428,43 @@ def test_reshape_and_cache_flash(
torch.testing.assert_close(value_cache_compact, cloned_value_cache)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@torch.inference_mode()
def test_reshape_and_cache_flash_unaligned_rows(
kv_cache_factory_flashinfer,
dtype: torch.dtype,
kv_cache_dtype: str,
kv_cache_layout: str,
implementation: str,
) -> None:
"""Regression test for https://github.com/vllm-project/vllm/issues/41257.
head_size=46 with num_heads=13 places KV-cache rows at byte offsets
that are not a multiple of the vector width (NHD row pitch
13*46*itemsize, HND head pitch 46*itemsize), unlike HEAD_SIZES above
which are all 16-byte multiples. The CUDA kernel used to issue
vectorized stores to those rows -> CUDA misaligned address.
"""
test_reshape_and_cache_flash(
kv_cache_factory_flashinfer,
num_tokens=42,
num_heads=13,
head_size=46,
block_size=16,
num_blocks=128,
dtype=dtype,
seed=0,
device=CUDA_DEVICES[0],
kv_cache_dtype=kv_cache_dtype,
kv_cache_layout=kv_cache_layout,
kv_scale_type="tensor",
implementation=implementation,
)
@pytest.mark.parametrize("direction", COPYING_DIRECTION)
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
+17 -9
View File
@@ -5,10 +5,12 @@ import math
import random
import time
from collections.abc import Callable
from contextlib import nullcontext
import pytest
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from vllm.platforms import current_platform
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
@@ -557,15 +559,21 @@ def test_contexted_kv_attention_alibi(
query_len, seq_len, alibi_slopes, device, dtype
)
# Compute attention
out = F.scaled_dot_product_attention(
q_sdpa,
k_sdpa,
v_sdpa,
attn_mask=alibi_mask,
dropout_p=0.0,
scale=scale,
)
# Compute attention. On ROCm we force use of the Math SDPA backend rather than
# the Flash or Mem-Efficient backends for increased numerical accuracy
if current_platform.is_rocm():
sdpa_context = sdpa_kernel(SDPBackend.MATH)
else:
sdpa_context = nullcontext()
with sdpa_context:
out = F.scaled_dot_product_attention(
q_sdpa,
k_sdpa,
v_sdpa,
attn_mask=alibi_mask,
dropout_p=0.0,
scale=scale,
)
# Reshape output back to [query_len, num_heads, head_size]
out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)
@@ -90,7 +90,9 @@ def _ref_sparse_prefill_ragged(
return out.to(torch.bfloat16)
def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
def _pack_fp8_ds_mla_cache(
kv: torch.Tensor, block_size: int, is_extra: bool = False
) -> torch.Tensor:
assert kv.shape[-1] == HEAD_DIM
num_tokens = kv.shape[0]
num_blocks = (num_tokens + block_size - 1) // block_size
@@ -101,7 +103,9 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
)
cache_flat = cache.view(torch.uint8).flatten()
kv_nope_fp8 = (
kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
kv[:, :NOPE_HEAD_DIM]
.to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
.view(torch.uint8)
)
kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
@@ -120,7 +124,7 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
def _read_fp8_ds_mla_cache(
cache: torch.Tensor, slot: int, block_size: int
cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
) -> torch.Tensor:
cache_flat = cache.view(torch.uint8).flatten()
block_idx = slot // block_size
@@ -129,7 +133,9 @@ def _read_fp8_ds_mla_cache(
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)
nope = nope_u8.view(
torch.float8_e4m3fn if is_extra else 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
]
@@ -157,7 +163,9 @@ def _ref_sparse_decode_ragged(
]
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)
_read_fp8_ds_mla_cache(
extra_cache, int(slot), block_size, is_extra=True
)
for slot in extra_rows[query_idx]
)
@@ -326,7 +334,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
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)
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
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)
@@ -477,7 +485,7 @@ def test_sparse_attn_decode_split_k_kernel(
rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
extra_rows = rows
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
extra_indices, extra_indptr = _ragged_from_rows(rows, device)
attn_sink = (
@@ -18,11 +18,7 @@ HEAD_SIZES = [128, 256]
BLOCK_SIZES = [16]
DTYPES = [torch.bfloat16]
QDTYPES = (
[None, torch.float8_e4m3fn]
if not current_platform.is_rocm()
else [None, torch.float8_e4m3fnuz]
)
QDTYPES = [None, current_platform.fp8_dtype()]
FP8_DTYPE = current_platform.fp8_dtype()
# one value large enough to test overflow in index calculation.
+59 -3
View File
@@ -10,8 +10,12 @@ from torch.multiprocessing import spawn
from tests.kernels.utils import opcheck
from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
from vllm.model_executor.layers.minimax_rms_norm import (
MiniMaxText01RMSNormTP,
rms_norm_tp,
)
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON
from vllm.utils.network_utils import get_open_port
from vllm.utils.torch_utils import set_random_seed
@@ -54,8 +58,19 @@ def _worker_forward_qk(
torch.manual_seed(seed + 1000 + local_rank)
qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1)
ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref)
# Reference: eager all-reduce path. ``forward_qk`` no longer all-reduces
# the variance (it is the tp==1 / already-reduced building block), so the
# multi-rank reference must use the eager path that performs the global
# variance all-reduce, matching the fused kernel below.
ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager(
qkv.clone(),
q_norm.weight,
k_norm.weight,
hq,
hk,
world_size,
eps,
)
# Set up Lamport workspace.
from vllm.distributed.parallel_state import get_tp_group
@@ -150,3 +165,44 @@ def test_minimax_reduce_rms_qk(
nprocs=world_size,
join=True,
)
@pytest.mark.skipif(
not current_platform.is_cuda() or not HAS_TRITON,
reason="CUDA and Triton required",
)
@pytest.mark.parametrize("num_tokens", [1, 7, 128, 333, 2049])
@pytest.mark.parametrize("hidden_dims", [(3072, 512), (768, 256), (3000, 500)])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("tp_world", [1, 4, 8])
@pytest.mark.parametrize("eps", [1e-6])
@pytest.mark.parametrize("seed", [42])
def test_minimax_qk_norm_triton_fallback(
monkeypatch, num_tokens, hidden_dims, dtype, tp_world, eps, seed
):
"""Single-GPU check: Triton fallback kernels vs the pure-torch reference.
The all-reduce is a TP communication barrier, so it is monkeypatched to
identity here; both the Triton path and the reference see the same
(patched) reduction. This validates the kernel math and the folded
``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
are the per-rank q/k segment widths.
"""
monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
q_size, kv_size = hidden_dims
device = "cuda"
torch.manual_seed(seed)
qkv = torch.randn(num_tokens, q_size + 2 * kv_size, dtype=dtype, device=device)
q_weight = torch.randn(q_size, dtype=dtype, device=device)
k_weight = torch.randn(kv_size, dtype=dtype, device=device)
q_triton, k_triton = rms_norm_tp._minimax_qk_norm_tp_fallback(
qkv, q_weight, k_weight, q_size, kv_size, 0, tp_world, eps
)
q_ref, k_ref = rms_norm_tp._minimax_qk_norm_tp_eager(
qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
)
torch.testing.assert_close(q_triton, q_ref, atol=3e-2, rtol=3e-2)
torch.testing.assert_close(k_triton, k_ref, atol=3e-2, rtol=3e-2)
+8 -8
View File
@@ -9,6 +9,7 @@ import pytest
import torch
from packaging import version
from vllm._aiter_ops import is_aiter_found
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer
@@ -31,17 +32,15 @@ HOPPER_MXFP4_BF16_AVAILABLE = (
# ROCm platform and dependencies
ROCM_AVAILABLE = current_platform.is_rocm()
ROCM_TRITON_KERNELS_AVAILABLE = False
ROCM_AITER_AVAILABLE = False
ROCM_AITER_AVAILABLE = is_aiter_found()
ROCM_GFX950 = False
if ROCM_AVAILABLE:
from vllm._aiter_ops import rocm_aiter_ops
from vllm.platforms.rocm import on_gfx950
from vllm.utils.import_utils import has_triton_kernels
ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels()
ROCM_GFX950 = on_gfx950()
ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_enabled()
if ROCM_AITER_AVAILABLE:
from aiter.ops.triton.moe.quant_moe import upcast_from_mxfp
@@ -83,7 +82,7 @@ def enable_pickle(monkeypatch):
[
ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=2),
ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8),
ModelCase("fxmarty/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
ModelCase("mawong-amd/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=1),
ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=4),
],
@@ -102,6 +101,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
tensor_parallel_size=model_case.tp,
load_format="dummy",
compilation_config={"cudagraph_capture_sizes": [16]},
gpu_memory_utilization=0.8, # mxfp6 models use more scratch space
) as llm:
# Disabled as check_model is broken: https://github.com/vllm-project/vllm/pull/18465#issuecomment-3329880562
# def check_model(model):
@@ -1267,7 +1267,7 @@ def test_rocm_mxfp4_moe_oracle(
This test validates that the oracle functions work end-to-end:
- select_mxfp4_moe_backend() selects a valid backend
- convert_to_mxfp4_moe_kernel_format() converts weights without error
- convert_gpt_oss_weight_to_mxfp4_moe_kernel_format() converts weights without error
- make_mxfp4_moe_quant_config() builds a valid quant config
- make_mxfp4_moe_kernel() creates a kernel that runs without error
- The kernel output is within accuracy tolerance of reference
@@ -1287,7 +1287,7 @@ def test_rocm_mxfp4_moe_oracle(
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
Mxfp4MoeBackend,
backend_to_kernel_cls,
convert_to_mxfp4_moe_kernel_format,
convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
make_mxfp4_moe_kernel,
make_mxfp4_moe_quant_config,
)
@@ -1387,7 +1387,7 @@ def test_rocm_mxfp4_moe_oracle(
# Convert weights using oracle
w13_conv, w2_conv, w13_scale_conv, w2_scale_conv, w13_bias_conv, w2_bias_conv = (
convert_to_mxfp4_moe_kernel_format(
convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
mxfp4_backend=backend,
layer=layer, # type: ignore[arg-type]
w13_weight=w13_quant,
@@ -1423,7 +1423,7 @@ def test_rocm_mxfp4_moe_oracle(
mxfp4_backend=backend,
experts_cls=experts_cls,
routing_tables=None,
shared_experts=None,
layer=None,
)
# Create inputs
@@ -345,6 +345,63 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
)
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU",
)
def test_per_token_group_quant_fp8_packed_large_mn():
"""Regression test for https://github.com/vllm-project/vllm/issues/45099.
Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and
2^16 - 1, respectively.
Prior code introduced a bug where it incorrectly assumed grid.x and y both have
2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is
launched with a large mn that exceeds grid.y limit (2^16 - 1).
This issue doesn't surface often because each forward pass only processes a
bounded token batch, not the full context.
Quantizing tensors with more rows than that will fail at launch with
"CUDA error: invalid argument".
This is a differential test that compares fp8 output against Triton output
reference when token size sits just above the gridDim.y 2^16 - 1 limit.
"""
device = "cuda"
group_size = 128
# hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
# row, so any mn > 65535 overflowed grid.y before the fix.
num_tokens, hidden_dim = 65537, 2048
torch.manual_seed(42)
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
x,
group_size=group_size,
use_ue8m0=True,
)
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
x, group_size, use_ue8m0=True
)
assert torch.equal(out_q, ref_q), "Quantized output mismatch"
# Vectorized packed-scale check; the per-element loop used by the smaller
# tests is too slow at this size. groups_per_row is a multiple of 4 here,
# so there is no K padding and the packed view lines up.
mn = num_tokens
groups_per_row = hidden_dim // group_size
k_num_packed = (groups_per_row + 3) // 4
assert groups_per_row % 4 == 0
ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF
exp = ref_exponents.view(mn, k_num_packed, 4)
expected = (
exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24)
)
assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@@ -60,8 +60,10 @@ def test_rocm_compressed_tensors_w8a8(
vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
):
dtype = "bfloat16"
with vllm_runner(model_path, dtype=dtype) as vllm_model:
# Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179
with vllm_runner(
model_path, dtype=dtype, attention_backend="TRITON_ATTN"
) as vllm_model:
vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)
+25 -11
View File
@@ -70,17 +70,20 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason=(
"Mxfp4 LoRA on ROCm is blocked by a spawn compatibility issue. "
"The fused_moe_lora Triton kernel crashes in spawned subprocesses, "
"and vLLM forces spawn mode when HIP is initialized before "
"multiprocessing. Fixing this requires either making the LoRA "
"Triton kernel spawn-safe or pre-warming the kernel cache."
),
# TODO: make the Mxfp4MoeBackend.TRITON spawn-safe.
# For now just use TRITON_UNFUSED kernel
@pytest.mark.parametrize(
"mxfp4_use_marlin",
[
False,
pytest.param(
True,
marks=pytest.mark.skipif(
current_platform.is_rocm(), reason="marlin not supported"
),
),
],
)
@pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
@pytest.mark.parametrize("specialize_active_lora", [True, False])
def test_gpt_oss_lora(
gptoss20b_lora_files,
@@ -109,7 +112,18 @@ def test_gpt_oss_lora(
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize("fully_sharded_loras", [False, True])
@pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
@pytest.mark.parametrize(
"mxfp4_use_marlin",
[
False,
pytest.param(
True,
marks=pytest.mark.skipif(
current_platform.is_rocm(), reason="marlin not supported"
),
),
],
)
def test_gpt_oss_lora_tp2(
gptoss20b_lora_files,
fully_sharded_loras,
+124
View File
@@ -482,3 +482,127 @@ def test_kernels_hidden_size(
seq_length=128,
add_inputs=True,
)
@pytest.mark.parametrize("device", DEVICES)
def test_add_lora_fused_moe_early_exit(device):
"""
Ensures add_lora_fused_moe does not invoke the LoRA kernel or
modify the output tensor when no_lora_flag_cpu is True
"""
from types import SimpleNamespace
from vllm.lora.punica_wrapper.punica_gpu import PunicaWrapperGPU
torch.set_default_device(device)
torch.accelerator.set_device_index(device)
max_loras, num_tokens = 4, 16
num_experts, top_k, max_lora_rank = 8, 2, 16
K, N = 256, 128
# build PunicaWrapperGPU with minimal lora_config mock
lora_config = SimpleNamespace(
max_loras=max_loras,
specialize_active_lora=False,
)
wrapper = PunicaWrapperGPU(
max_num_batched_tokens=num_tokens,
max_batches=num_tokens,
device=device,
lora_config=lora_config,
)
# simulate a prior LoRA batch so the internal mapping is
# populated with stale LoRA IDs
lora_mapping = torch.zeros(
num_tokens,
dtype=torch.int32,
device=device,
)
lora_mapping[:8] = 1
lora_mapping[8:] = 2
wrapper.token_mapping_meta.prepare_tensors(lora_mapping)
# simulate a base-model batch (all -1)
base_mapping = torch.full(
(num_tokens,),
-1,
dtype=torch.int32,
device=device,
)
wrapper.token_mapping_meta.prepare_tensors(base_mapping)
assert wrapper.token_mapping_meta.no_lora_flag_cpu[0].item() is True
# dummy tensors for add_lora_fused_moe
y = torch.rand(num_tokens, top_k, N, dtype=torch.bfloat16, device=device)
y_snapshot = y.clone()
x = torch.rand(num_tokens, K, dtype=torch.bfloat16, device=device)
lora_a_stacked = (
torch.rand(
max_loras,
num_experts,
max_lora_rank,
K,
dtype=torch.bfloat16,
device=device,
),
)
lora_b_stacked = (
torch.rand(
max_loras,
num_experts,
N,
max_lora_rank,
dtype=torch.bfloat16,
device=device,
),
)
topk_weights = torch.ones(
num_tokens,
top_k,
dtype=torch.float32,
device=device,
)
adapter_enabled = torch.ones(
max_loras + 1,
dtype=torch.int32,
device=device,
)
shrink_config = expand_config = {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"NUM_WARPS": 4,
"NUM_STAGES": 3,
"SPLIT_K": 1,
}
# call add_lora_fused_moe - the early exit should prevent any
# modification to the output
wrapper.add_lora_fused_moe(
y=y,
x=x,
lora_a_stacked=lora_a_stacked,
lora_b_stacked=lora_b_stacked,
topk_weights=topk_weights,
sorted_token_ids=None,
expert_ids=torch.zeros(
num_tokens * top_k,
dtype=torch.int32,
device=device,
),
num_tokens_post_padded=None,
max_lora_rank=max_lora_rank,
top_k_num=top_k,
shrink_config=shrink_config,
expand_config=expand_config,
adapter_enabled=adapter_enabled,
)
assert torch.equal(y, y_snapshot), (
"add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True"
)
@@ -130,8 +130,12 @@ def test_models(
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
if model == "TitanML/tiny-mixtral":
# Untrained model: near-uniform logits make argmax sensitive to
# AITER's bfloat16 rounding error in plain rms_norm.
# AITER's bfloat16 rounding error. Route the plain rms_norm and the
# fused MoE (whose near-uniform router logits flip expert selection
# under ~1 ULP drift) through the native kernels for this model.
# See ROCm/aiter#3806 for the tracking issue and minimal repro.
monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
elif use_rocm_aiter and model not in AITER_MODEL_LIST:
# Skip model that are not using AITER tests.
# When more AITER kernels are added, this list will not be
@@ -25,7 +25,7 @@ TEST_IMAGE_NAMES = [
]
MAX_MODEL_LEN = 8192
REQUESTS_PER_ROUND = 4
WARMUP_ROUNDS = 1
WARMUP_ROUNDS = 2
MEASURED_ROUNDS = 16
GPU_GROWTH_THRESHOLD_MIB = 0
CPU_PEAK_GROWTH_THRESHOLD_MIB = 0
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
import transformers.utils
from PIL import Image
@@ -52,6 +53,7 @@ def _get_cherry_blossom_image() -> Image.Image:
)
@torch.inference_mode()
def _run_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
@@ -92,3 +92,95 @@ def test_processor_num_frames_timestamp(
assert len(video_phs) == 1, (
f"Expected exactly 1 video placeholder, got {len(video_phs)}"
)
@pytest.mark.parametrize("model_id", [MODEL_ID])
@pytest.mark.parametrize("num_videos", [2, 4])
def test_processor_multi_video(
model_id: str,
num_videos: int,
) -> None:
"""Verify that multi-video processing produces correct placeholders.
This exercises the token-level replacement path in
``_call_hf_processor`` which avoids the quadratic text-level
prompt expansion.
"""
ctx = build_model_context(
model_id,
limit_mm_per_prompt={"image": 0, "video": num_videos},
)
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
prompt = "<|vision_start|><|video_pad|><|vision_end|>" * num_videos
mm_data = {"video": [_build_video_mm_data(num_frames=8)["video"][0]] * num_videos}
processed = processor(
prompt,
mm_items=processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs={"num_frames": 8},
)
token_ids = processed["prompt_token_ids"]
assert len(token_ids) > 0
video_phs = processed["mm_placeholders"].get("video", [])
assert len(video_phs) == num_videos, (
f"Expected {num_videos} video placeholders, got {len(video_phs)}"
)
# All placeholders should have the same length (same video params)
# and must not overlap.
lengths = {ph.length for ph in video_phs}
assert len(lengths) == 1, f"Placeholder lengths differ: {lengths}"
for i in range(1, len(video_phs)):
prev_end = video_phs[i - 1].offset + video_phs[i - 1].length
assert video_phs[i].offset >= prev_end, (
f"Placeholder {i} overlaps with placeholder {i - 1}"
)
@pytest.mark.parametrize("model_id", [MODEL_ID])
@pytest.mark.parametrize(
"hf_mm_kwargs",
[{"num_frames": [8, 16]}, {"fps": [2.0, 4.0]}],
)
def test_processor_multi_video_list_kwargs(
model_id: str,
hf_mm_kwargs: dict[str, Any],
) -> None:
"""Regression test: a multi-video request with list-valued per-video
``mm_processor_kwargs`` (one ``fps``/``num_frames`` per video) must not
crash.
Before the fix, ``_call_hf_processor`` copied the whole kwargs to every
video without slicing, so ``_get_video_second_idx`` received the list
where a scalar was expected and raised ``TypeError``.
"""
ctx = build_model_context(
model_id,
limit_mm_per_prompt={"image": 0, "video": 2},
)
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
prompt = (
"<|vision_start|><|video_pad|><|vision_end|>"
"<|vision_start|><|video_pad|><|vision_end|>"
)
mm_data = {
"video": [
_build_video_mm_data(num_frames=16)["video"][0],
_build_video_mm_data(num_frames=32)["video"][0],
]
}
processed = processor(
prompt,
mm_items=processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs=hf_mm_kwargs,
)
video_phs = processed["mm_placeholders"].get("video", [])
assert len(video_phs) == 2, (
f"Expected exactly 2 video placeholders, got {len(video_phs)}"
)
+19
View File
@@ -1530,6 +1530,16 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
"Qwen/Qwen3-VL-8B-Instruct",
speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
),
"Eagle3Qwen3ForCausalLM": _HfExamplesInfo(
"Qwen/Qwen3-8B",
trust_remote_code=True,
speculative_model=(
"inference-optimization/"
"Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1"
),
tokenizer="Qwen/Qwen3-8B",
use_original_num_layers=True,
),
# [PEagle]
"PEagleDraftModel": _HfExamplesInfo(
"Qwen/Qwen3-8B",
@@ -1545,6 +1555,15 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
tokenizer="Qwen/Qwen3-8B",
use_original_num_layers=True,
),
"PeagleQwen3ForCausalLM": _HfExamplesInfo(
"Qwen/Qwen3-8B",
trust_remote_code=True,
speculative_model=(
"inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4"
),
tokenizer="Qwen/Qwen3-8B",
use_original_num_layers=True,
),
# [MTP]
"DeepSeekMTPModel": _HfExamplesInfo(
"luccafong/deepseek_mtp_main_random",
+25 -2
View File
@@ -15,6 +15,7 @@ from vllm.multimodal.video import (
DynamicVideoBackend,
GLM46VVideoBackend,
Molmo2VideoBackend,
Qwen2VLVideoBackend,
Qwen3VLVideoBackend,
VideoLoader,
VideoSourceMetadata,
@@ -70,11 +71,12 @@ def test_video_loader_type_doesnt_exist():
@pytest.mark.parametrize(
"model_repo, expected_loader_cls",
"model_repo, expected_loader_cls, hf_sample_kwargs",
[
pytest.param(
"allenai/Molmo2-4B",
Molmo2VideoBackend,
None,
marks=pytest.mark.skip(
reason="Video processor not aligned, investigate later.",
),
@@ -83,23 +85,44 @@ def test_video_loader_type_doesnt_exist():
pytest.param(
"zai-org/GLM-4.1V-9B-Thinking",
DynamicVideoBackend,
None,
id="glm4v",
),
pytest.param(
"zai-org/GLM-4.6V-Flash",
GLM46VVideoBackend,
None,
id="glm46v",
),
pytest.param(
"Qwen/Qwen3-VL-4B-Instruct",
Qwen3VLVideoBackend,
None,
id="qwen3vl",
),
# Qwen2-VL/Qwen2.5-VL ship no ``video_processor_type`` in their
# preprocessor config, so resolution relies on the model_type ->
# video processor fallback in get_video_processor_cls_name_from_config.
# They also ship no default fps/num_frames, so the HF sampler needs an
# explicit target rate; pass fps=2 to match the loader default.
pytest.param(
"Qwen/Qwen2-VL-7B-Instruct",
Qwen2VLVideoBackend,
{"fps": 2},
id="qwen2vl",
),
pytest.param(
"Qwen/Qwen2.5-VL-7B-Instruct",
Qwen2VLVideoBackend,
{"fps": 2},
id="qwen2_5_vl",
),
],
)
def test_video_processor_from_model_repo(
model_repo: str,
expected_loader_cls: type,
hf_sample_kwargs: dict[str, int | float] | None,
):
"""Test that a model repo resolves to the correct video loader backend.
@@ -143,7 +166,7 @@ def test_video_processor_from_model_repo(
fps=vllm_meta["fps"],
duration=vllm_meta["duration"],
)
hf_indices = processor.sample_frames(hf_metadata)
hf_indices = processor.sample_frames(hf_metadata, **(hf_sample_kwargs or {}))
vllm_indices = np.array(vllm_meta["frames_indices"])
np.testing.assert_array_equal(
hf_indices,
+3
View File
@@ -96,6 +96,9 @@ class MockTokenizer:
return "".join(parts)
CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
def make_mock_tokenizer(sample: Sample) -> MockTokenizer:
"""Build a mock tokenizer from a sample's vocab and token data."""
return MockTokenizer(
@@ -19,6 +19,7 @@ import pytest
from pydantic import TypeAdapter
from tests.parser.engine.replay_harness import (
CHUNK_SIZES,
MockTokenizer,
assert_parse_output,
collect_output,
@@ -113,8 +114,6 @@ _PAIRINGS = _discover_pairings()
_ALL_SAMPLES = [(p.parser_cls, s) for p in _PAIRINGS for s in p.samples]
CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
@pytest.mark.parametrize(
+18
View File
@@ -615,6 +615,24 @@ class TestArgConverter:
assert result["command"] == "ls -la"
assert result["desc"] == "\npartial value"
def test_partial_value_with_angle_bracket(self):
from vllm.parser.qwen3 import (
_qwen3_arg_converter,
)
raw = "<parameter=expr>x<5"
result = json.loads(_qwen3_arg_converter(raw, partial=True))
assert result == {"expr": "x<5"}
def test_partial_value_with_angle_bracket_and_complete_param(self):
from vllm.parser.qwen3 import (
_qwen3_arg_converter,
)
raw = "<parameter=city>Tokyo</parameter>\n<parameter=expr>x<5"
result = json.loads(_qwen3_arg_converter(raw, partial=True))
assert result == {"city": "Tokyo", "expr": "x<5"}
class TestSchemaAwareTypeCoercion:
"""Verify that _fix_arg_types corrects miscoerced values using the
@@ -0,0 +1,181 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression test for U+FFFD leak at reasoning→content transition.
When byte-fallback tokens span the reasoning/content boundary,
decoding isolated content-side token IDs via tokenizer.decode()
produces U+FFFD (Unicode replacement character). The fix flushes
the reasoning parser's engine lexer instead.
Reproduces the bug at various chunk sizes and validates that the
fix prevents U+FFFD from leaking into streamed content.
"""
from __future__ import annotations
import pytest
from tests.parser.engine.replay_harness import (
CHUNK_SIZES,
MockTokenizer,
collect_output,
replay_streaming,
)
from vllm.parser.abstract_parser import DelegatingParser
from vllm.parser.engine.registered_adapters import (
Glm47MoeParserReasoningAdapter,
Glm47MoeParserToolAdapter,
Qwen3ParserReasoningAdapter,
Qwen3ParserToolAdapter,
)
class ByteFallbackMockTokenizer(MockTokenizer):
"""MockTokenizer that returns U+FFFD for specified token IDs.
Simulates byte-fallback tokenizer behavior where isolated
partial-byte tokens decode to the Unicode replacement character.
"""
def __init__(
self,
vocab: dict[str, int],
tokens: list[tuple[int, str]],
ufffd_token_ids: set[int],
) -> None:
super().__init__(vocab, tokens)
self._ufffd_token_ids = frozenset(ufffd_token_ids)
def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
parts: list[str] = []
for tid in ids:
if skip_special_tokens and tid in self._special_ids:
continue
if tid in self._ufffd_token_ids:
parts.append("")
else:
text = self._token_decode_map.get(tid, f"?{tid}?")
parts.append(text)
return "".join(parts)
# ── Model-specific DelegatingParser subclasses ───────────────────────
class _Glm47Delegating(DelegatingParser):
reasoning_parser_cls = Glm47MoeParserReasoningAdapter
tool_parser_cls = Glm47MoeParserToolAdapter
class _Qwen3Delegating(DelegatingParser):
reasoning_parser_cls = Qwen3ParserReasoningAdapter
tool_parser_cls = Qwen3ParserToolAdapter
# ── Shared test data ─────────────────────────────────────────────────
_SHARED_TOKENS: list[tuple[int, str]] = [
(100, "Let me"),
(101, " think"),
(102, " about"),
(103, " Samsung."),
(51, "</think>"),
(200, "삼성"),
(201, "전자의"),
(202, " 주가를"),
(203, " 분석합니다."),
]
_SHARED_UFFFD_IDS: set[int] = {200}
EXPECTED_REASONING = "Let me think about Samsung."
EXPECTED_CONTENT = "삼성전자의 주가를 분석합니다."
_MODEL_CONFIGS = [
pytest.param(
{
"<think>": 50,
"</think>": 51,
"<tool_call>": 60,
"</tool_call>": 61,
"<arg_key>": 62,
"</arg_key>": 63,
"<arg_value>": 64,
"</arg_value>": 65,
},
_Glm47Delegating,
id="glm47",
),
pytest.param(
{
"<think>": 50,
"</think>": 51,
"<tool_call>": 60,
"</tool_call>": 61,
},
_Qwen3Delegating,
id="qwen3",
),
]
# ── Tests ────────────────────────────────────────────────────────────
class TestUfffdReasoningTransition:
"""U+FFFD must not appear at the reasoning→content transition."""
@pytest.mark.parametrize("vocab,delegating_cls", _MODEL_CONFIGS)
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
def test_no_ufffd(self, chunk_size, vocab, delegating_cls):
tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS)
parser = delegating_cls(tokenizer)
deltas = replay_streaming(
parser,
_SHARED_TOKENS,
chunk_size=chunk_size,
finished_on_last=True,
)
output = collect_output(deltas)
assert "" not in output.content, (
f"U+FFFD leaked into content: {output.content!r}"
)
assert output.content == EXPECTED_CONTENT
assert output.reasoning == EXPECTED_REASONING
def test_byte_fallback_tokenizer_produces_ufffd(self):
"""Validate the fixture: decode() returns U+FFFD for isolated
byte-fallback token IDs, proving the old code path would leak."""
vocab = dict(_MODEL_CONFIGS[0].values[0])
tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS)
assert tokenizer.decode([200]) == ""
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
def test_multiple_ufffd_tokens_at_boundary(self, chunk_size):
"""Multiple consecutive byte-fallback tokens at the boundary."""
tokens: list[tuple[int, str]] = [
(100, "Reasoning."),
(51, "</think>"),
(200, ""),
(201, ""),
(202, "전자"),
]
ufffd_ids: set[int] = {200, 201}
vocab = dict(_MODEL_CONFIGS[0].values[0])
tokenizer = ByteFallbackMockTokenizer(vocab, tokens, ufffd_ids)
parser = _Glm47Delegating(tokenizer)
deltas = replay_streaming(
parser,
tokens,
chunk_size=chunk_size,
finished_on_last=True,
)
output = collect_output(deltas)
assert "" not in output.content, (
f"U+FFFD leaked into content: {output.content!r}"
)
assert output.content == "삼성전자"
assert output.reasoning == "Reasoning."
+93 -2
View File
@@ -30,6 +30,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
)
from vllm.parser.engine.registered_adapters import (
Gemma4Parser,
Glm47MoeParser,
MinimaxM2Parser,
NemotronV3Parser,
Qwen3Parser,
@@ -142,6 +143,12 @@ SCENARIOS: list[Scenario] = [
tool_calls=[_READ_TOOL],
after_tool_response=True,
),
Scenario(
id="empty-tool-block",
description="Empty tool block followed by content (edge case recovery)",
content="Content after empty tools.",
tool_calls=[],
),
]
@@ -343,8 +350,11 @@ def _qwen3_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
if scenario.content is not None or scenario.tool_calls is not None:
segs.append(("</think>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<tool_call>", True))
segs.append(("</tool_call>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -436,8 +446,11 @@ def _minimax_m2_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
if scenario.content is not None or scenario.tool_calls is not None:
segs.append(("</think>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<minimax:tool_call>", True))
segs.append(("</minimax:tool_call>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -533,6 +546,9 @@ def _gemma4_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs.append((_GEMMA4_THOUGHT_PREFIX, False))
segs.append((scenario.reasoning, False))
segs.append(("<channel|>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<|tool_call>", True))
segs.append(("<tool_call|>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -571,6 +587,80 @@ def _build_nemotron_v3(scenario: Scenario, validate: bool = True) -> Sample:
)
# ── GLM-4.7 MoE (XML tool format, starts in REASONING) ──────────────
_GLM47_MOE_VOCAB: dict[str, int] = {
"<think>": 50,
"</think>": 51,
"<tool_call>": 60,
"</tool_call>": 61,
"<arg_key>": 62,
"</arg_key>": 63,
"<arg_value>": 64,
"</arg_value>": 65,
}
def _glm47_moe_arg_value(value: Any) -> str:
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, (int, float)):
return str(value)
if isinstance(value, str):
return value
return json.dumps(value, ensure_ascii=False)
def _glm47_moe_tool_segments(tc: ToolCallSpec) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = [
("<tool_call>", True),
(tc.name, False),
]
for key, value in tc.arguments.items():
segs.extend(
[
("<arg_key>", True),
(key, False),
("</arg_key>", True),
("<arg_value>", True),
(_glm47_moe_arg_value(value), False),
("</arg_value>", True),
]
)
segs.append(("</tool_call>", True))
return segs
def _glm47_moe_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
segs.append(("</think>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
for tc in scenario.tool_calls:
segs.extend(_glm47_moe_tool_segments(tc))
return segs
def _build_glm47_moe(scenario: Scenario, validate: bool = True) -> Sample:
sample = _make_sample(
sample_id=f"glm47_moe-{scenario.id}",
description=scenario.description,
vocab=_GLM47_MOE_VOCAB,
segments=_glm47_moe_segments(scenario),
expected_reasoning=scenario.reasoning if scenario.reasoning is not None else "",
expected_content=_qwen3_expected_content(scenario),
expected_tool_calls=_expected_tc(scenario),
tools=_expected_tools(scenario),
)
if validate:
_validate_sample(sample, Glm47MoeParser)
return sample
# ── Registry and public API ──────────────────────────────────────────
_BUILDERS: dict[str, Any] = {
@@ -578,6 +668,7 @@ _BUILDERS: dict[str, Any] = {
"gemma4": _build_gemma4,
"minimax_m2": _build_minimax_m2,
"nemotron_v3": _build_nemotron_v3,
"glm47_moe": _build_glm47_moe,
}
+53 -7
View File
@@ -61,8 +61,8 @@ class QuantConfig:
quant_max: float
quant_min: float
kv_quant_mode: KVQuantMode
# INT8 Triton stores truncate; FP8 hardware casts round.
uses_trunc: bool
# INT8 rounds explicitly; FP8 relies on dtype cast rounding.
rounds_before_store: bool
INT8_CONFIG = QuantConfig(
@@ -71,7 +71,7 @@ INT8_CONFIG = QuantConfig(
quant_max=127.0,
quant_min=-128.0,
kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
uses_trunc=True,
rounds_before_store=True,
)
FP8_CONFIG = QuantConfig(
cache_dtype=FP8_DTYPE,
@@ -79,7 +79,7 @@ FP8_CONFIG = QuantConfig(
quant_max=FP8_MAX,
quant_min=FP8_MIN,
kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
uses_trunc=False,
rounds_before_store=False,
)
QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG]
@@ -104,7 +104,7 @@ def _quantize_per_token_head_ref(
absmax = data.float().abs().amax(dim=2) # [num_tokens, num_heads]
scales = (absmax / cfg.quant_max).clamp(min=1e-6)
scaled = data.float() * (1.0 / scales[:, :, None])
if cfg.uses_trunc:
if cfg.rounds_before_store:
q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
else:
q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
@@ -255,7 +255,7 @@ def test_per_token_head_round_trip_accuracy(
):
"""Verify per-token-head round-trip: kernel dequant matches reference.
INT8: Triton truncates on float->int8 store.
INT8: round-to-nearest before int8 store.
FP8: hardware cast (clamp then cast).
"""
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
@@ -315,6 +315,52 @@ def test_per_token_head_round_trip_accuracy(
)
@torch.inference_mode()
def test_int8_per_token_head_raw_cache_matches_round_reference():
"""INT8 cache writes should match round-to-nearest quantization exactly."""
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
triton_reshape_and_cache_flash_per_token_head_quant,
)
torch.set_default_device(DEVICE_TYPE)
head_size = 8
block_size = 4
key = torch.tensor(
[[[-127.0, -2.6, -2.4, -1.6, -1.4, -0.6, -0.4, 127.0]]],
dtype=torch.bfloat16,
)
value = -key
key_cache = torch.zeros(1, block_size, 1, head_size, dtype=torch.int8)
value_cache = torch.zeros_like(key_cache)
k_scale_cache = torch.ones(1, block_size, 1, dtype=torch.float32)
v_scale_cache = torch.ones_like(k_scale_cache)
slot_mapping = torch.tensor([2], dtype=torch.long)
triton_reshape_and_cache_flash_per_token_head_quant(
key,
value,
key_cache,
value_cache,
k_scale_cache,
v_scale_cache,
slot_mapping,
)
ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, INT8_CONFIG)
ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, INT8_CONFIG)
slot = slot_mapping.item()
blk = slot // block_size
off = slot % block_size
assert torch.equal(key_cache[blk, off], ref_k_quant[0])
assert torch.equal(value_cache[blk, off], ref_v_quant[0])
torch.testing.assert_close(k_scale_cache[blk, off], ref_k_scales[0])
torch.testing.assert_close(v_scale_cache[blk, off], ref_v_scales[0])
# ===========================================================================
# 4. Negative slot mapping (padding tokens should be skipped)
# ===========================================================================
@@ -461,7 +507,7 @@ def test_triton_unified_attention_per_token_head_scale(
scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]
scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]
if qcfg.uses_trunc:
if qcfg.rounds_before_store:
key_cache_q = (
scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)
)
@@ -11,7 +11,7 @@ parser_name = "glm45"
start_token = "<think>"
end_token = "</think>"
REASONING_MODEL_NAME = "zai-org/GLM-4.5"
REASONING_MODEL_NAME = "zai-org/GLM-4.7"
@pytest.fixture(scope="module")
@@ -35,18 +35,32 @@ WITH_THINK_STREAM = {
WITHOUT_THINK = {
"output": "This is the rest",
"reasoning": None,
"content": "This is the rest",
"reasoning": "This is the rest",
"content": None,
"is_reasoning_end": False,
}
WITHOUT_THINK_STREAM = {
"output": "This is the rest",
"reasoning": None,
"content": "This is the rest",
"reasoning": "This is the rest",
"content": None,
"is_reasoning_end": False,
}
WITHOUT_OPEN_THINK = {
"output": "This is a reasoning section</think>This is the rest",
"reasoning": "This is a reasoning section",
"content": "This is the rest",
"is_reasoning_end": True,
}
WITHOUT_OPEN_THINK_STREAM = {
"output": "This is a reasoning section</think>This is the rest",
"reasoning": "This is a reasoning section",
"content": "This is the rest",
"is_reasoning_end": True,
}
COMPLETE_REASONING = {
"output": "<think>This is a reasoning section</think>",
"reasoning": "This is a reasoning section",
@@ -61,8 +75,8 @@ MULTILINE_REASONING = {
}
ONLY_OPEN_TAG = {
"output": "<think>This is a reasoning section",
"reasoning": None,
"content": "<think>This is a reasoning section",
"reasoning": "This is a reasoning section",
"content": None,
"is_reasoning_end": False,
}
@@ -94,6 +108,16 @@ TEST_CASES = [
WITHOUT_THINK_STREAM,
id="without_think_stream",
),
pytest.param(
False,
WITHOUT_OPEN_THINK,
id="without_open_think",
),
pytest.param(
True,
WITHOUT_OPEN_THINK_STREAM,
id="without_open_think_stream",
),
pytest.param(
False,
COMPLETE_REASONING,
@@ -18,6 +18,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
FP8ScaledMMLinearLayerConfig,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
get_fp8_min_max,
kFp8DynamicTokenSym,
kFp8StaticChannelSym,
kFp8StaticTensorSym,
@@ -309,7 +310,7 @@ def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel):
_check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens)
fp8_dtype = current_platform.fp8_dtype()
fp8_max = torch.finfo(fp8_dtype).max
fp8_max = get_fp8_min_max()[1]
device = torch.device("cuda")
# Build a bf16 weight and quantize per output channel (one scale per row).
+145
View File
@@ -0,0 +1,145 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Schema/aliasing tests for the AITER FP8 quantization custom ops.
#
# These use torch.library.opcheck, whose test_schema check catches custom ops
# whose implementation aliases an input that the registered schema declares as
# non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant
# regression (a returned scale that aliased the input scale).
#
# Skipped if AITER is not installed or the platform is not ROCm.
import importlib.util
import pytest
import torch
# this import statement is needed to ensure the ops are registered
from vllm._aiter_ops import rocm_aiter_ops
from vllm.platforms import current_platform
aiter_available = importlib.util.find_spec("aiter") is not None
pytestmark = pytest.mark.skipif(
not (current_platform.is_rocm() and aiter_available),
reason="AITER ops are only available on ROCm with aiter package installed",
)
FP8_DTYPE = current_platform.fp8_dtype()
def _x(M=128, N=4096):
return torch.randn((M, N), dtype=torch.float16, device="cuda")
# The in-place per-tensor op takes the fp8 output buffer as an input, which
# opcheck's test_schema cannot exercise ("mul_cuda" is unimplemented for fp8),
# so restrict to the utils that run on fp8 inputs. The aliasing contract for
# this op is instead covered by test_per_tensor_quant_torch_compile below.
_INPLACE_OPCHECK_UTILS = (
"test_faketensor",
"test_aot_dispatch_dynamic",
"test_autograd_registration",
)
def test_per_tensor_quant_static_schema():
"""Static per-tensor: caller provides scale (the aliasing regression)."""
x = _x()
out = torch.empty_like(x, dtype=FP8_DTYPE)
scale = torch.ones(1, dtype=torch.float32, device="cuda")
torch.library.opcheck(
torch.ops.vllm.rocm_aiter_per_tensor_quant,
(out, x, scale, False),
test_utils=_INPLACE_OPCHECK_UTILS,
)
def test_per_tensor_quant_dynamic_schema():
"""Dynamic per-tensor: op computes scale into the caller's buffer."""
x = _x()
out = torch.empty_like(x, dtype=FP8_DTYPE)
scale = torch.empty(1, dtype=torch.float32, device="cuda")
torch.library.opcheck(
torch.ops.vllm.rocm_aiter_per_tensor_quant,
(out, x, scale, True),
test_utils=_INPLACE_OPCHECK_UTILS,
)
def test_per_token_quant_dynamic_schema():
"""Dynamic per-token: op computes scale into a freshly allocated buffer."""
x = _x()
torch.library.opcheck(
torch.ops.vllm.rocm_aiter_per_token_quant,
(x, FP8_DTYPE, None),
)
def test_group_fp8_quant_schema():
"""Dynamic per-token-group quant."""
x = _x()
torch.library.opcheck(
torch.ops.vllm.rocm_aiter_group_fp8_quant,
(x, 128),
)
@pytest.mark.parametrize("dynamic", [True, False])
def test_per_tensor_quant_matches_native(dynamic):
"""Wrapper output matches the native scaled_fp8_quant reference."""
from vllm import _custom_ops as ops
torch.manual_seed(0)
x = _x()
if dynamic:
scale_in = None
else:
scale_in = torch.tensor([0.5], dtype=torch.float32, device="cuda")
out, scale = rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, scale_in)
ref_out, ref_scale = ops.scaled_fp8_quant(x, scale_in)
assert out.shape == x.shape
assert out.dtype == FP8_DTYPE
assert scale.shape == ref_scale.shape
if not dynamic:
# static scale is passed through unchanged
assert torch.equal(scale, scale_in)
# Compare dequantized values to be robust to 1-ULP fp8 boundary flips.
deq = out.to(torch.float32) * scale
ref_deq = ref_out.to(torch.float32) * ref_scale
torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2)
@pytest.mark.parametrize("dynamic", [True, False])
def test_per_tensor_quant_torch_compile(monkeypatch, dynamic):
"""per_tensor_quant compiles under inductor without an aliasing error.
Forces the custom-op aliasing check to error (it is otherwise only a
warning outside CI), so a regression that returns an input-aliasing
scale fails here regardless of the CI env var.
"""
aliasing_cfg = pytest.importorskip("torch._functorch.config")
monkeypatch.setattr(
aliasing_cfg, "error_on_custom_op_aliasing", True, raising=False
)
x = _x()
scale = None if dynamic else torch.tensor([0.5], dtype=torch.float32, device="cuda")
def fn(x, s):
return rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, s)
compiled = torch.compile(fn, fullgraph=True, backend="inductor", dynamic=False)
out_eager, scale_eager = fn(x, scale)
out_compiled, scale_compiled = compiled(x, scale)
assert out_compiled.shape == out_eager.shape
torch.testing.assert_close(
out_compiled.to(torch.float32) * scale_compiled,
out_eager.to(torch.float32) * scale_eager,
rtol=2e-2,
atol=2e-2,
)
+1 -1
View File
@@ -188,7 +188,7 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
is_moe=False,
is_quantized=True,
),
False,
True,
),
(
SimpleNamespace(
@@ -16,7 +16,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
from vllm.tokenizers import get_tokenizer
from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser
MODEL = "zai-org/GLM-4.5"
MODEL = "zai-org/GLM-4.7"
@pytest.fixture(scope="module")
@@ -136,9 +136,10 @@ class TestGlm47Streaming:
_reset(glm47_tool_parser)
chunks = ["<tool_call>", "get_current_date", "</tool_call>"]
current_text = ""
deltas = []
for chunk in chunks:
current_text += chunk
glm47_tool_parser.extract_tool_calls_streaming(
delta = glm47_tool_parser.extract_tool_calls_streaming(
previous_text="",
current_text=current_text,
delta_text=chunk,
@@ -147,7 +148,23 @@ class TestGlm47Streaming:
delta_token_ids=[],
request=mock_request,
)
assert len(glm47_tool_parser.prev_tool_call_arr) >= 1
if delta:
deltas.append(delta)
tool_calls = [
tool_call for delta in deltas for tool_call in (delta.tool_calls or [])
]
names = [
tool_call.function.name
for tool_call in tool_calls
if tool_call.function and tool_call.function.name
]
arguments = [
tool_call.function.arguments
for tool_call in tool_calls
if tool_call.function and tool_call.function.arguments
]
assert names == ["get_current_date"]
assert "".join(arguments) == "{}"
def test_with_args(self, glm47_tool_parser, mock_request):
_reset(glm47_tool_parser)
@@ -161,9 +178,10 @@ class TestGlm47Streaming:
"</tool_call>",
]
current_text = ""
deltas = []
for chunk in chunks:
current_text += chunk
glm47_tool_parser.extract_tool_calls_streaming(
delta = glm47_tool_parser.extract_tool_calls_streaming(
previous_text="",
current_text=current_text,
delta_text=chunk,
@@ -172,5 +190,13 @@ class TestGlm47Streaming:
delta_token_ids=[],
request=mock_request,
)
args = json.loads(glm47_tool_parser.prev_tool_call_arr[0]["arguments"])
if delta:
deltas.append(delta)
arguments = [
tool_call.function.arguments
for delta in deltas
for tool_call in (delta.tool_calls or [])
if tool_call.function and tool_call.function.arguments
]
args = json.loads("".join(arguments))
assert args["city"] == "Beijing"
File diff suppressed because it is too large Load Diff
@@ -1300,6 +1300,73 @@ def test_streaming_multi_param_single_chunk(qwen3_tool_parser, qwen3_tokenizer):
assert args["unit"] == "fahrenheit"
def test_streaming_complete_tool_call_single_delta(qwen3_tool_parser):
"""Regression: one delta may contain a complete tool call."""
request = ChatCompletionRequest(model=MODEL, messages=[])
from tests.tool_parsers.utils import (
run_tool_extraction_streaming,
)
reconstructor = run_tool_extraction_streaming(
qwen3_tool_parser,
[
(
"<tool_call>\n"
"<function=get_current_weather>\n"
"<parameter=city>\nDallas\n</parameter>\n"
"<parameter=state>\nTX\n</parameter>\n"
"</function>\n"
"</tool_call>"
)
],
request,
assert_one_tool_per_delta=False,
)
assert len(reconstructor.tool_calls) == 1
assert reconstructor.tool_calls[0].function.name == "get_current_weather"
args = json.loads(reconstructor.tool_calls[0].function.arguments)
assert args == {"city": "Dallas", "state": "TX"}
def test_streaming_next_tool_call_starts_in_close_delta(qwen3_tool_parser):
"""Regression: a close delta may also contain the next tool call."""
request = ChatCompletionRequest(model=MODEL, messages=[])
from tests.tool_parsers.utils import (
run_tool_extraction_streaming,
)
reconstructor = run_tool_extraction_streaming(
qwen3_tool_parser,
[
"<tool_call>\n",
"<function=get_current_weather>\n",
"<parameter=city>\nDallas\n</parameter>\n",
"<parameter=state>\nTX\n</parameter>\n",
"</function>",
(
"\n</tool_call>\n"
"<tool_call>\n"
"<function=get_current_weather>\n"
"<parameter=city>\nOrlando\n</parameter>\n"
"<parameter=state>\nFL\n</parameter>\n"
"</function>\n"
"</tool_call>"
),
],
request,
assert_one_tool_per_delta=False,
)
assert len(reconstructor.tool_calls) == 2
first_args = json.loads(reconstructor.tool_calls[0].function.arguments)
second_args = json.loads(reconstructor.tool_calls[1].function.arguments)
assert first_args == {"city": "Dallas", "state": "TX"}
assert second_args == {"city": "Orlando", "state": "FL"}
def test_no_double_serialization_string_args(qwen3_tool_parser):
"""Regression: string arguments must not be double-serialized (PR #35615)."""
tools = [
+67
View File
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import pytest
from vllm.tool_parsers.utils import (
@@ -91,6 +93,71 @@ class TestCoerceToSchemaType:
def test_invalid_number_fallback(self):
assert coerce_to_schema_type("abc", "number") == "abc"
class TestNonFiniteNumbers:
"""Non-finite numeric strings must not crash and must coerce to a
JSON-serializable value.
Regression: ``int(float("inf"))`` raised an uncaught ``OverflowError``
(only ``ValueError``/``TypeError`` were handled), and ``"1e999"``
round-tripped through ``json.loads`` to a float ``inf`` that
``json.dumps`` renders as invalid JSON ``Infinity``.
"""
@pytest.mark.parametrize(
"value", ["inf", "-inf", "Infinity", "1e999", "nan", "-nan"]
)
def test_non_finite_number_does_not_crash(self, value):
# Must not raise (previously OverflowError for inf/1e999/Infinity).
result = coerce_to_schema_type(value, "number")
# Result must serialize to valid, finite JSON and round-trip.
assert json.loads(json.dumps(result)) == result
@pytest.mark.parametrize("value", ["inf", "-inf", "1e999"])
def test_non_finite_number_preserved_as_string(self, value):
assert coerce_to_schema_type(value, "number") == value
@pytest.mark.parametrize("value", ["inf", "1e999", "Infinity"])
def test_non_finite_integer_not_float_inf(self, value):
result = coerce_to_schema_type(value, "integer")
assert isinstance(result, str)
assert result == value
class TestNonFiniteContainers:
"""Non-finite floats nested in object/array values must not produce
invalid JSON.
Regression: the ``object``/``array`` branch returned
``json.loads(value)`` directly, so ``"[1e999]"`` became ``[inf]`` and
``'{"x": Infinity}'`` became ``{"x": inf}`` -- values that
``json.dumps`` later renders as invalid JSON (``Infinity``/``NaN``).
"""
@pytest.mark.parametrize(
"value", ["[1e999]", "[1, 2, 1e999]", "[NaN]", "[-Infinity]"]
)
def test_array_with_non_finite_preserved_as_string(self, value):
result = coerce_to_schema_type(value, "array")
assert result == value
assert json.loads(json.dumps(result)) == result
@pytest.mark.parametrize(
"value", ['{"x": 1e999}', '{"x": Infinity}', '{"a": [1e999, 2]}']
)
def test_object_with_non_finite_preserved_as_string(self, value):
result = coerce_to_schema_type(value, "object")
assert result == value
assert json.loads(json.dumps(result)) == result
def test_finite_array_still_coerced(self):
assert coerce_to_schema_type("[1, 2, 3]", "array") == [1, 2, 3]
def test_finite_object_still_coerced(self):
assert coerce_to_schema_type('{"a": 1}', "object") == {"a": 1}
def test_unknown_type_non_finite_falls_back_to_string(self):
# Exercises the final json.loads fallback path.
assert coerce_to_schema_type("1e999", "unknown_type") == "1e999"
class TestBooleanType:
def test_true(self):
assert coerce_to_schema_type("true", "boolean") is True
+228
View File
@@ -0,0 +1,228 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for contiguous KV cache packing."""
from unittest.mock import MagicMock
import pytest
import torch
from vllm import envs
from vllm.v1.core.kv_cache_utils import (
_get_kv_cache_config_deepseek_v4,
get_kv_cache_config_from_groups,
)
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheGroupSpec,
KVCacheTensor,
MLAAttentionSpec,
SlidingWindowSpec,
UniformTypeKVCacheSpecs,
)
def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
return MLAAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=512,
dtype=torch.uint8,
page_size_padded=page_size,
cache_dtype_str="fp8_ds_mla",
model_version="deepseek_v4",
alignment=576,
)
def _make_full_spec() -> FullAttentionSpec:
return FullAttentionSpec(
block_size=16,
num_kv_heads=2,
head_size=64,
dtype=torch.float16,
)
def _make_sw_spec() -> SlidingWindowSpec:
return SlidingWindowSpec(
block_size=16,
num_kv_heads=2,
head_size=64,
dtype=torch.float16,
sliding_window=128,
)
def _make_groups(n_c4, n_c128, n_swa):
PS_C4_MLA = 37440
PS_C4_IDX = 8640
PS_C128 = 1728
PS_SWA = 37440
mla_specs = {}
for i in range(n_c4):
mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA)
mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX)
for i in range(n_c128):
mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128)
mla_group = KVCacheGroupSpec(
layer_names=list(mla_specs.keys()),
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs),
)
swa_specs = {}
for i in range(n_swa):
swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA)
swa_group = KVCacheGroupSpec(
layer_names=list(swa_specs.keys()),
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs),
)
return [mla_group, swa_group]
def _mock_vllm_config():
config = MagicMock()
config.cache_config.num_gpu_blocks_override = None
return config
def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
groups = _make_groups(n_c4, n_c128, n_swa)
return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem)
def _page_sizes_by_layer(
groups: list[KVCacheGroupSpec],
) -> dict[str, int]:
page_sizes = {}
for group in groups:
specs = group.kv_cache_spec.kv_cache_specs
for layer_name in group.layer_names:
page_sizes[layer_name] = specs[layer_name].page_size_bytes
return page_sizes
class TestInterleavedPacking:
def test_all_tensors_have_block_stride(self):
_, tensors = _run()
for t in tensors:
assert t.block_stride > 0
def test_all_tensors_share_same_size(self):
_, tensors = _run()
sizes = set(t.size for t in tensors)
assert len(sizes) == 1
assert sizes.pop() > 0
def test_offsets_within_one_block(self):
_, tensors = _run()
for t in tensors:
assert t.offset < t.block_stride
def test_all_layers_accounted_for(self):
n_c4, n_c128, n_swa = 5, 4, 7
_, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa)
all_names = set()
for t in tensors:
all_names.update(t.shared_by)
expected = n_c4 * 2 + n_c128 + n_swa
assert len(all_names) == expected
def test_strided_views_are_independent(self):
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
page_sizes = _page_sizes_by_layer(groups)
num_blocks, tensors = _get_kv_cache_config_deepseek_v4(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
views = []
for t in tensors:
page_size = page_sizes[t.shared_by[0]]
v = torch.as_strided(
backing,
size=(num_blocks, page_size),
stride=(t.block_stride, 1),
storage_offset=t.offset,
)
views.append(v)
for i, v in enumerate(views):
v.fill_(i + 1)
for i, v in enumerate(views):
assert (v == i + 1).all(), f"View {i} was corrupted"
def test_hma_attention_groups_keep_default_backing(self, monkeypatch):
monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False)
full = _make_full_spec()
sw = _make_sw_spec()
page_size = full.page_size_bytes
groups = [
KVCacheGroupSpec(["full.0", "full.1"], full),
KVCacheGroupSpec(["sw.0", "sw.2"], sw),
KVCacheGroupSpec(["sw.1", "sw.3"], sw),
]
config = get_kv_cache_config_from_groups(
_mock_vllm_config(), groups, available_memory=page_size * 2 * 32
)
assert config.num_blocks == 32
assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32
assert config.kv_cache_tensors == [
KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),
KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
]
def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch):
monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False)
full = _make_full_spec()
sw = _make_sw_spec()
page_size = full.page_size_bytes
groups = [
KVCacheGroupSpec(["full.0", "full.1"], full),
KVCacheGroupSpec(["sw.0", "sw.2"], sw),
KVCacheGroupSpec(["sw.1", "sw.3"], sw),
]
config = get_kv_cache_config_from_groups(
_mock_vllm_config(), groups, available_memory=page_size * 2 * 32
)
assert config.num_blocks == 32
assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}
assert config.kv_cache_tensors == [
KVCacheTensor(
size=page_size * 2 * 32,
shared_by=["full.0", "sw.0", "sw.1"],
offset=0,
block_stride=page_size * 2,
),
KVCacheTensor(
size=page_size * 2 * 32,
shared_by=["full.1", "sw.2", "sw.3"],
offset=page_size,
block_stride=page_size * 2,
),
]
def test_single_group_attention_keeps_unpacked_layout(self):
spec = _make_full_spec()
groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
config = get_kv_cache_config_from_groups(
_mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
)
assert sum(t.size for t in config.kv_cache_tensors) == (
spec.page_size_bytes * 2 * 32
)
assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]
if __name__ == "__main__":
pytest.main([__file__, "-v"])
+102 -7
View File
@@ -117,6 +117,7 @@ def new_kv_cache_spec(
page_size_padded=None,
sliding_window=None,
attention_chunk_size=None,
indexes_kv_by_block_stride=False,
):
return FullAttentionSpec(
block_size=block_size,
@@ -126,6 +127,7 @@ def new_kv_cache_spec(
page_size_padded=page_size_padded,
sliding_window=sliding_window,
attention_chunk_size=attention_chunk_size,
indexes_kv_by_block_stride=indexes_kv_by_block_stride,
)
@@ -136,6 +138,7 @@ def new_sliding_window_spec(
dtype=torch.float32,
page_size_padded=None,
sliding_window=1,
indexes_kv_by_block_stride=False,
):
return SlidingWindowSpec(
block_size=block_size,
@@ -144,6 +147,7 @@ def new_sliding_window_spec(
dtype=dtype,
page_size_padded=page_size_padded,
sliding_window=sliding_window,
indexes_kv_by_block_stride=indexes_kv_by_block_stride,
)
@@ -1799,16 +1803,38 @@ def test_get_kv_cache_config_one_worker():
],
)
# different hidden size that cannot be aligned by using different block size
# different hidden size that cannot be aligned by using different block size,
# but can be aligned by padding the smaller physical page.
swa_spec = new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True)
kv_cache_specs_hybrid = {
"layer_1": new_kv_cache_spec(head_size=64),
"layer_2": new_sliding_window_spec(head_size=96),
"layer_1": new_kv_cache_spec(head_size=64, indexes_kv_by_block_stride=True),
"layer_2": swa_spec,
}
with pytest.raises(NotImplementedError):
get_kv_cache_configs(
vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
)[0]
kv_cache_config_hybrid = get_kv_cache_configs(
vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
)[0]
padded_page_size = swa_spec.page_size_bytes
assert kv_cache_config_hybrid == KVCacheConfig(
num_blocks=42,
kv_cache_tensors=[
KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]),
],
kv_cache_groups=[
KVCacheGroupSpec(
["layer_1"],
new_kv_cache_spec(
head_size=64,
page_size_padded=padded_page_size,
indexes_kv_by_block_stride=True,
),
),
KVCacheGroupSpec(
["layer_2"],
new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True),
),
],
)
# Test num_gpu_blocks_override
vllm_config.cache_config.num_gpu_blocks_override = 16
@@ -2322,6 +2348,75 @@ def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
"""DFlash drafters can have a smaller head size than the target model.
For example, MiMo uses 192-dim target KV heads while its DFlash draft uses
128-dim KV heads. The resulting page sizes are 3:2 rather than an integer
block-size multiple, so the smaller page must be padded instead.
"""
# Both layers' backends opt into the padded-page strided view (e.g.
# FlashAttention / its DiffKV subclass), so padding is allowed.
target_spec = new_kv_cache_spec(
block_size=16,
num_kv_heads=1,
head_size=192,
dtype=torch.bfloat16,
indexes_kv_by_block_stride=True,
)
draft_spec = new_sliding_window_spec(
block_size=16,
num_kv_heads=1,
head_size=128,
dtype=torch.bfloat16,
sliding_window=1024,
indexes_kv_by_block_stride=True,
)
unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size(
{
"target_attn": target_spec,
"draft_attn": draft_spec,
}
)
assert unified_specs["target_attn"] == target_spec
unified_draft_spec = unified_specs["draft_attn"]
assert unified_draft_spec.block_size == draft_spec.block_size
assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes
assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes
assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes
def test_unify_kv_cache_page_size_padding_requires_backend_support():
"""Padding is gated on the backend declaring ``indexes_kv_by_block_stride``.
A backend that does not support the strided padded-page view must raise
rather than silently padding (and misreading KV at runtime).
"""
target_spec = new_kv_cache_spec(
block_size=16,
num_kv_heads=1,
head_size=192,
dtype=torch.bfloat16,
indexes_kv_by_block_stride=True,
)
# The non-divisible draft layer needs padding but its backend does not
# support the strided padded-page view -> must raise, not silently pad.
draft_spec = new_sliding_window_spec(
block_size=16,
num_kv_heads=1,
head_size=128,
dtype=torch.bfloat16,
sliding_window=1024,
indexes_kv_by_block_stride=False,
)
specs = {"target_attn": target_spec, "draft_attn": draft_spec}
with pytest.raises(NotImplementedError):
kv_cache_utils.unify_kv_cache_spec_page_size(specs)
def test_unify_hybrid_kv_cache_specs():
# 1. has_full_attention and has_sliding_window
before_spec_1 = new_kv_cache_spec()
+37
View File
@@ -144,6 +144,43 @@ def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders():
assert req.request_id in output.num_scheduled_tokens
def test_cached_request_data_resumed_all_token_ids_mrv1_only():
"""all_token_ids carries a resumed request's token ids to the connector
for the V1 model runner, but is skipped entirely for the V2 model runner.
"""
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
scheduler = create_scheduler()
(req,) = create_requests(num_requests=1, num_tokens=8)
req.append_output_token_ids([101, 102, 103])
# A resumed request was not scheduled in the previous step.
assert req.request_id not in scheduler.prev_step_scheduled_req_ids
empty_blocks = KVCacheBlocks(blocks=((),))
def make_cached():
return scheduler._make_cached_request_data(
running_reqs=[],
resumed_reqs=[req],
num_scheduled_tokens={req.request_id: 1},
spec_decode_tokens={},
req_to_new_blocks={req.request_id: empty_blocks},
)
# V1 model runner: the full token id list is propagated.
assert not scheduler.use_v2_model_runner
cached = make_cached()
assert req.request_id in cached.resumed_req_ids
assert cached.all_token_ids[req.request_id] == list(req.all_token_ids)
# V2 model runner: all_token_ids is skipped entirely.
scheduler.use_v2_model_runner = True
cached = make_cached()
assert req.request_id in cached.resumed_req_ids
assert cached.all_token_ids == {}
def test_schedule_partial_requests():
"""Test scheduling behavior with partial requests.
@@ -4,9 +4,16 @@
import pytest
from vllm import LLM, SamplingParams
from vllm.platforms import current_platform
from ....utils import create_new_process_for_each_test
if current_platform.is_rocm():
pytest.skip(
"Cascade attention backends FLASH_ATTN and FLASHINFER are notsupported on ROCm",
allow_module_level=True,
)
@create_new_process_for_each_test()
@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"])

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