Compare commits

...
Author SHA1 Message Date
Alexander Matveev a60418e6fb Sparse MLA on Hopper: Use SGLang's kernel for the sparse mla low latency runs
Signed-off-by: Alexander Matveev <amatveev@redhat.com>
2026-04-17 16:51:32 +00:00
Michael GoinandGitHub 1174723eba Fix TURBOQUANT backend selection in cuda.py (#40060)
Signed-off-by: Michael Goin <mgoin64@gmail.com>
2026-04-17 07:31:41 -07:00
sychen52andGitHub 6b2b7bd0eb Add nvfp4 support to reshape_and_cache_flash (#37332)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
2026-04-17 07:28:00 -07:00
Ben BrowningandGitHub 70770268c3 Add @bbrowning to CODEOWNERS (#40141)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-04-17 09:51:48 -04:00
ChaunceyandGitHub 7a51b3e415 [Bugfix] Fix empty delta detection in Qwen3XMLToolParser streaming (#40090)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-04-17 13:34:55 +00:00
Li, JiangandGitHub d02421a7db [CPU] Refactor CPU affinity and memory management (#39781)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
2026-04-17 21:01:08 +08:00
Lukas GeigerandGitHub b1dc87a098 [Models][Gemma4] Prevent GPU/CPU sync in embed_input_ids (#39234)
Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>
2026-04-17 12:37:21 +00:00
Or OzeriandGitHub 79a5b63253 [kv_offload]: Fix num CPU blocks for UniformTypeKVCacheSpecs (#39617)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
2026-04-17 15:13:55 +03:00
MaralandGitHub c0c98b8b9a [Bugfix] Add Marlin kernel in block scaled mm kernel selection. (#40105)
Signed-off-by: maral <maralbahari.98@gmail.com>
2026-04-17 10:20:32 +00:00
wang.yuqiandGitHub 8d2cff8140 [Examples] Resettle Observability examples. (#40123)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-04-17 03:13:31 -07:00
Cyrus LeungandGitHub 4f436782af [Misc] Improve new PR bot trigger condition (#40114)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
2026-04-17 16:56:22 +08:00
z1yingandGitHub bf45e6d0a5 [Doc] Add Gemma 4 to supported models list (#39607)
Signed-off-by: z1ying <tzzying@outlook.com>
Signed-off-by: Ziying Tao <tzzying@outlook.com>
2026-04-17 13:42:52 +08:00
978a4462bb [CI Failure] Fix Plugin Tests (2 GPUs) Failure (#40083)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
Co-authored-by: Michele Gazzetti <michele.gazzetti1@ibm.com>
2026-04-17 04:17:39 +00:00
Michael GoinandGitHub 1948d0c467 [UX] Defer some imports on CLI paths to save ~2s (#40056)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-04-16 19:48:37 -07:00
Shinichi HemmiandGitHub 4c47710bf7 [CI/Build] Apply ruff formatter to pass pre-commit (#40078)
Signed-off-by: Hemmi Shinichi <shemmi@preferred.jp>
2026-04-17 08:54:32 +08:00
Giancarlo DelfinandGitHub bf9a5ddb24 [MLA] Optimize mla indexer prepare uniform decode for MTP > 1 (#39458)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
2026-04-16 16:27:51 -07:00
bnellnmandGitHub 79e799ebbd [Bugfix] Temporarily disable B200 fp4 MoE layer tests (#40057)
Signed-off-by: Bill Nell <bnell@redhat.com>
2026-04-16 19:26:55 -04:00
Netanel HaberandGitHub c4e601c73c Bugfix: Parakeet: .conv.pointwise/depthwise_conv1/2.bias weigths can exist even if convolution_bias=False (#40007)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
2026-04-16 23:22:05 +00:00
BadrBasowidandGitHub 29057d3bee [Compilation] Add Unit Tests for VllmFusionPatternMatcherPass (#39692)
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
2026-04-16 22:57:16 +00:00
Matthew BonanniandGitHub 219bb5b8c0 [Misc] Update committers.md (#40058)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
2026-04-16 13:48:41 -07:00
Asaf GardinandGitHub ad2b1277f9 [Quantization] Consolidate experts_int8 with fp8 online quantization (#38463)
Signed-off-by: Josephasafg <ajgard7@gmail.com>
2026-04-16 13:12:20 -07:00
roikoren755andGitHub b897f00c9c Gate SSU dispatch setup (#40039)
Signed-off-by: Roi Koren <roik@nvidia.com>
2026-04-16 13:06:01 -07:00
adf9bb3c57 [CI] Add weight transfer tests to CI (#39821)
Signed-off-by: SumanthRH <sumanthrh99@gmail.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-04-16 15:51:45 -04:00
Flora FengandGitHub b16fda62b7 [Misc] Add @sfeng33 to CODEOWNERS (#40048)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-04-16 12:25:29 -07:00
Yufeng HeandGitHub de111f3246 [Bugfix] Fix bench_serve UTF-8 decode crash on split multi-byte chars (#38732) 2026-04-16 12:01:25 -07:00
Jared WenandGitHub afabb5f45a [bugfix] Normalize tool message content from array to string format (#39899)
Signed-off-by: JaredforReal <w13431838023@gmail.com>
2026-04-16 11:54:39 -07:00
Roger WangandGitHub 3abb7560c0 [Bugfix] Fix audioflamingo test (#40052)
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-04-16 11:53:58 -07:00
87 changed files with 5716 additions and 1041 deletions
@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
exit 1
fi
#echo "--- DP+TP"
#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
#server_pid=$!
#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
#vllm bench serve \
# --backend vllm \
# --dataset-name random \
# --model meta-llama/Llama-3.2-3B-Instruct \
# --num-prompts 20 \
# --result-dir ./test_results \
# --result-filename dp_pp.json \
# --save-result \
# --endpoint /v1/completions
#kill -s SIGTERM $server_pid; wait $server_pid || true
#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
#if [ "$failed_req" -ne 0 ]; then
# echo "Some requests were failed!"
# exit 1
#fi
echo "--- DP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename dp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/dp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
+2
View File
@@ -196,6 +196,8 @@ steps:
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
- pytest -v -s tests/distributed/test_packed_tensor.py
- label: Distributed Tests (2 GPUs)(B200)
device: b200
+7 -3
View File
@@ -44,8 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/vllm/pooling_params.py @noooop @DarkLight1337
/vllm/tokenizers @DarkLight1337 @njhill
/vllm/renderers @DarkLight1337 @njhill
/vllm/reasoning @aarnphm @chaunceyjiang
/vllm/tool_parsers @aarnphm @chaunceyjiang
/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning
# vLLM V1
/vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -91,7 +92,10 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
/tests/v1/kv_connector/nixl_integration @NickLucche
/tests/v1/kv_connector @ApostaC @orozery
/tests/v1/kv_offload @ApostaC @orozery
/tests/v1/determinism @yewentao256
/tests/v1/determinism @yewentao256
/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning
# Transformers modeling backend
/vllm/model_executor/models/transformers @hmellor
+1
View File
@@ -45,6 +45,7 @@ jobs:
- name: Smoke test vllm serve
run: |
# Start server in background
VLLM_CPU_KVCACHE_SPACE=1 \
vllm serve Qwen/Qwen3-0.6B \
--max-model-len=2K \
--load-format=dummy \
+5 -5
View File
@@ -62,14 +62,14 @@ jobs:
const prAuthor = context.payload.pull_request.user.login;
const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`,
per_page: 1,
});
const authorPRCount = searchResults.total_count;
console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
const mergedPRCount = searchResults.total_count;
console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`);
if (authorPRCount === 1) {
if (mergedPRCount === 0) {
console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
await github.rest.issues.createComment({
owner,
@@ -98,5 +98,5 @@ jobs:
].join('\n'),
});
} else {
console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`);
}
+15
View File
@@ -923,6 +923,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
# nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
# from cache_kernels.cu, so it belongs in _C rather than _C_stable.
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
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 NVFP4 for archs: ${FP4_ARCHS}")
@@ -949,6 +957,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
SRCS "${SRCS}"
CUDA_ARCHS "${FP4_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
set_gencode_flags_for_srcs(
SRCS "${NVFP4_KV_SRC}"
CUDA_ARCHS "${FP4_ARCHS}")
target_sources(_C PRIVATE ${NVFP4_KV_SRC})
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 NVFP4 for archs: ${FP4_ARCHS}")
@@ -1226,6 +1240,7 @@ endif()
if (VLLM_GPU_LANG STREQUAL "CUDA")
include(cmake/external_projects/deepgemm.cmake)
include(cmake/external_projects/flashmla.cmake)
include(cmake/external_projects/cutlass_fa3.cmake)
include(cmake/external_projects/qutlass.cmake)
# vllm-flash-attn should be last as it overwrites some CMake functions
+15 -14
View File
@@ -30,6 +30,21 @@ else()
list(APPEND CXX_COMPILE_FLAGS
"-fopenmp"
"-DVLLM_CPU_EXTENSION")
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
# and create a local shim dir with it
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
find_library(OPEN_MP
NAMES gomp
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
NO_DEFAULT_PATH
REQUIRED
)
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
if (OPEN_MP)
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
endif()
endif()
if (NOT MACOSX_FOUND)
@@ -175,20 +190,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
if(NOT NPROC)
set(NPROC 4)
endif()
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
# and create a local shim dir with it
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
find_library(OPEN_MP
NAMES gomp
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
NO_DEFAULT_PATH
REQUIRED
)
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
if (OPEN_MP)
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
endif()
# Fetch and populate ACL
if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
+163
View File
@@ -0,0 +1,163 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# CUTLASS FA3 MLA Sparse Attention — requires CUDA >= 12.4, SM90a
#
# Vendors the sgl-attn CUTLASS FlashAttention3 kernel from SGLang into vLLM
# as a self-contained extension (_cutlass_fa3_C). This provides a high-
# performance sparse MLA attention kernel for SM90 (Hopper) GPUs.
#
# Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
# CUTLASS: https://github.com/NVIDIA/cutlass (commit 57e3cfb4)
# Guard: CUDA >= 12.4 required for SM90a features used by FA3
if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.4")
message(STATUS "Skipping CUTLASS FA3: requires CUDA >= 12.4")
# Create empty target so setup.py doesn't fail on unsupported systems
add_custom_target(_cutlass_fa3_C)
return()
endif()
# Guard: SM90 architecture required
set(CUTLASS_FA3_SUPPORT_ARCHS)
list(APPEND CUTLASS_FA3_SUPPORT_ARCHS "9.0a")
cuda_archs_loose_intersection(
CUTLASS_FA3_ARCHS "${CUTLASS_FA3_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
if(NOT CUTLASS_FA3_ARCHS)
message(STATUS "Skipping CUTLASS FA3: requires SM90 (CUDA_ARCHS=${CUDA_ARCHS})")
add_custom_target(_cutlass_fa3_C)
return()
endif()
include(FetchContent)
# Fetch sgl-attn (Flash Attention 3 kernels from SGLang)
# We only need the source files, not the build system, so we use
# FetchContent_Populate to download without building.
if (DEFINED ENV{SGL_ATTN_SRC_DIR})
set(SGL_ATTN_SRC_DIR $ENV{SGL_ATTN_SRC_DIR})
endif()
if(SGL_ATTN_SRC_DIR)
FetchContent_Declare(cutlass_fa3
SOURCE_DIR ${SGL_ATTN_SRC_DIR})
else()
FetchContent_Declare(cutlass_fa3
GIT_REPOSITORY https://github.com/sgl-project/sgl-attn.git
GIT_TAG bcf72ccc6816b36a5fae2c5a3c027604629785e0
GIT_PROGRESS TRUE
GIT_SHALLOW FALSE)
endif()
FetchContent_GetProperties(cutlass_fa3)
if(NOT cutlass_fa3_POPULATED)
FetchContent_Populate(cutlass_fa3)
endif()
message(STATUS "CUTLASS FA3 sgl-attn source: ${cutlass_fa3_SOURCE_DIR}")
# Fetch CUTLASS for FA3 (headers only, separate from vLLM's main CUTLASS
# to avoid version conflicts). Use FetchContent_Populate to avoid running
# CUTLASS's own CMakeLists.txt which would create conflicting targets.
if (DEFINED ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
set(CUTLASS_FA3_CUTLASS_SRC_DIR $ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
endif()
if(CUTLASS_FA3_CUTLASS_SRC_DIR)
FetchContent_Declare(cutlass_for_fa3
SOURCE_DIR ${CUTLASS_FA3_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(cutlass_for_fa3
GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
GIT_TAG 57e3cfb47a2d9e0d46eb6335c3dc411498efa198
GIT_PROGRESS TRUE
GIT_SHALLOW FALSE)
endif()
FetchContent_GetProperties(cutlass_for_fa3)
if(NOT cutlass_for_fa3_POPULATED)
FetchContent_Populate(cutlass_for_fa3)
endif()
message(STATUS "CUTLASS FA3 cutlass source: ${cutlass_for_fa3_SOURCE_DIR}")
set(FA3_SRC "${cutlass_fa3_SOURCE_DIR}/hopper")
# flash_api.cpp dispatches to all head dimensions + dtypes (BF16, FP16, FP8)
# at compile time. With FLASHATTENTION_DISABLE_SM8x, only SM90 instantiations
# are needed. We exclude hdimall_* (fails on CUDA 13+) and backward files.
file(GLOB FA3_INSTANTIATION_SOURCES
# BF16 instantiations
"${FA3_SRC}/instantiations/flash_fwd_hdim64_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim96_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim128_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim192_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim256_bf16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdimdiff_bf16*_sm90.cu"
# FP16 instantiations
"${FA3_SRC}/instantiations/flash_fwd_hdim64_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim96_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim128_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim192_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim256_fp16*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdimdiff_fp16*_sm90.cu"
# FP8 (e4m3) instantiations
"${FA3_SRC}/instantiations/flash_fwd_hdim64_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim96_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim128_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim192_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdim256_e4m3*_sm90.cu"
"${FA3_SRC}/instantiations/flash_fwd_hdimdiff_e4m3*_sm90.cu")
set(FA3_CORE_SOURCES
"${FA3_SRC}/flash_api.cpp"
"${FA3_SRC}/flash_prepare_scheduler.cu"
"${FA3_SRC}/flash_fwd_combine.cu")
set(FA3_ALL_SOURCES
"${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_fa3_extension.cc"
${FA3_CORE_SOURCES}
${FA3_INSTANTIATION_SOURCES})
set(FA3_INCLUDE_DIRS
${FA3_SRC}
${cutlass_fa3_SOURCE_DIR}/include
${cutlass_for_fa3_SOURCE_DIR}/include
${cutlass_for_fa3_SOURCE_DIR}/tools/util/include
${CMAKE_CURRENT_SOURCE_DIR}/csrc)
# Set SM90a gencode flags for all FA3 CUDA sources
set_gencode_flags_for_srcs(
SRCS "${FA3_ALL_SOURCES}"
CUDA_ARCHS "${CUTLASS_FA3_ARCHS}")
define_extension_target(_cutlass_fa3_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${FA3_ALL_SOURCES}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${FA3_INCLUDE_DIRS}
USE_SABI 3
WITH_SOABI)
# FA3-specific compile options for CUDA and C++ source files:
# - C++17 required by CUTLASS
# - Fast math for performance
# - Relaxed constexpr for CUTLASS template metaprogramming
# - Disable backward pass, dropout, uneven K (not needed for inference)
# - Enable varlen-only mode (all our use cases are variable-length)
target_compile_options(_cutlass_fa3_C PRIVATE
$<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
$<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
$<$<COMPILE_LANGUAGE:CUDA>:-std=c++17>
$<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
$<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
target_compile_definitions(_cutlass_fa3_C PRIVATE
CUTE_USE_PACKED_TUPLE=1
CUTLASS_ENABLE_GDC_FOR_SM90
CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED
CUTLASS_ENABLE_TENSOR_CORE_MMA=1
FLASHATTENTION_DISABLE_BACKWARD
FLASHATTENTION_DISABLE_DROPOUT
FLASHATTENTION_DISABLE_UNEVEN_K
FLASHATTENTION_DISABLE_SM8x
FLASHATTENTION_VARLEN_ONLY)
message(STATUS "CUTLASS FA3 MLA Sparse: enabled for SM90 (${CUTLASS_FA3_ARCHS})")
+22 -2
View File
@@ -724,6 +724,28 @@ void reshape_and_cache_flash(
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (kv_cache_dtype == "nvfp4") {
#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
// NVFP4 dispatch is compiled separately for SM100+.
extern void reshape_and_cache_nvfp4_dispatch(
torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
torch::Tensor & value_cache, torch::Tensor & slot_mapping,
torch::Tensor & k_scale, torch::Tensor & v_scale);
reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
slot_mapping, k_scale, v_scale);
return;
#else
TORCH_CHECK(false,
"NVFP4 KV cache requires SM100+ (Blackwell). "
"Please rebuild vllm with a Blackwell-compatible CUDA target.");
#endif
}
// Original FP8/auto path.
int block_size = key_cache.size(1);
int64_t key_stride = key.stride(0);
@@ -741,8 +763,6 @@ void reshape_and_cache_flash(
dim3 grid(num_tokens);
dim3 block(std::min(num_heads * head_size, 512));
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
CALL_RESHAPE_AND_CACHE_FLASH);
+4
View File
@@ -141,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping,
const int64_t block_size);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
void eagle_prepare_inputs_padded_kernel_impl(
const torch::Tensor& cu_num_draft_tokens,
@@ -431,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"block_size) -> ()",
&compute_slot_mapping_kernel_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels
ops.def(
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
+73 -6
View File
@@ -13,13 +13,80 @@
#include "cpu/utils.hpp"
#ifdef VLLM_NUMA_DISABLED
std::string init_cpu_threads_env(const std::string& cpu_ids) {
return std::string(
"Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
"no effect to setup thread affinity.");
}
void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
#else
void init_cpu_memory_env(std::vector<int64_t> node_ids) {
// Memory node binding
if (numa_available() != -1) {
// Concatenate all node_ids into a single comma-separated string
if (!node_ids.empty()) {
std::string node_ids_str;
for (const int node_id : node_ids) {
if (!node_ids_str.empty()) {
node_ids_str += ",";
}
node_ids_str += std::to_string(node_id);
}
#endif
bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
bitmask* src_mask = numa_get_mems_allowed();
int pid = getpid();
if (mask && src_mask) {
// move all existing pages to the specified numa node.
*(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
int page_num = numa_migrate_pages(pid, src_mask, mask);
if (page_num == -1) {
TORCH_WARN("numa_migrate_pages failed. errno: " +
std::to_string(errno));
}
// Restrict memory allocation to the selected NUMA node(s).
// Enhances memory locality for the threads bound to those NUMA CPUs.
if (node_ids.size() > 1) {
errno = 0;
numa_set_interleave_mask(mask);
if (errno != 0) {
TORCH_WARN("numa_set_interleave_mask failed. errno: " +
std::to_string(errno));
} else {
TORCH_WARN(
"NUMA binding: Using INTERLEAVE policy for memory "
"allocation across multiple NUMA nodes (nodes: " +
node_ids_str +
"). Memory allocations will be "
"interleaved across the specified NUMA nodes.");
}
} else {
errno = 0;
numa_set_membind(mask);
if (errno != 0) {
TORCH_WARN("numa_set_membind failed. errno: " +
std::to_string(errno));
} else {
TORCH_WARN(
"NUMA binding: Using MEMBIND policy for memory "
"allocation on the NUMA nodes (" +
node_ids_str +
"). Memory allocations will be "
"strictly bound to these NUMA nodes.");
}
}
numa_set_strict(1);
numa_free_nodemask(mask);
numa_free_nodemask(src_mask);
} else {
TORCH_WARN(
"numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
std::to_string(errno));
}
}
}
}
#endif // VLLM_NUMA_DISABLED
namespace cpu_utils {
ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
+72
View File
@@ -0,0 +1,72 @@
/* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* Vendored CUTLASS FA3 MLA attention kernel binding for vLLM.
* Based on sgl-kernel/csrc/flash_extension.cc from SGLang.
*
* This registers the FA3 forward pass as a PyTorch C++ extension under
* the _cutlass_fa3_C namespace, enabling torch.ops._cutlass_fa3_C.fwd().
*
* Original source:
* https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
* sgl-kernel/csrc/flash_extension.cc
*/
#include <Python.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <torch/all.h>
#include <torch/library.h>
#include "sgl_flash_kernel_ops.h"
TORCH_LIBRARY_FRAGMENT(_cutlass_fa3_C, m) {
/*
* CUTLASS FA3 MLA forward pass.
* Signature matches sgl-attn's mha_fwd() exactly.
*/
m.def(
"fwd(Tensor q,"
" Tensor k,"
" Tensor v,"
" Tensor? k_new,"
" Tensor? v_new,"
" Tensor? q_v,"
" Tensor? out,"
" Tensor? cu_seqlens_q,"
" Tensor? cu_seqlens_k,"
" Tensor? cu_seqlens_k_new,"
" Tensor? seqused_q,"
" Tensor? seqused_k,"
" int? max_seqlen_q,"
" int? max_seqlen_k,"
" Tensor? page_table,"
" Tensor? kv_batch_idx,"
" Tensor? leftpad_k,"
" Tensor? rotary_cos,"
" Tensor? rotary_sin,"
" Tensor? seqlens_rotary,"
" Tensor? q_descale,"
" Tensor? k_descale,"
" Tensor? v_descale,"
" float? softmax_scale,"
" bool is_causal,"
" int window_size_left,"
" int window_size_right,"
" int attention_chunk,"
" float softcap,"
" bool is_rotary_interleaved,"
" Tensor? scheduler_metadata,"
" int num_splits,"
" bool? pack_gqa,"
" int sm_margin,"
" Tensor? sinks"
") -> (Tensor, Tensor, Tensor, Tensor)");
m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
}
// Python module initialization for _cutlass_fa3_C
PyMODINIT_FUNC PyInit__cutlass_fa3_C() {
static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, "_cutlass_fa3_C",
nullptr, 0, nullptr};
return PyModule_Create(&module);
}
+275
View File
@@ -0,0 +1,275 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
// NVFP4 KV cache store kernel.
// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them
// into the paged KV cache.
//
// Per page layout: [K_data | K_scale | V_data | V_scale]
// Both data and scale regions are contiguous per head, enabling direct
// TMA descriptor use.
//
// Reuses device functions from nvfp4_utils.cuh:
// - cvt_warp_fp16_to_fp4() for bf16 → fp4 quantization + block scale
// - pack_fp4() for packing float pairs to fp4
// - reciprocal_approximate_ftz() for fast reciprocal
#define NVFP4_ENABLE_ELTS16 1
#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include "dispatch_utils.h"
namespace vllm {
// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
// The swizzle pattern for HND layout is:
// [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
// where T = block_size (page_size), S = scale_dim = head_size // 16.
//
// For a linear (t, s) position, the swizzled position is:
// swizzled_t = (t / 4) * 4 + (s / (S / 4))
// swizzled_s = (s % (S / 4)) * 4 + (t % 4)
__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
int scale_dim) {
int s_group = scale_dim / 4;
int swizzled_t = (t / 4) * 4 + (s / s_group);
int swizzled_s = (s % s_group) * 4 + (t % 4);
return swizzled_t * scale_dim + swizzled_s;
}
// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
//
// Takes separate data and scale cache pointers for K and V.
// Within each KV side, data and scale are separate contiguous regions.
//
// Threading: one CUDA block per token, threads process heads and
// groups of 16 elements within each head.
template <typename scalar_t>
__global__ void reshape_and_cache_nvfp4_kernel(
const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size]
const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size]
uint8_t* __restrict__ key_data_cache, // data region for K
uint8_t* __restrict__ value_data_cache, // data region for V
uint8_t* __restrict__ key_scale_cache, // scale region for K
uint8_t* __restrict__ value_scale_cache, // scale region for V
const int64_t* __restrict__ slot_mapping, // [num_actual_tokens]
const float* __restrict__ k_scale_ptr, // pointer to checkpoint k_scale
const float* __restrict__ v_scale_ptr, // pointer to checkpoint v_scale
const int64_t key_stride, // key.stride(0) in elements
const int64_t value_stride, // value.stride(0) in elements
const int num_heads, const int head_size, const int block_size,
const int64_t data_block_stride, // data cache stride for dim 0
const int64_t data_head_stride, // data cache stride for heads
const int64_t data_block_offset_stride, // data cache stride for tokens
const int64_t scale_block_stride, // scale cache stride for dim 0
const int64_t scale_head_stride, // scale cache stride for heads
const int64_t scale_block_offset_stride // scale cache stride for tokens
) {
using CudaType = typename CUDATypeConverter<scalar_t>::Type;
using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD; // 16 or 8
static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;
const int64_t token_idx = blockIdx.x;
const int64_t slot_idx = slot_mapping[token_idx];
if (slot_idx < 0) return;
const int64_t block_idx = slot_idx / block_size;
const int block_offset = static_cast<int>(slot_idx % block_size);
const int scale_dim = head_size / 16;
const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
const int total_groups = num_heads * groups_per_head;
const int tid = threadIdx.x;
const int num_thread_groups = blockDim.x / THREADS_PER_SF;
const int tg_id = tid / THREADS_PER_SF;
const int tg_lane = tid % THREADS_PER_SF;
// Process both K (kv=0) and V (kv=1)
#pragma unroll
for (int kv = 0; kv < 2; kv++) {
const scalar_t* __restrict__ src = (kv == 0) ? key : value;
const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
uint8_t* __restrict__ data_cache =
(kv == 0) ? key_data_cache : value_data_cache;
uint8_t* __restrict__ sc_cache =
(kv == 0) ? key_scale_cache : value_scale_cache;
// Source pointer for this token (use actual stride, not assumed contiguous)
const CudaType* __restrict__ token_src =
reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;
// Destination bases in data and scale caches for this token's block
uint8_t* __restrict__ data_block =
data_cache + block_idx * data_block_stride;
uint8_t* __restrict__ scale_block =
sc_cache + block_idx * scale_block_stride;
for (int g = tg_id; g < total_groups; g += num_thread_groups) {
const int head = g / groups_per_head;
const int group_in_head = g % groups_per_head;
// Load 16 (or 8) bf16 elements from source
PVec in_vec;
const CudaType* __restrict__ src_ptr =
token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
tg_lane * ELTS;
#pragma unroll
for (int i = 0; i < ELTS / 2; i++) {
in_vec.elts[i] = reinterpret_cast<
const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
}
// Quantize: produces packed fp4 and writes scale factor.
uint8_t sf_val;
uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
in_vec, global_scale, sf_out_ptr);
// Write packed FP4 data to data cache
uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
block_offset * data_block_offset_stride;
#if CVT_FP4_PACK16
{
// 16 elements → 8 bytes (u32x2)
int data_byte_offset = group_in_head * 8;
reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
(uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
}
#else
{
// 8 elements → 4 bytes (uint32_t)
int data_byte_offset =
group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
}
#endif
// Write block scale to scale cache.
// K (kv==0): linear layout (no swizzle).
// V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel.
if (sf_out_ptr != nullptr) {
int scale_idx = group_in_head;
uint8_t* __restrict__ scale_dst;
if (kv == 0) {
scale_dst = scale_block + head * scale_head_stride +
block_offset * scale_block_offset_stride + scale_idx;
} else {
int swizzled_offset =
swizzle_scale_offset(block_offset, scale_idx, scale_dim);
int swizzled_t = swizzled_offset / scale_dim;
int swizzled_s = swizzled_offset % scale_dim;
scale_dst = scale_block + head * scale_head_stride +
swizzled_t * scale_block_offset_stride + swizzled_s;
}
*scale_dst = sf_val;
}
}
}
}
} // namespace vllm
// Non-template entry point callable from cache_kernels.cu.
// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
// Each KV side contains both data and scale:
// page = [K_data | K_scale | V_data | V_scale]
void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
torch::Tensor& k_scale,
torch::Tensor& v_scale) {
int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
int data_dim = head_size / 2;
int scale_dim = head_size / 16;
int full_dim = data_dim + scale_dim;
// key_cache is kv_cache[:, 0] with shape
// [num_blocks, block_size, num_heads, full_dim] in logical order.
// Strides encode the physical layout (HND or NHD).
TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
TORCH_CHECK(key_cache.size(3) == full_dim,
"key_cache last dim must be data_dim + scale_dim, got ",
key_cache.size(3), " expected ", full_dim);
int block_size = key_cache.size(1);
TORCH_CHECK(head_size % 16 == 0,
"head_size must be divisible by 16 for NVFP4 KV cache");
TORCH_CHECK(block_size % 4 == 0,
"block_size must be divisible by 4 for NVFP4 KV cache swizzle");
// Detect physical layout from strides (based on full_dim).
// HND: head stride > block_offset stride.
bool is_hnd = key_cache.stride(2) > key_cache.stride(1);
int64_t data_block_stride = key_cache.stride(0); // page_bytes
int64_t data_head_stride, data_block_offset_stride;
if (is_hnd) {
data_head_stride = (int64_t)block_size * data_dim;
data_block_offset_stride = data_dim;
} else {
data_head_stride = data_dim;
data_block_offset_stride = (int64_t)num_heads * data_dim;
}
// Page layout: [K_data | K_scale | V_data | V_scale]
// Scale follows data within each KV side.
int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
// Scale strides: same page stride, inner strides from layout.
int64_t scale_block_stride = data_block_stride;
int64_t scale_head_stride, scale_block_offset_stride;
if (is_hnd) {
scale_head_stride = (int64_t)block_size * scale_dim;
scale_block_offset_stride = scale_dim;
} else {
scale_head_stride = scale_dim;
scale_block_offset_stride = (int64_t)num_heads * scale_dim;
}
const float* k_scale_ptr = k_scale.data_ptr<float>();
const float* v_scale_ptr = v_scale.data_ptr<float>();
int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
int total_groups = num_heads * groups_per_head;
constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
num_threads = ((num_threads + 31) / 32) * 32;
dim3 grid(num_tokens);
dim3 block(num_threads);
const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_REDUCED_FLOATING_TYPES(
key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
<<<grid, block, 0, stream>>>(
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
key_scale_ptr, value_scale_ptr,
slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
key.stride(0), value.stride(0), num_heads, head_size,
block_size, data_block_stride, data_head_stride,
data_block_offset_stride, scale_block_stride, scale_head_stride,
scale_block_offset_stride);
});
}
+45
View File
@@ -0,0 +1,45 @@
/* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
*
* Vendored from sgl-kernel/include/sgl_flash_kernel_ops.h (commit bcf72ccc).
* Declares the mha_fwd() C++ function signature for CUTLASS FA3 kernels.
* NO MODIFICATIONS from the original (except removing unused macros).
*/
#pragma once
#include <ATen/ATen.h>
#include <ATen/Tensor.h>
#include <torch/library.h>
#include <torch/torch.h>
#include <vector>
#include "sgl_kernel_torch_shim.h"
/*
* From flash-attention (sgl-attn fork)
*/
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
at::Tensor q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
at::Tensor k, // (b_k, s_k, h_k, d) or (total_k, h_k, d) or paged
at::Tensor v, // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) or paged
std::optional<at::Tensor> k_new_, std::optional<at::Tensor> v_new_,
std::optional<at::Tensor> q_v_, // MLA value projection query
std::optional<at::Tensor> out_, std::optional<at::Tensor> cu_seqlens_q_,
std::optional<at::Tensor> cu_seqlens_k_,
std::optional<at::Tensor> cu_seqlens_k_new_,
std::optional<at::Tensor> seqused_q_, std::optional<at::Tensor> seqused_k_,
std::optional<int64_t> max_seqlen_q_, std::optional<int64_t> max_seqlen_k_,
std::optional<at::Tensor> page_table_,
std::optional<at::Tensor> kv_batch_idx_,
std::optional<at::Tensor> leftpad_k_, std::optional<at::Tensor> rotary_cos_,
std::optional<at::Tensor> rotary_sin_,
std::optional<at::Tensor> seqlens_rotary_,
std::optional<at::Tensor> q_descale_, std::optional<at::Tensor> k_descale_,
std::optional<at::Tensor> v_descale_, std::optional<double> softmax_scale_,
bool is_causal, int64_t window_size_left, int64_t window_size_right,
int64_t attention_chunk, double softcap, bool is_rotary_interleaved,
std::optional<at::Tensor> scheduler_metadata_, int64_t num_splits,
std::optional<bool> pack_gqa_, int64_t sm_margin,
std::optional<const at::Tensor>& sinks_);
+121
View File
@@ -0,0 +1,121 @@
/* Adapted from:
* https://github.com/neuralmagic/vllm-flash-attention/blob/90eacc1af2a7c3de62ea249e929ed5faccf38954/csrc/common/pytorch_shim.h
*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
*
* Vendored from sgl-kernel/include/sgl_kernel_torch_shim.h (commit bcf72ccc).
* Provides make_pytorch_shim() template for PyTorch op registration type
* conversion. NO MODIFICATIONS from the original.
*/
#pragma once
#include <torch/library.h>
/**
* Unfortunately, the type signatures of the flash_attn ops are not compatible
* with the PyTorch library bindings. To get around that we use
* `make_pytorch_shim` which creates a lambda that exposes the API using
* PyTorch compatible types to the types, then converts them to the types
* expected by the flash_attn ops. This shims allows us to make minimal changes
* to `flash_api.cpp` making it easier to synchronize with upstream changes.
*
* The `pytorch_library_compatible_type` struct is used to map from the
* flash_attn ops types to a PyTorch library compatible one. The main issues is
* that the following types are not support by PyTorch library bindings:
* - `int`
* - `float`
* - `std::optional<T> &`
* - `std::optional<const at::Tensor> &`
* So we convert them to (respectively):
* - `int64_t`
* - `double`
* - `const std::optional<T>&`
* - `const std::optional<at::Tensor>&`
*/
template <typename T>
struct pytorch_library_compatible_type {
using type = T;
static T convert_from_type(T arg) { return arg; }
};
template <typename T>
using pytorch_library_compatible_type_t =
typename pytorch_library_compatible_type<T>::type;
template <typename T>
T convert_from_pytorch_compatible_type(
pytorch_library_compatible_type_t<T> arg) {
return pytorch_library_compatible_type<T>::convert_from_type(arg);
}
// Map `c10::optional<T> &` -> `const c10::optional<T>&`
// (NOTE: this is bit unsafe but non of the ops in flash_attn mutate
// the optional container)
template <typename T>
struct pytorch_library_compatible_type<c10::optional<T>&> {
using type = const c10::optional<T>&;
static c10::optional<T>& convert_from_type(const c10::optional<T>& arg) {
return const_cast<c10::optional<T>&>(arg);
}
};
// Map `c10::optional<T>` ->
// `c10::optional<pytorch_library_compatible_type_t<T>>`
// (NOTE: tested for `c10::optional<int>` -> `c10::optional<int64_t>`)
template <typename T>
struct pytorch_library_compatible_type<c10::optional<T>> {
using type = c10::optional<pytorch_library_compatible_type_t<T>>;
static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(
c10::optional<T> arg) {
return arg;
}
};
// Map `c10::optional<const at::Tensor>&` -> `const c10::optional<at::Tensor>&`
template <>
struct pytorch_library_compatible_type<c10::optional<const at::Tensor>&> {
using type = const c10::optional<at::Tensor>&;
static c10::optional<const at::Tensor>& convert_from_type(
const c10::optional<at::Tensor>& arg) {
return const_cast<c10::optional<const at::Tensor>&>(
reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
}
};
// Map `int` -> `int64_t`
template <>
struct pytorch_library_compatible_type<int> {
using type = int64_t;
static int convert_from_type(int64_t arg) {
TORCH_CHECK(arg <= std::numeric_limits<int>::max(),
"int64_t value is too large to be converted to int");
TORCH_CHECK(arg >= std::numeric_limits<int>::min(),
"int64_t value is too small to be converted to int");
return arg;
}
};
// Map `float` -> `double`
template <>
struct pytorch_library_compatible_type<float> {
using type = double;
static float convert_from_type(double arg) {
TORCH_CHECK(std::abs(arg) <= std::numeric_limits<float>::max(),
"double value is too large to be converted to float");
return arg;
}
};
//
// Shim Utils
//
template <typename Ret, typename... Args>
auto make_pytorch_shim(Ret (*fun)(Args... args)) {
return [fun](pytorch_library_compatible_type_t<Args>... args) {
return fun(convert_from_pytorch_compatible_type<Args>(args)...);
};
}
+2 -1
View File
@@ -173,7 +173,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements/dev.txt && \
uv pip install -r requirements/lint.txt && \
uv pip install -r requirements/test/cpu.txt && \
pre-commit install --hook-type pre-commit --hook-type commit-msg
ENTRYPOINT ["bash"]
+3 -10
View File
@@ -106,6 +106,7 @@ Priority is **1 = highest** (tried first).
| 2 | `FLASH_ATTN` |
| 3 | `TRITON_ATTN` |
| 4 | `FLEX_ATTENTION` |
| 5 | `TURBOQUANT` |
**Ampere/Hopper (SM 8.x-9.x):**
@@ -115,6 +116,7 @@ Priority is **1 = highest** (tried first).
| 2 | `FLASHINFER` |
| 3 | `TRITON_ATTN` |
| 4 | `FLEX_ATTENTION` |
| 5 | `TURBOQUANT` |
### MLA Attention (DeepSeek-style)
@@ -130,16 +132,6 @@ Priority is **1 = highest** (tried first).
| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
| 7 | `FLASHMLA_SPARSE` |
**Ampere/Hopper (SM 8.x-9.x):**
| Priority | Backend |
| -------- | ------- |
| 1 | `FLASH_ATTN_MLA` |
| 2 | `FLASHMLA` |
| 3 | `FLASHINFER_MLA` |
| 4 | `TRITON_MLA` |
| 5 | `FLASHMLA_SPARSE` |
> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
>
> **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -207,6 +199,7 @@ configuration.
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
| `CUTLASS_FA3_MLA_SPARSE` | bf16 | `auto` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
+2 -2
View File
@@ -42,7 +42,7 @@ These are documented under [Inferencing and Serving -> Production Metrics](../us
### Grafana Dashboard
vLLM also provides [a reference example](../../examples/online_serving/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
vLLM also provides [a reference example](../../examples/observability/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
The subset of metrics exposed in the Grafana dashboard gives us an indication of which metrics are especially important:
@@ -657,7 +657,7 @@ vLLM has support for OpenTelemetry tracing:
- Added by <https://github.com/vllm-project/vllm/pull/4687> and reinstated by <https://github.com/vllm-project/vllm/pull/20372>
- Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces`
- [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/)
- [User-facing docs](../../examples/online_serving/opentelemetry/README.md)
- [User-facing docs](../../examples/observability/opentelemetry/README.md)
- [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f)
- [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview)
+9 -6
View File
@@ -14,6 +14,7 @@ Sorted alphabetically by GitHub handle:
- [@aarnphm](https://github.com/aarnphm): Structured output
- [@alexm-redhat](https://github.com/alexm-redhat): Performance
- [@ApostaC](https://github.com/ApostaC): Connectors, offloading
- [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
- [@benchislett](https://github.com/benchislett): Engine core and spec decode
- [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
- [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
@@ -31,6 +32,7 @@ Sorted alphabetically by GitHub handle:
- [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance
- [@luccafong](https://github.com/luccafong): Llama models, speculative decoding, distributed
- [@markmc](https://github.com/markmc): Observability
- [@MatthewBonanni](https://github.com/MatthewBonanni): Kernels and performance
- [@mgoin](https://github.com/mgoin): Quantization and performance
- [@NickLucche](https://github.com/NickLucche): KV connector
- [@njhill](https://github.com/njhill): Distributed, API server, engine core
@@ -41,6 +43,7 @@ Sorted alphabetically by GitHub handle:
- [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg
- [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support
- [@russellb](https://github.com/russellb): Structured output, engine core, security
- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser
- [@sighingnow](https://github.com/sighingnow): Qwen models, new model support
- [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community
- [@tdoublep](https://github.com/tdoublep): State space models
@@ -86,7 +89,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints
- @robertgshaw2-redhat, @njhill, @russellb
- ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg
- @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni
- KV Connector: Connector interface and implementation for KV cache offload and transfer
- @robertgshaw2-redhat, @njhill, @KuntaiDu, @NickLucche, @ApostaC
- Distributed, Parallelism, Process Management: Process launchers managing each worker, and assign them to the right DP/TP/PP/EP ranks
@@ -105,7 +108,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- Custom Layers: Utility layers in vLLM such as rotary embedding and rms norms
- @ProExpertProg
- Attention: Attention interface for paged attention
- @WoosukKwon, @LucasWilkinson, @heheda12345
- @WoosukKwon, @LucasWilkinson, @heheda12345, @MatthewBonanni
- FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB
- @tlrmchlsmth
- Quantization: Various quantization config, weight loading, and kernel.
@@ -119,7 +122,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
- State space models: The state space models implementation in vLLM
- @tdoublep, @tlrmchlsmth
- Reasoning and tool calling parsers
- @chaunceyjiang, @aarnphm
- @chaunceyjiang, @aarnphm, @sfeng33, @bbrowning
### Entrypoints
@@ -133,7 +136,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
### Features
- Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP.
- @WoosukKwon, @benchislett, @luccafong
- @WoosukKwon, @benchislett, @luccafong, @MatthewBonanni
- Structured Output: The structured output implementation
- @russellb, @aarnphm
- RL: The RL related features such as collective rpc, sleep mode, etc.
@@ -153,8 +156,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
### External Kernels Integration
- FlashAttention: @LucasWilkinson
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon
- FlashAttention: @LucasWilkinson, @MatthewBonanni
- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni
- Blackwell Kernels: @mgoin, @yewentao256
- DeepEP/DeepGEMM: @mgoin, @yewentao256
+8
View File
@@ -400,6 +400,7 @@ th {
| `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ |
| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ |
| `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ |
| `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
| `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ |
| `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ |
@@ -554,6 +555,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |
| `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>E+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |
| `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I<sup>+</sup> + V + A<sup>*</sup> | `google/gemma-4-E2B-it`, etc. | | ✅︎ |
| `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ |
| `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
@@ -633,6 +635,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
<sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
<sup>+</sup> Multiple items can be inputted per text prompt for this modality.
<sup>*</sup> Only specific variants of the model support this modality (see notes below).</br>
!!! note
`Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
@@ -643,6 +646,11 @@ Some models are supported only via the [Transformers modeling backend](#transfor
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
- There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
!!! note
For `Gemma4ForConditionalGeneration`:
- audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants.
- The model does not ingest videos directly. However, vLLMs Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model.
!!! note
For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
@@ -74,8 +74,8 @@ percli apply -f perses/performance_statistics.yaml
For detailed deployment instructions and platform-specific options, see:
- **[Grafana Documentation](./grafana)** - JSON dashboards, operator usage, manual import
- **[Perses Documentation](./perses)** - YAML specs, CLI usage, operator wrapping
- **[Grafana Documentation](grafana)** - JSON dashboards, operator usage, manual import
- **[Perses Documentation](perses)** - YAML specs, CLI usage, operator wrapping
## Contributing
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import vllm.config
from tests.compile.backend import TestBackend
from vllm.compilation.passes.vllm_inductor_pass import (
VllmFusionPatternMatcherPass,
VllmPatternMatcherPass,
VllmPatternReplacement,
)
from vllm.config import CompilationConfig, CompilationMode, VllmConfig
from vllm.platforms import current_platform
class ReluToAbsPattern(VllmPatternReplacement):
"""Replaces relu(x) with abs(x) — a minimal test fixture."""
@property
def pattern(self):
def _pattern(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.relu.default(x)
return _pattern
@property
def replacement(self):
def _replacement(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.abs.default(x)
return _replacement
def get_inputs(self) -> list[torch.Tensor]:
return [self.empty_fp32(4)]
class ExpToSqrtPattern(VllmPatternReplacement):
"""A second distinct pattern type — used to test uuid differentiation."""
@property
def pattern(self):
def _pattern(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.exp.default(x)
return _pattern
@property
def replacement(self):
def _replacement(x: torch.Tensor) -> torch.Tensor:
return torch.ops.aten.sqrt.default(x)
return _replacement
def get_inputs(self) -> list[torch.Tensor]:
return [self.empty_fp32(4)]
class ReluFusionPass(VllmFusionPatternMatcherPass):
def __init__(self, config: VllmConfig) -> None:
super().__init__(config, "test_relu_fusion")
self.register(ReluToAbsPattern())
class TwoPatternFusionPass(VllmFusionPatternMatcherPass):
def __init__(self, config: VllmConfig) -> None:
super().__init__(config, "test_two_pattern_fusion")
self.register(ReluToAbsPattern())
self.register(ExpToSqrtPattern())
@pytest.fixture
def vllm_config():
return VllmConfig(
compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE),
)
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
def test_register_tracks_patterns(vllm_config):
"""register() appends each VllmPatternReplacement to _pattern_replacements."""
with vllm.config.set_current_vllm_config(vllm_config):
single = ReluFusionPass(vllm_config)
two = TwoPatternFusionPass(vllm_config)
assert len(single._pattern_replacements) == 1
assert len(two._pattern_replacements) == 2
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
def test_uuid_stable(vllm_config):
"""Two instances of the same pass class produce identical uuids."""
with vllm.config.set_current_vllm_config(vllm_config):
p1 = ReluFusionPass(vllm_config)
p2 = ReluFusionPass(vllm_config)
p3 = TwoPatternFusionPass(vllm_config)
assert p1.uuid() == p2.uuid()
assert p1.uuid() != p3.uuid()
assert p2.uuid() != p3.uuid()
@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
@pytest.mark.parametrize("N", [1, 2, 4])
def test_matched_count_and_match_table(vllm_config, N):
"""matched_count and match_table reflect the number of matched patterns."""
class Model(torch.nn.Module):
def forward(self, *inputs):
# N independent relus
return sum(torch.relu(x) for x in inputs)
with vllm.config.set_current_vllm_config(vllm_config):
torch.set_default_device("cuda")
torch.set_default_dtype(torch.float32)
fusion_pass = ReluFusionPass(vllm_config)
backend = TestBackend(fusion_pass)
model = torch.compile(Model(), backend=backend)
inputs = [torch.rand(8) for _ in range(N)]
model(*inputs)
assert fusion_pass.matched_count == N
assert VllmPatternMatcherPass.match_table["test_relu_fusion"] >= N
@@ -41,6 +41,7 @@ def create_mock_parallel_config(
config.rank = rank
config.world_size = world_size
config.data_parallel_rank = dp_rank
config.data_parallel_index = dp_rank
return config
@@ -283,6 +284,7 @@ def inference_receive_tensor(
parallel_config.rank = 0
parallel_config.world_size = 1
parallel_config.data_parallel_rank = 0
parallel_config.data_parallel_index = 0
engine = NCCLWeightTransferEngine(config, parallel_config)
@@ -666,6 +668,7 @@ def inference_receive_ipc_tensor(
parallel_config.rank = 0
parallel_config.world_size = 1
parallel_config.data_parallel_rank = 0
parallel_config.data_parallel_index = 0
engine = IPCWeightTransferEngine(config, parallel_config)
+99 -21
View File
@@ -10,7 +10,7 @@ from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
from vllm.utils.torch_utils import nvfp4_kv_cache_split_views, set_random_seed
COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
DTYPES = [torch.bfloat16, torch.float]
@@ -172,7 +172,7 @@ def test_reshape_and_cache(
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE + ["nvfp4"])
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@@ -202,6 +202,25 @@ def test_reshape_and_cache_flash(
if kv_scale_type == "attn_head" and implementation != "cuda":
pytest.skip("Only CUDA implementation supports attn_head scaling.")
if kv_cache_dtype == "nvfp4":
if not current_platform.has_device_capability(100):
pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).")
if implementation != "cuda":
pytest.skip("NVFP4 only supports CUDA implementation.")
if kv_scale_type != "tensor":
pytest.skip("NVFP4 only supports per-tensor scaling.")
if head_size % 16 != 0:
pytest.skip("NVFP4 requires head_size divisible by 16.")
if (head_size // 16) % 4 != 0:
pytest.skip(
"NVFP4 requires (head_size // 16) divisible by 4 "
"for 4x4 block scale swizzle."
)
if block_size % 4 != 0:
pytest.skip("NVFP4 requires block_size divisible by 4.")
if dtype not in (torch.float16, torch.bfloat16):
pytest.skip("NVFP4 quantization only supports fp16/bf16 input.")
# fp8 conversion requires continugous memory buffer. Reduce the number of
# blocks and tokens to consume less memory.
num_tokens = num_tokens // 2
@@ -229,7 +248,23 @@ def test_reshape_and_cache_flash(
del key_caches
del value_caches
if kv_scale_type == "tensor":
# For nvfp4, the factory returns kv[:, 0] and kv[:, 1] like all dtypes.
# Split views are still needed for dequant verification.
key_scale_cache = None
value_scale_cache = None
nvfp4_key_data = None
nvfp4_value_data = None
if kv_cache_dtype == "nvfp4":
(nvfp4_key_data,), (key_scale_cache,) = nvfp4_kv_cache_split_views(key_cache)
(nvfp4_value_data,), (value_scale_cache,) = nvfp4_kv_cache_split_views(
value_cache
)
if kv_cache_dtype == "nvfp4":
# Global scale = amax / 448 (per-tensor)
k_scale = (key.abs().amax() / 448.0).to(torch.float32)
v_scale = (value.abs().amax() / 448.0).to(torch.float32)
elif kv_scale_type == "tensor":
k_scale = (key.amax() / 64.0).to(torch.float32)
v_scale = (value.amax() / 64.0).to(torch.float32)
else: # "attn_head"
@@ -240,8 +275,9 @@ def test_reshape_and_cache_flash(
y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
return y.contiguous()
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
if kv_cache_dtype != "nvfp4":
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
def convert_fp8_local(output, input, scale, kv_dtype):
fp8_input = input.view(current_platform.fp8_dtype())
@@ -257,7 +293,7 @@ def test_reshape_and_cache_flash(
result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
output.copy_(result)
# Clone the KV caches.
# Clone the KV caches (for non-nvfp4, used as reference baseline).
if kv_cache_dtype == "fp8":
cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
@@ -265,25 +301,27 @@ def test_reshape_and_cache_flash(
convert_fp8_local(
cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype
)
else:
elif kv_cache_dtype != "nvfp4":
cloned_key_cache = key_cache_compact.clone()
cloned_value_cache = value_cache_compact.clone()
# Call the reshape_and_cache kernel.
if implementation == "cuda":
opcheck(
torch.ops._C_cache_ops.reshape_and_cache_flash,
(
key,
value,
key_cache,
value_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
v_scale,
),
cond=(head_size == HEAD_SIZES[0]),
)
if kv_cache_dtype != "nvfp4":
opcheck(
torch.ops._C_cache_ops.reshape_and_cache_flash,
(
key,
value,
key_cache,
value_cache,
slot_mapping,
kv_cache_dtype,
k_scale,
v_scale,
),
cond=(head_size == HEAD_SIZES[0]),
)
ops.reshape_and_cache_flash(
key,
value,
@@ -309,6 +347,46 @@ def test_reshape_and_cache_flash(
k_scale,
v_scale,
)
if kv_cache_dtype == "nvfp4":
# Verify NVFP4 by dequantizing the entire cache and comparing
# the written positions against original bf16 values.
# Same pattern as FP8: dequant whole cache, then extract and compare.
from tests.kernels.quantization.nvfp4_utils import (
dequant_nvfp4_kv_cache,
)
def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
# data_cache: [N, T, H, data_dim] NHD (contiguous inner dims)
# scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims)
# Permute to HND layout for the dequant utility.
data_hnd = data_cache.permute(0, 2, 1, 3)
scale_hnd = scale_cache.permute(0, 2, 1, 3)
result_hnd = dequant_nvfp4_kv_cache(
data_hnd, scale_hnd, global_scale, head_size, block_size
)
return result_hnd.permute(0, 2, 1, 3) # back to [N, T, H, D]
result_key_cache = dequant_nvfp4_cache_nhd(
nvfp4_key_data, key_scale_cache, k_scale.item()
)
result_value_cache = dequant_nvfp4_cache_nhd(
nvfp4_value_data, value_scale_cache, v_scale.item()
)
# Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
num_slots = num_blocks * block_size
result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size)
result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size)
torch.testing.assert_close(
result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
)
torch.testing.assert_close(
result_value_flat[slot_mapping], value.float(), atol=1.5, rtol=0.5
)
return
key_cache_compact = permute_and_compact(key_cache)
value_cache_compact = permute_and_compact(value_cache)
File diff suppressed because it is too large Load Diff
+48 -4
View File
@@ -13,6 +13,11 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
selective_state_update,
)
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupSpec,
MambaSpec,
)
try:
import flashinfer.mamba # noqa: F401
@@ -22,22 +27,40 @@ except ImportError:
HAS_FLASHINFER = False
def _kv_cache_config_with_ssu(mamba_type: str = "mamba2") -> KVCacheConfig:
spec = MambaSpec(
block_size=16,
shapes=((16, 64),),
dtypes=(torch.float16,),
mamba_type=mamba_type,
)
return KVCacheConfig(
num_blocks=1,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(layer_names=["l0"], kv_cache_spec=spec)],
)
def test_default_backend_is_triton():
initialize_mamba_ssu_backend(MambaConfig())
initialize_mamba_ssu_backend(MambaConfig(), _kv_cache_config_with_ssu())
backend = get_mamba_ssu_backend()
assert isinstance(backend, TritonSSUBackend)
assert backend.name == "triton"
def test_explicit_triton_backend():
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
)
backend = get_mamba_ssu_backend()
assert isinstance(backend, TritonSSUBackend)
@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
def test_flashinfer_backend_init():
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.FLASHINFER), _kv_cache_config_with_ssu()
)
backend = get_mamba_ssu_backend()
assert isinstance(backend, FlashInferSSUBackend)
assert backend.name == "flashinfer"
@@ -53,6 +76,25 @@ def test_uninitialized_backend_raises():
mod._mamba_ssu_backend = old
@pytest.mark.parametrize(
"mamba_type", ["linear_attention", "gdn_attention", "short_conv"]
)
def test_init_is_noop_for_non_ssu_mamba_type(mamba_type):
import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
old = mod._mamba_ssu_backend
mod._mamba_ssu_backend = None
try:
initialize_mamba_ssu_backend(
MambaConfig(), _kv_cache_config_with_ssu(mamba_type)
)
assert mod._mamba_ssu_backend is None
with pytest.raises(RuntimeError, match="not been initialized"):
get_mamba_ssu_backend()
finally:
mod._mamba_ssu_backend = old
@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
def test_flashinfer_import_error():
with pytest.raises(ImportError, match="FlashInfer is required"):
@@ -61,7 +103,9 @@ def test_flashinfer_import_error():
def test_triton_basic_call():
set_random_seed(0)
initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
initialize_mamba_ssu_backend(
MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
)
device = "cuda"
batch_size = 2
dim = 64
+8
View File
@@ -465,6 +465,14 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
if config.enable_eplb and config.ep_size == 1:
return False, "EPLB only works with EP+DP"
# Disable fp4 tests until flashinfer is updated or the Dockerfile is
# modified to install cublasLt.h. See #39525.
if (
config.quantization == "modelopt_fp4"
and current_platform.is_device_capability_family(100)
):
return False, "Temporarily skip until #39525 is resolved"
return True, None
+54
View File
@@ -88,6 +88,60 @@ def break_fp4_bytes(a, dtype):
return values.reshape(m, n * 2).to(dtype=dtype)
def dequant_nvfp4_kv_cache(
fp4_data: torch.Tensor,
block_scale: torch.Tensor,
global_scale: float,
head_size: int,
block_size: int,
) -> torch.Tensor:
"""Dequantize an NVFP4 KV cache with 4x4-swizzled block scales.
The input must be in HND layout so that the last two dims are
(block_size, last_dim). For NHD caches, permute to HND first.
Args:
fp4_data: [..., num_heads, block_size, head_size//2] uint8 packed fp4.
block_scale: [..., num_heads, block_size, head_size//16] fp8 block
scales (as uint8 or float8_e4m3fn).
global_scale: checkpoint dequant scale (k_scale or v_scale).
head_size: head dimension.
block_size: page size.
Returns:
[..., num_heads, block_size, head_size] float32.
"""
data_dim = head_size // 2
scale_dim = head_size // 16
fp4_packed = fp4_data
sf_swizzled = block_scale.view(torch.uint8)
# Unswizzle 4x4 block scales on (block_size, scale_dim) plane.
# [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S]
batch_shape = sf_swizzled.shape[:-2]
T, S = block_size, scale_dim
sg = S // 4
sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4)
ndim = sf_reshape.ndim
# Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg)
perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32)
# Unpack fp4
shape = fp4_packed.shape # [..., T, data_dim]
fp4_flat = fp4_packed.reshape(-1, data_dim)
fp4_vals = break_fp4_bytes(fp4_flat, torch.float32)
fp4_vals = fp4_vals.reshape(*shape[:-1], head_size)
# Dequant: fp4_val * block_scale * global_scale per 16-element group
return (
fp4_vals.reshape(*shape[:-1], scale_dim, 16)
* (sf_f32 * global_scale).unsqueeze(-1)
).reshape(*shape[:-1], head_size)
def get_nvfp4_global_scale(a: torch.Tensor):
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32)
@@ -46,7 +46,7 @@ AITER_MODEL_LIST = [
),
pytest.param(
"openai-community/gpt2", # gpt2
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
marks=[pytest.mark.core_model],
),
pytest.param("Milos/slovak-gpt-j-405M"), # gptj
pytest.param("bigcode/tiny_starcoder_py"), # gpt_bigcode
@@ -143,11 +143,6 @@ def test_models(
# in parts of the operators
pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
if current_platform.is_cpu() and model in ("openai-community/gpt2",):
# These models are sensitive to the rounding error
# Fuse ops to reduce rounding
monkeypatch.setenv("VLLM_CPU_CI_ENV", "0")
with hf_runner(model) as hf_model:
hf_outputs = hf_model.generate_greedy_logprobs_limit(
example_prompts, max_tokens, num_logprobs
@@ -140,88 +140,3 @@ def test_audio_token_count_matches_hf_processor_math():
_count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 0) == 1499
)
assert _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 1) == 375
def test_audio_feature_pipeline_matches_hf_small_config():
from transformers.models.audioflamingo3 import (
modeling_audioflamingo3 as hf_audioflamingo3_modeling,
)
from transformers.models.audioflamingo3.configuration_audioflamingo3 import (
AudioFlamingo3Config,
)
from vllm.model_executor.models.audioflamingo3 import (
AudioFlamingo3Encoder,
AudioFlamingo3MultiModalProjector,
_build_audio_encoder_attention_mask,
_flatten_valid_audio_embeddings,
)
text_config = {
"model_type": "qwen2",
"intermediate_size": 64,
"initializer_range": 0.02,
"hidden_size": 32,
"max_position_embeddings": 1024,
"num_hidden_layers": 2,
"num_attention_heads": 4,
"num_key_value_heads": 2,
"vocab_size": 128,
"pad_token_id": 1,
"use_mrope": False,
}
audio_config = {
"hidden_size": 16,
"num_attention_heads": 4,
"intermediate_size": 32,
"num_hidden_layers": 2,
"num_mel_bins": 80,
"max_source_positions": 1500,
"dropout": 0.0,
"attention_dropout": 0.0,
"activation_dropout": 0.0,
"encoder_layerdrop": 0.0,
}
torch.manual_seed(0)
config = AudioFlamingo3Config(
text_config=text_config,
audio_config=audio_config,
audio_token_id=0,
)
hf_model = hf_audioflamingo3_modeling.AudioFlamingo3ForConditionalGeneration(
config
).eval()
vllm_encoder = AudioFlamingo3Encoder(config.audio_config).eval()
vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
vllm_projector = AudioFlamingo3MultiModalProjector(config).eval()
vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
input_features = torch.randn(3, 80, 3000)
feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
feature_attention_mask[0, :3000] = True
feature_attention_mask[1, :2500] = True
feature_attention_mask[2, :1500] = True
hf_output = hf_model.get_audio_features(
input_features,
feature_attention_mask,
return_dict=True,
).pooler_output
vllm_attention_mask = _build_audio_encoder_attention_mask(
feature_attention_mask,
dtype=vllm_encoder.conv1.weight.dtype,
device=vllm_encoder.conv1.weight.device,
)
vllm_hidden_states = vllm_encoder(
input_features,
attention_mask=vllm_attention_mask,
)
vllm_output, _ = _flatten_valid_audio_embeddings(
vllm_projector(vllm_hidden_states),
feature_attention_mask,
)
torch.testing.assert_close(vllm_output, hf_output)
-1
View File
@@ -38,6 +38,5 @@ def test_model_experts_int8_startup(
dtype=dtype,
enforce_eager=True,
quantization="experts_int8",
allow_deprecated_quantization=True,
) as vllm_model:
vllm_model.generate_greedy(example_prompts, max_tokens)
@@ -64,9 +64,6 @@ class TestQwen3xmlToolParser(ToolParserTests):
"test_empty_arguments": "Qwen3XML streaming has systematic issues",
"test_surrounding_text": "Qwen3XML streaming has systematic issues",
"test_escaped_strings": "Qwen3XML streaming has systematic issues",
"test_malformed_input": (
"Qwen3XML parser is lenient with malformed input"
),
"test_streaming_reconstruction": (
"Qwen3XML streaming reconstruction has known issues"
),
@@ -0,0 +1,968 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Backend integration tests for CUTLASS FA3 sparse MLA attention.
Tests verify:
- Backend class properties
- Metadata builder (decode, prefill, mixed, topk clipping)
- KV cache write/read consistency
- Backend registration and selection
"""
import pytest
import torch
from vllm.v1.attention.ops.cutlass_fa3 import is_cutlass_fa3_available
pytestmark = pytest.mark.skipif(
not is_cutlass_fa3_available(),
reason="CUTLASS FA3 not available (requires CUDA >= 12.4, SM90)",
)
# ─── TEST 2.1: Backend Class Properties ──────────────────────────────
def test_backend_class_properties():
"""Verify CutlassFA3MLASparseBackend class attributes."""
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseBackend,
)
assert CutlassFA3MLASparseBackend.get_name() == "CUTLASS_FA3_MLA_SPARSE"
assert CutlassFA3MLASparseBackend.is_mla() is True
assert CutlassFA3MLASparseBackend.is_sparse() is True
assert CutlassFA3MLASparseBackend.get_supported_head_sizes() == [576]
assert CutlassFA3MLASparseBackend.supported_kv_cache_dtypes == ["auto"]
assert CutlassFA3MLASparseBackend.get_supported_kernel_block_sizes() == [64]
def test_backend_compute_capability():
"""Verify SM90-only support."""
from vllm.platforms.interface import DeviceCapability
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseBackend,
)
assert CutlassFA3MLASparseBackend.supports_compute_capability(
DeviceCapability(major=9, minor=0)
)
assert not CutlassFA3MLASparseBackend.supports_compute_capability(
DeviceCapability(major=8, minor=0)
)
assert not CutlassFA3MLASparseBackend.supports_compute_capability(
DeviceCapability(major=10, minor=0)
)
def test_backend_kv_cache_shape():
"""Verify KV cache shape for BF16 format."""
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseBackend,
)
shape = CutlassFA3MLASparseBackend.get_kv_cache_shape(
num_blocks=100,
block_size=64,
num_kv_heads=1,
head_size=576,
cache_dtype_str="auto",
)
assert shape == (100, 64, 576)
# ─── TEST 2.2: Backend Registration ──────────────────────────────────
def test_backend_enum_registered():
"""Verify CUTLASS_FA3_MLA_SPARSE is in the backend enum."""
from vllm.v1.attention.backends.registry import AttentionBackendEnum
assert hasattr(AttentionBackendEnum, "CUTLASS_FA3_MLA_SPARSE")
backend_enum = AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE
assert "cutlass_fa3_sparse" in backend_enum.get_path()
def test_backend_class_loadable():
"""Verify the backend class can be loaded from the enum."""
from vllm.v1.attention.backends.registry import AttentionBackendEnum
backend_cls = AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE.get_class()
assert backend_cls.get_name() == "CUTLASS_FA3_MLA_SPARSE"
# ─── TEST 2.3: KV Cache Write/Read ───────────────────────────────────
def test_kv_cache_write_read_consistency():
"""Verify do_kv_cache_update writes match what forward_mqa would read."""
device = "cuda"
num_blocks = 4
block_size = 64
head_size = 576
kv_lora_rank = 512
qk_rope_head_dim = 64
# Create BF16 cache
cache = torch.zeros(
num_blocks, block_size, head_size, dtype=torch.bfloat16, device=device
)
# Write known values
T = 3
kv_c_normed = torch.randn(T, kv_lora_rank, dtype=torch.bfloat16, device=device)
k_pe = torch.randn(T, 1, qk_rope_head_dim, dtype=torch.bfloat16, device=device)
slot_mapping = torch.tensor([0, 1, 2], dtype=torch.int64, device=device)
k_scale = torch.ones(1, dtype=torch.float32, device=device)
from vllm import _custom_ops as ops
ops.concat_and_cache_mla(
kv_c_normed,
k_pe.squeeze(1),
cache,
slot_mapping,
kv_cache_dtype="auto",
scale=k_scale,
)
# Read back via flatten + split (same as forward_mqa does)
S = num_blocks * block_size
kv_flat = cache.reshape(S, head_size)
c_kv_read = kv_flat[:T, :kv_lora_rank]
k_rope_read = kv_flat[:T, kv_lora_rank:]
# Verify consistency
torch.testing.assert_close(c_kv_read, kv_c_normed, rtol=1e-3, atol=1e-3)
torch.testing.assert_close(k_rope_read, k_pe.squeeze(1), rtol=1e-3, atol=1e-3)
def test_kv_cache_dtype_auto():
"""Verify kv_cache_dtype='auto' uses BF16 direct copy."""
device = "cuda"
cache = torch.zeros(1, 64, 576, dtype=torch.bfloat16, device=device)
kv_c = torch.randn(1, 512, dtype=torch.bfloat16, device=device)
k_pe = torch.randn(1, 1, 64, dtype=torch.bfloat16, device=device)
slot_mapping = torch.tensor([0], dtype=torch.int64, device=device)
k_scale = torch.ones(1, dtype=torch.float32, device=device)
from vllm import _custom_ops as ops
ops.concat_and_cache_mla(
kv_c, k_pe.squeeze(1), cache, slot_mapping, kv_cache_dtype="auto", scale=k_scale
)
assert cache.dtype == torch.bfloat16
# ─── TEST 2.4: Edge Cases ────────────────────────────────────────────
def test_empty_kv_cache():
"""Verify do_kv_cache_update handles empty cache gracefully."""
kv_cache = torch.empty(0, device="cuda")
# Should return without error (numel() == 0 check)
# We call the static method from parent class directly
from vllm.v1.attention.backend import SparseMLAAttentionImpl
SparseMLAAttentionImpl.do_kv_cache_update(
None,
kv_c_normed=torch.empty(0),
k_pe=torch.empty(0),
kv_cache=kv_cache,
slot_mapping=torch.empty(0),
kv_cache_dtype="auto",
k_scale=torch.ones(1),
)
# ─── TEST 2.5: Valid Counts from Index Conversion ───────────────────
def test_triton_convert_valid_counts():
"""Verify triton_convert_req_index_to_global_index with return_valid_counts.
This tests the core fix mechanism: the Triton kernel atomically counts
valid (non -1) entries per row while converting indices.
"""
from vllm.v1.attention.backends.mla.sparse_utils import (
triton_convert_req_index_to_global_index,
)
device = "cuda"
T = 4
topk = 128
num_blocks = 16
block_size = 64
req_id = torch.zeros(T, dtype=torch.int32, device=device)
block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
0
) # [1, num_blocks]
# Create topk_indices with varying valid entries per token
topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
expected_valid = [1, 10, 50, 100]
for i in range(T):
nv = expected_valid[i]
# Use indices within the valid range
topk_indices[i, :nv] = torch.randint(
0,
num_blocks * block_size,
(nv,),
dtype=torch.int32,
device=device,
)
global_idx, valid_counts = triton_convert_req_index_to_global_index(
req_id,
block_table,
topk_indices,
BLOCK_SIZE=block_size,
NUM_TOPK_TOKENS=topk,
return_valid_counts=True,
)
# Verify valid counts match expected
for i in range(T):
assert valid_counts[i].item() == expected_valid[i], (
f"Token {i}: expected {expected_valid[i]} valid, "
f"got {valid_counts[i].item()}"
)
# Verify -1 propagation
for i in range(T):
nv = expected_valid[i]
# Entries beyond valid should be -1
assert (global_idx[i, nv:] == -1).all(), (
f"Token {i}: entries beyond valid count should be -1"
)
# ─── TEST 2.6: Prefill Metadata Correctness ─────────────────────────
def test_prefill_cache_seqlens_vs_valid_counts():
"""Verify metadata cache_seqlens = min(seq_len, topk) and that the
forward_mqa fix overrides with valid_counts.
The metadata builder computes cache_seqlens as min(seq_len, topk).
For prefill tokens, this can exceed the actual valid topk entries.
The fix in forward_mqa uses valid_counts instead.
"""
import numpy as np
device = "cuda"
# Simulate a prefill batch: 1 request, 4 tokens, seq_len=4
num_reqs = 1
T = 4
topk = 2048
seq_len = 4
# The metadata builder's logic (simplified):
starts = np.array([0, T], dtype=np.int32)
seg_lens = np.diff(starts) # [4]
seq_lens_np = np.array([seq_len], dtype=np.int32)
per_tok_seqlens = np.minimum(np.repeat(seq_lens_np, seg_lens), topk) # [4, 4, 4, 4]
# This is what the metadata builder produces:
assert all(per_tok_seqlens == 4), (
"Metadata cache_seqlens should be min(seq_len, topk) = 4"
)
# But the actual valid entries per token (with causal masking):
# Token 0: 1 valid entry, Token 1: 2, Token 2: 3, Token 3: 4
expected_valid = [1, 2, 3, 4]
# The fix in forward_mqa computes valid_counts from the page_table
# and uses those as cache_seqlens. Verify the fix produces correct
# valid counts:
from vllm.v1.attention.backends.mla.sparse_utils import (
triton_convert_req_index_to_global_index,
)
req_id = torch.zeros(T, dtype=torch.int32, device=device)
block_table = torch.arange(32, dtype=torch.int32, device=device).unsqueeze(0)
topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
for i in range(T):
nv = expected_valid[i]
topk_indices[i, :nv] = torch.arange(nv, dtype=torch.int32, device=device)
_, valid_counts = triton_convert_req_index_to_global_index(
req_id,
block_table,
topk_indices,
BLOCK_SIZE=64,
NUM_TOPK_TOKENS=topk,
return_valid_counts=True,
)
for i in range(T):
assert valid_counts[i].item() == expected_valid[i], (
f"Token {i}: valid_counts should be {expected_valid[i]}, "
f"got {valid_counts[i].item()}"
)
# ─── TEST 2.7: Clamp -1 to 0 Safety ─────────────────────────────────
def test_global_idx_clamp_safety():
"""Verify clamping -1 page indices to 0 prevents OOB access."""
device = "cuda"
# Create a page_table with -1 entries
page_table = torch.tensor(
[[5, 10, -1, -1], [3, -1, -1, -1]],
dtype=torch.int32,
device=device,
)
# Clamp -1 to 0
clamped = page_table.clamp(min=0)
# Verify
expected = torch.tensor(
[[5, 10, 0, 0], [3, 0, 0, 0]],
dtype=torch.int32,
device=device,
)
assert torch.equal(clamped, expected), (
f"Clamped page_table doesn't match expected: {clamped} vs {expected}"
)
# ─── TEST 2.8: In-place clamp correctness ───────────────────────────
def test_inplace_clamp_no_negative_indices():
"""Verify in-place clamp_(min=0) on global_idx leaves no -1 entries.
The review-fixed code uses clamp_() (in-place) instead of clamp()
to avoid unnecessary tensor allocations during CUDA graph capture.
"""
device = "cuda"
# Create a global_idx tensor with -1 entries
global_idx = torch.tensor(
[[100, 200, -1, -1, -1], [50, -1, -1, -1, -1]],
dtype=torch.int32,
device=device,
)
# In-place clamp
global_idx.clamp_(min=0)
# Verify no -1 entries remain
assert (global_idx >= 0).all(), (
f"In-place clamp should remove all -1 entries: {global_idx}"
)
# Verify valid entries are preserved
assert global_idx[0, 0].item() == 100
assert global_idx[0, 1].item() == 200
assert global_idx[1, 0].item() == 50
# ─── TEST 2.9: Full fix flow with index conversion ──────────────────
def test_full_fix_flow_valid_counts_and_clamp():
"""End-to-end test of the complete fix flow:
1. triton_convert_req_index_to_global_index with return_valid_counts=True
2. In-place clamp global_idx to replace -1 with 0
3. In-place clamp valid_counts to min=1
4. Use valid_counts as cache_seqlens
This simulates what forward_mqa does after the fix.
"""
from vllm.v1.attention.backends.mla.sparse_utils import (
triton_convert_req_index_to_global_index,
)
device = "cuda"
T = 4
topk = 128
num_blocks = 16
block_size = 64
req_id = torch.zeros(T, dtype=torch.int32, device=device)
block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
0
)
# Simulate causal prefill: token i has (i+1) valid entries
topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
expected_valid = [1, 2, 3, 4]
for i in range(T):
nv = expected_valid[i]
topk_indices[i, :nv] = torch.arange(nv, dtype=torch.int32, device=device)
# Step 1: Convert with valid counts
global_idx, valid_counts = triton_convert_req_index_to_global_index(
req_id,
block_table,
topk_indices,
BLOCK_SIZE=block_size,
NUM_TOPK_TOKENS=topk,
return_valid_counts=True,
)
# Step 2: In-place clamp global_idx (no -1 entries after)
global_idx.clamp_(min=0)
assert (global_idx >= 0).all(), "No -1 entries should remain after clamp_"
# Step 3: In-place clamp valid_counts to min=1
valid_counts.clamp_(min=1)
cache_seqlens = valid_counts
# Step 4: Verify valid counts match expected
for i in range(T):
assert cache_seqlens[i].item() == expected_valid[i], (
f"Token {i}: expected cache_seqlens={expected_valid[i]}, "
f"got {cache_seqlens[i].item()}"
)
# Step 5: Verify that for each token, entries 0..cache_seqlens-1 in
# global_idx are valid (non-zero, since we clamped -1 to 0 for the
# entries beyond valid_counts, the valid entries at positions 0..nv-1
# should be the actual converted indices)
for i in range(T):
nv = expected_valid[i]
valid_region = global_idx[i, :nv]
# Valid region should have specific converted values from block_table
# (not just zeros from clamping)
# For indices [0, 1, ..., nv-1] with block_size=64:
# block_id = index // 64, inblock_off = index % 64
# out = block_table[0, block_id] * 64 + inblock_off
for j in range(nv):
block_id = j // block_size
inblock_off = j % block_size
expected_val = block_table[0, block_id].item() * block_size + inblock_off
assert valid_region[j].item() == expected_val, (
f"Token {i}, position {j}: expected {expected_val}, "
f"got {valid_region[j].item()}"
)
# ─── TEST 2.10: CUDA Graph Padding Fix ─────────────────────────────
# These tests verify the fix for Issue 2: RuntimeError when
# num_actual_tokens (padded) != sum(seg_lens) (real tokens).
# This is the core bug that caused the crash during lm_eval with
# 32 concurrent requests on DeepSeek-V3.2.
def _make_mock_vllm_config(max_tokens=512):
"""Create a mock VllmConfig for metadata builder tests."""
from unittest.mock import MagicMock
vllm_config = MagicMock()
vllm_config.scheduler_config.max_num_batched_tokens = max_tokens
vllm_config.speculative_config = None
vllm_config.parallel_config.decode_context_parallel_size = 1
return vllm_config
def test_metadata_builder_cuda_graph_padding():
"""Verify build() handles CUDA graph padding (T > actual_tokens).
Reproduces the exact crash from Issue 2:
RuntimeError: The size of tensor a (32) must match the size
of tensor b (31) at non-singleton dimension 0
This happens when num_actual_tokens=32 (padded for CUDA graph)
but only 31 real tokens exist (one request completed mid-batch).
"""
from unittest.mock import MagicMock
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseMetadataBuilder,
)
device = "cuda"
max_tokens = 512
block_size = 64
topk = 2048
# Mock kv_cache_spec
kv_cache_spec = MagicMock()
kv_cache_spec.block_size = block_size
# Mock vllm_config
vllm_config = _make_mock_vllm_config(max_tokens)
builder = CutlassFA3MLASparseMetadataBuilder(
kv_cache_spec=kv_cache_spec,
layer_names=["layers.0.self_attn"],
vllm_config=vllm_config,
device=torch.device(device),
)
builder.topk_tokens = topk
# Simulate the crash scenario: 31 real tokens padded to 32
padded_T = 32
real_tokens = 31
num_reqs_padded = 32 # padded request count
# Accurately mock gpu_model_runner.py's padding behavior:
# query_start_loc.cpu[:num_reqs_padded+1] = [:33], 33 entries
# Real entries: [0,1,...,31], Padding: [31] (repeats last value)
query_start_loc_cpu = list(range(real_tokens + 1)) + [real_tokens]
# seq_lens_cpu[:num_reqs_padded] = [:32], 32 entries
# Real entries: [100]*31, Padding: [0] (stale/zero for padding slot)
seq_lens_cpu = [100] * real_tokens + [0]
# Build the mock CommonAttentionMetadata
cm = MagicMock()
cm.num_actual_tokens = padded_T # PADDED to 32
cm.query_start_loc_cpu = query_start_loc_cpu
cm.seq_lens_cpu = seq_lens_cpu
cm.num_reqs = num_reqs_padded # gpu_model_runner passes padded count
cm.max_query_len = 1
cm.max_seq_len = 100
cm.query_start_loc = torch.tensor(
query_start_loc_cpu, dtype=torch.int32, device=device
)
cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
cm.block_table_tensor = torch.zeros(
num_reqs_padded, 4, dtype=torch.int32, device=device
)
# This should NOT raise RuntimeError
metadata = builder.build(
common_prefix_len=0,
common_attn_metadata=cm,
)
# Verify metadata shapes match padded T
assert metadata.req_id_per_token.shape[0] == padded_T, (
f"req_id_per_token should have padded size {padded_T}, "
f"got {metadata.req_id_per_token.shape[0]}"
)
assert metadata.cache_seqlens.shape[0] == padded_T, (
f"cache_seqlens should have padded size {padded_T}, "
f"got {metadata.cache_seqlens.shape[0]}"
)
assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
# Verify real data portion is correct
for i in range(real_tokens):
assert metadata.req_id_per_token[i].item() == i, (
f"Token {i}: req_id should be {i}, "
f"got {metadata.req_id_per_token[i].item()}"
)
assert metadata.cache_seqlens[i].item() == 100, (
f"Token {i}: cache_seqlens should be 100, "
f"got {metadata.cache_seqlens[i].item()}"
)
# Verify padding tokens have safe defaults
assert metadata.req_id_per_token[real_tokens].item() == 0, (
"Padding token req_id should be 0"
)
assert metadata.cache_seqlens[real_tokens].item() >= 1, (
"Padding token cache_seqlens should be >= 1 (safe minimum)"
)
# Verify cu_seqlens_q is [0, 1, 2, ..., padded_T] (always correct)
for i in range(padded_T + 1):
assert metadata.cu_seqlens_q[i].item() == i, (
f"cu_seqlens_q[{i}] should be {i}, got {metadata.cu_seqlens_q[i].item()}"
)
# Verify cu_seqlens_k is monotonically non-decreasing
for i in range(padded_T):
assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item(), (
f"cu_seqlens_k must be non-decreasing at index {i}: "
f"{metadata.cu_seqlens_k[i].item()} -> {metadata.cu_seqlens_k[i + 1].item()}"
)
# Verify cu_seqlens_k at the real/padding boundary
assert metadata.cu_seqlens_k[real_tokens].item() == real_tokens * 100, (
f"cu_seqlens_k[{real_tokens}] should be {real_tokens * 100}, "
f"got {metadata.cu_seqlens_k[real_tokens].item()}"
)
@pytest.mark.parametrize(
"real_tokens,padded_T",
[
(1, 2), # minimal padding
(3, 32), # large padding gap
(7, 8), # small batch
(15, 16), # medium batch
(31, 32), # the exact crash scenario
(100, 104), # larger padding gap
],
)
def test_metadata_builder_cuda_graph_padding_various(real_tokens, padded_T):
"""Verify build() handles various CUDA graph padding scenarios.
Uses accurate mock that matches gpu_model_runner.py's padding behavior:
- query_start_loc_cpu has num_reqs_padded+1 entries (with padded suffix)
- seq_lens_cpu has num_reqs_padded entries (with stale padding entries)
"""
from unittest.mock import MagicMock
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseMetadataBuilder,
)
device = "cuda"
max_tokens = max(512, padded_T + 1) # ensure buffer large enough
block_size = 64
topk = 2048
num_reqs_padded = padded_T # For decode-only, padded_T == num_reqs_padded
kv_cache_spec = MagicMock()
kv_cache_spec.block_size = block_size
vllm_config = _make_mock_vllm_config(max_tokens)
builder = CutlassFA3MLASparseMetadataBuilder(
kv_cache_spec=kv_cache_spec,
layer_names=["layers.0.self_attn"],
vllm_config=vllm_config,
device=torch.device(device),
)
builder.topk_tokens = topk
# Accurate mock: query_start_loc_cpu[:num_reqs_padded+1]
# Real entries [0,1,...,real_tokens], then (num_reqs_padded - real_tokens)
# padding entries all equal to real_tokens (flat, non-decreasing)
query_start_loc_cpu = list(range(real_tokens + 1))
num_padding_reqs = num_reqs_padded - real_tokens
query_start_loc_cpu += [real_tokens] * num_padding_reqs
# seq_lens_cpu[:num_reqs_padded] — padding entries are stale (zero)
seq_lens_cpu = [200] * real_tokens + [0] * num_padding_reqs
cm = MagicMock()
cm.num_actual_tokens = padded_T
cm.query_start_loc_cpu = query_start_loc_cpu
cm.seq_lens_cpu = seq_lens_cpu
cm.num_reqs = num_reqs_padded
cm.max_query_len = 1
cm.max_seq_len = 200
cm.query_start_loc = torch.tensor(
query_start_loc_cpu, dtype=torch.int32, device=device
)
cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
cm.block_table_tensor = torch.zeros(
max(num_reqs_padded, 1), 4, dtype=torch.int32, device=device
)
# Should NOT raise any errors
metadata = builder.build(
common_prefix_len=0,
common_attn_metadata=cm,
)
# Verify shapes match padded T
assert metadata.req_id_per_token.shape[0] == padded_T
assert metadata.cache_seqlens.shape[0] == padded_T
assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
assert metadata.num_actual_tokens == padded_T
# Verify real portion
for i in range(real_tokens):
assert metadata.req_id_per_token[i].item() == i
assert metadata.cache_seqlens[i].item() == 200
# Verify padding
for i in range(real_tokens, padded_T):
assert metadata.req_id_per_token[i].item() == 0
assert metadata.cache_seqlens[i].item() >= 1
# Verify cu_seqlens_q is [0, 1, ..., padded_T]
for i in range(padded_T + 1):
assert metadata.cu_seqlens_q[i].item() == i
# Verify cu_seqlens_k monotonicity
for i in range(padded_T):
assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item()
# Verify cu_seqlens_k at boundary
assert metadata.cu_seqlens_k[real_tokens].item() == real_tokens * 200
def test_metadata_builder_no_padding():
"""Verify build() still works correctly when T == actual_tokens (no padding)."""
from unittest.mock import MagicMock
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseMetadataBuilder,
)
device = "cuda"
max_tokens = 512
block_size = 64
kv_cache_spec = MagicMock()
kv_cache_spec.block_size = block_size
vllm_config = _make_mock_vllm_config(max_tokens)
builder = CutlassFA3MLASparseMetadataBuilder(
kv_cache_spec=kv_cache_spec,
layer_names=["layers.0.self_attn"],
vllm_config=vllm_config,
device=torch.device(device),
)
builder.topk_tokens = 2048
# No padding: T == real tokens
T = 4
query_start_loc_cpu = [0, 1, 2, 3, 4] # 4 decode tokens
seq_lens_cpu = [50, 100, 150, 200]
cm = MagicMock()
cm.num_actual_tokens = T
cm.query_start_loc_cpu = query_start_loc_cpu
cm.seq_lens_cpu = seq_lens_cpu
cm.num_reqs = 4
cm.max_query_len = 1
cm.max_seq_len = 200
cm.query_start_loc = torch.tensor(
query_start_loc_cpu, dtype=torch.int32, device=device
)
cm.slot_mapping = torch.zeros(T, dtype=torch.int64, device=device)
cm.block_table_tensor = torch.zeros(4, 4, dtype=torch.int32, device=device)
metadata = builder.build(
common_prefix_len=0,
common_attn_metadata=cm,
)
assert metadata.req_id_per_token.shape[0] == T
assert metadata.cache_seqlens.shape[0] == T
assert metadata.num_actual_tokens == T
# Verify exact values
assert metadata.req_id_per_token[0].item() == 0
assert metadata.req_id_per_token[1].item() == 1
assert metadata.req_id_per_token[2].item() == 2
assert metadata.req_id_per_token[3].item() == 3
assert metadata.cache_seqlens[0].item() == 50
assert metadata.cache_seqlens[1].item() == 100
assert metadata.cache_seqlens[2].item() == 150
assert metadata.cache_seqlens[3].item() == 200
def test_metadata_builder_mixed_prefill_decode_with_padding():
"""Verify build() handles mixed prefill+decode with CUDA graph padding.
This tests a more complex scenario: 2 decode tokens + 3 prefill tokens
from 3 requests, padded from 5 to 8 tokens.
"""
from unittest.mock import MagicMock
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseMetadataBuilder,
)
device = "cuda"
max_tokens = 512
block_size = 64
kv_cache_spec = MagicMock()
kv_cache_spec.block_size = block_size
vllm_config = _make_mock_vllm_config(max_tokens)
builder = CutlassFA3MLASparseMetadataBuilder(
kv_cache_spec=kv_cache_spec,
layer_names=["layers.0.self_attn"],
vllm_config=vllm_config,
device=torch.device(device),
)
builder.topk_tokens = 2048
# 3 real requests: req0 (1 decode token), req1 (1 decode token),
# req2 (3 prefill tokens)
# Total: 5 real tokens, padded to 8 tokens, 8 padded request slots
real_tokens = 5
num_real_reqs = 3
padded_T = 8
num_reqs_padded = 8 # padded request count
# Accurate: query_start_loc_cpu[:num_reqs_padded+1] = 9 entries
# Real: [0, 1, 2, 5], Padding: [5, 5, 5, 5, 5]
query_start_loc_cpu = [0, 1, 2, 5] + [5] * (num_reqs_padded - num_real_reqs)
# seq_lens_cpu[:num_reqs_padded] = 8 entries
seq_lens_cpu = [100, 200, 3] + [0] * (num_reqs_padded - num_real_reqs)
cm = MagicMock()
cm.num_actual_tokens = padded_T
cm.query_start_loc_cpu = query_start_loc_cpu
cm.seq_lens_cpu = seq_lens_cpu
cm.num_reqs = num_reqs_padded
cm.max_query_len = 3
cm.max_seq_len = 200
cm.query_start_loc = torch.tensor(
query_start_loc_cpu, dtype=torch.int32, device=device
)
cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
cm.block_table_tensor = torch.zeros(
num_reqs_padded, 4, dtype=torch.int32, device=device
)
metadata = builder.build(
common_prefix_len=0,
common_attn_metadata=cm,
)
# Verify shapes
assert metadata.req_id_per_token.shape[0] == padded_T
assert metadata.cache_seqlens.shape[0] == padded_T
# Verify req_id mapping
assert metadata.req_id_per_token[0].item() == 0 # req0, decode
assert metadata.req_id_per_token[1].item() == 1 # req1, decode
assert metadata.req_id_per_token[2].item() == 2 # req2, prefill tok0
assert metadata.req_id_per_token[3].item() == 2 # req2, prefill tok1
assert metadata.req_id_per_token[4].item() == 2 # req2, prefill tok2
# Padding tokens
assert metadata.req_id_per_token[5].item() == 0
assert metadata.req_id_per_token[6].item() == 0
assert metadata.req_id_per_token[7].item() == 0
# Verify cache_seqlens
assert metadata.cache_seqlens[0].item() == 100 # req0 seq_len
assert metadata.cache_seqlens[1].item() == 200 # req1 seq_len
assert metadata.cache_seqlens[2].item() == 3 # req2 seq_len
assert metadata.cache_seqlens[3].item() == 3 # req2 seq_len
assert metadata.cache_seqlens[4].item() == 3 # req2 seq_len
# Padding (default = 1)
assert metadata.cache_seqlens[5].item() >= 1
assert metadata.cache_seqlens[6].item() >= 1
assert metadata.cache_seqlens[7].item() >= 1
# ─── TEST 2.11: Zero Real Tokens Edge Case (Review Issue #3) ─────
# Tests the edge case where ALL tokens are padding (actual_tokens=0).
# This can happen during CUDA graph warmup/capture with dummy batches.
def test_metadata_builder_zero_real_tokens():
"""Verify build() handles the case where all tokens are padding.
This edge case can occur during CUDA graph warmup or capture where
dummy batches may have zero real tokens but T > 0 (padded size).
"""
from unittest.mock import MagicMock
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseMetadataBuilder,
)
device = "cuda"
max_tokens = 512
block_size = 64
kv_cache_spec = MagicMock()
kv_cache_spec.block_size = block_size
vllm_config = _make_mock_vllm_config(max_tokens)
builder = CutlassFA3MLASparseMetadataBuilder(
kv_cache_spec=kv_cache_spec,
layer_names=["layers.0.self_attn"],
vllm_config=vllm_config,
device=torch.device(device),
)
builder.topk_tokens = 2048
# Zero real tokens, padded to 4
# This happens when query_start_loc = [0] only (1 entry, no requests)
# and num_actual_tokens is still the padded count.
padded_T = 4
real_tokens = 0
# query_start_loc_cpu with a single entry means 0 requests
query_start_loc_cpu = [0]
seq_lens_cpu = []
cm = MagicMock()
cm.num_actual_tokens = padded_T
cm.query_start_loc_cpu = query_start_loc_cpu
cm.seq_lens_cpu = seq_lens_cpu
cm.num_reqs = 0
cm.max_query_len = 0
cm.max_seq_len = 0
cm.query_start_loc = torch.tensor(
query_start_loc_cpu, dtype=torch.int32, device=device
)
cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
cm.block_table_tensor = torch.zeros(1, 4, dtype=torch.int32, device=device)
# Should NOT raise any errors
metadata = builder.build(
common_prefix_len=0,
common_attn_metadata=cm,
)
# Verify shapes match padded T
assert metadata.req_id_per_token.shape[0] == padded_T
assert metadata.cache_seqlens.shape[0] == padded_T
assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
# All tokens are padding — verify safe defaults
for i in range(padded_T):
assert metadata.req_id_per_token[i].item() == 0
assert metadata.cache_seqlens[i].item() >= 1
# cu_seqlens_k should be monotonically non-decreasing
for i in range(padded_T):
assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item()
# ─── TEST 2.12: Batch Size Gating Constant ───────────────────────────
def test_batch_size_gating_threshold():
"""Verify MAX_BATCH_SIZE_FOR_FA3 is 16 and controls routing."""
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
MAX_BATCH_SIZE_FOR_FA3,
_flashmla_sparse_available,
)
assert MAX_BATCH_SIZE_FOR_FA3 == 16
# On SM90 builds, FlashMLA fallback should be available
# (unless FlashMLA was explicitly excluded from the build)
assert isinstance(_flashmla_sparse_available, bool)
# ─── TEST 2.13: FlashMLA Fallback Head Padding ──────────────────────
def test_flashmla_fallback_head_padding():
"""Verify FlashMLA fallback head padding constant is 64 for SM90."""
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
_FLASHMLA_SM90_HEAD_PADDING,
)
assert _FLASHMLA_SM90_HEAD_PADDING == 64, (
f"SM90 head padding should be 64, got {_FLASHMLA_SM90_HEAD_PADDING}"
)
# ─── TEST 2.14: Forward MQA Dispatch Verification ────────────────────
def test_forward_mqa_has_fa3_and_fallback_methods():
"""Verify CutlassFA3MLASparseImpl has both kernel dispatch methods."""
from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
CutlassFA3MLASparseImpl,
)
assert hasattr(CutlassFA3MLASparseImpl, "_forward_fa3"), (
"CutlassFA3MLASparseImpl should have _forward_fa3 method"
)
assert hasattr(CutlassFA3MLASparseImpl, "_forward_flashmla_bf16_fallback"), (
"CutlassFA3MLASparseImpl should have _forward_flashmla_bf16_fallback method"
)
+3 -1
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""The request function for API endpoints."""
import codecs
import io
import json
import os
@@ -25,11 +26,12 @@ class StreamedResponseHandler:
def __init__(self):
self.buffer = ""
self._decoder = codecs.getincrementaldecoder("utf-8")()
def add_chunk(self, chunk_bytes: bytes) -> list[str]:
"""Add a chunk of bytes to the buffer and return any complete
messages."""
chunk_str = chunk_bytes.decode("utf-8")
chunk_str = self._decoder.decode(chunk_bytes)
self.buffer += chunk_str
messages = []
+16 -14
View File
@@ -8,7 +8,7 @@ from dataclasses import dataclass
from functools import partial
from pathlib import Path
from types import TracebackType
from typing import ClassVar
from typing import TYPE_CHECKING, ClassVar
from typing_extensions import Self, override
@@ -17,20 +17,8 @@ from vllm.utils.import_utils import PlaceholderModule
from .utils import sanitize_filename
try:
import matplotlib.pyplot as plt
except ImportError:
plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
try:
if TYPE_CHECKING:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
try:
import seaborn as sns
except ImportError:
seaborn = PlaceholderModule("seaborn")
@dataclass
@@ -265,6 +253,20 @@ def _plot_fig(
fig_height: float,
fig_dpi: int,
):
# Lazy-import matplotlib/pandas/seaborn
try:
import matplotlib.pyplot as plt
except ImportError:
plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
try:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
try:
import seaborn as sns
except ImportError:
sns = PlaceholderModule("seaborn")
fig_group, fig_data = fig_group_data
row_groups = full_groupby(
+16 -14
View File
@@ -6,7 +6,7 @@ from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from typing import ClassVar
from typing import TYPE_CHECKING, ClassVar
from vllm.utils.collection_utils import full_groupby
from vllm.utils.import_utils import PlaceholderModule
@@ -14,20 +14,8 @@ from vllm.utils.import_utils import PlaceholderModule
from .plot import DummyExecutor, _json_load_bytes
from .utils import sanitize_filename
try:
import matplotlib.pyplot as plt
except ImportError:
plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
try:
if TYPE_CHECKING:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
try:
import seaborn as sns
except ImportError:
seaborn = PlaceholderModule("seaborn")
def _first_present(run_data: dict[str, object], keys: list[str]):
@@ -195,6 +183,20 @@ def _plot_fig(
print("[END FIGURE]")
return
# Lazy-import matplotlib/pandas/seaborn
try:
import matplotlib.pyplot as plt
except ImportError:
plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
try:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
try:
import seaborn as sns
except ImportError:
sns = PlaceholderModule("seaborn")
df = pd.DataFrame.from_records(fig_data)
df = df.dropna(subset=["tokens_per_user", "tokens_per_gpu"])
+1 -3
View File
@@ -30,6 +30,7 @@ CacheDType = Literal[
"turboquant_3bit_nc",
"int8_per_token_head",
"fp8_per_token_head",
"nvfp4",
]
MambaDType = Literal["auto", "float32", "float16"]
MambaCacheMode = Literal["all", "align", "none"]
@@ -101,8 +102,6 @@ class CacheConfig:
kv_cache_dtype_skip_layers: list[str] = field(default_factory=list)
"""Layer patterns to skip KV cache quantization. Accepts layer indices
(e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window')."""
cpu_kvcache_space_bytes: int | None = None
"""(CPU backend only) CPU key-value cache space."""
mamba_page_size_padded: int | None = None
""" Optional override for mamba page size; used by hybrid mamba/attention
models to ensure exact alignment with attention page size."""
@@ -183,7 +182,6 @@ class CacheConfig:
"num_gpu_blocks_override",
"enable_prefix_caching",
"prefix_caching_hash_algo",
"cpu_kvcache_space_bytes",
"mamba_page_size_padded",
"user_specified_block_size",
"user_specified_mamba_block_size",
+4
View File
@@ -19,6 +19,10 @@ class OnlineQuantScheme(Enum):
# blocks of 128x128 elements (popularized by DeepSeek)
FP8_PER_BLOCK = "fp8_per_block"
# int8, weight-only per-channel quantization for MoE expert weights.
# Linear layers remain unquantized.
INT8_PER_CHANNEL_WEIGHT_ONLY = "int8_per_channel_weight_only"
# TODO(future PRs): add more online quant schemes here: mxfp8, etc
+12
View File
@@ -1550,6 +1550,18 @@ def _parse_chat_message_content(
parsed_msg = _ToolParser(message)
if "tool_call_id" in parsed_msg:
result_msg["tool_call_id"] = parsed_msg["tool_call_id"]
# Normalize tool message content from OpenAI array format to plain
# string. Clients like Claude Code / Cursor send tool results as
# [{"type": "text", "text": "..."}], but most chat templates only
# handle string content for tool messages.
msg_content = result_msg.get("content")
if isinstance(msg_content, list):
texts = [
item.get("text", "")
for item in msg_content
if isinstance(item, dict) and item.get("type") == "text"
]
result_msg["content"] = "\n".join(texts) if texts else ""
if "name" in message and isinstance(message["name"], str):
result_msg["name"] = message["name"]
-17
View File
@@ -1,19 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.entrypoints.cli.benchmark.latency import BenchmarkLatencySubcommand
from vllm.entrypoints.cli.benchmark.mm_processor import (
BenchmarkMMProcessorSubcommand,
)
from vllm.entrypoints.cli.benchmark.serve import BenchmarkServingSubcommand
from vllm.entrypoints.cli.benchmark.startup import BenchmarkStartupSubcommand
from vllm.entrypoints.cli.benchmark.sweep import BenchmarkSweepSubcommand
from vllm.entrypoints.cli.benchmark.throughput import BenchmarkThroughputSubcommand
__all__: list[str] = [
"BenchmarkLatencySubcommand",
"BenchmarkMMProcessorSubcommand",
"BenchmarkServingSubcommand",
"BenchmarkStartupSubcommand",
"BenchmarkSweepSubcommand",
"BenchmarkThroughputSubcommand",
]
+34 -12
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import sys
import typing
from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
@@ -14,6 +15,17 @@ else:
FlexibleArgumentParser = argparse.ArgumentParser
def _import_bench_subcommand_modules() -> None:
# Imported lazily so `BenchmarkSubcommandBase` subclasses register only
# when `vllm bench` is actually invoked.
import vllm.entrypoints.cli.benchmark.latency # noqa: F401
import vllm.entrypoints.cli.benchmark.mm_processor # noqa: F401
import vllm.entrypoints.cli.benchmark.serve # noqa: F401
import vllm.entrypoints.cli.benchmark.startup # noqa: F401
import vllm.entrypoints.cli.benchmark.sweep # noqa: F401
import vllm.entrypoints.cli.benchmark.throughput # noqa: F401
class BenchmarkSubcommand(CLISubcommand):
"""The `bench` subcommand for the vLLM CLI."""
@@ -38,18 +50,28 @@ class BenchmarkSubcommand(CLISubcommand):
)
bench_subparsers = bench_parser.add_subparsers(required=True, dest="bench_type")
for cmd_cls in BenchmarkSubcommandBase.__subclasses__():
cmd_subparser = bench_subparsers.add_parser(
cmd_cls.name,
help=cmd_cls.help,
description=cmd_cls.help,
usage=f"vllm {self.name} {cmd_cls.name} [options]",
)
cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd)
cmd_cls.add_cli_args(cmd_subparser)
cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(
subcmd=f"{self.name} {cmd_cls.name}"
)
# Only build the nested bench subparsers when the user is actually
# invoking `bench`; otherwise we'd drag in imports
# unnecessarily on every `vllm --help` and `vllm serve`.
# Scan for the first positional arg so global flags (e.g. `-v`)
# before the subcommand don't break detection.
first_positional = next(
(arg for arg in sys.argv[1:] if not arg.startswith("-")), None
)
if first_positional == self.name:
_import_bench_subcommand_modules()
for cmd_cls in BenchmarkSubcommandBase.__subclasses__():
cmd_subparser = bench_subparsers.add_parser(
cmd_cls.name,
help=cmd_cls.help,
description=cmd_cls.help,
usage=f"vllm {self.name} {cmd_cls.name} [options]",
)
cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd)
cmd_cls.add_cli_args(cmd_subparser)
cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(
subcmd=f"{self.name} {cmd_cls.name}"
)
return bench_parser
@@ -97,6 +97,11 @@ class IOProcessorRequest(PoolingBasicRequestMixin, EncodingRequestMixin, Generic
max_total_tokens_param="max_model_len",
)
def to_pooling_params(self):
return PoolingParams(
task=self.task,
)
class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
request_id: str | None = None
+2 -3
View File
@@ -100,10 +100,9 @@ logger = init_logger(__name__)
# it avoids unintentional cuda initialization from torch.cuda.is_available()
os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
# see https://github.com/vllm-project/vllm/issues/10480
# see https://github.com/vllm-project/vllm/issues/10480 and
# https://github.com/vllm-project/vllm/issues/10619.
os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
# see https://github.com/vllm-project/vllm/issues/10619
torch._inductor.config.compile_threads = 1
# Enable Triton autotuning result caching to disk by default.
# Without this, Triton re-runs autotuning on every process restart,
+16 -2
View File
@@ -186,12 +186,13 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
# in priority/performance order (when available)
_POSSIBLE_FP8_BLOCK_KERNELS: dict[
PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel]]
PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]]
] = {
PlatformEnum.CUDA: [
FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
DeepGemmFp8BlockScaledMMKernel,
CutlassFp8BlockScaledMMKernel,
MarlinFP8ScaledMMLinearKernel,
TritonFp8BlockScaledMMKernel,
],
PlatformEnum.ROCM: [
@@ -392,6 +393,19 @@ def init_fp8_linear_kernel(
scope="global",
)
# TODO make scaled_mm kernels inherit from MMLinearKernel
# only MarlinFP8ScaledMMLinearKernel is a type of FP8ScaledMMLinearKernel.
if issubclass(kernel_type, FP8ScaledMMLinearKernel):
return kernel_type(
scaled_mm_linear_kernel_config,
layer_param_names=[
"weight",
"weight_scale",
"input_scale",
"input_scale_ub",
],
)
return kernel_type(
scaled_mm_linear_kernel_config,
)
@@ -399,7 +413,7 @@ def init_fp8_linear_kernel(
else:
kernel_type = choose_scaled_mm_linear_kernel(
config=scaled_mm_linear_kernel_config,
possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[misc]
possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type]
force_kernel=force_kernel,
)
if module_name:
@@ -387,7 +387,9 @@ class Attention(nn.Module, AttentionLayerBase):
self.query_quant = None
if (
self.impl.supports_quant_query_input
and self.kv_cache_dtype.startswith("fp8")
and (
self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4"
)
and not self.kv_cache_dtype.endswith("per_token_head")
):
is_per_head = (
@@ -492,7 +494,7 @@ class Attention(nn.Module, AttentionLayerBase):
# which reduces overheads during decoding.
# Otherwise queries are quantized using custom ops
# which causes decoding overheads
assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"}
assert self.kv_cache_dtype in {"fp8", "fp8_e4m3", "nvfp4"}
# check if query quantization is supported
if self.impl.supports_quant_query_input:
@@ -367,6 +367,16 @@ class MLAAttention(nn.Module, AttentionLayerBase):
"KV cache format, please set `--attention-backend FLASHMLA_SPARSE`"
)
# CUTLASS FA3 MLA Sparse requires BF16 KV cache — force "auto" dtype
if self.attn_backend.get_name() == "CUTLASS_FA3_MLA_SPARSE":
if cache_config is not None:
cache_config.cache_dtype = "auto"
kv_cache_dtype = "auto"
logger.info_once(
"CUTLASS FA3 MLA Sparse backend requires BF16 KV cache. "
"Setting kv_cache_dtype to 'auto' (BF16)."
)
# Initialize KV cache quantization attributes
self.kv_cache_dtype = kv_cache_dtype
self.calculate_kv_scales = calculate_kv_scales
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.all2all_utils import (
maybe_make_prepare_finalize,
)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEQuantConfig,
int8_w8a16_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
SharedExperts,
)
logger = init_logger(__name__)
def select_int8_moe_backend(
config: FusedMoEConfig,
) -> type[mk.FusedMoEExperts]:
from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
supported, reason = TritonExperts.is_supported_config(
TritonExperts,
config,
None,
None,
mk.FusedMoEActivationFormat.Standard,
)
if not supported:
raise ValueError(
f"INT8 Triton MoE backend does not support the "
f"deployment configuration: {reason}"
)
logger.info_once("Using Triton INT8 MoE backend", scope="local")
return TritonExperts
def make_int8_moe_quant_config(
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
) -> FusedMoEQuantConfig:
return int8_w8a16_moe_quant_config(
w1_scale=w1_scale,
w2_scale=w2_scale,
w1_zp=None,
w2_zp=None,
)
def make_int8_moe_kernel(
moe_quant_config: FusedMoEQuantConfig,
moe_config: FusedMoEConfig,
experts_cls: type[mk.FusedMoEExperts],
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
shared_experts: SharedExperts | None = None,
) -> mk.FusedMoEKernel:
prepare_finalize = maybe_make_prepare_finalize(
moe=moe_config,
quant_config=moe_quant_config,
routing_tables=routing_tables,
allow_new_interface=True,
)
assert prepare_finalize is not None
logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
experts = experts_cls(
moe_config=moe_config,
quant_config=moe_quant_config,
)
return mk.FusedMoEKernel(
prepare_finalize,
experts,
shared_experts=shared_experts,
inplace=not moe_config.disable_inplace,
)
@@ -15,6 +15,7 @@ import torch
from vllm.config.mamba import MambaBackendEnum, MambaConfig
from vllm.logger import init_logger
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
logger = init_logger(__name__)
@@ -188,12 +189,22 @@ _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
_mamba_ssu_backend: MambaSSUBackend | None = None
def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None:
def initialize_mamba_ssu_backend(
mamba_config: MambaConfig,
kv_cache_config: KVCacheConfig,
) -> None:
"""Initialize the global Mamba SSU backend.
Args:
mamba_config: Mamba configuration.
No-op if `kv_cache_config` contains no specs that call
selective_state_update.
"""
if not any(
isinstance(g.kv_cache_spec, MambaSpec)
and g.kv_cache_spec.mamba_type in ("mamba1", "mamba2")
for g in kv_cache_config.kv_cache_groups
):
return
global _mamba_ssu_backend
backend = mamba_config.backend
@@ -203,7 +214,11 @@ def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None:
f"Valid options: {list(_BACKEND_REGISTRY.keys())}"
)
_mamba_ssu_backend = _BACKEND_REGISTRY[backend](mamba_config)
backend_cls = _BACKEND_REGISTRY[backend]
if isinstance(_mamba_ssu_backend, backend_cls):
return
_mamba_ssu_backend = backend_cls(mamba_config)
logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
@@ -40,6 +40,7 @@ QuantizationMethods = Literal[
# shorthand for creating a more complicated online quant config object
"fp8_per_tensor",
"fp8_per_block",
"int8_per_channel_weight_only",
]
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
@@ -47,7 +48,6 @@ DEPRECATED_QUANTIZATION_METHODS = [
"tpu_int8",
"fbgemm_fp8",
"fp_quant",
"experts_int8",
]
# The customized quantization methods which will be added to this dict.
@@ -5,27 +5,25 @@ from typing import Any
import torch
from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
from vllm.model_executor.layers.fused_moe import (
FusedMoE,
FusedMoEConfig,
FusedMoEMethodBase,
)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
int8_w8a16_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm.model_executor.layers.quantization import QuantizationMethods
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
QuantizeMethodBase,
)
from vllm.model_executor.utils import set_weight_attrs
from vllm.model_executor.layers.quantization.online.int8 import (
Int8OnlineMoEMethod,
)
class ExpertsInt8Config(QuantizationConfig):
"""Config class for Int8 experts quantization."""
"""Online int8 quantization for MoE expert weights.
Linear layers are left unquantized.
Backward-compatible config for ``--quantization experts_int8``.
Prefer ``--quantization int8_per_channel``
"""
def __init__(self) -> None:
super().__init__()
@@ -56,149 +54,5 @@ class ExpertsInt8Config(QuantizationConfig):
if isinstance(layer, LinearBase):
return UnquantizedLinearMethod()
elif isinstance(layer, FusedMoE):
return ExpertsInt8MoEMethod(self, layer.moe_config)
return Int8OnlineMoEMethod(layer=layer)
return None
class ExpertsInt8MoEMethod(FusedMoEMethodBase):
def __init__(
self,
quant_config: ExpertsInt8Config,
moe: FusedMoEConfig,
):
super().__init__(moe)
self.quant_config = quant_config
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
int8_dtype = torch.int8
assert "weight_loader" in extra_weight_attrs
weight_loader = extra_weight_attrs["weight_loader"]
wrapped_weight_loader = ExpertsInt8MoEMethod.quantizing_weight_loader(
layer, weight_loader
)
extra_weight_attrs["weight_loader"] = wrapped_weight_loader
# Fused gate_up_proj (column parallel)
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
dtype=int8_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
# down_proj (row parallel)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
dtype=int8_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
w13_scale = torch.nn.Parameter(
torch.zeros(
num_experts, 2 * intermediate_size_per_partition, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_scale", w13_scale)
w2_scale = torch.nn.Parameter(
torch.zeros(num_experts, hidden_size, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_scale", w2_scale)
def get_fused_moe_quant_config(
self, layer: torch.nn.Module
) -> FusedMoEQuantConfig | None:
return int8_w8a16_moe_quant_config(
w1_scale=layer.w13_scale, w2_scale=layer.w2_scale, w1_zp=None, w2_zp=None
)
def apply(
self,
layer: FusedMoE,
x: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
shared_experts_input: torch.Tensor | None,
) -> torch.Tensor:
from vllm.model_executor.layers.fused_moe import fused_experts
return fused_experts(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights=topk_weights,
topk_ids=topk_ids,
inplace=not self.moe.disable_inplace,
activation=layer.activation,
apply_router_weight_on_input=layer.apply_router_weight_on_input,
global_num_experts=layer.global_num_experts,
expert_map=layer.expert_map,
quant_config=self.moe_quant_config,
)
@staticmethod
def quantizing_weight_loader(layer, weight_loader):
def quantize_and_call_weight_loader(
param: torch.nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: int,
expert_id: int,
):
tp_rank = get_tensor_model_parallel_rank()
shard_size = layer.intermediate_size_per_partition
shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
device = get_tp_group().device
loaded_weight = loaded_weight.to(device)
# w1, gate_proj case: Load into first shard of w13.
if shard_id == "w1":
scales = quantize_in_place_and_get_scales(loaded_weight[shard, :])
layer.w13_scale.data[expert_id, 0:shard_size].copy_(scales[:, 0])
# w3, up_proj case: Load into second shard of w13.
elif shard_id == "w3":
scales = quantize_in_place_and_get_scales(loaded_weight[shard, :])
layer.w13_scale.data[expert_id, shard_size : 2 * shard_size].copy_(
scales[:, 0]
)
# w2, down_proj case: Load into only shard of w2.
elif shard_id == "w2":
scales = quantize_in_place_and_get_scales(loaded_weight[:, shard])
layer.w2_scale.data[expert_id, :].copy_(scales[:, 0])
else:
raise ValueError(f"Shard id must be in [0,1,2] but got {shard_id}")
weight_loader(param, loaded_weight, weight_name, shard_id, expert_id)
return quantize_and_call_weight_loader
def quantize_in_place_and_get_scales(weight: torch.Tensor) -> torch.Tensor:
vmax = torch.iinfo(torch.int8).max
scales = torch.max(torch.abs(weight), dim=1, keepdim=True)[0] / vmax
weight.div_(scales)
weight.round_()
weight.clamp_(-vmax, vmax)
return scales
@@ -9,6 +9,7 @@ from vllm.config.quantization import (
OnlineQuantizationConfigArgs,
OnlineQuantScheme,
)
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe import (
FusedMoE,
)
@@ -33,6 +34,11 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
Fp8PerTensorOnlineLinearMethod,
Fp8PerTensorOnlineMoEMethod,
)
from vllm.model_executor.layers.quantization.online.int8 import (
Int8OnlineMoEMethod,
)
logger = init_logger(__name__)
class OnlineQuantizationConfig(QuantizationConfig):
@@ -96,7 +102,13 @@ class OnlineQuantizationConfig(QuantizationConfig):
return UnquantizedLinearMethod()
linear_scheme = self.args.linear_scheme_override or self.args.global_scheme
if linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
if linear_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
logger.warning_once(
"INT8 online quantization only quantizes MoE expert "
"weights. linear layers remain in full precision."
)
return UnquantizedLinearMethod()
elif linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
return Fp8PerBlockOnlineLinearMethod()
else:
return Fp8PerTensorOnlineLinearMethod()
@@ -109,7 +121,9 @@ class OnlineQuantizationConfig(QuantizationConfig):
return UnquantizedFusedMoEMethod(layer.moe_config)
moe_scheme = self.args.moe_scheme_override or self.args.global_scheme
if moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
if moe_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
return Int8OnlineMoEMethod(layer=layer)
elif moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
return Fp8PerBlockOnlineMoEMethod(layer=layer)
else:
return Fp8PerTensorOnlineMoEMethod(layer=layer)
@@ -10,7 +10,6 @@ if TYPE_CHECKING:
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEQuantConfig,
)
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
@@ -19,15 +18,15 @@ import vllm.envs as envs
from vllm import _custom_ops as ops
from vllm.config import get_current_vllm_config
from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
from vllm.model_executor.layers.fused_moe import (
FusedMoEMethodBase,
)
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
select_fp8_moe_backend,
)
from vllm.model_executor.layers.linear import (
LinearMethodBase,
)
from vllm.model_executor.layers.quantization.online.moe_base import (
OnlineMoEMethodBase,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
create_fp8_quant_key,
@@ -44,7 +43,7 @@ from vllm.model_executor.model_loader.reload.layerwise import (
initialize_online_processing,
)
from vllm.model_executor.parameter import ModelWeightParameter
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
from vllm.model_executor.utils import replace_parameter
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import per_block_cast_to_fp8
@@ -268,21 +267,15 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
# ---------------------------------------------------------------------------
class _Fp8OnlineMoEBase(FusedMoEMethodBase):
class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
"""Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint
weights onto meta device and materializes them just-in-time."""
uses_meta_device: bool = True
# Declared here for mypy; actual values are set in __init__.
fp8_backend: "Fp8MoeBackend"
experts_cls: "type[mk.FusedMoEExperts] | None"
weight_scale_name: str
weight_block_size: list[int] | None
moe: "FusedMoEConfig"
is_monolithic: bool
moe_quant_config: "FusedMoEQuantConfig | None"
moe_kernel: "mk.FusedMoEKernel | None"
def __init__(
self,
@@ -313,77 +306,6 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
allow_vllm_cutlass=False,
)
def create_weights(
self,
layer: Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
layer.num_experts = num_experts
layer.orig_dtype = params_dtype
layer.weight_block_size = None
# WEIGHTS
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
device="meta",
dtype=params_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
device="meta", # materialized and processed during loading
dtype=params_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# BIASES (for models like GPT-OSS that have biased MoE)
if self.moe.has_bias:
w13_bias = torch.nn.Parameter(
torch.zeros(
num_experts,
2 * intermediate_size_per_partition,
device="meta", # materialized and processed during loading
dtype=layer.orig_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_bias", w13_bias)
set_weight_attrs(w13_bias, extra_weight_attrs)
w2_bias = torch.nn.Parameter(
torch.zeros(
num_experts,
hidden_size,
device="meta", # materialized and processed during loading
dtype=layer.orig_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_bias", w2_bias)
set_weight_attrs(w2_bias, extra_weight_attrs)
layer.w13_input_scale = None
layer.w2_input_scale = None
initialize_online_processing(layer)
def _setup_kernel(
self,
layer: "FusedMoE",
@@ -430,15 +352,6 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
shared_experts=layer.shared_experts,
)
def maybe_make_prepare_finalize(
self,
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> "mk.FusedMoEPrepareAndFinalizeModular | None":
raise ValueError(
f"{self.__class__.__name__} uses the new modular kernel "
"initialization logic. This function should not be called."
)
def get_fused_moe_quant_config(
self, layer: torch.nn.Module
) -> "FusedMoEQuantConfig":
@@ -460,68 +373,9 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
block_shape=self.weight_block_size,
)
# Inject biases into the quant config if the model has them
# (e.g. GPT-OSS biased MoE)
if quant_config is not None and self.moe.has_bias:
w13_bias = getattr(layer, "w13_bias", None)
w2_bias = getattr(layer, "w2_bias", None)
if w13_bias is not None:
quant_config._w1.bias = w13_bias
if w2_bias is not None:
quant_config._w2.bias = w2_bias
self._maybe_inject_biases(quant_config, layer)
return quant_config
@property
def supports_eplb(self) -> bool:
return True
def apply_monolithic(
self,
layer: "FusedMoE",
x: torch.Tensor,
router_logits: torch.Tensor,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
assert self.is_monolithic
assert self.moe_kernel is not None
return self.moe_kernel.apply_monolithic(
x,
layer.w13_weight,
layer.w2_weight,
router_logits,
activation=layer.activation,
global_num_experts=layer.global_num_experts,
expert_map=layer.expert_map,
apply_router_weight_on_input=layer.apply_router_weight_on_input,
num_expert_group=layer.num_expert_group,
topk_group=layer.topk_group,
e_score_correction_bias=layer.e_score_correction_bias,
routed_scaling_factor=layer.routed_scaling_factor,
)
def apply(
self,
layer: "FusedMoE",
x: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
shared_experts_input: torch.Tensor | None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
assert not self.is_monolithic
assert self.moe_kernel is not None
return self.moe_kernel.apply(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_ids,
activation=layer.activation,
global_num_experts=layer.global_num_experts,
expert_map=layer.expert_map,
apply_router_weight_on_input=layer.apply_router_weight_on_input,
shared_experts_input=shared_experts_input,
)
class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
"""Online tensorwise FP8 MoE quantization.
@@ -0,0 +1,109 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import TYPE_CHECKING
import torch
from torch.nn import Module
if TYPE_CHECKING:
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
)
from vllm.model_executor.layers.fused_moe.oracle.int8 import (
make_int8_moe_kernel,
make_int8_moe_quant_config,
select_int8_moe_backend,
)
from vllm.model_executor.layers.quantization.online.moe_base import (
OnlineMoEMethodBase,
)
from vllm.model_executor.utils import replace_parameter
class Int8OnlineMoEMethod(OnlineMoEMethodBase):
"""Online per-channel INT8 MoE quantization.
Loads fp16/bf16 weights and quantizes them per-row to int8 during loading.
"""
def __init__(
self,
*,
layer: torch.nn.Module,
):
super().__init__(layer.moe_config)
self.experts_cls: type[mk.FusedMoEExperts] = select_int8_moe_backend(
config=self.moe,
)
def process_weights_after_loading(self, layer: Module) -> None:
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
self._quantize_weights(layer)
self._setup_kernel(layer)
layer._already_called_process_weights_after_loading = True
def _quantize_weights(self, layer: Module) -> None:
vmax = torch.iinfo(torch.int8).max
w13 = torch.empty_like(layer.w13_weight, dtype=torch.int8)
w2 = torch.empty_like(layer.w2_weight, dtype=torch.int8)
w13_scale = torch.zeros(
layer.num_experts,
layer.w13_weight.shape[1],
device=w13.device,
dtype=torch.float32,
)
w2_scale = torch.zeros(
layer.num_experts,
layer.w2_weight.shape[1],
device=w2.device,
dtype=torch.float32,
)
for expert in range(layer.local_num_experts):
# w13: per-row quantization over hidden_size dim
w = layer.w13_weight[expert, :, :]
scales = w.abs().amax(dim=1) / vmax
q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
w13[expert, :, :] = q.to(torch.int8)
w13_scale[expert, :] = scales
# w2: per-row quantization over intermediate_size dim
w = layer.w2_weight[expert, :, :]
scales = w.abs().amax(dim=1) / vmax
q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
w2[expert, :, :] = q.to(torch.int8)
w2_scale[expert, :] = scales
replace_parameter(layer, "w13_weight", w13)
replace_parameter(layer, "w2_weight", w2)
replace_parameter(layer, "w13_scale", w13_scale)
replace_parameter(layer, "w2_scale", w2_scale)
def _setup_kernel(self, layer: "FusedMoE") -> None:
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
assert self.moe_quant_config is not None
assert self.experts_cls is not None
self.moe_kernel = make_int8_moe_kernel(
moe_quant_config=self.moe_quant_config,
moe_config=self.moe,
experts_cls=self.experts_cls,
routing_tables=layer._maybe_init_expert_routing_tables(),
shared_experts=layer.shared_experts,
)
def get_fused_moe_quant_config(
self, layer: torch.nn.Module
) -> "FusedMoEQuantConfig | None":
quant_config = make_int8_moe_quant_config(
w1_scale=layer.w13_scale,
w2_scale=layer.w2_scale,
)
self._maybe_inject_biases(quant_config, layer)
return quant_config
@@ -0,0 +1,172 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import abstractmethod
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.model_executor.layers.fused_moe import FusedMoEMethodBase
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
from vllm.model_executor.model_loader.reload.layerwise import (
initialize_online_processing,
)
from vllm.model_executor.utils import set_weight_attrs
class OnlineMoEMethodBase(FusedMoEMethodBase):
"""Base for MoE methods that load full-precision weights on meta device
and quantize them after loading via the QeRL layerwise processing system.
"""
uses_meta_device: bool = True
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
layer.num_experts = num_experts
layer.orig_dtype = params_dtype
layer.weight_block_size = None
# Fused gate_up_proj (column parallel) — full precision on meta device
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
device="meta",
dtype=params_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
# down_proj (row parallel) — full precision on meta device
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
device="meta",
dtype=params_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# BIASES (for models like GPT-OSS that have biased MoE)
if self.moe.has_bias:
w13_bias = torch.nn.Parameter(
torch.zeros(
num_experts,
2 * intermediate_size_per_partition,
device="meta",
dtype=layer.orig_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_bias", w13_bias)
set_weight_attrs(w13_bias, extra_weight_attrs)
w2_bias = torch.nn.Parameter(
torch.zeros(
num_experts,
hidden_size,
device="meta",
dtype=layer.orig_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_bias", w2_bias)
set_weight_attrs(w2_bias, extra_weight_attrs)
layer.w13_input_scale = None
layer.w2_input_scale = None
initialize_online_processing(layer)
@abstractmethod
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
pass
def _maybe_inject_biases(
self,
quant_config: FusedMoEQuantConfig,
layer: torch.nn.Module,
) -> None:
"""Inject biases into the quant config if the model has them
(e.g. GPT-OSS biased MoE)."""
if self.moe.has_bias:
w13_bias = getattr(layer, "w13_bias", None)
w2_bias = getattr(layer, "w2_bias", None)
if w13_bias is not None:
quant_config._w1.bias = w13_bias
if w2_bias is not None:
quant_config._w2.bias = w2_bias
def maybe_make_prepare_finalize(
self,
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> mk.FusedMoEPrepareAndFinalizeModular | None:
raise ValueError(
f"{self.__class__.__name__} uses the new modular kernel "
"initialization logic. This function should not be called."
)
@property
def supports_eplb(self) -> bool:
return True
def apply_monolithic(
self,
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
x: torch.Tensor,
router_logits: torch.Tensor,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
assert self.is_monolithic
assert self.moe_kernel is not None
return self.moe_kernel.apply_monolithic(
x,
layer.w13_weight,
layer.w2_weight,
router_logits,
activation=layer.activation,
global_num_experts=layer.global_num_experts,
expert_map=layer.expert_map,
apply_router_weight_on_input=layer.apply_router_weight_on_input,
num_expert_group=layer.num_expert_group,
topk_group=layer.topk_group,
e_score_correction_bias=layer.e_score_correction_bias,
routed_scaling_factor=layer.routed_scaling_factor,
)
def apply(
self,
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
x: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
shared_experts_input: torch.Tensor | None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
assert not self.is_monolithic
assert self.moe_kernel is not None
return self.moe_kernel.apply(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_ids,
activation=layer.activation,
global_num_experts=layer.global_num_experts,
expert_map=layer.expert_map,
apply_router_weight_on_input=layer.apply_router_weight_on_input,
shared_experts_input=shared_experts_input,
)
+3 -2
View File
@@ -1254,9 +1254,10 @@ class Gemma4ForConditionalGeneration(
# computation (using token_type_ids == 0 as text_mask).
# Replicate this: map image token positions to token 0.
if is_multimodal is not None:
is_multimodal = is_multimodal.to(input_ids.device)
ple_input_ids = torch.where(
is_multimodal, torch.zeros_like(input_ids), input_ids
is_multimodal.to(input_ids.device, non_blocking=True),
torch.zeros_like(input_ids),
input_ids,
)
else:
ple_input_ids = input_ids
+23
View File
@@ -99,6 +99,8 @@ class ProjectedParakeet(nn.Module):
if target is None:
target = buffers_dict.get(target_name)
if target is None:
if self._can_skip_missing_named_param(target_name):
continue
raise ValueError(f"Unknown weight: {name}")
weight_loader = getattr(target, "weight_loader", default_weight_loader)
with torch.no_grad():
@@ -107,6 +109,27 @@ class ProjectedParakeet(nn.Module):
return loaded_params
def _can_skip_missing_named_param(self, target_name: str) -> bool:
if self.config.convolution_bias:
return False
# In transformers v5 (not v4), `convolution_bias=False` is
# propagated from parakeet config. If `False`, torch.conv1d will
# *skip registering the param*, thus it will be missing in the
# module's named params. *If* you happen to also have the bias
# tensors in the weights, it will cause a mismatch between the
# weights and the params.
# This allows us to have `convolution_bias=False` in the sound config,
# but still allow for the weights to exist.
return target_name.endswith(
(
".conv.pointwise_conv1.bias",
".conv.depthwise_conv.bias",
".conv.pointwise_conv2.bias",
)
)
EPSILON = 1e-5
LOG_ZERO_GUARD_VALUE = 2**-24
+51 -157
View File
@@ -6,15 +6,16 @@ import os
import platform
import subprocess
import sys
from dataclasses import dataclass
from typing import TYPE_CHECKING
import psutil
import torch
from vllm import envs
from vllm.logger import init_logger
from vllm.utils.ompmultiprocessing import OMPProcessManager
from vllm.utils.cpu_resource_utils import (
DEVICE_CONTROL_ENV_VAR,
get_memory_node_info,
)
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.torch_utils import is_quantized_kv_cache
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -38,49 +39,13 @@ def get_max_threads(pid=0):
raise NotImplementedError("Unsupported OS")
@dataclass
class LogicalCPUInfo:
id: int = -1
physical_core: int = -1
numa_node: int = -1
@classmethod
def _int(cls, value: str) -> int:
try:
int_value = int(value)
except Exception:
int_value = -1
return int_value
@staticmethod
def json_decoder(obj_dict: dict):
id = obj_dict.get("cpu")
physical_core = obj_dict.get("core")
numa_node = obj_dict.get("node")
if not (id is None or physical_core is None or numa_node is None):
return LogicalCPUInfo(
id=LogicalCPUInfo._int(id),
physical_core=LogicalCPUInfo._int(physical_core),
numa_node=LogicalCPUInfo._int(numa_node),
)
else:
return obj_dict
class CpuPlatform(Platform):
_enum = PlatformEnum.CPU
device_name: str = "cpu"
device_type: str = "cpu"
dispatch_key: str = "CPU"
dist_backend: str = "gloo"
device_control_env_var = "CPU_VISIBLE_MEMORY_NODES"
omp_process_manager = None
# Simultaneous Multithreading (SMT) level for OpenMP:
# 4 on PowerPC, 1 on non-PowerPC architectures
smt = 1
global_cpu_mask = None
simulate_numa = int(os.environ.get("_SIM_MULTI_NUMA", 0))
device_control_env_var = DEVICE_CONTROL_ENV_VAR
@property
def supported_dtypes(self) -> list[torch.dtype]:
@@ -123,29 +88,9 @@ class CpuPlatform(Platform):
@classmethod
def get_device_total_memory(cls, device_id: int = 0) -> int:
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.mem_utils import format_gib
meminfo = get_memory_node_info(device_id)
kv_cache_space = envs.VLLM_CPU_KVCACHE_SPACE
node_dir = "/sys/devices/system/node"
if kv_cache_space is None:
nodes = (
[d for d in os.listdir(node_dir) if d.startswith("node")]
if os.path.exists(node_dir)
else []
)
num_numa_nodes = len(nodes) or 1
free_cpu_memory = psutil.virtual_memory().total // num_numa_nodes
DEFAULT_CPU_MEM_UTILIZATION = 0.5
kv_cache_space = int(free_cpu_memory * DEFAULT_CPU_MEM_UTILIZATION)
logger.warning_once(
"VLLM_CPU_KVCACHE_SPACE not set. Using %s GiB for KV cache.",
format_gib(kv_cache_space),
)
else:
kv_cache_space *= GiB_bytes
return kv_cache_space
return meminfo.total_memory
@classmethod
def set_device(cls, device: torch.device) -> None:
@@ -180,6 +125,12 @@ class CpuPlatform(Platform):
"otherwise the performance is not optimized."
)
# Lagecy setting
env_key = "VLLM_CPU_KVCACHE_SPACE"
if env_key in os.environ and os.environ[env_key] != "":
kv_cache_space = int(os.environ[env_key])
cache_config.kv_cache_memory_bytes = kv_cache_space * GiB_bytes
scheduler_config = vllm_config.scheduler_config
# async scheduling is not required on CPU
scheduler_config.async_scheduling = False
@@ -198,8 +149,6 @@ class CpuPlatform(Platform):
)
cache_config.cache_dtype = "auto"
cache_config.cpu_kvcache_space_bytes = CpuPlatform.get_device_total_memory()
parallel_config = vllm_config.parallel_config
# OMP requires the MP executor to function correctly, UniProc is not
# supported as it is not possible to set the OMP environment correctly
@@ -278,21 +227,45 @@ class CpuPlatform(Platform):
os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"
ld_preload_str = os.getenv("LD_PRELOAD", "")
# Intel and CLANG OpenMP setting
if "libiomp5.so" in ld_preload_str or "libomp5" in ld_preload_str:
# The time(milliseconds) that a thread should wait after
# completing the execution of a parallel region, before sleeping.
os.environ["KMP_BLOCKTIME"] = "1"
# Prevents the CPU to run into low performance state
os.environ["KMP_TPAUSE"] = "0"
# Provides fine granularity parallelism
os.environ["KMP_FORKJOIN_BARRIER_PATTERN"] = "dist,dist"
os.environ["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist"
os.environ["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist"
cpu_architecture = Platform.get_cpu_architecture()
if (
platform.system() == "Linux"
and cpu_architecture
in (CpuArchEnum.ARM, CpuArchEnum.POWERPC, CpuArchEnum.X86)
and not (
"libomp" in ld_preload_str
or "libgomp" in ld_preload_str
or "libiomp" in ld_preload_str
)
):
# We need to LD_PRELOAD PyTorch's libgomp, otherwise only
# one core will be properly utilized when we thread-bind
# See: https://github.com/vllm-project/vllm/issues/27369
# TODO: Remove once:
# https://github.com/pytorch/pytorch/issues/166087 is fixed
# We need to find the location of PyTorch's libgomp
torch_pkg = os.path.dirname(torch.__file__)
site_root = os.path.dirname(torch_pkg)
# Search both torch.libs and torch/lib - See:
# https://github.com/vllm-project/vllm/issues/30470
torch_libs_paths = [
os.path.join(site_root, "torch.libs"),
os.path.join(torch_pkg, "lib"),
]
pytorch_libgomp_so_candidates = []
for torch_libs in torch_libs_paths:
pytorch_libgomp_so_candidates.extend(
glob.glob(os.path.join(torch_libs, "libgomp*.so*"))
)
if pytorch_libgomp_so_candidates:
pytorch_libgomp_so = pytorch_libgomp_so_candidates[0]
if ld_preload_str:
ld_preload_str += ":"
ld_preload_str += pytorch_libgomp_so
os.environ["LD_PRELOAD"] = ld_preload_str
# LD_PRELOAD libtcmalloc, bundled under vllm/libs to reduce
# memory allocation overhead
if (
@@ -331,13 +304,6 @@ class CpuPlatform(Platform):
vllm_config.model_config.max_model_len,
vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,
)
# CI specific "quick" NUMA simulation - split all available CPUs
# into a fake NUMA topology
if os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", None) is not None:
os.environ["_SIM_MULTI_NUMA"] = str(
vllm_config.parallel_config.world_size
* vllm_config.parallel_config._api_process_count
)
@classmethod
def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
@@ -345,78 +311,6 @@ class CpuPlatform(Platform):
# Move that logic here so block_size is chosen by the backend.
pass
@classmethod
def get_omp_manager(cls) -> OMPProcessManager:
# initialise the OMP resource management if need be and return the manager
if cls.omp_process_manager is None:
if cls.get_cpu_architecture() == CpuArchEnum.POWERPC:
cls.smt = 4
cls.omp_process_manager = OMPProcessManager(
affinity=cls.get_global_cpu_mask(), smt=cls.smt
)
# we need to fix up the topology returned by the OMP Manager for
# simulated NUMA environments in CI
if cls.simulate_numa > 0:
logger.info(
"Adjusting numa topology to resemble at least %d nodes",
int(cls.simulate_numa),
)
om = cls.omp_process_manager
while len(om.omp_places) < cls.simulate_numa:
new_omp_places = []
touched = False
for omp_place in om.omp_places:
if len(omp_place["mask"]) > 1:
touched = True
cpu_list = sorted(list(omp_place["mask"]))
new_omp_places.append(
{
"mask": set(cpu_list[0 : int(len(cpu_list) / 2)]),
"available": True,
}
)
new_omp_places.append(
{
"mask": set(cpu_list[int(len(cpu_list) / 2) :]),
"available": True,
}
)
if touched:
om.omp_places = new_omp_places
else:
raise ValueError(
"Cannot split the existing NUMA topology to match "
"simulation requirements"
)
return cls.omp_process_manager
@classmethod
def get_global_cpu_mask(cls) -> set[int]:
# get global cpu mask
if cls.global_cpu_mask is None:
if hasattr(os, "sched_getaffinity"):
cls.global_cpu_mask = os.sched_getaffinity(0)
else:
# macOS does not support sched_getaffinity
cpu_count = os.cpu_count() or 1
cls.global_cpu_mask = set(range(cpu_count))
return cls.global_cpu_mask
@classmethod
def reserve_cpus(cls, reserve: set[int]) -> bool:
# remove CPUs from global mask, for now there is no "release" mechanism
if cls.omp_process_manager is not None:
for place in cls.omp_process_manager.omp_places:
if not place["available"]:
return False
cls.global_cpu_mask = cls.get_global_cpu_mask() - reserve
# reinitialize OMP resource management
cls.omp_process_manager = OMPProcessManager(
affinity=cls.global_cpu_mask, smt=cls.smt
)
return True
@classmethod
def discover_numa_topology(cls) -> list[list[int]]:
"""
+18 -5
View File
@@ -108,6 +108,22 @@ def _get_backend_priorities(
AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
]
return [
AttentionBackendEnum.FLASHINFER_MLA,
AttentionBackendEnum.CUTLASS_MLA,
AttentionBackendEnum.FLASH_ATTN_MLA,
AttentionBackendEnum.FLASHMLA,
AttentionBackendEnum.TRITON_MLA,
*sparse_backends,
]
elif device_capability.major == 9:
# Hopper (SM90) — CUTLASS FA3 is highest priority for sparse MLA
# with BF16 KV cache. Falls back to FlashMLA Sparse for FP8.
sparse_backends = [
AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE,
AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
AttentionBackendEnum.FLASHMLA_SPARSE,
]
return [
AttentionBackendEnum.FLASHINFER_MLA,
AttentionBackendEnum.CUTLASS_MLA,
@@ -131,6 +147,7 @@ def _get_backend_priorities(
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TRITON_ATTN,
AttentionBackendEnum.FLEX_ATTENTION,
AttentionBackendEnum.TURBOQUANT,
]
else:
return [
@@ -138,6 +155,7 @@ def _get_backend_priorities(
AttentionBackendEnum.FLASHINFER,
AttentionBackendEnum.TRITON_ATTN,
AttentionBackendEnum.FLEX_ATTENTION,
AttentionBackendEnum.TURBOQUANT,
]
@@ -255,11 +273,6 @@ class CudaPlatformBase(Platform):
valid_backends_priorities = []
invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}
# TurboQuant KV cache: route directly to TQ backend
kv_cache_dtype = attn_selector_config.kv_cache_dtype
if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"):
return [(AttentionBackendEnum.TURBOQUANT, 0)], {}
backend_priorities = _get_backend_priorities(
attn_selector_config.use_mla,
device_capability,
+7 -4
View File
@@ -1258,11 +1258,11 @@ class Qwen3XMLToolParser(ToolParser):
return None
# Parse the delta text and get the result
result = self.parser.parse_single_streaming_chunks(delta_text)
delta = self.parser.parse_single_streaming_chunks(delta_text)
# Update tool call tracking arrays based on incremental parsing results
if result and result.tool_calls:
for tool_call in result.tool_calls:
if delta and delta.tool_calls:
for tool_call in delta.tool_calls:
if tool_call.function:
tool_index = (
tool_call.index
@@ -1292,4 +1292,7 @@ class Qwen3XMLToolParser(ToolParser):
self.streamed_args_for_tool[tool_index] += (
tool_call.function.arguments
)
return result
if delta.content is None and not delta.tool_calls and delta.reasoning is None:
# If no content and no tool calls, return None to indicate no update
return None
return delta
+173
View File
@@ -0,0 +1,173 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import os
import platform
import subprocess
from dataclasses import dataclass
from functools import cache
import psutil
import regex as re
DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES"
@dataclass
class LogicalCPUInfo:
id: int = -1
physical_core: int = -1
numa_node: int = -1
@classmethod
def _int(cls, value: str) -> int:
try:
int_value = int(value)
except Exception:
int_value = -1
return int_value
@staticmethod
def json_decoder(obj_dict: dict):
id = obj_dict.get("cpu")
physical_core = obj_dict.get("core")
numa_node = obj_dict.get("node")
if not (id is None or physical_core is None or numa_node is None):
return LogicalCPUInfo(
id=LogicalCPUInfo._int(id),
physical_core=LogicalCPUInfo._int(physical_core),
numa_node=LogicalCPUInfo._int(numa_node),
)
else:
return obj_dict
@dataclass
class MemoryNodeInfo:
total_memory: int = -1
available_memory: int = -1
def get_memory_affinity(pid: int = 0) -> list[int]:
pid = os.getpid() if pid == 0 else pid
path = f"/proc/{pid}/status"
with open(path) as f:
for line in f:
if line.startswith("Mems_allowed_list:"):
# Extract the string part (e.g., "0-1,3")
raw_list = line.split(":")[1].strip()
return parse_id_list(raw_list)
return []
def parse_id_list(raw_str: str) -> list[int]:
"""Parses strings like '0-2,4,7-8' into [0, 1, 2, 4, 7, 8]"""
result: list[int] = []
if not raw_str:
return result
for part in raw_str.split(","):
if "-" in part:
start, end = map(int, part.split("-"))
result.extend(range(start, end + 1))
else:
result.append(int(part))
return sorted(list(set(result)))
def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:
if platform.system() == "Darwin":
# MacOS has no memory node
return MemoryNodeInfo(
total_memory=psutil.virtual_memory().total,
available_memory=psutil.virtual_memory().available,
)
meminfo_path = f"/sys/devices/system/node/node{node_id}/meminfo"
if not os.path.exists(meminfo_path):
raise RuntimeError(f"{meminfo_path} doesn't exit.")
meminfo = {}
with open(meminfo_path) as f:
for line in f:
# Each line looks like: "Node 0 MemTotal: 97421888 kB"
parts = line.split()
key = parts[2].rstrip(":")
# convert to Bytes
value = int(parts[3]) * 1024
meminfo[key] = value
total_memory = meminfo["MemTotal"]
free_memory = meminfo["MemFree"]
active_file_memory = meminfo["Active(file)"]
inactive_file_memory = meminfo["Inactive(file)"]
reclaimable_memory = meminfo["SReclaimable"]
available_memory = (
free_memory + active_file_memory + inactive_file_memory + reclaimable_memory
)
return MemoryNodeInfo(
total_memory=total_memory,
available_memory=available_memory,
)
def get_allowed_cpu_list() -> list[LogicalCPUInfo]:
cpu_list = _get_cpu_list()
if platform.system() == "Darwin":
return cpu_list
global_allowed_cpu_id_list = os.sched_getaffinity(0)
logical_cpu_list = [x for x in cpu_list if x.id in global_allowed_cpu_id_list]
return logical_cpu_list
def get_visible_memory_node() -> list[int]:
if platform.system() == "Darwin":
return [0]
allowed_memory_node_list = get_memory_affinity()
env_key = DEVICE_CONTROL_ENV_VAR
if (
("VLLM_CPU_SIM_MULTI_NUMA" not in os.environ)
and env_key in os.environ
and os.environ[env_key] != ""
):
visible_nodes = [int(s) for s in os.environ[env_key].split(",")]
visible_nodes = [
node for node in visible_nodes if node in allowed_memory_node_list
]
return visible_nodes
return allowed_memory_node_list
@cache
def _get_cpu_list() -> list[LogicalCPUInfo]:
if platform.system() == "Darwin":
# For MacOS, no user-level CPU affinity and SMT, return all CPUs
cpu_count = os.cpu_count()
assert cpu_count
return [LogicalCPUInfo(i, i, 0) for i in range(cpu_count)]
lscpu_output = subprocess.check_output(
"lscpu -J -e=CPU,CORE,NODE", shell=True, text=True
)
# For platform without NUMA, replace '-' to '0'
lscpu_output = re.sub(r'"node":\s*-\s*(,|\n)', r'"node": 0\1', lscpu_output)
logical_cpu_list: list[LogicalCPUInfo] = json.loads(
lscpu_output, object_hook=LogicalCPUInfo.json_decoder
)["cpus"]
# Filter CPUs with invalid attributes
logical_cpu_list = [
x for x in logical_cpu_list if -1 not in (x.id, x.physical_core, x.numa_node)
]
return logical_cpu_list
+269 -185
View File
@@ -5,196 +5,280 @@ Copyright (c) 2026 Red Hat Inc
Copyright (c) 2026 Cambridge Greys Ltd
"""
import json
import os
import platform
import subprocess
from collections.abc import Callable
from contextlib import contextmanager
from typing import TYPE_CHECKING
import vllm.utils.cpu_resource_utils as cr_utils
from vllm import envs
from vllm.logger import init_logger
from vllm.platforms import CpuArchEnum, current_platform
from vllm.utils.cpu_resource_utils import LogicalCPUInfo
if TYPE_CHECKING:
from vllm.config import VllmConfig
logger = init_logger(__name__)
def _int(arg):
"""Relaxed parsing of ints which handles a - instead of a number.
The lscpu json may contain that for nodes in some cases. If that
is the case we parse it to zero
"""
try:
if int(arg) >= 0:
return int(arg)
except ValueError:
pass
return 0
def parse_mask(mask):
"""Expand a X-Y,Z list"""
result = []
for token in mask.split(","):
try:
start, finish = token.split("-")
if int(start) > int(finish):
raise IndexError("Invalid Indexes for cpu ranges")
for cpu in range(int(start), int(finish) + 1):
result.append(cpu)
except ValueError:
result.append(int(token))
return set(result)
def _get_default_affinity() -> set[int]:
"""Get the set of CPUs the process is allowed to run on."""
if hasattr(os, "sched_getaffinity"):
return os.sched_getaffinity(0)
# macOS does not support sched_getaffinity; fall back to cpu_count
cpu_count = os.cpu_count() or 1
return set(range(cpu_count))
def _get_cpu_topology_json() -> bytes:
"""Get CPU topology as JSON.
On Linux this uses ``lscpu -Je``. On other platforms (e.g. macOS) we
synthesize a simple topology where every logical CPU is its own core
on NUMA node 0, which is sufficient for the OMP place-list builder.
"""
if platform.system() == "Linux":
return subprocess.run(["lscpu", "-Je"], check=True, capture_output=True).stdout
# Fallback for non-Linux (macOS, etc.)
cpu_count = os.cpu_count() or 1
cpus = []
for i in range(cpu_count):
cpus.append({"cpu": str(i), "core": str(i), "node": "0"})
return json.dumps({"cpus": cpus}).encode()
def enumerate_resources(resource_map, mask=None, allowed=None):
"""Enumerate system resources"""
if allowed is None:
allowed = _get_default_affinity()
if mask is not None:
allowed = allowed & mask
try:
allowed_nodes = parse_mask(os.environ["CPU_VISIBLE_MEMORY_NODES"])
except KeyError:
allowed_nodes = None
lscpu: dict[str, dict] = {"cpus": {}, "cores": {}, "nodes": {}}
for cpu in resource_map["cpus"]:
cpunum = int(cpu["cpu"])
if (
cpunum in allowed
and cpunum >= 0
and (allowed_nodes is None or _int(cpu["node"]) in allowed_nodes)
):
lscpu["cpus"][cpunum] = [cpu]
core = _int(cpu["core"])
if lscpu["cores"].get(core, None) is None:
lscpu["cores"][core] = [cpu]
else:
lscpu["cores"][core].append(cpu)
node = _int(cpu["node"])
if lscpu["nodes"].get(node, None) is None:
lscpu["nodes"][node] = [cpu]
else:
lscpu["nodes"][node].append(cpu)
return lscpu
def produce_cpu_list(cpus, smt=1):
"""Produce a CPU list with/without SMT pairs - main cpu list case"""
mask: list[int] = []
for key, value in cpus.items():
exists = 0
for cpu in mask:
if cpu == value[0]["core"]:
exists += 1
break
if exists < smt:
mask.append(int(key))
return {"mask": set(mask), "available": True}
def produce_cpu_sublist(scpus, smt=1):
"""Produce a CPU list with/without SMT pairs - resource leaf case"""
cpu_list: list[dict] = []
for value in scpus:
exists = 0
for cpu in cpu_list:
if int(cpu["core"]) == int(value["core"]):
exists += 1
break
if exists < smt:
cpu_list.append(value)
mask = []
for cpu in cpu_list:
mask.append(int(cpu["cpu"]))
return {"mask": set(mask), "available": True}
def create_omp_places(resources, strategy, smt=True):
"""Parse CPU topology and generate possible CPU masks"""
omp_places = []
if strategy == "all":
omp_places.append(produce_cpu_list(resources["cpus"], smt))
elif strategy == "cores":
for value in resources["cores"].values():
omp_places.append(produce_cpu_sublist(value, smt))
elif strategy == "nodes":
for value in resources["nodes"].values():
omp_places.append(produce_cpu_sublist(value, smt))
else:
raise NotImplementedError("Unknown strategy")
return omp_places
# pylint: disable=too-few-public-methods
class OMPProcessManager:
"""OMP aware wrapper to run mp Process()"""
def __init__(self, config: "VllmConfig"):
if not current_platform.is_cpu():
return
def __init__(self, strategy="nodes", smt=1, mock=None, affinity=None):
self.strategy = strategy
self.smt = smt
self.omp_places = []
vllm_mask = os.environ.get("VLLM_CPU_OMP_THREADS_BIND", None)
self.setup_omp = vllm_mask != "nobind"
if self.setup_omp:
omp_places = []
if vllm_mask is not None:
masks = []
for spec in vllm_mask.split("|"):
masks.append(parse_mask(spec))
self.local_world_size = config.parallel_config.local_world_size
self.local_dp_rank = config.parallel_config.data_parallel_rank_local
# This is a bit tricky because the internal DP size
# is always 1 for non-MoE models
self.internal_dp_size = config.parallel_config._api_process_count
self.simulate_multi_node = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0"
ld_preload_str = os.getenv("LD_PRELOAD", "")
self.use_iomp = "libiomp" in ld_preload_str or "libomp" in ld_preload_str
self.use_gomp = "libgomp" in ld_preload_str
assert not (self.use_iomp and self.use_gomp)
# at least reserve 1/local_world_size(for ARM) core for scheduler
# proc as always use MP executor
# TODO: make scheduler proc sleep when idle
self.reserve_cpu_num = (
self.local_world_size
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM
else 1
)
# reserve at one more core for nixl_connector under p/d case
if config.kv_transfer_config:
self.reserve_cpu_num += 1
if envs.VLLM_CPU_NUM_OF_RESERVED_CPU is not None:
if self.reserve_cpu_num > envs.VLLM_CPU_NUM_OF_RESERVED_CPU:
msg = (
f"VLLM_CPU_NUM_OF_RESERVED_CPU is less than "
"the minimum requirement"
f": {self.reserve_cpu_num} cores"
)
logger.warning(msg=msg)
self.reserve_cpu_num = envs.VLLM_CPU_NUM_OF_RESERVED_CPU
self._parse_omp_threads_bind_env()
assert not self.simulate_multi_node or self.auto_setup
@contextmanager
def configure_omp_envs(self, rank: int, local_rank: int):
if not current_platform.is_cpu() or self.skip_setup:
yield
return
envs_dict = {}
cpu_list = [str(i) for i in self.cpu_lists[local_rank]]
envs_dict["OMP_NUM_THREADS"] = str(len(cpu_list))
if self.use_iomp:
# set IOMP envs
cpu_list_str = ",".join(cpu_list)
envs_dict["KMP_AFFINITY"] = (
f"granularity=fine,explicit,proclist=[{cpu_list_str}]"
)
# The time(milliseconds) that a thread should wait after
# completing the execution of a parallel region, before sleeping.
envs_dict["KMP_BLOCKTIME"] = "1"
# Prevents the CPU to run into low performance state
envs_dict["KMP_TPAUSE"] = "0"
# Provides fine granularity parallelism
envs_dict["KMP_FORKJOIN_BARRIER_PATTERN"] = "dist,dist"
envs_dict["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist"
envs_dict["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist"
elif self.use_gomp:
# set GOMP envs
# likes '0 1 2 ...'
cpu_list_str = " ".join(cpu_list)
envs_dict["GOMP_CPU_AFFINITY"] = cpu_list_str
else:
# set OMP envs
# likes '{0,1,2,...}'
cpu_list_str = ",".join(cpu_list)
envs_dict["OMP_PLACES"] = f"{{{cpu_list_str}}}"
envs_dict["OMP_PROC_BIND"] = "true"
# backup envs
old_envs_dict = {}
for k in envs_dict:
old_envs_dict[k] = os.environ.get(k)
try:
# set envs
for k, v in envs_dict.items():
os.environ[k] = v
yield
finally:
# restore old envs
for k, v in old_envs_dict.items(): # type: ignore
if v is None:
os.environ.pop(k, None)
else:
os.environ[k] = v
def _parse_omp_threads_bind_env(self):
vllm_mask = envs.VLLM_CPU_OMP_THREADS_BIND
self.skip_setup = vllm_mask == "nobind"
self.auto_setup = vllm_mask == "auto"
self.reserved_cpu_list = []
self.cpu_lists = []
if self.auto_setup:
# auto generate CPU lists
cpu_arch = current_platform.get_cpu_architecture()
if cpu_arch == CpuArchEnum.POWERPC:
# For POWERPC SMT-8/4/2
cpu_list, reserve_list = self._get_autobind_cpu_ids(
lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4]
)
elif cpu_arch in (CpuArchEnum.X86, CpuArchEnum.S390X):
# For x86/S390X SMT-2, use 1 logical CPU per physical core
cpu_list, reserve_list = self._get_autobind_cpu_ids(
lambda cpus: cpus[-1:]
)
elif cpu_arch == CpuArchEnum.ARM:
# For AArch64, no SMT, use all logical CPU
cpu_list, reserve_list = self._get_autobind_cpu_ids(lambda cpus: cpus)
else:
masks = [None]
if mock is None:
data = _get_cpu_topology_json()
else:
with open(mock, mode="rb") as jf:
data = jf.read()
lscpu = json.loads(data)
for mask in masks:
resources = enumerate_resources(lscpu, mask, affinity)
omp_places.extend(create_omp_places(resources, strategy, smt))
self.omp_places = sorted(
omp_places,
key=lambda p: "{:04d}-{:04d}".format(len(p["mask"]), max(p["mask"])),
reverse=True,
cpu_list, reserve_list = [], []
raise RuntimeError(f"{cpu_arch} doesn't support auto CPU binding.")
for item in cpu_list:
self.cpu_lists.append([x.id for x in item])
self.reserved_cpu_list = [x.id for x in reserve_list]
elif not self.skip_setup:
# user defined CPU lists
omp_cpuids_list = vllm_mask.split("|")
if self.local_dp_rank is not None:
local_dp_rank = self.local_dp_rank
world_size = self.local_world_size
# Rank mapping [DP, PP, TP]
omp_cpuids_list = omp_cpuids_list[
local_dp_rank * world_size : (local_dp_rank + 1) * world_size
]
assert len(omp_cpuids_list) == self.local_world_size, (
"Given "
f"number of CPU id list {omp_cpuids_list} doesn't match "
f"local world size {self.local_world_size}."
)
def run(self, what, *args, **kwargs):
"""Run arg with correct OMP environment"""
if self.setup_omp:
for place in self.omp_places:
if place["available"]:
reserve = int(os.environ.get("VLLM_CPU_NUM_OF_RESERVED_CPU", 0))
place["available"] = False
# pylint: disable=consider-using-f-string
os.environ["OMP_PLACES"] = "{}".format(place["mask"])
os.environ["OMP_NUM_THREADS"] = "{}".format(
len(place["mask"]) - reserve
)
os.environ["OMP_PROC_BIND"] = "TRUE"
return what(*args, **kwargs)
raise IndexError("Out of OMP places")
return what(*args, **kwargs)
# parse CPU list strings like "5,2-4" to [5, 2, 3, 4]
self.cpu_lists = [cr_utils.parse_id_list(s) for s in omp_cpuids_list]
else:
# skip
self.cpu_lists = []
msg = "OpenMP thread binding info: \n"
for i in range(self.local_world_size):
msg += f"\tlocal_rank={i}, core ids={self.cpu_lists[i]}\n"
msg += f"\treserved_cpus={self.reserved_cpu_list}"
logger.info(msg)
def _get_autobind_cpu_ids(
self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]]
) -> tuple[list[list[LogicalCPUInfo]], list[LogicalCPUInfo]]:
"""
Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for
other processes.
Currently for rank N, only CPU ids on the N-th node in available NUMA
node list will be selected.
Args:
cpu_selector: a callable object to select CPUs from a CPU list
of a physical core. The input is a LogicalCPUInfo list contains
logical CPUs of a physical CPU, sorted by the LogicalCPUInfo.id.
A selected LogicalCPUInfo list should be returned.
"""
# this memory node list has been sliced for DP offset
allowed_numa_nodes = cr_utils.get_visible_memory_node()
logical_cpu_list = cr_utils.get_allowed_cpu_list()
local_world_size = self.local_world_size
assert (
len(allowed_numa_nodes) >= local_world_size or self.simulate_multi_node
), (
f"Not enough allowed NUMA nodes to bind threads of "
f"{local_world_size} local CPUWorkers. "
f"Allowed NUMA nodes are {allowed_numa_nodes}. "
"Please try to bind threads manually or decrease DP/TP/PP."
)
# Generate OMP CPU list for each rank
cpu_lists_of_ranks = []
reserved_cpu_list = []
total_cpu_num = 0
for local_rank in range(self.local_world_size):
if not self.simulate_multi_node:
selected_numa_node = allowed_numa_nodes[local_rank]
selected_logical_cpu_list = [
x for x in logical_cpu_list if x.numa_node == selected_numa_node
]
else:
world_size_across_dp = self.local_world_size * self.internal_dp_size
assert len(logical_cpu_list) >= world_size_across_dp
selected_logical_cpu_list = sorted(
logical_cpu_list, key=lambda x: x.numa_node
)
sim_cpu_num_per_node = (
len(selected_logical_cpu_list) // world_size_across_dp
)
assert self.local_dp_rank is not None
start_idx = (
local_rank + self.local_world_size * self.local_dp_rank
) * sim_cpu_num_per_node
selected_logical_cpu_list = selected_logical_cpu_list[
start_idx : (start_idx + sim_cpu_num_per_node)
]
# Select logical CPUs on same physical cores via cpu_selector
core_to_cpus: dict[int, list[LogicalCPUInfo]] = {}
for cpu_info in selected_logical_cpu_list:
if cpu_info.physical_core not in core_to_cpus:
core_to_cpus[cpu_info.physical_core] = []
core_to_cpus[cpu_info.physical_core].append(cpu_info)
selected_logical_cpu_list = []
for cpu_list in core_to_cpus.values():
cpu_list = sorted(cpu_list, key=lambda x: x.id)
selected_logical_cpu_list.extend(cpu_selector(cpu_list))
# sort selected cores based on core id
selected_logical_cpu_list = sorted(
selected_logical_cpu_list, key=lambda x: x.id
)
cpu_lists_of_ranks.append(selected_logical_cpu_list)
total_cpu_num += len(selected_logical_cpu_list)
# Reserve CPUs for other processes
if total_cpu_num <= self.reserve_cpu_num:
logger.warning(
"Selected CPU core number (%s) "
"should be greater than reserved CPU core "
"number (%s).",
total_cpu_num,
self.reserve_cpu_num,
)
return cpu_lists_of_ranks, []
reserve_num_per_rank = [
self.reserve_cpu_num // self.local_world_size
] * self.local_world_size
# last rank first
for i in range(
self.local_world_size - 1,
self.local_world_size - 1 - self.reserve_cpu_num % self.local_world_size,
-1,
):
reserve_num_per_rank[i] += 1
for i in range(self.local_world_size):
num = reserve_num_per_rank[i]
if num > 0:
reserved_cpu_list.extend(cpu_lists_of_ranks[i][-num:])
cpu_lists_of_ranks[i] = cpu_lists_of_ranks[i][:-num]
return cpu_lists_of_ranks, reserved_cpu_list
+122 -12
View File
@@ -46,6 +46,7 @@ STR_DTYPE_TO_TORCH_DTYPE = {
"turboquant_4bit_nc": torch.uint8,
"turboquant_k3v4_nc": torch.uint8,
"turboquant_3bit_nc": torch.uint8,
"nvfp4": torch.uint8,
}
TORCH_DTYPE_TO_NUMPY_DTYPE = {
@@ -59,17 +60,19 @@ TORCH_DTYPE_TO_NUMPY_DTYPE = {
MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {
# TODO: Add more modelopt kv cache dtype
# mappings here when it supported by some attention backend
# (for example supports nvfp4).
"fp8": "fp8_e4m3",
"nvfp4": "nvfp4",
}
T = TypeVar("T")
def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
return kv_cache_dtype.startswith("fp8") or kv_cache_dtype.endswith("per_token_head")
return (
kv_cache_dtype.startswith("fp8")
or kv_cache_dtype.endswith("per_token_head")
or kv_cache_dtype == "nvfp4"
)
def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
@@ -299,6 +302,8 @@ def get_kv_cache_quant_algo_string(quant_cfg: dict[str, Any]) -> str | None:
and kv_algo.get("type") == "float"
):
kv_algo = "fp8"
elif kv_algo.get("num_bits") == 4 and kv_algo.get("type") == "float":
kv_algo = "nvfp4"
else:
# Unknown/unsupported format - return "auto" as safe fallback
logger.warning(
@@ -375,6 +380,95 @@ def set_random_seed(seed: int | None) -> None:
current_platform.manual_seed_all(seed)
def nvfp4_kv_cache_full_dim(head_size: int) -> int:
"""Packed last dim for NVFP4 KV cache: fp4 data + fp8 block scales."""
return head_size // 2 + head_size // 16
def _nvfp4_split_data_scale(
kv_side: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Split a single NVFP4 KV-side buffer into data and scale views.
The input is a 4D tensor for one KV side (K or V) whose last
dimension is ``full_dim = data_dim + scale_dim``. The physical
layout within each side is [data | scale], both packed contiguously.
Args:
kv_side: 4D uint8 tensor with shape
``(num_pages, dim_1, dim_2, full_dim)``.
May be in any permutation order (NHD or HND).
Returns:
``(data, scale)`` where
``data`` is a uint8 view with shape
``(num_pages, dim_1, dim_2, data_dim)``.
``scale`` is a float8_e4m3fn view with shape
``(num_pages, dim_1, dim_2, scale_dim)``.
"""
num_pages = kv_side.shape[0]
dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2]
full_dim = kv_side.shape[3]
data_dim = full_dim * 8 // 9
scale_dim = full_dim - data_dim
data_per_kv = dim_1 * dim_2 * data_dim
page_bytes = kv_side.stride(0)
# Derive inner strides from the kv_side strides, scaling by the
# ratio of the target dim to full_dim. This preserves the physical
# layout (NHD vs HND) encoded in the input tensor's strides.
s1 = kv_side.stride(1) * data_dim // full_dim
s2 = kv_side.stride(2) * data_dim // full_dim
data_shape = (num_pages, dim_1, dim_2, data_dim)
data_strides = (page_bytes, s1, s2, 1)
s1_s = kv_side.stride(1) * scale_dim // full_dim
s2_s = kv_side.stride(2) * scale_dim // full_dim
scale_shape = (num_pages, dim_1, dim_2, scale_dim)
scale_strides = (page_bytes, s1_s, s2_s, 1)
base = kv_side.storage_offset()
data = torch.as_strided(kv_side, data_shape, data_strides, storage_offset=base)
scale = torch.as_strided(
kv_side, scale_shape, scale_strides, storage_offset=base + data_per_kv
).view(torch.float8_e4m3fn)
return data, scale
def nvfp4_kv_cache_split_views(kv_cache: torch.Tensor) -> tuple[tuple, tuple]:
"""Split an NVFP4 KV cache tensor into data and scale views.
Accepts either a 5D tensor ``(num_pages, 2, dim_2, dim_3, full_dim)``
or a 4D single-side tensor ``(num_pages, dim_2, dim_3, full_dim)``.
Per-page layout: [K_data | K_scale | V_data | V_scale].
Each KV side is self-contained (data followed by its scale), so the
5D case simply splits each side independently.
The returned views are in the same dim order as the input (NHD or
HND), so callers get views matching whichever order they passed in.
Args:
kv_cache: 5D or 4D uint8 tensor where the last dimension is
``full_dim = data_dim + scale_dim = 9 * head_size / 16``.
Returns:
For 5D input:
``(k_data, v_data), (k_scale, v_scale)``
For 4D input (single KV side):
``(data,), (scale,)``
"""
if kv_cache.dim() == 4:
data, scale = _nvfp4_split_data_scale(kv_cache)
return (data,), (scale,)
k_data, k_scale = _nvfp4_split_data_scale(kv_cache[:, 0])
v_data, v_scale = _nvfp4_split_data_scale(kv_cache[:, 1])
return (k_data, v_data), (k_scale, v_scale)
def create_kv_caches_with_random_flash(
num_blocks: int,
block_size: int,
@@ -401,15 +495,31 @@ def create_kv_caches_with_random_flash(
value_caches: list[torch.Tensor] = []
for _ in range(num_layers):
key_value_cache = torch.empty(
size=kv_cache_allocation_shape, dtype=dtype, device=device
).permute(*stride_order)
if cache_dtype in ["auto", "half", "bfloat16", "float"]:
key_value_cache.uniform_(-scale, scale)
elif cache_dtype == "fp8":
_generate_random_fp8(key_value_cache, -scale, scale)
if cache_dtype == "nvfp4":
# Full page dim: fp4 data + fp8 block scales per head.
# Per page layout: [K_data | K_scale | V_data | V_scale]
# Returns [:, 0] and [:, 1] like all other dtypes.
full_dim = nvfp4_kv_cache_full_dim(head_size)
nvfp4_shape = (num_blocks, 2, block_size, num_heads, full_dim)
nvfp4_phys = tuple(nvfp4_shape[i] for i in stride_order)
inv = [stride_order.index(i) for i in range(len(stride_order))]
key_value_cache = torch.randint(
0,
256,
nvfp4_phys,
dtype=dtype,
device=device,
).permute(*inv)
else:
raise ValueError(f"Does not support key cache of type {cache_dtype}")
key_value_cache = torch.empty(
size=kv_cache_allocation_shape, dtype=dtype, device=device
).permute(*stride_order)
if cache_dtype in ["auto", "half", "bfloat16", "float"]:
key_value_cache.uniform_(-scale, scale)
elif cache_dtype == "fp8":
_generate_random_fp8(key_value_cache, -scale, scale)
else:
raise ValueError(f"Does not support key cache of type {cache_dtype}")
key_caches.append(key_value_cache[:, 0])
value_caches.append(key_value_cache[:, 1])
return key_caches, value_caches
+61 -11
View File
@@ -42,7 +42,12 @@ from vllm.utils.flashinfer import (
)
from vllm.utils.math_utils import cdiv
from vllm.utils.platform_utils import is_pin_memory_available
from vllm.utils.torch_utils import is_quantized_kv_cache, is_strictly_contiguous
from vllm.utils.torch_utils import (
is_quantized_kv_cache,
is_strictly_contiguous,
nvfp4_kv_cache_full_dim,
nvfp4_kv_cache_split_views,
)
from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
@@ -355,6 +360,10 @@ class FlashInferBackend(AttentionBackend):
head_size: int,
cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
if cache_dtype_str == "nvfp4":
# Packed layout: fp4 data + fp8 block scales in last dim
last_dim = nvfp4_kv_cache_full_dim(head_size)
return (num_blocks, 2, block_size, num_kv_heads, last_dim)
return (num_blocks, 2, block_size, num_kv_heads, head_size)
@staticmethod
@@ -608,11 +617,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
self.cache_dtype = self.cache_config.cache_dtype
# Cannot use self.kv_cache_spec.dtype here because kv_cache_spec
# storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
self.cache_dtype
)
self.is_kvcache_nvfp4 = self.cache_dtype == "nvfp4"
if self.is_kvcache_nvfp4:
# For NVFP4, kv_cache_dtype stays as the string "nvfp4"
# which is passed to FlashInferImpl
self.kv_cache_dtype = self.cache_dtype
raise NotImplementedError("nvfp4 KV cache is not yet supported")
else:
self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
self.cache_dtype
)
else:
self.cache_dtype = "auto"
self.is_kvcache_nvfp4 = False
assert self.kv_cache_spec.dtype == self.model_config.dtype
self.kv_cache_dtype = self.kv_cache_spec.dtype
@@ -626,7 +643,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
can_use_trtllm
and not vllm_config.attention_config.disable_flashinfer_q_quantization
):
self.q_data_type = self.kv_cache_dtype
if self.is_kvcache_nvfp4:
# NVFP4 KV cache uses FP8 quantized queries
self.q_data_type = FlashInferBackend.get_fp8_dtype_for_flashinfer(
"fp8_e4m3"
)
else:
self.q_data_type = self.kv_cache_dtype
else:
self.q_data_type = self.model_config.dtype
@@ -1228,6 +1251,8 @@ class FlashInferImpl(AttentionImpl):
self.sliding_window[0] if self.sliding_window is not None else -1
)
self.kv_cache_dtype = kv_cache_dtype
self.is_kvcache_nvfp4 = kv_cache_dtype == "nvfp4"
self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0
self.logits_soft_cap = logits_soft_cap
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
@@ -1406,7 +1431,16 @@ class FlashInferImpl(AttentionImpl):
num_prefill_tokens = attn_metadata.num_prefill_tokens
stride_order = FlashInferBackend.get_kv_cache_stride_order()
kv_cache_permute = kv_cache.permute(*stride_order)
kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous
# For NVFP4, the kv_cache last dim is full_dim (data + scale packed).
# Split into correctly-strided data and scale views.
nvfp4_kv_data = None
nvfp4_kv_block_scales = None
if self.is_kvcache_nvfp4:
nvfp4_kv_data, nvfp4_kv_block_scales = nvfp4_kv_cache_split_views(
kv_cache_permute
)
use_dcp = self.dcp_world_size > 1
@@ -1490,8 +1524,20 @@ class FlashInferImpl(AttentionImpl):
assert self.o_sf_scale is None
out = output[num_decode_tokens:]
if attn_metadata.q_data_type != FP8_DTYPE and is_quantized_kv_cache(
self.kv_cache_dtype
prefill_kv_block_scales = None
if self.is_kvcache_nvfp4:
# NVFP4 trtllm-gen kernel requires FP8 query.
assert attn_metadata.q_data_type == FP8_DTYPE, (
"NVFP4 KV cache requires FP8 quantized queries for "
"trtllm-gen prefill. Set "
"disable_flashinfer_q_quantization=False."
)
mock_kv_cache = nvfp4_kv_data
mock_block_table = block_tables_prefill
prefill_kv_block_scales = nvfp4_kv_block_scales # noqa: F841
elif (
attn_metadata.q_data_type != FP8_DTYPE
and self.kv_cache_dtype.startswith("fp8")
):
# TRTLLM prefill attention does not support BF16 Q
# and fp8 kv cache. So to enable prefill attention
@@ -1636,7 +1682,9 @@ class FlashInferImpl(AttentionImpl):
trtllm_batch_decode_with_kv_cache(
query=decode_query,
kv_cache=kv_cache_permute,
kv_cache=nvfp4_kv_data
if self.is_kvcache_nvfp4
else kv_cache_permute,
workspace_buffer=workspace_buffer,
block_tables=block_tables_decode,
seq_lens=seq_lens_decode,
@@ -1667,11 +1715,13 @@ class FlashInferImpl(AttentionImpl):
# and value[:num_actual_tokens] because the reshape_and_cache_flash
# op uses the slot_mapping's shape to determine the number of
# actual tokens.
k_cache = kv_cache[:, 0]
v_cache = kv_cache[:, 1]
torch.ops._C_cache_ops.reshape_and_cache_flash(
key,
value,
kv_cache[:, 0],
kv_cache[:, 1],
k_cache,
v_cache,
slot_mapping,
self.kv_cache_dtype,
layer._k_scale,
@@ -0,0 +1,633 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CUTLASS FA3 Sparse MLA Attention Backend for vLLM.
This backend uses the vendored CUTLASS FlashAttention3 Sm90 kernel from
sgl-attn to implement sparse MLA attention for DeepSeek-V3.2 and similar
models on SM90 (Hopper) GPUs.
Key differences from FlashMLASparseBackend:
- Uses BF16 KV cache (576 bytes/token) instead of FP8 (656 bytes/token)
- No head padding needed (FA3 handles arbitrary head counts natively)
- Accepts Q_rope and Q_nope (qv) separately (no ConcatMLAQ kernel)
- 3 sub-kernels: scheduler + main attention + combine
- ~4x faster per transformer block (~16us vs ~64us)
All execution modes (decode, prefill, mixed) are handled identically:
each token is treated as an independent batch element with seqlen=1.
This simplifies metadata building and CUDA graph support.
Backend priority: Highest for SM90 with kv_cache_dtype="auto".
Graceful fallback to FlashMLA Sparse when FP8 cache requested or non-SM90.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, ClassVar
import numpy as np
import torch
from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
AttentionLayer,
AttentionMetadata,
AttentionMetadataBuilder,
CommonAttentionMetadata,
SparseMLAAttentionImpl,
)
from vllm.v1.attention.ops.cutlass_fa3 import is_cutlass_fa3_available
logger = logging.getLogger(__name__)
# Maximum batch size (number of tokens) for which CUTLASS FA3 is used.
# For larger batch sizes, fall back to FlashMLA BF16 sparse prefill kernel.
# FA3 is ~4x faster than FlashMLA for small batches (bs<=16) but regresses
# for larger batches due to higher per-token overhead from the 3-kernel
# launch pattern (scheduler + main + combine) and page_size=1 layout.
MAX_BATCH_SIZE_FOR_FA3 = 16
# FlashMLA sparse prefill kernel requires num_heads padded to this multiple
# on SM90 (Hopper). SM100 (Blackwell) requires 128.
_FLASHMLA_SM90_HEAD_PADDING = 64
# Check if FlashMLA BF16 sparse kernel is available for fallback
_flashmla_sparse_available = False
try:
from vllm.v1.attention.ops.flashmla import flash_mla_sparse_fwd
_flashmla_sparse_available = True
except (ImportError, Exception):
pass
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.config.cache import CacheDType
from vllm.model_executor.layers.linear import ColumnParallelLinear
from vllm.platforms.interface import DeviceCapability
from vllm.v1.kv_cache_interface import AttentionSpec
# ─── Backend Class ────────────────────────────────────────────────────
class CutlassFA3MLASparseBackend(AttentionBackend):
"""CUTLASS FA3 sparse MLA for SM90 (Hopper). BF16 KV cache only.
When FP8 cache is requested, vLLM's backend selection falls back to
FlashMLASparseBackend automatically since this backend only supports
kv_cache_dtype="auto" (which maps to BF16 for MLA).
"""
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
@staticmethod
def get_supported_kernel_block_sizes() -> list[int]:
return [64]
@staticmethod
def get_name() -> str:
return "CUTLASS_FA3_MLA_SPARSE"
@staticmethod
def get_builder_cls() -> type[CutlassFA3MLASparseMetadataBuilder]:
return CutlassFA3MLASparseMetadataBuilder
@staticmethod
def get_impl_cls() -> type[CutlassFA3MLASparseImpl]:
return CutlassFA3MLASparseImpl
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
return [576]
@classmethod
def is_mla(cls) -> bool:
return True
@classmethod
def is_sparse(cls) -> bool:
return True
@classmethod
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
return capability.major == 9
@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
# BF16 cache: 576 bf16 elements per token = 1152 bytes
# Layout per token: [kv_c_normed(512 bf16) | k_pe(64 bf16)]
return (num_blocks, block_size, head_size)
@classmethod
def validate_configuration(
cls,
head_size: int,
dtype: torch.dtype,
kv_cache_dtype: CacheDType | None,
block_size: int | None,
use_mla: bool,
has_sink: bool,
use_sparse: bool,
use_mm_prefix: bool,
use_per_head_quant_scales: bool,
device_capability: DeviceCapability,
attn_type: str,
use_non_causal: bool = False,
) -> list[str]:
invalid = super().validate_configuration(
head_size=head_size,
dtype=dtype,
kv_cache_dtype=kv_cache_dtype,
block_size=block_size,
use_mla=use_mla,
has_sink=has_sink,
use_sparse=use_sparse,
use_mm_prefix=use_mm_prefix,
use_per_head_quant_scales=use_per_head_quant_scales,
device_capability=device_capability,
attn_type=attn_type,
use_non_causal=use_non_causal,
)
if not is_cutlass_fa3_available():
invalid.append("_cutlass_fa3_C not available (requires CUDA >= 12.4, SM90)")
return invalid
# ─── Metadata ─────────────────────────────────────────────────────────
@dataclass
class CutlassFA3MLASparseMetadata(AttentionMetadata):
"""Flat metadata for CUTLASS FA3 sparse MLA attention.
ALL tokens (decode/prefill/mixed) are treated as independent batch
elements with seqlen=1. There are no nested Decode/Prefill sub-objects.
This simplification is valid because:
- Sparse MLA always routes through forward_mqa (not forward_mha)
- Each token independently selects its top-K KV positions
- The FA3 kernel handles variable-length sequences via cu_seqlens
"""
num_reqs: int
max_query_len: int
max_seq_len: int
num_actual_tokens: int
query_start_loc: torch.Tensor
slot_mapping: torch.Tensor
block_table: torch.Tensor # [num_reqs, max_blocks_per_req] int32
req_id_per_token: torch.Tensor # [T] int32
block_size: int = 64
topk_tokens: int = 2048
# FA3-specific metadata (pre-allocated for CUDA graph safety)
cache_seqlens: torch.Tensor | None = None # [T] int32
cu_seqlens_q: torch.Tensor | None = None # [T+1] int32
cu_seqlens_k: torch.Tensor | None = None # [T+1] int32
# For MLAAttention.forward_impl() routing: sparse -> all MQA
# Setting num_decodes = num_reqs ensures all tokens go through
# the forward_mqa path (no MHA prefill path).
num_decodes: int | None = 0
num_decode_tokens: int | None = 0
num_prefills: int | None = 0
num_prefill_tokens: int | None = 0
# ─── Metadata Builder ─────────────────────────────────────────────────
class CutlassFA3MLASparseMetadataBuilder(
AttentionMetadataBuilder[CutlassFA3MLASparseMetadata]
):
"""Builds CutlassFA3MLASparseMetadata from CommonAttentionMetadata.
Key design choices:
- Pre-allocates GPU buffers in __init__ for CUDA graph compatibility
- All tokens (decode + prefill) treated as independent seqlen=1 elements
- Uses in-place .copy_() for buffer updates (safe for CUDA graph replay)
"""
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
def __init__(
self,
kv_cache_spec: AttentionSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
) -> None:
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
self.topk_tokens = 2048
max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
self.block_size = kv_cache_spec.block_size
# Enable speculative decoding support
self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
# Pre-allocate GPU buffers (persist across CUDA graph replays).
# These are updated in-place via .copy_() before each replay.
self.req_id_buf = torch.zeros(max_tokens, dtype=torch.int32, device=device)
self.cache_seqlens_buf = torch.ones(
max_tokens, dtype=torch.int32, device=device
)
self.cu_seqlens_q_buf = torch.arange(
0, max_tokens + 1, dtype=torch.int32, device=device
)
self.cu_seqlens_k_buf = torch.zeros(
max_tokens + 1, dtype=torch.int32, device=device
)
def build(
self,
common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False,
) -> CutlassFA3MLASparseMetadata:
"""Build metadata from common attention metadata.
Converts the request-level metadata into per-token flat metadata:
- req_id_per_token: maps each token to its request index
- cache_seqlens: min(seq_len, topk) per token (topk-clipped)
- cu_seqlens_q: [0, 1, 2, ..., T] (each token = seqlen 1)
- cu_seqlens_k: cumsum of cache_seqlens
"""
cm = common_attn_metadata
T = cm.num_actual_tokens
starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
seg_lens = np.diff(starts)
# req_id_per_token: map each token -> request index
req_ids = np.repeat(np.arange(len(seg_lens), dtype=np.int32), seg_lens)
# CUDA graph padding fix: T = cm.num_actual_tokens may include
# padding tokens (e.g., T=32 when only 31 real tokens exist).
# The computed req_ids array has sum(seg_lens) elements which
# equals the real (unpadded) token count. We must:
# 1) Zero-fill the entire buffer first (safe default for padding)
# 2) Copy only the actual data using req_ids.shape[0]
# 3) Slice to padded T for the metadata return
# This matches the pattern used by FlashMLASparseMetadataBuilder,
# FlashInferMLASparseMetadataBuilder, and all other sparse backends.
actual_tokens = req_ids.shape[0]
self.req_id_buf.fill_(0)
self.req_id_buf[:actual_tokens].copy_(
torch.from_numpy(req_ids), non_blocking=True
)
# cache_seqlens: UPPER BOUND = min(seq_len, topk) per token.
# NOTE: This is a per-REQUEST uniform value, NOT the correct
# per-token causal seqlen. For prefill, token i at position p
# can only attend to min(p+1, topk) entries, but this gives
# all tokens min(seq_len, topk). The actual per-token
# cache_seqlens is computed in forward_mqa() using valid_counts
# from the index conversion kernel, which correctly reflects
# the number of valid KV entries per token.
seq_lens_np = np.asarray(cm.seq_lens_cpu, dtype=np.int32)
per_tok_seqlens = np.minimum(np.repeat(seq_lens_np, seg_lens), self.topk_tokens)
# Same CUDA graph padding fix: zero-fill then copy actual data.
# Default to 1 (safe minimum seqlen for FA3 kernel).
self.cache_seqlens_buf.fill_(1)
self.cache_seqlens_buf[:actual_tokens].copy_(
torch.from_numpy(per_tok_seqlens), non_blocking=True
)
# cu_seqlens_q: [0, 1, 2, ..., T] — each token is seqlen=1
cu_q = self.cu_seqlens_q_buf[: T + 1]
# cu_seqlens_k: cumsum(cache_seqlens)
self.cu_seqlens_k_buf[0] = 0
self.cu_seqlens_k_buf[1 : T + 1].copy_(
torch.cumsum(self.cache_seqlens_buf[:T], dim=0)
)
cu_k = self.cu_seqlens_k_buf[: T + 1]
return CutlassFA3MLASparseMetadata(
num_reqs=cm.num_reqs,
max_query_len=cm.max_query_len,
max_seq_len=cm.max_seq_len,
num_actual_tokens=T,
query_start_loc=cm.query_start_loc,
slot_mapping=cm.slot_mapping,
block_table=cm.block_table_tensor,
req_id_per_token=self.req_id_buf[:T],
block_size=self.block_size,
topk_tokens=self.topk_tokens,
cache_seqlens=self.cache_seqlens_buf[:T],
cu_seqlens_q=cu_q,
cu_seqlens_k=cu_k,
# Route ALL tokens through MQA in forward_impl
num_decodes=cm.num_reqs,
num_decode_tokens=T,
num_prefills=0,
num_prefill_tokens=0,
)
# ─── Implementation ───────────────────────────────────────────────────
class CutlassFA3MLASparseImpl(SparseMLAAttentionImpl[CutlassFA3MLASparseMetadata]):
"""CUTLASS FA3 sparse MLA attention implementation.
This implementation replaces the FlashMLA C sparse_attn_fwd_kernel
with the CUTLASS FA3 Sm90 kernel from sgl-attn, providing ~4x speedup
per transformer block on Hopper GPUs.
Key advantages over FlashMLASparseImpl:
- No head padding (FA3 handles arbitrary head counts natively)
- No Q concatenation kernel (FA3 accepts q_rope and qv separately)
- BF16 KV cache (smaller footprint, no dequantization overhead)
- SM90 warpgroup MMA + TMA for higher compute efficiency
"""
supports_quant_query_input: bool = False
def __init__(
self,
num_heads: int,
head_size: int,
scale: float,
num_kv_heads: int,
alibi_slopes: list[float] | None,
sliding_window: int | None,
kv_cache_dtype: str,
logits_soft_cap: float | None,
attn_type: str,
kv_sharing_target_layer_name: str | None,
# MLA Specific Arguments
q_lora_rank: int | None = None,
kv_lora_rank: int = 512,
qk_nope_head_dim: int = 128,
qk_rope_head_dim: int = 64,
qk_head_dim: int = 192,
v_head_dim: int = 128,
kv_b_proj: ColumnParallelLinear | None = None,
indexer: object | None = None,
q_pad_num_heads: int | None = None,
**kwargs,
) -> None:
self.num_heads = num_heads # 16 (per GPU for TP=8)
self.head_size = head_size # 576 (kv_lora_rank + qk_rope_head_dim)
self.scale = float(scale) # 192**-0.5
self.num_kv_heads = num_kv_heads # 1 (MQA)
self.kv_cache_dtype = kv_cache_dtype # "auto" (maps to BF16)
self.kv_lora_rank = kv_lora_rank
self.qk_rope_head_dim = qk_rope_head_dim
self.softmax_scale = scale
self.topk_tokens = 2048
self.num_splits = 0 # auto; CUDA-graph safe (deterministic per bs)
self.logits_soft_cap = float(logits_soft_cap) if logits_soft_cap else 0.0
# The indexer provides topk_indices_buffer shared across layers
assert indexer is not None, (
"CutlassFA3MLASparseImpl requires an indexer "
"for sparse top-K index selection"
)
self.topk_indices_buffer = indexer.topk_indices_buffer
# DCP (Decode Context Parallelism) requires softmax LSE from the
# attention kernel. FA3's return_softmax_lse=True is not yet wired
# through this backend. When DCP is needed, fall back to FlashMLA.
# TODO: Wire return_softmax_lse=True through forward_mqa for DCP.
def forward_mqa(
self,
q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
kv_c_and_k_pe_cache: torch.Tensor,
attn_metadata: CutlassFA3MLASparseMetadata,
layer: AttentionLayer,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""FA3 sparse MLA attention with batch size gating.
For batch sizes <= MAX_BATCH_SIZE_FOR_FA3 (16), uses the fast
CUTLASS FA3 kernel. For larger batch sizes, falls back to the
FlashMLA BF16 sparse prefill kernel which handles larger batches
more efficiently.
All execution modes (decode/prefill/mixed) are handled identically:
each token is an independent batch element with seqlen=1.
Input: q = tuple(ql_nope[T, N, 512], q_pe[T, N, 64])
Output: (attn_out[T, N, 512], None)
The _v_up_proj in MLAAttention.forward_impl() handles the
subsequent .view(-1, N, kv_lora_rank) correctly for 3D output.
"""
# FA3 does not yet return LSE; DCP requires it.
assert self.dcp_world_size <= 1, (
"CutlassFA3MLASparseImpl does not support DCP (dcp_world_size > 1). "
"Use FlashMLA Sparse instead."
)
# 1) Unpack Q components
if isinstance(q, tuple):
ql_nope, q_pe = q # [T, N, 512], [T, N, 64]
else:
ql_nope = q[..., : self.kv_lora_rank] # [T, N, 512]
q_pe = q[..., self.kv_lora_rank :] # [T, N, 64]
T = ql_nope.shape[0]
# 2) Convert topk_indices -> global cache slot indices
# Reuses vLLM's existing Triton kernel (no changes needed)
from vllm.v1.attention.backends.mla.sparse_utils import (
triton_convert_req_index_to_global_index,
)
global_idx, valid_counts = triton_convert_req_index_to_global_index(
attn_metadata.req_id_per_token, # [T] int32
attn_metadata.block_table, # [R, max_blocks] int32
self.topk_indices_buffer[:T], # [T, 2048] int32
BLOCK_SIZE=attn_metadata.block_size, # 64
NUM_TOPK_TOKENS=self.topk_tokens, # 2048
return_valid_counts=True,
)
# global_idx: [T, 2048] int32 — flat cache slot IDs
# valid_counts: [T] int32 — number of valid (non -1) entries per token
# Replace -1 (invalid) page indices with 0 (a safe, valid page index)
# IN-PLACE for CUDA graph friendliness (no extra allocation).
global_idx.clamp_(min=0)
# Use valid_counts as cache_seqlens instead of metadata.cache_seqlens.
# CRITICAL FIX (Issue #1): metadata cache_seqlens = min(seq_len, topk)
# can exceed actual valid topk entries for prefill tokens.
valid_counts.clamp_(min=1) # in-place; min=1 for seqlen_k safety
cache_seqlens = valid_counts
# 3) Route to FA3 or FlashMLA based on batch size
# FA3 is faster for small batches (bs<=16) but regresses for
# larger batches. FlashMLA BF16 sparse prefill handles larger
# batches more efficiently.
use_fa3 = (T <= MAX_BATCH_SIZE_FOR_FA3) or not _flashmla_sparse_available
if use_fa3:
attn_out = self._forward_fa3(
ql_nope,
q_pe,
kv_c_and_k_pe_cache,
global_idx,
cache_seqlens,
attn_metadata,
)
else:
attn_out = self._forward_flashmla_bf16_fallback(
ql_nope,
q_pe,
kv_c_and_k_pe_cache,
global_idx,
cache_seqlens,
)
# Output: [T, N, 512] — already 3D
return attn_out, None
def _forward_fa3(
self,
ql_nope: torch.Tensor, # [T, N, 512]
q_pe: torch.Tensor, # [T, N, 64]
kv_c_and_k_pe_cache: torch.Tensor,
global_idx: torch.Tensor, # [T, 2048]
cache_seqlens: torch.Tensor, # [T]
attn_metadata: CutlassFA3MLASparseMetadata,
) -> torch.Tensor:
"""CUTLASS FA3 kernel path — fast for small batch sizes (bs<=16).
Accepts Q_rope and Q_nope (qv) separately, no head padding needed.
Uses page_size=1 paged KV format with split-KV parallelism.
"""
T = ql_nope.shape[0]
S = kv_c_and_k_pe_cache.shape[0] * kv_c_and_k_pe_cache.shape[1]
kv_flat = kv_c_and_k_pe_cache.reshape(S, self.head_size) # [S, 576]
# Split NoPE and RoPE, reshape for FA3 paged format (page_size=1)
c_kv = kv_flat[:, : self.kv_lora_rank].reshape(
S, 1, 1, self.kv_lora_rank
) # [S, 1, 1, 512]
k_rope = kv_flat[:, self.kv_lora_rank :].reshape(
S, 1, 1, self.qk_rope_head_dim
) # [S, 1, 1, 64]
from vllm.v1.attention.ops.cutlass_fa3 import flash_attn_with_kvcache
attn_out = flash_attn_with_kvcache(
q=q_pe, # [T, N, 64]
k_cache=k_rope, # [S, 1, 1, 64]
v_cache=c_kv, # [S, 1, 1, 512]
qv=ql_nope, # [T, N, 512]
page_table=global_idx, # [T, 2048]
cache_seqlens=cache_seqlens, # [T]
cu_seqlens_q=attn_metadata.cu_seqlens_q, # [T+1]
cu_seqlens_k_new=None,
max_seqlen_q=1,
softmax_scale=self.softmax_scale, # 192**-0.5
causal=True,
window_size=(-1, -1),
softcap=self.logits_soft_cap,
num_splits=self.num_splits,
)
return attn_out # [T, N, 512]
def _forward_flashmla_bf16_fallback(
self,
ql_nope: torch.Tensor, # [T, N, 512]
q_pe: torch.Tensor, # [T, N, 64]
kv_c_and_k_pe_cache: torch.Tensor,
global_idx: torch.Tensor, # [T, 2048]
cache_seqlens: torch.Tensor, # [T]
) -> torch.Tensor:
"""FlashMLA BF16 sparse prefill fallback — for larger batch sizes.
Used when T > MAX_BATCH_SIZE_FOR_FA3 (16). The FlashMLA BF16 sparse
prefill kernel handles larger batches more efficiently than FA3's
3-kernel launch pattern (scheduler + main + combine).
This path:
1. Concatenates Q components: [ql_nope | q_pe] -> [T, N, 576]
2. Pads heads to 64 (SM90 FlashMLA requirement)
3. Reshapes KV cache to [S, 1, 576] (flattened, MQA format)
4. Reshapes indices to [T, 1, topk] (MQA format)
5. Calls flash_mla_sparse_fwd with topk_length for valid bounds
6. Unpads output heads back to N
The BF16 KV cache format [kv_c_normed(512) | k_pe(64)] is identical
between FA3 and FlashMLA, so no cache format conversion is needed.
"""
T = ql_nope.shape[0]
N = self.num_heads
# 1) Concatenate Q: [ql_nope(512) | q_pe(64)] -> [T, N, 576]
q_concat = torch.cat([ql_nope, q_pe], dim=-1) # [T, N, 576]
# 2) Pad heads to _FLASHMLA_SM90_HEAD_PADDING (64 for SM90)
padded_heads = _FLASHMLA_SM90_HEAD_PADDING
if padded_heads > N:
q_padded = q_concat.new_zeros((T, padded_heads, q_concat.shape[-1]))
q_padded[:, :N, :] = q_concat
q_concat = q_padded
# 3) Reshape KV cache: (num_blocks, block_size, 576) -> (S, 1, 576)
S = kv_c_and_k_pe_cache.shape[0] * kv_c_and_k_pe_cache.shape[1]
kv = kv_c_and_k_pe_cache.reshape(S, 1, self.head_size) # [S, 1, 576]
# 4) Reshape indices for MQA: (T, 2048) -> (T, 1, 2048)
indices = global_idx.unsqueeze(1) # [T, 1, 2048]
# 5) Call FlashMLA BF16 sparse prefill kernel
# NOTE: Unlike FlashMLASparseImpl._bf16_flash_mla_kernel which does
# not pass topk_length (it relies on all indices being valid), we
# pass topk_length=valid_counts because our indices have been
# clamped (global_idx.clamp_(min=0)), so entries beyond valid_counts
# are 0 (valid but irrelevant data). topk_length prevents the kernel
# from processing these clamped entries, saving compute and ensuring
# correctness.
output = flash_mla_sparse_fwd(
q_concat, # [T, padded_heads, 576]
kv, # [S, 1, 576]
indices, # [T, 1, 2048]
self.softmax_scale, # 192**-0.5
d_v=self.kv_lora_rank, # 512
topk_length=cache_seqlens, # [T] valid entry counts
)[0] # extract output tensor from (output, max_logits, lse) tuple
# 6) Unpad heads: (T, padded_heads, 512) -> (T, N, 512)
return output[:, :N, :]
def do_kv_cache_update(
self,
kv_c_normed: torch.Tensor,
k_pe: torch.Tensor,
kv_cache: torch.Tensor,
slot_mapping: torch.Tensor,
kv_cache_dtype: str,
k_scale: torch.Tensor,
) -> None:
"""BF16 KV cache write using existing vLLM kernel.
kv_cache_dtype MUST be "auto" which maps to Fp8KVCacheDataType::kAuto
in the C++ dispatch, performing a direct BF16 copy (no quantization).
Passing "bfloat16" would crash because concat_and_cache_mla expects
the "auto" string for the non-quantized path.
"""
if kv_cache.numel() == 0:
return
from vllm import _custom_ops as ops
ops.concat_and_cache_mla(
kv_c_normed,
k_pe.squeeze(1),
kv_cache,
slot_mapping.flatten(),
kv_cache_dtype="auto",
scale=k_scale,
)
+100 -43
View File
@@ -8,6 +8,7 @@ import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.deep_gemm import (
get_paged_mqa_logits_metadata,
has_deep_gemm,
@@ -30,6 +31,40 @@ from vllm.v1.worker.cp_utils import get_total_cp_world_size
logger = init_logger(__name__)
@triton.jit
def _prepare_uniform_decode_kernel(
seq_lens_ptr,
decode_seq_lens_ptr,
block_table_ptr,
block_table_stride,
expanded_block_table_ptr,
expanded_bt_stride,
decode_lens_ptr,
max_decode_len,
BLOCK_SIZE: tl.constexpr,
):
idx = tl.program_id(0)
req_id = idx // max_decode_len
local_idx = idx % max_decode_len
# Compute number of KVs attended to by this token.
seq_len = tl.load(seq_lens_ptr + req_id)
per_token_seq_len = seq_len - max_decode_len + local_idx + 1
tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
# Copy block table row.
src = block_table_ptr + req_id * block_table_stride
dst = expanded_block_table_ptr + idx * expanded_bt_stride
for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
off = i + tl.arange(0, BLOCK_SIZE)
mask = off < expanded_bt_stride
src_block = tl.load(src + off, mask=mask)
tl.store(dst + off, src_block, mask=mask)
# All reqs now have decode_len = 1.
tl.store(decode_lens_ptr + idx, 1)
def split_indexer_prefill_chunks(
seq_lens_cpu: torch.Tensor,
query_lens_cpu: torch.Tensor,
@@ -405,52 +440,75 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).
seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, next_n) for native MTP.
"""
min_decode_len = int(decode_lens_cpu.min().item())
if not use_native and max_decode_len > 1:
assert self.decode_seq_lens_buffer.dim() == 1
# Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
# padding) and decode_lens [3, 1, 4, 0] in the below example comments.
# The context lengths are therefore
# [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
# 3 + 1 + 4 + 0 = 8
actual_expanded = int(decode_lens_cpu.sum().item())
# Fuse expanded_base and expanded_starts into a single repeat_interleave:
# seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
# where context_start[b] = seq_lens[b] - decode_lens[b].
# Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
# expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4]
# result = [8, 9, 10, 7, 9, 10, 11, 12]
expanded_offsets = torch.repeat_interleave(
seq_lens - decode_lens - query_start_loc,
decode_lens,
output_size=actual_expanded,
)
# [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
self.decode_seq_lens_buffer[:actual_expanded] = (
expanded_offsets + self.arange_buffer[:actual_expanded] + 1
)
self.decode_seq_lens_buffer[actual_expanded:] = 0
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
# Give each of the flattened entries the same block table row as the
# original request.
self.expanded_block_table_buffer[:actual_expanded] = (
torch.repeat_interleave(
block_table, decode_lens, dim=0, output_size=actual_expanded
if min_decode_len == max_decode_len:
# Uniform decode lengths.
num_decode_tokens = num_decodes * max_decode_len
_prepare_uniform_decode_kernel[(num_decode_tokens,)](
seq_lens,
self.decode_seq_lens_buffer,
block_table,
block_table.stride(0),
self.expanded_block_table_buffer,
self.expanded_block_table_buffer.stride(0),
self.decode_lens_buffer,
max_decode_len,
BLOCK_SIZE=1024,
)
)
if actual_expanded < num_decode_tokens:
self.expanded_block_table_buffer[
actual_expanded:num_decode_tokens, 0
] = 0
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
self.decode_seq_lens_buffer[num_decode_tokens:] = 0
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
return seq_lens, block_table, decode_lens, num_decode_tokens, False
else:
# Variable decode lengths.
# Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
# padding) and decode_lens [3, 1, 4, 0] in the below example comments.
# The context lengths are therefore
# [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
# All reqs now have decode_len=1
self.decode_lens_buffer[:num_decode_tokens] = 1
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
return seq_lens, block_table, decode_lens, num_decode_tokens, False
# 3 + 1 + 4 + 0 = 8
actual_expanded = int(decode_lens_cpu.sum().item())
# Fuse expanded_base and expanded_starts into a single
# repeat_interleave:
# seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
# where context_start[b] = seq_lens[b] - decode_lens[b].
# Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
# expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4]
# result = [8, 9, 10, 7, 9, 10, 11, 12]
expanded_offsets = torch.repeat_interleave(
seq_lens - decode_lens - query_start_loc,
decode_lens,
output_size=actual_expanded,
)
# [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
self.decode_seq_lens_buffer[:actual_expanded] = (
expanded_offsets + self.arange_buffer[:actual_expanded] + 1
)
self.decode_seq_lens_buffer[actual_expanded:] = 0
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
# Give each of the flattened entries the same block table row as the
# original request.
self.expanded_block_table_buffer[:actual_expanded] = (
torch.repeat_interleave(
block_table, decode_lens, dim=0, output_size=actual_expanded
)
)
if actual_expanded < num_decode_tokens:
self.expanded_block_table_buffer[
actual_expanded:num_decode_tokens, 0
] = 0
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
# All reqs now have decode_len=1
self.decode_lens_buffer[:num_decode_tokens] = 1
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
return seq_lens, block_table, decode_lens, num_decode_tokens, False
else:
# Native path: plain decode (next_n==1) or spec decode
# with 2D per-token context lengths (next_n > 1).
@@ -459,7 +517,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
# decode_len < next_n due to padding or short prefills), the simple
# reshape in sparse_attn_indexer won't work. Use pack_seq_triton
# (requires_padding) instead.
min_decode_len = int(decode_lens_cpu.min().item())
requires_padding = min_decode_len != max_decode_len
if use_native and next_n > 1:
assert self.decode_seq_lens_buffer.dim() == 2
+3
View File
@@ -73,6 +73,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
FLASHMLA_SPARSE = (
"vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend"
)
CUTLASS_FA3_MLA_SPARSE = (
"vllm.v1.attention.backends.mla.cutlass_fa3_sparse.CutlassFA3MLASparseBackend"
)
FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
+162
View File
@@ -0,0 +1,162 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Vendored CUTLASS FA3 MLA attention kernel wrapper.
This module wraps the CUTLASS FlashAttention3 Sm90 kernel from sgl-attn,
providing a Python interface compatible with vLLM's sparse MLA attention
backend. The kernel is vendored as a self-contained C++ extension
(_cutlass_fa3_C) and does NOT depend on sglang, sgl_kernel, or any sgl*
modules.
The FA3 kernel supports MLA (Multi-head Latent Attention) with:
- Separate Q_rope and QV (Q_nope) components
- Paged KV cache with page_size=1
- Variable-length sequences via cu_seqlens
- Split-KV parallelism with automatic split count
- SM90 (Hopper) CUTLASS warpgroup MMA + TMA
Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
"""
import torch
from vllm.platforms import current_platform
_cutlass_fa3_available = False
if current_platform.is_cuda():
try:
import vllm._cutlass_fa3_C # noqa: F401
_cutlass_fa3_available = True
except ImportError:
pass
def is_cutlass_fa3_available() -> bool:
"""Check if the CUTLASS FA3 extension is available.
Requires CUDA >= 12.4 and SM90 (Hopper) GPU.
"""
return _cutlass_fa3_available
def flash_attn_with_kvcache(
q: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
k: torch.Tensor | None = None,
v: torch.Tensor | None = None,
qv: torch.Tensor | None = None,
rotary_cos: torch.Tensor | None = None,
rotary_sin: torch.Tensor | None = None,
cache_seqlens: torch.Tensor | None = None,
cache_batch_idx: torch.Tensor | None = None,
cache_leftpad: torch.Tensor | None = None,
page_table: torch.Tensor | None = None,
cu_seqlens_q: torch.Tensor | None = None,
cu_seqlens_k_new: torch.Tensor | None = None,
max_seqlen_q: int | None = None,
rotary_seqlens: torch.Tensor | None = None,
q_descale: torch.Tensor | None = None,
k_descale: torch.Tensor | None = None,
v_descale: torch.Tensor | None = None,
softmax_scale: float | None = None,
causal: bool = False,
window_size: tuple[int, int] = (-1, -1),
attention_chunk: int | None = None,
softcap: float = 0.0,
rotary_interleaved: bool = True,
scheduler_metadata: torch.Tensor | None = None,
num_splits: int = 0,
pack_gqa: bool | None = None,
sm_margin: int = 0,
return_softmax_lse: bool = False,
sinks: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""CUTLASS FA3 attention with paged KV cache for MLA.
MLA mode shapes (DeepSeek-V3.2, varlen mode with cu_seqlens_q):
q: [T, N, 64] -- RoPE query (total_q=T, heads=N, dim=64)
qv: [T, N, 512] -- NoPE query (total_q=T, heads=N, dim_v=512)
k_cache: [S, 1, 1, 64] -- Paged RoPE keys (pages, pg_sz=1, kv_h=1, d)
v_cache: [S, 1, 1, 512] -- Paged NoPE latent (pages, pg_sz=1, kv_h=1, dv)
page_table: [T, topk] -- Global cache slot indices per token
The FA3 kernel internally computes:
score = Q_rope @ K_rope^T + QV(Q_nope) @ V_cache(C_KV)^T
output = softmax(score * scale) @ V_cache(C_KV)
FA3 MLA constraints (from flash_api.cpp):
headdim_qk <= 64, headdim_v >= 256, SM90 only, BF16/FP16 only
FA3 produces 3 sub-kernels:
1. prepare_varlen_num_blocks_kernel (scheduler)
2. FlashAttnFwdSm90 (main attention, TMA+WGMMA)
3. FlashAttnFwdCombine (split-KV merge, when num_splits > 1)
Args:
q: Query tensor for RoPE component.
k_cache: Paged K cache (RoPE keys).
v_cache: Paged V cache (NoPE latent).
qv: Query tensor for NoPE/value component (MLA specific).
page_table: Page table mapping tokens to cache slots.
cache_seqlens: Number of valid KV entries per batch element.
cu_seqlens_q: Cumulative query sequence lengths.
cu_seqlens_k_new: Cumulative KV sequence lengths.
max_seqlen_q: Maximum query sequence length.
softmax_scale: Softmax scale factor (default: q.shape[-1]**-0.5).
causal: Whether to apply causal masking.
window_size: (left, right) attention window sizes.
softcap: Logits soft cap value (0.0 = disabled).
num_splits: Number of split-KV splits (0 = auto).
return_softmax_lse: Whether to return log-sum-exp values.
Returns:
Attention output tensor, or tuple of (output, softmax_lse).
"""
assert _cutlass_fa3_available, (
"CUTLASS FA3 requires CUDA >= 12.4 and SM90 (Hopper) GPU. "
"The _cutlass_fa3_C extension was not compiled or could not be loaded."
)
if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5)
attention_chunk_val = 0 if attention_chunk is None else int(attention_chunk)
out, softmax_lse, *rest = torch.ops._cutlass_fa3_C.fwd.default(
q, # 0: q
k_cache, # 1: k (paged KV cache)
v_cache, # 2: v (paged KV cache)
k, # 3: k_new
v, # 4: v_new
qv, # 5: q_v (MLA NoPE query)
None, # 6: out buffer
cu_seqlens_q, # 7: cu_seqlens_q
None, # 8: cu_seqlens_k
cu_seqlens_k_new, # 9: cu_seqlens_k_new
None, # 10: seqused_q
cache_seqlens, # 11: seqused_k
max_seqlen_q, # 12: max_seqlen_q
None, # 13: max_seqlen_k
page_table, # 14: page_table
cache_batch_idx, # 15: kv_batch_idx
cache_leftpad, # 16: leftpad_k
rotary_cos, # 17: rotary_cos
rotary_sin, # 18: rotary_sin
rotary_seqlens, # 19: seqlens_rotary
q_descale, # 20: q_descale
k_descale, # 21: k_descale
v_descale, # 22: v_descale
softmax_scale, # 23: softmax_scale
causal, # 24: is_causal
window_size[0], # 25: window_size_left
window_size[1], # 26: window_size_right
attention_chunk_val, # 27: attention_chunk
softcap, # 28: softcap
rotary_interleaved, # 29: is_rotary_interleaved
scheduler_metadata, # 30: scheduler_metadata
num_splits, # 31: num_splits
pack_gqa, # 32: pack_gqa
sm_margin, # 33: sm_margin
sinks, # 34: sinks
)
return (out, softmax_lse) if return_softmax_lse else out
+6 -15
View File
@@ -51,6 +51,7 @@ from vllm.utils.network_utils import (
get_loopback_ip,
get_open_port,
)
from vllm.utils.ompmultiprocessing import OMPProcessManager
from vllm.utils.system_utils import (
_maybe_force_spawn,
decorate_logs,
@@ -169,24 +170,14 @@ class MultiprocExecutor(Executor):
[] if context.get_start_method() == "fork" else None
)
# For CPU backend only, to setup OpenMP threads affinity
cpu_omp_manager = OMPProcessManager(self.vllm_config)
for local_rank in range(self.local_world_size):
global_rank = global_start_rank + local_rank
is_driver_worker = self._is_driver_worker(global_rank)
if current_platform.is_cpu():
om = current_platform.get_omp_manager()
logger.info("Configured OMP PLACES %s", str(om.omp_places))
unready_worker_handle = om.run(
WorkerProc.make_worker_process,
vllm_config=self.vllm_config,
local_rank=local_rank,
rank=global_rank,
distributed_init_method=distributed_init_method,
input_shm_handle=scheduler_output_handle,
shared_worker_lock=shared_worker_lock,
is_driver_worker=is_driver_worker,
inherited_fds=inherited_fds,
)
else:
with cpu_omp_manager.configure_omp_envs(
rank=global_rank, local_rank=local_rank
):
unready_worker_handle = WorkerProc.make_worker_process(
vllm_config=self.vllm_config,
local_rank=local_rank,
+27 -3
View File
@@ -17,7 +17,7 @@ from vllm.logger import init_logger
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import get_dtype_size
from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
logger = init_logger(__name__)
@@ -38,11 +38,20 @@ class KVQuantMode(IntEnum):
FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path)
INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8
FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8
NVFP4 = 4 # packed fp4 data + fp8 block scales
@property
def is_per_token_head(self) -> bool:
"""True for any per-token-head quantization mode."""
return self >= 2
return self in (
KVQuantMode.INT8_PER_TOKEN_HEAD,
KVQuantMode.FP8_PER_TOKEN_HEAD,
)
@property
def is_nvfp4(self) -> bool:
"""True for NVFP4 packed quantization mode."""
return self == KVQuantMode.NVFP4
def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
@@ -51,7 +60,9 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
return KVQuantMode.INT8_PER_TOKEN_HEAD
if kv_cache_dtype == "fp8_per_token_head":
return KVQuantMode.FP8_PER_TOKEN_HEAD
if kv_cache_dtype.startswith("fp8"):
if kv_cache_dtype == "nvfp4":
return KVQuantMode.NVFP4
if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):
return KVQuantMode.FP8_PER_TENSOR
return KVQuantMode.NONE
@@ -237,6 +248,19 @@ class FullAttentionSpec(AttentionSpec):
@property
def real_page_size_bytes(self) -> int:
if self.kv_quant_mode.is_nvfp4:
# Packed layout per head: fp4 data + fp8 block scales.
# fp4 data: head_size//2 bytes (2 fp4 values per byte)
# fp8 block scale: head_size//16 bytes (1 scale per 16 elements)
last_dim = nvfp4_kv_cache_full_dim(
self.head_size
) + nvfp4_kv_cache_full_dim(self.head_size_v)
return (
self.block_size
* self.num_kv_heads
* last_dim
* get_dtype_size(self.dtype)
)
return (
self.block_size
* self.num_kv_heads
+7 -11
View File
@@ -26,17 +26,13 @@ class CPUOffloadingSpec(OffloadingSpec):
# calculate kv_bytes_per_offloaded_block
assert kv_cache_config is not None
page_sizes = {
kv_cache_group.kv_cache_spec.page_size_bytes
for kv_cache_group in kv_cache_config.kv_cache_groups
}
assert len(page_sizes) == 1
page_size_bytes = page_sizes.pop()
kv_bytes_per_block = (
page_size_bytes
* len(kv_cache_config.kv_cache_tensors)
* vllm_config.parallel_config.world_size
)
if kv_cache_config.num_blocks > 0:
total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
kv_bytes_per_block = (
total_gpu_kv_bytes // kv_cache_config.num_blocks
) * vllm_config.parallel_config.world_size
else:
kv_bytes_per_block = 0
kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
self.num_blocks = (
+1 -15
View File
@@ -116,21 +116,7 @@ class CPUModelRunner(GPUModelRunner):
logger.info("Warming up model for the compilation...")
# Only generate graph for the generic shape
with _set_global_compilation_settings(self.vllm_config):
self._dummy_run(
min(
max(16, self.max_num_reqs),
self.scheduler_config.max_num_batched_tokens,
)
)
# Warm up drafter for speculative decoding
if self.speculative_config and (self.speculative_config.uses_draft_model()):
from vllm.v1.spec_decode.draft_model import DraftModelProposer
if isinstance(self.drafter, (DraftModelProposer)):
logger.info("Warming up drafter model...")
self.drafter.dummy_run(max(16, self.max_num_reqs))
self.profile_run()
logger.info("Warming up done.")
def initialize_kv_cache(
+106 -2
View File
@@ -1,15 +1,23 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
import os
import sys
from typing import Any
import psutil
import torch
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.platforms import CpuArchEnum, current_platform
from vllm.profiler.wrapper import TorchProfilerWrapper
from vllm.utils.cpu_resource_utils import (
get_allowed_cpu_list,
get_memory_node_info,
get_visible_memory_node,
)
from vllm.utils.mem_utils import format_gib
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.worker.cpu_model_runner import CPUModelRunner
from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
@@ -27,6 +35,46 @@ class CPUWorker(Worker):
distributed_init_method: str,
is_driver_worker: bool = False,
):
# TODO: use numactl for process setup
# TODO: optimize for `interleaved` policy
# Bind memory node
allowed_memory_nodes = get_visible_memory_node()
allowed_cpu_list = get_allowed_cpu_list()
cpu_core = allowed_cpu_list[0]
# TODO: some CI hosts are not correctly set, change to assertion
# after fix
if cpu_core.numa_node not in allowed_memory_nodes:
logger.warning(
"Node %s is not in available memory nodes %s.",
cpu_core.numa_node,
allowed_memory_nodes,
)
torch.ops._C.init_cpu_memory_env([cpu_core.numa_node])
memory_status = get_memory_node_info(cpu_core.numa_node)
memory_fraction = vllm_config.cache_config.gpu_memory_utilization
self.requested_cpu_memory = math.ceil(
memory_status.total_memory * memory_fraction
)
available_memory = memory_status.available_memory
if (
vllm_config.cache_config.kv_cache_memory_bytes is None
and self.requested_cpu_memory > available_memory
):
raise ValueError(
f"Available memory on node {cpu_core.numa_node} "
f"({format_gib(available_memory)}/"
f"{format_gib(memory_status.total_memory)} GiB) on startup "
f"is less than desired CPU memory utilization "
f"({vllm_config.cache_config.gpu_memory_utilization}, "
f"{format_gib(self.requested_cpu_memory)} GiB). "
"Decrease --gpu-memory-utilization"
f" or reduce CPU memory used by other processes."
)
super().__init__(
vllm_config,
local_rank,
@@ -103,13 +151,69 @@ class CPUWorker(Worker):
pass
def determine_available_memory(self) -> int:
return self.cache_config.cpu_kvcache_space_bytes or 0
self.model_runner.warming_up_model()
allowed_cpu_list = get_allowed_cpu_list()
cpu_core = allowed_cpu_list[0]
memory_status = get_memory_node_info(cpu_core.numa_node)
available_memory = memory_status.available_memory
explicit_kv_cache_size = self.cache_config.kv_cache_memory_bytes
kv_cache_size = None
msg = None
if explicit_kv_cache_size is not None:
if explicit_kv_cache_size > available_memory:
raise ValueError(
f"Available memory on node {cpu_core.numa_node} "
f"({format_gib(available_memory)}/"
f"{format_gib(memory_status.total_memory)} GiB) on kv cache"
f" allocation is less than requested memory for kv "
f"({format_gib(explicit_kv_cache_size)} GiB). "
"Decrease --kv-cache-memory-bytes, VLLM_CPU_KVCACHE_SPACE, "
"or reduce CPU memory used by other processes."
)
kv_cache_size = explicit_kv_cache_size
msg = (
f"Explicitly set ({format_gib(kv_cache_size)}/"
f"{format_gib(memory_status.total_memory)}) GiB for KV cache "
f"on node {cpu_core.numa_node}."
)
else:
consumed_memory = psutil.Process(os.getpid()).memory_info().rss
requested_memory_for_kv = int(self.requested_cpu_memory - consumed_memory)
if (
requested_memory_for_kv <= 0
or requested_memory_for_kv > available_memory
):
raise ValueError(
f"Available memory on node {cpu_core.numa_node} "
f"({format_gib(available_memory)}/"
f"{format_gib(memory_status.total_memory)} GiB) on kv cache"
f" allocation is less than requested memory for kv "
f"({format_gib(requested_memory_for_kv)}/"
f"{format_gib(self.requested_cpu_memory)} GiB). "
"Reduce CPU memory used by other processes."
)
kv_cache_size = requested_memory_for_kv
msg = (
f"Auto set ({format_gib(kv_cache_size)}/"
f"{format_gib(memory_status.total_memory)}) GiB for KV cache "
f"on node {cpu_core.numa_node}, with "
f"{format_gib(self.requested_cpu_memory)} GiB requested memory"
f" for the worker. {format_gib(consumed_memory)} GiB"
f" memory was consumed by non-kv usages."
)
logger.info(msg)
return kv_cache_size
def compile_or_warm_up_model(self) -> CompilationTimes:
# Reset the seed to ensure that the random state is not affected by
# the model initialization and profiling.
set_random_seed(self.model_config.seed)
self.model_runner.warming_up_model()
# Note: the model has been compiled in determine_available_memory()
return CompilationTimes(
language_model=self.compilation_config.compilation_time,
encoder=self.compilation_config.encoder_compilation_time,
+3 -1
View File
@@ -363,7 +363,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend(
self.kv_cache_config, self.vllm_config, self.device
)
initialize_mamba_ssu_backend(self.vllm_config.mamba_config)
initialize_mamba_ssu_backend(
self.vllm_config.mamba_config, self.kv_cache_config
)
cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
attn_cg_support.min_cg_support,
attn_cg_support.min_cg_attn_backend,
+3 -1
View File
@@ -6738,7 +6738,9 @@ class GPUModelRunner(
self.may_add_encoder_only_layers_to_kv_cache_config()
self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
initialize_mamba_ssu_backend(self.vllm_config.mamba_config)
initialize_mamba_ssu_backend(
self.vllm_config.mamba_config, self.kv_cache_config
)
# The kernel block size for all KV cache groups. For example, if
# kv_cache_manager uses block_size 256 for a given group, but the attention
# backends for that group only supports block_size 64, we will return