forked from Karylab-cklius/vllm
Compare commits
27
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a60418e6fb | ||
|
|
1174723eba | ||
|
|
6b2b7bd0eb | ||
|
|
70770268c3 | ||
|
|
7a51b3e415 | ||
|
|
d02421a7db | ||
|
|
b1dc87a098 | ||
|
|
79a5b63253 | ||
|
|
c0c98b8b9a | ||
|
|
8d2cff8140 | ||
|
|
4f436782af | ||
|
|
bf45e6d0a5 | ||
|
|
978a4462bb | ||
|
|
1948d0c467 | ||
|
|
4c47710bf7 | ||
|
|
bf9a5ddb24 | ||
|
|
79e799ebbd | ||
|
|
c4e601c73c | ||
|
|
29057d3bee | ||
|
|
219bb5b8c0 | ||
|
|
ad2b1277f9 | ||
|
|
b897f00c9c | ||
|
|
adf9bb3c57 | ||
|
|
b16fda62b7 | ||
|
|
de111f3246 | ||
|
|
afabb5f45a | ||
|
|
3abb7560c0 |
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -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)`);
|
||||
}
|
||||
|
||||
@@ -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
@@ -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}")
|
||||
|
||||
@@ -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
@@ -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);
|
||||
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
}
|
||||
@@ -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_);
|
||||
@@ -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)...);
|
||||
};
|
||||
}
|
||||
@@ -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"]
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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, vLLM’s 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.
|
||||
|
||||
|
||||
+2
-2
@@ -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)
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
@@ -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 = []
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"])
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user