iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Mini-VLLM] Remove the ROCm device path, debug routes and arch-gated tests

Browse Source 
- Delete csrc/rocm, hipify, the _rocm_C extension and the setup.py rocm
  target path; drop the rocm requirement files
- Delete the aiter, rdna3, skinny-gemm and ROCm fused-MoE kernel modules
  and retarget the linear kernel registries away from them
- Reduce platforms/rocm.py to the GCN architecture predicates that live
  code still imports, and delete the RocmPlatform class and plugin
- Drop the VLLM_SERVER_DEV_MODE debug surface: reset, collective_rpc,
  server_info and sleep routes plus the log_error_stack flag
- Delete lm-format-enforcer, lark and otel-semantic-conventions-ai pins;
  the live server process maps none of them
- Delete tests that only build architectures this fork no longer
  registers, and test directories left without any test file
- Restore the tests/v1/kv_connector mock helpers that scheduler tests
  import at module level

Co-authored-by: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-19 19:31:07 +07:00
co-authored by Hinata
parent 2066d3143f
commit 5f3c7bd8ab
162 changed files with 599 additions and 27490 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+1 -101
View File
@@ -166,8 +166,7 @@ define_extension_target(
#
# Forward the non-CUDA device extensions to external CMake scripts.
#
if (NOT VLLM_TARGET_DEVICE STREQUAL "cuda" AND
NOT VLLM_TARGET_DEVICE STREQUAL "rocm")
if (NOT VLLM_TARGET_DEVICE STREQUAL "cuda")
return()
endif()
@@ -352,27 +351,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
WITH_SOABI)
endif()
#
# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
#
if(VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_EXT_SRC
"csrc/rocm/quickreduce/torch_bindings.cpp"
"csrc/rocm/quickreduce/custom_quickreduce.cu")
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
USE_SABI 3
WITH_SOABI)
endif() # _C HIP endif
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
#
# _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
@@ -1416,84 +1394,6 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
endif()
endif()
if(VLLM_GPU_LANG STREQUAL "HIP")
#
# _rocm_C extension
#
set(VLLM_ROCM_EXT_SRC
"csrc/rocm/torch_bindings.cpp"
"csrc/rocm/skinny_gemms.cu"
"csrc/rocm/skinny_gemms_int4.cu"
"csrc/rocm/attention.cu")
set(VLLM_ROCM_EXT_FLAGS ${VLLM_GPU_FLAGS})
# skinny_gemms*.cu are built on gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy
# s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude them from the
# gfx1250 build and disable their op registrations (VLLM_SKIP_SKINNY_GEMMS);
# vLLM falls back to default/Triton GEMM for those ops on gfx1250.
list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu")
list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms_int4.cu")
set(VLLM_SKINNY_ARCHES ${VLLM_GPU_ARCHES})
list(FILTER VLLM_SKINNY_ARCHES EXCLUDE REGEX "gfx1250")
if(VLLM_SKINNY_ARCHES)
message(STATUS "Building skinny_gemms for archs: ${VLLM_SKINNY_ARCHES}")
set(VLLM_SKINNY_SRC
"csrc/rocm/skinny_gemms.cu"
"csrc/rocm/skinny_gemms_int4.cu")
hipify_sources_target(VLLM_SKINNY_HIP_SRCS _rocm_C_skinny "${VLLM_SKINNY_SRC}")
unset(_VLLM_LAST_HIPIFY_TARGET)
add_library(_rocm_C_skinny OBJECT ${VLLM_SKINNY_HIP_SRCS})
add_dependencies(_rocm_C_skinny hipify_all)
set_source_files_properties(${VLLM_SKINNY_HIP_SRCS} PROPERTIES LANGUAGE ${VLLM_GPU_LANG})
set_target_properties(_rocm_C_skinny PROPERTIES
${VLLM_GPU_LANG}_ARCHITECTURES "${VLLM_SKINNY_ARCHES}"
POSITION_INDEPENDENT_CODE ON)
target_include_directories(_rocm_C_skinny PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc)
target_compile_options(_rocm_C_skinny PRIVATE
$<$<COMPILE_LANGUAGE:${VLLM_GPU_LANG}>:${VLLM_ROCM_EXT_FLAGS}>)
target_compile_definitions(_rocm_C_skinny PRIVATE "-DTORCH_EXTENSION_NAME=_rocm_C")
target_link_libraries(_rocm_C_skinny PRIVATE torch)
else()
message(STATUS "Only gfx1250, skipping skinny_gemms")
list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS")
endif()
set(VLLM_ROCM_HAS_GFX1100 OFF)
if(VLLM_GPU_ARCHES MATCHES "gfx1100")
set(VLLM_ROCM_HAS_GFX1100 ON)
list(APPEND VLLM_ROCM_EXT_SRC
"csrc/rocm/q_gemm_rdna3.cu"
"csrc/rocm/q_gemm_rdna3_wmma.cu"
"csrc/rocm/moe_q_gemm_rdna3.cu")
endif()
define_extension_target(
_rocm_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_ROCM_EXT_SRC}
COMPILE_FLAGS ${VLLM_ROCM_EXT_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
USE_SABI 3
WITH_SOABI)
if(TARGET _rocm_C_skinny)
target_link_libraries(_rocm_C PRIVATE _rocm_C_skinny)
else()
target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS)
endif()
if(VLLM_ROCM_HAS_GFX1100)
target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
endif()
endif()
# Must run after the last HIP `define_extension_target` so every extension
# has registered its sources.
if (VLLM_GPU_LANG STREQUAL "HIP")
vllm_finalize_hipify_target()
endif()
# For CUDA and HIP builds also build the triton_kernels external package.
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
include(cmake/external_projects/triton_kernels.cmake)
cmake/hipify.py
-101
View File
@@ -1,101 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# A command line tool for running pytorch's hipify preprocessor on CUDA
# source files.
#
# See https://github.com/ROCm/hipify_torch
# and <torch install dir>/utils/hipify/hipify_python.py
#
import argparse
import os
import shutil
from torch.utils.hipify.hipify_python import get_hip_file_path, hipify
def _expected_hip_build_path(source_abs: str, output_directory: str) -> str:
"""Match torch.utils.hipify.hipify_python.preprocessor fout_path naming."""
rel = os.path.relpath(source_abs, output_directory)
return os.path.abspath(
os.path.join(
output_directory, get_hip_file_path(rel, is_pytorch_extension=True)
)
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Project directory where all the source + include files live.
parser.add_argument(
"-p",
"--project_dir",
help="The project directory.",
)
# Directory where hipified files are written.
parser.add_argument(
"-o",
"--output_dir",
help="The output directory.",
)
# Source files to convert.
parser.add_argument(
"sources", help="Source files to hipify.", nargs="*", default=[]
)
args = parser.parse_args()
# Limit include scope to project_dir only
includes = [os.path.join(args.project_dir, "*")]
# Get absolute path for all source files.
extra_files = [os.path.abspath(s) for s in args.sources]
# Copy sources from project directory to output directory.
# The directory might already exist to hold object files so we ignore that.
shutil.copytree(args.project_dir, args.output_dir, dirs_exist_ok=True)
hipify_result = hipify(
project_directory=args.project_dir,
output_directory=args.output_dir,
# Hipify resolves quoted includes next to the including file first; vLLM
# uses paths relative to csrc/ (e.g. "torch_utils.h"
# from quantization/w8a8/fp8/*.cu). Without an include root here, those
# headers are never found and are not hipified or rewritten in dependents.
header_include_dirs=["."],
includes=includes,
extra_files=extra_files,
show_detailed=True,
is_pytorch_extension=True,
hipify_extra_files_only=True,
)
hipified_sources = []
for source in args.sources:
s_abs = os.path.abspath(source)
if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None:
path = hipify_result[s_abs].hipified_path
# PyTorch skips writing when is_pytorch_extension and text unchanged;
# hipified_path then stays *.cu. CMake expects *.hip under output_dir.
if s_abs.endswith(".cu") and path.endswith(".cu"):
dest = _expected_hip_build_path(s_abs, args.output_dir)
if os.path.normpath(path) != os.path.normpath(dest):
os.makedirs(os.path.dirname(dest), exist_ok=True)
shutil.copy2(path, dest)
hipified_s_abs = dest
else:
hipified_s_abs = path
else:
hipified_s_abs = s_abs
hipified_sources.append(hipified_s_abs)
assert len(hipified_sources) == len(args.sources)
# Print hipified source files.
print("\n".join(hipified_sources))
cmake/utils.cmake
+2 -89
View File
@@ -47,80 +47,6 @@ macro (append_cmake_prefix_path PKG EXPR)
list(APPEND CMAKE_PREFIX_PATH ${_PREFIX_PATH})
endmacro()
# Resolve hipified output paths for `NAME` into `OUT_SRCS` and register the
# `.cu` sources with the shared `hipify_all` target. Per-extension hipify
# targets are unsafe to run in parallel against a shared csrc/ output dir, so
# accumulation here is paired with a single finalize step.
function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
if (TARGET hipify_all)
message(FATAL_ERROR
"hipify_sources_target(${NAME}) called after vllm_finalize_hipify_target. "
"Add the new HIP extension before the finalizer call in CMakeLists.txt.")
endif()
#
# Split into C++ and non-C++ (i.e. CUDA) sources.
#
set(SRCS ${ORIG_SRCS})
set(CXX_SRCS ${ORIG_SRCS})
list(FILTER SRCS EXCLUDE REGEX "\.(cc)|(cpp)|(hip)$")
list(FILTER CXX_SRCS INCLUDE REGEX "\.(cc)|(cpp)|(hip)$")
#
# Generate ROCm/HIP source file names from CUDA file names.
# Since HIP files are generated code, they will appear in the build area
# `CMAKE_CURRENT_BINARY_DIR` directory rather than the original csrc dir.
#
set(HIP_SRCS)
foreach (SRC ${SRCS})
string(REGEX REPLACE "\.cu$" "\.hip" SRC ${SRC})
string(REGEX REPLACE "cuda" "hip" SRC ${SRC})
list(APPEND HIP_SRCS "${CMAKE_CURRENT_BINARY_DIR}/${SRC}")
endforeach()
set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS})
set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS})
# Chain hipify targets so they run sequentially. Parallel hipify
# invocations race on shutil.copytree, overwriting .hip files
# produced by another target back to .cu originals.
if (DEFINED _VLLM_LAST_HIPIFY_TARGET)
add_dependencies(hipify${NAME} ${_VLLM_LAST_HIPIFY_TARGET})
endif()
set(_VLLM_LAST_HIPIFY_TARGET "hipify${NAME}" PARENT_SCOPE)
# Swap out original extension sources with hipified sources.
list(APPEND HIP_SRCS ${CXX_SRCS})
set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE)
endfunction()
# Define the single shared `hipify_all` custom target that runs hipify once
# on the union of every HIP extension's sources. Call after the last HIP
# `define_extension_target`.
function (vllm_finalize_hipify_target)
if (TARGET hipify_all)
return()
endif()
get_property(ALL_SRCS GLOBAL PROPERTY VLLM_HIPIFY_ALL_SRCS)
get_property(ALL_BYPRODUCTS GLOBAL PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS)
if (NOT ALL_SRCS)
return()
endif()
list(REMOVE_DUPLICATES ALL_SRCS)
list(REMOVE_DUPLICATES ALL_BYPRODUCTS)
set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc)
add_custom_target(
hipify_all
COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${ALL_SRCS}
DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${ALL_SRCS}
BYPRODUCTS ${ALL_BYPRODUCTS}
COMMENT "Running hipify on all extension source files.")
endfunction()
#
# Get additional GPU compiler flags from torch.
#
@@ -602,11 +528,6 @@ function (define_extension_target MOD_NAME)
"DESTINATION;LANGUAGE;USE_SABI"
"SOURCES;ARCHITECTURES;COMPILE_FLAGS;INCLUDE_DIRECTORIES;LIBRARIES")
# Add hipify preprocessing step when building with HIP/ROCm.
if (ARG_LANGUAGE STREQUAL "HIP")
hipify_sources_target(ARG_SOURCES ${MOD_NAME} "${ARG_SOURCES}")
endif()
if (ARG_WITH_SOABI)
set(SOABI_KEYWORD WITH_SOABI)
else()
@@ -625,16 +546,8 @@ function (define_extension_target MOD_NAME)
Python_add_library(${MOD_NAME} MODULE ${SOABI_KEYWORD} "${ARG_SOURCES}")
endif()
if (ARG_LANGUAGE STREQUAL "HIP")
# Make this target dependent on the hipify preprocessor step.
add_dependencies(${MOD_NAME} hipify_all)
# Make sure we include the hipified versions of the headers, and avoid conflicts with the ones in the original source folder
target_include_directories(${MOD_NAME} PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc
${ARG_INCLUDE_DIRECTORIES})
else()
target_include_directories(${MOD_NAME} PRIVATE csrc
${ARG_INCLUDE_DIRECTORIES})
endif()
target_include_directories(${MOD_NAME} PRIVATE csrc
${ARG_INCLUDE_DIRECTORIES})
if (ARG_ARCHITECTURES)
set_target_properties(${MOD_NAME} PROPERTIES
csrc/quantization/w8a8/fp8/common.cuh
+2 -4
View File
@@ -5,11 +5,9 @@
#include <cmath>
// This header is shared between _rocm_C and _C_stable_libtorch targets.
// This header is used by the _C_stable_libtorch target.
// torch_utils.h provides get_device_prop(). We need to pass USE_CUDA
// to the .so to expose some of the shims used by torch_utils.h. For now
// this is only done for _C_stable_libtorch and not for _rocm_C, so we use
// the non stable at::cuda::getCurrentDeviceProperties for _rocm_C for now.
// to the .so to expose the shims used by torch_utils.h.
#ifdef TORCH_TARGET_VERSION
#include "torch_utils.h"
#else
csrc/rocm/attention.cu
-3936
View File
File diff suppressed because it is too large Load Diff
csrc/rocm/dispatch_utils.h
-24
View File
@@ -1,24 +0,0 @@
/*
* Adapted from
* https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h
*/
#pragma once
#include <torch/all.h>
// Need a special dispatch case macro since we will nest the FP8 dispatch.
// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'.
#define AT_DISPATCH_FP8_CASE(enum_type, ...) \
AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
// ROCm devices might use either fn or fnuz, so set up dispatch table for both.
// A host-based check at runtime will create a preferred FP8 type for ROCm
// such that the correct kernel is dispatched.
#define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__)
// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'.
// See AT_DISPATCH_FP8_CASE above.
#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
csrc/rocm/moe_q_gemm_rdna3.cu
-646
View File
@@ -1,646 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Fused MoE W4A16 GPTQ kernel for RDNA3 (gfx1100).
//
// Combines expert routing (sorted_token_ids / expert_ids) with the RDNA3
// W4A16 dequant+dot from q_gemm_rdna3.cu into a single kernel launch.
// Each block processes BLOCK_SIZE_M tokens assigned to one expert, covering
// a tile of N output columns and K input positions.
//
// Weight format: same as the dense kernel — [E, K/8, N] uint32 shuffled,
// [E, groups, N] scales, [E, groups, N/8] packed zeros.
//
// Design: THREADS_X=256 (8 waves on wave32), BLOCK_KN_SIZE=256, each thread
// handles 4 N columns. Output via 64-bit packed CAS atomic-add directly to
// the pre-zeroed output tensor (no FP32 scratch buffer).
#include <cstdint>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include <ATen/cuda/CUDAContext.h>
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
#include "qdq_4_rdna3.cuh"
#if defined(__HIPCC__) && defined(__gfx1100__)
#define __HIP__RDNA3__
#endif
namespace vllm {
namespace moe_gptq_rdna3 {
#define BLOCK_KN_SIZE 256
#define THREADS_X 256
#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
using gptq_rdna3::bf162_t;
using gptq_rdna3::bf16_t;
// --- Helpers (same as q_gemm_rdna3.cu) ---
template <typename T>
__forceinline__ __device__ T tzero();
template <>
__forceinline__ __device__ half tzero<half>() {
return __float2half_rn(0.0f);
}
template <>
__forceinline__ __device__ bf16_t tzero<bf16_t>() {
return __float2bfloat16(0.0f);
}
__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
half2 v23) {
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
unsigned long long old = *addr_u;
while (true) {
union {
unsigned long long u;
half2 h2[2];
} cur, sum;
cur.u = old;
sum.h2[0] = __hadd2(cur.h2[0], v01);
sum.h2[1] = __hadd2(cur.h2[1], v23);
unsigned long long prev = atomicCAS(addr_u, old, sum.u);
if (prev == old) break;
old = prev;
}
}
__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
bf162_t v23) {
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
unsigned long long old = *addr_u;
while (true) {
union {
unsigned long long u;
bf162_t b2[2];
} cur, sum;
cur.u = old;
sum.b2[0] = __hadd2(cur.b2[0], v01);
sum.b2[1] = __hadd2(cur.b2[1], v23);
unsigned long long prev = atomicCAS(addr_u, old, sum.u);
if (prev == old) break;
old = prev;
}
}
__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
int (&zeros)[4]) {
int qcol = n / 8;
int shift = (n & 0x07) * 4;
uint32_t d = qzeros_row[qcol] >> shift;
zeros[0] = (int)(d & 0xF);
zeros[1] = (int)((d >> 4) & 0xF);
zeros[2] = (int)((d >> 8) & 0xF);
zeros[3] = (int)((d >> 12) & 0xF);
}
template <typename T>
__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
T (&scales)[4]) {
scales[0] = scales_row[n + 0];
scales[1] = scales_row[n + 1];
scales[2] = scales_row[n + 2];
scales[3] = scales_row[n + 3];
}
// ---------------------------------------------------------------------------
// Fused MoE kernel.
// ---------------------------------------------------------------------------
template <typename T, int BLOCK_SIZE_M>
__global__ void moe_gemm_q4_kernel_rdna3(
const T* __restrict__ a, // [size_m, size_k] or [M*topk, K]
T* __restrict__ c, // [M*topk, size_n] pre-zeroed
const uint32_t* __restrict__ b_q_weight, // [E, K/8, N] packed
const T* __restrict__ b_scales, // [E, groups, N]
const uint32_t* __restrict__ b_qzeros, // [E, groups, N/8] packed
const float* __restrict__ topk_weights, // [M*topk] or nullptr
const int32_t* __restrict__ sorted_token_ids,
const int32_t* __restrict__ expert_ids,
const int32_t* __restrict__ num_tokens_post_padded,
const int size_m, // total tokens (original M, or M*topk for w2)
const int size_n, // output features per expert
const int size_k, // input features
const int groups, // K / group_size
const int top_k, // routing top-k (1 for w2 pass)
// Per-expert strides (in elements, not bytes)
const int expert_weight_stride, // (K/8) * N
const int expert_scales_stride, // groups * N
const int expert_zeros_stride, // groups * (N/8)
const bool mul_topk_weight,
const int output_topk) { // >0: reduce output by token_id/output_topk
const int t = threadIdx.x;
const int token_block = blockIdx.x;
const int offset_n = blockIdx.y * BLOCK_KN_SIZE * 4;
const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
const int n = offset_n + t * 4;
// Early exit for padding blocks or invalid experts (expert_map = -1)
if (token_block * BLOCK_SIZE_M >= num_tokens_post_padded[0]) return;
const int expert_id = expert_ids[token_block];
if (expert_id == -1) return;
// Expert-specific pointers
const uint32_t* expert_weights =
b_q_weight + (int64_t)expert_id * expert_weight_stride;
const T* expert_scales = b_scales + (int64_t)expert_id * expert_scales_stride;
const uint32_t* expert_qzeros =
b_qzeros + (int64_t)expert_id * expert_zeros_stride;
// LDS for activations
constexpr int LDS_PAD = 8;
__shared__ T block_a[BLOCK_SIZE_M][BLOCK_KN_SIZE + LDS_PAD];
static_assert(BLOCK_KN_SIZE == THREADS_X,
"BLOCK_KN_SIZE must equal THREADS_X");
// For bf16 M=1, we can skip LDS and read A from global (same as dense).
// fp16 always needs LDS: the exact-factored inner loop indexes
// block_a[m][a_off] directly.
constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same<T, half>::value;
const int offset_m_base = token_block * BLOCK_SIZE_M;
if constexpr (USE_LDS_A) {
if (offset_k + t < end_k) {
#pragma unroll
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
int32_t token_id = sorted_token_ids[offset_m_base + m];
int token_row = token_id / top_k;
T av;
if (token_row < size_m) {
av = a[(int64_t)token_row * size_k + offset_k + t];
} else {
av = tzero<T>();
}
block_a[m][t] = av;
}
}
__syncthreads();
}
if (n >= size_n) return;
// Group bookkeeping
const int groupsize = size_k / groups;
int group = offset_k / groupsize;
int nextgroup = (group + 1) * groupsize;
// Weight pointer for this expert
int qk = offset_k / 8;
const uint32_t* b_ptr = expert_weights + qk * size_n + n;
// Per-column dequant constants (4 columns per thread). fp16 uses the
// exact factored form (fp32 y/z pairs, no baked bias — see
// qdq_4_rdna3.cuh); bf16 keeps fp32 scalars as before.
float yf_h[4][2], zf_h[4][2];
float z_b_f[4], y_b_f[4];
// GPTQv1: zero_offset = 1
constexpr int zero_offset = 1;
auto refresh_group = [&](int g) {
const uint32_t* qz_row = expert_qzeros + g * (size_n / 8);
const T* sc_row = expert_scales + g * size_n;
int zeros[4];
T scales[4];
load4_zeros(qz_row, n, zeros);
load4_scales<T>(sc_row, n, scales);
if constexpr (std::is_same<T, half>::value) {
#pragma unroll
for (int i = 0; i < 4; ++i) {
gptq_rdna3::prep_zero_scale_fp16_f32((uint32_t)(zeros[i] + zero_offset),
scales[i], yf_h[i], zf_h[i]);
}
} else {
#pragma unroll
for (int i = 0; i < 4; ++i) {
gptq_rdna3::prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset),
scales[i], z_b_f[i], y_b_f[i]);
}
}
};
refresh_group(group);
float block_c[BLOCK_SIZE_M][4];
#pragma unroll
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
#pragma unroll
for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
}
// --- Main K-loop ---
int k = offset_k;
while (k < end_k) {
if (k == nextgroup) {
group++;
nextgroup += groupsize;
refresh_group(group);
}
// Prefetch 4 weight words (128 bytes)
int4 b_w[4];
#pragma unroll
for (int j = 0; j < 4; ++j) {
b_w[j] = *(const int4*)(b_ptr + j * size_n);
}
b_ptr += 4 * size_n;
// fp16 exact factored form: activation sums per row, split by nibble
// slot, accumulated across the 32-K round; the z correction is applied
// once after the j loop (see qdq_4_rdna3.cuh).
float sum_lo[BLOCK_SIZE_M] = {}, sum_hi[BLOCK_SIZE_M] = {};
#pragma unroll
for (int j = 0; j < 4; ++j) {
const int a_off = (k - offset_k) + 8 * j;
if constexpr (std::is_same<T, half>::value) {
// fp16 path: magic values without bias, dots via v_dot2_f32_f16,
// the (y, z) correction applied in fp32 (qdq_4_rdna3.cuh).
half2 qm[4][4];
gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].x, qm[0]);
gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].y, qm[1]);
gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].z, qm[2]);
gptq_rdna3::magic_4bit_8_fp16((uint32_t)b_w[j].w, qm[3]);
const half2 ones = gptq_rdna3::ones_half2_fp16();
#pragma unroll
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
const half2* a2 = reinterpret_cast<const half2*>(&block_a[m][a_off]);
sum_lo[m] = __builtin_amdgcn_fdot2(a2[0], ones, sum_lo[m], false);
sum_hi[m] = __builtin_amdgcn_fdot2(a2[1], ones, sum_hi[m], false);
sum_lo[m] = __builtin_amdgcn_fdot2(a2[2], ones, sum_lo[m], false);
sum_hi[m] = __builtin_amdgcn_fdot2(a2[3], ones, sum_hi[m], false);
#pragma unroll
for (int c = 0; c < 4; ++c) {
float pl = 0.0f, ph = 0.0f;
pl = __builtin_amdgcn_fdot2(qm[c][0], a2[0], pl, /*clamp=*/false);
ph = __builtin_amdgcn_fdot2(qm[c][1], a2[1], ph, /*clamp=*/false);
pl = __builtin_amdgcn_fdot2(qm[c][2], a2[2], pl, /*clamp=*/false);
ph = __builtin_amdgcn_fdot2(qm[c][3], a2[3], ph, /*clamp=*/false);
block_c[m][c] += yf_h[c][0] * pl + yf_h[c][1] * ph;
}
}
} else if constexpr (BLOCK_SIZE_M == 1) {
// bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity
typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
constexpr uint32_t BF16_MAGIC = 0x43004300u;
constexpr uint32_t BF16_ONES = 0x3F803F80u;
union pack4 {
float f[4];
uint32_t u[4];
};
uint32_t w[4];
__builtin_memcpy(w, &b_w[j], sizeof(int4));
// Load activations — read from global (no LDS for bf16 M=1)
pack4 a_pack;
{
int32_t token_id = sorted_token_ids[offset_m_base];
int token_row = token_id / top_k;
if (token_row < size_m) {
const uint32_t* a_words = reinterpret_cast<const uint32_t*>(
a + (int64_t)token_row * size_k + offset_k + a_off);
a_pack.u[0] = a_words[0];
a_pack.u[1] = a_words[1];
a_pack.u[2] = a_words[2];
a_pack.u[3] = a_words[3];
} else {
a_pack.u[0] = 0;
a_pack.u[1] = 0;
a_pack.u[2] = 0;
a_pack.u[3] = 0;
}
}
// sum_a for bias correction
float sum_a = 0.0f;
#pragma unroll
for (int b = 0; b < 4; ++b) {
sum_a = __builtin_amdgcn_fdot2_f32_bf16(
*((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
sum_a, /*clamp=*/false);
}
#pragma unroll 1
for (int col = 0; col < 4; ++col) {
pack4 q_pack;
const uint32_t qa = w[col];
q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
float partial = 0.0f;
#pragma unroll
for (int b = 0; b < 4; ++b) {
partial = __builtin_amdgcn_fdot2_f32_bf16(
*((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
partial, /*clamp=*/false);
}
block_c[0][col] =
__fmaf_rn(y_b_f[col], partial,
__fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
}
} else {
// bf16 M>1: v_dot2_f32_bf16 with LDS-staged activations
typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
constexpr uint32_t BF16_MAGIC = 0x43004300u;
constexpr uint32_t BF16_ONES = 0x3F803F80u;
union pack4 {
float f[4];
uint32_t u[4];
};
uint32_t w[4];
__builtin_memcpy(w, &b_w[j], sizeof(int4));
pack4 a_pack[BLOCK_SIZE_M];
#pragma unroll
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
const uint32_t* a_words =
reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
a_pack[m].u[0] = a_words[0];
a_pack[m].u[1] = a_words[1];
a_pack[m].u[2] = a_words[2];
a_pack[m].u[3] = a_words[3];
}
float sum_a[BLOCK_SIZE_M];
#pragma unroll
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
float s = 0.0f;
#pragma unroll
for (int b = 0; b < 4; ++b) {
s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
*((const bf16x2_t*)&BF16_ONES),
s, /*clamp=*/false);
}
sum_a[m] = s;
}
#pragma unroll 1
for (int col = 0; col < 4; ++col) {
pack4 q_pack;
const uint32_t qa = w[col];
q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
#pragma unroll
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
float partial = 0.0f;
#pragma unroll
for (int b = 0; b < 4; ++b) {
partial = __builtin_amdgcn_fdot2_f32_bf16(
*((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
partial, /*clamp=*/false);
}
block_c[m][col] =
__fmaf_rn(y_b_f[col], partial,
__fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
}
}
}
}
if constexpr (std::is_same<T, half>::value) {
// z is a group constant and sum_a is complete for this round: apply
// the fp16 z correction once per round (see qdq_4_rdna3.cuh).
#pragma unroll
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
#pragma unroll
for (int c = 0; c < 4; ++c) {
block_c[m][c] += zf_h[c][0] * sum_lo[m] + zf_h[c][1] * sum_hi[m];
}
}
}
k += 32;
}
// --- Epilogue: apply topk_weight and atomic-add to output ---
#pragma unroll
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
int32_t token_id = sorted_token_ids[offset_m_base + m];
if (token_id / top_k >= size_m) continue;
// Apply router weight
if (mul_topk_weight && topk_weights != nullptr) {
float tw = topk_weights[token_id];
#pragma unroll
for (int j = 0; j < 4; ++j) block_c[m][j] *= tw;
}
// output_topk > 0: reduce by mapping token_id back to original token
// (multiple experts write to the same row via atomics)
int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk)
: (int64_t)token_id;
T* out = c + out_row * size_n + n;
if constexpr (std::is_same<T, half>::value) {
half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
__float2half_rn(block_c[m][1]));
half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
__float2half_rn(block_c[m][3]));
atomic_add_pk4_f16(out, r01, r23);
} else {
bf162_t r01;
r01.x = __float2bfloat16(block_c[m][0]);
r01.y = __float2bfloat16(block_c[m][1]);
bf162_t r23;
r23.x = __float2bfloat16(block_c[m][2]);
r23.y = __float2bfloat16(block_c[m][3]);
atomic_add_pk4_bf16(out, r01, r23);
}
}
}
#else // non-RDNA3: empty stub for symbol parity
template <typename T, int BLOCK_SIZE_M>
__global__ void moe_gemm_q4_kernel_rdna3(
const T*, T*, const uint32_t*, const T*, const uint32_t*, const float*,
const int32_t*, const int32_t*, const int32_t*, const int, const int,
const int, const int, const int, const int, const int, const int,
const bool, const int) {}
#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
// ---------------------------------------------------------------------------
// Launcher
// ---------------------------------------------------------------------------
template <typename T, int BLOCK_SIZE_M>
void launch_moe_gemm_q4(
const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
const uint32_t* b_qzeros, const float* topk_weights,
const int32_t* sorted_token_ids, const int32_t* expert_ids,
const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
int size_n, int size_k, int groups, int top_k, int expert_weight_stride,
int expert_scales_stride, int expert_zeros_stride, bool mul_topk_weight,
int output_topk, cudaStream_t stream) {
dim3 block(THREADS_X);
dim3 grid(num_token_blocks,
(size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
(size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
moe_gemm_q4_kernel_rdna3<T, BLOCK_SIZE_M><<<grid, block, 0, stream>>>(
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
expert_ids, num_tokens_post_padded, size_m, size_n, size_k, groups, top_k,
expert_weight_stride, expert_scales_stride, expert_zeros_stride,
mul_topk_weight, output_topk);
}
template <typename T>
void dispatch_moe_gemm_q4(
const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
const uint32_t* b_qzeros, const float* topk_weights,
const int32_t* sorted_token_ids, const int32_t* expert_ids,
const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
int size_n, int size_k, int groups, int top_k, int block_size_m,
int expert_weight_stride, int expert_scales_stride, int expert_zeros_stride,
bool mul_topk_weight, int output_topk, cudaStream_t stream) {
// Dispatch to template instantiation based on block_size_m
switch (block_size_m) {
case 1:
launch_moe_gemm_q4<T, 1>(
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
expert_zeros_stride, mul_topk_weight, output_topk, stream);
break;
case 2:
launch_moe_gemm_q4<T, 2>(
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
expert_zeros_stride, mul_topk_weight, output_topk, stream);
break;
case 4:
launch_moe_gemm_q4<T, 4>(
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
expert_zeros_stride, mul_topk_weight, output_topk, stream);
break;
case 8:
launch_moe_gemm_q4<T, 8>(
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
expert_zeros_stride, mul_topk_weight, output_topk, stream);
break;
default:
TORCH_CHECK(false,
"moe_gptq_gemm_rdna3: block_size_m must be 1, 2, 4, or 8, "
"got ",
block_size_m);
}
}
} // namespace moe_gptq_rdna3
} // namespace vllm
// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------
//
// Inputs:
// a [M, K] or [M*top_k, K] half or bfloat16
// c [M*top_k, N] same dtype (pre-zeroed!)
// b_q_weight [E, K/8, N] uint32 (shuffled)
// b_scales [E, groups, N] same dtype as a
// b_qzeros [E, groups, N/8] uint32 (packed 4-bit)
// topk_weights [M*top_k] or empty float32
// sorted_token_ids [num_blocks * block_m] int32
// expert_ids [num_blocks] int32
// num_tokens_post_padded [1] int32
// top_k int
// block_size_m int (1, 2, 4, or 8)
// mul_topk_weight bool
void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
torch::Tensor b_q_weight, torch::Tensor b_scales,
torch::Tensor b_qzeros, torch::Tensor topk_weights,
torch::Tensor sorted_token_ids,
torch::Tensor expert_ids,
torch::Tensor num_tokens_post_padded, int64_t top_k,
int64_t block_size_m, bool mul_topk_weight,
int64_t output_topk) {
TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
TORCH_CHECK(c.is_cuda(), "c must be a CUDA/HIP tensor");
TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
TORCH_CHECK(a.dim() == 2, "a must be 2D");
TORCH_CHECK(c.dim() == 2, "c must be 2D");
TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [E, K/8, N]");
TORCH_CHECK(b_scales.dim() == 3, "b_scales must be 3D [E, groups, N]");
TORCH_CHECK(b_qzeros.dim() == 3, "b_qzeros must be 3D [E, groups, N/8]");
TORCH_CHECK(
a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
"a must be half or bfloat16");
TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
"b_scales dtype must match a");
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
auto stream = at::cuda::getCurrentCUDAStream();
int size_m = (int)a.size(0);
int size_k = (int)a.size(1);
int size_n = (int)b_q_weight.size(2);
int groups = (int)b_scales.size(1);
TORCH_CHECK(size_k % (groups * 32) == 0,
"group size (K/groups = ", size_k / groups,
") must be a multiple of 32: the kernel checks group "
"transitions at 32-K granularity");
// Per-expert strides
int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2));
int expert_scales_stride = (int)(b_scales.size(1) * b_scales.size(2));
int expert_zeros_stride = (int)(b_qzeros.size(1) * b_qzeros.size(2));
int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m);
const float* topk_w_ptr =
(topk_weights.numel() > 0) ? topk_weights.data_ptr<float>() : nullptr;
// Manual dtype dispatch using HIP native types (c10::Half/BFloat16 don't
// implicitly convert to half/__hip_bfloat16 in device code).
using vllm::gptq_rdna3::bf16_t;
auto dispatch = [&](auto* a_ptr, auto* c_ptr, const auto* s_ptr) {
using T = std::remove_const_t<std::remove_pointer_t<decltype(a_ptr)>>;
vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4<T>(
a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr<int32_t>(), s_ptr,
(const uint32_t*)b_qzeros.data_ptr<int32_t>(), topk_w_ptr,
sorted_token_ids.data_ptr<int32_t>(), expert_ids.data_ptr<int32_t>(),
num_tokens_post_padded.data_ptr<int32_t>(), num_token_blocks, size_m,
size_n, size_k, groups, (int)top_k, (int)block_size_m,
expert_weight_stride, expert_scales_stride, expert_zeros_stride,
mul_topk_weight, (int)output_topk, stream);
};
if (a.scalar_type() == torch::kHalf) {
dispatch((const half*)a.data_ptr(), (half*)c.data_ptr(),
(const half*)b_scales.data_ptr());
} else {
dispatch((const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(),
(const bf16_t*)b_scales.data_ptr());
}
}
csrc/rocm/ops.h
-53
View File
@@ -1,53 +0,0 @@
#pragma once
#include <torch/all.h>
torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
const int64_t rows_per_block);
torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
const std::optional<at::Tensor>& in_bias,
const int64_t CuCount);
torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b,
const at::Tensor& in_scale,
const std::optional<at::Tensor>& in_zero_points,
const std::optional<at::Tensor>& in_bias,
const int64_t CuCount, const int64_t group_size);
torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
const std::optional<at::Tensor>& in_bias,
const int64_t CuCount);
void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
const at::Tensor& scale_a, const at::Tensor& scale_b,
const int64_t CuCount);
torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
torch::Tensor b_qzeros, torch::Tensor b_scales,
bool use_v2_format);
torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
torch::Tensor b_qzeros,
torch::Tensor b_scales, bool use_v2_format);
void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
torch::Tensor b_q_weight, torch::Tensor b_scales,
torch::Tensor b_qzeros, torch::Tensor topk_weights,
torch::Tensor sorted_token_ids,
torch::Tensor expert_ids,
torch::Tensor num_tokens_post_padded, int64_t top_k,
int64_t block_size_m, bool mul_topk_weight,
int64_t output_topk);
void paged_attention(
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens,
const std::optional<torch::Tensor>& query_start_loc, int64_t block_size,
int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale,
const std::string& mfma_type, bool use_interleaved_v_cache);
csrc/rocm/q_gemm_rdna3.cu
-1040
View File
File diff suppressed because it is too large Load Diff
csrc/rocm/q_gemm_rdna3_wmma.cu
-2405
View File
File diff suppressed because it is too large Load Diff
csrc/rocm/qdq_4_rdna3.cuh
-142
View File
@@ -1,142 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// W4A16 dequant primitives for RDNA3 (gfx1100/gfx1101/gfx1102), templated on
// the activation/scale dtype (half or __hip_bfloat16). The fp16 path reuses
// the classic exllamav2 bit-trick:
//
// (qa & 0x000F000F) | 0x64006400 -> half2(1024+q_lo, 1024+q_hi)
// (qa & 0x00F000F0) | 0x64006400 -> half2(1024+q_lo*16, 1024+q_hi*16)
//
// The "*16 then divide by 16 in the FMA" trick for the upper-nibble pairs
// works in fp16 because the mantissa (10 bits) is wide enough to hold a value
// shifted by 4 bits. In bf16 the mantissa is only 7 bits, so shifting an upper
// nibble into bits [7:4] would spill into the exponent. To avoid that, the
// bf16 path shifts each pair of nibbles down to bits [3:0]/[19:16] with a
// single right-shift before the OR with 0x43004300 (= bf162(128, 128)).
#ifndef _qdq_4_rdna3_cuh
#define _qdq_4_rdna3_cuh
#include <cstdint>
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
namespace vllm {
namespace gptq_rdna3 {
using bf16_t = __hip_bfloat16;
using bf162_t = __hip_bfloat162;
// Bit-shuffle for an int32 holding 8 sequential 4-bit weights q[0..7]:
// in: q[7] q[6] q[5] q[4] q[3] q[2] q[1] q[0] (LSB first)
// out: q[7] q[5] q[3] q[1] q[6] q[4] q[2] q[0] (even/odd interleaved)
//
// After shuffle, q[2k] sits at bits [4k : 4k+3] (lower 16)
// q[2k+1] sits at bits [16+4k: 16+4k+3] (upper 16)
// so a single mask 0x000F000F selects the matching even/odd pair, ready to
// bitcast to half2 / bfloat162 after OR-ing with the magic constant.
__forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) {
uint32_t qa = q[0];
uint32_t qb = 0;
#pragma unroll
for (int i = 0; i < 4; i++) {
uint32_t qa0 = qa & 0x0F;
uint32_t qa1 = (qa & 0xF0) >> 4;
qa >>= 8;
qb |= (qa1 << (i * 4 + 16));
qb |= (qa0 << (i * 4));
}
q[0] = qb;
}
// ---------------------------------------------------------------------------
// fp16 path: exact factored form.
// ---------------------------------------------------------------------------
// The pre-#54706-v2 kernels baked the zero/scale bias into every weight,
// which forced z = scale * (-1024 - zero) to be *stored* as fp16. At
// scale ~0.02 that costs ~0.008 abs per (group, column) constant and the
// error then accumulates along the K/groups axis (measured max abs
// error ~1.0 vs an FP32 dequantized reference at K=4096). Computing the
// product in fp32 does not help: the error is in the fp16 *representation*
// of z, not in the multiply.
//
// Instead use the identity, exact for any grouping:
//
// sum_k a_k * w_k = y * sum_k a_k * (1024 + q_k) + z * sum_k a_k
//
// with one (y, z) pair per nibble slot (the high pairs carry a factor of
// 16): y_lo = s, z_lo = s*(-1024 - zero), y_hi = s/16, z_hi = s*(-64 -
// zero). The y/z constants live in fp32 and are never narrowed; the magic
// values (1024 + q) / (1024 + 16q) stay in fp16 where they are exact. The
// z * sum_a correction is applied once per 32-K round in the kernel (z is
// constant within a group, so per-round applications sum to the per-group
// application). sum_k a_k does not depend on the output column, so the
// extra v_dot2 against half2(1,1) is shared by all four columns a thread
// owns. This is the same structure the bf16 branches of q_gemm_rdna3.cu
// already use; fp16 needs two (y, z) pairs and two activation sums because
// its magic values carry the upper-nibble *16 trick.
// ---------------------------------------------------------------------------
// fp32 (y, z) pairs for the exact factored form above.
__forceinline__ __device__ void prep_zero_scale_fp16_f32(uint32_t zero,
half scale,
float (&y)[2],
float (&z)[2]) {
const float s = __half2float(scale);
const float zf = (float)(int)zero;
y[0] = s; // low pairs: (1024 + q) * s
z[0] = s * (-1024.0f - zf);
y[1] = s * (1.0f / 16.0f); // high pairs: (1024 + 16q) * s/16
z[1] = s * (-64.0f - zf);
}
// Magic values only, with no y/z folded in: low pairs land in [0] and [2],
// high pairs in [1] and [3]. The lane mapping assumes the host-side
// gptq_shuffle interleave (see shuffle_4bit_8 above): even elements sit in
// the low nibbles, odd elements in the high nibbles.
__forceinline__ __device__ void magic_4bit_8_fp16(uint32_t qa, half2 (&qm)[4]) {
const uint32_t c0 = 0x64006400;
union {
uint32_t u;
half2 h2;
} t;
t.u = (qa & 0x000F000F) | c0; // half2(1024 + q[0], 1024 + q[1])
qm[0] = t.h2;
t.u = (qa & 0x00F000F0) | c0; // half2(1024 + q[2]*16, 1024 + q[3]*16)
qm[1] = t.h2;
const uint32_t qa_hi = qa >> 8;
t.u = (qa_hi & 0x000F000F) | c0; // half2(1024 + q[4], 1024 + q[5])
qm[2] = t.h2;
t.u = (qa_hi & 0x00F000F0) | c0; // half2(1024 + q[6]*16, 1024 + q[7]*16)
qm[3] = t.h2;
}
// half2(1.0, 1.0) for the activation-sum v_dot2 in the exact factored form.
__forceinline__ __device__ half2 ones_half2_fp16() {
union {
uint32_t u;
half2 h2;
} t;
t.u = 0x3C003C00u;
return t.h2;
}
// ---------------------------------------------------------------------------
// bf16 path
// ---------------------------------------------------------------------------
__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
bf16_t scale,
float& z_prep,
float& y_prep) {
float scale_f = __bfloat162float(scale);
z_prep = -(128.0f + (float)zero) * scale_f;
y_prep = scale_f;
}
} // namespace gptq_rdna3
} // namespace vllm
#endif // _qdq_4_rdna3_cuh
csrc/rocm/quickreduce/base.h
-372
View File
@@ -1,372 +0,0 @@
#pragma once
#include <cstdint>
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include <hip/hip_bf16.h>
#define __quickreduce_device_inline__ __device__ __forceinline__
#define __quickreduce_launch_bounds_two_shot__ __launch_bounds__(256, 4)
#define __quickreduce_launch_bounds_one_shot__ __launch_bounds__(512, 4)
namespace quickreduce {
typedef __hip_bfloat16 nv_bfloat16;
typedef __hip_bfloat162 nv_bfloat162;
using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
// Setup acquire-release semantics for vector memory reads (mubuf instruction)
// as per architecture.
#if defined(__gfx942__)
// CDNA3: Scope bits sc0, sc1
#define MUBUF_ACQUIRE 16
#define MUBUF_RELEASE 16
#elif (defined(__gfx908__) || defined(__gfx90a__))
// CDNA1 and CDNA2 - glc bit
#define MUBUF_ACQUIRE 1
#define MUBUF_RELEASE 0
#endif
static constexpr int kNegOne = 0xBC00BC00; // {-1, -1}, fp16x2_t
// Number of atoms (4xf16x2_t) processed by a single thread
static constexpr int kAtoms = 8;
// We use a workgroup of 256 threads
static constexpr int kBlockSize = 256;
static constexpr int kAtomStride = kBlockSize;
// Size and atom stride of source/destination data that the block will
// process.
// Workgroup scope = Tile = (256 threads x 8 atoms x 16B)
static constexpr int kTileSize = kBlockSize * kAtoms * sizeof(int32x4_t);
// Max number of blocks. 304 CUs on MI300
static constexpr int kMaxNumBlocks = 304 * 4;
// Standard CDNA wavefront size.
static constexpr int kWavefront = 64;
// 256 thread, 4 wavefronts.
static dim3 constexpr kBlockTwoShot = {kWavefront, kBlockSize / kWavefront, 1};
// Number of threads in a group for quantization
// It corresponds to 32 F16 elements in quantization block
static constexpr int kThreadGroupSize = 8;
// Methods
__quickreduce_device_inline__ __host__ unsigned long divceil(unsigned long x,
unsigned long y) {
return ((x + y - 1) / y);
}
union BufferResource {
__quickreduce_device_inline__ constexpr BufferResource()
: config(0x00020000U) {}
__quickreduce_device_inline__ constexpr BufferResource(void* buffer_address,
uint32_t buffer_size)
: address(buffer_address), range(buffer_size), config(0x00020000U) {}
int32x4_t descriptor;
struct {
void* address; // 8B, out of which first 48b is address, and 16b is stride
// (unused)
uint32_t range; // Byte range for the buffer resource
uint32_t config; // Constant, DFMT=32b
};
};
#if !defined(__gfx1250__)
__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
int32x4_t srsrc, int32_t voffset, int32_t soffset,
int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
__quickreduce_device_inline__ static void buffer_store_dwordx4(
int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset,
int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
#else
__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {}
__quickreduce_device_inline__ static void buffer_store_dwordx4(int32x4_t data,
int32x4_t srsrc,
int32_t voffset,
int32_t soffset,
int32_t aux) {}
#endif
__quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
#if defined(__gfx942__)
if (value) {
asm volatile("s_setreg_imm32_b32 0xdc1, 1;" ::);
} else {
asm volatile("s_setreg_imm32_b32 0xdc1, 0;" ::);
}
#endif
}
union bf162_int_union {
int i;
nv_bfloat162 bf2;
};
template <typename T>
__quickreduce_device_inline__ void packed_assign_add(int32x4_t* A,
int32x4_t* B);
template <>
__quickreduce_device_inline__ void packed_assign_add<half>(int32x4_t* A,
int32x4_t* B) {
int32x4_t& tR_fragment = A[0];
int32x4_t& tA_fragment = B[0];
asm volatile("v_pk_add_f16 %0, %1, %2"
: "=v"(tR_fragment[0])
: "v"(tR_fragment[0]), "v"(tA_fragment[0]));
asm volatile("v_pk_add_f16 %0, %1, %2"
: "=v"(tR_fragment[1])
: "v"(tR_fragment[1]), "v"(tA_fragment[1]));
asm volatile("v_pk_add_f16 %0, %1, %2"
: "=v"(tR_fragment[2])
: "v"(tR_fragment[2]), "v"(tA_fragment[2]));
asm volatile("v_pk_add_f16 %0, %1, %2"
: "=v"(tR_fragment[3])
: "v"(tR_fragment[3]), "v"(tA_fragment[3]));
}
template <>
__quickreduce_device_inline__ void packed_assign_add<nv_bfloat16>(
int32x4_t* A, int32x4_t* B) {
nv_bfloat162* tA = reinterpret_cast<nv_bfloat162*>(A);
nv_bfloat162* tB = reinterpret_cast<nv_bfloat162*>(B);
#pragma unroll
for (int i = 0; i < 4; i++) {
tA[i] = __hadd2(tA[i], tB[i]);
}
}
template <typename T>
__quickreduce_device_inline__ int packed_max(int a, int b);
template <>
__quickreduce_device_inline__ int packed_max<half>(int a, int b) {
int result;
asm volatile("v_pk_max_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
return result;
}
template <>
__quickreduce_device_inline__ int packed_max<nv_bfloat16>(int a, int b) {
bf162_int_union A, B, R;
A.i = a;
B.i = b;
R.bf2 = __hmax2(A.bf2, B.bf2);
return R.i;
}
template <typename T>
__quickreduce_device_inline__ int packed_min(int a, int b);
template <>
__quickreduce_device_inline__ int packed_min<half>(int a, int b) {
int result;
asm volatile("v_pk_min_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
return result;
}
template <>
__quickreduce_device_inline__ int packed_min<nv_bfloat16>(int a, int b) {
bf162_int_union A, B, R;
A.i = a;
B.i = b;
R.bf2 = __hmin2(A.bf2, B.bf2);
return R.i;
}
template <typename T>
__quickreduce_device_inline__ int packed_abs_max(int a, int b);
template <>
__quickreduce_device_inline__ int packed_abs_max<half>(int a, int b) {
half2 wmaxh2 = __builtin_bit_cast(half2, a);
half2 wminh2 = __builtin_bit_cast(half2, b);
half2 wblockmaxh2;
wblockmaxh2.x =
__hgt(__habs(wmaxh2.x), __habs(wminh2.x)) ? wmaxh2.x : wminh2.x;
wblockmaxh2.y =
__hgt(__habs(wmaxh2.y), __habs(wminh2.y)) ? wmaxh2.y : wminh2.y;
return __builtin_bit_cast(int, wblockmaxh2);
}
template <>
__quickreduce_device_inline__ int packed_abs_max<nv_bfloat16>(int a, int b) {
bf162_int_union A, B, R;
A.i = a;
B.i = b;
R.bf2.x = __hgt(__habs(A.bf2.x), __habs(B.bf2.x)) ? A.bf2.x : B.bf2.x;
R.bf2.y = __hgt(__habs(A.bf2.y), __habs(B.bf2.y)) ? A.bf2.y : B.bf2.y;
return R.i;
}
template <typename T>
__quickreduce_device_inline__ int packed_add(int a, int b);
template <>
__quickreduce_device_inline__ int packed_add<half>(int a, int b) {
int result;
asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
return result;
}
template <>
__quickreduce_device_inline__ int packed_add<nv_bfloat16>(int a, int b) {
bf162_int_union A, B, R;
A.i = a;
B.i = b;
R.bf2 = __hadd2(A.bf2, B.bf2);
return R.i;
}
template <>
__quickreduce_device_inline__ int packed_add<int16_t>(int a, int b) {
int result;
asm volatile("v_pk_add_i16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
return result;
}
template <typename T>
__quickreduce_device_inline__ int packed_sub(int a, int b);
template <>
__quickreduce_device_inline__ int packed_sub<half>(int a, int b) {
int result;
// MI300 lacks packed fp16 sub instruction. So we do -1 * min + max
asm volatile("v_pk_fma_f16 %0, %1, %2 %3"
: "=v"(result)
: "v"(kNegOne), "v"(b), "v"(a));
return result;
}
template <>
__quickreduce_device_inline__ int packed_sub<nv_bfloat16>(int a, int b) {
bf162_int_union A, B, R;
A.i = a;
B.i = b;
R.bf2 = __hsub2(A.bf2, B.bf2);
return R.i;
}
template <typename T>
__quickreduce_device_inline__ int packed_mul(int a, int b);
template <>
__quickreduce_device_inline__ int packed_mul<half>(int a, int b) {
int result;
asm volatile("v_pk_mul_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
return result;
}
template <>
__quickreduce_device_inline__ int packed_mul<nv_bfloat16>(int a, int b) {
nv_bfloat162* tA = reinterpret_cast<nv_bfloat162*>(&a);
nv_bfloat162* tB = reinterpret_cast<nv_bfloat162*>(&b);
nv_bfloat162 tR = __hmul2(*tA, *tB);
return *(reinterpret_cast<int*>(&tR));
}
template <typename T>
__quickreduce_device_inline__ int packed_rcp(int a);
template <>
__quickreduce_device_inline__ int packed_rcp<half>(int a) {
return __builtin_bit_cast(int, h2rcp(__builtin_bit_cast(half2, a)));
}
template <>
__quickreduce_device_inline__ int packed_rcp<nv_bfloat16>(int a) {
bf162_int_union A, R;
A.i = a;
R.bf2 = h2rcp(A.bf2);
return R.i;
}
template <typename T>
__quickreduce_device_inline__ int packed_from_int16_pair(int16_t low,
int16_t high);
template <>
__quickreduce_device_inline__ int packed_from_int16_pair<half>(int16_t low,
int16_t high) {
// Convert two signed integers to one fp16x2 packed 32-bit lane.
half2 h = __halves2half2(__int2half_rn(static_cast<int>(low)),
__int2half_rn(static_cast<int>(high)));
return __builtin_bit_cast(int, h);
}
template <>
__quickreduce_device_inline__ int packed_from_int16_pair<nv_bfloat16>(
int16_t low, int16_t high) {
// Convert two signed integers to one bf16x2 packed 32-bit lane.
nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
return *reinterpret_cast<int*>(&bf2);
}
// changes dtype
__quickreduce_device_inline__ float T2float_cast(half a) {
return __half2float(a);
}
__quickreduce_device_inline__ float T2float_cast(nv_bfloat16 a) {
return __bfloat162float(a);
}
template <typename T>
__quickreduce_device_inline__ int group_abs_max(int32x4_t atom) {
const int group_leader = (threadIdx.x / kThreadGroupSize) * kThreadGroupSize;
int wmax, wmin, wblockmax;
int a, b;
a = packed_max<T>(atom[0], atom[1]);
b = packed_max<T>(atom[2], atom[3]);
wmax = packed_max<T>(a, b);
a = packed_min<T>(atom[0], atom[1]);
b = packed_min<T>(atom[2], atom[3]);
wmin = packed_min<T>(a, b);
// Reduce the max among a group of threads
// Note: This is basically 2 blocks of values setup as the
// upper/lower halves of the f16x2_t
for (int i = 1; i < kThreadGroupSize; i <<= 1) {
int x = __shfl_down(wmax, i);
wmax = packed_max<T>(wmax, x);
int y = __shfl_down(wmin, i);
wmin = packed_min<T>(wmin, y);
}
wblockmax = packed_abs_max<T>(wmax, wmin);
// Share with the cohort
wblockmax = __shfl(wblockmax, group_leader);
return wblockmax;
}
__quickreduce_device_inline__ void set_sync_flag(uint32_t* flag_ptr,
uint32_t flag) {
__atomic_store_n(flag_ptr, flag, __ATOMIC_RELEASE);
}
__quickreduce_device_inline__ void wait_sync_flag(uint32_t* flag_ptr,
uint32_t flag) {
while (__atomic_load_n(flag_ptr, __ATOMIC_RELAXED) != flag) {
}
}
} // namespace quickreduce
csrc/rocm/quickreduce/custom_quickreduce.cu
-124
View File
@@ -1,124 +0,0 @@
#include <ATen/cuda/Exceptions.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/all.h>
#ifdef USE_ROCM
#include "rocm/quickreduce/quick_reduce.h"
quickreduce::fptr_t init_custom_qr(int64_t rank, int64_t world_size,
std::optional<int64_t> qr_max_size) {
if (world_size > 8)
throw std::invalid_argument("world size > 8 is not supported");
if (world_size == 6)
throw std::invalid_argument("world size == 6 is not supported");
if (world_size % 2 != 0)
throw std::invalid_argument("Odd num gpus is not supported for now");
if (rank < 0 || rank >= world_size)
throw std::invalid_argument("invalid rank passed in");
quickreduce::DeviceComms* fptr = new quickreduce::DeviceComms();
fptr->init(world_size, rank, qr_max_size);
return (quickreduce::fptr_t)fptr;
}
void qr_destroy(quickreduce::fptr_t _fa) {
if (_fa) {
auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
fa->destroy();
delete fa;
}
}
torch::Tensor qr_get_handle(quickreduce::fptr_t _fa) {
auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
hipIpcMemHandle_t handle = fa->get_handle();
auto options =
torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
auto data_handle =
torch::empty({static_cast<int64_t>(sizeof(hipIpcMemHandle_t))}, options);
std::memcpy(data_handle.data_ptr(), &handle, sizeof(hipIpcMemHandle_t));
return data_handle;
}
void qr_open_handles(quickreduce::fptr_t _fa,
const std::vector<torch::Tensor>& handles) {
auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
std::vector<hipIpcMemHandle_t> ipc_handles;
ipc_handles.reserve(handles.size());
for (auto& handle : handles) {
// Ensure the tensor is on the same device as the current device.
hipIpcMemHandle_t ipc_handle;
std::memcpy(&ipc_handle, handle.data_ptr(), sizeof(hipIpcMemHandle_t));
ipc_handles.push_back(ipc_handle);
}
fa->open_ipc_handles(ipc_handles);
}
void qr_all_reduce(quickreduce::fptr_t _fa, torch::Tensor& inp,
torch::Tensor& out, int64_t quant_level, bool cast_bf2half) {
auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
auto stream = at::cuda::getCurrentHIPStreamMasqueradingAsCUDA();
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
TORCH_CHECK_EQ(inp.numel(), out.numel());
TORCH_CHECK_LE(out.numel(), fa->kMaxProblemSize);
if (out.scalar_type() == at::ScalarType::Half) {
fa->allreduce<half, false>(reinterpret_cast<half*>(inp.data_ptr()),
reinterpret_cast<half*>(out.data_ptr()),
out.numel(), quant_level, stream);
} else if (out.scalar_type() == at::ScalarType::BFloat16) {
if (cast_bf2half) {
fa->allreduce<half, true>(reinterpret_cast<half*>(inp.data_ptr()),
reinterpret_cast<half*>(out.data_ptr()),
out.numel(), quant_level, stream);
} else {
fa->allreduce<quickreduce::nv_bfloat16, false>(
reinterpret_cast<quickreduce::nv_bfloat16*>(inp.data_ptr()),
reinterpret_cast<quickreduce::nv_bfloat16*>(out.data_ptr()),
out.numel(), quant_level, stream);
}
} else {
throw std::runtime_error(
"quick allreduce only supports float16 and bfloat16");
}
}
int64_t qr_max_size() {
// The default is 2GB (2,147,483,648 bytes)
return static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
}
#define INSTANTIATE_FOR_WORLDSIZE(T, Codec, cast_bf2half) \
template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, \
cast_bf2half>; \
template struct quickreduce::AllReduceTwoshot<T, Codec<T, 4>, \
cast_bf2half>; \
template struct quickreduce::AllReduceTwoshot<T, Codec<T, 8>, cast_bf2half>;
// INT3 (CodecQ3) is restricted to TP2 only, so we only instantiate the
// world_size == 2 kernel for it.
#define INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(T, Codec, cast_bf2half) \
template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, cast_bf2half>;
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false)
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false)
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false)
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false)
INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
quickreduce::CodecQ3, false)
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true)
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true)
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true)
INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true)
INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
quickreduce::CodecQ3, true)
INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false)
INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false)
INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false)
INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false)
INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(half, quickreduce::CodecQ3, false)
#endif // USE_ROCM
csrc/rocm/quickreduce/quick_reduce.h
-241
View File
@@ -1,241 +0,0 @@
#pragma once
#include <vector>
#include <hip/hip_runtime.h>
#include "quick_reduce_impl.cuh"
#define HIP_CHECK(err) \
do { \
hipError_t err_ = (err); \
if (err_ != hipSuccess) { \
std::printf("HIP error %d at %s:%d. %s\n", err_, __FILE__, __LINE__, \
hipGetErrorString(err_)); \
throw std::runtime_error("HIP error"); \
} \
} while (0)
namespace quickreduce {
using fptr_t = int64_t;
static_assert(sizeof(void*) == sizeof(fptr_t));
template <typename AllReduceKernel, typename T>
__global__ __quickreduce_launch_bounds_two_shot__ static void
allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
int rank, uint8_t** dbuffer_list,
uint32_t data_offset, uint32_t* d_flag_counters,
int64_t data_size_per_phase) {
int block = blockIdx.x;
int grid = gridDim.x;
// Load this block's counter from device memory and advance it on-device,
// so the color keeps changing across graph replays instead of being frozen.
uint32_t flag_color = d_flag_counters[blockIdx.x];
while (block < num_blocks) {
AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
flag_color, data_size_per_phase);
block += grid;
flag_color++;
}
// All threads compute the same final value; one writer per block is enough.
if (threadIdx.x == 0 && threadIdx.y == 0) {
d_flag_counters[blockIdx.x] = flag_color;
}
}
#define TWOSHOT_DISPATCH(__codec) \
if (world_size == 2) { \
using LineCodec = __codec<T, 2>; \
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
num_blocks, rank, dbuffer_list, data_offset, \
d_flag_counters, this->kMaxProblemSize); \
} else if (world_size == 4) { \
using LineCodec = __codec<T, 4>; \
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
num_blocks, rank, dbuffer_list, data_offset, \
d_flag_counters, this->kMaxProblemSize); \
} else if (world_size == 8) { \
using LineCodec = __codec<T, 8>; \
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
num_blocks, rank, dbuffer_list, data_offset, \
d_flag_counters, this->kMaxProblemSize); \
}
// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
// the 3-bit codec's pack/unpack overhead outweighs the reduced communication
// volume, so INT3 is restricted to a TP2-only dispatch here.
#define TWOSHOT_DISPATCH_TP2_ONLY(__codec) \
if (world_size == 2) { \
using LineCodec = __codec<T, 2>; \
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
num_blocks, rank, dbuffer_list, data_offset, \
d_flag_counters, this->kMaxProblemSize); \
} else { \
throw std::runtime_error( \
"INT3 quick all-reduce is only supported for world_size == 2 " \
"(TP2); use INT4/NONE for larger world sizes."); \
}
enum QuickReduceQuantLevel {
// Keep these ids in sync with Python QuickReduceRegime enum.
F16 = 0, // full-precision fp16/bf16 communication
INT8 = 1, // symmetric int8 + per-block scale
INT6 = 2, // symmetric int6 + per-block scale
INT4 = 3, // symmetric int4 + per-block scale
INT3 = 4, // symmetric int3 + per-block scale (TP2 only)
};
struct DeviceComms {
// Max problem size is 2GB (in bytes) or half of uint32_t max value.
int64_t kMaxProblemSize =
static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
// Max TP-8
static int constexpr kMaxWorldSize = 8;
bool initialized = false;
uint32_t* d_flag_counters = nullptr;
int world_size;
int rank;
uint8_t* dbuffer;
uint8_t** dbuffer_list;
hipIpcMemHandle_t buffer_ipc_handle;
std::vector<hipIpcMemHandle_t> all_buffer_ipc_handles;
std::vector<uint8_t*> buffer_list;
uint32_t data_offset;
DeviceComms() : initialized(false), world_size(1), rank(0) {}
~DeviceComms() { destroy(); }
void init(int world_size, int rank,
std::optional<int64_t> max_problem_size = std::nullopt) {
destroy();
this->world_size = world_size;
this->rank = rank;
if (max_problem_size.has_value() && max_problem_size.value() > 0) {
this->kMaxProblemSize = max_problem_size.value();
}
// Allocate buffer size for worst case: F16 2-stage buffer.
uint32_t flags_buffer_size =
2 * world_size * kMaxNumBlocks * sizeof(uint32_t);
static int64_t data_buffer_size = 2 * this->kMaxProblemSize;
int64_t total_buffer_size = flags_buffer_size + data_buffer_size;
data_offset = flags_buffer_size;
HIP_CHECK(hipExtMallocWithFlags((void**)&dbuffer, total_buffer_size,
hipDeviceMallocUncached));
// Clear the flags buffer.
HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
// One flag-color counter per block, advanced by the kernel. Start at 1
// to stay clear of the flags buffer we just zeroed.
HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
{
std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(),
kMaxNumBlocks * sizeof(uint32_t),
hipMemcpyHostToDevice));
}
// Device-side list of IPC buffers.
buffer_list.resize(world_size);
HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
// Create IPC handles for rank's communication buffer.
all_buffer_ipc_handles.resize(world_size);
HIP_CHECK(hipIpcGetMemHandle(&buffer_ipc_handle, dbuffer));
initialized = true;
}
int get_world_size() { return world_size; }
int get_rank() { return rank; }
bool status() { return initialized; }
hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; }
void destroy() {
// Allocated before `initialized` flips true, so free it on its own guard
// to avoid a leak if init fails partway through.
if (d_flag_counters) {
HIP_CHECK(hipFree(d_flag_counters));
d_flag_counters = nullptr;
}
if (initialized) {
for (int i = 0; i < world_size; i++) {
if (i != rank) {
HIP_CHECK(hipIpcCloseMemHandle(dbuffer_list[i]));
}
}
HIP_CHECK(hipFree(dbuffer));
HIP_CHECK(hipFree(dbuffer_list));
initialized = false;
}
}
void open_ipc_handles(std::vector<hipIpcMemHandle_t> const& ipc_handles) {
assert(ipc_handles.size() == all_buffer_ipc_handles.size());
for (int i = 0; i < world_size; i++) {
all_buffer_ipc_handles[i] = ipc_handles[i];
}
// Open device memory access to the IPC communication buffers.
// Note: For our own rank, we do not need to open a handle.
for (int i = 0; i < world_size; i++) {
if (i != rank) {
HIP_CHECK(hipIpcOpenMemHandle((void**)&buffer_list[i],
all_buffer_ipc_handles[i],
hipIpcMemLazyEnablePeerAccess));
} else {
buffer_list[i] = dbuffer;
}
}
HIP_CHECK(hipMemcpy(dbuffer_list, buffer_list.data(),
world_size * sizeof(uint8_t*), hipMemcpyHostToDevice));
}
template <typename T, bool cast_bf2half>
void allreduce(T const* A, T* B, uint32_t N, int quant_level,
hipStream_t stream) {
if (world_size != 2 && world_size != 4 && world_size != 8) {
throw std::runtime_error("All Reduce not supported for world_size = " +
std::to_string(world_size));
}
// Configuration.
uint32_t msg_size = N * sizeof(T);
uint32_t num_blocks = divceil(msg_size, kTileSize);
uint32_t grid = min(kMaxNumBlocks, num_blocks);
auto quant_level_ = static_cast<QuickReduceQuantLevel>(quant_level);
switch (quant_level_) {
case QuickReduceQuantLevel::INT8:
TWOSHOT_DISPATCH(CodecQ8)
break;
case QuickReduceQuantLevel::INT6:
TWOSHOT_DISPATCH(CodecQ6)
break;
case QuickReduceQuantLevel::INT4:
TWOSHOT_DISPATCH(CodecQ4)
break;
case QuickReduceQuantLevel::INT3:
TWOSHOT_DISPATCH_TP2_ONLY(CodecQ3)
break;
default:
TWOSHOT_DISPATCH(CodecFP)
break;
}
HIP_CHECK(cudaGetLastError());
}
};
} // namespace quickreduce
csrc/rocm/quickreduce/quick_reduce_impl.cuh
-857
View File
@@ -1,857 +0,0 @@
#pragma once
#include <hip/hip_runtime.h>
#include "base.h"
#include <hip/hip_fp16.h>
namespace quickreduce {
struct CodecBase {
const int thread;
const int rank;
const int group_leader;
__quickreduce_device_inline__ CodecBase(int thread, int rank)
: thread(thread),
rank(rank),
group_leader((threadIdx.x / kThreadGroupSize) * kThreadGroupSize) {
set_fp16_ovfl(true);
}
};
// Default full precision codec.
template <typename T, int world_size>
struct CodecFP : public CodecBase {
static constexpr int kWorldSize = world_size;
static constexpr int kRankAtoms = kAtoms / kWorldSize;
// Codec tile size process by this workgroup.
// Each thread processes atoms of f16x8_t (16B).
static constexpr int kRankTransmittedTileSize =
kBlockSize * kRankAtoms * sizeof(int32x4_t);
static_assert(kRankTransmittedTileSize % 16 == 0,
"kRankTransmittedTileSize must be 16B aligned.");
// Total tile size for the collective communication.
static constexpr int kTransmittedTileSize =
kRankTransmittedTileSize * kWorldSize;
__quickreduce_device_inline__ CodecFP(int thread, int rank)
: CodecBase(thread, rank) {}
__quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
const int32x4_t* __restrict__ data) {
for (int i = 0; i < kRankAtoms; i++) {
__builtin_nontemporal_store(data[i], send_buffer + thread);
send_buffer += kAtomStride;
}
}
__quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
int32x4_t* __restrict__ data) {
for (int i = 0; i < kRankAtoms; i++) {
data[i] = __builtin_nontemporal_load(*recv_buffer + thread);
*recv_buffer += kAtomStride;
}
}
};
// Int4 symmetric quantization codec.
// We quantize the FP16 data to block-scaled Int4 in blocks of 4 *
// kThreadGroupSize.
template <typename T, int world_size>
struct CodecQ4 : public CodecBase {
static constexpr int kWorldSize = world_size;
// Codec tile size process by this workgroup.
// Each threads processes a fragment of fp16x8_t (16B),
// into a int4x8_t (4B) and a fp16 scale shared among 32 values.
static constexpr int kRankAtoms = kAtoms / kWorldSize;
static constexpr int kRankTileStride = 1152;
static constexpr int kRankTileScaleOffset = 1024;
static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
static_assert(kRankTransmittedTileSize % 16 == 0,
"kRankTransmittedTileSize must be 16B aligned.");
static constexpr int kRankBufferTileStride =
kRankTileStride / sizeof(int32x4_t);
// Total tile size for the collective communication.
static constexpr int kTransmittedTileSize =
kRankTransmittedTileSize * kWorldSize;
// Constants configuration
// {-1/8.0h, -1/8.0h}, f16x2_t
static constexpr int kScaleFactor =
std::is_same<T, half>::value ? 0xB000B000 : 0xBE00BE00;
// {1e-7, 1e-7}, f16x2_t
static constexpr int kScaleEpsilon =
std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
// {-8, -8}, f16x2_t
static constexpr int kRangeMin =
std::is_same<T, half>::value ? 0xC800C800 : 0xC100C100;
// {+7, +7}, f16x2_t
static constexpr int kRangeMax =
std::is_same<T, half>::value ? 0x47004700 : 0x40E040E0;
// {+8, +8}, int16x2_t
static constexpr int kRangeBias = 0x00080008;
__quickreduce_device_inline__ CodecQ4(int thread, int rank)
: CodecBase(thread, rank) {}
__quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
const int32x4_t* __restrict__ data) {
for (int k = 0; k < kRankAtoms; k++) {
int32x4_t const atom = data[k];
// Compute the absolute maximum of the atom in the thread group
// In 2 blocks of values, upper/lower halves of the f16x2_t
int wblockmax = group_abs_max<T>(atom);
// Derive scales
int decoding_scale;
int encoding_scale;
decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
encoding_scale = packed_rcp<T>(encoding_scale);
// Apply scales to get quantized values
int32x4_t w;
for (int i = 0; i < 4; i++) {
w[i] = packed_mul<T>(atom[i], encoding_scale);
w[i] = packed_max<T>(w[i], kRangeMin);
w[i] = packed_min<T>(w[i], kRangeMax);
}
// Convert from f16x2_t to uint16x2_t
int32x4_t q;
{
int16_t* qi = reinterpret_cast<int16_t*>(&q);
T* wh = reinterpret_cast<T*>(&w);
for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
for (int i = 0; i < 4; i++) {
q[i] = packed_add<int16_t>(q[i], kRangeBias);
}
}
// Pack 8 x q4 into int32_t
int qw = q[0] | (q[1] << 4) | (q[2] << 8) | (q[3] << 12);
// Write quantized atom to send_buffer
// note: only the group leader stores the scale
uint8_t* atom_ptr =
reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
int32_t* qw_ptr = reinterpret_cast<int32_t*>(atom_ptr) + thread;
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
(thread / 8);
__builtin_nontemporal_store(qw, qw_ptr);
if (threadIdx.x == group_leader) {
__builtin_nontemporal_store(decoding_scale, qs_ptr);
}
}
}
__quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
int32x4_t* __restrict__ data) {
for (int k = 0; k < kRankAtoms; k++) {
// Directly read quantized atom from recv_buffer
uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
int32_t* qw_ptr = reinterpret_cast<int32_t*>(atom_ptr) + thread;
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
(thread / 8);
int32_t qw = __builtin_nontemporal_load(qw_ptr);
int qs = __builtin_nontemporal_load(qs_ptr);
*recv_buffer += kRankBufferTileStride;
// Unpack q4 into f16x8_t
int32x4_t w;
{
static constexpr uint kMask000F = 0x000F000F;
static constexpr uint kHalf2_1024 =
0x64006400; // {1024.0, 1024.0}, fp16x2_t
static uint constexpr kHalf2_1032 =
0xE408E408; // {-1032.0, -1032.0}, fp16x2_t
for (int i = 0; i < 4; i++) {
if constexpr (std::is_same<T, half>::value) {
int32_t q4 = ((qw >> (i * 4)) & kMask000F) | kHalf2_1024;
w[i] = packed_add<half>(q4, kHalf2_1032);
} else {
int32_t int16_2 = (qw >> (i * 4)) & kMask000F;
int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
}
}
}
// Apply decoding scales
for (int i = 0; i < 4; i++) {
w[i] = packed_mul<T>(w[i], qs);
}
data[k] = w;
}
}
};
// Int3 symmetric quantization codec.
// We quantize the FP16 data to block-scaled Int3 in blocks of 4 *
// kThreadGroupSize. Uniform symmetric quantization (round-to-int + clip),
// matching the structure of CodecQ4. Signed range is [-4, +3].
template <typename T, int world_size>
struct CodecQ3 : public CodecBase {
static constexpr int kWorldSize = world_size;
// Layout per quantization block (32 values = 8 threads * 4 fp16x2 lanes):
// - each thread owns 8 values and writes:
// * q2 payload : 8 * 2 bits -> uint16 (2 bytes)
// * q1 payload : 8 * 1 bit -> uint8 (1 byte)
// - one scale is shared per 32 values and written by group leader.
//
// kRankTileStride is split as:
// [0 .. 511] : q2 payload region (256 threads * 2 bytes)
// [512 .. 767] : q1 payload region (256 threads * 1 byte)
// [768 .. 895] : scale region (32 groups * 4 bytes)
static constexpr int kRankAtoms = kAtoms / kWorldSize;
static constexpr int kRankTileStride = 896;
static constexpr int kRankTileQ1Offset = 512;
static constexpr int kRankTileScaleOffset = 768;
static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
static_assert(kRankTransmittedTileSize % 16 == 0,
"kRankTransmittedTileSize must be 16B aligned.");
static constexpr int kRankBufferTileStride =
kRankTileStride / sizeof(int32x4_t);
static constexpr int kTransmittedTileSize =
kRankTransmittedTileSize * kWorldSize;
// {-1/4.0h, -1/4.0h}, f16x2_t / bf16x2_t. Sign-flipped so absmax maps
// to -4; the sign cancels with decoding_scale on the recv side.
static constexpr int kScaleFactor =
std::is_same<T, half>::value ? 0xB400B400 : 0xBE80BE80;
// {1e-7, 1e-7}, f16x2_t
static constexpr int kScaleEpsilon =
std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
// {-4, -4}, f16x2_t / bf16x2_t
static constexpr int kRangeMin =
std::is_same<T, half>::value ? 0xC400C400 : 0xC080C080;
// {+3, +3}, f16x2_t / bf16x2_t
static constexpr int kRangeMax =
std::is_same<T, half>::value ? 0x42004200 : 0x40404040;
// {+4, +4}, int16x2_t -- shifts signed [-4, +3] to unsigned [0, 7].
static constexpr int kRangeBias = 0x00040004;
__quickreduce_device_inline__ CodecQ3(int thread, int rank)
: CodecBase(thread, rank) {}
__quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
const int32x4_t* __restrict__ data) {
for (int k = 0; k < kRankAtoms; k++) {
int32x4_t const atom = data[k];
// 1) Per-group dynamic scale (shared across 32 values).
int wblockmax = group_abs_max<T>(atom);
int decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
int encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
encoding_scale = packed_rcp<T>(encoding_scale);
// 2) Scale + clip to signed int3 range [-4, +3].
int32x4_t w;
for (int i = 0; i < 4; i++) {
w[i] = packed_mul<T>(atom[i], encoding_scale);
w[i] = packed_max<T>(w[i], kRangeMin);
w[i] = packed_min<T>(w[i], kRangeMax);
}
// 3) Round to integer and bias to unsigned domain [0, 7].
int32x4_t q;
{
int16_t* qi = reinterpret_cast<int16_t*>(&q);
T* wh = reinterpret_cast<T*>(&w);
for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
for (int i = 0; i < 4; i++) {
q[i] = packed_add<int16_t>(q[i], kRangeBias);
}
}
// 4) Split each 3-bit unsigned value into low-2-bit and high-1-bit
// halves, packed into one uint16 (low 2 bits per value) plus one
// uint8 (high 1 bit per value).
uint16_t q2w = 0;
uint8_t q1w = 0;
{
int16_t* tw = reinterpret_cast<int16_t*>(&q);
#pragma unroll
for (int i = 0; i < 8; i++) {
uint32_t v = static_cast<uint32_t>(tw[i]) & 0x7u;
q2w |= static_cast<uint16_t>((v & 0x3u) << (i * 2));
q1w |= static_cast<uint8_t>(((v >> 2) & 0x1u) << i);
}
}
uint8_t* atom_ptr =
reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
uint16_t* q2w_ptr = reinterpret_cast<uint16_t*>(atom_ptr) + thread;
uint8_t* q1w_ptr =
reinterpret_cast<uint8_t*>(atom_ptr + kRankTileQ1Offset) + thread;
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
(thread / 8);
__builtin_nontemporal_store(q2w, q2w_ptr);
*q1w_ptr = q1w;
if (threadIdx.x == group_leader) {
__builtin_nontemporal_store(decoding_scale, qs_ptr);
}
}
}
__quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
int32x4_t* __restrict__ data) {
for (int k = 0; k < kRankAtoms; k++) {
uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
uint16_t* q2w_ptr = reinterpret_cast<uint16_t*>(atom_ptr) + thread;
uint8_t* q1w_ptr =
reinterpret_cast<uint8_t*>(atom_ptr + kRankTileQ1Offset) + thread;
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
(thread / 8);
uint16_t q2w = __builtin_nontemporal_load(q2w_ptr);
uint8_t q1w = *q1w_ptr;
int qs = __builtin_nontemporal_load(qs_ptr);
*recv_buffer += kRankBufferTileStride;
// Unpack unsigned values [0, 7] then shift back to signed domain
// [-4, +3] by adding kRangeMin.
int32x4_t w;
{
int16_t qv[8];
#pragma unroll
for (int i = 0; i < 8; i++) {
uint32_t low2 = (q2w >> (2 * i)) & 0x3u;
uint32_t high1 = (q1w >> i) & 0x1u;
qv[i] = static_cast<int16_t>(low2 | (high1 << 2));
}
#pragma unroll
for (int i = 0; i < 4; i++) {
int qpack = packed_from_int16_pair<T>(qv[2 * i], qv[2 * i + 1]);
w[i] = packed_add<T>(qpack, kRangeMin);
}
}
// Apply decode scale to reconstruct fp16/bf16 lanes.
for (int i = 0; i < 4; i++) {
w[i] = packed_mul<T>(w[i], qs);
}
data[k] = w;
}
}
};
// Int6 symmetric quantization codec.
// We quantize the FP16 data to block-scaled Int6 in blocks of 4 *
// kThreadGroupSize.
template <typename T, int world_size>
struct CodecQ6 : public CodecBase {
static constexpr int kWorldSize = world_size;
// Codec tile size process by this workgroup.
// Each threads processes a fragment of fp16x8_t (16B),
// into a int6x8_t (4B + 2B) and a fp16 scale shared among 32 values.
static constexpr int kRankAtoms = kAtoms / kWorldSize;
static constexpr int kRankTileStride = 1664;
static constexpr int kRankTileQ2Offset = 1024;
static constexpr int kRankTileScaleOffset = 1536;
static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
static_assert(kRankTransmittedTileSize % 16 == 0,
"kRankTransmittedTileSize must be 16B aligned.");
static constexpr int kRankBufferTileStride =
kRankTileStride / sizeof(int32x4_t);
// Total tile size for the collective communication.
static constexpr int kTransmittedTileSize =
kRankTransmittedTileSize * kWorldSize;
// Constants configuration
// {-1/32.0h, -1/32.0h}, fp16x2_t
static constexpr int kScaleFactor =
std::is_same<T, half>::value ? 0xA800A800 : 0xBD00BD00;
// {1e-7, 1e-7}, fp16x2_t
static constexpr int kScaleEpsilon =
std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
// {-32, -32}, fp16x2_t
static constexpr int kRangeMin =
std::is_same<T, half>::value ? 0xD000D000 : 0xC200C200;
// {+31, +31}, fp16x2_t
static constexpr int kRangeMax =
std::is_same<T, half>::value ? 0x4FC04FC0 : 0x41F841F8;
// {+32, +32}, int16x2_t
static constexpr int kRangeBias = 0x00200020;
__quickreduce_device_inline__ CodecQ6(int thread, int rank)
: CodecBase(thread, rank) {}
__quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
const int32x4_t* __restrict__ data) {
for (int k = 0; k < kRankAtoms; k++) {
int32x4_t const atom = data[k];
// Compute the absolute maximum of the atom in the thread group
// In 2 blocks of values, upper/lower halves of the f16x2_t
int wblockmax = group_abs_max<T>(atom);
// Derive scales
int decoding_scale;
int encoding_scale;
decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
encoding_scale = packed_rcp<T>(encoding_scale);
// Apply scales to get quantized values
int32x4_t w;
for (int i = 0; i < 4; i++) {
w[i] = packed_mul<T>(atom[i], encoding_scale);
w[i] = packed_max<T>(w[i], kRangeMin);
w[i] = packed_min<T>(w[i], kRangeMax);
}
// Convert from f16x2_t to uint16x2_t
int32x4_t q;
{
int16_t* qi = reinterpret_cast<int16_t*>(&q);
T* wh = reinterpret_cast<T*>(&w);
for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
for (int i = 0; i < 4; i++) {
q[i] = packed_add<int16_t>(q[i], kRangeBias);
}
}
// Pack 8 x q6 into int32_t + int16_t
uint32_t q4w;
uint16_t q2w = 0;
q4w = (q[0] & 0x000F000F) | ((q[1] & 0x000F000F) << 4) |
((q[2] & 0x000F000F) << 8) | ((q[3] & 0x000F000F) << 12);
{
int16_t* tw = reinterpret_cast<int16_t*>(&q);
#pragma unroll
for (int i = 0; i < 8; i++) {
q2w |= (tw[i] >> 4) << (i * 2);
}
}
// Write quantized atom to send_buffer
// note: only the group leader stores the scale
uint8_t* atom_ptr =
reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
uint32_t* q4w_ptr = reinterpret_cast<uint32_t*>(atom_ptr) + thread;
uint16_t* q2w_ptr =
reinterpret_cast<uint16_t*>(atom_ptr + kRankTileQ2Offset) + thread;
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
(thread / 8);
__builtin_nontemporal_store(q4w, q4w_ptr);
__builtin_nontemporal_store(q2w, q2w_ptr);
if (threadIdx.x == group_leader) {
__builtin_nontemporal_store(decoding_scale, qs_ptr);
}
}
}
__quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
int32x4_t* __restrict__ data) {
for (int k = 0; k < kRankAtoms; k++) {
// Directly read quantized atom from recv_buffer
uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
uint32_t* q4w_ptr = reinterpret_cast<uint32_t*>(atom_ptr) + thread;
uint16_t* q2w_ptr =
reinterpret_cast<uint16_t*>(atom_ptr + kRankTileQ2Offset) + thread;
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
(thread / 8);
uint32_t q4w = __builtin_nontemporal_load(q4w_ptr);
uint16_t q2w = __builtin_nontemporal_load(q2w_ptr);
int qs = __builtin_nontemporal_load(qs_ptr);
*recv_buffer += kRankBufferTileStride;
// Unpack q6 into fp16x8_t
int32x4_t w;
{
static uint constexpr kMask000F = 0x000F000F;
static uint constexpr kHalf2_1024 =
0x64006400; // {1024.0, 1024.0}, fp16x2_t
static uint constexpr kHalf2_1056 =
0xE420E420; // {-1056.0, -1056.0}, fp16x2_t
#pragma unroll
for (int i = 0; i < 4; i++) {
int32_t q4 = q4w & kMask000F;
int32_t q2 = (q2w & 0x3) | ((q2w & 0xC) << 14);
q4w >>= 4;
q2w >>= 4;
if constexpr (std::is_same<T, half>::value) {
int32_t q6 = q4 | (q2 << 4) | kHalf2_1024;
asm volatile("v_pk_add_f16 %0, %1, %2"
: "=v"(w[i])
: "v"(q6), "v"(kHalf2_1056));
} else {
int32_t int16_2 = q4 | (q2 << 4);
int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
}
}
}
// Apply decoding scales
for (int i = 0; i < 4; i++) {
w[i] = packed_mul<T>(w[i], qs);
}
data[k] = w;
}
}
};
// Int8 symmetric quantization codec.
// We quantize the FP16 data to block-scaled Int8 in blocks of 4 *
// kThreadGroupSize.
template <typename T, int world_size>
struct CodecQ8 : public CodecBase {
static constexpr int kWorldSize = world_size;
// Codec tile size process by this workgroup.
// Each threads processes a fragment of f16x8_t (16B),
// into a int8x8_t (8B) and a f16 scale shared among 32 values.
static constexpr int kRankAtoms = kAtoms / kWorldSize;
static constexpr int kRankTileStride = 2176;
static constexpr int kRankTileScaleOffset = 2048;
static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
static_assert(kRankTransmittedTileSize % 16 == 0,
"kRankTileSize must be 16B aligned.");
static constexpr int kRankBufferTileStride =
kRankTileStride / sizeof(int32x4_t);
// Total tile size for the collective communication.
static constexpr int kTransmittedTileSize =
kRankTransmittedTileSize * kWorldSize;
// Constants configuration
// {-1/128.0h, -1/128.0h}, f16x2_t
static constexpr int kScaleFactor =
std::is_same<T, half>::value ? 0xA000A000 : 0xBC00BC00;
// {1e-7, 1e-7}, f16x2_t
static constexpr int kScaleEpsilon =
std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
// {-128, -128}, f16x2_t
static constexpr int kRangeMin =
std::is_same<T, half>::value ? 0xD800D800 : 0xC300C300;
// {+127, +127}, f16x2_t
static constexpr int kRangeMax =
std::is_same<T, half>::value ? 0x57F057F0 : 0x42FE42FE;
// {+128, +128}, int16x2_t
static constexpr int kRangeBias = 0x00800080;
__quickreduce_device_inline__ CodecQ8(int thread, int rank)
: CodecBase(thread, rank) {}
__quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
int32x4_t const* __restrict__ data) {
for (int k = 0; k < kRankAtoms; k++) {
int32x4_t const atom = data[k];
// Compute the absolute maximum of the atom in the thread group
// In 2 blocks of values, upper/lower halves of the f16x2_t
int wblockmax = group_abs_max<T>(atom);
// Derive scales
int decoding_scale;
int encoding_scale;
decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
encoding_scale = packed_rcp<T>(encoding_scale);
// Apply scales to get quantized values
int32x4_t w;
for (int i = 0; i < 4; i++) {
w[i] = packed_mul<T>(atom[i], encoding_scale);
w[i] = packed_max<T>(w[i], kRangeMin);
w[i] = packed_min<T>(w[i], kRangeMax);
}
// Convert from f16x2_t to uint16x2_t
int32x4_t q;
{
int16_t* qi = reinterpret_cast<int16_t*>(&q);
T* wh = reinterpret_cast<T*>(&w);
for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
for (int i = 0; i < 4; i++) {
q[i] = packed_add<int16_t>(q[i], kRangeBias);
}
}
// Pack 8 x q8 into int32x2_t
int32x2_t qw;
qw[0] = q[0] | (q[1] << 8);
qw[1] = q[2] | (q[3] << 8);
// Write quantized atom to send_buffer
// note: only the group leader stores the scale
uint8_t* atom_ptr =
reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
int32x2_t* qw_ptr = reinterpret_cast<int32x2_t*>(atom_ptr) + thread;
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
(thread / 8);
__builtin_nontemporal_store(qw, qw_ptr);
if (threadIdx.x == group_leader) {
__builtin_nontemporal_store(decoding_scale, qs_ptr);
}
}
}
__quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
int32x4_t* __restrict__ data) {
for (int k = 0; k < kRankAtoms; k++) {
// Directly read quantized atom from recv_buffer
uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
int32x2_t* qw_ptr = reinterpret_cast<int32x2_t*>(atom_ptr) + thread;
int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
(thread / 8);
int32x2_t qw = __builtin_nontemporal_load(qw_ptr);
int qs = __builtin_nontemporal_load(qs_ptr);
*recv_buffer += kRankBufferTileStride;
// Unpack q8 into fp16x8_t
int32x4_t w;
{
static uint constexpr kMask00FF = 0x00FF00FF;
// {1024.0, 1024.0}, fp16x2_t
static uint constexpr kHalf2_1024 = 0x64006400;
// {-1152.0, -1152.0}, fp16x2_t
static uint constexpr kHalf2_1152 = 0xE480E480;
#pragma unroll
for (int i = 0; i < 4; i++) {
if constexpr (std::is_same<T, half>::value) {
int32_t q8 =
((qw[i / 2] >> ((i % 2) * 8)) & kMask00FF) | kHalf2_1024;
w[i] = packed_add<half>(q8, kHalf2_1152);
} else {
int32_t int16_2 = (qw[i / 2] >> ((i % 2) * 8)) & kMask00FF;
int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
}
}
}
// Apply decoding scales
for (int i = 0; i < 4; i++) {
w[i] = packed_mul<T>(w[i], qs);
}
data[k] = w;
}
}
};
// Twoshot All Reduce
template <typename T, class Codec, bool cast_bf2half>
struct AllReduceTwoshot {
static_assert(sizeof(T) == 2);
static constexpr int kWorldSize = Codec::kWorldSize;
__device__ static void run(
T const* __restrict__ input, T* __restrict__ output,
uint32_t const N, // number of elements
int const block, // block index
int const rank, // rank index
uint8_t** __restrict__ buffer_list, // communication buffers
uint32_t const data_offset, // offset to start of the data buffer
uint32_t flag_color, int64_t data_size_per_phase) {
// Topology
int thread = threadIdx.x + threadIdx.y * kWavefront;
uint8_t* rank_buffer = buffer_list[rank];
Codec codec(thread, rank);
int block_id = blockIdx.x;
// --------------------------------------------------------
// Read input into registers
int32x4_t tA[kAtoms];
BufferResource src_buffer(const_cast<T*>(input), N * sizeof(T));
uint32_t src_offset = block * kTileSize + thread * sizeof(int32x4_t);
for (int i = 0; i < kAtoms; i++) {
tA[i] = buffer_load_dwordx4(src_buffer.descriptor, src_offset, 0, 0);
src_offset += kAtomStride * sizeof(int32x4_t);
if constexpr (cast_bf2half) {
const nv_bfloat162* bf_buf =
reinterpret_cast<const nv_bfloat162*>(&tA[i]);
half2 half_buf[4];
#pragma unroll
for (int j = 0; j < 4; ++j) {
float2 f = __bfloat1622float2(bf_buf[j]);
half_buf[j] = __float22half2_rn(f);
}
tA[i] = *reinterpret_cast<const int32x4_t*>(half_buf);
}
}
// --------------------------------------------------------
// Phase-1A: Write segment data into the communication buffer of the target
// rank responsible for this segment.
uint32_t comm_data0_offset =
data_offset + block_id * Codec::kTransmittedTileSize;
uint32_t comm_data1_offset = data_size_per_phase + comm_data0_offset;
uint32_t comm_flags0_offset = block_id * (kWorldSize * sizeof(uint32_t));
uint32_t comm_flags1_offset = (data_offset / 2) + comm_flags0_offset;
for (int r = 0; r < kWorldSize; r++) {
int32x4_t* send_buffer =
reinterpret_cast<int32x4_t*>(buffer_list[r] + comm_data0_offset +
rank * Codec::kRankTransmittedTileSize);
codec.send(send_buffer, &tA[r * Codec::kRankAtoms]);
}
__syncthreads();
if (thread < kWorldSize) {
int r = thread;
uint32_t* flag_ptr = reinterpret_cast<uint32_t*>(
buffer_list[r] + comm_flags0_offset + rank * sizeof(uint32_t));
set_sync_flag(flag_ptr, flag_color);
}
// --------------------------------------------------------
// Phase-1B: Reduce the segment data from the communication buffers.
int32x4_t tR[Codec::kRankAtoms] = {};
{
// Read the data from the communication buffer.
int32x4_t* recv_buffer =
reinterpret_cast<int32x4_t*>(rank_buffer + comm_data0_offset);
uint32_t* flag_ptr =
reinterpret_cast<uint32_t*>(rank_buffer + comm_flags0_offset);
for (int r = 0; r < kWorldSize; r++) {
// Wait for the flags to be set.
if (thread == 0) {
wait_sync_flag(&flag_ptr[r], flag_color);
}
__syncthreads();
// note: we reuse tA as temp buffer here
codec.recv(&recv_buffer, tA);
for (int i = 0; i < Codec::kRankAtoms; i++) {
packed_assign_add<T>(&tR[i], &tA[i]);
}
}
}
// Phase-2: Write the reduced segment to every other rank
for (int r = 0; r < kWorldSize; r++) {
int32x4_t* send_buffer =
reinterpret_cast<int32x4_t*>(buffer_list[r] + comm_data1_offset +
rank * Codec::kRankTransmittedTileSize);
codec.send(send_buffer, tR);
}
__syncthreads();
if (thread < kWorldSize) {
int r = thread;
uint32_t* flag_ptr = reinterpret_cast<uint32_t*>(
buffer_list[r] + comm_flags1_offset + rank * sizeof(uint32_t));
set_sync_flag(flag_ptr, flag_color);
}
// Phase-2: Read the gather segments from the rank's communication buffer.
{
// Read the data from the communication buffer.
int32x4_t* recv_buffer =
reinterpret_cast<int32x4_t*>(rank_buffer + comm_data1_offset);
uint32_t* flag_ptr =
reinterpret_cast<uint32_t*>(rank_buffer + comm_flags1_offset);
for (int r = 0; r < kWorldSize; r++) {
// Wait for the flags to be set.
if (thread == 0) {
wait_sync_flag(&flag_ptr[r], flag_color);
}
__syncthreads();
// Gather all reduced and final rank segments into tA.
codec.recv(&recv_buffer, &tA[r * Codec::kRankAtoms]);
}
}
// --------------------------------------------------------
// Write the result to output.
BufferResource dst_buffer(output, N * sizeof(T));
uint32_t dst_offset = block * kTileSize + thread * sizeof(int32x4_t);
for (int i = 0; i < kAtoms; i++) {
if constexpr (cast_bf2half) {
const half2* half_buf = reinterpret_cast<const half2*>(&tA[i]);
nv_bfloat162 bf16_buf[4];
#pragma unroll
for (int j = 0; j < 4; ++j) {
float2 f = __half22float2(half_buf[j]);
bf16_buf[j] = __float22bfloat162_rn(f);
}
buffer_store_dwordx4(*reinterpret_cast<const int32x4_t*>(bf16_buf),
dst_buffer.descriptor, dst_offset, 0, 0);
} else {
buffer_store_dwordx4(tA[i], dst_buffer.descriptor, dst_offset, 0, 0);
}
dst_offset += kAtomStride * sizeof(int32x4_t);
}
}
};
} // namespace quickreduce
csrc/rocm/quickreduce/torch_bindings.cpp
-51
View File
@@ -1,51 +0,0 @@
#include <torch/all.h>
#include "core/registration.h"
#include <torch/library.h>
#include <torch/version.h>
#include <optional>
#include <vector>
#ifdef USE_ROCM
// Defined in custom_quickreduce.cu.
using fptr_t = int64_t;
fptr_t init_custom_qr(int64_t rank, int64_t world_size,
std::optional<int64_t> qr_max_size = std::nullopt);
void qr_destroy(fptr_t _fa);
torch::Tensor qr_get_handle(fptr_t _fa);
void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
int64_t quant_level, bool cast_bf2half = false);
int64_t qr_max_size();
#endif
// Note on op signatures:
// The X_meta signatures are for the meta functions corresponding to op X.
// They must be kept in sync with the signature for X. Generally, only
// functions that return Tensors require a meta function.
//
// See the following links for detailed docs on op registration and function
// schemas.
// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9
// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations
#ifdef USE_ROCM
TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
// Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C).
custom_ar.def(
"qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
"cast_bf2half) -> ()");
custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
custom_ar.def("init_custom_qr", &init_custom_qr);
custom_ar.def("qr_destroy", &qr_destroy);
custom_ar.def("qr_get_handle", &qr_get_handle);
custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
custom_ar.def("qr_max_size", &qr_max_size);
}
#endif
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
csrc/rocm/skinny_gemms.cu
-2544
View File
File diff suppressed because it is too large Load Diff
csrc/rocm/skinny_gemms_int4.cu
-1170
View File
File diff suppressed because it is too large Load Diff
csrc/rocm/torch_bindings.cpp
-100
View File
@@ -1,100 +0,0 @@
#include "core/registration.h"
#include "rocm/ops.h"
// Note on op signatures:
// The X_meta signatures are for the meta functions corresponding to op X.
// They must be kept in sync with the signature for X. Generally, only
// functions that return Tensors require a meta function.
//
// See the following links for detailed docs on op registration and function
// schemas.
// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9
// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
// vLLM custom ops for rocm
// skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250
// (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid
// undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250.
#ifndef VLLM_SKIP_SKINNY_GEMMS
// Custom gemm op for matrix-vector multiplication
rocm_ops.def(
"LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> "
"Tensor");
rocm_ops.impl("LLMM1", torch::kCUDA, &LLMM1);
// Custom gemm op for skinny matrix-matrix multiplication
rocm_ops.def(
"wvSplitK(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> "
"Tensor");
rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
// W4A16 grouped skinny GEMM: packed int4 weights, per-group scales,
// optional zero points [M/8, K/group_size] int32 for asymmetric
// quantization
rocm_ops.def(
"wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, "
"Tensor? in_zero_points, Tensor? in_bias, int CuCount, "
"int group_size) -> Tensor");
rocm_ops.impl("wvSplitK_int4_g", torch::kCUDA, &wvSplitK_int4_g);
// Custom gemm op for skinny matrix-matrix multiplication
rocm_ops.def(
"wvSplitKrc(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> "
"Tensor");
rocm_ops.impl("wvSplitKrc", torch::kCUDA, &wvSplitKrc);
// wvSplitK for fp8
rocm_ops.def(
"wvSplitKQ(Tensor in_a, Tensor in_b, Tensor? in_bias, Tensor! out_c, "
"Tensor scale_a, "
" Tensor scale_b, int CuCount) -> ()");
rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
#endif // VLLM_SKIP_SKINNY_GEMMS
#ifdef VLLM_ROCM_GFX1100
// W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
rocm_ops.def(
"gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
"Tensor b_scales, bool use_v2_format) -> Tensor");
rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3);
rocm_ops.def(
"gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
"Tensor b_scales, bool use_v2_format) -> Tensor");
rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
rocm_ops.def(
"moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, "
"Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, "
"Tensor sorted_token_ids, Tensor expert_ids, "
"Tensor num_tokens_post_padded, "
"int top_k, int block_size_m, bool mul_topk_weight, "
"int output_topk) -> ()");
rocm_ops.impl("moe_gptq_gemm_rdna3", torch::kCUDA, &moe_gptq_gemm_rdna3);
#endif
// Custom attention op
// Compute the attention between an input query and the cached
// keys/values using PagedAttention.
rocm_ops.def(
"paged_attention(Tensor! out, Tensor exp_sums,"
" Tensor max_logits, Tensor tmp_out,"
" Tensor query, Tensor key_cache,"
" Tensor value_cache, int num_kv_heads,"
" float scale, Tensor block_tables,"
" Tensor seq_lens,"
" Tensor? query_start_loc,"
" int block_size,"
" int max_seq_len,"
" Tensor? alibi_slopes,"
" str kv_cache_dtype,"
" Tensor k_scale, Tensor v_scale,"
" Tensor? fp8_out_scale,"
" str mfma_type,"
" bool use_interleaved_v_cache=False) -> ()");
rocm_ops.impl("paged_attention", torch::kCUDA, &paged_attention);
}
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
requirements/common.txt
-3
View File
@@ -22,10 +22,8 @@ prometheus_client >= 0.18.0
pillow >= 12.3.0 # Required for image processing
prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0
tiktoken >= 0.6.0 # Required for DBRX tokenizer
lm-format-enforcer == 0.11.3
llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" or platform_machine == "s390x"
outlines_core == 0.2.14
lark == 1.2.2
xgrammar == 0.2.8; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
typing_extensions >= 4.10
filelock >= 3.20.3 # need to contain https://github.com/tox-dev/filelock/pull/317; GHSA floor
@@ -53,4 +51,3 @@ mcp >= 2.0.0, < 3.0.0
opentelemetry-sdk >= 1.27.0
opentelemetry-api >= 1.27.0
opentelemetry-exporter-otlp >= 1.27.0
opentelemetry-semantic-conventions-ai >= 0.4.1
requirements/kv_connectors_rocm.txt
-2
View File
@@ -1,2 +0,0 @@
tblib
lm_eval[api]
requirements/rocm.txt
-30
View File
@@ -1,30 +0,0 @@
# Common dependencies
-r common.txt
# The version of gRPC libraries should be consistent with each other
grpcio==1.78.0
grpcio-reflection==1.78.0
numba == 0.65.0 # Required for N-gram speculative decoding
# Dependencies for AMD GPUs
datasets
peft
pytest-asyncio
tensorizer==2.10.1
packaging>=24.2
setuptools>=78.1.1,<80.0.0
setuptools-scm>=8
runai-model-streamer[s3,gcs,azure]==0.15.7
conch-triton-kernels==1.2.1
timm>=1.0.17
# amd-quark: required for Quark quantization on ROCm
# To be consistent with test_quark.py
amd-quark==0.12.post1
tilelang==0.1.10
# Required apache-tvm-ffi matching tilelang version
apache-tvm-ffi==0.1.11
# Required for faster safetensors model loading
fastsafetensors >= 0.3.3
# Mooncake transfer engine for KV cache offloading/transfer (ROCm build, published on PyPI)
mooncake-transfer-engine-rocm >= 0.3.13
requirements/test/rocm.in
-109
View File
@@ -1,109 +0,0 @@
-r ../common.txt
# testing
pytest
tensorizer==2.10.1
pytest-forked
pytest-asyncio
pytest-rerunfailures
pytest-shard
pytest-timeout
pytest-cov
# testing utils
albumentations # required for Nemotron Parse in test_common.py
av # required for audio_in_video tests
backoff # required for phi4mm test
blobfile # required for kimi-vl test
httpx
librosa # required for audio tests
vector_quantize_pytorch # required for minicpmo_26 test
vocos # required for minicpmo_26 test
peft>=0.20.0 # required for phi-4-mm test
pqdm
ray[default]==2.55.1 # Includes worker startup getenv/setenv race fix
sentence-transformers>=5.2.0 # required for embedding tests
soundfile # required for audio tests
jiwer # required for audio tests
tblib # for pickling test exceptions
timm>=1.0.17 # required for internvl and gemma3n-mm test
transformers_stream_generator # required for qwen-vl test
matplotlib # required for qwen-vl test
mistral_common[image,audio]>=1.11.6 # required for voxtral test
num2words # required for smolvlm test
open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
datamodel_code_generator>=0.64.0 # required for minicpm3 test
nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs
lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.19.0
schemathesis>=4.0.0 # Required for openai schema test
# quantization
bitsandbytes==0.50.0
buildkite-test-collector==0.1.9
tilelang==0.1.10
genai_perf>=0.0.8
tritonclient>=2.51.0
# The version of gRPC libraries should be consistent with each other
grpcio==1.78.0
grpcio-reflection==1.78.0
arctic-inference==0.1.1 # Required for suffix decoding test
numba==0.65.0 # Required for N-gram speculative decoding
runai-model-streamer[s3,gcs,azure]==0.15.7
fastsafetensors>=0.3.3
instanttensor>=0.1.9
decord==0.6.0
# Prithvi tests
# terratorch is temporarily disabled while PyPI has the `lightning` package
# in `quarantined` status (every published terratorch version transitively
# requires `lightning`, so the resolver fails with "no versions of lightning").
# Re-enable once PyPI lifts the quarantine. Tracked in #41376.
# terratorch>=1.2.2
imagehash # Required for Prithvi tests
segmentation-models-pytorch>0.4.0 # Required for Prithvi tests
gpt-oss[eval]>=0.0.7; python_version > '3.11'
perceptron # required for isaac test
kaldi-native-fbank>=1.18.7 # required for fireredasr2 test
# Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
# Older versions are in conflict with terratorch requirements.
datasets>=3.3.0,<=3.6.0
openpyxl # required for perf comparison excel report
plotly # required for perf comparison html report
# Dependabot GHSA floors for transitive test deps (excl. dependabot.yml ignores).
# Skipped (no patched release yet): sqlitedict, accelerate, lmcache, nltk GHSA-8mgp-746c-j5xp.
urllib3>=2.7.0
cryptography>=50.0.0
pyjwt>=2.14.0
pyasn1>=0.6.4
python-multipart>=0.0.31
anyio>=4.14.2
httpx2>=2.12.0
werkzeug>=3.1.6
mako>=1.3.12
black>=26.3.1
msgpack>=1.2.1
lxml>=6.1.0
pyarrow>=23.0.1
orjson>=3.11.6
hpack>=4.2.0
h2>=4.4.1
pydantic-settings>=2.14.2
virtualenv>=20.36.1
fonttools>=4.60.2
pygments>=2.20.0
# ROCm-specific extras (not in CUDA cuda.in)
rapidfuzz
# torchgeo also pulled in `lightning` transitively; disabled for the same
# quarantine reason as terratorch above. Restore once the quarantine clears.
# torchgeo==0.7.0
multiprocess==0.70.16
requirements/test/rocm.txt
-1455
View File
File diff suppressed because it is too large Load Diff
setup.py
+18 -489
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import ctypes
import importlib.util
import json
import logging
@@ -19,7 +18,7 @@ from packaging.version import Version, parse
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
from setuptools_scm import get_version
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
from torch.utils.cpp_extension import CUDA_HOME
# Select vLLM release tags, excluding crate tags such as "proto-v0.3.0".
VLLM_GIT_DESCRIBE_COMMAND = [
@@ -60,7 +59,7 @@ def is_metadata_only_build() -> bool:
if sys.platform.startswith("darwin"):
logger.warning(
"This build only supports CUDA and ROCm; building without GPU extensions"
"This build only supports CUDA; building without GPU extensions"
)
VLLM_TARGET_DEVICE = "empty"
elif not (sys.platform.startswith("linux") or sys.platform.startswith("darwin")):
@@ -72,10 +71,7 @@ elif not (sys.platform.startswith("linux") or sys.platform.startswith("darwin"))
)
VLLM_TARGET_DEVICE = "empty"
elif sys.platform.startswith("linux") and os.getenv("VLLM_TARGET_DEVICE") is None:
if torch.version.hip is not None:
VLLM_TARGET_DEVICE = "rocm"
logger.info("Auto-detected ROCm")
elif torch.version.xpu is not None:
if torch.version.xpu is not None:
VLLM_TARGET_DEVICE = "xpu"
logger.info("Auto-detected XPU")
elif torch.version.cuda is not None:
@@ -84,6 +80,9 @@ elif sys.platform.startswith("linux") and os.getenv("VLLM_TARGET_DEVICE") is Non
else:
VLLM_TARGET_DEVICE = "empty"
if VLLM_TARGET_DEVICE == "rocm":
VLLM_TARGET_DEVICE = "empty"
def is_sccache_available() -> bool:
return which("sccache") is not None and not bool(
@@ -231,8 +230,6 @@ class cmake_build_ext(build_ext):
# Make sure we use the nvcc from CUDA_HOME
if _is_cuda() and CUDA_HOME is not None:
cmake_args += [f"-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc"]
elif _is_hip() and ROCM_HOME is not None:
cmake_args += [f"-DROCM_PATH={ROCM_HOME}"]
other_cmake_args = os.environ.get("CMAKE_ARGS")
if other_cmake_args:
@@ -322,7 +319,7 @@ class cmake_build_ext(build_ext):
os.makedirs(os.path.dirname(dst_file), exist_ok=True)
self.copy_file(file, dst_file)
if _is_cuda() or _is_hip():
if _is_cuda():
# copy vllm/third_party/triton_kernels/**/*.py from self.build_lib
# to current directory so that they can be included in the editable
# build
@@ -393,84 +390,21 @@ class precompiled_wheel_utils:
def fetch_metadata_for_variant(
commit: str,
variant: str | None,
*,
rocm: bool = False,
) -> tuple[list[dict], str]:
"""Fetches metadata for a specific variant of the precompiled wheel.
For non-ROCm, fetches vllm metadata.
For ROCm, discovers all first-level packages and combines their
metadata into a single list.
"""
"""Fetches metadata for a specific variant of the precompiled wheel."""
import json
from html.parser import HTMLParser
from urllib.request import urlopen
variant_dir = f"{variant}/" if variant is not None else ""
if not rocm:
# Keep original behavior
repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/"
meta_url = repo_url + "metadata.json"
print(f"Trying to fetch nightly build metadata from {meta_url}")
with urlopen(meta_url) as resp:
wheels = json.loads(resp.read().decode("utf-8"))
repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/"
meta_url = repo_url + "metadata.json"
print(f"Trying to fetch nightly build metadata from {meta_url}")
with urlopen(meta_url) as resp:
wheels = json.loads(resp.read().decode("utf-8"))
return wheels, repo_url
# ROCm: discover all packages under the variant directory.
repo_url = f"https://wheels.vllm.ai/rocm/{commit}/{variant_dir}"
class LinkParser(HTMLParser):
def __init__(self):
super().__init__()
self.links: list[str] = []
def handle_starttag(self, tag, attrs):
if tag == "a":
href = dict(attrs).get("href")
if href:
self.links.append(href)
with urlopen(repo_url) as resp:
parser = LinkParser()
parser.feed(resp.read().decode("utf-8"))
packages = [
href.rstrip("/")
for href in parser.links
if href.endswith("/")
and href not in ("../", "./", "/")
and "/" not in href.rstrip("/")
]
wheels: list[dict] = []
for package in packages:
meta_url = f"{repo_url}{package}/metadata.json"
print(f"Trying to fetch nightly build metadata from {meta_url}")
with urlopen(meta_url) as resp:
package_wheels = json.loads(resp.read().decode("utf-8"))
wheels.extend(package_wheels)
return wheels, repo_url
@staticmethod
def is_rocm_system() -> bool:
"""Detect ROCm without relying on torch (for build environment)."""
if os.getenv("ROCM_PATH"):
return True
if os.path.isdir("/opt/rocm"):
return True
if which("rocminfo") is not None:
return True
try:
import torch
return torch.version.hip is not None
except ImportError:
return False
@staticmethod
def detect_system_cuda_variant() -> str:
"""Auto-detect CUDA variant from torch, nvidia-smi, or env default."""
@@ -513,261 +447,6 @@ class precompiled_wheel_utils:
print(f"Detected CUDA {cuda_version}, using variant {variant}")
return variant
@staticmethod
def rocm_version_to_variant(rocm_version: str) -> str:
"""Convert a ROCm version string to a wheel variant from its major and
minor version, e.g. 7.2.3 -> rocm72, 10.0.0 -> rocm100.
"""
return "rocm" + "".join(rocm_version.split(".")[:2])
@staticmethod
def detect_system_rocm_variant() -> str | None:
"""Auto-detect the ROCm wheel variant from the installed ROCm stack."""
rocm_version = get_rocm_version()
if not rocm_version:
try:
import torch
rocm_version = torch.version.hip
except Exception:
pass
if not rocm_version:
return None
variant = precompiled_wheel_utils.rocm_version_to_variant(rocm_version)
print(f"Detected ROCm {rocm_version}, using variant {variant}")
return variant
@staticmethod
def fetch_available_rocm_variants(commit: str) -> list[str]:
"""List ROCm wheel variants published for a commit on wheels.vllm.ai."""
from urllib.request import urlopen
index_url = f"https://wheels.vllm.ai/rocm/{commit}/"
print(f"Fetching available ROCm variants from {index_url}")
try:
with urlopen(index_url) as resp:
html = resp.read().decode("utf-8")
except Exception as e:
logger.warning(
"Failed to fetch ROCm variant index for commit %s: %s", commit, e
)
return []
variants = sorted(set(re.findall(r"rocm\d+", html)))
print(f"Available ROCm variants for commit {commit}: {variants}")
return variants
@staticmethod
def resolve_rocm_wheel_variant(
commit: str, variant_override: str | None
) -> str | None:
"""Resolve a ROCm wheel variant for a commit from wheels.vllm.ai."""
env_variant = precompiled_wheel_utils.detect_system_rocm_variant()
available = precompiled_wheel_utils.fetch_available_rocm_variants(commit)
if variant_override is not None:
if env_variant and variant_override != env_variant:
logger.warning(
"VLLM_PRECOMPILED_WHEEL_VARIANT=%s does not match the "
"detected environment ROCm variant %s; refusing to use a "
"different ROCm patch wheel",
variant_override,
env_variant,
)
return None
if available and variant_override not in available:
logger.warning(
"Requested ROCm variant %s is not available for commit %s "
"(available: %s)",
variant_override,
commit,
available,
)
return None
return variant_override
if env_variant is None:
if available:
print(
"Could not detect ROCm version from the environment; "
f"available variants for commit {commit}: {available}"
)
return None
if env_variant in available:
return env_variant
logger.warning(
"Environment ROCm variant %s is not available for commit %s "
"(available: %s). Precompiled wheels may not be compatible.",
env_variant,
commit,
available,
)
return None
@staticmethod
def warn_if_rocm_torch_version_mismatch(
wheels: list[dict], repo_url: str, arch: str
) -> None:
"""Warn if installed torch differs from the custom ROCm build on
wheels.vllm.ai and suggest the correct install command."""
try:
installed = torch.__version__
except Exception:
return
def _wheel_version(pkg: str) -> str | None:
for w in wheels:
if w.get("package_name") == pkg and arch in w.get("platform_tag", ""):
v = w["version"]
if w.get("variant") and "+" not in v:
v = f"{v}+{w['variant']}"
return v
return None
expected = _wheel_version("torch")
if expected is None or installed == expected:
return
pkgs = f"torch=={expected}"
triton_ver = _wheel_version("triton")
if triton_ver:
pkgs += f" triton=={triton_ver}"
logger.warning(
"Installed PyTorch %s does not match the custom build %s "
"shipped with vLLM ROCm wheels. The ABI may differ from "
"official releases. If you hit extension load errors, "
"reinstall from the vLLM index:\n"
" pip install %s --extra-index-url %s\n"
" uv pip install %s --extra-index-url %s "
"--index-strategy unsafe-best-match",
installed,
expected,
pkgs,
repo_url,
pkgs,
repo_url,
)
@staticmethod
def find_local_rocm_wheel() -> str | None:
"""Search for a local vllm wheel in common locations."""
import glob
for pattern in ["/vllm-workspace/dist/vllm-*.whl", "./dist/vllm-*.whl"]:
wheels = glob.glob(pattern)
if wheels:
return sorted(wheels)[-1]
return None
@staticmethod
def fetch_wheel_from_pypi_index(index_url: str, package: str = "vllm") -> str:
"""Fetch the latest wheel URL from a PyPI-style simple index."""
import platform
from html.parser import HTMLParser
from urllib.parse import urljoin
from urllib.request import urlopen
arch = platform.machine()
class WheelLinkParser(HTMLParser):
def __init__(self):
super().__init__()
self.wheels = []
def handle_starttag(self, tag, attrs):
if tag == "a":
for name, value in attrs:
if name == "href" and value.endswith(".whl"):
self.wheels.append(value)
simple_url = f"{index_url.rstrip('/')}/{package}/"
print(f"Fetching wheel list from {simple_url}")
with urlopen(simple_url) as resp:
html = resp.read().decode("utf-8")
parser = WheelLinkParser()
parser.feed(html)
for wheel in reversed(parser.wheels):
if arch in wheel:
if wheel.startswith("http"):
return wheel
return urljoin(simple_url, wheel)
raise ValueError(f"No compatible wheel found for {arch} at {simple_url}")
@staticmethod
def determine_wheel_url_rocm() -> tuple[str, str | None]:
"""Determine the precompiled wheel for ROCm."""
# Search for local wheel first
local_wheel = precompiled_wheel_utils.find_local_rocm_wheel()
if local_wheel is not None:
print(f"Found local ROCm wheel: {local_wheel}")
return local_wheel, None
import platform
arch = platform.machine()
commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower()
if not commit or len(commit) != 40:
print(
f"VLLM_PRECOMPILED_WHEEL_COMMIT not valid: {commit}"
", trying to fetch base commit in main branch"
)
commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
variant = precompiled_wheel_utils.resolve_rocm_wheel_variant(
commit, os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None)
)
print(f"Using precompiled ROCm wheel commit {commit} with variant {variant}")
wheels, repo_url = None, None
if variant is not None:
try:
wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant(
commit, variant, rocm=True
)
precompiled_wheel_utils.warn_if_rocm_torch_version_mismatch(
wheels, repo_url, arch
)
except Exception as e:
logger.warning(
"Failed to fetch ROCm wheel metadata for variant %s: %s",
variant,
e,
)
if wheels is not None and repo_url is not None:
from urllib.parse import urljoin
for wheel in wheels:
if wheel.get("package_name") == "vllm" and arch in wheel.get(
"platform_tag", ""
):
print(f"Found precompiled wheel metadata: {wheel}")
if "path" not in wheel:
raise ValueError(f"Wheel metadata missing path: {wheel}")
wheel_url = urljoin(f"{repo_url}vllm/", wheel["path"])
download_filename = wheel.get("filename")
print(f"Using precompiled wheel URL: {wheel_url}")
return wheel_url, download_filename
logger.warning(
"No precompiled vllm wheel found for architecture %s "
"from repo %s. All available wheels: %s",
arch,
repo_url,
wheels,
)
# Fall back to AMD's PyPI index
index_url = os.getenv(
"VLLM_ROCM_WHEEL_INDEX", "https://pypi.amd.com/vllm-rocm/simple"
)
print(f"Fetching ROCm precompiled wheel from {index_url}")
wheel_url = precompiled_wheel_utils.fetch_wheel_from_pypi_index(index_url)
download_filename = wheel_url.split("/")[-1].split("#")[0]
print(f"Using ROCm precompiled wheel: {wheel_url}")
return wheel_url, download_filename
# Publication lags main by roughly 5-10 commits; anything much older is
# unlikely to still match the checkout's compiled sources.
WHEEL_SEARCH_DEPTH = 20
@@ -873,10 +552,6 @@ class precompiled_wheel_utils:
print(f"Using user-specified precompiled wheel location: {wheel_location}")
return wheel_location, None
else:
# ROCm: resolve wheels from wheels.vllm.ai with environment-matched variant
if precompiled_wheel_utils.is_rocm_system():
return precompiled_wheel_utils.determine_wheel_url_rocm()
import platform
arch = platform.machine()
@@ -998,8 +673,6 @@ class precompiled_wheel_utils:
"vllm/cumem_allocator.abi3.so",
"vllm/spinloop.abi3.so",
"vllm/fs_io_C.abi3.so",
# ROCm-specific libraries
"vllm/_rocm_C.abi3.so",
}
)
@@ -1134,12 +807,6 @@ def _is_cuda() -> bool:
return VLLM_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu()
def _is_hip() -> bool:
return (
VLLM_TARGET_DEVICE == "cuda" or VLLM_TARGET_DEVICE == "rocm"
) and torch.version.hip is not None
def _is_tpu() -> bool:
return VLLM_TARGET_DEVICE == "tpu"
@@ -1149,41 +816,7 @@ def _is_xpu() -> bool:
def _build_custom_ops() -> bool:
return _is_cuda() or _is_hip()
def get_rocm_version():
# Get the Rocm version from the ROCM_HOME/bin/librocm-core.so
# see https://github.com/ROCm/rocm-core/blob/d11f5c20d500f729c393680a01fa902ebf92094b/rocm_version.cpp#L21
try:
if ROCM_HOME is None:
return None
librocm_core_file = Path(ROCM_HOME) / "lib" / "librocm-core.so"
if not librocm_core_file.is_file():
return None
librocm_core = ctypes.CDLL(librocm_core_file)
VerErrors = ctypes.c_uint32
get_rocm_core_version = librocm_core.getROCmVersion
get_rocm_core_version.restype = VerErrors
get_rocm_core_version.argtypes = [
ctypes.POINTER(ctypes.c_uint32),
ctypes.POINTER(ctypes.c_uint32),
ctypes.POINTER(ctypes.c_uint32),
]
major = ctypes.c_uint32()
minor = ctypes.c_uint32()
patch = ctypes.c_uint32()
if (
get_rocm_core_version(
ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)
)
== 0
):
return f"{major.value}.{minor.value}.{patch.value}"
return None
except Exception:
return None
return _is_cuda()
def get_nvcc_cuda_version() -> Version:
@@ -1235,13 +868,6 @@ def get_vllm_version() -> str:
# skip this for source tarball, required for pypi
if "sdist" not in sys.argv:
version += f"{sep}cu{cuda_version_str}"
elif _is_hip():
# Get the Rocm Version
rocm_version = get_rocm_version() or torch.version.hip
if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION:
version += sep + precompiled_wheel_utils.rocm_version_to_variant(
rocm_version
)
elif _is_tpu():
version += f"{sep}tpu"
elif _is_xpu():
@@ -1252,74 +878,6 @@ def get_vllm_version() -> str:
return version
def _check_requirements_preinstalled(requirements: list[str]) -> None:
"""`python setup.py develop` satisfies missing `install_requires` via
setuptools' legacy easy_install path instead of pip's wheel-aware
resolver. That path has repeatedly tried (and failed) to compile old
sdists of fast-moving C-extension packages such as aiohttp from source
against the running Python's headers, e.g.
https://github.com/vllm-project/vllm/issues/34073. Fail fast with an
actionable message instead of a cryptic compiler error.
Skipped under `--no-deps`, which makes setuptools skip `install_requires`
processing entirely, so the easy_install path this guards against is
unreachable anyway.
"""
from importlib.metadata import PackageNotFoundError
from importlib.metadata import requires as installed_requires
from importlib.metadata import version as installed_version
from packaging.requirements import Requirement
missing = []
checked = set()
def _check(req_str: str, extra: str | None = None) -> None:
# Strip inline "# comment" suffixes (as pip does when reading
# requirements files); packaging.Requirement cannot parse them.
req_str = re.split(r"\s+#", req_str, maxsplit=1)[0].strip()
try:
req = Requirement(req_str)
except Exception:
return
# `extra` is only set when recursing into an extra's own deps
# below; it selects which `; extra == "..."` markers apply.
if req.marker is not None and not req.marker.evaluate(
{"extra": extra} if extra else None
):
return
if (req.name, extra) in checked:
return
checked.add((req.name, extra))
try:
installed = installed_version(req.name)
except PackageNotFoundError:
missing.append(req_str)
return
if req.specifier and not req.specifier.contains(installed, prereleases=True):
missing.append(f"{req_str} (found {req.name}=={installed})")
return
# A requested extra (e.g. "fastapi[standard]") pulls in packages
# that easy_install resolves too, so check those as well.
for extra_name in req.extras:
for dep in installed_requires(req.name) or []:
_check(dep, extra=extra_name)
for req_str in requirements:
_check(req_str)
if missing:
raise RuntimeError(
"The following dependencies are missing or outdated:\n "
+ "\n ".join(missing)
+ "\n\n`python setup.py develop` cannot reliably install these "
"itself (see https://github.com/vllm-project/vllm/issues/34073). "
"Install them with pip first, e.g.:\n"
" pip install -r requirements/rocm.txt\n"
"then re-run `python setup.py develop`."
)
def get_requirements() -> list[str]:
"""Get Python package dependencies from requirements.txt."""
requirements_dir = ROOT_DIR / "requirements"
@@ -1362,41 +920,18 @@ def get_requirements() -> list[str]:
req = req.replace("humming-kernels[cu13]", "humming-kernels[cu12]")
modified_requirements.append(req)
requirements = modified_requirements
elif _is_hip():
requirements = _read_requirements("rocm.txt")
if "develop" in sys.argv[1:] and "--no-deps" not in sys.argv[1:]:
_check_requirements_preinstalled(requirements)
requirements += get_rocm_device_requirements()
elif _is_tpu():
requirements = _read_requirements("tpu.txt")
elif _is_xpu():
requirements = _read_requirements("xpu.txt")
else:
raise ValueError("Unsupported platform, please use CUDA or ROCm.")
raise ValueError("Unsupported platform, please use CUDA.")
return requirements
def get_rocm_device_requirements() -> list[str]:
"""GPU kernel packages for TheRock builds.
TheRock ships GPU kernels as separate per-arch wheels selected through
``[device-<arch>]`` extras on torch, torchvision and rocm. Require them for
every arch this wheel was compiled for, so ``pip install vllm`` works on
any supported GPU.
"""
if not _is_hip() or importlib.util.find_spec("rocm_sdk") is None:
return []
arches = [a for a in os.getenv("PYTORCH_ROCM_ARCH", "").split(";") if a]
return [
f"{pkg}[device-{arch}]"
for pkg in ("torch", "torchvision", "rocm")
for arch in arches
]
ext_modules = []
if _is_cuda() or _is_hip():
if _is_cuda():
ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
# Optional since this doesn't get built (produce an .so file). This is just
# copying the relevant .py files from the source repository.
@@ -1406,9 +941,6 @@ if not _is_xpu():
ext_modules.append(CMakeExtension(name="vllm.spinloop"))
ext_modules.append(CMakeExtension(name="vllm.fs_io_C"))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
if _is_cuda():
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C"))
if USE_PRECOMPILED_EXTENSIONS or (
@@ -1451,11 +983,8 @@ if _is_cuda():
ext_modules.append(CMakeExtension(name="vllm.tml_fa4", optional=True))
if _build_custom_ops():
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
package_data = {
"vllm": [
tests/basic_correctness/memory/sleep_mode/__init__.py
View File
tests/basic_correctness/models/__init__.py
View File
tests/benchmarks/sweep/__init__.py
View File
tests/compile/correctness_e2e/__init__.py
View File
tests/compile/dynamic_shapes/__init__.py
View File
tests/compile/fusions_e2e/__init__.py
View File
tests/compile/fusions_e2e/common.py
-144
View File
@@ -1,144 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
from collections.abc import Callable, Iterable
from typing import Any, NamedTuple
import pytest
import regex as re
from vllm.platforms import current_platform
from vllm.v1.attention.backends.registry import AttentionBackendEnum
class Matches(NamedTuple):
# simple pointwise
aiter_rms_quant_fusion: int = 0
rms_quant_fusion: int = 0
act_quant_fusion: int = 0
norm_rope_fusion: int = 0
attn_quant_fusion: int = 0
# distributed
ar_rms_fusion: int = 0
aiter_ar_rms_fusion: int = 0
class ModelFusionInfo(NamedTuple):
model_name: str
matches: Callable[[int], Matches]
"""Given number of hidden layers, produces the matches object"""
model_kwargs: dict[str, Any] = {}
hf_overrides: Callable[[int], dict] = lambda n: {"num_hidden_layers": n}
class AttentionBackendCase(NamedTuple):
backend: AttentionBackendEnum
model_kwargs: dict[str, Any] = {}
"""Additional args required for attn+quant fusion"""
is_blackwell = lambda: current_platform.is_device_capability_family(100)
"""Are we running on Blackwell, a lot of tests depend on it"""
def nvfp4_kernel_exposes_input_quant_key() -> bool:
"""Check if the NVFP4 kernel selected on this platform exposes input_quant_key.
FlashInferCuteDslNvFp4LinearKernel does not expose input_quant_key() due to
layout incompatibility between the manual fusion kernel (silu_and_mul_nvfp4_quant)
output format and the cutedsl backend's expected input format.
FlashInferCutlassNvFp4LinearKernel does expose input_quant_key() and supports
manual fusion.
"""
if not current_platform.is_cuda():
return False
try:
from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
FlashInferCuteDslNvFp4LinearKernel,
)
is_supported, _ = FlashInferCuteDslNvFp4LinearKernel.is_supported()
if is_supported:
return False
except ImportError:
pass
try:
from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
FlashInferCutlassNvFp4LinearKernel,
)
is_supported, _ = FlashInferCutlassNvFp4LinearKernel.is_supported()
if is_supported:
return True
except ImportError:
pass
# Fallback: assume no manual fusion support
return False
def custom_ops_combos(*custom_ops: str) -> Iterable[str]:
"""Generate all combinations of custom ops for parametrization."""
custom_ops_lists = [[f"-{op}", f"+{op}"] for op in custom_ops]
for op_list in itertools.product(*custom_ops_lists):
yield ",".join(op_list)
# Quick inline validation
assert list(custom_ops_combos("silu_and_mul")) == ["-silu_and_mul", "+silu_and_mul"]
assert list(custom_ops_combos("quant_fp8", "rms_norm")) == [
"-quant_fp8,-rms_norm",
"-quant_fp8,+rms_norm",
"+quant_fp8,-rms_norm",
"+quant_fp8,+rms_norm",
]
def has_cuda_graph_wrapper_metadata() -> bool:
from importlib import import_module
try:
module = import_module("torch._inductor.utils")
module.CUDAGraphWrapperMetadata # noqa B018
except AttributeError:
return False
return True
INDUCTOR_GRAPH_PARTITION = [
pytest.param(
True,
marks=pytest.mark.skipif(
not has_cuda_graph_wrapper_metadata(),
reason="torch version does not support Inductor partition",
),
id="inductor_partition",
),
pytest.param(False, id="dynamo_partition"),
]
FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {
"aiter_rms_quant_fusion": re.compile(
r"RocmAiterRMSNormQuantFusionPass Replaced (\d+) patterns"
),
"rms_quant_fusion": re.compile(r"rms_quant_fusion.py:\d+] Replaced (\d+) patterns"),
"act_quant_fusion": re.compile(r"act_quant_fusion.py:\d+] Replaced (\d+) patterns"),
# The combined pass can consume these sites before the standalone pass.
"norm_rope_fusion": re.compile(
r"(?:qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on|"
r"qk_norm_rope_kvcache_fusion.py:\d+] "
r"QK-Norm\+RoPE(?:/MRoPE)?\+KVCache fusion: replaced) (\d+)"
),
"attn_quant_fusion": re.compile(
r"attn_quant_fusion.py:\d+] Fused quant onto (\d+) attention nodes"
),
"ar_rms_fusion": re.compile(
r"allreduce_rms_fusion.py:\d+] Replaced (\d+) patterns"
),
"aiter_ar_rms_fusion": re.compile(
r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns"
),
}
tests/compile/fusions_e2e/conftest.py
-252
View File
@@ -1,252 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
from collections import defaultdict
import pytest
import regex as re
from vllm import SamplingParams
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
from .common import FUSION_LOG_PATTERNS, AttentionBackendCase, Matches
def run_model(
vllm_runner,
compile_config: CompilationMode | CompilationConfig,
model: str,
**model_kwargs,
):
"""Run a model with the given compilation config for E2E fusion tests."""
compilation_config = (
compile_config
if isinstance(compile_config, CompilationConfig)
else CompilationConfig(mode=compile_config)
)
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
sampling_params = SamplingParams(temperature=0)
# Allow override from model_kwargs
model_kwargs = {"tensor_parallel_size": 1, **model_kwargs}
model_kwargs = {"disable_custom_all_reduce": True, **model_kwargs}
# No cudagraphs by default
if compilation_config.cudagraph_mode is None:
compilation_config.cudagraph_mode = CUDAGraphMode.NONE
# trust_remote_code/enable_chunked_prefill/block_size/kernel_config are
# pinned to LLM()'s own defaults since VllmRunner overrides them for
# other test needs (e.g. it forces block_size=16, which some attention
# backends like FLASHINFER_MLA reject, and disables JIT warmup).
model_kwargs.setdefault("kernel_config", {"enable_jit_warmup": True})
with vllm_runner(
model,
compilation_config=compilation_config,
trust_remote_code=False,
enable_chunked_prefill=None,
block_size=None,
**model_kwargs,
) as vllm_model:
llm = vllm_model.llm
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
# Get the compile ranges endpoints after vllm config post init
# in order to compute compile ranges correctly
compilation_config.compile_ranges_endpoints = (
llm.llm_engine.vllm_config.compilation_config.compile_ranges_endpoints
)
# Fetch match table from each worker via RPC and sum across workers.
worker_tables: list[dict[str, int]] = llm.llm_engine.engine_core.collective_rpc(
"get_compilation_match_table"
)
combined: defaultdict[str, int] = defaultdict(int)
for table in worker_tables:
for k, v in table.items():
combined[k] += v
return dict(combined)
@pytest.fixture
def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn, vllm_runner):
def run(
model_name: str,
matches: Matches,
model_kwargs: dict,
attn_backend: AttentionBackendCase,
compilation_config: dict,
matches_check: list[str],
use_deepgemm: bool = False,
use_aiter: bool = False,
tp_size: int = 1,
):
monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1" if use_aiter else "0")
from vllm._aiter_ops import rocm_aiter_ops
rocm_aiter_ops.refresh_env_variables()
# Filter here to reduce code duplication
backend_name = attn_backend.backend.name.lower()
requires_mla = "deepseek" in model_name.lower()
is_mla = "mla" in backend_name
if requires_mla != is_mla:
pytest.skip(
f"Incompatible model '{model_name}' and "
f"attention backend '{attn_backend.backend.name}'"
)
if backend_name == "rocm_attn" and model_name == "openai/gpt-oss-20b":
pytest.skip(
"ROCM_ATTN does not support attention sinks (required by gpt-oss-20b)"
)
if attn_backend.backend.name == "FLASHINFER":
from vllm.utils.flashinfer import supports_trtllm_attention
if not supports_trtllm_attention():
matches = matches._replace(attn_quant_fusion=0)
# Disable, compile cache to make sure custom passes run.
# Otherwise, we can't verify fusion happened through the logs.
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
# To capture subprocess logs, we need to know whether spawn or fork is used.
# Force spawn as it is more general.
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
model_kwargs = {**attn_backend.model_kwargs, **model_kwargs}
model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
model_kwargs["tensor_parallel_size"] = tp_size
# Cap warmup memory: tests use small max_model_len (1024) but the
# engine default max_num_batched_tokens is 16384. Warming up large
# models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM.
model_kwargs.setdefault("max_num_batched_tokens", 8192)
# Always compile the full graph instead of piecewise
if not compilation_config["use_inductor_graph_partition"]:
compilation_config["splitting_ops"] = []
full_compilation_config = CompilationConfig(
cudagraph_mode=CUDAGraphMode.NONE,
mode=CompilationMode.VLLM_COMPILE,
inductor_compile_config={"force_disable_caches": True},
**compilation_config,
)
with caplog_mp_spawn(logging.DEBUG) as log_holder:
match_table = run_model(
vllm_runner, full_compilation_config, model_name, **model_kwargs
)
num_compile_ranges = len(full_compilation_config.get_compile_ranges())
assert num_compile_ranges in [1, 2, 3]
print(f"Compile ranges: {full_compilation_config.get_compile_ranges()}")
print("Fusion results:")
# Iterate through all so printing happens before asserting
log_matches_dict = {}
for match_name, pattern in FUSION_LOG_PATTERNS.items():
log_matches_dict[match_name] = list(pattern.findall(log_holder.text))
print(f"- {match_name}={','.join(log_matches_dict[match_name])}")
# Now check the matches
for match_name in matches_check:
log_matches = list(int(ms) for ms in log_matches_dict[match_name])
# AR+RMS skips the largest range.
if match_name == "ar_rms_fusion" and num_compile_ranges >= 2:
num_ranges_activated = num_compile_ranges - 1
else:
num_ranges_activated = num_compile_ranges
# TODO: Remove log counting in unit tests
# once all matchers implement VllmFusionPatternMatcherPass
n_expected = tp_size * num_ranges_activated
if match_name not in (
"attn_quant_fusion",
"act_quant_fusion",
"norm_rope_fusion",
):
assert len(log_matches) == n_expected, (
f"Could not find {n_expected} {match_name} "
f"(found {len(log_matches)}) in:\n {log_holder.text}"
)
expected_matches = getattr(matches, match_name)
if match_name == "norm_rope_fusion":
# Opaque LayerName lets the combined pass consume short-range
# sites before the standalone pass; count either owner.
assert len(log_matches) >= n_expected
assert all(m in (0, expected_matches) for m in log_matches)
assert sum(log_matches) == expected_matches * n_expected, (
f"{match_name} expected {expected_matches * n_expected} "
f"sites across both passes, found: {log_matches}"
)
elif match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:
# AR+rms+quant takes precedence over rms+quant if activated.
# That means we get full matching where ar+rms+quant was not
# activated, and less where it was (only the smallest range).
assert sum(m == expected_matches for m in log_matches) == tp_size * (
num_ranges_activated - 1
), "Expecting full rms+quant fusion where ar+rms+quant not activated"
assert all(
expected_matches - matches.ar_rms_fusion <= m <= expected_matches
for m in log_matches
), (
f"Expecting at least {expected_matches - matches.ar_rms_fusion} "
f"where ar+rms+quant was activated"
)
elif match_name == "act_quant_fusion":
actual_match = match_table.get("activation_quant_fusion_pass", 0)
assert actual_match == expected_matches * n_expected, (
f"Could not find {expected_matches * n_expected} "
f"{match_name} (found {actual_match})."
)
elif match_name == "attn_quant_fusion":
actual_match = match_table.get(
"attn_quant_fusion", 0
) + match_table.get("mla_attn_quant_fusion", 0)
assert actual_match == expected_matches * n_expected, (
f"Could not find {expected_matches * n_expected} "
f"{match_name} (found {actual_match})."
)
else:
expected_matches_list = [expected_matches] * n_expected
assert sorted(log_matches) == expected_matches_list, (
f"{match_name} expected: {expected_matches_list}, "
f"found: {sorted(log_matches)}"
)
if match_name == "ar_rms_fusion" and num_compile_ranges >= 2:
log_matches = re.findall(
r"pass_manager.py:\d+] Skipping "
r".*AllReduceFusionPass.* with compile range",
log_holder.text,
)
n_expected = tp_size * (num_compile_ranges - num_ranges_activated)
assert len(log_matches) == n_expected, (
f'Could not find {n_expected} "Skipping AllReduceFusionPass" '
f"(found {len(log_matches)}) in:\n {log_holder.text}"
)
return run
tests/compile/passes/distributed/__init__.py
View File
tests/config/test_mp_reducer.py
-51
View File
@@ -1,51 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
from unittest.mock import patch
from vllm.config import VllmConfig
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.v1.engine.async_llm import AsyncLLM
def test_mp_reducer():
"""Test that _reduce_config reducer is registered when AsyncLLM is instantiated
without transformers_modules. This is a regression test for
https://github.com/vllm-project/vllm/pull/18640.
"""
# Ensure transformers_modules is not in sys.modules
if "transformers_modules" in sys.modules:
del sys.modules["transformers_modules"]
with patch("multiprocessing.reducer.register") as mock_register:
engine_args = AsyncEngineArgs(
model="facebook/opt-125m",
max_model_len=32,
gpu_memory_utilization=0.1,
disable_log_stats=True,
)
async_llm = AsyncLLM.from_engine_args(
engine_args,
start_engine_loop=False,
)
assert mock_register.called, (
"multiprocessing.reducer.register should have been called"
)
vllm_config_registered = False
for call_args in mock_register.call_args_list:
# Verify that a reducer for VllmConfig was registered
if len(call_args[0]) >= 2 and call_args[0][0] == VllmConfig:
vllm_config_registered = True
reducer_func = call_args[0][1]
assert callable(reducer_func), "Reducer function should be callable"
break
assert vllm_config_registered, (
"VllmConfig should have been registered to multiprocessing.reducer"
)
async_llm.shutdown()
tests/detokenizer/test_disable_detokenization.py
-31
View File
@@ -1,31 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.entrypoints.llm import LLM
from vllm.sampling_params import SamplingParams
@pytest.mark.parametrize("model", ["distilbert/distilgpt2"])
def test_computed_prefix_blocks(model: str):
# This test checks if the engine generates completions both with and
# without optional detokenization, that detokenization includes text
# and no-detokenization doesn't, and that both completions have the same
# token_ids.
prompt = (
"You are a helpful assistant. How do I build a car from cardboard and "
"paper clips? Is there an easy to follow video tutorial available "
"online for free?"
)
llm = LLM(model=model)
sampling_params = SamplingParams(max_tokens=10, temperature=0.0, detokenize=False)
outputs_no_detokenization = llm.generate(prompt, sampling_params)[0].outputs[0]
sampling_params.detokenize = True
outputs_with_detokenization = llm.generate(prompt, sampling_params)[0].outputs[0]
assert outputs_no_detokenization.text == ""
assert outputs_with_detokenization.text != ""
assert outputs_no_detokenization.token_ids == outputs_with_detokenization.token_ids
tests/distributed/test_ray_v2_executor_e2e.py
-205
View File
@@ -1,205 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Orchestration-level integration tests for RayExecutorV2."""
import gc
import os
import pathlib
import pytest
import ray
pytestmark = pytest.mark.usefixtures("enable_ray_v2_backend")
MODEL = "facebook/opt-125m"
def _get_env_var(worker, name):
return os.environ.get(name)
def _ray_init():
"""Start Ray with the project root on workers' PYTHONPATH.
Without this, workers cannot unpickle actor classes defined in the
``tests`` package, causing FunctionActorManager to fall back to
TemporaryActor which drops async method signatures."""
project_root = str(pathlib.Path(__file__).resolve().parents[2])
ray.init(
ignore_reinit_error=True,
runtime_env={"env_vars": {"PYTHONPATH": project_root}},
)
@pytest.fixture
def ray_init():
_ray_init()
class _AsyncLLMActor:
def start(self, pg, bundle_indices=None, ray_runtime_env=None):
# Needed so collective_rpc can pickle _get_env_var over the
# AsyncLLM -> EngineCore ZMQ boundary.
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
if bundle_indices is not None:
os.environ["VLLM_RAY_BUNDLE_INDICES"] = bundle_indices
else:
os.environ.pop("VLLM_RAY_BUNDLE_INDICES", None)
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.executor.abstract import Executor
engine_args = AsyncEngineArgs(
model=MODEL,
tensor_parallel_size=2,
distributed_executor_backend="ray",
enforce_eager=True,
max_model_len=256,
gpu_memory_utilization=0.8,
)
vllm_config = engine_args.create_engine_config()
vllm_config.parallel_config.placement_group = pg
if ray_runtime_env is not None:
vllm_config.parallel_config.ray_runtime_env = ray_runtime_env
executor_class = Executor.get_class(vllm_config)
self.engine = AsyncLLM(
vllm_config=vllm_config,
executor_class=executor_class,
log_stats=False,
log_requests=False,
)
async def generate(self, prompt):
from vllm.sampling_params import SamplingParams
params = SamplingParams(max_tokens=16)
result = None
async for output in self.engine.generate(
prompt, params, request_id="test_request_id"
):
result = output
assert result is not None
return result.outputs[0].text
async def generate_and_get_worker_envs(self, prompt, env_names):
from vllm.sampling_params import SamplingParams
params = SamplingParams(max_tokens=16)
result = None
async for output in self.engine.generate(
prompt, params, request_id="test_request_id"
):
result = output
assert result is not None
text = result.outputs[0].text
env_results = {}
for name in env_names:
vals = await self.engine.collective_rpc(
_get_env_var, timeout=10, args=(name,)
)
env_results[name] = vals
return text, env_results
def shutdown(self):
if engine := getattr(self, "engine", None):
engine.shutdown()
del self.engine
gc.collect()
AsyncLLMActor = ray.remote(num_cpus=0, max_concurrency=1)(_AsyncLLMActor)
def test_multi_replicas(ray_init):
pg1 = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
pg2 = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
ray.get([pg1.ready(), pg2.ready()])
actor1 = AsyncLLMActor.remote()
actor2 = AsyncLLMActor.remote()
ray.get(actor1.start.remote(pg1))
ray.get(actor2.start.remote(pg2))
out1, out2 = ray.get(
[
actor1.generate.remote("Hello world"),
actor2.generate.remote("Hello world"),
]
)
assert len(out1) > 0
assert len(out2) > 0
def test_multi_replicas_with_bundle_indices(ray_init):
pg = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 4, strategy="PACK")
ray.get(pg.ready())
actor1 = AsyncLLMActor.remote()
actor2 = AsyncLLMActor.remote()
ray.get(actor1.start.remote(pg, bundle_indices="2,1"))
ray.get(actor2.start.remote(pg, bundle_indices="0,3"))
out1, out2 = ray.get(
[
actor1.generate.remote("Hello world"),
actor2.generate.remote("Hello world"),
]
)
assert len(out1) > 0
assert len(out2) > 0
def test_env_var_and_runtime_env_propagation():
"""Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env
propagate to RayWorkerProc actors.
"""
sentinel_vars = {
"NCCL_DEBUG": "INFO",
"HF_TOKEN": "test_sentinel_token",
}
for k, v in sentinel_vars.items():
os.environ[k] = v
try:
# Called directly (not via the ray_init fixture) because sentinel
# env vars must be in os.environ before ray.init() so that Ray
# worker processes inherit them.
_ray_init()
pg = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
ray.get(pg.ready())
# Include the project root so that RayWorkerProc actors can
# unpickle _get_env_var.
project_root = str(pathlib.Path(__file__).resolve().parents[2])
ray_runtime_env = {
"env_vars": {
"RAY_RUNTIME_ENV_TEST": "ray_runtime_env",
"PYTHONPATH": project_root,
},
}
actor = AsyncLLMActor.remote()
ray.get(actor.start.remote(pg, ray_runtime_env=ray_runtime_env))
all_env_names = list(sentinel_vars) + ["RAY_RUNTIME_ENV_TEST"]
text, env_results = ray.get(
actor.generate_and_get_worker_envs.remote("Hello world", all_env_names)
)
assert len(text) > 0
for name, expected in sentinel_vars.items():
for val in env_results[name]:
assert val == expected
for val in env_results["RAY_RUNTIME_ENV_TEST"]:
assert val == "ray_runtime_env"
finally:
for k in sentinel_vars:
os.environ.pop(k, None)
tests/entrypoints/launchers/api_server/test_multi_api_servers.py
-175
View File
@@ -1,175 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import os
import openai # use the official client for correctness check
import pytest
import pytest_asyncio
from tests.utils import RemoteOpenAIServer
from tests.v1.utils import check_request_balancing
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
DP_SIZE = os.getenv("DP_SIZE", "1")
@pytest.fixture(scope="module")
def default_server_args():
return [
# use half precision for speed and memory savings in CI environment
"--dtype",
"bfloat16",
"--max-model-len",
"2048",
"--max-num-seqs",
"128",
"--enforce-eager",
"--api-server-count",
"4",
"--data_parallel_size",
DP_SIZE,
]
@pytest.fixture(scope="module")
def server(default_server_args):
with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client(server):
async with server.get_async_client() as async_client:
yield async_client
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model_name",
[MODEL_NAME],
)
async def test_single_completion(
client: openai.AsyncOpenAI, server: RemoteOpenAIServer, model_name: str
) -> None:
async def make_request():
completion = await client.completions.create(
model=model_name, prompt="Hello, my name is", max_tokens=10, temperature=1.0
)
assert completion.id is not None
assert completion.choices is not None and len(completion.choices) == 1
choice = completion.choices[0]
# The exact number of tokens can vary slightly with temperature=1.0,
# so we check for a reasonable minimum length.
assert len(choice.text) >= 1
# Finish reason might not always be 'length' if the model finishes early
# or due to other reasons, especially with high temperature.
# So, we'll accept 'length' or 'stop'.
assert choice.finish_reason in ("length", "stop")
# Token counts can also vary, so we check they are positive.
assert completion.usage.completion_tokens > 0
assert completion.usage.prompt_tokens > 0
assert completion.usage.total_tokens > 0
return completion
# Test single request
result = await make_request()
assert result is not None
await asyncio.sleep(0.5)
# Send two bursts of requests
num_requests = 100
tasks = [make_request() for _ in range(num_requests)]
results = await asyncio.gather(*tasks)
assert len(results) == num_requests
assert all(completion is not None for completion in results)
await asyncio.sleep(0.5)
tasks = [make_request() for _ in range(num_requests)]
results = await asyncio.gather(*tasks)
assert len(results) == num_requests
assert all(completion is not None for completion in results)
# Check request balancing via Prometheus metrics if DP_SIZE > 1
check_request_balancing(server, int(DP_SIZE))
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model_name",
[MODEL_NAME],
)
async def test_completion_streaming(
client: openai.AsyncOpenAI, server: RemoteOpenAIServer, model_name: str
) -> None:
prompt = "What is an LLM?"
async def make_streaming_request():
# Perform a non-streaming request to get the expected full output
single_completion = await client.completions.create(
model=model_name,
prompt=prompt,
max_tokens=5,
temperature=0.0,
)
single_output = single_completion.choices[0].text
# Perform the streaming request
stream = await client.completions.create(
model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
)
chunks: list[str] = []
finish_reason_count = 0
last_chunk = None
async for chunk in stream:
chunks.append(chunk.choices[0].text)
if chunk.choices[0].finish_reason is not None:
finish_reason_count += 1
last_chunk = chunk # Keep track of the last chunk
# finish reason should only return in the last block for OpenAI API
assert finish_reason_count == 1, "Finish reason should appear exactly once."
assert last_chunk is not None, "Stream should have yielded at least one chunk."
assert last_chunk.choices[0].finish_reason == "length", (
"Finish reason should be 'length'."
)
# Check that the combined text matches the non-streamed version.
assert "".join(chunks) == single_output, (
"Streamed output should match non-streamed output."
)
return True # Indicate success for this request
# Test single request
result = await make_streaming_request()
assert result is not None
await asyncio.sleep(0.5)
# Send two bursts of requests
num_requests = 100
tasks = [make_streaming_request() for _ in range(num_requests)]
results = await asyncio.gather(*tasks)
assert len(results) == num_requests, (
f"Expected {num_requests} results, got {len(results)}"
)
assert all(results), "Not all streaming requests completed successfully."
await asyncio.sleep(0.5)
tasks = [make_streaming_request() for _ in range(num_requests)]
results = await asyncio.gather(*tasks)
assert len(results) == num_requests, (
f"Expected {num_requests} results, got {len(results)}"
)
assert all(results), "Not all streaming requests completed successfully."
# Check request balancing via Prometheus metrics if DP_SIZE > 1
check_request_balancing(server, int(DP_SIZE))
tests/entrypoints/multimodal/llm/__init__.py
View File
tests/entrypoints/multimodal/llm/test_chat.py
-43
View File
@@ -1,43 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
import pytest
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
@pytest.fixture(scope="function")
def vision_llm(vllm_runner):
with vllm_runner(
"microsoft/Phi-3.5-vision-instruct",
max_model_len=4096,
max_num_seqs=5,
enforce_eager=True,
trust_remote_code=True,
limit_mm_per_prompt={"image": 2},
seed=0,
) as runner:
# pytest caches yielded fixtures until after teardown, so use a proxy to
# avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory.
yield weakref.proxy(runner.llm)
@pytest.mark.parametrize(
"image_urls", [[TEST_IMAGE_ASSETS[0], TEST_IMAGE_ASSETS[1]]], indirect=True
)
def test_chat_multi_image(vision_llm, image_urls: list[str]):
messages = [
{
"role": "user",
"content": [
*(
{"type": "image_url", "image_url": {"url": image_url}}
for image_url in image_urls
),
{"type": "text", "text": "What's in this image?"},
],
}
]
outputs = vision_llm.chat(messages)
assert len(outputs) >= 0
tests/entrypoints/openai/chat_completion/test_chat_error.py
-643
View File
@@ -1,643 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from dataclasses import dataclass, field
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from vllm.config.multimodal import MultiModalConfig
from vllm.entrypoints.openai.chat_completion.protocol import (
BatchChatCompletionRequest,
ChatCompletionRequest,
)
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.models.protocol import BaseModelPath
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.exceptions import GenerationError, VLLMValidationError
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.renderers.hf import HfRenderer
from vllm.renderers.online_renderer import OnlineRenderer
from vllm.tokenizers.registry import cached_tokenizer_from_config
from vllm.v1.engine.async_llm import AsyncLLM
MODEL_NAME = "openai-community/gpt2"
MODEL_NAME_SHORT = "gpt2"
BASE_MODEL_PATHS = [
BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
BaseModelPath(name=MODEL_NAME_SHORT, model_path=MODEL_NAME_SHORT),
]
@dataclass
class MockHFConfig:
model_type: str = "any"
@dataclass
class MockModelConfig:
task = "generate"
runner_type = "generate"
model = MODEL_NAME
tokenizer = MODEL_NAME
trust_remote_code = False
tokenizer_mode = "auto"
max_model_len = 100
revision = None
code_revision = None
tokenizer_revision = None
multimodal_config = MultiModalConfig()
hf_config = MockHFConfig()
hf_text_config = MockHFConfig()
logits_processors: list[str] | None = None
diff_sampling_param: dict | None = None
allowed_local_media_path: str = ""
allowed_media_domains: list[str] | None = None
encoder_config = None
generation_config: str = "auto"
media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
skip_tokenizer_init: bool = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
supports_multimodal_inputs: bool = False
renderer_num_workers: int = 1
def get_diff_sampling_param(self):
return self.diff_sampling_param or {}
@dataclass
class MockParallelConfig:
_api_process_rank: int = 0
@dataclass
class MockVllmConfig:
model_config: MockModelConfig
parallel_config: MockParallelConfig
def _build_renderer(model_config: MockModelConfig):
return HfRenderer(
MockVllmConfig(model_config, parallel_config=MockParallelConfig()),
cached_tokenizer_from_config(model_config),
)
def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
models = OpenAIServingModels(
engine_client=engine,
base_model_paths=BASE_MODEL_PATHS,
)
online_renderer = OnlineRenderer(
model_config=engine.model_config,
renderer=engine.renderer,
request_logger=None,
chat_template=None,
chat_template_content_format="auto",
)
serving_chat = OpenAIServingChat(
engine,
models,
response_role="assistant",
online_renderer=online_renderer,
request_logger=None,
chat_template=None,
chat_template_content_format="auto",
)
async def _fake_preprocess_chat(*args, **kwargs):
# return conversation, engine_inputs
return (
[{"role": "user", "content": "Test"}],
[{"prompt_token_ids": [1, 2, 3]}],
)
serving_chat.online_renderer.preprocess_chat = AsyncMock(
side_effect=_fake_preprocess_chat
)
return serving_chat
@pytest.mark.asyncio
async def test_chat_error_non_stream():
"""Test finish_reason='error' returns 500 InternalServerError (non-streaming)."""
mock_engine = MagicMock(spec=AsyncLLM)
mock_engine.errored = False
mock_engine.model_config = MockModelConfig()
mock_engine.input_processor = MagicMock()
mock_engine.renderer = _build_renderer(mock_engine.model_config)
serving_chat = _build_serving_chat(mock_engine)
completion_output = CompletionOutput(
index=0,
text="",
token_ids=[],
cumulative_logprob=None,
logprobs=None,
finish_reason="error",
)
request_output = RequestOutput(
request_id="test-id",
prompt="Test prompt",
prompt_token_ids=[1, 2, 3],
prompt_logprobs=None,
outputs=[completion_output],
finished=True,
metrics=None,
lora_request=None,
encoder_prompt=None,
encoder_prompt_token_ids=None,
)
async def mock_generate(*args, **kwargs):
yield request_output
mock_engine.generate = MagicMock(side_effect=mock_generate)
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "Test prompt"}],
max_tokens=10,
stream=False,
)
with pytest.raises(GenerationError):
await serving_chat.create_chat_completion(request)
@pytest.mark.asyncio
async def test_openai_chat_keeps_mm_cache_for_engine_execution():
mock_engine = MagicMock(spec=AsyncLLM)
mock_engine.errored = False
mock_engine.model_config = MockModelConfig()
mock_engine.input_processor = MagicMock()
mock_engine.renderer = _build_renderer(mock_engine.model_config)
serving_chat = _build_serving_chat(mock_engine)
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "Test prompt"}],
)
result = await serving_chat.render_chat_request(request)
assert isinstance(result, tuple)
assert (
serving_chat.online_renderer.preprocess_chat.call_args.kwargs["skip_mm_cache"]
is False
)
@pytest.mark.asyncio
@pytest.mark.parametrize("include_usage", [False, True])
async def test_chat_error_stream(include_usage: bool):
"""Test finish_reason='error' returns 500 InternalServerError (streaming)."""
mock_engine = MagicMock(spec=AsyncLLM)
mock_engine.errored = False
mock_engine.model_config = MockModelConfig()
mock_engine.input_processor = MagicMock()
mock_engine.renderer = _build_renderer(mock_engine.model_config)
serving_chat = _build_serving_chat(mock_engine)
completion_output_1 = CompletionOutput(
index=0,
text="Hello",
token_ids=[100],
cumulative_logprob=None,
logprobs=None,
finish_reason=None,
)
request_output_1 = RequestOutput(
request_id="test-id",
prompt="Test prompt",
prompt_token_ids=[1, 2, 3],
prompt_logprobs=None,
outputs=[completion_output_1],
finished=False,
metrics=None,
lora_request=None,
encoder_prompt=None,
encoder_prompt_token_ids=None,
)
completion_output_2 = CompletionOutput(
index=0,
text="Hello",
token_ids=[100],
cumulative_logprob=None,
logprobs=None,
finish_reason="error",
)
request_output_2 = RequestOutput(
request_id="test-id",
prompt="Test prompt",
prompt_token_ids=[1, 2, 3],
prompt_logprobs=None,
outputs=[completion_output_2],
finished=True,
metrics=None,
lora_request=None,
encoder_prompt=None,
encoder_prompt_token_ids=None,
)
async def mock_generate(*args, **kwargs):
yield request_output_1
yield request_output_2
mock_engine.generate = MagicMock(side_effect=mock_generate)
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "Test prompt"}],
max_tokens=10,
stream=True,
stream_options={"include_usage": include_usage},
)
response = await serving_chat.create_chat_completion(request)
chunks = []
async for chunk in response:
chunks.append(chunk)
assert len(chunks) >= 2
assert any("Internal server error" in chunk for chunk in chunks), (
f"Expected error message in chunks: {chunks}"
)
assert chunks[-1] == "data: [DONE]\n\n"
@pytest.mark.asyncio
@pytest.mark.parametrize("include_usage", [False, True])
async def test_chat_error_stream_first_chunk(include_usage: bool):
"""An error before the first token must reach the streaming client."""
mock_engine = MagicMock(spec=AsyncLLM)
mock_engine.errored = False
mock_engine.model_config = MockModelConfig()
mock_engine.input_processor = MagicMock()
mock_engine.renderer = _build_renderer(mock_engine.model_config)
serving_chat = _build_serving_chat(mock_engine)
completion_output = CompletionOutput(
index=0,
text="",
token_ids=[],
cumulative_logprob=None,
logprobs=None,
finish_reason="error",
)
request_output = RequestOutput(
request_id="test-id",
prompt="Test prompt",
prompt_token_ids=[1, 2, 3],
prompt_logprobs=None,
outputs=[completion_output],
finished=True,
metrics=None,
lora_request=None,
encoder_prompt=None,
encoder_prompt_token_ids=None,
)
async def mock_generate(*args, **kwargs):
yield request_output
mock_engine.generate = MagicMock(side_effect=mock_generate)
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "Test prompt"}],
max_tokens=10,
stream=True,
stream_options={"include_usage": include_usage},
)
response = await serving_chat.create_chat_completion(request)
chunks = []
async for chunk in response:
chunks.append(chunk)
assert len(chunks) >= 2
assert any("Internal server error" in chunk for chunk in chunks), (
f"Expected error message in chunks: {chunks}"
)
assert chunks[-1] == "data: [DONE]\n\n"
@pytest.mark.parametrize(
"image_content",
[
[{"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}}],
[{"image_url": {"url": "https://example.com/image.jpg"}}],
],
)
def test_system_message_warns_on_image(image_content):
"""Test that system messages with image content trigger a warning."""
with patch(
"vllm.entrypoints.openai.chat_completion.protocol.logger"
) as mock_logger:
ChatCompletionRequest(
model=MODEL_NAME,
messages=[
{
"role": "system",
"content": image_content,
}
],
)
mock_logger.warning_once.assert_called()
call_args = str(mock_logger.warning_once.call_args)
assert "System messages should only contain text" in call_args
assert "image_url" in call_args
def test_system_message_accepts_text():
"""Test that system messages can contain text content."""
# Should not raise an exception
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[
{"role": "system", "content": "You are a helpful assistant."},
],
)
assert request.messages[0]["role"] == "system"
def test_system_message_accepts_text_array():
"""Test that system messages can contain an array with text content."""
# Should not raise an exception
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[
{
"role": "system",
"content": [{"type": "text", "text": "You are a helpful assistant."}],
},
],
)
assert request.messages[0]["role"] == "system"
def test_user_message_accepts_image():
"""Test that user messages can still contain image content."""
# Should not raise an exception
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "What's in this image?"},
{
"type": "image_url",
"image_url": {"url": "https://example.com/image.jpg"},
},
],
},
],
)
assert request.messages[0]["role"] == "user"
@pytest.mark.parametrize(
"audio_content",
[
[
{
"type": "input_audio",
"input_audio": {"data": "base64data", "format": "wav"},
}
],
[{"input_audio": {"data": "base64data", "format": "wav"}}],
],
)
def test_system_message_warns_on_audio(audio_content):
"""Test that system messages with audio content trigger a warning."""
with patch(
"vllm.entrypoints.openai.chat_completion.protocol.logger"
) as mock_logger:
ChatCompletionRequest(
model=MODEL_NAME,
messages=[
{
"role": "system",
"content": audio_content,
}
],
)
mock_logger.warning_once.assert_called()
call_args = str(mock_logger.warning_once.call_args)
assert "System messages should only contain text" in call_args
assert "input_audio" in call_args
@pytest.mark.parametrize(
"video_content",
[
[{"type": "video_url", "video_url": {"url": "https://example.com/video.mp4"}}],
[{"video_url": {"url": "https://example.com/video.mp4"}}],
],
)
def test_system_message_warns_on_video(video_content):
"""Test that system messages with video content trigger a warning."""
with patch(
"vllm.entrypoints.openai.chat_completion.protocol.logger"
) as mock_logger:
ChatCompletionRequest(
model=MODEL_NAME,
messages=[
{
"role": "system",
"content": video_content,
}
],
)
mock_logger.warning_once.assert_called()
call_args = str(mock_logger.warning_once.call_args)
assert "System messages should only contain text" in call_args
assert "video_url" in call_args
def test_json_schema_response_format_missing_schema():
"""When response_format type is 'json_schema' but the json_schema field
is not provided, request construction should raise a validation error
so the API returns 400 instead of 500."""
with pytest.raises(Exception, match="json_schema.*must be provided"):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
response_format={"type": "json_schema"},
)
@pytest.mark.asyncio
async def test_online_renderer_rejects_mm_processor_kwargs_by_default():
model_config = MockModelConfig()
model_config.multimodal_config = MultiModalConfig()
online_renderer = OnlineRenderer(
model_config=model_config,
renderer=MagicMock(),
request_logger=None,
chat_template=None,
chat_template_content_format="auto",
)
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
mm_processor_kwargs={
"patch_size": 1,
"vision_min_num_patches": 3_000_000_000,
},
)
with pytest.raises(
VLLMValidationError,
match="Per-request mm_processor_kwargs are disabled",
):
await online_renderer.preprocess_chat(
request,
request.messages,
default_template=None,
default_template_content_format="auto",
default_template_kwargs=None,
)
@pytest.mark.asyncio
async def test_online_renderer_rejects_media_io_kwargs_by_default():
model_config = MockModelConfig()
model_config.multimodal_config = MultiModalConfig()
online_renderer = OnlineRenderer(
model_config=model_config,
renderer=MagicMock(),
request_logger=None,
chat_template=None,
chat_template_content_format="auto",
)
request = ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
media_io_kwargs={"video": {"num_frames": 1_000_000}},
)
with pytest.raises(
VLLMValidationError,
match="Per-request media_io_kwargs are disabled",
):
await online_renderer.preprocess_chat(
request,
request.messages,
default_template=None,
default_template_content_format="auto",
default_template_kwargs=None,
)
def test_trust_request_mm_kwargs_opt_in_allows_overrides():
from vllm.entrypoints.generate.base.protocol import validate_request_mm_kwargs
validate_request_mm_kwargs(
mm_processor_kwargs={"use_audio_in_video": True},
media_io_kwargs={"video": {"num_frames": 4}},
trust_request_mm_kwargs=True,
)
@pytest.mark.parametrize("format_value", [None, {}])
def test_structural_tag_response_format_invalid(format_value):
"""Malformed structural tags should be rejected during request validation."""
with pytest.raises(
VLLMValidationError,
match="Invalid response_format structural_tag",
):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
response_format={"type": "structural_tag", "format": format_value},
)
@pytest.mark.parametrize("format_value", [None, {}])
def test_batch_structural_tag_response_format_invalid(format_value):
"""Batch chat should reject malformed structural tags at request parsing."""
with pytest.raises(
VLLMValidationError,
match="Invalid response_format structural_tag",
):
BatchChatCompletionRequest(
model=MODEL_NAME,
messages=[[{"role": "user", "content": "hello"}]],
response_format={"type": "structural_tag", "format": format_value},
)
@pytest.mark.parametrize(
"structural_tag",
[
"not json",
"",
# json.loads raises a plain ValueError, not JSONDecodeError
'{"type": "structural_tag", "x": ' + "1" * 5000 + "}",
],
)
def test_structured_outputs_structural_tag_invalid(structural_tag):
"""Malformed direct structured_outputs structural tags should be rejected."""
with pytest.raises(
VLLMValidationError,
match="Invalid structured_outputs structural_tag",
):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
structured_outputs={"structural_tag": structural_tag},
)
@pytest.mark.parametrize("field", ["response_format", "structured_outputs"])
def test_deeply_nested_structural_tag_rejected(field):
"""Structural tags are compiled during request parsing, so a deeply nested
one must be rejected there instead of blocking the server while converting,
and reported as too deep rather than malformed."""
schema = '{"type": "array", "items": ' * 200 + "{}" + "}" * 200
tag = (
'{"type": "structural_tag", "format": {"type": "json_schema", '
f'"json_schema": {schema}}}}}'
)
value = json.loads(tag) if field == "response_format" else {"structural_tag": tag}
with pytest.raises(VLLMValidationError, match="nested too deeply"):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
**{field: value},
)
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"])
def test_non_numeric_logprobs_rejected(field_name):
"""A non-numeric logprobs value must be a clean 400 validation error, not a
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
**{field_name: "2"},
)
tests/entrypoints/openai/completion/test_completion_with_prompt_embeds.py
-304
View File
@@ -1,304 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import io
import json
import openai # use the official client for correctness check
import pybase64 as base64
import pytest
import pytest_asyncio
import torch
# downloading lora to test lora requests
from openai import BadRequestError
from transformers import AutoConfig
from tests.utils import RemoteOpenAIServer
# any model with a chat template should work here
MODEL_NAME = "facebook/opt-125m"
LORA_SERVING_MODEL_NAME = "opt125m-lora"
CONFIG = AutoConfig.from_pretrained(MODEL_NAME)
@pytest.fixture(scope="module", params=["use-lora"])
def default_server_args(
request: pytest.FixtureRequest, opt125_lora_files: str
) -> list[str]:
args = [
# use half precision for speed and memory savings in CI environment
"--dtype",
"bfloat16",
"--max-model-len",
"2048",
"--max-num-seqs",
"128",
"--enforce-eager",
# Prompt Embeds server args
"--enable-prompt-embeds",
]
if request.param == "use-lora":
lora_module_1 = {
"name": LORA_SERVING_MODEL_NAME,
"path": opt125_lora_files,
"base_model_name": MODEL_NAME,
}
args.extend(
[
"--enable-lora",
"--lora-module",
json.dumps(lora_module_1),
"--max-lora-rank",
"64",
"--max-cpu-loras",
"2",
]
)
return args
EXAMPLE_PROMPTS = [
"Hello, my name is",
"What is an LLM?",
]
def _encode_embeds(embeds: torch.Tensor):
buffer = io.BytesIO()
torch.save(embeds, buffer)
return base64.b64encode(buffer.getvalue()).decode("utf-8")
@pytest.fixture(scope="module")
def example_prompt_embeds(hf_runner):
"""Create example embeddings and return them as base64 encoded string."""
with hf_runner(MODEL_NAME) as hf_model:
example_embeddings = hf_model.get_prompt_embeddings(EXAMPLE_PROMPTS)
return [_encode_embeds(item) for item in example_embeddings]
@pytest.fixture(scope="module")
def server_with_prompt_embeds(default_server_args):
with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client_with_prompt_embeds(server_with_prompt_embeds):
async with server_with_prompt_embeds.get_async_client() as async_client:
yield async_client
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
async def test_completions_with_prompt_embeds(
example_prompt_embeds,
client_with_prompt_embeds: openai.AsyncOpenAI,
model_name: str,
):
encoded_embeds, encoded_embeds2 = example_prompt_embeds
# Test case: Single prompt embeds input
completion = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt=None,
max_tokens=5,
temperature=0.0,
extra_body={"prompt_embeds": encoded_embeds},
)
assert len(completion.choices[0].text) >= 1
assert completion.choices[0].prompt_logprobs is None
# Test case: batch completion with prompt_embeds
completion = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt=None,
max_tokens=5,
temperature=0.0,
extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]},
)
assert len(completion.choices) == 2
assert len(completion.choices[0].text) >= 1
assert len(completion.choices[1].text) >= 1
# Test case: streaming with prompt_embeds
single_completion = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt=None,
max_tokens=5,
temperature=0.0,
extra_body={"prompt_embeds": encoded_embeds},
)
single_output = single_completion.choices[0].text
stream = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt=None,
max_tokens=5,
temperature=0.0,
stream=True,
extra_body={"prompt_embeds": encoded_embeds},
)
chunks = []
finish_reason_count = 0
async for chunk in stream:
chunks.append(chunk.choices[0].text)
if chunk.choices[0].finish_reason is not None:
finish_reason_count += 1
assert finish_reason_count == 1
assert chunk.choices[0].finish_reason == "length"
assert chunk.choices[0].text
assert "".join(chunks) == single_output
# Test case: batch streaming with prompt_embeds
stream = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt=None,
max_tokens=5,
temperature=0.0,
stream=True,
extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]},
)
chunks_stream_embeds: list[list[str]] = [[], []]
finish_reason_count = 0
async for chunk in stream:
chunks_stream_embeds[chunk.choices[0].index].append(chunk.choices[0].text)
if chunk.choices[0].finish_reason is not None:
finish_reason_count += 1
assert finish_reason_count == 2
assert chunk.choices[0].finish_reason == "length"
assert chunk.choices[0].text
assert len(chunks_stream_embeds[0]) > 0
assert len(chunks_stream_embeds[1]) > 0
# Test case: mixed text and prompt_embeds
completion_mixed = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt="This is a prompt",
max_tokens=5,
temperature=0.0,
extra_body={"prompt_embeds": encoded_embeds},
)
assert len(completion.choices) == 2
completion_text_only = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt="This is a prompt",
max_tokens=5,
temperature=0.0,
)
completion_embeds_only = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt=None,
max_tokens=5,
temperature=0.0,
extra_body={"prompt_embeds": encoded_embeds},
)
# Embeddings responses should be handled first
assert completion_mixed.choices[0].text == completion_embeds_only.choices[0].text
assert completion_mixed.choices[1].text == completion_text_only.choices[0].text
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
async def test_completions_errors_with_prompt_embeds(
client_with_prompt_embeds: openai.AsyncOpenAI, model_name: str
):
# Test error case: invalid prompt_embeds
with pytest.raises(BadRequestError):
await client_with_prompt_embeds.completions.create(
prompt=None,
model=model_name,
max_tokens=5,
temperature=0.0,
extra_body={"prompt_embeds": "invalid_base64"},
)
@pytest.mark.asyncio
@pytest.mark.parametrize("logprobs_arg", [1, 0])
@pytest.mark.parametrize("model_name", [MODEL_NAME, LORA_SERVING_MODEL_NAME])
async def test_completions_with_logprobs_and_prompt_embeds(
example_prompt_embeds,
client_with_prompt_embeds: openai.AsyncOpenAI,
logprobs_arg: int,
model_name: str,
):
encoded_embeds, encoded_embeds2 = example_prompt_embeds
# Test case: Logprobs using prompt_embeds
completion = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt=None,
max_tokens=5,
temperature=0.0,
echo=False,
logprobs=logprobs_arg,
extra_body={"prompt_embeds": encoded_embeds},
)
logprobs = completion.choices[0].logprobs
assert logprobs is not None
assert len(logprobs.text_offset) == 5
assert len(logprobs.token_logprobs) == 5
assert len(logprobs.top_logprobs) == 5
for top_logprobs in logprobs.top_logprobs[1:]:
assert max(logprobs_arg, 1) <= len(top_logprobs) <= logprobs_arg + 1
assert len(logprobs.tokens) == 5
# Test case: Log probs with batch completion and prompt_embeds
completion = await client_with_prompt_embeds.completions.create(
model=model_name,
prompt=None,
max_tokens=5,
temperature=0.0,
echo=False,
logprobs=logprobs_arg,
extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]},
)
assert len(completion.choices) == 2
for choice in completion.choices:
logprobs = choice.logprobs
assert logprobs is not None
assert len(logprobs.text_offset) == 5
assert len(logprobs.token_logprobs) == 5
assert len(logprobs.top_logprobs) == 5
for top_logprobs in logprobs.top_logprobs[1:]:
assert max(logprobs_arg, 1) <= len(top_logprobs) <= logprobs_arg + 1
assert len(logprobs.tokens) == 5
@pytest.mark.asyncio
async def test_prompt_logprobs_raises_error(
example_prompt_embeds,
client_with_prompt_embeds: openai.AsyncOpenAI,
):
encoded_embeds, _ = example_prompt_embeds
with pytest.raises(BadRequestError, match="not compatible"):
await client_with_prompt_embeds.completions.create(
model=MODEL_NAME,
prompt=None,
max_tokens=5,
temperature=0.0,
extra_body={"prompt_embeds": encoded_embeds, "prompt_logprobs": True},
)
@pytest.mark.asyncio
async def test_empty_prompt_embeds(
client_with_prompt_embeds: openai.AsyncOpenAI,
) -> None:
await client_with_prompt_embeds.completions.create(
model=MODEL_NAME,
prompt="Hello",
max_tokens=5,
temperature=0.0,
extra_body={"prompt_embeds": []},
)
tests/entrypoints/openai/completion/test_lora_resolvers.py
-256
View File
@@ -1,256 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from contextlib import suppress
from dataclasses import dataclass, field
from http import HTTPStatus
from unittest.mock import AsyncMock, MagicMock
import pytest
from vllm.config.multimodal import MultiModalConfig
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
from vllm.entrypoints.openai.models.protocol import BaseModelPath
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.serve.engine.protocol import ErrorResponse
from vllm.lora.request import LoRARequest
from vllm.lora.resolver import LoRAResolver, LoRAResolverRegistry
from vllm.renderers.hf import HfRenderer
from vllm.renderers.online_renderer import OnlineRenderer
from vllm.tokenizers.registry import cached_tokenizer_from_config
from vllm.v1.engine.async_llm import AsyncLLM
MODEL_NAME = "openai-community/gpt2"
BASE_MODEL_PATHS = [BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME)]
MOCK_RESOLVER_NAME = "mock_test_resolver"
@dataclass
class MockHFConfig:
model_type: str = "any"
@dataclass
class MockModelConfig:
"""Minimal mock ModelConfig for testing."""
model: str = MODEL_NAME
runner_type = "generate"
tokenizer: str = MODEL_NAME
trust_remote_code: bool = False
tokenizer_mode: str = "auto"
max_model_len: int = 100
tokenizer_revision: str | None = None
multimodal_config: MultiModalConfig = field(default_factory=MultiModalConfig)
hf_config: MockHFConfig = field(default_factory=MockHFConfig)
logits_processors: list[str] | None = None
diff_sampling_param: dict | None = None
allowed_local_media_path: str = ""
allowed_media_domains: list[str] | None = None
encoder_config = None
generation_config: str = "auto"
skip_tokenizer_init: bool = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
supports_multimodal_inputs: bool = False
renderer_num_workers: int = 1
def get_diff_sampling_param(self):
return self.diff_sampling_param or {}
@dataclass
class MockParallelConfig:
_api_process_rank: int = 0
@dataclass
class MockVllmConfig:
model_config: MockModelConfig
parallel_config: MockParallelConfig
class MockLoRAResolver(LoRAResolver):
async def resolve_lora(
self, base_model_name: str, lora_name: str
) -> LoRARequest | None:
if lora_name == "test-lora":
return LoRARequest(
lora_name="test-lora",
lora_int_id=1,
lora_path="/fake/path/test-lora",
)
elif lora_name == "invalid-lora":
return LoRARequest(
lora_name="invalid-lora",
lora_int_id=2,
lora_path="/fake/path/invalid-lora",
)
return None
@pytest.fixture(autouse=True)
def register_mock_resolver():
"""Fixture to register and unregister the mock LoRA resolver."""
resolver = MockLoRAResolver()
LoRAResolverRegistry.register_resolver(MOCK_RESOLVER_NAME, resolver)
yield
# Cleanup: remove the resolver after the test runs
if MOCK_RESOLVER_NAME in LoRAResolverRegistry.resolvers:
del LoRAResolverRegistry.resolvers[MOCK_RESOLVER_NAME]
def _build_renderer(model_config: MockModelConfig):
return HfRenderer(
MockVllmConfig(model_config, parallel_config=MockParallelConfig()),
cached_tokenizer_from_config(model_config),
)
@pytest.fixture
def mock_serving_setup():
"""Provides a mocked engine and serving completion instance."""
mock_engine = MagicMock(spec=AsyncLLM)
mock_engine.errored = False
async def mock_add_lora_side_effect(lora_request: LoRARequest):
"""Simulate engine behavior when adding LoRAs."""
if lora_request.lora_name == "test-lora":
# Simulate successful addition
return True
if lora_request.lora_name == "invalid-lora":
# Simulate failure during addition (e.g. invalid format)
raise ValueError(f"Simulated failure adding LoRA: {lora_request.lora_name}")
return True
mock_engine.add_lora = AsyncMock(side_effect=mock_add_lora_side_effect)
async def mock_generate(*args, **kwargs):
for _ in []:
yield _
mock_engine.generate = MagicMock(spec=AsyncLLM.generate, side_effect=mock_generate)
mock_engine.generate.reset_mock()
mock_engine.add_lora.reset_mock()
mock_engine.model_config = MockModelConfig()
mock_engine.input_processor = MagicMock()
mock_engine.renderer = _build_renderer(mock_engine.model_config)
models = OpenAIServingModels(
engine_client=mock_engine,
base_model_paths=BASE_MODEL_PATHS,
)
online_renderer = OnlineRenderer(
model_config=mock_engine.model_config,
renderer=mock_engine.renderer,
request_logger=None,
chat_template=None,
chat_template_content_format="auto",
)
serving_completion = OpenAIServingCompletion(
mock_engine, models, online_renderer=online_renderer, request_logger=None
)
return mock_engine, serving_completion
@pytest.mark.asyncio
async def test_serving_completion_with_lora_resolver(mock_serving_setup, monkeypatch):
monkeypatch.setenv("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "true")
mock_engine, serving_completion = mock_serving_setup
lora_model_name = "test-lora"
req_found = CompletionRequest(
model=lora_model_name,
prompt="Generate with LoRA",
)
# Suppress potential errors during the mocked generate call,
# as we are primarily checking for add_lora and generate calls
with suppress(Exception):
await serving_completion.create_completion(req_found)
mock_engine.add_lora.assert_awaited_once()
called_lora_request = mock_engine.add_lora.call_args[0][0]
assert isinstance(called_lora_request, LoRARequest)
assert called_lora_request.lora_name == lora_model_name
mock_engine.generate.assert_called_once()
called_lora_request = mock_engine.generate.call_args[1]["lora_request"]
assert isinstance(called_lora_request, LoRARequest)
assert called_lora_request.lora_name == lora_model_name
@pytest.mark.asyncio
async def test_serving_completion_resolver_not_found(mock_serving_setup, monkeypatch):
monkeypatch.setenv("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "true")
mock_engine, serving_completion = mock_serving_setup
non_existent_model = "non-existent-lora-adapter"
req = CompletionRequest(
model=non_existent_model,
prompt="what is 1+1?",
)
response = await serving_completion.create_completion(req)
mock_engine.add_lora.assert_not_awaited()
mock_engine.generate.assert_not_called()
assert isinstance(response, ErrorResponse)
assert response.error.code == HTTPStatus.NOT_FOUND.value
assert non_existent_model in response.error.message
@pytest.mark.asyncio
async def test_serving_completion_resolver_add_lora_fails(
mock_serving_setup, monkeypatch
):
monkeypatch.setenv("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "true")
mock_engine, serving_completion = mock_serving_setup
invalid_model = "invalid-lora"
req = CompletionRequest(
model=invalid_model,
prompt="what is 1+1?",
)
response = await serving_completion.create_completion(req)
# Assert add_lora was called before the failure
mock_engine.add_lora.assert_awaited_once()
called_lora_request = mock_engine.add_lora.call_args[0][0]
assert isinstance(called_lora_request, LoRARequest)
assert called_lora_request.lora_name == invalid_model
# Assert generate was *not* called due to the failure
mock_engine.generate.assert_not_called()
# Assert the correct error response
assert isinstance(response, ErrorResponse)
assert response.error.code == HTTPStatus.BAD_REQUEST.value
assert invalid_model in response.error.message
@pytest.mark.asyncio
async def test_serving_completion_flag_not_set(mock_serving_setup):
mock_engine, serving_completion = mock_serving_setup
lora_model_name = "test-lora"
req_found = CompletionRequest(
model=lora_model_name,
prompt="Generate with LoRA",
)
await serving_completion.create_completion(req_found)
mock_engine.add_lora.assert_not_called()
mock_engine.generate.assert_not_called()
tests/entrypoints/openai/correctness/__init__.py
View File
tests/entrypoints/openai/responses/conftest.py
-1
View File
@@ -396,7 +396,6 @@ def server_with_store(default_server_args):
default_server_args,
env_dict={
"VLLM_ENABLE_RESPONSES_API_STORE": "1",
"VLLM_SERVER_DEV_MODE": "1",
},
) as remote_server:
yield remote_server
tests/entrypoints/serve/dev/__init__.py
View File
tests/entrypoints/serve/dev/rpc/__init__.py
View File
tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
-84
View File
@@ -1,84 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
import pytest
import requests
from tests.utils import RemoteOpenAIServer
MODEL_NAME = "Qwen/Qwen3-0.6B"
class TestWorkerExtension:
def get_model_name(self) -> str:
"""Test non-pydantic return type."""
return MODEL_NAME
def echo_args_kwargs(self, *args, **kwargs) -> dict[str, Any]:
"""Echo back both args and kwargs."""
return dict(
args=list(args),
kwargs=kwargs,
total_items=len(args) + len(kwargs),
)
def return_none(self, *args, **kwargs) -> None:
"""Test method that does not return anything."""
return
@pytest.fixture(scope="module")
def server():
args = [
"--max-model-len",
"8192",
"--max-num-seqs",
"128",
"--worker-extension-cls",
"tests.entrypoints.serve.dev.rpc.test_collective_rpc.TestWorkerExtension",
]
with RemoteOpenAIServer(
MODEL_NAME,
args,
env_dict={"VLLM_SERVER_DEV_MODE": "1", "CUDA_VISIBLE_DEVICES": "0"},
) as remote_server:
yield remote_server
def test_get_model_name(server):
"""Test basic response."""
response = requests.post(
server.url_for("collective_rpc"), json={"method": "get_model_name"}
)
assert response.status_code == 200
results = response.json()
assert "results" in results
assert results["results"] == [MODEL_NAME]
def test_return_none(server):
"""Test return none."""
response = requests.post(
server.url_for("collective_rpc"), json={"method": "return_none"}
)
assert response.status_code == 200
results = response.json()
assert results["results"] == [None]
def test_echo_args_kwargs(server):
"""Test args, kwargs, and dict response."""
args = ["arg1", "arg2"]
kwargs = {"key1": "value1", "key2": "value2"}
response = requests.post(
server.url_for("collective_rpc"),
json={"method": "echo_args_kwargs", "args": args, "kwargs": kwargs},
)
assert response.status_code == 200
results = response.json()
result = results["results"][0]
assert result["args"] == args
assert result["kwargs"] == kwargs
assert result["total_items"] == len(args) + len(kwargs)
tests/entrypoints/serve/middleware/test_authentication_middleware.py
+2 -4
View File
@@ -70,19 +70,17 @@ class MockModelConfig:
@pytest.fixture(params=get_args(SupportedTask))
def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]:
def task_routes(request) -> tuple[str, list[tuple[str, list[str]]]]:
"""For each supported task, build an app with only that task's routers,
extract all routes, and return the task name and routes."""
task = request.param
# Enable development mode to register all routes (including dev-only routes).
monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1")
app = FastAPI()
args = Namespace()
app.state = Namespace()
app.state.args = args
# Register routers for this specific task (development mode already enabled).
# Register routers for this specific task.
register_api_routers(
args,
app,
tests/entrypoints/serve/tokenize/test_serving_tokenization.py
-250
View File
@@ -1,250 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from argparse import Namespace
from dataclasses import dataclass, field
from http import HTTPStatus
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from vllm.config.multimodal import MultiModalConfig
from vllm.entrypoints.openai.models.protocol import BaseModelPath
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.serve.exception_handling.register import init_exception_handler
from vllm.entrypoints.serve.tokenize.api_router import attach_router
from vllm.entrypoints.serve.tokenize.protocol import (
TokenizeChatRequest,
TokenizeCompletionRequest,
TokenizeResponse,
)
from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
from vllm.exceptions import VLLMNotFoundError, VLLMValidationError
from vllm.renderers.online_renderer import OnlineRenderer
from vllm.v1.engine.async_llm import AsyncLLM
MODEL_NAME = "openai-community/gpt2"
BASE_MODEL_PATHS = [
BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
]
@dataclass
class MockHFConfig:
model_type: str = "any"
@dataclass
class MockModelConfig:
task = "generate"
runner_type = "generate"
model = MODEL_NAME
tokenizer = MODEL_NAME
trust_remote_code = False
tokenizer_mode = "auto"
max_model_len = 100
tokenizer_revision = None
multimodal_config = MultiModalConfig()
hf_config = MockHFConfig()
hf_text_config = MockHFConfig()
logits_processors: list[str] | None = None
diff_sampling_param: dict | None = None
allowed_local_media_path: str = ""
allowed_media_domains: list[str] | None = None
encoder_config = None
generation_config: str = "auto"
media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
skip_tokenizer_init: bool = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
supports_multimodal_inputs: bool = False
renderer_num_workers: int = 1
def get_diff_sampling_param(self):
return self.diff_sampling_param or {}
def _build_serving_tokenization(engine: AsyncLLM) -> ServingTokenization:
models = OpenAIServingModels(
engine_client=engine,
base_model_paths=BASE_MODEL_PATHS,
)
online_renderer = OnlineRenderer(
model_config=engine.model_config,
renderer=engine.renderer,
request_logger=None,
chat_template=None,
chat_template_content_format="auto",
)
return ServingTokenization(
models,
online_renderer=online_renderer,
chat_template=None,
chat_template_content_format="auto",
)
@pytest.mark.asyncio
async def test_tokenize_chat_skips_mm_cache_for_renderer_only_path():
mock_engine = MagicMock(spec=AsyncLLM)
mock_engine.errored = False
mock_engine.model_config = MockModelConfig()
mock_engine.input_processor = MagicMock()
mock_engine.renderer = MagicMock()
serving = _build_serving_tokenization(mock_engine)
serving.online_renderer.preprocess_chat = AsyncMock( # type: ignore[method-assign]
return_value=(
[{"role": "user", "content": "Test"}],
[{"prompt_token_ids": [1, 2, 3]}],
)
)
request = TokenizeChatRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "Test prompt"}],
)
response = await serving.create_tokenize(request, MagicMock(headers={}))
assert isinstance(response, TokenizeResponse) and response.tokens == [1, 2, 3]
assert (
serving.online_renderer.preprocess_chat.call_args.kwargs["skip_mm_cache"]
is True
)
@pytest.mark.asyncio
async def test_tokenize_completion_skips_mm_cache_for_renderer_only_path():
mock_engine = MagicMock(spec=AsyncLLM)
mock_engine.errored = False
mock_engine.model_config = MockModelConfig()
mock_engine.input_processor = MagicMock()
mock_engine.renderer = MagicMock()
serving = _build_serving_tokenization(mock_engine)
serving.online_renderer.preprocess_completion = AsyncMock( # type: ignore[method-assign]
return_value=[{"prompt_token_ids": [1, 2, 3]}]
)
request = TokenizeCompletionRequest(
model=MODEL_NAME,
prompt="Test prompt",
)
response = await serving.create_tokenize(request, MagicMock(headers={}))
assert isinstance(response, TokenizeResponse) and response.tokens == [1, 2, 3]
assert (
serving.online_renderer.preprocess_completion.call_args.kwargs["skip_mm_cache"]
is True
)
class TestDetokenizeClientErrorResponses:
"""Client-caused errors from /detokenize are 4xx, not 500.
/detokenize used to wrap ``create_detokenize`` in a blanket
``except Exception`` that reported every failure as 500, while its
sibling /tokenize let exceptions propagate to the global exception
handlers, which map client-caused errors to 4xx and keep genuine
server errors at 500. These tests pin the shared contract for both
endpoints. See #52246 for the same fix in the Anthropic entrypoint.
"""
@staticmethod
def _make_api_app(handler: MagicMock) -> FastAPI:
app = FastAPI()
app.state.args = Namespace(log_error_stack=False)
app.state.serving_tokenization = handler
attach_router(app)
init_exception_handler(app)
return app
def _post(self, handler: MagicMock, path: str = "/detokenize"):
app = self._make_api_app(handler)
body: dict[str, Any] = {"model": "test-model"}
if path == "/detokenize":
body["tokens"] = [1, 2, 3]
else:
body["prompt"] = "Hello"
with TestClient(app, raise_server_exceptions=False) as client:
return client.post(path, json=body)
def test_vllm_validation_error_returns_bad_request(self):
handler = MagicMock(spec=ServingTokenization)
handler.create_detokenize.side_effect = VLLMValidationError(
"invalid token ids", parameter="tokens"
)
response = self._post(handler)
assert response.status_code == HTTPStatus.BAD_REQUEST
error = response.json()["error"]
assert error["type"] == "BadRequestError"
assert error["param"] == "tokens"
def test_vllm_not_found_error_returns_not_found(self):
handler = MagicMock(spec=ServingTokenization)
handler.create_detokenize.side_effect = VLLMNotFoundError(
"LoRA adapter nonexistent-lora not found"
)
response = self._post(handler)
assert response.status_code == HTTPStatus.NOT_FOUND
assert response.json()["error"]["type"] == "NotFoundError"
def test_value_error_returns_bad_request(self):
handler = MagicMock(spec=ServingTokenization)
handler.create_detokenize.side_effect = ValueError("bad input")
response = self._post(handler)
assert response.status_code == HTTPStatus.BAD_REQUEST
assert response.json()["error"]["type"] == "BadRequestError"
def test_overflow_error_returns_bad_request(self):
"""Fast tokenizers raise OverflowError for out-of-range token ids.
This used to be special-cased into a RequestValidationError; the
global OverflowError handler keeps the same 400 status code.
"""
handler = MagicMock(spec=ServingTokenization)
handler.create_detokenize.side_effect = OverflowError(
"out of range integral type conversion attempted"
)
response = self._post(handler)
assert response.status_code == HTTPStatus.BAD_REQUEST
assert response.json()["error"]["type"] == "BadRequestError"
def test_server_error_still_returns_internal_server_error(self):
"""Genuine server bugs keep the existing 500 behaviour."""
handler = MagicMock(spec=ServingTokenization)
handler.create_detokenize.side_effect = RuntimeError("boom")
response = self._post(handler)
assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
assert response.json()["error"]["type"] == "InternalServerError"
def test_tokenize_and_detokenize_agree_on_client_errors(self):
"""The two sibling endpoints report the same client error alike."""
statuses = {}
for path, method in (
("/tokenize", "create_tokenize"),
("/detokenize", "create_detokenize"),
):
handler = MagicMock(spec=ServingTokenization)
getattr(handler, method).side_effect = VLLMValidationError(
"invalid input", parameter="model"
)
statuses[path] = self._post(handler, path=path).status_code
assert statuses["/tokenize"] == HTTPStatus.BAD_REQUEST
assert statuses["/detokenize"] == HTTPStatus.BAD_REQUEST
tests/entrypoints/speech_to_text/correctness/__init__.py
View File
tests/entrypoints/speech_to_text/realtime/__init__.py
View File
tests/entrypoints/speech_to_text/translation/__init__.py
View File
tests/fusion/__init__.py
-2
View File
@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/jit_monitor/__init__.py
View File
tests/jit_monitor/conftest.py
-27
View File
@@ -1,27 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.utils import jit_monitor
@pytest.fixture(autouse=True)
def _reset_monitor():
"""Reset global monitor state between tests.
``activate()`` installs process-global hooks and flips module globals, so
without this every test would inherit the previous test's monitor state.
"""
def reset():
jit_monitor._active = False
jit_monitor._mode = "warn"
jit_monitor._verbose = False
jit_monitor._cutedsl_hook_installed = False
jit_monitor._tilelang_hook_installed = False
jit_monitor._tilelang_jitimpl_compile_depth = 0
reset()
yield
reset()
tests/kernels/attention/conftest.py
-19
View File
@@ -1,19 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.utils.torch_utils import (
create_kv_caches_with_random,
create_kv_caches_with_random_flash,
)
@pytest.fixture()
def kv_cache_factory():
return create_kv_caches_with_random
@pytest.fixture()
def kv_cache_factory_flashinfer():
return create_kv_caches_with_random_flash
tests/kernels/fla/__init__.py
View File
tests/kernels/mhc/__init__.py
View File
tests/kernels/moe/modular_kernel_tools/__init__.py
View File
tests/kernels/moe/modular_kernel_tools/cli_args.py
-158
View File
@@ -1,158 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
from .common import Config
from .mk_objects import (
MK_ALL_PREPARE_FINALIZE_TYPES,
MK_FUSED_EXPERT_TYPES,
MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES,
)
def make_config_arg_parser(description: str):
def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalizeModular:
for pf in MK_ALL_PREPARE_FINALIZE_TYPES:
if pf.__name__ == s:
return pf
raise ValueError(f"Cannot find a PrepareFinalize type that matches {s}")
def to_experts_class_type(s: str) -> mk.FusedMoEExpertsModular:
for fe in MK_FUSED_EXPERT_TYPES:
if fe.__name__ == s:
return fe
raise ValueError(f"Cannot find a FusedExperts type that matches {s}")
def to_quant_torch_dtype(s: str) -> torch.dtype:
if s == "torch.float8_e4m3fn":
return torch.float8_e4m3fn
raise ValueError(f"Unsupported quant type {s}")
parser = argparse.ArgumentParser(description=description)
parser.add_argument(
"--world-size",
type=int,
default=2,
help="Number of ranks that participate in all2all",
)
parser.add_argument(
"--pf-type",
type=to_pf_class_type,
required=True,
help=(
"Choose a PrepareFinalize Type : "
f"{[x.__name__ for x in MK_ALL_PREPARE_FINALIZE_TYPES]}"
),
)
parser.add_argument(
"--experts-type",
type=to_experts_class_type,
required=True,
help=(
f"Choose a FusedExpert type : {[x.__name__ for x in MK_FUSED_EXPERT_TYPES]}"
),
)
parser.add_argument(
"-m",
nargs="+",
type=int,
default=[64],
help="num tokens per rank",
)
parser.add_argument(
"-k",
type=int,
default=7168,
help="hidden-size",
)
parser.add_argument(
"-n",
type=int,
default=1024,
help="N dimension of the first fused-moe matmul",
)
parser.add_argument(
"--num-experts", type=int, default=32, help="Global num experts"
)
parser.add_argument("--topk", nargs="+", type=int, default=[4, 1], help="num topk")
# Quant args
parser.add_argument(
"--quant-dtype", type=to_quant_torch_dtype, help="Quant datatype"
)
parser.add_argument(
"--per-token-quantized-activations",
action="store_true",
help=("The input activations must be per-token quantized"),
)
parser.add_argument(
"--per-channel-quantized-weights",
action="store_true",
help="The weights must be per-channel quantized.",
)
parser.add_argument(
"--block-shape", nargs="+", type=int, help="Quantization block shape"
)
# Torch trace profile generation args
parser.add_argument(
"--torch-trace-dir-path",
type=str,
default=None,
help="Get torch trace for single execution",
)
return parser
def _validate_args(args: argparse.Namespace):
if args.quant_dtype is not None:
assert args.quant_dtype == torch.float8_e4m3fn
if args.block_shape is not None:
assert len(args.block_shape) == 2, (
f"block shape must have 2 elements. got {args.block_shape}"
)
if args.experts_type in MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES:
assert args.world_size == 1, "Single GPU objects need world size set to 1"
if args.torch_trace_dir_path is not None:
from pathlib import Path
assert Path(args.torch_trace_dir_path).is_dir(), (
f"Please create {args.torch_trace_dir_path}"
)
def make_config(args: argparse.Namespace) -> Config:
_validate_args(args)
quant_config = None
if args.quant_dtype is not None:
quant_config = FusedMoEQuantConfig.make(
quant_dtype=args.quant_dtype,
per_act_token_quant=args.per_token_quantized_activations,
per_out_ch_quant=args.per_channel_quantized_weights,
block_shape=args.block_shape,
)
return Config(
Ms=args.m,
K=args.k,
N=args.n,
E=args.num_experts,
topks=args.topk,
dtype=torch.bfloat16, # hard-code
quant_config=quant_config,
prepare_finalize_type=args.pf_type,
fused_experts_type=args.experts_type,
world_size=args.world_size,
torch_trace_dir_path=args.torch_trace_dir_path,
)
tests/kernels/moe/modular_kernel_tools/make_feature_matrix.py
-189
View File
@@ -1,189 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import copy
from enum import Enum
from itertools import product
import torch
from tqdm import tqdm
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.config import FUSED_MOE_UNQUANTIZED_CONFIG
from vllm.utils.torch_utils import set_random_seed
from .common import (
Config,
RankTensors,
WeightTensors,
reference_moe_impl,
run_modular_kernel,
)
from .mk_objects import (
MK_FUSED_EXPERT_TYPES,
MK_MULTI_GPU_PREPARE_FINALIZE_TYPES,
MK_QUANT_CONFIGS,
)
from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
class Result(Enum):
PASS = 1
FAIL = 2
SKIP = 3
def rank_worker(
pgi: ProcessGroupInfo,
vllm_config: VllmConfig,
cpu_group,
config: Config,
weights: WeightTensors,
):
set_random_seed(pgi.rank)
# get weights to this device
weights.to_current_device()
Ms = config.Ms
assert isinstance(Ms, list)
TOPKs = config.topks
assert isinstance(TOPKs, list)
for m, topk in product(Ms, TOPKs):
print(f"Running m={m}, topk={topk} ...")
# override m and topk
cfgx = copy.deepcopy(config)
cfgx.Ms = m
cfgx.topks = topk
# inputs for rank
rank_tensors = RankTensors.make(cfgx, pgi)
# modular kernel out
mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights, rank_tensors)
with set_current_vllm_config(vllm_config):
ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
torch.testing.assert_close(ref_out, mk_out, atol=3e-2, rtol=3e-2)
def make_feature_matrix(csv_file_path: str):
from dataclasses import asdict
import pandas as pd
def add_to_results(
config: Config, success: Result, results_df: pd.DataFrame | None = None
):
config_dict = asdict(config)
config_dict["prepare_finalize_type"] = config_dict[
"prepare_finalize_type"
].__name__
config_dict["fused_experts_type"] = config_dict["fused_experts_type"].__name__
config_dict["per_tensor_act_quant"] = config.is_per_tensor_act_quant
quant_config_dict = config_dict["quant_config"]
del config_dict["quant_config"]
if quant_config_dict is None:
quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
quant_config_dict = asdict(quant_config)
config_dict |= quant_config_dict
result_dict = config_dict | {"success": success.name}
result_df = pd.DataFrame([result_dict])
if results_df is None:
results_df = result_df
else:
results_df = pd.concat([results_df, result_df], ignore_index=True)
return results_df
Ms = [64]
Ks = [7168] # hidden sizes
Ns = [2048]
TOPKs = [[4, 1]]
Es = [32]
DTYPEs = [torch.bfloat16]
PF_TYPES = MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
FE_TYPES = MK_FUSED_EXPERT_TYPES
Q_TYPES = MK_QUANT_CONFIGS
combinations = list(
product(Ms, Ks, Ns, Es, TOPKs, DTYPEs, PF_TYPES, FE_TYPES, Q_TYPES)
)
results_df: pd.DataFrame | None = None
for m, k, n, e, topks, dtype, pf_type, experts_type, quant_config in tqdm(
combinations
):
config = Config(
Ms=[m],
K=k,
N=n,
E=e,
topks=topks,
dtype=dtype,
prepare_finalize_type=pf_type,
fused_experts_type=experts_type,
quant_config=quant_config,
world_size=2,
)
success = None
if config.is_valid()[0]:
print(f"Running config : {config.describe()} ...")
try:
weights: WeightTensors = WeightTensors.make(config)
vllm_config, env_dict = config.make_env_data()
parallel_launch_with_config(
config.world_size,
rank_worker,
vllm_config,
env_dict,
config,
weights,
)
success = Result.PASS
except Exception as _:
success = Result.FAIL
else:
success = Result.SKIP
results_df = add_to_results(config, success, results_df)
if results_df is not None:
results_df.to_csv(f"{csv_file_path}")
if __name__ == "__main__":
import argparse
from pathlib import Path
parser = argparse.ArgumentParser(
description=(
"Make ModularKernel feature matrix \n"
"Example : python3 -m tests.kernels.moe.modular_kernel_tools.make_feature_matrix " # noqa: E501
"-f ./feature_matrices/feature_matrix.csv"
)
)
parser.add_argument(
"-f",
"--feature-matrix-csv-file-path",
type=str,
required=True,
help="File name to Generate a .csv file",
)
args = parser.parse_args()
csv_path = args.feature_matrix_csv_file_path
assert csv_path.endswith("csv"), (
f"Need a file path ending with .csv, got {csv_path}"
)
assert Path(csv_path).parent.is_dir(), (
f"Cannot find parent directory for {Path(csv_path).parent}"
)
make_feature_matrix(args.feature_matrix_csv_file_path)
tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py
-182
View File
@@ -1,182 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import copy
from collections.abc import Callable
from itertools import product
from typing import Any
import torch
from vllm.config import VllmConfig
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.worker.workspace import init_workspace_manager
from .common import (
Config,
RankTensors,
WeightTensors,
_make_gscale,
make_modular_kernel,
)
from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
def do_profile(
fn: Callable,
fn_kwargs: dict[Any, Any],
pgi: ProcessGroupInfo,
config: Config,
num_warmups: int = 5,
):
for _ in range(num_warmups):
fn(**fn_kwargs)
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA,
],
with_stack=True,
record_shapes=True,
) as tprof:
fn(**fn_kwargs)
device = torch.accelerator.current_device_index()
torch.accelerator.synchronize(device)
# TODO (varun): Add a descriptive trace file name
tprof.export_chrome_trace(
f"{config.torch_trace_dir_path}/m{config.M}_{pgi.rank}_trace.json"
)
def profile_modular_kernel(
pgi: ProcessGroupInfo,
vllm_config: VllmConfig,
config: Config,
weights: WeightTensors,
rank_tensors: RankTensors,
) -> None:
assert isinstance(config.Ms, int)
assert isinstance(config.topks, int)
# weights for rank
rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
if config.quant_dtype == "nvfp4":
gscale = _make_gscale(config.num_local_experts)
else:
gscale = None
quant_config = FusedMoEQuantConfig.make(
config.quant_dtype,
w1_scale=rank_weights.w1_scale,
w2_scale=rank_weights.w2_scale,
a1_scale=rank_tensors.hidden_states_scale,
g1_alphas=(1 / rank_weights.w1_gs) if rank_weights.w1_gs is not None else None,
g2_alphas=(1 / rank_weights.w2_gs) if rank_weights.w2_gs is not None else None,
a1_gscale=gscale,
a2_gscale=gscale,
block_shape=config.quant_block_shape,
per_act_token_quant=config.is_per_act_token_quant,
per_out_ch_quant=config.is_per_out_ch_quant,
)
# make modular kernel
mk = make_modular_kernel(config, vllm_config, quant_config)
topk_ids = rank_tensors.topk_ids.to(
mk.prepare_finalize.topk_indices_dtype() or rank_tensors.topk_ids.dtype
)
# impls might update the tensor in place
hidden_states = rank_tensors.hidden_states.clone()
mk_kwargs = {
"hidden_states": hidden_states,
"w1": rank_weights.w1,
"w2": rank_weights.w2,
"topk_weights": rank_tensors.topk_weights,
"topk_ids": topk_ids,
"activation": MoEActivation.SILU,
"expert_map": rank_tensors.expert_map,
"global_num_experts": config.E,
"apply_router_weight_on_input": config.topk == 1
and config.supports_apply_weight_on_input(),
}
num_tokens = hidden_states.shape[0]
num_tokens_across_dp = torch.tensor(
[num_tokens] * config.world_size, device="cpu", dtype=torch.int
)
with set_forward_context(
None,
vllm_config,
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp,
):
do_profile(mk.apply, mk_kwargs, pgi, config)
def rank_worker(
pgi: ProcessGroupInfo,
vllm_config: VllmConfig,
cpu_group,
config: Config,
weights: WeightTensors,
):
set_random_seed(pgi.rank)
# workspace manager is normally initialized by GPUModelRunner; we initialize
# it here for the standalone benchmark process.
init_workspace_manager(torch.device(f"cuda:{pgi.local_rank}"))
# get weights to this device
weights.to_current_device()
Ms = config.Ms
assert isinstance(Ms, list)
TOPKs = config.topks
assert isinstance(TOPKs, list)
for m, topk in product(Ms, TOPKs):
print(f"Running m={m}, topk={topk} ...")
# override m and topk
cfgx = copy.deepcopy(config)
cfgx.Ms = m
cfgx.topks = topk
# inputs for rank
rank_tensors = RankTensors.make(cfgx, pgi)
profile_modular_kernel(pgi, vllm_config, cfgx, weights, rank_tensors)
def run(config: Config):
weights: WeightTensors = WeightTensors.make(config)
vllm_config, env_dict = config.make_env_data()
parallel_launch_with_config(
config.world_size, rank_worker, vllm_config, env_dict, config, weights
)
if __name__ == "__main__":
from .cli_args import make_config, make_config_arg_parser
parser = make_config_arg_parser(
description=(
"Run single prepare-finalize & fused-experts combination test"
"Example : python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel " # noqa: E501
"--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts"
)
)
args = parser.parse_args()
assert args.torch_trace_dir_path is not None, (
"Please pass in a directory to store torch traces"
)
config = make_config(args)
run(config)
tests/kernels/test_gate_linear_rocm_dispatch.py
-259
View File
@@ -1,259 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Platform-agnostic eligibility tests for GateLinear's ROCm router GEMMs.
These assert the dispatch flags directly, so they run device-free by mocking
the platform predicates. ``allow_cublas_router_gemm`` selects the
bf16xbf16->fp32 ``torch.mm`` epilogue, while ``allow_fp32_router_gemm`` selects
the gfx950 low-M kernel with fp32 weights and output.
The ROCm branch is guarded on ``not bias`` because ``torch.mm`` has no bias
term; a biased gate must fall back so the bias is not silently dropped.
"""
import inspect
import pytest
import torch
import vllm.model_executor.layers.fused_moe.router.gate_linear as gate_linear_mod
from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
def _make_gate(
monkeypatch,
*,
is_rocm: bool,
is_cuda: bool = False,
bias: bool = False,
params_dtype: torch.dtype = torch.bfloat16,
out_dtype: torch.dtype | None = torch.float32,
input_size: int = 2048,
output_size: int = 64,
on_gfx950: bool = False,
parallel_world_size: int = 1,
) -> GateLinear:
"""Build a GateLinear with platform predicates mocked, no GPU needed."""
for target in (
"vllm.model_executor.layers.linear",
"vllm.model_executor.parameter",
):
monkeypatch.setattr(
f"{target}.get_tensor_model_parallel_rank",
lambda: 0,
)
monkeypatch.setattr(
f"{target}.get_tensor_model_parallel_world_size",
lambda: parallel_world_size,
)
platform = gate_linear_mod.current_platform
monkeypatch.setattr(platform, "is_cuda", lambda: is_cuda)
monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm)
monkeypatch.setattr(platform, "is_device_capability", lambda *a, **k: False)
monkeypatch.setattr(platform, "is_device_capability_family", lambda *a, **k: False)
if is_rocm:
import vllm.platforms.rocm as rocm_platform
monkeypatch.setattr(rocm_platform, "on_gfx950", lambda: on_gfx950)
return GateLinear(
input_size=input_size,
output_size=output_size,
bias=bias,
out_dtype=out_dtype,
params_dtype=params_dtype,
)
def test_rocm_no_bias_bf16_fp32_enables_fused_gemm(monkeypatch):
gate = _make_gate(monkeypatch, is_rocm=True, bias=False)
assert not gate.allow_specialized_router_gemm
assert gate.allow_cublas_router_gemm
def test_rocm_bias_disables_fused_gemm(monkeypatch):
# torch.mm cannot add a bias, so a biased gate must not take the fused path.
gate = _make_gate(monkeypatch, is_rocm=True, bias=True)
assert not gate.allow_cublas_router_gemm
def test_rocm_fp32_weight_disables_fused_gemm(monkeypatch):
gate = _make_gate(monkeypatch, is_rocm=True, params_dtype=torch.float32)
assert not gate.allow_cublas_router_gemm
def test_rocm_non_fp32_out_dtype_disables_fused_gemm(monkeypatch):
gate = _make_gate(monkeypatch, is_rocm=True, out_dtype=torch.bfloat16)
assert not gate.allow_cublas_router_gemm
def test_non_rocm_non_cuda_disables_fused_gemm(monkeypatch):
# Neither the CUDA specialized path nor the ROCm branch applies.
gate = _make_gate(monkeypatch, is_rocm=False, is_cuda=False)
assert not gate.allow_cublas_router_gemm
def test_rocm_set_out_dtype_enables_fused_gemm(monkeypatch):
gate = _make_gate(monkeypatch, is_rocm=True, bias=False, out_dtype=None)
assert not gate.allow_cublas_router_gemm
gate.set_out_dtype(torch.float32)
assert gate.allow_cublas_router_gemm
def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch):
gate = _make_gate(monkeypatch, is_rocm=True, bias=True, out_dtype=None)
gate.set_out_dtype(torch.float32)
assert not gate.allow_cublas_router_gemm
@pytest.mark.parametrize(
("input_size", "output_size"),
[(3072, 256), (4096, 8), (6144, 128), (2048, 64), (2560, 512), (7168, 384)],
)
def test_rocm_gfx950_enables_fp32_router_gemm(
monkeypatch, input_size: int, output_size: int
) -> None:
gate = _make_gate(
monkeypatch,
is_rocm=True,
params_dtype=torch.float32,
input_size=input_size,
output_size=output_size,
on_gfx950=True,
)
assert gate.allow_fp32_router_gemm
assert not gate.allow_cublas_router_gemm
@pytest.mark.parametrize("ep_size", [2, 4, 8])
def test_rocm_fp32_router_gemm_is_replicated_across_ep_sizes(
monkeypatch, ep_size: int
) -> None:
gate = _make_gate(
monkeypatch,
is_rocm=True,
params_dtype=torch.float32,
input_size=6144,
output_size=128,
on_gfx950=True,
parallel_world_size=ep_size,
)
assert gate.weight.shape == (128, 6144)
assert gate.allow_fp32_router_gemm
@pytest.mark.parametrize(
("input_size", "output_size", "params_dtype", "bias", "on_gfx950"),
[
pytest.param(6144, 128, torch.float32, False, False, id="gfx942"),
pytest.param(6144, 128, torch.bfloat16, False, True, id="bf16-weight"),
pytest.param(6144, 128, torch.float32, True, True, id="bias"),
],
)
def test_rocm_fp32_router_gemm_rejects_unsupported_configs(
monkeypatch,
input_size: int,
output_size: int,
params_dtype: torch.dtype,
bias: bool,
on_gfx950: bool,
) -> None:
gate = _make_gate(
monkeypatch,
is_rocm=True,
params_dtype=params_dtype,
bias=bias,
input_size=input_size,
output_size=output_size,
on_gfx950=on_gfx950,
)
assert not gate.allow_fp32_router_gemm
# ---------------------------------------------------------------------------
# ROCm BF16x3 tier (fp32 router weights)
# ---------------------------------------------------------------------------
def _make_bf16x3_gate(monkeypatch, *, supported=True, **kwargs):
"""Build a gate with the ROCm bf16x3 platform check mocked."""
from vllm.model_executor.layers.fused_moe.router import (
bf16x3_router_gemm_rocm as rocm_bf16x3,
)
monkeypatch.setattr(rocm_bf16x3, "platform_supported", lambda: supported)
kwargs.setdefault("params_dtype", torch.float32)
return _make_gate(monkeypatch, is_rocm=True, **kwargs)
def test_rocm_bf16x3_enabled_for_fp32_weight(monkeypatch):
gate = _make_bf16x3_gate(monkeypatch)
assert gate.allow_rocm_bf16x3_router_gemm
# Built by process_weights_after_loading, not lazily in forward().
# Registered up front but still None, so named_buffers() skips it.
assert gate._bf16x3_weight is None
assert "_bf16x3_weight" in gate._buffers
def test_rocm_bf16x3_disabled_by_bias(monkeypatch):
# torch.mm has no bias term.
gate = _make_bf16x3_gate(monkeypatch, bias=True)
assert not gate.allow_rocm_bf16x3_router_gemm
def test_rocm_bf16x3_disabled_for_bf16_weight(monkeypatch):
# A bf16 weight belongs to the cuBLAS tier; there is nothing to split.
gate = _make_bf16x3_gate(monkeypatch, params_dtype=torch.bfloat16)
assert not gate.allow_rocm_bf16x3_router_gemm
def test_rocm_bf16x3_disabled_for_non_fp32_out_dtype(monkeypatch):
gate = _make_bf16x3_gate(monkeypatch, out_dtype=torch.bfloat16)
assert not gate.allow_rocm_bf16x3_router_gemm
def test_rocm_bf16x3_disabled_on_unsupported_arch(monkeypatch):
# Only gfx950 is benchmarked; gfx942 has a different fp32:bf16 MFMA ratio.
gate = _make_bf16x3_gate(monkeypatch, supported=False)
assert not gate.allow_rocm_bf16x3_router_gemm
assert not gate._rocm_bf16x3_weight_eligible
def test_rocm_bf16x3_no_split_built_when_tier_disabled(monkeypatch):
"""An ineligible out_dtype must not allocate a split that is never read;
it is 1.5x the router weight on every layer."""
gate = _make_bf16x3_gate(monkeypatch, out_dtype=torch.bfloat16)
gate.quant_method.process_weights_after_loading(gate)
assert gate._bf16x3_weight is None
def test_rocm_bf16x3_set_out_dtype_enables_tier(monkeypatch):
gate = _make_bf16x3_gate(monkeypatch, out_dtype=None)
assert not gate.allow_rocm_bf16x3_router_gemm
gate.set_out_dtype(torch.float32)
assert gate.allow_rocm_bf16x3_router_gemm
def test_rocm_bf16x3_dispatch_is_a_custom_op():
"""The num_tokens branch must stay opaque to Dynamo; see the op's impl."""
assert hasattr(torch.ops.vllm, "rocm_bf16x3_router_gemm_dispatch")
x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device="meta")
w = torch.zeros(128, 6144, dtype=torch.float32, device="meta")
split = torch.zeros(3, 128, 6144, dtype=torch.bfloat16, device="meta")
out = torch.ops.vllm.rocm_bf16x3_router_gemm_dispatch(x, w, split)
assert out.shape == (8192, 128)
assert out.dtype == torch.float32
def test_rocm_bf16x3_is_reachable_for_low_m_shapes(monkeypatch):
"""The gfx950 low-M fp32 tier claims these shapes too and returns
unconditionally, so bf16x3 is only reachable if checked first."""
gate = _make_bf16x3_gate(
monkeypatch, input_size=6144, output_size=128, on_gfx950=True
)
assert gate.allow_fp32_router_gemm
assert gate.allow_rocm_bf16x3_router_gemm
source = inspect.getsource(type(gate).forward)
assert source.index("allow_rocm_bf16x3_router_gemm") < source.index(
"allow_fp32_router_gemm"
)
tests/kernels/turboquant/__init__.py
View File
tests/lora/test_add_lora.py
-112
View File
@@ -1,112 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import time
import pytest
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.entrypoints.launchers.api_server.entry import (
build_async_engine_client_from_engine_args,
)
from vllm.inputs import TextPrompt
from vllm.lora.request import LoRARequest
from vllm.sampling_params import SamplingParams
from vllm.utils.async_utils import merge_async_iterators
MODEL_PATH = "zai-org/chatglm3-6b"
LORA_RANK = 64
DEFAULT_MAX_LORAS = 4 * 3
def get_lora_requests(lora_path) -> list[LoRARequest]:
lora_requests: list[LoRARequest] = [
LoRARequest(lora_name=f"{i}", lora_int_id=i, lora_path=lora_path)
for i in range(1, DEFAULT_MAX_LORAS + 1)
]
return lora_requests
async def requests_processing_time(llm, lora_requests: list[LoRARequest]) -> float:
sampling_params = SamplingParams(
n=1, temperature=0.0, top_p=1.0, ignore_eos=True, max_tokens=1
)
generators = []
start = time.perf_counter()
for lora_request in lora_requests:
lora_int_id = lora_request.lora_int_id
generator = llm.generate(
prompt=TextPrompt(prompt=f"hello {lora_int_id}", multi_modal_data=None), # type: ignore
sampling_params=sampling_params,
lora_request=lora_request,
request_id=f"test{lora_int_id}",
)
generators.append(generator)
all_gens = merge_async_iterators(*generators)
async for i, res in all_gens:
pass
end = time.perf_counter()
return end - start
@pytest.mark.asyncio
async def test_add_lora(chatglm3_lora_files):
"""The add_lora function is used to preload some LoRA adapters into the
engine in anticipation of future requests using these adapters. To test
this functionality, we use the async engine to process some requests - We
do it twice, once with add_lora() preloading and once without.
We measure the request processing time in both cases and expect the time
to be lesser in the case with add_lora() calls.
"""
lora_requests: list[LoRARequest] = get_lora_requests(chatglm3_lora_files)
max_loras = len(set([lr.lora_int_id for lr in lora_requests]))
# Create engine in eager-mode. Due to high max_loras, the CI can
# OOM during cuda-graph capture.
engine_args = AsyncEngineArgs(
model=MODEL_PATH,
enable_lora=True,
max_loras=max_loras,
max_lora_rank=LORA_RANK,
max_model_len=128,
gpu_memory_utilization=0.8, # avoid OOM
trust_remote_code=True,
enforce_eager=True,
)
# split lora_requests into 3 parts
part_size = len(lora_requests) // 3
dummy_run_requests = lora_requests[:part_size]
warmup_run_requests = lora_requests[part_size : part_size * 2]
cold_run_requests = lora_requests[part_size * 2 :]
async with build_async_engine_client_from_engine_args(engine_args) as llm:
# Dummy run - So any 1-time functionality like triton kernel compilation
# is complete here.
await requests_processing_time(llm, dummy_run_requests)
# Run with warmup
add_lora_tasks = [llm.add_lora(lr) for lr in warmup_run_requests]
add_lora_results = await asyncio.gather(*add_lora_tasks)
# Test that all all_lora calls are successful.
assert all(add_lora_results)
time_with_add_lora = await requests_processing_time(llm, warmup_run_requests)
# Run without any warmup
time_cold_start = await requests_processing_time(llm, cold_run_requests)
print(f"time hot-start {time_with_add_lora} vs time cold-start {time_cold_start} ")
assert time_with_add_lora < time_cold_start, (
f"time_with_add_lora={time_with_add_lora}, "
f"time_cold_start={time_cold_start}"
"The engine request processing time with LoRA pre-loading "
"must be less than the version that does on-demand LoRA loading."
)
tests/model_executor/model_loader/fastsafetensors_loader/__init__.py
View File
tests/model_executor/model_loader/instanttensor_loader/__init__.py
View File
tests/model_executor/model_loader/tensorizer_loader/__init__.py
View File
tests/model_executor/model_loader/tensorizer_loader/conftest.py
-92
View File
@@ -1,92 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
import pytest
from vllm import LLM, EngineArgs
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.model_executor.model_loader import tensorizer as tensorizer_mod
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
from vllm.utils.network_utils import get_file_store_init_method
from vllm.v1.executor import UniProcExecutor
from vllm.v1.worker.worker_base import WorkerWrapperBase
MODEL_REF = "facebook/opt-125m"
@pytest.fixture()
def model_ref():
return MODEL_REF
@pytest.fixture(autouse=True)
def allow_insecure_serialization(monkeypatch):
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
@pytest.fixture(autouse=True)
def cleanup():
cleanup_dist_env_and_memory(shutdown_ray=True)
@pytest.fixture()
def just_serialize_model_tensors(model_ref, monkeypatch, tmp_path):
def noop(*args, **kwargs):
return None
args = EngineArgs(model=model_ref)
tc = TensorizerConfig(tensorizer_uri=f"{tmp_path}/model.tensors")
monkeypatch.setattr(tensorizer_mod, "serialize_extra_artifacts", noop)
tensorizer_mod.tensorize_vllm_model(args, tc)
yield tmp_path
@pytest.fixture(autouse=True)
def tensorizer_config():
config = TensorizerConfig(tensorizer_uri="vllm")
return config
@pytest.fixture()
def model_path(model_ref, tmp_path):
yield tmp_path / model_ref / "model.tensors"
def assert_from_collective_rpc(engine: LLM, closure: Callable, closure_kwargs: dict):
res = engine.collective_rpc(method=closure, kwargs=closure_kwargs)
return all(res)
# This is an object pulled from tests/v1/engine/test_engine_core.py
# Modified to strip the `load_model` method from its `_init_executor`
# method. It's purely used as a dummy utility to run methods that test
# Tensorizer functionality
class DummyExecutor(UniProcExecutor):
def _init_executor(self) -> None:
"""Initialize the worker and load the model."""
self.driver_worker = WorkerWrapperBase(rpc_rank=0)
distributed_init_method = get_file_store_init_method()
local_rank = 0
# set local rank as the device index if specified
device_info = self.vllm_config.device_config.device.__str__().split(":")
if len(device_info) > 1:
local_rank = int(device_info[1])
rank = 0
is_driver_worker = True
kwargs = dict(
vllm_config=self.vllm_config,
local_rank=local_rank,
rank=rank,
distributed_init_method=distributed_init_method,
is_driver_worker=is_driver_worker,
)
self.mm_receiver_cache = None
self.collective_rpc("init_worker", args=([kwargs],))
self.collective_rpc("init_device")
def shutdown(self):
if hasattr(self, "thread_pool"):
self.thread_pool.shutdown(wait=False)
tests/models/glm5next/__init__.py
-2
View File
@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/models/inkling/rocm/conftest.py
-17
View File
@@ -1,17 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.platforms import current_platform
@pytest.fixture(autouse=True)
def require_rocm_cdna4() -> None:
if not current_platform.is_rocm():
pytest.skip("requires ROCm")
from vllm.platforms.rocm import on_gfx950
if not on_gfx950():
pytest.skip("requires CDNA4")
tests/models/language/generation/__init__.py
View File
tests/models/quantization/__init__.py
View File
tests/models/quantization/conftest.py
-41
View File
@@ -1,41 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared BF16 baselines for the quantization comparison tests.
Several tests here quantize the same model and compare against an
unquantized run of it. Computing that reference once per session avoids
booting an identical engine for every quantization scheme under test.
Only the outputs are cached; each engine is still torn down as usual, so
tests keep the isolation ``VllmRunner.__exit__`` provides.
"""
import pytest
@pytest.fixture(scope="session")
def quant_baseline_logprobs(vllm_runner):
"""Greedy logprobs for an unquantized model, computed once per config."""
cache: dict = {}
def baseline(
model: str,
prompts: list[str],
*,
max_model_len: int,
max_tokens: int,
num_logprobs: int,
):
key = (model, max_model_len, max_tokens, num_logprobs, tuple(prompts))
if key not in cache:
with vllm_runner(
model,
max_model_len=max_model_len,
enforce_eager=True,
) as vllm_model:
cache[key] = vllm_model.generate_greedy_logprobs(
prompts, max_tokens, num_logprobs
)
return cache[key]
return baseline
tests/models/transformers/__init__.py
View File
tests/models/transformers/fusers/__init__.py
View File
tests/plugins/vllm_add_dummy_model/setup.py
-13
View File
@@ -1,13 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from setuptools import setup
setup(
name="vllm_add_dummy_model",
version="0.1",
packages=["vllm_add_dummy_model"],
entry_points={
"vllm.general_plugins": ["register_dummy_model = vllm_add_dummy_model:register"]
},
)
tests/plugins/vllm_add_dummy_platform/setup.py
-18
View File
@@ -1,18 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from setuptools import setup
setup(
name="vllm_add_dummy_platform",
version="0.1",
packages=["vllm_add_dummy_platform"],
entry_points={
"vllm.platform_plugins": [
"dummy_platform_plugin = vllm_add_dummy_platform:dummy_platform_plugin" # noqa
],
"vllm.general_plugins": [
"dummy_custom_ops = vllm_add_dummy_platform:register_ops"
],
},
)
tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py
-13
View File
@@ -1,13 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
def dummy_platform_plugin() -> str | None:
# Real platform plugins often import vLLM modules while registering.
from vllm.utils.torch_utils import direct_register_custom_op # noqa: F401
return "vllm_add_dummy_platform.dummy_platform.DummyPlatform"
def register_ops():
import vllm_add_dummy_platform.dummy_custom_ops # noqa
tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_attention_backend.py
-10
View File
@@ -1,10 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
class DummyAttentionBackend(TritonAttentionBackend):
@staticmethod
def get_name() -> str:
return "Dummy_Backend"
tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py
-32
View File
@@ -1,32 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import TYPE_CHECKING
from vllm.platforms.interface import Platform, PlatformEnum
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.attention.selector import AttentionSelectorConfig
else:
VllmConfig = None
class DummyPlatform(Platform):
_enum = PlatformEnum.OOT
device_name = "DummyDevice"
device_type: str = "privateuseone"
dispatch_key: str = "PrivateUse1"
@classmethod
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
vllm_config.compilation_config.custom_ops = ["all"]
@classmethod
def get_attn_backend_cls(
cls,
selected_backend: "AttentionBackendEnum",
attn_selector_config: "AttentionSelectorConfig",
num_heads: int | None = None,
) -> str:
return "vllm_add_dummy_platform.dummy_attention_backend.DummyAttentionBackend" # noqa E501
tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py
-30
View File
@@ -1,30 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.v1.metrics.loggers import StatLoggerBase
class DummyStatLogger(StatLoggerBase):
"""A dummy stat logger for testing purposes.
Implements the minimal interface expected by StatLoggerManager.
"""
def __init__(self, vllm_config, engine_idx=0):
self.vllm_config = vllm_config
self.engine_idx = engine_idx
self.recorded = []
self.logged = False
self.engine_initialized = False
def record(
self, scheduler_stats, iteration_stats, mm_cache_stats=None, engine_idx=0
):
self.recorded.append(
(scheduler_stats, iteration_stats, mm_cache_stats, engine_idx)
)
def log(self):
self.logged = True
def log_engine_initialized(self):
self.engine_initialized = True
tests/plugins/vllm_add_dummy_stat_logger/setup.py
-15
View File
@@ -1,15 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from setuptools import setup
setup(
name="dummy_stat_logger",
version="0.1",
packages=["dummy_stat_logger"],
entry_points={
"vllm.stat_logger_plugins": [
"dummy_stat_logger = dummy_stat_logger.dummy_stat_logger:DummyStatLogger" # noqa
]
},
)
tests/plugins_tests/bitsandbytes/__init__.py
-2
View File
@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/plugins_tests/test_scheduler_plugins.py
-36
View File
@@ -1,36 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.engine.arg_utils import EngineArgs
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.engine.llm_engine import LLMEngine
class DummyV1Scheduler(Scheduler):
def schedule(self, throttle_prefills: bool = False):
raise Exception("Exception raised by DummyV1Scheduler")
def test_scheduler_plugins_v1(monkeypatch: pytest.MonkeyPatch):
with monkeypatch.context() as m:
# Explicitly turn off engine multiprocessing so
# that the scheduler runs in this process
m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
with pytest.raises(Exception) as exception_info:
engine_args = EngineArgs(
model="facebook/opt-125m",
enforce_eager=True, # reduce test time
scheduler_cls=DummyV1Scheduler,
)
engine = LLMEngine.from_engine_args(engine_args=engine_args)
sampling_params = SamplingParams(max_tokens=1)
engine.add_request("0", "foo", sampling_params)
engine.step()
assert str(exception_info.value) == "Exception raised by DummyV1Scheduler"
tests/plugins_tests/test_stats_logger_plugins.py
-78
View File
@@ -1,78 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from dummy_stat_logger.dummy_stat_logger import DummyStatLogger
from vllm.config import VllmConfig
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.metrics.loggers import load_stat_logger_plugin_factories
def test_stat_logger_plugin_is_discovered(monkeypatch: pytest.MonkeyPatch):
with monkeypatch.context() as m:
m.setenv("VLLM_PLUGINS", "dummy_stat_logger")
factories = load_stat_logger_plugin_factories()
assert len(factories) == 1, f"Expected 1 factory, got {len(factories)}"
assert factories[0] is DummyStatLogger, (
f"Expected DummyStatLogger class, got {factories[0]}"
)
# instantiate and confirm the right type
vllm_config = VllmConfig()
instance = factories[0](vllm_config)
# The `is` check above narrows DummyStatLogger to a Callable
assert isinstance(instance, DummyStatLogger)
def test_no_plugins_loaded_if_env_empty(monkeypatch: pytest.MonkeyPatch):
with monkeypatch.context() as m:
m.setenv("VLLM_PLUGINS", "")
factories = load_stat_logger_plugin_factories()
assert factories == []
def test_invalid_stat_logger_plugin_raises(monkeypatch: pytest.MonkeyPatch):
def fake_plugin_loader(group: str):
assert group == "vllm.stat_logger_plugins"
return {"bad": object()}
with monkeypatch.context() as m:
m.setattr(
"vllm.v1.metrics.loggers.load_plugins_by_group",
fake_plugin_loader,
)
with pytest.raises(
TypeError,
match="Stat logger plugin 'bad' must be a subclass of StatLoggerBase",
):
load_stat_logger_plugin_factories()
@pytest.mark.asyncio
async def test_stat_logger_plugin_integration_with_engine(
monkeypatch: pytest.MonkeyPatch,
):
with monkeypatch.context() as m:
m.setenv("VLLM_PLUGINS", "dummy_stat_logger")
engine_args = AsyncEngineArgs(
model="facebook/opt-125m",
enforce_eager=True, # reduce test time
disable_log_stats=True, # disable default loggers
)
engine = AsyncLLM.from_engine_args(engine_args=engine_args)
assert engine.logger_manager is not None
assert len(engine.logger_manager.stat_loggers) == 2
assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1 # type: ignore[attr-defined]
assert isinstance(
engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0], # type: ignore[attr-defined]
DummyStatLogger,
)
engine.shutdown()
tests/samplers/__init__.py
View File
tests/samplers/conftest.py
-21
View File
@@ -1,21 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm import envs
from vllm.platforms import current_platform
@pytest.fixture(autouse=True)
def skip_unsupported_flashinfer_sampler():
# CI runs this suite with FlashInfer both explicitly enabled and disabled.
if (
current_platform.is_cuda()
and envs.is_set("VLLM_USE_FLASHINFER_SAMPLER")
and envs.VLLM_USE_FLASHINFER_SAMPLER
and current_platform.num_compute_units(torch.accelerator.current_device_index())
<= 16
):
pytest.skip("FlashInfer top-k masking requires more than 16 SMs")
tests/test_config.py
-39
View File
@@ -109,45 +109,6 @@ def _write_json(path: Path, value: object) -> None:
path.write_text(json.dumps(value), encoding="utf-8")
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test")
@pytest.mark.parametrize(
("is_mm_prefix_lm", "is_multimodal_model", "expected"),
[
pytest.param(True, True, True, id="multimodal-prefix-lm"),
pytest.param(False, True, False, id="multimodal-causal"),
pytest.param(True, False, False, id="text-prefix-lm"),
pytest.param(None, True, False, id="missing-model-config"),
],
)
def test_rocm_mm_prefix_lm_disables_chunked_mm_input(
is_mm_prefix_lm: bool | None,
is_multimodal_model: bool,
expected: bool,
) -> None:
from vllm.platforms.rocm import RocmPlatform
config = SimpleNamespace(
compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE),
parallel_config=SimpleNamespace(
prefill_context_parallel_size=1,
worker_cls="test-worker",
),
model_config=(
None
if is_mm_prefix_lm is None
else SimpleNamespace(is_mm_prefix_lm=is_mm_prefix_lm)
),
scheduler_config=SimpleNamespace(
is_multimodal_model=is_multimodal_model,
disable_chunked_mm_input=False,
),
)
RocmPlatform.check_and_update_config(config)
assert config.scheduler_config.disable_chunked_mm_input is expected
def _sampling_replay_config(
*,
return_sampling_mask: bool = True,
tests/tokenizers_/__init__.py
-4
View File
@@ -1,4 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# NOTE: Since CI runs the tests from the `tests` directory, it is necessary to rename
# this module to avoid conflicting with HF's `tokenizers` package
tests/tokenizers_/fixtures/deepseek_v4/test_input_1.json
-81
View File
@@ -1,81 +0,0 @@
{
"tools": [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the weather for a specific location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city name"
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
"description": "Temperature unit"
}
},
"required": ["location"]
}
}
},
{
"type": "function",
"function": {
"name": "search",
"description": "Search the web for information",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search query"
},
"num_results": {
"type": "integer",
"description": "Number of results to return"
}
},
"required": ["query"]
}
}
}
],
"messages": [
{
"role": "system",
"content": "You are a helpful assistant."
},
{
"role": "user",
"content": "What's the weather in Beijing?"
},
{
"role": "assistant",
"reasoning": "The user wants to know the weather in Beijing. I should use the get_weather tool.",
"tool_calls": [
{
"id": "call_001",
"type": "function",
"function": {
"name": "get_weather",
"arguments": "{\"location\": \"Beijing\", \"unit\": \"celsius\"}"
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_001",
"content": "{\"temperature\": 22, \"condition\": \"sunny\", \"humidity\": 45}"
},
{
"role": "assistant",
"reasoning": "Got the weather data. Let me format a nice response.",
"content": "The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity."
}
]
}
tests/tokenizers_/fixtures/deepseek_v4/test_input_2.json
-24
View File
@@ -1,24 +0,0 @@
[
{
"role": "system",
"content": "You are a helpful assistant."
},
{
"role": "user",
"content": "Hello"
},
{
"role": "assistant",
"reasoning": "The user said hello, I should greet back.",
"content": "Hi there! How can I help you?"
},
{
"role": "user",
"content": "What is the capital of France?"
},
{
"role": "assistant",
"reasoning": "The user asks about the capital of France. It is Paris.",
"content": "The capital of France is Paris."
}
]
tests/tokenizers_/fixtures/deepseek_v4/test_input_3.json
-159
View File
@@ -1,159 +0,0 @@
[
{
"role": "system",
"content": "该助手为DeepSeek，由深度求索公司创造。"
},
{
"role": "latest_reminder",
"content": "2026-02-21,星期六,广州,App,中文"
},
{
"role": "developer",
"content": "小柴胡冲剂和布洛芬能一起吃吗？\n\nCITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】",
"tools": [
{
"type": "function",
"function": {
"name": "search",
"description": "Web search. Split multiple queries with '||'.",
"parameters": {
"type": "object",
"properties": {
"queries": {
"type": "string",
"description": "query1||query2"
}
},
"required": [
"queries"
],
"additionalProperties": false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
}
},
{
"type": "function",
"function": {
"name": "open",
"description": "Batch open IDs (format 【{id}†...】) or URLs.",
"parameters": {
"type": "object",
"properties": {
"open_list": {
"type": "array",
"items": {
"type": "object",
"properties": {
"id": {
"description": "ID or URL",
"anyOf": [
{
"type": "integer"
},
{
"type": "string"
}
],
"default": -1
},
"cursor": {
"type": "integer",
"description": "",
"default": -1
},
"loc": {
"type": "integer",
"description": "Start line",
"default": -1
},
"num_lines": {
"type": "integer",
"description": "",
"default": -1
},
"view_source": {
"type": "boolean",
"description": "",
"default": false
}
},
"additionalProperties": false
},
"description": ""
}
},
"required": [
"open_list"
],
"additionalProperties": false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
}
},
{
"type": "function",
"function": {
"name": "find",
"description": "Find exact text pattern in pages.",
"parameters": {
"type": "object",
"properties": {
"find_list": {
"type": "array",
"items": {
"type": "object",
"properties": {
"pattern": {
"type": "string",
"description": ""
},
"cursor": {
"type": "integer",
"description": "",
"default": -1
}
},
"required": [
"pattern"
],
"additionalProperties": false
},
"description": ""
}
},
"required": [
"find_list"
],
"additionalProperties": false,
"$schema": "http://json-schema.org/draft-07/schema#"
}
}
}
]
},
{
"role": "assistant",
"content": "",
"reasoning": "用户想知道小柴胡冲剂和布洛芬能否一起服用。",
"tool_calls": [
{
"type": "function",
"function": {
"name": "search",
"arguments": "{\"queries\": \"小柴胡冲剂 布洛芬 相互作用 一起吃\"}"
}
}
]
},
{
"role": "tool",
"content": "[0]"
},
{
"role": "assistant",
"content": "请及时就医。",
"reasoning": "现在开始组织回答。",
"tool_calls": []
}
]
tests/tokenizers_/fixtures/deepseek_v4/test_input_4.json
-28
View File
@@ -1,28 +0,0 @@
[
{
"role": "system",
"content": "该助手为DeepSeek-V3，由深度求索公司创造。\n今天是2025年10月17日，星期五。"
},
{
"role": "latest_reminder",
"content": "2024-11-15,上海市,App,中文"
},
{
"role": "user",
"content": "热海大滚锅是世界著名温泉吗"
},
{
"role": "assistant",
"content": "关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。\n\n为了让你快速了解，我整理了一个简要的表格：\n\n| 维度 | 热海大滚锅的地位与特点 |\n| :--- | :--- |\n| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |\n| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |\n| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |\n| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |\n\n### 💡 游玩攻略与温馨提示\n\n如果你计划前往热海大滚锅，这里有一些实用信息供你参考：\n\n- **门票与开放时间**：\n - **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。\n - **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。\n\n- **特色体验**：\n - **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。\n - **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。\n\n- **注意事项**：\n - **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。\n - **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。\n\n希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。",
"mask": 1
},
{
"role": "user",
"content": "世界著名温泉有哪些",
"task": "action"
},
{
"role": "assistant",
"content": "Search"
}
]
tests/tokenizers_/fixtures/deepseek_v4/test_input_5.json
-39
View File
@@ -1,39 +0,0 @@
{
"thinking_mode": "thinking",
"reasoning_effort": "max",
"messages": [
{
"role": "system",
"content": "You are a helpful vision assistant."
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "请按“第一张、第二张”的顺序回答：第一张图"
},
{
"type": "image_url",
"image_url": {
"url": "examples/images/carrots.jpeg"
}
},
{
"type": "text",
"text": "和第二张图"
},
{
"type": "image_url",
"image_url": {
"url": "examples/images/corn.jpeg"
}
},
{
"type": "text",
"text": "中分别是什么食材？它们通常食用的部位分别是什么？"
}
]
}
]
}
tests/tokenizers_/fixtures/deepseek_v4/test_output_1.txt
-36
View File
@@ -1,36 +0,0 @@
<｜begin▁of▁sentence｜>You are a helpful assistant.
## Tools
You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
<｜DSML｜tool_calls>
<｜DSML｜invoke name="$TOOL_NAME">
<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
...
</｜DSML｜invoke>
<｜DSML｜invoke name="$TOOL_NAME2">
...
</｜DSML｜invoke>
</｜DSML｜tool_calls>
String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
Otherwise, output directly after </think> with tool calls or final response.
### Available Tool Schemas
{"name": "get_weather", "description": "Get the weather for a specific location", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city name"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "Temperature unit"}}, "required": ["location"]}}
{"name": "search", "description": "Search the web for information", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Search query"}, "num_results": {"type": "integer", "description": "Number of results to return"}}, "required": ["query"]}}
You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
<｜User｜>What's the weather in Beijing?<｜Assistant｜><think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>
<｜DSML｜tool_calls>
<｜DSML｜invoke name="get_weather">
<｜DSML｜parameter name="location" string="true">Beijing</｜DSML｜parameter>
<｜DSML｜parameter name="unit" string="true">celsius</｜DSML｜parameter>
</｜DSML｜invoke>
</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>{"temperature": 22, "condition": "sunny", "humidity": 45}</tool_result><｜Assistant｜><think>Got the weather data. Let me format a nice response.</think>The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.<｜end▁of▁sentence｜>
tests/tokenizers_/fixtures/deepseek_v4/test_output_2.txt
-1
View File
@@ -1 +0,0 @@
<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>Hello<｜Assistant｜></think>Hi there! How can I help you?<｜end▁of▁sentence｜><｜User｜>What is the capital of France?<｜Assistant｜><think>The user asks about the capital of France. It is Paris.</think>The capital of France is Paris.<｜end▁of▁sentence｜>
tests/tokenizers_/fixtures/deepseek_v4/test_output_3.txt
-38
View File
@@ -1,38 +0,0 @@
<｜begin▁of▁sentence｜>该助手为DeepSeek，由深度求索公司创造。<｜latest_reminder｜>2026-02-21,星期六,广州,App,中文<｜User｜>小柴胡冲剂和布洛芬能一起吃吗？
CITATION FORMAT: 【{cursor_id}†L{start_line_id}(-L{end_line_id})?】
## Tools
You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:
<｜DSML｜tool_calls>
<｜DSML｜invoke name="$TOOL_NAME">
<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>
...
</｜DSML｜invoke>
<｜DSML｜invoke name="$TOOL_NAME2">
...
</｜DSML｜invoke>
</｜DSML｜tool_calls>
String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
If thinking_mode is enabled (triggered by <think>), you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
Otherwise, output directly after </think> with tool calls or final response.
### Available Tool Schemas
{"name": "search", "description": "Web search. Split multiple queries with '||'.", "parameters": {"type": "object", "properties": {"queries": {"type": "string", "description": "query1||query2"}}, "required": ["queries"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
{"name": "open", "description": "Batch open IDs (format 【{id}†...】) or URLs.", "parameters": {"type": "object", "properties": {"open_list": {"type": "array", "items": {"type": "object", "properties": {"id": {"description": "ID or URL", "anyOf": [{"type": "integer"}, {"type": "string"}], "default": -1}, "cursor": {"type": "integer", "description": "", "default": -1}, "loc": {"type": "integer", "description": "Start line", "default": -1}, "num_lines": {"type": "integer", "description": "", "default": -1}, "view_source": {"type": "boolean", "description": "", "default": false}}, "additionalProperties": false}, "description": ""}}, "required": ["open_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
{"name": "find", "description": "Find exact text pattern in pages.", "parameters": {"type": "object", "properties": {"find_list": {"type": "array", "items": {"type": "object", "properties": {"pattern": {"type": "string", "description": ""}, "cursor": {"type": "integer", "description": "", "default": -1}}, "required": ["pattern"], "additionalProperties": false}, "description": ""}}, "required": ["find_list"], "additionalProperties": false, "$schema": "http://json-schema.org/draft-07/schema#"}}
You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
<｜Assistant｜><think>用户想知道小柴胡冲剂和布洛芬能否一起服用。</think>
<｜DSML｜tool_calls>
<｜DSML｜invoke name="search">
<｜DSML｜parameter name="queries" string="true">小柴胡冲剂 布洛芬 相互作用 一起吃</｜DSML｜parameter>
</｜DSML｜invoke>
</｜DSML｜tool_calls><｜end▁of▁sentence｜><｜User｜><tool_result>[0]</tool_result><｜Assistant｜><think>现在开始组织回答。</think>请及时就医。<｜end▁of▁sentence｜>
tests/tokenizers_/fixtures/deepseek_v4/test_output_4.txt
-29
View File
@@ -1,29 +0,0 @@
<｜begin▁of▁sentence｜>该助手为DeepSeek-V3，由深度求索公司创造。
今天是2025年10月17日，星期五。<｜latest_reminder｜>2024-11-15,上海市,App,中文<｜User｜>热海大滚锅是世界著名温泉吗<｜Assistant｜></think>关于热海大滚锅是否是世界著名温泉，可以这样总结：它在**中国乃至全球的地热奇观中占有重要地位**，但“世界著名”的称号更侧重于它作为**独特的地质现象和旅游景点**，而非普遍的温泉疗养体验。
为了让你快速了解，我整理了一个简要的表格：
| 维度 | 热海大滚锅的地位与特点 |
| :--- | :--- |
| **地质奇观** | **世界罕见**的高温地热系统，泉眼多、水温高、形态多样。 |
| **国内声誉** | **中国三大地热区之一**，**国家5A级旅游景区**，自古闻名（徐霞客曾游历并记载）。 |
| **国际知名度** | 在一些旅游平台被提及为“世界六大温泉”之一，但此说法流传不广，其国际声誉更多建立在地质独特性上。 |
| **核心体验** | **观赏地热奇观**（如97℃沸腾的“大滚锅”）、**体验温泉煮鸡蛋**。 |
### 💡 游玩攻略与温馨提示
如果你计划前往热海大滚锅，这里有一些实用信息供你参考：
- **门票与开放时间**：
- **门票**：景区门票约为**50元/人**。如果选择包含温泉沐浴的套餐，价格会更高，例如约**288元**。
- **开放时间**：景区一般**08:00-18:00**开放，但具体时间可能变动，建议提前核实。
- **特色体验**：
- **温泉煮鸡蛋**：这几乎是必试项目。可以在景区门口购买用草绳串起的生鸡蛋（约5-8元/串），然后到“大滚锅”旁的指定区域蒸煮，几分钟便可熟食，趣味十足。
- **金汤足浴**：可以直接用从“大滚锅”流出的温泉水泡脚，缓解旅途疲劳。
- **注意事项**：
- **安全第一**：“大滚锅”水温极高，务必遵守游览规则，在指定区域内观赏，切勿随意触碰泉水。
- **规划行程**：建议为热海景区预留**3-4小时**的游览时间。景区内步道不走回头路，出入口有观光车接送。
希望这些信息能帮助你更好地了解热海大滚锅。如果你对腾冲的其他景点或者行程规划有更多疑问，我很乐意提供进一步的信息。<｜end▁of▁sentence｜><｜User｜>世界著名温泉有哪些<｜Assistant｜></think><｜action｜>Search<｜end▁of▁sentence｜>

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