forked from Karylab-cklius/vllm
Compare commits
45
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
aeee7ef939 | ||
|
|
cda19ecf4d | ||
|
|
80b18230e0 | ||
|
|
d0697cc7b6 | ||
|
|
b0755523dc | ||
|
|
993859ceb0 | ||
|
|
48a65ccb02 | ||
|
|
55842a8d69 | ||
|
|
1f45e83756 | ||
|
|
a8bffaa133 | ||
|
|
5cdddddd4a | ||
|
|
6ef1efd51f | ||
|
|
58da4ee047 | ||
|
|
1ae11e2bfc | ||
|
|
251c18d1f8 | ||
|
|
512765d52d | ||
|
|
640cc9dd7d | ||
|
|
ceade1952c | ||
|
|
747256bb5d | ||
|
|
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
|
||||
|
||||
@@ -2613,6 +2613,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=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
|
||||
|
||||
|
||||
@@ -3601,7 +3602,6 @@ steps:
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- 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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -141,6 +141,7 @@ steps:
|
||||
- pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
|
||||
- pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
|
||||
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer.py
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
|
||||
|
||||
+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)`);
|
||||
}
|
||||
|
||||
+17
-1
@@ -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}")
|
||||
@@ -944,11 +952,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
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}")
|
||||
|
||||
+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}")
|
||||
|
||||
+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) {
|
||||
|
||||
@@ -134,4 +134,27 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor& input,
|
||||
torch::stable::Tensor& input_global_scale);
|
||||
|
||||
void mxfp4_experts_quant(
|
||||
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& input_offset_by_experts,
|
||||
torch::stable::Tensor const& output_scale_offset_by_experts,
|
||||
int64_t n_experts);
|
||||
|
||||
void silu_and_mul_mxfp4_experts_quant(
|
||||
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& input_offset_by_experts,
|
||||
torch::stable::Tensor const& output_scale_offset_by_experts,
|
||||
int64_t n_experts);
|
||||
|
||||
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
|
||||
const torch::stable::Tensor& a,
|
||||
const torch::stable::Tensor& b,
|
||||
const torch::stable::Tensor& a_blockscale,
|
||||
const torch::stable::Tensor& b_blockscales,
|
||||
const torch::stable::Tensor& problem_sizes,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets);
|
||||
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,468 @@
|
||||
/*
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
*
|
||||
* MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
|
||||
* Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
|
||||
*/
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include <cutlass/arch/arch.h>
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/epilogue/collective/default_epilogue.hpp"
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/gemm/group_array_problem_shape.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
#include <cassert>
|
||||
|
||||
using namespace cute;
|
||||
|
||||
// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
|
||||
template <typename ElementAB, typename ElementC, typename ElementSF,
|
||||
typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
|
||||
__global__ void __mxfp4_get_group_gemm_starts(
|
||||
ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
|
||||
ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
|
||||
LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
|
||||
ElementAB* a_base_as_int, ElementAB* b_base_as_int,
|
||||
ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
|
||||
ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
|
||||
const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
|
||||
int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
|
||||
const int64_t a_stride_val, const int64_t b_stride_val,
|
||||
const int64_t c_stride_val, const int K, const int N) {
|
||||
int64_t expert_id = threadIdx.x;
|
||||
if (expert_id >= gridDim.x * blockDim.x) {
|
||||
return;
|
||||
}
|
||||
int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
|
||||
int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
|
||||
int64_t group_size = 32;
|
||||
int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
|
||||
int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
|
||||
int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
|
||||
assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
|
||||
"unexpected problem sizes");
|
||||
|
||||
int64_t half_k = static_cast<int64_t>(k / 2);
|
||||
int64_t group_k = static_cast<int64_t>(k / group_size);
|
||||
// Shape of A as uint8/byte = [M, K // 2]
|
||||
a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
|
||||
// Shape of B as uint8/byte = [E, N, K // 2]
|
||||
b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
|
||||
// Shape of C = [M, N]
|
||||
out_offsets[expert_id] = out_base_as_int + expert_offset * n;
|
||||
// Shape of a_scale = [sum(sf_sizes), K // group_size]
|
||||
a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
|
||||
|
||||
assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
|
||||
0 &&
|
||||
"TMA requires 128-byte alignment");
|
||||
|
||||
// Shape of B scale = [E, N, K // group_size]
|
||||
b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
|
||||
assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
|
||||
0 &&
|
||||
"TMA requires 128-byte alignment");
|
||||
|
||||
// Initialize strides
|
||||
a_strides[expert_id] = a_stride_val;
|
||||
b_strides[expert_id] = b_stride_val;
|
||||
c_strides[expert_id] = c_stride_val;
|
||||
|
||||
LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
|
||||
LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
|
||||
|
||||
*layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
|
||||
static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
|
||||
*layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
|
||||
static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
|
||||
}
|
||||
|
||||
#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE, \
|
||||
TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
|
||||
LayoutSFB, ScaleConfig) \
|
||||
else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
|
||||
__mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
|
||||
LayoutSFB, ScaleConfig> \
|
||||
<<<1, num_experts, 0, stream>>>( \
|
||||
static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \
|
||||
static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()), \
|
||||
static_cast<C_TYPE**>(out_starts.data_ptr()), \
|
||||
static_cast<SF_TYPE**>(a_scales_starts.data_ptr()), \
|
||||
static_cast<SF_TYPE**>(b_scales_starts.data_ptr()), \
|
||||
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \
|
||||
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \
|
||||
static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()), \
|
||||
static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()), \
|
||||
static_cast<C_TYPE*>(out_tensors.data_ptr()), \
|
||||
static_cast<SF_TYPE*>(a_scales.data_ptr()), \
|
||||
static_cast<SF_TYPE*>(b_scales.data_ptr()), \
|
||||
static_cast<int32_t*>(expert_offsets.data_ptr()), \
|
||||
static_cast<int32_t*>(sf_offsets.data_ptr()), \
|
||||
static_cast<int32_t*>(problem_sizes.data_ptr()), \
|
||||
static_cast<int64_t*>(a_strides.data_ptr()), \
|
||||
static_cast<int64_t*>(b_strides.data_ptr()), \
|
||||
static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val, \
|
||||
b_stride_val, c_stride_val, K, N); \
|
||||
}
|
||||
|
||||
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
|
||||
void mxfp4_run_get_group_gemm_starts(
|
||||
const torch::stable::Tensor& a_starts,
|
||||
const torch::stable::Tensor& b_starts,
|
||||
const torch::stable::Tensor& out_starts,
|
||||
const torch::stable::Tensor& a_scales_starts,
|
||||
const torch::stable::Tensor& b_scales_starts,
|
||||
const torch::stable::Tensor& layout_sfa,
|
||||
const torch::stable::Tensor& layout_sfb,
|
||||
const torch::stable::Tensor& a_strides,
|
||||
const torch::stable::Tensor& b_strides,
|
||||
const torch::stable::Tensor& c_strides, int64_t a_stride_val,
|
||||
int64_t b_stride_val, int64_t c_stride_val,
|
||||
torch::stable::Tensor const& a_tensors,
|
||||
torch::stable::Tensor const& b_tensors,
|
||||
torch::stable::Tensor const& out_tensors,
|
||||
torch::stable::Tensor const& a_scales,
|
||||
torch::stable::Tensor const& b_scales,
|
||||
torch::stable::Tensor const& expert_offsets,
|
||||
torch::stable::Tensor const& sf_offsets,
|
||||
torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
|
||||
int num_experts = (int)expert_offsets.size(0);
|
||||
auto stream = get_current_cuda_stream(a_tensors.get_device_index());
|
||||
|
||||
STD_TORCH_CHECK(out_tensors.size(1) == N,
|
||||
"Output tensor shape doesn't match expected shape");
|
||||
STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
|
||||
"b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
|
||||
" dimension must match");
|
||||
if (false) {
|
||||
}
|
||||
// MXFP4 uses E8M0 (float_ue8m0_t) scale factors
|
||||
__CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
|
||||
torch::headeronly::ScalarType::BFloat16,
|
||||
cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
|
||||
ScaleConfig)
|
||||
__CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
|
||||
torch::headeronly::ScalarType::Half, half,
|
||||
LayoutSFA, LayoutSFB, ScaleConfig)
|
||||
else {
|
||||
STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
|
||||
}
|
||||
}
|
||||
|
||||
template <typename OutType>
|
||||
void run_mxfp4_blockwise_scaled_group_mm_sm100(
|
||||
torch::stable::Tensor& output, const torch::stable::Tensor& a,
|
||||
const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
|
||||
const torch::stable::Tensor& b_blockscales,
|
||||
const torch::stable::Tensor& problem_sizes,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
|
||||
using ProblemShape =
|
||||
cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
|
||||
using ElementType = cutlass::float_e2m1_t;
|
||||
using ElementSFType = cutlass::float_ue8m0_t;
|
||||
using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
|
||||
using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
|
||||
|
||||
using ElementC = OutType;
|
||||
using ElementD = ElementC;
|
||||
using ElementAccumulator = float;
|
||||
// Layout definitions
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
using LayoutD = LayoutC;
|
||||
|
||||
static constexpr int AlignmentA = 32;
|
||||
static constexpr int AlignmentB = 32;
|
||||
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
|
||||
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
|
||||
|
||||
// Architecture definitions
|
||||
using ArchTag = cutlass::arch::Sm100;
|
||||
using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
|
||||
using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
|
||||
using StageCountType = cutlass::gemm::collective::StageCountAuto;
|
||||
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
struct MMA1SMConfig {
|
||||
using MmaTileShape = Shape<_128, _128, _128>;
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
|
||||
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
|
||||
};
|
||||
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
|
||||
ClusterShape, Shape<_128, _64>, ElementAccumulator,
|
||||
ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
|
||||
LayoutC*, AlignmentD,
|
||||
typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
|
||||
ElementB, LayoutB*, AlignmentB, ElementAccumulator,
|
||||
typename MMA1SMConfig::MmaTileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
|
||||
sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
|
||||
|
||||
using GemmKernel =
|
||||
cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
|
||||
CollectiveEpilogue>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using StrideA = typename Gemm::GemmKernel::InternalStrideA;
|
||||
using StrideB = typename Gemm::GemmKernel::InternalStrideB;
|
||||
using StrideC = typename Gemm::GemmKernel::InternalStrideC;
|
||||
using StrideD = typename Gemm::GemmKernel::InternalStrideD;
|
||||
|
||||
using LayoutSFA =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
|
||||
using LayoutSFB =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
|
||||
using ScaleConfig =
|
||||
typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
|
||||
|
||||
using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
|
||||
int num_experts = static_cast<int>(expert_offsets.size(0));
|
||||
|
||||
torch::stable::Tensor a_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor b_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor out_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor a_scales_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor b_scales_ptrs =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor layout_sfa = torch::stable::empty(
|
||||
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
|
||||
a.device());
|
||||
torch::stable::Tensor layout_sfb = torch::stable::empty(
|
||||
{num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
|
||||
a.device());
|
||||
torch::stable::Tensor a_strides1 =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor b_strides1 =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
torch::stable::Tensor c_strides1 =
|
||||
torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
|
||||
std::nullopt, a.device());
|
||||
|
||||
mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
|
||||
a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
|
||||
layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
|
||||
b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
|
||||
b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
|
||||
|
||||
// Create an instance of the GEMM
|
||||
Gemm gemm_op;
|
||||
|
||||
UnderlyingProblemShape* problem_sizes_as_shapes =
|
||||
static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
|
||||
|
||||
// Set the Scheduler info
|
||||
cutlass::KernelHardwareInfo hw_info;
|
||||
using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
|
||||
PersistentTileSchedulerSm100GroupParams<
|
||||
typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
|
||||
typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
|
||||
scheduler.raster_order = RasterOrderOptions::AlongM;
|
||||
hw_info.device_id = a.get_device_index();
|
||||
static std::unordered_map<int, int> cached_sm_counts;
|
||||
if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
|
||||
cached_sm_counts[hw_info.device_id] =
|
||||
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
|
||||
hw_info.device_id);
|
||||
}
|
||||
hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
|
||||
|
||||
// Mainloop Arguments
|
||||
typename GemmKernel::MainloopArguments mainloop_args{
|
||||
static_cast<const ElementType**>(a_ptrs.data_ptr()),
|
||||
static_cast<StrideA*>(a_strides1.data_ptr()),
|
||||
static_cast<const ElementType**>(b_ptrs.data_ptr()),
|
||||
static_cast<StrideB*>(b_strides1.data_ptr()),
|
||||
static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
|
||||
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
|
||||
static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
|
||||
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
|
||||
|
||||
// Epilogue Arguments
|
||||
typename GemmKernel::EpilogueArguments epilogue_args{
|
||||
{}, // epilogue.thread
|
||||
nullptr,
|
||||
static_cast<StrideC*>(c_strides1.data_ptr()),
|
||||
static_cast<ElementD**>(out_ptrs.data_ptr()),
|
||||
static_cast<StrideC*>(c_strides1.data_ptr())};
|
||||
auto& fusion_args = epilogue_args.thread;
|
||||
// Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
|
||||
fusion_args.alpha_ptr = nullptr;
|
||||
fusion_args.beta_ptr = nullptr;
|
||||
fusion_args.alpha = 1.0f;
|
||||
fusion_args.alpha_ptr_array = nullptr;
|
||||
fusion_args.dAlpha = {_0{}, _0{}, 0};
|
||||
fusion_args.beta = 0.0f;
|
||||
fusion_args.beta_ptr_array = nullptr;
|
||||
fusion_args.dBeta = {_0{}, _0{}, 0};
|
||||
|
||||
// Gemm Arguments
|
||||
typename GemmKernel::Arguments args{
|
||||
cutlass::gemm::GemmUniversalMode::kGrouped,
|
||||
{num_experts, problem_sizes_as_shapes, nullptr},
|
||||
mainloop_args,
|
||||
epilogue_args,
|
||||
hw_info,
|
||||
scheduler};
|
||||
|
||||
size_t workspace_size = Gemm::get_workspace_size(args);
|
||||
auto workspace =
|
||||
torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
|
||||
std::nullopt, a.device());
|
||||
const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
|
||||
|
||||
auto can_implement_status = gemm_op.can_implement(args);
|
||||
STD_TORCH_CHECK(
|
||||
can_implement_status == cutlass::Status::kSuccess,
|
||||
"Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
|
||||
|
||||
// Run the GEMM
|
||||
auto status = gemm_op.initialize(args, workspace.data_ptr());
|
||||
STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
|
||||
"Failed to initialize MXFP4 GEMM: status=", (int)status,
|
||||
" workspace_size=", workspace_size,
|
||||
" num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
|
||||
|
||||
status = gemm_op.run(args, workspace.data_ptr(), stream);
|
||||
STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
|
||||
"Failed to run MXFP4 GEMM");
|
||||
}
|
||||
|
||||
template <typename OutType>
|
||||
void run_mxfp4_blockwise_scaled_group_mm(
|
||||
torch::stable::Tensor& output, const torch::stable::Tensor& a,
|
||||
const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
|
||||
const torch::stable::Tensor& b_blockscales,
|
||||
const torch::stable::Tensor& problem_sizes,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
|
||||
if (version_num >= 100 && version_num < 120) {
|
||||
run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
|
||||
output, a, b, a_blockscale, b_blockscales, problem_sizes,
|
||||
expert_offsets, sf_offsets, M, N, K);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
STD_TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false,
|
||||
"No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
|
||||
version_num, ". Required capability: 100");
|
||||
}
|
||||
|
||||
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
|
||||
constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
|
||||
// E8M0 scale factors stored as uint8
|
||||
constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
|
||||
#endif
|
||||
|
||||
#define CHECK_TYPE(x, st, m) \
|
||||
STD_TORCH_CHECK(x.scalar_type() == st, \
|
||||
": Inconsistency of torch::stable::Tensor type:", m)
|
||||
#define CHECK_TH_CUDA(x, m) \
|
||||
STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
|
||||
#define CHECK_CONTIGUOUS(x, m) \
|
||||
STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
|
||||
#define CHECK_INPUT(x, st, m) \
|
||||
CHECK_TH_CUDA(x, m); \
|
||||
CHECK_CONTIGUOUS(x, m); \
|
||||
CHECK_TYPE(x, st, m)
|
||||
|
||||
void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
|
||||
const torch::stable::Tensor& a,
|
||||
const torch::stable::Tensor& b,
|
||||
const torch::stable::Tensor& a_blockscale,
|
||||
const torch::stable::Tensor& b_blockscales,
|
||||
const torch::stable::Tensor& problem_sizes,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets) {
|
||||
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
|
||||
// Input validation
|
||||
CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
|
||||
CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
|
||||
// MXFP4 uses E8M0 scale factors (stored as uint8)
|
||||
CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
|
||||
CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
|
||||
|
||||
STD_TORCH_CHECK(
|
||||
a_blockscale.dim() == 2,
|
||||
"expected a_blockscale to be of shape [num_experts, rounded_m,"
|
||||
" k // group_size], observed rank: ",
|
||||
a_blockscale.dim())
|
||||
STD_TORCH_CHECK(b_blockscales.dim() == 3,
|
||||
"expected b_blockscale to be of shape: "
|
||||
" [num_experts, n, k // group_size], observed rank: ",
|
||||
b_blockscales.dim())
|
||||
STD_TORCH_CHECK(problem_sizes.dim() == 2,
|
||||
"problem_sizes must be a 2D tensor");
|
||||
STD_TORCH_CHECK(problem_sizes.size(1) == 3,
|
||||
"problem_sizes must have the shape (num_experts, 3)");
|
||||
STD_TORCH_CHECK(
|
||||
problem_sizes.size(0) == expert_offsets.size(0),
|
||||
"Number of experts in problem_sizes must match expert_offsets");
|
||||
STD_TORCH_CHECK(
|
||||
problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"problem_sizes must be int32.");
|
||||
|
||||
int M = static_cast<int>(a.size(0));
|
||||
int N = static_cast<int>(b.size(1));
|
||||
int E = static_cast<int>(b.size(0));
|
||||
int K = static_cast<int>(2 * b.size(2));
|
||||
|
||||
if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
|
||||
output, a, b, a_blockscale, b_blockscales, problem_sizes,
|
||||
expert_offsets, sf_offsets, M, N, K);
|
||||
} else {
|
||||
run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
|
||||
output, a, b, a_blockscale, b_blockscales, problem_sizes,
|
||||
expert_offsets, sf_offsets, M, N, K);
|
||||
}
|
||||
#else
|
||||
STD_TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false,
|
||||
"No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
|
||||
"SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
|
||||
#endif
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
||||
m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
|
||||
}
|
||||
@@ -0,0 +1,422 @@
|
||||
/*
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
*
|
||||
* MXFP4 activation quantization kernel for MoE experts.
|
||||
* Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
|
||||
* over 32-element groups.
|
||||
*
|
||||
* Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
|
||||
* - Block size 32 (2 threads per SF in PACK16 mode)
|
||||
* - E8M0 (power-of-two) scale factors
|
||||
* - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
|
||||
*/
|
||||
|
||||
// MXFP4 requires PACK16 mode (16 elements per thread) so that
|
||||
// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
|
||||
// Must be defined before any header that (transitively) includes
|
||||
// nvfp4_utils.cuh.
|
||||
#define NVFP4_ENABLE_ELTS16 1
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime_api.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp8.h>
|
||||
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
#include "libtorch_stable/dispatch_utils.h"
|
||||
#include "cuda_vec_utils.cuh"
|
||||
#include "cuda_utils.h"
|
||||
|
||||
#include "nvfp4_utils.cuh"
|
||||
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
|
||||
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
|
||||
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// MXFP4 block size constants
|
||||
static constexpr int MXFP4_SF_VEC_SIZE = 32;
|
||||
|
||||
// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
|
||||
// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
|
||||
static constexpr int MXFP4_NUM_THREADS_PER_SF =
|
||||
MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
// MXFP4 quantization kernel for experts.
|
||||
// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
|
||||
// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
|
||||
// SiLU(gate)*up before quantization.
|
||||
template <class Type, bool FUSE_SILU_MUL = false,
|
||||
bool SMALL_NUM_EXPERTS = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
fp4_packed_t* out, uint32_t* SFout,
|
||||
uint32_t* input_offset_by_experts,
|
||||
uint32_t* output_scale_offset_by_experts,
|
||||
int n_experts, bool low_latency) {
|
||||
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
"Vec size is not matched.");
|
||||
|
||||
// MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
|
||||
int32_t const numKTiles = (numCols + 127) / 128;
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
globalIdx += gridDim.x * blockDim.x) {
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
int colIdx = globalIdx % colsPerRow;
|
||||
|
||||
int rowIdx_in_expert = 0;
|
||||
int expert_idx = 0;
|
||||
|
||||
if constexpr (SMALL_NUM_EXPERTS) {
|
||||
for (int i = 0; i < n_experts; i++) {
|
||||
uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
|
||||
uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
|
||||
if (rowIdx >= current_offset && rowIdx < next_offset) {
|
||||
rowIdx_in_expert = rowIdx - current_offset;
|
||||
expert_idx = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
uint32_t local_offsets[17];
|
||||
for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
|
||||
*reinterpret_cast<int4*>(local_offsets) =
|
||||
__ldca(reinterpret_cast<const int4*>(
|
||||
&input_offset_by_experts[chunk_start]));
|
||||
*reinterpret_cast<int4*>(local_offsets + 4) =
|
||||
__ldca(reinterpret_cast<const int4*>(
|
||||
&input_offset_by_experts[chunk_start + 4]));
|
||||
*reinterpret_cast<int4*>(local_offsets + 8) =
|
||||
__ldca(reinterpret_cast<const int4*>(
|
||||
&input_offset_by_experts[chunk_start + 8]));
|
||||
*reinterpret_cast<int4*>(local_offsets + 12) =
|
||||
__ldca(reinterpret_cast<const int4*>(
|
||||
&input_offset_by_experts[chunk_start + 12]));
|
||||
local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; i++) {
|
||||
if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
|
||||
rowIdx_in_expert = rowIdx - local_offsets[i];
|
||||
expert_idx = chunk_start + i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Load input and optionally apply fused SiLU+Mul
|
||||
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
PackedVec quant_input;
|
||||
if constexpr (FUSE_SILU_MUL) {
|
||||
PackedVec in_vec_up =
|
||||
reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
|
||||
quant_input = compute_silu_mul(in_vec, in_vec_up);
|
||||
} else {
|
||||
quant_input = in_vec;
|
||||
}
|
||||
|
||||
// In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
|
||||
int64_t outOffset = rowIdx * colsPerRow + colIdx;
|
||||
auto& out_pos = out[outOffset];
|
||||
|
||||
uint32_t* SFout_in_expert =
|
||||
SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
|
||||
|
||||
// Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
|
||||
auto sf_out =
|
||||
cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
// Block E8M0 scales only; no extra tensor-level scale in this path
|
||||
constexpr float SFScaleVal = 1.0f;
|
||||
// UE8M0_SF=true for MXFP4 E8M0 scale factors
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
|
||||
quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
// Large M_topk variant using shared memory for expert offsets
|
||||
template <class Type, bool FUSE_SILU_MUL = false,
|
||||
bool SMALL_NUM_EXPERTS = false>
|
||||
__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
fp4_packed_t* out, uint32_t* SFout,
|
||||
uint32_t* input_offset_by_experts,
|
||||
uint32_t* output_scale_offset_by_experts,
|
||||
int n_experts) {
|
||||
using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
"Vec size is not matched.");
|
||||
|
||||
// MXFP4: numKTiles = ceil(numCols / 128)
|
||||
int32_t const numKTiles = (numCols + 127) / 128;
|
||||
|
||||
extern __shared__ uint32_t shared_input_offsets[];
|
||||
|
||||
if constexpr (SMALL_NUM_EXPERTS) {
|
||||
for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
|
||||
shared_input_offsets[i] = input_offset_by_experts[i];
|
||||
}
|
||||
} else {
|
||||
for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
|
||||
*reinterpret_cast<int4*>(&shared_input_offsets[i]) =
|
||||
*reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
|
||||
|
||||
for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
|
||||
globalIdx += gridDim.x * blockDim.x) {
|
||||
int rowIdx = globalIdx / colsPerRow;
|
||||
int colIdx = globalIdx % colsPerRow;
|
||||
|
||||
int rowIdx_in_expert = 0;
|
||||
int expert_idx = 0;
|
||||
|
||||
// Binary search through experts using shared memory
|
||||
int left = 0, right = n_experts - 1;
|
||||
while (left <= right) {
|
||||
int mid = (left + right) / 2;
|
||||
uint32_t mid_offset = shared_input_offsets[mid];
|
||||
uint32_t next_offset = shared_input_offsets[mid + 1];
|
||||
|
||||
if (rowIdx >= mid_offset && rowIdx < next_offset) {
|
||||
rowIdx_in_expert = rowIdx - mid_offset;
|
||||
expert_idx = mid;
|
||||
break;
|
||||
} else if (rowIdx < mid_offset) {
|
||||
right = mid - 1;
|
||||
} else {
|
||||
left = mid + 1;
|
||||
}
|
||||
}
|
||||
|
||||
int64_t inOffset = rowIdx * inColsPerRow + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
PackedVec quant_input;
|
||||
if constexpr (FUSE_SILU_MUL) {
|
||||
PackedVec in_vec_up =
|
||||
reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
|
||||
quant_input = compute_silu_mul(in_vec, in_vec_up);
|
||||
} else {
|
||||
quant_input = in_vec;
|
||||
}
|
||||
|
||||
int64_t outOffset = rowIdx * colsPerRow + colIdx;
|
||||
auto& out_pos = out[outOffset];
|
||||
|
||||
// MXFP4 has no global scale - only block-level E8M0 scale factors
|
||||
constexpr float SFScaleVal = 1.0f;
|
||||
|
||||
uint32_t* SFout_in_expert =
|
||||
SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
|
||||
|
||||
auto sf_out =
|
||||
cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
|
||||
quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, bool FUSE_SILU_MUL = false>
|
||||
void mxfp4_quant_impl(void* output, void* output_scale, void* input,
|
||||
void* input_offset_by_experts,
|
||||
void* output_scale_offset_by_experts, int m_topk, int k,
|
||||
int n_experts, cudaStream_t stream) {
|
||||
int multiProcessorCount =
|
||||
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
|
||||
|
||||
int const workSizePerRow = k / ELTS_PER_THREAD;
|
||||
int const totalWorkSize = m_topk * workSizePerRow;
|
||||
dim3 block(std::min(workSizePerRow, 512));
|
||||
int const numBlocksPerSM =
|
||||
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
|
||||
dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
|
||||
multiProcessorCount * numBlocksPerSM));
|
||||
while (grid.x <= multiProcessorCount && block.x > 64) {
|
||||
grid.x *= 2;
|
||||
block.x = (block.x + 1) / 2;
|
||||
}
|
||||
|
||||
int const blockRepeat =
|
||||
(totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
|
||||
if (blockRepeat > 1) {
|
||||
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
|
||||
if (n_experts >= 4) {
|
||||
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<fp4_packed_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts);
|
||||
} else {
|
||||
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<fp4_packed_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts);
|
||||
}
|
||||
} else {
|
||||
if (n_experts >= 16) {
|
||||
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<fp4_packed_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts, /* bool low_latency */ true);
|
||||
} else {
|
||||
mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
reinterpret_cast<fp4_packed_t*>(output),
|
||||
reinterpret_cast<uint32_t*>(output_scale),
|
||||
reinterpret_cast<uint32_t*>(input_offset_by_experts),
|
||||
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
|
||||
n_experts, /* bool low_latency */ true);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
/*Quantization entry for mxfp4 experts quantization*/
|
||||
#define CHECK_TH_CUDA(x, m) \
|
||||
STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x, m) \
|
||||
STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
|
||||
#define CHECK_INPUT(x, m) \
|
||||
CHECK_TH_CUDA(x, m); \
|
||||
CHECK_CONTIGUOUS(x, m);
|
||||
|
||||
constexpr auto HALF = torch::headeronly::ScalarType::Half;
|
||||
constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
|
||||
constexpr auto INT = torch::headeronly::ScalarType::Int;
|
||||
constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
|
||||
|
||||
static constexpr int MXFP4_BLOCK_SIZE = 32;
|
||||
|
||||
static void validate_mxfp4_experts_quant_inputs(
|
||||
torch::stable::Tensor const& output,
|
||||
torch::stable::Tensor const& output_scale,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& input_offset_by_experts,
|
||||
torch::stable::Tensor const& output_scale_offset_by_experts,
|
||||
int64_t n_experts, int64_t m_topk, int64_t k) {
|
||||
CHECK_INPUT(output, "output");
|
||||
CHECK_INPUT(output_scale, "output_scale");
|
||||
CHECK_INPUT(input, "input");
|
||||
CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
|
||||
CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
|
||||
|
||||
STD_TORCH_CHECK(output.dim() == 2);
|
||||
STD_TORCH_CHECK(output_scale.dim() == 2);
|
||||
STD_TORCH_CHECK(input.dim() == 2);
|
||||
STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
|
||||
STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
|
||||
|
||||
STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
|
||||
STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
|
||||
STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
|
||||
// output is uint8 (two mxfp4 values packed into one uint8)
|
||||
// output_scale is int32 (four E8M0 values packed into one int32)
|
||||
STD_TORCH_CHECK(output.scalar_type() == UINT8);
|
||||
STD_TORCH_CHECK(output_scale.scalar_type() == INT);
|
||||
|
||||
STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
|
||||
STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
|
||||
STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
|
||||
STD_TORCH_CHECK(output.size(0) == m_topk);
|
||||
STD_TORCH_CHECK(output.size(1) == k / 2);
|
||||
int scales_k = k / MXFP4_BLOCK_SIZE;
|
||||
// K-dimension scale columns padded to a multiple of 4 for swizzle layout
|
||||
int padded_k = (scales_k + (4 - 1)) / 4 * 4;
|
||||
// 4 = 4 E8M0 values packed into one int32
|
||||
STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
|
||||
}
|
||||
|
||||
void mxfp4_experts_quant(
|
||||
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& input_offset_by_experts,
|
||||
torch::stable::Tensor const& output_scale_offset_by_experts,
|
||||
int64_t n_experts) {
|
||||
auto m_topk = input.size(0);
|
||||
auto k = input.size(1);
|
||||
|
||||
validate_mxfp4_experts_quant_inputs(
|
||||
output, output_scale, input, input_offset_by_experts,
|
||||
output_scale_offset_by_experts, n_experts, m_topk, k);
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
|
||||
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
|
||||
input_offset_by_experts.data_ptr(),
|
||||
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
|
||||
stream);
|
||||
});
|
||||
}
|
||||
|
||||
void silu_and_mul_mxfp4_experts_quant(
|
||||
torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& input_offset_by_experts,
|
||||
torch::stable::Tensor const& output_scale_offset_by_experts,
|
||||
int64_t n_experts) {
|
||||
auto m_topk = input.size(0);
|
||||
auto k_times_2 = input.size(1);
|
||||
STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
|
||||
auto k = k_times_2 / 2;
|
||||
|
||||
validate_mxfp4_experts_quant_inputs(
|
||||
output, output_scale, input, input_offset_by_experts,
|
||||
output_scale_offset_by_experts, n_experts, m_topk, k);
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
|
||||
output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
|
||||
input_offset_by_experts.data_ptr(),
|
||||
output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
|
||||
stream);
|
||||
});
|
||||
}
|
||||
@@ -116,6 +116,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
" Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
|
||||
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
|
||||
|
||||
// cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
|
||||
ops.def(
|
||||
"cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
|
||||
" Tensor a_blockscale, Tensor b_blockscales,"
|
||||
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
|
||||
|
||||
// Compute NVFP4 block quantized tensor.
|
||||
ops.def(
|
||||
"scaled_fp4_quant(Tensor input,"
|
||||
@@ -149,6 +155,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
|
||||
"Tensor output_scale_offset_by_experts) -> ()");
|
||||
|
||||
// Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
|
||||
ops.def(
|
||||
"mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
|
||||
"Tensor input, Tensor input_offset_by_experts,"
|
||||
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
|
||||
|
||||
// Fused SiLU+Mul+MXFP4 experts quantization.
|
||||
ops.def(
|
||||
"silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
|
||||
"output_scale,"
|
||||
"Tensor input, Tensor input_offset_by_experts,"
|
||||
"Tensor output_scale_offset_by_experts, int n_experts) -> ()");
|
||||
|
||||
// Fused SiLU+Mul+NVFP4 quantization.
|
||||
ops.def(
|
||||
"silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
|
||||
@@ -233,6 +252,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("silu_and_mul_scaled_fp4_experts_quant",
|
||||
TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
|
||||
ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
|
||||
ops.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
|
||||
ops.impl("silu_and_mul_mxfp4_experts_quant",
|
||||
TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
|
||||
|
||||
// W4A8 ops: impl registrations are in the source files
|
||||
// (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
}
|
||||
@@ -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"]
|
||||
|
||||
+7
-32
@@ -192,9 +192,10 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
|
||||
FROM base AS build_deep
|
||||
ARG ROCSHMEM_BRANCH="ba0bf0f3"
|
||||
ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
|
||||
ARG DEEPEP_BRANCH="e84464ec"
|
||||
ARG DEEPEP_BRANCH="5d90af8b"
|
||||
ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
|
||||
ARG DEEPEP_NIC="cx7"
|
||||
ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
|
||||
ENV ROCSHMEM_DIR=/opt/rocshmem
|
||||
|
||||
RUN git clone ${ROCSHMEM_REPO} \
|
||||
@@ -202,13 +203,11 @@ RUN git clone ${ROCSHMEM_REPO} \
|
||||
&& git checkout ${ROCSHMEM_BRANCH} \
|
||||
&& mkdir -p projects/rocshmem/build \
|
||||
&& cd projects/rocshmem/build \
|
||||
&& cmake .. \
|
||||
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
|
||||
-DROCM_PATH=/opt/rocm \
|
||||
-DCMAKE_POSITION_INDEPENDENT_CODE=ON \
|
||||
-DUSE_EXTERNAL_MPI=OFF \
|
||||
&& make -j \
|
||||
&& make install
|
||||
&& bash ../scripts/build_configs/all_backends \
|
||||
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
|
||||
-DROCM_PATH=/opt/rocm \
|
||||
-DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
|
||||
-DUSE_EXTERNAL_MPI=OFF
|
||||
|
||||
# Build DeepEP wheel.
|
||||
# DeepEP looks for rocshmem at ROCSHMEM_DIR.
|
||||
@@ -262,30 +261,6 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \
|
||||
&& echo "Detected vLLM version: ${VLLM_VERSION}" \
|
||||
&& echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
|
||||
|
||||
# Fail if git-based package dependencies are found in requirements files
|
||||
# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
|
||||
# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
|
||||
RUN echo "Checking for git-based packages in requirements files..." \
|
||||
&& echo "Checking common.txt for git-based packages:" \
|
||||
&& if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
|
||||
echo "ERROR: Git-based packages found in common.txt:"; \
|
||||
grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
|
||||
echo "Please publish these packages to PyPI instead of using git dependencies."; \
|
||||
exit 1; \
|
||||
else \
|
||||
echo " ✓ No git-based packages found in common.txt"; \
|
||||
fi \
|
||||
&& echo "Checking rocm.txt for git-based packages:" \
|
||||
&& if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
|
||||
echo "ERROR: Git-based packages found in rocm.txt:"; \
|
||||
grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
|
||||
echo "Please publish these packages to PyPI instead of using git dependencies."; \
|
||||
exit 1; \
|
||||
else \
|
||||
echo " ✓ No git-based packages found in rocm.txt"; \
|
||||
fi \
|
||||
&& echo "All requirements files are clean - no git-based packages found"
|
||||
|
||||
# Pin vLLM dependencies to exact versions of custom ROCm wheels
|
||||
# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
|
||||
COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -674,6 +682,24 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
|
||||
!!! note
|
||||
`VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed.
|
||||
|
||||
#### Realtime Transcription
|
||||
|
||||
Speech models that support streaming transcription via the
|
||||
[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api)
|
||||
WebSocket endpoint.
|
||||
|
||||
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|
||||
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
|
||||
| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | |
|
||||
| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | |
|
||||
|
||||
!!! note
|
||||
`VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`.
|
||||
|
||||
`Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`.
|
||||
You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'`
|
||||
when serving.
|
||||
|
||||
## Pooling Models
|
||||
|
||||
See [this page](pooling_models/README.md) for more information on how to use pooling models.
|
||||
|
||||
@@ -60,7 +60,7 @@ We currently support the following OpenAI APIs:
|
||||
- [Translation API](#translations-api) (`/v1/audio/translations`)
|
||||
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
|
||||
- [Realtime API](#realtime-api) (`/v1/realtime`)
|
||||
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
|
||||
- Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription).
|
||||
|
||||
In addition, we have the following custom APIs:
|
||||
|
||||
|
||||
+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,117 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Disaggregated multimodal serving: render → generate round-trip.
|
||||
|
||||
Demonstrates the two-phase disaggregated flow:
|
||||
1. /v1/chat/completions/render – preprocesses a multimodal chat request
|
||||
into token IDs and serialized tensor features.
|
||||
2. /inference/v1/generate – runs inference on the preprocessed tokens.
|
||||
|
||||
The render response is passed *directly* to generate with only
|
||||
``sampling_params`` added, showing that the two endpoints compose with
|
||||
zero client-side transformation.
|
||||
|
||||
Launch the server first:
|
||||
|
||||
vllm serve Qwen/Qwen3-VL-2B-Instruct \
|
||||
--dtype bfloat16 --max-model-len 4096 --enforce-eager
|
||||
|
||||
Then run this script:
|
||||
|
||||
python example_mm_serve.py
|
||||
"""
|
||||
|
||||
import io
|
||||
|
||||
import pybase64 as base64
|
||||
import requests
|
||||
from PIL import Image
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
BASE_URL = "http://localhost:8000"
|
||||
MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
|
||||
|
||||
|
||||
def make_data_url(image: Image.Image) -> str:
|
||||
"""Encode a PIL image as a base64 data URL."""
|
||||
buf = io.BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
b64 = base64.b64encode(buf.getvalue()).decode()
|
||||
return f"data:image/png;base64,{b64}"
|
||||
|
||||
|
||||
def main():
|
||||
# -- Step 1: Create a test image (solid red) -------------------------
|
||||
image = Image.new("RGB", (224, 224), color=(255, 0, 0))
|
||||
data_url = make_data_url(image)
|
||||
print("Created 224x224 red test image")
|
||||
|
||||
# -- Step 2: Render (preprocess) -------------------------------------
|
||||
render_payload = {
|
||||
"model": MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What color is this image? Answer in one word.",
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
print("\n--- Render ---")
|
||||
render_resp = requests.post(
|
||||
f"{BASE_URL}/v1/chat/completions/render", json=render_payload
|
||||
)
|
||||
render_resp.raise_for_status()
|
||||
render_data = render_resp.json()
|
||||
|
||||
print(f"Response keys: {list(render_data.keys())}")
|
||||
print(f"Number of token_ids: {len(render_data['token_ids'])}")
|
||||
|
||||
features = render_data.get("features")
|
||||
if features and features.get("kwargs_data"):
|
||||
print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}")
|
||||
for modality, items in features["kwargs_data"].items():
|
||||
print(
|
||||
f" {modality}: {len(items)} item(s), "
|
||||
f"first item type: {type(items[0])} length: {len(items[0])}"
|
||||
if items
|
||||
else "First item: (empty)"
|
||||
)
|
||||
else:
|
||||
print("WARNING: no kwargs_data in render response")
|
||||
|
||||
# -- Step 3: Generate (inference) ------------------------------------
|
||||
# Pass the render output directly — only add sampling_params.
|
||||
generate_payload = render_data
|
||||
generate_payload["sampling_params"] = {
|
||||
"max_tokens": 20,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
print("\n--- Generate ---")
|
||||
gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload)
|
||||
gen_resp.raise_for_status()
|
||||
gen_data = gen_resp.json()
|
||||
|
||||
# -- Step 4: Decode & print ------------------------------------------
|
||||
output_ids = gen_data["choices"][0]["token_ids"]
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
||||
text = tokenizer.decode(output_ids, skip_special_tokens=True)
|
||||
|
||||
print(f"Output token count: {len(output_ids)}")
|
||||
print(f"Generated text: {text!r}")
|
||||
|
||||
if "red" in text.lower():
|
||||
print("\nModel correctly identified the red image.")
|
||||
else:
|
||||
print(f"\nWARNING: Expected 'red' in output, got: {text!r}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -22,4 +22,4 @@ timm>=1.0.17
|
||||
# To be consistent with test_quark.py
|
||||
amd-quark>=0.8.99
|
||||
# Required for faster safetensors model loading
|
||||
fastsafetensors >= 0.2.2
|
||||
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2
|
||||
@@ -55,7 +55,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
|
||||
numba==0.61.2 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
|
||||
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
decord==0.6.0
|
||||
|
||||
@@ -277,7 +277,7 @@ fastar==0.10.0
|
||||
# via fastapi-cloud-cli
|
||||
fastparquet==2026.3.0
|
||||
# via genai-perf
|
||||
fastsafetensors==0.2.2
|
||||
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
|
||||
# via
|
||||
# -c requirements/rocm.txt
|
||||
# -r requirements/test/rocm.in
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Roundtrip tests for multimodal serde used by the disagg generate endpoint."""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.entrypoints.serve.disagg.mm_serde import (
|
||||
decode_mm_kwargs_item,
|
||||
encode_mm_kwargs_item,
|
||||
)
|
||||
from vllm.entrypoints.serve.disagg.protocol import (
|
||||
MultiModalFeatures,
|
||||
PlaceholderRangeInfo,
|
||||
)
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalBatchedField,
|
||||
MultiModalFieldElem,
|
||||
MultiModalFlatField,
|
||||
MultiModalKwargsItem,
|
||||
MultiModalSharedField,
|
||||
)
|
||||
|
||||
|
||||
def test_mm_kwargs_item_roundtrip():
|
||||
"""Full roundtrip test with all three field types and multiple dtypes."""
|
||||
e1 = MultiModalFieldElem(
|
||||
data=torch.zeros(1000, dtype=torch.bfloat16),
|
||||
field=MultiModalBatchedField(),
|
||||
)
|
||||
e2 = MultiModalFieldElem(
|
||||
data=torch.ones(100, dtype=torch.int32),
|
||||
field=MultiModalSharedField(batch_size=4),
|
||||
)
|
||||
e3 = MultiModalFieldElem(
|
||||
data=torch.randn(20, dtype=torch.float32),
|
||||
field=MultiModalFlatField(slices=[slice(0, 10), slice(10, 20)], dim=0),
|
||||
)
|
||||
|
||||
item = MultiModalKwargsItem({"pixel_values": e1, "grid_thw": e2, "embeds": e3})
|
||||
encoded = encode_mm_kwargs_item(item)
|
||||
|
||||
# Encoded result is a base64 string
|
||||
assert isinstance(encoded, str)
|
||||
|
||||
decoded = decode_mm_kwargs_item(encoded)
|
||||
|
||||
assert set(decoded.keys()) == {"pixel_values", "grid_thw", "embeds"}
|
||||
assert torch.equal(item["pixel_values"].data, decoded["pixel_values"].data)
|
||||
assert torch.equal(item["grid_thw"].data, decoded["grid_thw"].data)
|
||||
assert torch.equal(item["embeds"].data, decoded["embeds"].data)
|
||||
assert isinstance(decoded["pixel_values"].field, MultiModalBatchedField)
|
||||
assert isinstance(decoded["grid_thw"].field, MultiModalSharedField)
|
||||
assert isinstance(decoded["embeds"].field, MultiModalFlatField)
|
||||
|
||||
|
||||
def test_mm_kwargs_item_none_data():
|
||||
"""Roundtrip with None data field."""
|
||||
elem = MultiModalFieldElem(
|
||||
data=None,
|
||||
field=MultiModalSharedField(batch_size=2),
|
||||
)
|
||||
item = MultiModalKwargsItem({"empty": elem})
|
||||
encoded = encode_mm_kwargs_item(item)
|
||||
decoded = decode_mm_kwargs_item(encoded)
|
||||
|
||||
assert decoded["empty"].data is None
|
||||
assert isinstance(decoded["empty"].field, MultiModalSharedField)
|
||||
|
||||
|
||||
def test_mm_kwargs_item_nested_tensors():
|
||||
"""Roundtrip with nested tensor data."""
|
||||
nested = [torch.randn(3, 4), torch.randn(5, 4)]
|
||||
elem = MultiModalFieldElem(
|
||||
data=nested,
|
||||
field=MultiModalBatchedField(),
|
||||
)
|
||||
item = MultiModalKwargsItem({"nested": elem})
|
||||
encoded = encode_mm_kwargs_item(item)
|
||||
decoded = decode_mm_kwargs_item(encoded)
|
||||
|
||||
decoded_data = decoded["nested"].data
|
||||
assert len(decoded_data) == 2
|
||||
assert torch.equal(nested[0], decoded_data[0])
|
||||
assert torch.equal(nested[1], decoded_data[1])
|
||||
|
||||
|
||||
def test_mm_features_with_kwargs_data():
|
||||
"""Test that MultiModalFeatures can carry serialized tensor data."""
|
||||
elem = MultiModalFieldElem(
|
||||
data=torch.randn(5, 3, dtype=torch.float32),
|
||||
field=MultiModalBatchedField(),
|
||||
)
|
||||
item = MultiModalKwargsItem({"pixel_values": elem})
|
||||
encoded = encode_mm_kwargs_item(item)
|
||||
|
||||
features = MultiModalFeatures(
|
||||
mm_hashes={"image": ["abc123"]},
|
||||
mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=10)]},
|
||||
kwargs_data={"image": [encoded]},
|
||||
)
|
||||
|
||||
# JSON roundtrip
|
||||
json_str = features.model_dump_json()
|
||||
features2 = MultiModalFeatures.model_validate_json(json_str)
|
||||
|
||||
assert features2.mm_hashes == {"image": ["abc123"]}
|
||||
assert features2.kwargs_data is not None
|
||||
assert len(features2.kwargs_data["image"]) == 1
|
||||
|
||||
decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0])
|
||||
assert torch.equal(elem.data, decoded["pixel_values"].data)
|
||||
@@ -0,0 +1,158 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for multimodal features through the /inference/v1/generate endpoint.
|
||||
|
||||
Mirrors test_serving_tokens.py but exercises the multimodal piping
|
||||
using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render ->
|
||||
/generate -> /detokenize path. Intentionally avoids running the HF
|
||||
processor in the pytest parent process to keep os.fork() in sibling
|
||||
tests (e.g. test_weight_transfer_llm.py) deadlock-free.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from PIL import Image
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.multimodal.utils import encode_image_url
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
|
||||
GEN_ENDPOINT = "/inference/v1/generate"
|
||||
RENDER_ENDPOINT = "/v1/chat/completions/render"
|
||||
DETOKENIZE_ENDPOINT = "/detokenize"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def test_image():
|
||||
return Image.new("RGB", (224, 224), color=(255, 0, 0))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--enforce-eager",
|
||||
"--no-enable-prefix-caching",
|
||||
]
|
||||
|
||||
envs = os.environ.copy()
|
||||
envs["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=envs) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server: RemoteOpenAIServer):
|
||||
transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
|
||||
headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport,
|
||||
base_url=server.url_root,
|
||||
timeout=600,
|
||||
headers=headers,
|
||||
) as c:
|
||||
yield c
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_render_to_generate_roundtrip(client, test_image):
|
||||
"""End-to-end: render a multimodal chat -> feed into generate -> decode.
|
||||
|
||||
All preprocessing and detokenization happens in the server subprocess;
|
||||
the pytest parent never imports transformers or touches torch tensors.
|
||||
"""
|
||||
data_url = encode_image_url(test_image, format="PNG")
|
||||
|
||||
render_payload = {
|
||||
"model": MODEL_NAME,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What color is this image? Answer in one word.",
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
render_resp = await client.post(RENDER_ENDPOINT, json=render_payload)
|
||||
render_resp.raise_for_status()
|
||||
render_data = render_resp.json()
|
||||
|
||||
# Validate render output structure: keys exist and values are non-empty
|
||||
# and well-typed.
|
||||
assert "token_ids" in render_data
|
||||
assert isinstance(render_data["token_ids"], list)
|
||||
assert len(render_data["token_ids"]) > 0
|
||||
assert all(isinstance(t, int) for t in render_data["token_ids"])
|
||||
|
||||
assert "features" in render_data
|
||||
features = render_data["features"]
|
||||
assert features is not None
|
||||
assert isinstance(features, dict)
|
||||
|
||||
assert "mm_hashes" in features
|
||||
assert "image" in features["mm_hashes"]
|
||||
image_hashes = features["mm_hashes"]["image"]
|
||||
assert isinstance(image_hashes, list)
|
||||
assert len(image_hashes) > 0
|
||||
assert all(isinstance(h, str) and h for h in image_hashes)
|
||||
|
||||
assert "mm_placeholders" in features
|
||||
assert "image" in features["mm_placeholders"]
|
||||
image_placeholders = features["mm_placeholders"]["image"]
|
||||
assert isinstance(image_placeholders, list)
|
||||
assert len(image_placeholders) > 0
|
||||
for p in image_placeholders:
|
||||
assert isinstance(p.get("offset"), int)
|
||||
assert isinstance(p.get("length"), int)
|
||||
assert p["length"] > 0
|
||||
|
||||
assert "kwargs_data" in features
|
||||
assert "image" in features["kwargs_data"]
|
||||
assert len(features["kwargs_data"]["image"]) > 0
|
||||
|
||||
# Build generate request from render output
|
||||
generate_payload = render_data
|
||||
generate_payload["sampling_params"] = {
|
||||
"max_tokens": 10,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
gen_resp = await client.post(GEN_ENDPOINT, json=generate_payload)
|
||||
gen_resp.raise_for_status()
|
||||
gen_data = gen_resp.json()
|
||||
|
||||
assert "choices" in gen_data
|
||||
assert isinstance(gen_data["choices"], list)
|
||||
assert len(gen_data["choices"]) >= 1
|
||||
choice = gen_data["choices"][0]
|
||||
assert "token_ids" in choice
|
||||
assert isinstance(choice["token_ids"], list)
|
||||
assert len(choice["token_ids"]) > 0
|
||||
assert all(isinstance(t, int) for t in choice["token_ids"])
|
||||
|
||||
detok_resp = await client.post(
|
||||
DETOKENIZE_ENDPOINT,
|
||||
json={"model": MODEL_NAME, "tokens": choice["token_ids"]},
|
||||
)
|
||||
detok_resp.raise_for_status()
|
||||
detok_data = detok_resp.json()
|
||||
assert "prompt" in detok_data
|
||||
text = detok_data["prompt"]
|
||||
assert isinstance(text, str)
|
||||
assert len(text) > 0
|
||||
assert "red" in text.lower(), (
|
||||
f"Expected model to identify the red image, got: {text!r}"
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+45
-166
@@ -14,8 +14,6 @@ import pytest
|
||||
import torch
|
||||
from torch.nn import Parameter
|
||||
from torch.nn import functional as F
|
||||
from transformers import MixtralConfig
|
||||
from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
|
||||
|
||||
import vllm.model_executor.layers.fused_moe # noqa
|
||||
from tests.kernels.moe.utils import (
|
||||
@@ -24,10 +22,7 @@ from tests.kernels.moe.utils import (
|
||||
modular_triton_fused_moe,
|
||||
)
|
||||
from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed.parallel_state import init_distributed_environment
|
||||
from vllm.forward_context import get_forward_context, set_forward_context
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
MoEActivation,
|
||||
fused_topk,
|
||||
@@ -56,12 +51,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
|
||||
marlin_quantize,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
|
||||
from vllm.model_executor.models.mixtral import MixtralMoE
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import ScalarType, scalar_types
|
||||
from vllm.utils.math_utils import next_power_of_2
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.worker.workspace import init_workspace_manager
|
||||
|
||||
|
||||
def iterative_moe(
|
||||
@@ -150,12 +143,14 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [
|
||||
{
|
||||
"a_type": [scalar_types.bfloat16],
|
||||
"b_type": scalar_types.float4_e2m1f,
|
||||
"c_type": [scalar_types.bfloat16],
|
||||
"group_blocks": [2],
|
||||
},
|
||||
# MXFP8
|
||||
{
|
||||
"a_type": [scalar_types.bfloat16],
|
||||
"b_type": scalar_types.float8_e4m3fn,
|
||||
"c_type": [scalar_types.bfloat16],
|
||||
"group_blocks": [2],
|
||||
},
|
||||
# AWQ-INT4 with INT8 activation
|
||||
@@ -681,154 +676,35 @@ def test_fused_moe_wn16(
|
||||
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("padding", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_mixtral_moe(
|
||||
default_vllm_config,
|
||||
dist_init,
|
||||
dtype: torch.dtype,
|
||||
padding: bool,
|
||||
use_rocm_aiter: bool,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Make sure our Mixtral MoE implementation agrees with the one from
|
||||
huggingface."""
|
||||
|
||||
# Explicitly set AITER env var based on test parameter to ensure
|
||||
# consistent behavior regardless of external environment
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
if use_rocm_aiter and dtype == torch.float32:
|
||||
pytest.skip("AITER ROCm test skip for float32")
|
||||
|
||||
monkeypatch.setenv("RANK", "0")
|
||||
monkeypatch.setenv("LOCAL_RANK", "0")
|
||||
monkeypatch.setenv("WORLD_SIZE", "1")
|
||||
monkeypatch.setenv("MASTER_ADDR", "localhost")
|
||||
monkeypatch.setenv("MASTER_PORT", "12345")
|
||||
init_distributed_environment()
|
||||
init_workspace_manager(torch.accelerator.current_device_index())
|
||||
|
||||
# Instantiate our and huggingface's MoE blocks
|
||||
vllm_config.compilation_config.static_forward_context = dict()
|
||||
with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
|
||||
config = MixtralConfig()
|
||||
hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
|
||||
vllm_moe = MixtralMoE(
|
||||
num_experts=config.num_local_experts,
|
||||
top_k=config.num_experts_per_tok,
|
||||
hidden_size=config.hidden_size,
|
||||
intermediate_size=config.intermediate_size,
|
||||
params_dtype=dtype,
|
||||
tp_size=1,
|
||||
dp_size=1,
|
||||
).cuda()
|
||||
|
||||
# Load the weights
|
||||
vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data
|
||||
if isinstance(hf_moe.experts, torch.nn.ModuleList):
|
||||
# Transformers v4
|
||||
for i in range(config.num_local_experts):
|
||||
weights = (
|
||||
hf_moe.experts[i].w1.weight.data,
|
||||
hf_moe.experts[i].w3.weight.data,
|
||||
)
|
||||
vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
|
||||
vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
|
||||
else:
|
||||
# Transformers v5
|
||||
vllm_moe.experts.w13_weight.data[:] = hf_moe.experts.gate_up_proj.data
|
||||
vllm_moe.experts.w2_weight.data[:] = hf_moe.experts.down_proj.data
|
||||
# TODO: remove this line after https://github.com/huggingface/transformers/pull/43622
|
||||
hf_moe.experts.config._experts_implementation = "eager"
|
||||
|
||||
# Generate input batch of dimensions [batch_size, seq_len, hidden_dim]
|
||||
hf_inputs = torch.randn((1, 64, config.hidden_size)).to(dtype).to("cuda")
|
||||
# vLLM uses 1D query [num_tokens, hidden_dim]
|
||||
vllm_inputs = hf_inputs.flatten(0, 1)
|
||||
|
||||
# Pad the weight if moe padding is enabled
|
||||
if padding:
|
||||
vllm_moe.experts.w13_weight = Parameter(
|
||||
F.pad(vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[
|
||||
..., 0:-128
|
||||
],
|
||||
requires_grad=False,
|
||||
)
|
||||
vllm_moe.experts.w2_weight = Parameter(
|
||||
F.pad(vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128],
|
||||
requires_grad=False,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
# FIXME (zyongye) fix this after we move self.kernel
|
||||
# assignment in FusedMoE.__init__
|
||||
|
||||
vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts)
|
||||
|
||||
# need to override the forward context for unittests, otherwise it assumes
|
||||
# we're running the model forward pass (the model specified in vllm_config)
|
||||
get_forward_context().all_moe_layers = None
|
||||
|
||||
# Run forward passes for both MoE blocks
|
||||
hf_states = hf_moe.forward(hf_inputs)
|
||||
if isinstance(hf_states, tuple):
|
||||
# Transformers v4
|
||||
hf_states = hf_states[0]
|
||||
vllm_states = vllm_moe.forward(vllm_inputs)
|
||||
|
||||
mixtral_moe_tol = {
|
||||
torch.float32: 1e-3,
|
||||
torch.float16: 1e-3,
|
||||
torch.bfloat16: 1e-2,
|
||||
}
|
||||
|
||||
if use_rocm_aiter:
|
||||
# The values of rtol and atol are set based on the tests in ROCM AITER package.
|
||||
# https://github.com/ROCm/aiter/blob/dfed377f4be7da96ca2d75ac0761f569676f7240/op_tests/test_moe.py#L174
|
||||
torch.testing.assert_close(
|
||||
hf_states.flatten(0, 1), vllm_states, rtol=0.01, atol=100
|
||||
)
|
||||
else:
|
||||
torch.testing.assert_close(
|
||||
hf_states.flatten(0, 1),
|
||||
vllm_states,
|
||||
rtol=mixtral_moe_tol[dtype],
|
||||
atol=mixtral_moe_tol[dtype],
|
||||
)
|
||||
MARLIN_MOE_SCENARIOS = [
|
||||
# (m, n, k, e, topk, ep_size, act_order, is_k_full)
|
||||
# No act_order: is_k_full=True matches usual case (marlin_is_k_full).
|
||||
# N>=256 required for Marlin kernel thread config for MXFP8.
|
||||
# Single token, small matrices
|
||||
(1, 128, 256, 5, 2, 1, False, True),
|
||||
# Single token, large matrices
|
||||
(1, 1024, 2048, 5, 2, 1, False, True),
|
||||
# Unaligned m, small matrices
|
||||
(133, 256, 256, 5, 2, 1, False, True),
|
||||
# Unaligned m, large matrices
|
||||
(133, 1024, 2048, 12, 3, 1, False, True),
|
||||
# Aligned batch, small matrices
|
||||
(128, 256, 256, 5, 2, 1, False, True),
|
||||
# Aligned batch, large matrices
|
||||
(128, 1024, 2048, 12, 3, 1, False, True),
|
||||
# Expert parallelism
|
||||
(64, 1024, 2048, 12, 3, 4, False, True),
|
||||
# Act order with is_k_full=True (no tensor parallelism)
|
||||
(1, 1024, 2048, 5, 2, 1, True, True),
|
||||
# Act order with is_k_full=False (tensor parallelism)
|
||||
(133, 256, 256, 5, 2, 1, True, False),
|
||||
]
|
||||
|
||||
|
||||
def marlin_moe_generate_valid_test_cases():
|
||||
import itertools
|
||||
|
||||
m_list = [1, 123, 666]
|
||||
n_list = [128, 1024]
|
||||
k_list = [256, 2048]
|
||||
e_list = [5, 12]
|
||||
topk_list = [2, 3]
|
||||
ep_size_list = [1, 4]
|
||||
act_order_list = [True, False]
|
||||
is_k_full_list = [True, False]
|
||||
|
||||
all_combinations = itertools.product(
|
||||
MOE_MARLIN_QUANT_TEST_CONFIGS,
|
||||
m_list,
|
||||
n_list,
|
||||
k_list,
|
||||
e_list,
|
||||
topk_list,
|
||||
ep_size_list,
|
||||
act_order_list,
|
||||
is_k_full_list,
|
||||
)
|
||||
|
||||
def is_invalid(
|
||||
def is_valid(
|
||||
a_type,
|
||||
b_type,
|
||||
c_type,
|
||||
@@ -845,39 +721,42 @@ def marlin_moe_generate_valid_test_cases():
|
||||
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
|
||||
if group_size > 0 and k % group_size != 0:
|
||||
return False
|
||||
|
||||
if act_order and group_size in [-1, k, n]:
|
||||
return False
|
||||
if group_size in [k, n]:
|
||||
return False
|
||||
if not act_order and is_k_full:
|
||||
if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
|
||||
return False
|
||||
|
||||
return a_type.size_bits < 16 or a_type is c_type
|
||||
|
||||
cases = []
|
||||
for case in all_combinations:
|
||||
quant_test_config, m, n, k, _, _, _, act_order, *_ = case
|
||||
if act_order and not quant_test_config.get("support_act_order", False):
|
||||
continue
|
||||
|
||||
for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS:
|
||||
f16_types = [scalar_types.float16]
|
||||
inner_combinations = itertools.product(
|
||||
quant_test_config.get("a_type", f16_types),
|
||||
[quant_test_config["b_type"]],
|
||||
quant_test_config.get("c_type", f16_types),
|
||||
quant_test_config["group_blocks"],
|
||||
inner_combinations = list(
|
||||
itertools.product(
|
||||
quant_test_config.get("a_type", f16_types),
|
||||
[quant_test_config["b_type"]],
|
||||
quant_test_config.get("c_type", f16_types),
|
||||
quant_test_config["group_blocks"],
|
||||
)
|
||||
)
|
||||
|
||||
supports_act_order = quant_test_config.get("support_act_order", False)
|
||||
|
||||
for sub_case in inner_combinations:
|
||||
if (
|
||||
sub_case[0] == scalar_types.float8_e4m3fn
|
||||
and current_platform.get_device_capability() not in [89, 120]
|
||||
):
|
||||
continue
|
||||
args = sub_case + (m, n, k) + case[4:]
|
||||
if is_invalid(*args):
|
||||
cases.append(args)
|
||||
|
||||
for scenario in MARLIN_MOE_SCENARIOS:
|
||||
m, n, k, e, topk, ep_size, act_order, is_k_full = scenario
|
||||
if act_order and not supports_act_order:
|
||||
continue
|
||||
args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full)
|
||||
if is_valid(*args):
|
||||
cases.append(args)
|
||||
return cases
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels."""
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import torch_moe_single
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
random.seed(42)
|
||||
set_random_seed(42)
|
||||
|
||||
MXFP4_BLOCK_SIZE = 32
|
||||
|
||||
|
||||
def align(val: int, alignment: int = 128) -> int:
|
||||
return int((val + alignment - 1) // alignment * alignment)
|
||||
|
||||
|
||||
def calc_diff(x, y):
|
||||
x, y = x.double(), y.double()
|
||||
denominator = (x * x + y * y).sum()
|
||||
sim = 2 * (x * y).sum() / denominator
|
||||
return 1 - sim
|
||||
|
||||
|
||||
def is_sm100_supported() -> bool:
|
||||
return current_platform.is_cuda() and current_platform.is_device_capability_family(
|
||||
100
|
||||
)
|
||||
|
||||
|
||||
def compute_ref_output(
|
||||
input_tensor: torch.Tensor,
|
||||
weight_list: list[torch.Tensor],
|
||||
expert_offsets: list[int],
|
||||
expert_offset: int,
|
||||
num_experts: int,
|
||||
) -> torch.Tensor:
|
||||
"""Reference output using torch_moe_single with top-1 routing."""
|
||||
score = torch.full(
|
||||
(expert_offset, num_experts),
|
||||
-1e9,
|
||||
device=input_tensor.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
for g in range(num_experts):
|
||||
start = expert_offsets[g]
|
||||
end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset
|
||||
score[start:end, g] = 0.0
|
||||
|
||||
return torch_moe_single(
|
||||
input_tensor, torch.stack(weight_list, dim=0), score, topk=1
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_sm100_supported(),
|
||||
reason="cutlass_mxfp4_group_mm requires CUDA SM100",
|
||||
)
|
||||
@pytest.mark.parametrize("num_experts", [8, 16, 32])
|
||||
@pytest.mark.parametrize("out_dtype", [torch.bfloat16])
|
||||
def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype):
|
||||
"""
|
||||
Test the MXFP4 grouped GEMM kernel by:
|
||||
1. Creating random per-expert inputs and weights
|
||||
2. Quantizing both to MXFP4 using the CUDA kernel
|
||||
3. Running the CUTLASS grouped GEMM
|
||||
4. Comparing against BF16 reference
|
||||
"""
|
||||
device = "cuda"
|
||||
alignment = 128
|
||||
# N and K must be multiples of 128 for clean swizzle layout
|
||||
n_g = random.randint(1, 16) * alignment
|
||||
k_g = random.randint(1, 16) * alignment
|
||||
|
||||
expert_offset = 0
|
||||
expert_offsets_input = []
|
||||
problem_sizes = []
|
||||
input_list = []
|
||||
weight_list = []
|
||||
|
||||
for g in range(num_experts):
|
||||
m_g = random.randint(1, 256)
|
||||
expert_offsets_input.append(expert_offset)
|
||||
expert_offset += m_g
|
||||
problem_sizes.append([m_g, n_g, k_g])
|
||||
|
||||
input_list.append(
|
||||
torch.normal(0.0, std=0.5, size=(m_g, k_g), device=device, dtype=out_dtype)
|
||||
)
|
||||
weight_list.append(
|
||||
torch.normal(0.0, std=0.5, size=(n_g, k_g), device=device, dtype=out_dtype)
|
||||
)
|
||||
|
||||
input_tensor = torch.concat(input_list, dim=0) # [M_total, K]
|
||||
|
||||
# --- Quantize INPUTS via mxfp4_experts_quant ---
|
||||
input_bs_offsets = []
|
||||
tot = 0
|
||||
for g in range(num_experts):
|
||||
input_bs_offsets.append(tot)
|
||||
tot += align(problem_sizes[g][0], 128)
|
||||
input_bs_offsets.append(tot)
|
||||
|
||||
_inp_expert_offsets = torch.tensor(
|
||||
expert_offsets_input + [expert_offset], device=device, dtype=torch.int32
|
||||
)
|
||||
_inp_bs_offsets = torch.tensor(input_bs_offsets, device=device, dtype=torch.int32)
|
||||
|
||||
input_quant, input_sf = ops.mxfp4_experts_quant(
|
||||
input_tensor,
|
||||
_inp_expert_offsets,
|
||||
_inp_bs_offsets,
|
||||
num_experts,
|
||||
topk=1,
|
||||
)
|
||||
|
||||
# --- Quantize WEIGHTS via mxfp4_experts_quant ---
|
||||
# Treat each expert's N weight rows as an "expert" with N tokens
|
||||
weight_tensor = torch.concat(weight_list, dim=0) # [E*N, K]
|
||||
weight_expert_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
|
||||
# N is always multiple of 128, so blockscale offsets are clean
|
||||
weight_bs_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
|
||||
|
||||
_wt_expert_offsets = torch.tensor(
|
||||
weight_expert_offsets, device=device, dtype=torch.int32
|
||||
)
|
||||
_wt_bs_offsets = torch.tensor(weight_bs_offsets, device=device, dtype=torch.int32)
|
||||
|
||||
weight_quant, weight_sf = ops.mxfp4_experts_quant(
|
||||
weight_tensor,
|
||||
_wt_expert_offsets,
|
||||
_wt_bs_offsets,
|
||||
num_experts,
|
||||
topk=1,
|
||||
)
|
||||
|
||||
# Reshape weight quantized data to [E, N, K//2]
|
||||
weight_quant = weight_quant[: num_experts * n_g].view(num_experts, n_g, k_g // 2)
|
||||
|
||||
# Reshape weight scale factors to [E, N, K//32]
|
||||
# The quant kernel produces uint8 SF buffer. Each row has K//32 SFs.
|
||||
scales_per_row = k_g // MXFP4_BLOCK_SIZE
|
||||
weight_sf_flat = weight_sf.view(-1)[: num_experts * n_g * scales_per_row]
|
||||
weight_sf_3d = weight_sf_flat.view(num_experts, n_g, scales_per_row)
|
||||
|
||||
# Output
|
||||
output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
|
||||
|
||||
_problem_sizes = torch.tensor(problem_sizes, device=device, dtype=torch.int32)
|
||||
_expert_offsets = torch.tensor(
|
||||
expert_offsets_input, device=device, dtype=torch.int32
|
||||
)
|
||||
_input_bs = torch.tensor(input_bs_offsets[:-1], device=device, dtype=torch.int32)
|
||||
|
||||
# Run the MXFP4 grouped GEMM
|
||||
ops.cutlass_mxfp4_moe_mm(
|
||||
output,
|
||||
input_quant,
|
||||
weight_quant,
|
||||
input_sf,
|
||||
weight_sf_3d,
|
||||
_problem_sizes,
|
||||
_expert_offsets,
|
||||
_input_bs,
|
||||
)
|
||||
|
||||
# Reference: BF16 matmul
|
||||
ref_output = compute_ref_output(
|
||||
input_tensor=input_tensor,
|
||||
weight_list=weight_list,
|
||||
expert_offsets=expert_offsets_input,
|
||||
expert_offset=expert_offset,
|
||||
num_experts=num_experts,
|
||||
)
|
||||
|
||||
# Compare per-expert
|
||||
for g in range(num_experts):
|
||||
start = expert_offsets_input[g]
|
||||
end = expert_offsets_input[g + 1] if g + 1 < num_experts else expert_offset
|
||||
if start == end:
|
||||
continue
|
||||
baseline = ref_output[start:end]
|
||||
actual = output[start:end]
|
||||
diff = calc_diff(actual, baseline)
|
||||
print(
|
||||
f"m_g={end - start} n_g={n_g} k_g={k_g} "
|
||||
f"num_experts={num_experts}, "
|
||||
f"out_dtype={out_dtype}, diff={diff:.5f}"
|
||||
)
|
||||
# FP4 quantization is very lossy (~4 bits precision)
|
||||
# Comparing quantized vs full-precision gives cosine diff of 0.05-0.15
|
||||
assert diff < 0.15, f"Expert {g}: diff={diff:.5f} exceeds threshold"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_sm100_supported(),
|
||||
reason="mxfp4_experts_quant requires CUDA SM100",
|
||||
)
|
||||
def test_mxfp4_experts_quant_basic():
|
||||
"""
|
||||
Basic smoke test for the MXFP4 experts quantization kernel.
|
||||
"""
|
||||
device = "cuda"
|
||||
num_experts = 4
|
||||
k = 256
|
||||
tokens_per_expert = 16
|
||||
|
||||
total_tokens = tokens_per_expert * num_experts
|
||||
input_tensor = torch.randn(total_tokens, k, device=device, dtype=torch.bfloat16) / 5
|
||||
|
||||
expert_offsets = [i * tokens_per_expert for i in range(num_experts + 1)]
|
||||
blockscale_offsets = [
|
||||
align(i * tokens_per_expert, 128) for i in range(num_experts + 1)
|
||||
]
|
||||
|
||||
_expert_offsets = torch.tensor(expert_offsets, device=device, dtype=torch.int32)
|
||||
_blockscale_offsets = torch.tensor(
|
||||
blockscale_offsets, device=device, dtype=torch.int32
|
||||
)
|
||||
|
||||
output, output_sf = ops.mxfp4_experts_quant(
|
||||
input_tensor,
|
||||
_expert_offsets,
|
||||
_blockscale_offsets,
|
||||
num_experts,
|
||||
topk=1,
|
||||
)
|
||||
|
||||
assert output.shape == (total_tokens, k // 2)
|
||||
assert output.dtype == torch.uint8
|
||||
assert output_sf.dtype == torch.uint8
|
||||
assert output.any(), "Quantized output is all zeros"
|
||||
print(
|
||||
f"MXFP4 experts quant: output shape={output.shape}, sf shape={output_sf.shape}"
|
||||
)
|
||||
print("PASSED")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
@@ -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)
|
||||
|
||||
@@ -26,11 +26,8 @@ def test_placeholder_range_get_num_embeds(is_embed, expected):
|
||||
"is_embed,expected",
|
||||
[
|
||||
(None, None),
|
||||
(
|
||||
torch.tensor([False, True, False, True, True]),
|
||||
torch.tensor([0, 1, 1, 2, 3]),
|
||||
),
|
||||
(torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
|
||||
(torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]),
|
||||
(torch.tensor([True, True, True]), [1, 2, 3]),
|
||||
],
|
||||
)
|
||||
def test_placeholder_range_embeds_cumsum(is_embed, expected):
|
||||
@@ -41,6 +38,6 @@ def test_placeholder_range_embeds_cumsum(is_embed, expected):
|
||||
assert pr.embeds_cumsum is None
|
||||
return
|
||||
|
||||
assert torch.equal(pr.embeds_cumsum, expected)
|
||||
assert pr.embeds_cumsum == expected
|
||||
# cached_property should return the same object on repeated access
|
||||
assert pr.embeds_cumsum is pr.embeds_cumsum
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.turboquant.config import (
|
||||
from vllm.model_executor.layers.quantization.turboquant.quantizer import (
|
||||
generate_wht_signs,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import next_power_of_2
|
||||
|
||||
# ============================================================================
|
||||
@@ -345,7 +346,8 @@ class TestLloydMax:
|
||||
# Rotation matrix tests (GPU required)
|
||||
# ============================================================================
|
||||
|
||||
CUDA_AVAILABLE = torch.cuda.is_available()
|
||||
GPGPU_AVAILABLE = torch.cuda.is_available() or torch.xpu.is_available()
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
|
||||
def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
|
||||
@@ -360,16 +362,16 @@ def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Te
|
||||
return Q.to(device)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
|
||||
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
|
||||
class TestRotationMatrix:
|
||||
"""Tests for the QR-based rotation (standalone benchmarks only)."""
|
||||
|
||||
@pytest.mark.parametrize("dim", [64, 96, 128, 256])
|
||||
def test_rotation_matrix_shape_and_orthogonal(self, dim):
|
||||
Pi = generate_rotation_matrix(dim, seed=42, device="cuda")
|
||||
Pi = generate_rotation_matrix(dim, seed=42, device=DEVICE_TYPE)
|
||||
assert Pi.shape == (dim, dim)
|
||||
eye = Pi @ Pi.T
|
||||
assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
|
||||
assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
|
||||
f"Pi not orthogonal for dim={dim}"
|
||||
)
|
||||
|
||||
@@ -385,7 +387,7 @@ class TestRotationMatrix:
|
||||
|
||||
def test_rotation_matrix_det_is_pm1(self):
|
||||
"""Orthogonal matrix determinant must be +1 or -1."""
|
||||
Pi = generate_rotation_matrix(128, seed=42, device="cuda")
|
||||
Pi = generate_rotation_matrix(128, seed=42, device=DEVICE_TYPE)
|
||||
det = torch.linalg.det(Pi)
|
||||
assert abs(abs(det.item()) - 1.0) < 1e-4
|
||||
|
||||
@@ -403,31 +405,31 @@ def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
|
||||
return (H / math.sqrt(d)).to(torch.device(device))
|
||||
|
||||
|
||||
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
|
||||
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
|
||||
class TestWHTRotation:
|
||||
"""Tests for the WHT rotation actually used in serving."""
|
||||
|
||||
@pytest.mark.parametrize("dim", [64, 128, 256])
|
||||
def test_wht_orthonormal(self, dim):
|
||||
"""signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
|
||||
signs = generate_wht_signs(dim, seed=42, device="cuda")
|
||||
H = _build_hadamard(dim, "cuda")
|
||||
signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE)
|
||||
H = _build_hadamard(dim, DEVICE_TYPE)
|
||||
PiT = (signs.unsqueeze(1) * H).contiguous()
|
||||
eye = PiT @ PiT.T
|
||||
assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
|
||||
assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
|
||||
f"WHT rotation not orthonormal for dim={dim}"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("dim", [64, 128, 256])
|
||||
def test_wht_self_inverse(self, dim):
|
||||
"""PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
|
||||
signs = generate_wht_signs(dim, seed=42, device="cuda")
|
||||
H = _build_hadamard(dim, "cuda")
|
||||
signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE)
|
||||
H = _build_hadamard(dim, DEVICE_TYPE)
|
||||
PiT = (signs.unsqueeze(1) * H).contiguous()
|
||||
Pi = PiT.T.contiguous()
|
||||
# Pi @ PiT should be identity (rotation then inverse)
|
||||
result = Pi @ PiT
|
||||
assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
|
||||
assert torch.allclose(result, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
|
||||
f"WHT rotation not self-inverse for dim={dim}"
|
||||
)
|
||||
|
||||
@@ -454,7 +456,7 @@ class TestWHTRotation:
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
|
||||
@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
|
||||
class TestStoreDecodeRoundTrip:
|
||||
"""End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""
|
||||
|
||||
@@ -487,11 +489,11 @@ class TestStoreDecodeRoundTrip:
|
||||
block_size = 16
|
||||
num_blocks = 1
|
||||
|
||||
device = torch.device("cuda")
|
||||
device = torch.device(DEVICE_TYPE)
|
||||
|
||||
# Generate rotation
|
||||
signs = generate_wht_signs(D, seed=42, device=device)
|
||||
H = _build_hadamard(D, "cuda")
|
||||
H = _build_hadamard(D, DEVICE_TYPE)
|
||||
PiT = (signs.unsqueeze(1) * H).contiguous().float()
|
||||
Pi = PiT.T.contiguous()
|
||||
|
||||
|
||||
@@ -17,22 +17,6 @@ from vllm import LLM, SamplingParams
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="In V1, we reject tokens > max_seq_len")
|
||||
def test_duplicated_ignored_sequence_group():
|
||||
"""https://github.com/vllm-project/vllm/issues/1655"""
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.01, top_p=0.1, max_tokens=256)
|
||||
llm = LLM(
|
||||
model="distilbert/distilgpt2",
|
||||
max_num_batched_tokens=4096,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
prompts = ["This is a short prompt", "This is a very long prompt " * 1000]
|
||||
outputs = llm.generate(prompts, sampling_params=sampling_params)
|
||||
|
||||
assert len(prompts) == len(outputs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
|
||||
@@ -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"
|
||||
),
|
||||
|
||||
@@ -7,17 +7,39 @@ from tests.models.utils import check_embeddings_close
|
||||
from vllm.utils.serial_utils import (
|
||||
EMBED_DTYPES,
|
||||
ENDIANNESS,
|
||||
MM_METADATA_DTYPES,
|
||||
EmbedDType,
|
||||
Endianness,
|
||||
MmMetadataDType,
|
||||
binary2tensor,
|
||||
tensor2binary,
|
||||
)
|
||||
|
||||
FLOAT_EMBED_DTYPES = tuple(EMBED_DTYPES.keys())
|
||||
INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys())
|
||||
|
||||
|
||||
def _build_integer_tensor(
|
||||
embed_dtype: MmMetadataDType, shape: tuple[int, ...]
|
||||
) -> torch.Tensor:
|
||||
torch_dtype = MM_METADATA_DTYPES[embed_dtype].torch_dtype
|
||||
|
||||
if torch_dtype is torch.bool:
|
||||
return torch.randint(0, 2, shape, dtype=torch.int32).to(torch.bool)
|
||||
if torch_dtype is torch.uint8:
|
||||
return torch.randint(0, 256, shape, dtype=torch.uint8)
|
||||
if torch_dtype is torch.int32:
|
||||
return torch.randint(-(2**20), 2**20, shape, dtype=torch.int32)
|
||||
if torch_dtype is torch.int64:
|
||||
return torch.randint(-(2**62), 2**62, shape, dtype=torch.int64)
|
||||
|
||||
raise AssertionError(f"Unsupported non-floating embed dtype: {embed_dtype}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endianness", ENDIANNESS)
|
||||
@pytest.mark.parametrize("embed_dtype", EMBED_DTYPES.keys())
|
||||
@pytest.mark.parametrize("embed_dtype", FLOAT_EMBED_DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
|
||||
def test_encode_and_decode_floats(embed_dtype: EmbedDType, endianness: Endianness):
|
||||
for i in range(10):
|
||||
tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32)
|
||||
shape = tensor.shape
|
||||
@@ -40,3 +62,20 @@ def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
|
||||
name_1="new",
|
||||
tol=1e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endianness", ENDIANNESS)
|
||||
@pytest.mark.parametrize("embed_dtype", INTEGER_EMBED_DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_encode_and_decode_integers(
|
||||
embed_dtype: MmMetadataDType, endianness: Endianness
|
||||
):
|
||||
shape = (2, 3, 5, 7, 11, 13)
|
||||
|
||||
for i in range(10):
|
||||
tensor = _build_integer_tensor(embed_dtype, shape)
|
||||
binary = tensor2binary(tensor, embed_dtype, endianness)
|
||||
new_tensor = binary2tensor(binary, shape, embed_dtype, endianness)
|
||||
|
||||
assert new_tensor.dtype == MM_METADATA_DTYPES[embed_dtype].torch_dtype
|
||||
torch.testing.assert_close(tensor, new_tensor, atol=0, rtol=0)
|
||||
|
||||
@@ -1150,6 +1150,38 @@ def cutlass_fp4_moe_mm(
|
||||
)
|
||||
|
||||
|
||||
def cutlass_mxfp4_moe_mm(
|
||||
out_tensors: torch.Tensor,
|
||||
a_tensors: torch.Tensor,
|
||||
b_tensors: torch.Tensor,
|
||||
a_scales: torch.Tensor,
|
||||
b_scales: torch.Tensor,
|
||||
problem_sizes: torch.Tensor,
|
||||
expert_offsets: torch.Tensor,
|
||||
sf_offsets: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
|
||||
|
||||
Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
|
||||
- a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
|
||||
- a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout)
|
||||
- Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales).
|
||||
- expert_offsets/sf_offsets: expert boundary indices
|
||||
- problem_sizes: (num_experts, 3) with (M, N, K) per expert
|
||||
"""
|
||||
return torch.ops._C.cutlass_mxfp4_group_mm(
|
||||
out_tensors,
|
||||
a_tensors,
|
||||
b_tensors,
|
||||
a_scales,
|
||||
b_scales,
|
||||
problem_sizes,
|
||||
expert_offsets,
|
||||
sf_offsets,
|
||||
)
|
||||
|
||||
|
||||
def mxfp8_experts_quant(
|
||||
input_tensor: torch.Tensor,
|
||||
problem_sizes: torch.Tensor,
|
||||
@@ -1848,6 +1880,109 @@ def silu_and_mul_scaled_fp4_experts_quant(
|
||||
return output, output_scales
|
||||
|
||||
|
||||
def mxfp4_experts_quant(
|
||||
input_tensor: torch.Tensor,
|
||||
expert_offsets: torch.Tensor,
|
||||
blockscale_offsets: torch.Tensor,
|
||||
n_experts: int,
|
||||
topk: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Quantize input tensor to MXFP4 for packed MoE inputs.
|
||||
Uses 32-element blocks with E8M0 (power-of-two) scale factors.
|
||||
MXFP4 has no global scale - only block-level E8M0 scale factors.
|
||||
|
||||
Args:
|
||||
input_tensor: [m_topk, k] BF16/FP16 activations
|
||||
expert_offsets: [n_experts+1] token boundaries per expert
|
||||
blockscale_offsets: [n_experts+1] SF row boundaries per expert
|
||||
n_experts: number of experts
|
||||
topk: number of top-k experts
|
||||
Returns:
|
||||
output: [m_topk, k//2] packed E2M1 values (uint8)
|
||||
output_scales: E8M0 blockscales in swizzled layout (uint8 view)
|
||||
"""
|
||||
assert not current_platform.is_rocm()
|
||||
assert input_tensor.ndim == 2
|
||||
|
||||
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
|
||||
m_numtopk, k = input_tensor.shape
|
||||
|
||||
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
|
||||
f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
|
||||
f"{MAX_TOKENS_PER_EXPERT})"
|
||||
f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use"
|
||||
f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
|
||||
)
|
||||
scales_k = k // 32
|
||||
padded_k = (scales_k + (4 - 1)) // 4
|
||||
|
||||
output = torch.empty(
|
||||
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
|
||||
)
|
||||
output_scales = torch.empty(
|
||||
MAX_TOKENS_PER_EXPERT * topk,
|
||||
padded_k,
|
||||
dtype=torch.int32,
|
||||
device=input_tensor.device,
|
||||
)
|
||||
torch.ops._C.mxfp4_experts_quant(
|
||||
output,
|
||||
output_scales,
|
||||
input_tensor,
|
||||
expert_offsets,
|
||||
blockscale_offsets,
|
||||
n_experts,
|
||||
)
|
||||
# E8M0 SFs are stored as uint8
|
||||
output_scales = output_scales.view(torch.uint8)
|
||||
return output, output_scales
|
||||
|
||||
|
||||
def silu_and_mul_mxfp4_experts_quant(
|
||||
input_tensor: torch.Tensor,
|
||||
expert_offsets: torch.Tensor,
|
||||
blockscale_offsets: torch.Tensor,
|
||||
n_experts: int,
|
||||
topk: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
|
||||
MXFP4 has no global scale - only block-level E8M0 scale factors.
|
||||
"""
|
||||
assert not current_platform.is_rocm()
|
||||
assert input_tensor.ndim == 2
|
||||
|
||||
MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
|
||||
m_numtopk, k_times_2 = input_tensor.shape
|
||||
assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
|
||||
k = k_times_2 // 2
|
||||
|
||||
assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk
|
||||
scales_k = k // 32
|
||||
padded_k = (scales_k + (4 - 1)) // 4
|
||||
|
||||
output = torch.empty(
|
||||
m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
|
||||
)
|
||||
output_scales = torch.empty(
|
||||
MAX_TOKENS_PER_EXPERT * topk,
|
||||
padded_k,
|
||||
dtype=torch.int32,
|
||||
device=input_tensor.device,
|
||||
)
|
||||
torch.ops._C.silu_and_mul_mxfp4_experts_quant(
|
||||
output,
|
||||
output_scales,
|
||||
input_tensor,
|
||||
expert_offsets,
|
||||
blockscale_offsets,
|
||||
n_experts,
|
||||
)
|
||||
output_scales = output_scales.view(torch.uint8)
|
||||
return output, output_scales
|
||||
|
||||
|
||||
# fp8
|
||||
def scaled_fp8_quant(
|
||||
input: torch.Tensor,
|
||||
|
||||
@@ -22,6 +22,23 @@ else:
|
||||
except ImportError:
|
||||
from torch.library import impl_abstract as register_fake
|
||||
|
||||
if hasattr(torch.ops._xpu_C, "fp8_gemm"):
|
||||
|
||||
@register_fake("_xpu_C::fp8_gemm")
|
||||
def _fp8_gemm_fake(
|
||||
q_input: torch.Tensor,
|
||||
q_weight: torch.Tensor,
|
||||
out_dtype: torch.dtype,
|
||||
input_scales: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
input_2d = q_input.view(-1, q_input.shape[-1])
|
||||
M = input_2d.size(0)
|
||||
N = q_weight.size(1)
|
||||
return torch.empty((M, N), dtype=out_dtype, device=q_input.device)
|
||||
|
||||
|
||||
if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"):
|
||||
|
||||
@register_fake("_xpu_C::fp8_gemm_w8a16")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -47,9 +47,10 @@ class XpuCommunicator(DeviceCommunicatorBase):
|
||||
self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
|
||||
logger.info("Using AgRs manager on XPU device.")
|
||||
|
||||
def all_reduce(self, input_) -> torch.Tensor:
|
||||
dist.all_reduce(input_, group=self.device_group)
|
||||
return input_
|
||||
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
|
||||
output = input_.clone() if torch.compiler.is_compiling() else input_
|
||||
dist.all_reduce(output, group=self.device_group)
|
||||
return output
|
||||
|
||||
def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
|
||||
world_size = self.world_size
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from
|
||||
# https://github.com/vllm/vllm/entrypoints/openai/serving_chat.py
|
||||
# https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/chat_completion/serving.py
|
||||
|
||||
"""Anthropic Messages API serving handler"""
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -557,6 +557,20 @@ class OpenAIServingChat(OpenAIServing):
|
||||
and self._should_stream_with_auto_tool_parsing(request)
|
||||
)
|
||||
|
||||
# Determine whether required/named tool_choice should fall back to
|
||||
# the auto tool_parser path instead of the standard JSON-based parsing.
|
||||
# This happens when the parser declares supports_required_and_named=False
|
||||
# (e.g. GLM models that output XML instead of JSON).
|
||||
tool_choice_uses_parser = (
|
||||
self.tool_parser is not None
|
||||
and not self.tool_parser.supports_required_and_named
|
||||
and request.tools
|
||||
and (
|
||||
request.tool_choice == "required"
|
||||
or isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
|
||||
)
|
||||
)
|
||||
|
||||
all_previous_token_ids: list[list[int]] | None
|
||||
function_name_returned = [False] * num_choices
|
||||
if self.tool_call_id_type == "kimi_k2":
|
||||
@@ -569,7 +583,12 @@ class OpenAIServingChat(OpenAIServing):
|
||||
|
||||
# Only one of these will be used, thus previous_texts and
|
||||
# all_previous_token_ids will not be used twice in the same iteration.
|
||||
if is_mistral_grammar_path or tool_choice_auto or reasoning_parser:
|
||||
if (
|
||||
is_mistral_grammar_path
|
||||
or tool_choice_auto
|
||||
or tool_choice_uses_parser
|
||||
or reasoning_parser
|
||||
):
|
||||
# These are only required in "auto" tool choice case
|
||||
all_previous_token_ids = [[] for _ in range(num_choices)]
|
||||
reasoning_end_arr = [False] * num_choices
|
||||
@@ -764,7 +783,12 @@ class OpenAIServingChat(OpenAIServing):
|
||||
delta_message: DeltaMessage | None
|
||||
|
||||
# just update previous_texts and previous_token_ids
|
||||
if is_mistral_grammar_path or tool_choice_auto or reasoning_parser:
|
||||
if (
|
||||
is_mistral_grammar_path
|
||||
or tool_choice_auto
|
||||
or tool_choice_uses_parser
|
||||
or reasoning_parser
|
||||
):
|
||||
assert previous_texts is not None
|
||||
assert all_previous_token_ids is not None
|
||||
previous_text = previous_texts[i]
|
||||
@@ -813,7 +837,9 @@ class OpenAIServingChat(OpenAIServing):
|
||||
if result.tools_called:
|
||||
tools_streamed[i] = True
|
||||
# handle streaming deltas for tools with named tool_choice
|
||||
elif tool_choice_function_name:
|
||||
# Skip when tool_choice_uses_parser so it falls through
|
||||
# to the auto tool_parser branches below.
|
||||
elif tool_choice_function_name and not tool_choice_uses_parser:
|
||||
# When encountering think end id in prompt_token_ids
|
||||
# i.e {"enable_thinking": False},
|
||||
# check BEFORE calling the parser to avoid a spurious
|
||||
@@ -851,7 +877,6 @@ class OpenAIServingChat(OpenAIServing):
|
||||
):
|
||||
reasoning_end_arr[i] = True
|
||||
if delta_message and delta_message.content:
|
||||
# This need to be added to next `delta_text`
|
||||
current_text = delta_message.content
|
||||
delta_message.content = None
|
||||
else:
|
||||
@@ -896,7 +921,12 @@ class OpenAIServingChat(OpenAIServing):
|
||||
)
|
||||
tools_streamed[i] = True
|
||||
|
||||
elif request.tool_choice == "required":
|
||||
# Skip when tool_choice_uses_parser so it falls through
|
||||
# to the auto tool_parser branches below.
|
||||
elif (
|
||||
request.tool_choice == "required"
|
||||
and not tool_choice_uses_parser
|
||||
):
|
||||
assert previous_texts is not None
|
||||
previous_text = previous_texts[i]
|
||||
current_text = previous_text + delta_text
|
||||
@@ -966,7 +996,10 @@ class OpenAIServingChat(OpenAIServing):
|
||||
|
||||
# update the previous values for the next iteration
|
||||
if (
|
||||
is_mistral_grammar_path or tool_choice_auto or reasoning_parser
|
||||
is_mistral_grammar_path
|
||||
or tool_choice_auto
|
||||
or tool_choice_uses_parser
|
||||
or reasoning_parser
|
||||
) and not self.use_harmony:
|
||||
assert previous_texts is not None
|
||||
assert all_previous_token_ids is not None
|
||||
|
||||
@@ -627,7 +627,7 @@ class OpenAIServing:
|
||||
and isinstance(request.tool_choice, ToolChoiceFunction)
|
||||
):
|
||||
assert content is not None
|
||||
# Forced Function Call
|
||||
# Forced Function Call (Responses API)
|
||||
function_calls.append(
|
||||
FunctionCall(name=request.tool_choice.name, arguments=content)
|
||||
)
|
||||
@@ -636,14 +636,20 @@ class OpenAIServing:
|
||||
not use_mistral_tool_parser
|
||||
and request.tool_choice
|
||||
and isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
|
||||
and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named)
|
||||
):
|
||||
# Named function with standard JSON-based parsing
|
||||
assert content is not None
|
||||
# Forced Function Call
|
||||
function_calls.append(
|
||||
FunctionCall(name=request.tool_choice.function.name, arguments=content)
|
||||
)
|
||||
content = None # Clear content since tool is called.
|
||||
elif not use_mistral_tool_parser and request.tool_choice == "required":
|
||||
elif (
|
||||
not use_mistral_tool_parser
|
||||
and request.tool_choice == "required"
|
||||
and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named)
|
||||
):
|
||||
# "required" with standard JSON-based parsing
|
||||
tool_calls = []
|
||||
with contextlib.suppress(ValidationError):
|
||||
content = content or ""
|
||||
@@ -662,15 +668,30 @@ class OpenAIServing:
|
||||
use_mistral_tool_parser
|
||||
or (
|
||||
enable_auto_tools
|
||||
and (request.tool_choice == "auto" or request.tool_choice is None)
|
||||
and (
|
||||
request.tool_choice == "auto"
|
||||
or request.tool_choice is None
|
||||
or (
|
||||
not tool_parser_cls.supports_required_and_named
|
||||
and request.tools
|
||||
and (
|
||||
request.tool_choice == "required"
|
||||
or isinstance(
|
||||
request.tool_choice,
|
||||
ChatCompletionNamedToolChoiceParam,
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
):
|
||||
# Automatic Tool Call Parsing (also used as fallback for
|
||||
# required/named when supports_required_and_named=False)
|
||||
if tokenizer is None:
|
||||
raise ValueError(
|
||||
"Tokenizer not available when `skip_tokenizer_init=True`"
|
||||
)
|
||||
|
||||
# Automatic Tool Call Parsing
|
||||
try:
|
||||
tool_parser = tool_parser_cls(tokenizer, request.tools)
|
||||
except RuntimeError as e:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Encode/decode utilities for multimodal tensors and field metadata
|
||||
over JSON/HTTP, used by the disaggregated generate endpoint."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pybase64
|
||||
|
||||
from vllm.multimodal.inputs import MultiModalKwargsItem
|
||||
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
|
||||
|
||||
_encoder = MsgpackEncoder(size_threshold=2**62) # force all tensors inline
|
||||
_decoder = MsgpackDecoder(t=MultiModalKwargsItem)
|
||||
|
||||
|
||||
def encode_mm_kwargs_item(item: MultiModalKwargsItem) -> str:
|
||||
"""Serialize a MultiModalKwargsItem to a base64 string."""
|
||||
bufs = _encoder.encode(item)
|
||||
assert len(bufs) == 1, "All tensors should be inline"
|
||||
return pybase64.b64encode(bufs[0]).decode("ascii")
|
||||
|
||||
|
||||
def decode_mm_kwargs_item(data: str) -> MultiModalKwargsItem:
|
||||
"""Deserialize a base64 string back to a MultiModalKwargsItem."""
|
||||
raw = pybase64.b64decode(data)
|
||||
return _decoder.decode(raw)
|
||||
@@ -35,14 +35,6 @@ class MultiModalFeatures(BaseModel):
|
||||
Carries hashes (for cache lookup / identification) and placeholder
|
||||
positions so the downstream `/generate` service knows *where* in
|
||||
the token sequence each multimodal item lives.
|
||||
|
||||
Note:
|
||||
Phase 1 — metadata only.
|
||||
Phase 2 should add `mm_kwargs` (processed tensor data) using a
|
||||
binary transport so the ``/generate` side can skip re-processing.
|
||||
The `/generate` endpoint must also be updated to inject these
|
||||
features into `EngineInput` before passing to
|
||||
`InputProcessor.process_inputs`.
|
||||
"""
|
||||
|
||||
mm_hashes: dict[str, list[str]]
|
||||
@@ -51,6 +43,15 @@ class MultiModalFeatures(BaseModel):
|
||||
mm_placeholders: dict[str, list[PlaceholderRangeInfo]]
|
||||
"""Per-modality placeholder ranges in the token sequence."""
|
||||
|
||||
kwargs_data: dict[str, list[str | None]] | None = None
|
||||
"""Per-modality serialized tensor data.
|
||||
|
||||
Each value is a list parallel to ``mm_hashes[modality]``. A ``str``
|
||||
entry is a base64-encoded ``MultiModalKwargsItem``; ``None`` means
|
||||
the item should be resolved from cache. The entire field is
|
||||
``None`` for metadata-only (cache-hit) responses.
|
||||
"""
|
||||
|
||||
|
||||
class GenerateRequest(BaseModel):
|
||||
request_id: str = Field(
|
||||
|
||||
@@ -25,6 +25,7 @@ from vllm.entrypoints.openai.engine.protocol import (
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.serving import OpenAIServing, clamp_prompt_logprobs
|
||||
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
||||
from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item
|
||||
from vllm.entrypoints.serve.disagg.protocol import (
|
||||
GenerateRequest,
|
||||
GenerateResponse,
|
||||
@@ -34,8 +35,14 @@ from vllm.entrypoints.serve.disagg.protocol import (
|
||||
)
|
||||
from vllm.entrypoints.serve.render.serving import OpenAIServingRender
|
||||
from vllm.entrypoints.utils import should_include_usage
|
||||
from vllm.inputs import EngineInput, mm_input
|
||||
from vllm.logger import init_logger
|
||||
from vllm.logprobs import Logprob
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalKwargsItem,
|
||||
MultiModalKwargsItems,
|
||||
PlaceholderRange,
|
||||
)
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.utils.collection_utils import as_list
|
||||
@@ -103,11 +110,42 @@ class ServingTokens(OpenAIServing):
|
||||
if raw_request:
|
||||
raw_request.state.request_metadata = request_metadata
|
||||
|
||||
(engine_input,) = await self.openai_serving_render.preprocess_completion(
|
||||
request,
|
||||
prompt_input=request.token_ids,
|
||||
prompt_embeds=None,
|
||||
)
|
||||
engine_input: EngineInput
|
||||
if features := request.features:
|
||||
# Convert PlaceholderRangeInfo → PlaceholderRange per modality.
|
||||
mm_placeholders: dict[str, list[PlaceholderRange]] = {
|
||||
modality: [
|
||||
PlaceholderRange(offset=p.offset, length=p.length) for p in ranges
|
||||
]
|
||||
for modality, ranges in features.mm_placeholders.items()
|
||||
}
|
||||
|
||||
# Deserialize tensor data when present; None → cache hit.
|
||||
mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {}
|
||||
if features.kwargs_data is not None:
|
||||
for modality, items in features.kwargs_data.items():
|
||||
mm_kwargs[modality] = [
|
||||
decode_mm_kwargs_item(item) if item is not None else None
|
||||
for item in items
|
||||
]
|
||||
else:
|
||||
for modality, hashes in features.mm_hashes.items():
|
||||
mm_kwargs[modality] = [None] * len(hashes)
|
||||
|
||||
engine_input = mm_input(
|
||||
prompt_token_ids=request.token_ids,
|
||||
mm_kwargs=MultiModalKwargsItems(mm_kwargs),
|
||||
mm_hashes=features.mm_hashes,
|
||||
mm_placeholders=mm_placeholders,
|
||||
cache_salt=request.cache_salt,
|
||||
)
|
||||
else:
|
||||
(engine_input,) = await self.openai_serving_render.preprocess_completion(
|
||||
request,
|
||||
prompt_input=request.token_ids,
|
||||
prompt_embeds=None,
|
||||
skip_mm_cache=True,
|
||||
)
|
||||
|
||||
# Schedule the request and get the result generator.
|
||||
result_generator: AsyncGenerator[RequestOutput, None] | None = None
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from collections.abc import Sequence
|
||||
from http import HTTPStatus
|
||||
from typing import Any
|
||||
from typing import Any, cast
|
||||
|
||||
from openai_harmony import Message as OpenAIMessage
|
||||
|
||||
@@ -25,6 +25,7 @@ from vllm.entrypoints.openai.parser.harmony_utils import (
|
||||
render_for_completion,
|
||||
)
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
from vllm.entrypoints.serve.disagg.mm_serde import encode_mm_kwargs_item
|
||||
from vllm.entrypoints.serve.disagg.protocol import (
|
||||
GenerateRequest,
|
||||
MultiModalFeatures,
|
||||
@@ -37,6 +38,7 @@ from vllm.entrypoints.utils import (
|
||||
from vllm.inputs import (
|
||||
EngineInput,
|
||||
MultiModalHashes,
|
||||
MultiModalInput,
|
||||
MultiModalPlaceholders,
|
||||
PromptType,
|
||||
SingletonPrompt,
|
||||
@@ -251,6 +253,7 @@ class OpenAIServingRender:
|
||||
default_template_kwargs=self.default_chat_template_kwargs,
|
||||
tool_dicts=tool_dicts,
|
||||
tool_parser=tool_parser,
|
||||
skip_mm_cache=True,
|
||||
reasoning_parser=self.reasoning_parser,
|
||||
)
|
||||
else:
|
||||
@@ -342,6 +345,7 @@ class OpenAIServingRender:
|
||||
request,
|
||||
prompt_input=request.prompt,
|
||||
prompt_embeds=request.prompt_embeds,
|
||||
skip_mm_cache=True,
|
||||
)
|
||||
|
||||
return engine_inputs
|
||||
@@ -357,9 +361,10 @@ class OpenAIServingRender:
|
||||
if engine_input.get("type") != "multimodal":
|
||||
return None
|
||||
|
||||
# At this point engine_input is a MultiModalInputs TypedDict.
|
||||
mm_hashes: MultiModalHashes = engine_input["mm_hashes"] # type: ignore[typeddict-item]
|
||||
raw_placeholders: MultiModalPlaceholders = engine_input["mm_placeholders"] # type: ignore[typeddict-item]
|
||||
# At this point engine_input is a MultiModalInput TypedDict.
|
||||
mm_engine_input = cast(MultiModalInput, engine_input)
|
||||
mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
|
||||
raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
|
||||
|
||||
mm_placeholders = {
|
||||
modality: [
|
||||
@@ -368,9 +373,20 @@ class OpenAIServingRender:
|
||||
for modality, ranges in raw_placeholders.items()
|
||||
}
|
||||
|
||||
# Serialize tensor data per modality.
|
||||
kwargs_data: dict[str, list[str | None]] | None = None
|
||||
if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"):
|
||||
kwargs_data = {}
|
||||
for modality, items in raw_mm_kwargs.items():
|
||||
kwargs_data[modality] = [
|
||||
encode_mm_kwargs_item(item) if item is not None else None
|
||||
for item in items
|
||||
]
|
||||
|
||||
return MultiModalFeatures(
|
||||
mm_hashes=mm_hashes,
|
||||
mm_placeholders=mm_placeholders,
|
||||
kwargs_data=kwargs_data,
|
||||
)
|
||||
|
||||
def _make_request_with_harmony(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -53,14 +53,15 @@ if not has_helion():
|
||||
)
|
||||
|
||||
import helion
|
||||
from helion._compat import requires_torch_version
|
||||
from helion.autotuner.base_search import BaseAutotuner
|
||||
from helion.runtime.config import Config
|
||||
from helion.runtime.settings import default_autotuner_fn
|
||||
|
||||
# TODO(gmagogsfm): Remove CustomOp fallback path (_get_or_register_custom_op,
|
||||
# vllm_helion_lib, direct_register_custom_op) once vLLM requires PyTorch >= 2.11.
|
||||
_HOP_AVAILABLE = requires_torch_version("2.11")
|
||||
# FIXME(gmagogsfm): Re-enable HOP path once performance regression is fixed.
|
||||
# _HOP_AVAILABLE = requires_torch_version("2.11")
|
||||
_HOP_AVAILABLE = False
|
||||
|
||||
if _HOP_AVAILABLE:
|
||||
from helion._compat import supports_torch_compile_fusion
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -762,6 +762,25 @@ def nvfp4_moe_quant_config(
|
||||
)
|
||||
|
||||
|
||||
def mxfp4_moe_quant_config(
|
||||
w1_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
) -> FusedMoEQuantConfig:
|
||||
"""
|
||||
Construct a quant config for MXFP4 x MXFP4 MoE.
|
||||
MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no
|
||||
separate alphas / global activation scales in this config.
|
||||
"""
|
||||
return FusedMoEQuantConfig.make(
|
||||
"mxfp4",
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
per_act_token_quant=False,
|
||||
per_out_ch_quant=False,
|
||||
block_shape=None,
|
||||
)
|
||||
|
||||
|
||||
def nvfp4_w4a16_moe_quant_config(
|
||||
g1_alphas: torch.Tensor,
|
||||
g2_alphas: torch.Tensor,
|
||||
|
||||
@@ -36,6 +36,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
kMxfp4Dynamic,
|
||||
kMxfp4Static,
|
||||
kNvfp4Dynamic,
|
||||
kNvfp4Static,
|
||||
)
|
||||
@@ -795,6 +797,299 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
|
||||
)
|
||||
|
||||
|
||||
def run_cutlass_moe_mxfp4(
|
||||
output: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
w1_fp4: torch.Tensor,
|
||||
w1_blockscale: torch.Tensor,
|
||||
w2_fp4: torch.Tensor,
|
||||
w2_blockscale: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
activation: MoEActivation,
|
||||
workspace13: torch.Tensor,
|
||||
workspace2: torch.Tensor,
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
e: int,
|
||||
device: torch.device,
|
||||
apply_router_weight_on_input: bool = False,
|
||||
) -> None:
|
||||
"""MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM."""
|
||||
is_gated = activation.is_gated
|
||||
w1_n = n * 2 if is_gated else n
|
||||
|
||||
assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
|
||||
assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8"
|
||||
assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8"
|
||||
assert (
|
||||
w1_fp4.ndim == 3
|
||||
and w2_fp4.ndim == 3
|
||||
and w1_blockscale.ndim == 3
|
||||
and w2_blockscale.ndim == 3
|
||||
), "All Weights must be of rank 3 for cutlass_moe_mxfp4"
|
||||
m_a, k_a = a.shape
|
||||
e_w1, w1_n_actual, half_k_w1 = w1_fp4.shape
|
||||
e_w2, k_w2, half_n_w2 = w2_fp4.shape
|
||||
|
||||
assert e_w1 == e_w2 and e_w1 == e
|
||||
assert k_a == half_k_w1 * 2 and k == k_w2
|
||||
assert w1_n_actual == w1_n and half_n_w2 * 2 == n
|
||||
assert m == m_a
|
||||
assert 2 * half_k_w1 == k_w2
|
||||
assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype"
|
||||
assert topk_weights.size(0) == m and topk_ids.size(0) == m
|
||||
|
||||
topk = topk_ids.size(1)
|
||||
out_dtype = a.dtype
|
||||
num_topk = topk_ids.size(1)
|
||||
|
||||
expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
|
||||
blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
|
||||
problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device)
|
||||
problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device)
|
||||
|
||||
a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
|
||||
c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
|
||||
|
||||
if apply_router_weight_on_input:
|
||||
assert num_topk == 1, (
|
||||
"apply_router_weight_on_input is only implemented for topk=1"
|
||||
)
|
||||
a.mul_(topk_weights.to(out_dtype))
|
||||
|
||||
ops.get_cutlass_moe_mm_data(
|
||||
topk_ids,
|
||||
expert_offsets,
|
||||
problem_sizes1,
|
||||
problem_sizes2,
|
||||
a_map,
|
||||
c_map,
|
||||
e,
|
||||
n,
|
||||
k,
|
||||
blockscale_offsets,
|
||||
is_gated=is_gated,
|
||||
)
|
||||
|
||||
a = ops.shuffle_rows(a, a_map)
|
||||
rep_a_fp4, rep_a_blockscale = ops.mxfp4_experts_quant(
|
||||
a,
|
||||
expert_offsets,
|
||||
blockscale_offsets,
|
||||
e,
|
||||
num_topk,
|
||||
)
|
||||
c1 = _resize_cache(workspace13, (m * topk, w1_n))
|
||||
c2 = _resize_cache(workspace2, (m * topk, n))
|
||||
c3 = _resize_cache(workspace13, (m * topk, k))
|
||||
|
||||
ops.cutlass_mxfp4_moe_mm(
|
||||
c1,
|
||||
rep_a_fp4,
|
||||
w1_fp4,
|
||||
rep_a_blockscale,
|
||||
w1_blockscale,
|
||||
problem_sizes1,
|
||||
expert_offsets[:-1],
|
||||
blockscale_offsets[:-1],
|
||||
)
|
||||
del rep_a_fp4, rep_a_blockscale
|
||||
if activation == MoEActivation.SILU:
|
||||
int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant(
|
||||
c1, expert_offsets, blockscale_offsets, e, num_topk
|
||||
)
|
||||
else:
|
||||
apply_moe_activation(activation, c2, c1)
|
||||
int_fp4, int_blockscale = ops.mxfp4_experts_quant(
|
||||
c2, expert_offsets, blockscale_offsets, e, num_topk
|
||||
)
|
||||
|
||||
ops.cutlass_mxfp4_moe_mm(
|
||||
c3,
|
||||
int_fp4,
|
||||
w2_fp4,
|
||||
int_blockscale,
|
||||
w2_blockscale,
|
||||
problem_sizes2,
|
||||
expert_offsets[:-1],
|
||||
blockscale_offsets[:-1],
|
||||
)
|
||||
del int_fp4, int_blockscale
|
||||
|
||||
c3 = ops.shuffle_rows(c3, c_map)
|
||||
|
||||
assert output.dtype == out_dtype
|
||||
if not apply_router_weight_on_input:
|
||||
output.copy_(
|
||||
(
|
||||
c3.view(m, num_topk, k)
|
||||
* topk_weights.view(m, num_topk, 1).to(out_dtype)
|
||||
).sum(dim=1),
|
||||
non_blocking=True,
|
||||
)
|
||||
else:
|
||||
output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True)
|
||||
return
|
||||
|
||||
|
||||
def swizzle_mxfp4_scales(
|
||||
scales: torch.Tensor,
|
||||
N: int,
|
||||
K: int,
|
||||
) -> torch.Tensor:
|
||||
"""Swizzle flat [N, K//32] E8M0 scales to CUTLASS tiled layout.
|
||||
|
||||
CUTLASS expects MX scale factors in a tiled layout:
|
||||
[numMTiles, numKTiles, 32, 4, 4]
|
||||
where numMTiles = ceil(N/128), numKTiles = ceil(K/128),
|
||||
and the inner dimensions correspond to the swizzle pattern:
|
||||
mTileIdx = mIdx / 128
|
||||
outerMIdx = mIdx % 32
|
||||
innerMIdx = (mIdx / 32) % 4
|
||||
kTileIdx = kIdx / 4
|
||||
innerKIdx = kIdx % 4
|
||||
with kIdx = col_in_scale_space (i.e., index into K//32).
|
||||
"""
|
||||
assert scales.dtype == torch.uint8
|
||||
num_scale_cols = K // 32 # number of E8M0 scale values per row
|
||||
|
||||
num_m_tiles = (N + 127) // 128
|
||||
num_k_tiles = (num_scale_cols + 3) // 4
|
||||
|
||||
# Pad N to multiple of 128 and scale_cols to multiple of 4
|
||||
padded_N = num_m_tiles * 128
|
||||
padded_scale_cols = num_k_tiles * 4
|
||||
|
||||
# Start with flat scales, pad if needed
|
||||
padded = torch.zeros(
|
||||
padded_N, padded_scale_cols, dtype=torch.uint8, device=scales.device
|
||||
)
|
||||
padded[:N, :num_scale_cols] = scales
|
||||
|
||||
# Reshape to tile structure:
|
||||
# [numMTiles, 4, 32, numKTiles, 4]
|
||||
# mTileIdx, innerMIdx, outerMIdx, kTileIdx, innerKIdx
|
||||
tiled = padded.reshape(num_m_tiles, 4, 32, num_k_tiles, 4)
|
||||
# Permute to [numMTiles, numKTiles, 32, 4, 4]
|
||||
# (outerMIdx, innerMIdx, innerKIdx)
|
||||
tiled = tiled.permute(0, 3, 2, 1, 4).contiguous()
|
||||
return tiled.reshape(-1)
|
||||
|
||||
|
||||
class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
|
||||
"""CUTLASS MXFP4 x MXFP4 fused MoE expert implementation."""
|
||||
|
||||
@property
|
||||
def expects_unquantized_inputs(self) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _supports_current_device() -> bool:
|
||||
p = current_platform
|
||||
return p.is_cuda() and p.is_device_capability_family(100)
|
||||
|
||||
@staticmethod
|
||||
def _supports_no_act_and_mul() -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _supports_quant_scheme(
|
||||
weight_key: QuantKey | None,
|
||||
activation_key: QuantKey | None,
|
||||
) -> bool:
|
||||
return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic)
|
||||
|
||||
@staticmethod
|
||||
def _supports_activation(activation: MoEActivation) -> bool:
|
||||
return activation in [
|
||||
MoEActivation.SILU,
|
||||
MoEActivation.GELU,
|
||||
MoEActivation.SWIGLUOAI,
|
||||
MoEActivation.SWIGLUSTEP,
|
||||
MoEActivation.SILU_NO_MUL,
|
||||
MoEActivation.GELU_NO_MUL,
|
||||
MoEActivation.RELU2_NO_MUL,
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _supports_parallel_config(
|
||||
moe_parallel_config: FusedMoEParallelConfig,
|
||||
) -> bool:
|
||||
return moe_parallel_config.ep_size == 1
|
||||
|
||||
@staticmethod
|
||||
def activation_format() -> mk.FusedMoEActivationFormat:
|
||||
return mk.FusedMoEActivationFormat.Standard
|
||||
|
||||
def supports_expert_map(self) -> bool:
|
||||
return False
|
||||
|
||||
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
|
||||
return TopKWeightAndReduceNoOP()
|
||||
|
||||
def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
|
||||
return act_dtype
|
||||
|
||||
def workspace_shapes(
|
||||
self,
|
||||
M: int,
|
||||
N: int,
|
||||
K: int,
|
||||
topk: int,
|
||||
global_num_experts: int,
|
||||
local_num_experts: int,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
activation: MoEActivation,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
|
||||
workspace1 = (M * topk, max(2 * N, K))
|
||||
workspace2 = (M * topk, N)
|
||||
output = (M, K)
|
||||
return (workspace1, workspace2, output)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
activation: MoEActivation,
|
||||
global_num_experts: int,
|
||||
expert_map: torch.Tensor | None,
|
||||
a1q_scale: torch.Tensor | None,
|
||||
a2_scale: torch.Tensor | None,
|
||||
workspace13: torch.Tensor | None,
|
||||
workspace2: torch.Tensor | None,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
apply_router_weight_on_input: bool,
|
||||
):
|
||||
e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
|
||||
n = w2.shape[2] * 2
|
||||
|
||||
run_cutlass_moe_mxfp4(
|
||||
output=output,
|
||||
a=hidden_states,
|
||||
w1_fp4=w1,
|
||||
w1_blockscale=self.w1_scale,
|
||||
w2_fp4=w2,
|
||||
w2_blockscale=self.w2_scale,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=activation,
|
||||
workspace13=workspace13,
|
||||
workspace2=workspace2,
|
||||
m=m,
|
||||
n=n,
|
||||
k=k,
|
||||
e=e,
|
||||
device=hidden_states.device,
|
||||
apply_router_weight_on_input=apply_router_weight_on_input,
|
||||
)
|
||||
|
||||
|
||||
# W4A8
|
||||
def run_cutlass_moe_w4a8_fp8(
|
||||
output: torch.Tensor,
|
||||
|
||||
@@ -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.
|
||||
|
||||
+65
-13
@@ -4,6 +4,7 @@
|
||||
|
||||
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 import (
|
||||
FusedMoE,
|
||||
@@ -11,6 +12,10 @@ from vllm.model_executor.layers.fused_moe import (
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
mxfp4_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
|
||||
CutlassExpertsMxfp4,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
|
||||
MarlinExperts,
|
||||
@@ -36,7 +41,14 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
|
||||
super().__init__(moe)
|
||||
self.group_size = 32
|
||||
self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
|
||||
self.experts_cls = MarlinExperts
|
||||
self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
|
||||
self.experts_cls: type[mk.FusedMoEExperts]
|
||||
if self.use_cutlass_mxfp4:
|
||||
logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE", scope="local")
|
||||
self.experts_cls = CutlassExpertsMxfp4
|
||||
else:
|
||||
logger.info_once("Using MarlinExperts for MXFP4 MoE", scope="local")
|
||||
self.experts_cls = MarlinExperts
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
@@ -109,11 +121,19 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
|
||||
def get_fused_moe_quant_config(
|
||||
self, layer: torch.nn.Module
|
||||
) -> FusedMoEQuantConfig | None:
|
||||
return make_mxfp4_moe_quant_config(
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
)
|
||||
if self.use_cutlass_mxfp4:
|
||||
# W4A4: both weights and activations quantized to MXFP4
|
||||
return mxfp4_moe_quant_config(
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
)
|
||||
else:
|
||||
# W4A16: weight-only via Marlin
|
||||
return make_mxfp4_moe_quant_config(
|
||||
mxfp4_backend=self.mxfp4_backend,
|
||||
w1_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: FusedMoE) -> None:
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
@@ -126,13 +146,45 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
|
||||
)
|
||||
delattr(layer, "w2_weight_packed")
|
||||
|
||||
logger.warning_once(
|
||||
"Your GPU does not have native support for FP4 computation but "
|
||||
"FP4 quantization is being used. Weight-only FP4 compression "
|
||||
"will be used leveraging the Marlin kernel. This may degrade "
|
||||
"performance for compute-heavy workloads."
|
||||
)
|
||||
prepare_moe_fp4_layer_for_marlin(layer)
|
||||
if self.use_cutlass_mxfp4:
|
||||
# Swizzle weight scales from flat checkpoint layout [E, N, K//32]
|
||||
# to CUTLASS tiled layout [E, numMTiles*numKTiles*512].
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
|
||||
swizzle_mxfp4_scales,
|
||||
)
|
||||
|
||||
E = layer.w13_weight_scale.shape[0]
|
||||
w13_N = layer.w13_weight_scale.shape[1]
|
||||
w13_scale_K = layer.w13_weight_scale.shape[2]
|
||||
w13_K = w13_scale_K * 32
|
||||
|
||||
w2_M = layer.w2_weight_scale.shape[1]
|
||||
w2_scale_N = layer.w2_weight_scale.shape[2]
|
||||
w2_N = w2_scale_N * 32
|
||||
|
||||
swizzled_w13 = []
|
||||
swizzled_w2 = []
|
||||
for e_idx in range(E):
|
||||
s13 = layer.w13_weight_scale[e_idx]
|
||||
sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K)
|
||||
swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K))
|
||||
s2 = layer.w2_weight_scale[e_idx]
|
||||
sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N)
|
||||
swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N))
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
torch.stack(swizzled_w13), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
torch.stack(swizzled_w2), requires_grad=False
|
||||
)
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Your GPU does not have native support for FP4 computation "
|
||||
"but FP4 quantization is being used. Weight-only FP4 "
|
||||
"compression will be used leveraging the Marlin kernel. "
|
||||
"This may degrade performance for compute-heavy workloads."
|
||||
)
|
||||
prepare_moe_fp4_layer_for_marlin(layer)
|
||||
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
if self.moe_quant_config is not None:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -67,6 +67,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
|
||||
from .interfaces import (
|
||||
MultiModalEmbeddings,
|
||||
SupportsEagle3,
|
||||
SupportsLoRA,
|
||||
SupportsMultiModal,
|
||||
SupportsPP,
|
||||
)
|
||||
@@ -880,6 +881,7 @@ class Gemma4ForConditionalGeneration(
|
||||
nn.Module,
|
||||
SupportsMultiModal,
|
||||
SupportsPP,
|
||||
SupportsLoRA,
|
||||
SupportsEagle3,
|
||||
):
|
||||
packed_modules_mapping = {
|
||||
@@ -1254,9 +1256,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
|
||||
@@ -1357,10 +1360,16 @@ class Gemma4ForConditionalGeneration(
|
||||
|
||||
def get_mm_mapping(self) -> MultiModelKeys:
|
||||
"""Get the module prefix mapping for multimodal models."""
|
||||
connectors = ["embed_vision"]
|
||||
tower_models = ["vision_tower"]
|
||||
if self.audio_tower is not None:
|
||||
connectors.append("embed_audio")
|
||||
tower_models.append("audio_tower")
|
||||
|
||||
return MultiModelKeys.from_string_field(
|
||||
language_model="language_model",
|
||||
connector=["embed_vision", "embed_audio"],
|
||||
tower_model=["vision_tower", "audio_tower"],
|
||||
connector=connectors,
|
||||
tower_model=tower_models,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -66,7 +66,7 @@ from .interfaces import (
|
||||
SupportsTranscription,
|
||||
)
|
||||
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
|
||||
from .whisper import ISO639_1_SUPPORTED_LANGS
|
||||
from .whisper import ISO639_1_SUPPORTED_LANGS, _create_fake_bias_for_k_proj
|
||||
|
||||
|
||||
class GlmAsrEncoderRotaryEmbedding(nn.Module):
|
||||
@@ -499,6 +499,8 @@ class GlmAsrEncoder(nn.Module):
|
||||
"""Custom weight loading to handle q_proj/k_proj/v_proj -> qkv_proj mapping."""
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
|
||||
|
||||
stacked_params_mapping = [
|
||||
# (param_name, shard_name, shard_id)
|
||||
("qkv_proj", "q_proj", "q"),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -145,14 +145,15 @@ class PlaceholderRange:
|
||||
"""
|
||||
|
||||
@cached_property
|
||||
def embeds_cumsum(self) -> torch.Tensor | None:
|
||||
return None if self.is_embed is None else self.is_embed.cumsum(dim=0)
|
||||
def embeds_cumsum(self) -> list[int] | None:
|
||||
# python list so python indexing avoids torch C++ overhead/conversions/deallocs
|
||||
return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist()
|
||||
|
||||
def get_num_embeds(self) -> int:
|
||||
if self.embeds_cumsum is None:
|
||||
return self.length
|
||||
|
||||
return int(self.embeds_cumsum[-1])
|
||||
return self.embeds_cumsum[-1] if self.embeds_cumsum else 0
|
||||
|
||||
def get_embeds_indices_in_range(
|
||||
self, start_idx: int, end_idx: int
|
||||
@@ -170,10 +171,8 @@ class PlaceholderRange:
|
||||
if self.embeds_cumsum is None:
|
||||
return start_idx, end_idx
|
||||
|
||||
embeds_start_idx = (
|
||||
int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0
|
||||
)
|
||||
embeds_end_idx = int(self.embeds_cumsum[end_idx - 1])
|
||||
embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0
|
||||
embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0
|
||||
|
||||
return embeds_start_idx, embeds_end_idx
|
||||
|
||||
|
||||
+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]]:
|
||||
"""
|
||||
|
||||
@@ -131,6 +131,7 @@ def _get_backend_priorities(
|
||||
AttentionBackendEnum.FLASH_ATTN,
|
||||
AttentionBackendEnum.TRITON_ATTN,
|
||||
AttentionBackendEnum.FLEX_ATTENTION,
|
||||
AttentionBackendEnum.TURBOQUANT,
|
||||
]
|
||||
else:
|
||||
return [
|
||||
@@ -138,6 +139,7 @@ def _get_backend_priorities(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
AttentionBackendEnum.TRITON_ATTN,
|
||||
AttentionBackendEnum.FLEX_ATTENTION,
|
||||
AttentionBackendEnum.TURBOQUANT,
|
||||
]
|
||||
|
||||
|
||||
@@ -255,11 +257,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,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user