minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Kernel] Move sgl-kernel under sglang.kernels.aot (#32648)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-07-29 17:25:00 +08:00
committed by GitHub
parent 1b9dfa14e6
commit c32c4ef79c
370 changed files with 300 additions and 269 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/MANIFEST.in
+1
View File
@@ -0,0 +1 @@
prune sglang/kernels/aot
python/pyproject.toml
+8
View File
@@ -199,6 +199,12 @@ killall_sglang = "sglang.cli.killall:main"
"multimodal_gen/apps/realtime_webui/**/*"
]
[tool.setuptools.exclude-package-data]
"sglang" = [
"kernels/aot/*",
"kernels/aot/**/*",
]
[tool.setuptools.packages.find]
exclude = [
"assets*",
@@ -207,6 +213,7 @@ exclude = [
"dist*",
"playground*",
"scripts*",
"sglang.kernels.aot*",
"tests*",
]
@@ -218,6 +225,7 @@ exclude = [
"dist*",
"playground*",
"scripts*",
"sglang/kernels/aot*",
"tests*",
]
python/sglang/kernels/aot/.clang-format
+15
View File
@@ -0,0 +1,15 @@
BasedOnStyle: Google
IndentWidth: 2
ColumnLimit: 120
AllowShortFunctionsOnASingleLine: Empty
DerivePointerAlignment: false
PointerAlignment: Left
NamespaceIndentation: None
SortIncludes: true
AllowShortLoopsOnASingleLine: false
BinPackParameters: false # Prevents packing parameters in declarations
BinPackArguments: false # Prevents packing arguments in function calls
AlignAfterOpenBracket: AlwaysBreak # Forces a break after the opening parenthesis
AlignOperands: Align # Aligns arguments vertically
PenaltyBreakBeforeFirstCallParameter: 1 # Encourages breaking before the first argument
PenaltyReturnTypeOnItsOwnLine: 100 # Keeps return type with function name
python/sglang/kernels/aot/CMakeLists.txt
+557
View File
@@ -0,0 +1,557 @@
cmake_minimum_required(VERSION 3.26 FATAL_ERROR)
project(sgl-kernel LANGUAGES CXX CUDA)
# utils
include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake)
include(FetchContent)
# CMake
cmake_policy(SET CMP0169 OLD)
cmake_policy(SET CMP0177 NEW)
set(CMAKE_COLOR_DIAGNOSTICS ON)
set(CMAKE_VERBOSE_MAKEFILE ON CACHE BOOL "ON")
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
set(CMAKE_SHARED_LIBRARY_PREFIX "")
# GitHub Artifactory
set(GITHUB_ARTIFACTORY "github.com" CACHE STRING "GitHub mirror URL")
# Python
find_package(Python COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT} REQUIRED)
# CXX
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3")
# CUDA
enable_language(CUDA)
find_package(CUDAToolkit REQUIRED)
set_property(GLOBAL PROPERTY CUDA_SEPARABLE_COMPILATION ON)
message(STATUS "Detected CUDA_VERSION=${CUDA_VERSION}")
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "13.0")
message("CUDA_VERSION ${CUDA_VERSION} >= 13.0")
elseif ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8")
message("CUDA_VERSION ${CUDA_VERSION} >= 12.8")
elseif ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4")
message("CUDA_VERSION ${CUDA_VERSION} >= 12.4")
elseif ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.1")
message("CUDA_VERSION ${CUDA_VERSION} >= 12.1")
elseif ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "11.8")
message("CUDA_VERSION ${CUDA_VERSION} >= 11.8")
endif()
# Torch
find_package(Torch REQUIRED)
clear_cuda_arches(CMAKE_FLAG)
# Third Party repos
# cutlass
FetchContent_Declare(
repo-cutlass
URL https://${GITHUB_ARTIFACTORY}/NVIDIA/cutlass/archive/57e3cfb47a2d9e0d46eb6335c3dc411498efa198.tar.gz
URL_HASH SHA256=09237099a70f80bff1dc8bb80c843a674bb4fdcb46e43cc6993e711c5ca89bb5
)
FetchContent_Populate(repo-cutlass)
# fmt
FetchContent_Declare(
repo-fmt
URL https://${GITHUB_ARTIFACTORY}/fmtlib/fmt/archive/553ec11ec06fbe0beebfbb45f9dc3c9eabd83d28.tar.gz
URL_HASH SHA256=c314292789d28c3c3b420e75a7b2d1706f685f7fb63289128d46aeaea2c6be71
)
FetchContent_Populate(repo-fmt)
# Triton kernel
FetchContent_Declare(
repo-triton
URL https://${GITHUB_ARTIFACTORY}/triton-lang/triton/archive/v3.6.0.tar.gz
URL_HASH SHA256=be270ed11ca5a8fbd9d7941c5bbe9a23a9f6e2ffd372c8398346928bee464774
)
FetchContent_Populate(repo-triton)
# flashinfer
FetchContent_Declare(
repo-flashinfer
URL https://${GITHUB_ARTIFACTORY}/flashinfer-ai/flashinfer/archive/bc29697ba20b7e6bdb728ded98f04788e16ee021.tar.gz
URL_HASH SHA256=931dfd118f4b6de8c7d98702153c7c03840139170af21a07607693bd9749744d
)
FetchContent_Populate(repo-flashinfer)
# flash-attention
FetchContent_Declare(
repo-flash-attention
URL https://${GITHUB_ARTIFACTORY}/sgl-project/sgl-attn/archive/f89bc2306632d1ec5f97b014dded4254f5b4a907.tar.gz
URL_HASH SHA256=418b5681584dc3efff496a1cab5ffd58d2728d89dcfe0ea16e6985d6ef35c68c
)
FetchContent_Populate(repo-flash-attention)
# ccache option
option(ENABLE_CCACHE "Whether to use ccache" ON)
find_program(CCACHE_FOUND ccache)
if(CCACHE_FOUND AND ENABLE_CCACHE AND DEFINED ENV{CCACHE_DIR})
message(STATUS "Building with CCACHE enabled")
set_property(GLOBAL PROPERTY RULE_LAUNCH_COMPILE "ccache")
set_property(GLOBAL PROPERTY RULE_LAUNCH_LINK "ccache")
endif()
# Configure gencode below SM90
if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
set(DEFAULT_ENABLE_BELOW_SM90 OFF)
message(STATUS "For aarch64, disable gencode below SM90 by default")
else()
set(DEFAULT_ENABLE_BELOW_SM90 ON)
endif()
option(ENABLE_BELOW_SM90 "Enable gencode below SM90" ${DEFAULT_ENABLE_BELOW_SM90})
set(DEFAULT_SGL_KERNEL_ENABLE_FA3 OFF)
if (CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$")
message(STATUS "For aarch64, disable FA3 by default")
endif()
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4" AND NOT CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$")
set(DEFAULT_SGL_KERNEL_ENABLE_FA3 ON)
endif()
include_directories(
${PROJECT_SOURCE_DIR}/include
${PROJECT_SOURCE_DIR}/csrc
)
set(SGL_KERNEL_CUDA_FLAGS
"-DNDEBUG"
"-DOPERATOR_NAMESPACE=sgl-kernel"
"-O3"
"-Xcompiler"
"-fPIC"
"-gencode=arch=compute_90,code=sm_90"
"-std=c++17"
"-DFLASHINFER_ENABLE_F16"
"-DCUTE_USE_PACKED_TUPLE=1"
"-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1"
"-DCUTLASS_VERSIONS_GENERATED"
"-DCUTLASS_TEST_LEVEL=0"
"-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1"
"-DCUTLASS_DEBUG_TRACE_LEVEL=0"
"--expt-relaxed-constexpr"
"--expt-extended-lambda"
# The following flag leads to the CMAKE_BUILD_PARALLEL_LEVEL breaking,
# it triggers OOM with low memory host. Extract the threads number to
# option named SGL_KERNEL_COMPILE_THREADS, default value 32.
# "--threads=32"
# Supress warnings
"-Xcompiler=-Wno-clang-format-violations"
"-Xcompiler=-Wno-conversion"
"-Xcompiler=-Wno-deprecated-declarations"
"-Xcompiler=-Wno-terminate"
"-Xcompiler=-Wfatal-errors"
"-Xcompiler=-ftemplate-backtrace-limit=1"
"-Xcudafe=--diag_suppress=177" # variable was declared but never referenced
"-Xcudafe=--diag_suppress=2361" # invalid narrowing conversion from "char" to "signed char"
# uncomment to debug
# "--ptxas-options=-v"
# "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage"
)
set(SGL_KERNEL_COMPILE_THREADS 32 CACHE STRING "Set compilation threads, default 32")
# When SGL_KERNEL_COMPILE_THREADS value is less than 1, set it to 1
if (NOT SGL_KERNEL_COMPILE_THREADS MATCHES "^[0-9]+$")
message(FATAL_ERROR "SGL_KERNEL_COMPILE_THREADS must be an integer, but was set to '${SGL_KERNEL_COMPILE_THREADS}'.")
elseif (SGL_KERNEL_COMPILE_THREADS LESS 1)
message(STATUS "SGL_KERNEL_COMPILE_THREADS was set to a value less than 1. Using 1 instead.")
set(SGL_KERNEL_COMPILE_THREADS 1)
endif()
list(APPEND SGL_KERNEL_CUDA_FLAGS
"--threads=${SGL_KERNEL_COMPILE_THREADS}"
)
option(SGL_KERNEL_ENABLE_BF16 "Enable BF16" ON)
option(SGL_KERNEL_ENABLE_FP8 "Enable FP8" ON)
option(SGL_KERNEL_ENABLE_FP4 "Enable FP4" OFF)
option(SGL_KERNEL_ENABLE_FA3 "Enable FA3" ${DEFAULT_SGL_KERNEL_ENABLE_FA3})
option(SGL_KERNEL_ENABLE_FA3_SPARSE_MASK "Enable FA3 sparse mask kernels" OFF)
option(SGL_KERNEL_ENABLE_SM90A "Enable SM90A" OFF)
option(SGL_KERNEL_ENABLE_SM100A "Enable SM100A" OFF)
if (SGL_KERNEL_ENABLE_BF16)
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-DFLASHINFER_ENABLE_BF16"
)
endif()
if (SGL_KERNEL_ENABLE_FP8)
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-DFLASHINFER_ENABLE_FP8"
"-DFLASHINFER_ENABLE_FP8_E4M3"
"-DFLASHINFER_ENABLE_FP8_E5M2"
)
endif()
if (ENABLE_BELOW_SM90)
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_80,code=sm_80"
"-gencode=arch=compute_89,code=sm_89"
)
if (CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64")
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_87,code=sm_87"
)
endif()
endif()
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A)
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_100a,code=sm_100a"
"-gencode=arch=compute_120a,code=sm_120a"
)
# refer sm_121, sm_110 and sm_101 description https://github.com/pytorch/pytorch/pull/156176
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "13.0")
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_103a,code=sm_103a"
"--compress-mode=size"
)
if (CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64")
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_110a,code=sm_110a"
"-gencode=arch=compute_121a,code=sm_121a"
)
endif()
else()
if (CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64")
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_101a,code=sm_101a"
)
endif()
endif()
endif()
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4" AND SGL_KERNEL_ENABLE_FA3)
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_90a,code=sm_90a"
)
endif()
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_FP4)
list(APPEND SGL_KERNEL_CUDA_FLAGS
"-DENABLE_NVFP4=1"
)
endif()
# All source files
# NOTE: Please sort the filenames alphabetically
set(SOURCES
"csrc/allreduce/custom_all_reduce.cu"
"csrc/attention/cutlass_mla_kernel.cu"
"csrc/attention/merge_attn_states.cu"
"csrc/attention/vertical_slash_index.cu"
"csrc/common_extension.cc"
"csrc/elementwise/activation.cu"
"csrc/elementwise/concat_mla.cu"
"csrc/elementwise/copy.cu"
"csrc/elementwise/dsv4_norm_rope.cu"
"csrc/elementwise/fused_add_rms_norm_kernel.cu"
"csrc/elementwise/pos_enc.cu"
"csrc/elementwise/topk.cu"
"csrc/expert_specialization/es_fp8_blockwise.cu"
"csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu"
"csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu"
"csrc/gemm/awq_kernel.cu"
"csrc/gemm/fp8_gemm_kernel.cu"
"csrc/gemm/int8_gemm_kernel.cu"
"csrc/gemm/per_token_group_quant_8bit.cu"
"csrc/gemm/per_token_group_quant_8bit_v2.cu"
"csrc/gemm/per_token_quant_fp8.cu"
"csrc/gemm/gptq/gptq_kernel.cu"
"csrc/grammar/apply_token_bitmask_inplace_cuda.cu"
"csrc/infllm_v2/max_pooling.cu"
"csrc/kvcacheio/transfer.cu"
"csrc/mamba/causal_conv1d.cu"
"csrc/memory/weak_ref_tensor.cpp"
"csrc/moe/cutlass_moe/w4a8/scaled_mm_entry.cu"
"csrc/moe/cutlass_moe/w4a8/w4a8_moe_data.cu"
"csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cu"
"csrc/moe/moe_align_kernel.cu"
"csrc/moe/fused_qknorm_rope_kernel.cu"
"csrc/moe/moe_sum.cu"
"csrc/moe/moe_sum_reduce.cu"
"csrc/moe/moe_topk_softmax_kernels.cu"
"csrc/moe/moe_topk_sigmoid_kernels.cu"
"csrc/moe/fp8_blockwise_moe_kernel.cu"
"csrc/moe/prepare_moe_input.cu"
"csrc/quantization/gguf/gguf_kernel.cu"
"csrc/speculative/eagle_utils.cu"
"csrc/speculative/ngram_utils.cu"
"csrc/speculative/packbit.cu"
"csrc/speculative/speculative_sampling.cu"
"${repo-flashinfer_SOURCE_DIR}/csrc/norm.cu"
"${repo-flashinfer_SOURCE_DIR}/csrc/renorm.cu"
"${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_causal_sm80.cu"
"${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_sm80.cu"
"${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_fp16_causal_sm80.cu"
"${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_fp16_sm80.cu"
"${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/flash_sparse_api.cpp"
)
set(INCLUDES
${repo-cutlass_SOURCE_DIR}/include
${repo-cutlass_SOURCE_DIR}/tools/util/include
${repo-flashinfer_SOURCE_DIR}/include
${repo-flashinfer_SOURCE_DIR}/csrc
${repo-cutlass_SOURCE_DIR}/examples/77_blackwell_fmha
${repo-cutlass_SOURCE_DIR}/examples/common
${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src
)
# =========================== Common SM90 Build ============================= #
# Build SM90 library with fast math optimization (same namespace, different directory)
Python_add_library(common_ops_sm90_build MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES})
target_compile_options(common_ops_sm90_build PRIVATE
$<$<COMPILE_LANGUAGE:CUDA>:${SGL_KERNEL_CUDA_FLAGS} -use_fast_math>
)
target_include_directories(common_ops_sm90_build PRIVATE ${INCLUDES})
# Set output name and separate build directory to avoid conflicts
set_target_properties(common_ops_sm90_build PROPERTIES
OUTPUT_NAME "common_ops"
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/sm90"
)
# =========================== Common SM100+ Build ============================= #
# Build SM100+ library with precise math (same namespace, different directory)
Python_add_library(common_ops_sm100_build MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES})
target_compile_options(common_ops_sm100_build PRIVATE
$<$<COMPILE_LANGUAGE:CUDA>:${SGL_KERNEL_CUDA_FLAGS}>
)
target_include_directories(common_ops_sm100_build PRIVATE ${INCLUDES})
# Set output name and separate build directory to avoid conflicts
set_target_properties(common_ops_sm100_build PROPERTIES
OUTPUT_NAME "common_ops"
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/sm100"
)
find_package(Python3 COMPONENTS Interpreter REQUIRED)
execute_process(
COMMAND ${Python3_EXECUTABLE} -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))"
OUTPUT_VARIABLE TORCH_CXX11_ABI
OUTPUT_STRIP_TRAILING_WHITESPACE
)
if(TORCH_CXX11_ABI STREQUAL "0")
message(STATUS "Using old C++ ABI (-D_GLIBCXX_USE_CXX11_ABI=0)")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0")
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0")
else()
message(STATUS "Using new C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI=1)")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=1")
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=1")
endif()
target_link_libraries(common_ops_sm90_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt)
target_link_libraries(common_ops_sm100_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt)
# sparse flash attention
target_compile_definitions(common_ops_sm90_build PRIVATE
FLASHATTENTION_DISABLE_BACKWARD
FLASHATTENTION_DISABLE_DROPOUT
FLASHATTENTION_DISABLE_UNEVEN_K
)
target_compile_definitions(common_ops_sm100_build PRIVATE
FLASHATTENTION_DISABLE_BACKWARD
FLASHATTENTION_DISABLE_DROPOUT
FLASHATTENTION_DISABLE_UNEVEN_K
)
# Install to different subdirectories
# CMake will find the built libraries in their respective LIBRARY_OUTPUT_DIRECTORY locations
# and install them to the specified destinations
install(TARGETS common_ops_sm90_build LIBRARY DESTINATION sgl_kernel/sm90)
install(TARGETS common_ops_sm100_build LIBRARY DESTINATION sgl_kernel/sm100)
# ============================ Optional Install: FA3 ============================= #
# set flash-attention sources file
# Now FA3 support sm80/sm86/sm90
if (SGL_KERNEL_ENABLE_FA3)
set(SGL_FLASH_KERNEL_CUDA_FLAGS
"-DNDEBUG"
"-DOPERATOR_NAMESPACE=sgl-kernel"
"-O3"
"-Xcompiler"
"-fPIC"
"-gencode=arch=compute_90a,code=sm_90a"
"-std=c++17"
"-DCUTE_USE_PACKED_TUPLE=1"
"-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1"
"-DCUTLASS_VERSIONS_GENERATED"
"-DCUTLASS_TEST_LEVEL=0"
"-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1"
"-DCUTLASS_DEBUG_TRACE_LEVEL=0"
"-DCUTLASS_ENABLE_GDC_FOR_SM90" # For PDL
"-DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED" # Necessary for the WGMMA shapes that we use
"--expt-relaxed-constexpr"
"--expt-extended-lambda"
"--use_fast_math"
"-Xcompiler=-Wconversion"
"-Xcompiler=-fno-strict-aliasing"
)
if (ENABLE_BELOW_SM90)
list(APPEND SGL_FLASH_KERNEL_CUDA_FLAGS
"-gencode=arch=compute_80,code=sm_80"
"-gencode=arch=compute_86,code=sm_86"
)
# SM8X Logic
file(GLOB FA3_SM8X_GEN_SRCS
"${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdim*_sm80.cu")
endif()
file(GLOB FA3_BF16_GEN_SRCS
"${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdim[0-9]*_bf16*_sm90.cu")
file(GLOB FA3_BF16_GEN_SRCS_
"${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdimdiff_bf16*_sm90.cu")
list(APPEND FA3_BF16_GEN_SRCS ${FA3_BF16_GEN_SRCS_})
# FP16 source files - use individual hdim files instead of hdimall to avoid ptxas crash
file(GLOB FA3_FP16_GEN_SRCS
"${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdim[0-9]*_fp16*_sm90.cu")
file(GLOB FA3_FP16_GEN_SRCS_
"${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdimdiff_fp16*_sm90.cu")
list(APPEND FA3_FP16_GEN_SRCS ${FA3_FP16_GEN_SRCS_})
# FP8 source files
file(GLOB FA3_FP8_GEN_SRCS
"${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdim[0-9]*_e4m3*_sm90.cu")
file(GLOB FA3_FP8_GEN_SRCS_
"${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdimdiff_e4m3*_sm90.cu")
list(APPEND FA3_FP8_GEN_SRCS ${FA3_FP8_GEN_SRCS_})
set(FA3_GEN_SRCS ${FA3_BF16_GEN_SRCS} ${FA3_FP16_GEN_SRCS} ${FA3_FP8_GEN_SRCS} ${FA3_SM8X_GEN_SRCS})
set(FLASH_SOURCES
"csrc/flash_extension.cc"
"${repo-flash-attention_SOURCE_DIR}/hopper/flash_prepare_scheduler.cu"
"${repo-flash-attention_SOURCE_DIR}/hopper/flash_api.cpp"
"${repo-flash-attention_SOURCE_DIR}/hopper/flash_fwd_combine.cu"
"${FA3_GEN_SRCS}"
)
Python_add_library(flash_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${FLASH_SOURCES})
target_compile_options(flash_ops PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:${SGL_FLASH_KERNEL_CUDA_FLAGS}>)
target_include_directories(flash_ops PRIVATE
${repo-cutlass_SOURCE_DIR}/include
${repo-cutlass_SOURCE_DIR}/tools/util/include
${repo-flash-attention_SOURCE_DIR}/hopper
)
target_link_libraries(flash_ops PRIVATE ${TORCH_LIBRARIES} c10 cuda)
install(TARGETS flash_ops LIBRARY DESTINATION "sgl_kernel")
set(FLASH_OPS_COMPILE_DEFS
FLASHATTENTION_DISABLE_BACKWARD
FLASHATTENTION_DISABLE_DROPOUT
FLASHATTENTION_DISABLE_UNEVEN_K
FLASHATTENTION_VARLEN_ONLY
)
if(NOT ENABLE_BELOW_SM90)
list(APPEND FLASH_OPS_COMPILE_DEFS FLASHATTENTION_DISABLE_SM8x)
endif()
if(NOT SGL_KERNEL_ENABLE_FA3_SPARSE_MASK)
list(APPEND FLASH_OPS_COMPILE_DEFS FLASHATTENTION_DISABLE_SPARSE_MASK)
endif()
target_compile_definitions(flash_ops PRIVATE ${FLASH_OPS_COMPILE_DEFS})
endif()
# ===================== InfLLM-V2 FlashAttention backend ===================== #
# Standalone pybind extension `infllm_ops`, vendored from
# 3rdparty/infllmv2_cuda_impl. Kept as its own module so its `flash::` symbols
# stay isolated from sgl-kernel's own flash attention. Mirrors the original
# setup.py: only hdim 64/128 bf16 forward instantiations are compiled (the
# vendored static_switch.h forces bf16 and dispatches headdim to {64, 128}
# only). Backward kernels are intentionally omitted because SGLang only uses
# these ops for inference.
set(INFLLM_FLASH_CUDA_FLAGS
"-DNDEBUG"
"-O3"
"-std=c++17"
"-Xcompiler"
"-fPIC"
"-U__CUDA_NO_HALF_OPERATORS__"
"-U__CUDA_NO_HALF_CONVERSIONS__"
"-U__CUDA_NO_HALF2_OPERATORS__"
"-U__CUDA_NO_BFLOAT16_CONVERSIONS__"
"--expt-relaxed-constexpr"
"--expt-extended-lambda"
"--use_fast_math"
"-DFLASHATTENTION_DISABLE_DROPOUT"
"-DFLASHATTENTION_DISABLE_ALIBI"
"-DFLASHATTENTION_DISABLE_SOFTCAP"
"-DFLASHATTENTION_DISABLE_UNEVEN_K"
"-DFLASHATTENTION_DISABLE_LOCAL"
"--threads=${SGL_KERNEL_COMPILE_THREADS}"
)
# Arch gencodes: match the original setup.py auto-detection
# (80 always; 90 for CUDA>=11.8; 120 for CUDA>=12.8).
if (ENABLE_BELOW_SM90)
list(APPEND INFLLM_FLASH_CUDA_FLAGS "-gencode=arch=compute_80,code=sm_80")
endif()
list(APPEND INFLLM_FLASH_CUDA_FLAGS "-gencode=arch=compute_90,code=sm_90")
if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A)
list(APPEND INFLLM_FLASH_CUDA_FLAGS "-gencode=arch=compute_120a,code=sm_120a")
endif()
set(INFLLM_FLASH_SOURCES
"csrc/infllm_v2/flash_extension.cc"
"csrc/infllm_v2/flash_attn/flash_api.cpp"
"csrc/infllm_v2/flash_attn/src/flash_fwd_split_hdim64_bf16_sm80.cu"
"csrc/infllm_v2/flash_attn/src/flash_fwd_split_hdim128_bf16_sm80.cu"
"csrc/infllm_v2/flash_attn/src/flash_fwd_split_hdim64_bf16_causal_sm80.cu"
"csrc/infllm_v2/flash_attn/src/flash_fwd_split_hdim128_bf16_causal_sm80.cu"
)
Python_add_library(infllm_ops MODULE WITH_SOABI ${INFLLM_FLASH_SOURCES})
target_compile_options(infllm_ops PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:${INFLLM_FLASH_CUDA_FLAGS}>)
target_include_directories(infllm_ops PRIVATE
${repo-cutlass_SOURCE_DIR}/include
${repo-cutlass_SOURCE_DIR}/tools/util/include
${CMAKE_CURRENT_LIST_DIR}/csrc/infllm_v2/flash_attn
${CMAKE_CURRENT_LIST_DIR}/csrc/infllm_v2/flash_attn/src
)
# The pybind module binds functions taking at::Generator, which pulls in
# THPGeneratorClass from libtorch_python (not part of TORCH_LIBRARIES).
find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib" REQUIRED)
target_link_libraries(infllm_ops PRIVATE ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY} c10 cuda)
install(TARGETS infllm_ops LIBRARY DESTINATION "sgl_kernel")
# Build spatial_ops as a separate, optional extension for green contexts
set(SPATIAL_SOURCES
"csrc/spatial/greenctx_stream.cu"
"csrc/spatial_extension.cc"
)
Python_add_library(spatial_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SPATIAL_SOURCES})
target_compile_options(spatial_ops PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:${SGL_KERNEL_CUDA_FLAGS}>)
target_link_libraries(spatial_ops PRIVATE ${TORCH_LIBRARIES} c10 cuda)
install(TARGETS spatial_ops LIBRARY DESTINATION sgl_kernel)
# ============================ Extra Install: FLashMLA ============================= #
include(${CMAKE_CURRENT_LIST_DIR}/cmake/flashmla.cmake)
# ============================ Extra Install: triton kernels ============================= #
install(DIRECTORY "${repo-triton_SOURCE_DIR}/python/triton_kernels/triton_kernels/"
DESTINATION "triton_kernels"
PATTERN ".git*" EXCLUDE
PATTERN "__pycache__" EXCLUDE)
python/sglang/kernels/aot/Dockerfile
+199
View File
@@ -0,0 +1,199 @@
ARG BASE_IMG=pytorch/manylinux2_28-builder
ARG CUDA_VERSION=13.0
# Dependency stage: install system deps, CMake, ccache, Python deps (including torch)
FROM ${BASE_IMG}:cuda${CUDA_VERSION} AS deps
# Overridable build arguments
ARG ARCH=x86_64
ARG CUDA_VERSION=13.0
ARG PYTHON_VERSION=3.10
# Manylinux python path tag, e.g. cp310-cp310 / cp312-cp312
ARG PYTHON_TAG=cp310-cp310
ARG CMAKE_VERSION_MAJOR=3.31
ARG CMAKE_VERSION_MINOR=1
# Install ccache 4.12.1 from source for CUDA support (yum provides old 3.7.7)
ARG USE_CCACHE=1
ARG CCACHE_VERSION=4.12.1
ARG GITHUB_ARTIFACTORY=github.com
ARG PYTORCH_INDEX_BASE=https://download.pytorch.org/whl
ARG PIP_DEFAULT_INDEX=https://pypi.python.org/simple
# Optional mirror for the manylinux base image's yum repos (AlmaLinux 8).
# Pass scheme + host (and optional path prefix), e.g.
# --build-arg YUM_MIRROR=https://mirrors.aliyun.com
# Empty (default) keeps upstream repo.almalinux.org.
ARG YUM_MIRROR=
ENV PYTHON_ROOT_PATH=/opt/python/${PYTHON_TAG}
ENV PATH=/opt/cmake/bin:${PATH}
ENV LD_LIBRARY_PATH=/lib64:${LD_LIBRARY_PATH}
ENV NINJA_STATUS="[%f/%t %es] "
ENV FLASHINFER_CUDA_ARCH_LIST="8.0 8.9 9.0a 10.0a 12.0a"
# CUDA headers path
ENV CPLUS_INCLUDE_PATH=/usr/local/cuda/include/cccl${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}
ENV C_INCLUDE_PATH=/usr/local/cuda/include/cccl${C_INCLUDE_PATH:+:${C_INCLUDE_PATH}}
RUN if [ -n "${YUM_MIRROR}" ]; then \
set -eux; \
sed -i \
-e 's|^mirrorlist=|#mirrorlist=|g' \
-e 's|^# *baseurl=https://repo.almalinux.org|baseurl='"${YUM_MIRROR}"'|g' \
/etc/yum.repos.d/almalinux*.repo; \
sed -i 's|^enabled=1|enabled=0|g' /etc/yum.repos.d/epel*.repo; \
fi
# Install build dependencies. libzstd-devel + xxhash-devel let ccache's
# FindZstd.cmake / FindXxhash.cmake skip their hardcoded github.com FetchContent
# fallbacks — critical when github.com is flaky/blocked. xxhash-devel is in
# PowerTools (disabled by default on AlmaLinux 8); libzstd-devel is in BaseOS.
RUN yum install gcc gcc-c++ make wget tar numactl-devel libibverbs libzstd-devel -y --nogpgcheck \
&& yum --enablerepo=powertools install xxhash-devel -y --nogpgcheck \
&& ln -sv /usr/lib64/libibverbs.so.1 /usr/lib64/libibverbs.so \
&& yum clean all && rm -rf /var/cache/yum
# Install CMake (cached download)
RUN --mount=type=cache,id=sgl-kernel-cmake,target=/cmake-downloads \
set -eux; \
CMAKE_TARBALL=cmake-${CMAKE_VERSION_MAJOR}.${CMAKE_VERSION_MINOR}-linux-${ARCH}.tar.gz; \
# Check if CMake is already cached
if [ -f /cmake-downloads/${CMAKE_TARBALL} ]; then \
echo "Using cached CMake from /cmake-downloads/${CMAKE_TARBALL}"; \
cp /cmake-downloads/${CMAKE_TARBALL} .; \
else \
CMAKE_TARBALL_URL=https://${GITHUB_ARTIFACTORY}/Kitware/CMake/releases/download/v${CMAKE_VERSION_MAJOR}.${CMAKE_VERSION_MINOR}/${CMAKE_TARBALL}; \
echo "Downloading CMake from: ${CMAKE_TARBALL_URL}"; \
wget --progress=dot ${CMAKE_TARBALL_URL}; \
# Cache the downloaded file
cp ${CMAKE_TARBALL} /cmake-downloads/; \
fi; \
tar -xzf ${CMAKE_TARBALL}; \
mv cmake-${CMAKE_VERSION_MAJOR}.${CMAKE_VERSION_MINOR}-linux-${ARCH} /opt/cmake; \
rm -f ${CMAKE_TARBALL}; \
cmake --version
# Install ccache
RUN if [ "${USE_CCACHE}" = "1" ]; then \
set -eux && \
cd /tmp && \
wget --progress=dot https://${GITHUB_ARTIFACTORY}/ccache/ccache/releases/download/v${CCACHE_VERSION}/ccache-${CCACHE_VERSION}.tar.xz && \
tar -xf ccache-${CCACHE_VERSION}.tar.xz && \
cd ccache-${CCACHE_VERSION} && \
mkdir build && cd build && \
cmake -D CMAKE_BUILD_TYPE=Release -D CMAKE_INSTALL_PREFIX=/usr -D ENABLE_TESTING=OFF -D REDIS_STORAGE_BACKEND=OFF -D HTTP_STORAGE_BACKEND=OFF -D ENABLE_DOCUMENTATION=OFF .. && \
make -j"$(nproc)" && \
make install && \
ccache --version && \
rm -rf /tmp/ccache-${CCACHE_VERSION}*; \
else \
echo "Skipping ccache build (USE_CCACHE=${USE_CCACHE})"; \
fi
RUN set -eux; \
if [ "${ARCH}" = "aarch64" ]; then _LIB=sbsa; else _LIB="${ARCH}"; fi; \
mkdir -p /usr/lib/${ARCH}-linux-gnu/; \
ln -sf /usr/local/cuda-${CUDA_VERSION}/targets/${_LIB}-linux/lib/stubs/libcuda.so /usr/lib/${ARCH}-linux-gnu/libcuda.so
# Install Python dependencies (torch + build tools)
RUN --mount=type=cache,id=sgl-kernel-pip,target=/root/.cache/pip \
set -eux; \
case "${CUDA_VERSION}" in \
13.0) TORCH_VER=2.11.0; CU_TAG=cu130 ;; \
12.9) TORCH_VER=2.11.0; CU_TAG=cu129 ;; \
12.8) TORCH_VER=2.11.0; CU_TAG=cu128 ;; \
*) TORCH_VER=2.11.0; CU_TAG=cu126 ;; \
esac; \
${PYTHON_ROOT_PATH}/bin/pip install torch==${TORCH_VER} --index-url ${PYTORCH_INDEX_BASE}/${CU_TAG}; \
${PYTHON_ROOT_PATH}/bin/pip install ninja setuptools==75.0.0 wheel==0.41.0 numpy uv scikit-build-core --index-url ${PIP_DEFAULT_INDEX}
# Build stage: copy source and build wheel
FROM deps AS build
WORKDIR /sgl-kernel
# Only copy sgl-kernel source so code changes only affect later layers
COPY . /sgl-kernel/
# Optional: enable CMake/Ninja profiling (pass non-empty via --build-arg ENABLE_*)
ARG ENABLE_CMAKE_PROFILE
ARG ENABLE_BUILD_PROFILE
ARG ARCH=x86_64
ARG USE_CCACHE=1
# Parallelism knobs (override via --build-arg)
# BUILD_JOBS: number of parallel compilation units (ninja -j)
# NVCC_THREADS: per-compilation-unit NVCC --threads (multi-arch PTXAS)
ARG BUILD_JOBS=0
ARG NVCC_THREADS=32
# Redeclare so CMake third-party FetchContent uses the same mirror as deps stage
ARG GITHUB_ARTIFACTORY=github.com
RUN --mount=type=cache,id=sgl-kernel-ccache,target=/ccache \
--mount=type=cache,id=sgl-kernel-pip,target=/root/.cache/pip \
set -eux; \
if [ "${USE_CCACHE}" = "1" ]; then \
export CCACHE_DIR=/ccache; \
export CCACHE_BASEDIR=/sgl-kernel; \
export CCACHE_MAXSIZE=10G; \
export CCACHE_COMPILERCHECK=content; \
export CCACHE_COMPRESS=true; \
export CCACHE_SLOPPINESS=file_macro,time_macros,include_file_mtime,include_file_ctime; \
export CMAKE_C_COMPILER_LAUNCHER=ccache; \
export CMAKE_CXX_COMPILER_LAUNCHER=ccache; \
export CMAKE_CUDA_COMPILER_LAUNCHER=ccache; \
ccache -sV; \
fi; \
# Setting these flags to reduce OOM chance only on ARM
if [ "${ARCH}" = "aarch64" ]; then \
export CUDA_NVCC_FLAGS="-Xcudafe --threads=2"; \
export MAKEFLAGS="-j2"; \
export CMAKE_BUILD_PARALLEL_LEVEL=2; \
export NINJAFLAGS="-j2"; \
echo "ARM detected: Using extra conservative settings (2 parallel jobs)"; \
elif [ "${BUILD_JOBS}" -gt 0 ] 2>/dev/null; then \
export CMAKE_BUILD_PARALLEL_LEVEL=${BUILD_JOBS}; \
else \
export CMAKE_BUILD_PARALLEL_LEVEL=$(echo "$(( $(nproc) * 2 / 3 )) 64" | awk '{print ($1 < $2) ? $1 : $2}'); \
fi; \
export CMAKE_ARGS="${CMAKE_ARGS:-} -DSGL_KERNEL_COMPILE_THREADS=${NVCC_THREADS} -DGITHUB_ARTIFACTORY=${GITHUB_ARTIFACTORY}"; \
if [ -n "${ENABLE_CMAKE_PROFILE:-}" ]; then \
echo "CMake profiling enabled - will save to /sgl-kernel/cmake-profile.json"; \
export CMAKE_ARGS="${CMAKE_ARGS} --profiling-output=/sgl-kernel/cmake-profile.json --profiling-format=google-trace"; \
fi; \
echo "Build parallelism: CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL}, NVCC_THREADS=${NVCC_THREADS}"; \
echo "CMAKE_ARGS=${CMAKE_ARGS}"; \
${PYTHON_ROOT_PATH}/bin/python -m uv build --wheel -Cbuild-dir=build . --color=always --no-build-isolation; \
./rename_wheels.sh; \
if [ -n "${ENABLE_BUILD_PROFILE:-}" ] && [ -f /sgl-kernel/build/.ninja_log ]; then \
echo "Ninja build profiling enabled - will save to /sgl-kernel/build-trace.json"; \
wget --progress=dot https://raw.githubusercontent.com/cradleapps/ninjatracing/084212eaf68f25c70579958a2ed67fb4ec2a9ca4/ninjatracing -O /tmp/ninjatracing; \
if [ -f /tmp/ninjatracing ]; then \
${PYTHON_ROOT_PATH}/bin/python /tmp/ninjatracing /sgl-kernel/build/.ninja_log > /sgl-kernel/build-trace.json; \
fi; \
if [ -f /sgl-kernel/build-trace.json ]; then \
gzip -9 -k /sgl-kernel/build-trace.json 2>/dev/null || true; \
echo "Build trace saved to: build-trace.json"; \
if [ -f /sgl-kernel/build-trace.json.gz ]; then \
ORIGINAL_SIZE=$(stat -f%z /sgl-kernel/build-trace.json 2>/dev/null || stat -c%s /sgl-kernel/build-trace.json); \
COMPRESSED_SIZE=$(stat -f%z /sgl-kernel/build-trace.json.gz 2>/dev/null || stat -c%s /sgl-kernel/build-trace.json.gz); \
RATIO=$(( (ORIGINAL_SIZE - COMPRESSED_SIZE) * 100 / ORIGINAL_SIZE )); \
echo "Compressed to: build-trace.json.gz (${RATIO}% smaller)"; \
fi; \
echo ""; \
echo "View in browser:"; \
echo " - chrome://tracing (load JSON file)"; \
echo " - ui.perfetto.dev (recommended, supports .gz files)"; \
echo ""; \
echo "Shows:"; \
echo " - Compilation time per file"; \
echo " - Parallelism utilization"; \
echo " - Critical path (longest dependency chain)"; \
echo " - Where the 2-hour build time went"; \
fi; \
fi; \
if [ "${USE_CCACHE}" = "1" ]; then \
echo "ccache Statistics"; \
ccache -s; \
else \
echo "ccache disabled (USE_CCACHE=${USE_CCACHE})"; \
fi
# Artifact stage (for --output to export wheel)
FROM scratch AS artifact
COPY --from=build /sgl-kernel/dist/*.whl /
python/sglang/kernels/aot/LICENSE
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright 2023-2024 SGLang Team
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
python/sglang/kernels/aot/Makefile
+93
View File
@@ -0,0 +1,93 @@
.PHONY: help check-deps install-deps tree ln submodule install build clean rebuild test format update
# ---------------------------
# Build resource controls
# ---------------------------
# By default, build uses all available CPU cores, but users can override:
# make build MAX_JOBS=2 CMAKE_BUILD_PARALLEL_LEVEL=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"
NPROC ?= $(shell nproc 2>/dev/null || echo 1)
MAX_JOBS ?= $(NPROC)
CMAKE_BUILD_PARALLEL_LEVEL ?= $(MAX_JOBS)
UV_BUILD_DIR ?= build
CMAKE_POLICY_VERSION_MINIMUM ?= 3.5
# Optional GitHub mirror for FetchContent (e.g. internal artifactory).
# Empty -> CMakeLists.txt default (github.com).
GITHUB_ARTIFACTORY ?=
ifneq ($(strip $(GITHUB_ARTIFACTORY)),)
CMAKE_ARGS += -DGITHUB_ARTIFACTORY=$(GITHUB_ARTIFACTORY)
endif
# Show help for each target
help: ## Show this help message
@echo "Available targets:"
@grep -E '^[a-zA-Z0-9_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}'
check-deps: ## Check and install required Python formatting dependencies
@command -v isort >/dev/null 2>&1 || (echo "Installing isort..." && pip install isort)
@command -v black >/dev/null 2>&1 || (echo "Installing black..." && pip install black)
install-deps: ## Install Python formatting tools (isort and black)
pip install scikit-build-core isort black
tree: ## Show project directory structure
@tree --prune -I "__pycache__|*.egg-info|*.so|build|3rdparty|dist"
submodule: ## Initialize and update git submodules
@git submodule update --init --recursive
ln: submodule ## Create compilation database
@rm -rf build && mkdir build && cd build && cmake .. -DCMAKE_EXPORT_COMPILE_COMMANDS=YES -DCMAKE_POLICY_VERSION_MINIMUM=3.5
install: submodule ## Install package in development mode
@pip install -e . --no-build-isolation
build: install-deps submodule ## Build and install wheel package
@rm -rf dist/* || true && \
CMAKE_POLICY_VERSION_MINIMUM=$(CMAKE_POLICY_VERSION_MINIMUM) \
MAX_JOBS=$(MAX_JOBS) \
CMAKE_BUILD_PARALLEL_LEVEL=$(CMAKE_BUILD_PARALLEL_LEVEL) \
CMAKE_ARGS="$(CMAKE_ARGS)" \
uv build --wheel -Cbuild-dir=$(UV_BUILD_DIR) . --verbose --color=always --no-build-isolation && \
pip3 install dist/*whl --force-reinstall --no-deps
clean: ## Remove build artifacts
@rm -rf build dist *.egg-info
rebuild: clean submodule build ## Clean and rebuild the project
@echo "Succeed to rebuild"
test: ## Run all tests
@find tests -name "test_*.py" | xargs -n 1 python3
format: check-deps ## Format all source files
@echo "Formatting source files..."
@find csrc tests -name '*.cc' -o -name '*.cu' -o -name '*.cuh' -o -name '*.h' -o -name '*.hpp' | xargs clang-format -i
@find python tests -name '*.py' | xargs isort
@find python tests -name '*.py' | xargs black
@pre-commit run --all-files
FILES_TO_UPDATE = python/sgl_kernel/version.py \
pyproject.toml \
pyproject_rocm.toml \
pyproject_cpu.toml
update: ## Update version numbers across project files. Usage: make update <new_version>
@if [ -z "$(filter-out $@,$(MAKECMDGOALS))" ]; then \
echo "Version required. Usage: make update <new_version>"; \
exit 1; \
fi
@OLD_VERSION=$$(grep "version" python/sgl_kernel/version.py | cut -d '"' -f2); \
NEW_VERSION=$(filter-out $@,$(MAKECMDGOALS)); \
echo "Updating version from $$OLD_VERSION to $$NEW_VERSION"; \
for file in $(FILES_TO_UPDATE); do \
if [ "$(shell uname)" = "Darwin" ]; then \
sed -i '' -e "s/$$OLD_VERSION/$$NEW_VERSION/g" $$file; \
else \
sed -i -e "s/$$OLD_VERSION/$$NEW_VERSION/g" $$file; \
fi \
done; \
echo "Version update complete"
%:
@:
python/sglang/kernels/aot/README.md
+143
View File
@@ -0,0 +1,143 @@
# sglang-kernel (prior sgl-kernel)
[Kernel Library](https://github.com/sgl-project/sglang/tree/main/python/sglang/kernels/aot) for LLM inference engines
<div align="center">
[![License: Apache-2.0](https://img.shields.io/badge/License-Apache--2.0-blue.svg)](https://github.com/sgl-project/sglang/blob/main/LICENSE)
[![PyPI](https://img.shields.io/pypi/v/sglang-kernel)](https://pypi.org/project/sglang-kernel)
</div>
`sglang-kernel` provides optimized compute primitives for LLM inference engines, enabling efficient inference for large language models and vision-language models through custom kernel operations. The source tree lives under the `python/sglang/kernels/aot/` directory and the Python import path remains `sgl_kernel`.
## Installation
Requires torch == 2.11.0
```bash
# Latest version
pip3 install sglang-kernel --upgrade
```
## Building from Source
Requires
- CMake ≥3.31,
- Python ≥3.10
- scikit-build-core
- ninja(optional)
### Use Makefile to build from the sgl-kernel source tree
```bash
make build
```
### Limit build resource usage (CPU / parallelism)
By default, `make build` uses all available CPU cores. You can override build parallelism and NVCC compile threads:
```bash
# Limit parallel jobs (controls both make and cmake parallelism)
make build MAX_JOBS=2
# Additionally limit NVCC internal threads (reduces CPU and peak memory)
make build MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"
```
## Contribution
### Steps to add a new kernel:
1. Implement the kernel in [csrc](https://github.com/sgl-project/sglang/tree/main/python/sglang/kernels/aot/csrc)
2. Expose the interface in [include/sgl_kernel_ops.h](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/include/sgl_kernel_ops.h)
3. Create torch extension in [csrc/common_extension.cc](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/common_extension.cc)
4. Update [CMakeLists.txt](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/CMakeLists.txt) to include new CUDA source
5. Expose Python interface in [python](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/python/sgl_kernel)
6. Add test and benchmark
### Development Tips
1. When creating torch extensions, add the function definition with `m.def`, and device binding with `m.impl`:
- How to write schema: [Schema reference](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#func)
```cpp
// We need def with schema here for torch.compile
m.def(
"bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, "
"int cublas_handle) -> ()");
m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8);
```
### Adapting C++ Native Types for Torch Compatibility
Third-party C++ libraries often use int and float, but PyTorch bindings require int64_t and double due to Python's type mapping.
Use make_pytorch_shim from sgl_kernel_torch_shim.h to handle conversions automatically:
```cpp
// Add type conversion for int -> int64_t
template <>
struct pytorch_library_compatible_type<int> {
using type = int64_t;
static int convert_from_type(int64_t arg) {
TORCH_CHECK(arg <= std::numeric_limits<int>::max(), "value too large");
TORCH_CHECK(arg >= std::numeric_limits<int>::min(), "value too small");
return arg;
}
};
```
```cpp
// Wrap your function
m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
```
### Testing & Benchmarking
1. Add pytest tests in [tests/](https://github.com/sgl-project/sglang/tree/main/python/sglang/kernels/aot/tests), if you need to skip some test, please use `@pytest.mark.skipif`
```python
@pytest.mark.skipif(
skip_condition, reason="Nvfp4 Requires compute capability of 10 or above."
)
```
2. Add benchmarks using [triton benchmark](https://triton-lang.org/main/python-api/generated/triton.testing.Benchmark.html) in [benchmark/](https://github.com/sgl-project/sglang/tree/main/python/sglang/kernels/aot/benchmark)
**We recommend using `triton.testing.do_bench_cudagraph` for kernel benchmarking**:
Compared to `triton.testing.do_bench`, `do_bench_cudagraph` provides:
- Reduced CPU overhead impact for more accurate kernel performance measurements
- Incorporation of PDL (Programmatic Dependent Launch) effects into individual kernel results
- More realistic performance data on PDL-supported architectures (SM >= 90)
3. Run test suite
## Kernel Size Analysis
Analyze CUDA kernel sizes in compiled wheel files to identify oversized kernels and template-instantiation bloat:
This tool requires `cubloaty` (install with `pip install cubloaty`) to work.
```bash
# Install cubloaty
pip install cubloaty
# Analyze a wheel file
python analyze_whl_kernel_sizes.py path/to/sglang_kernel-*.whl
# Custom output file
python analyze_whl_kernel_sizes.py path/to/sglang_kernel-*.whl --output my_analysis.txt
```
The tool generates:
- A text report with:
- Kernel groups (by name prefix)
- Individual kernel sizes (sorted by size)
Use this to identify large kernels and potential template instantiation bloat.
## FAQ
- Q: Segmentation fault with CUDA 12.6
- A: Update ptxas to 12.8, reference: [segment fault error](https://github.com/Dao-AILab/flash-attention/issues/1453)
python/sglang/kernels/aot/THIRDPARTYNOTICES.txt
+488
View File
@@ -0,0 +1,488 @@
Notice for flashinfer-ai/flashinfer
-------------------------------
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-------------------------------------------------------------------------------------------------
Some of the code in this project are adapted from other open-source projects with different
licenses. This product also bundles some third-party components under other open source licenses.
This section summarizes those components and their licenses.
See licenses/ for text of these licenses.
BSD 3-Clause License
--------------------
include/flashinfer/attention/hopper/epilogue.cuh
include/flashinfer/attention/hopper/mainloop.cuh
include/flashinfer/attention/hopper/kernel_traits.cuh
include/flashinfer/attention/hopper/named_barrier.cuh
include/flashinfer/attention/hopper/tile_scheduler.cuh
include/flashinfer/attention/hopper/utils.cuh
BSD 3-Clause "New" License
--------------------------
3rdparty/cutlass
include/flashinfer/attention/hopper/block_sparse_gather.cuh
Notice for NVIDIA/TensorRT-LLM
-------------------------------
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
Notice for deepseek-ai/DeepGEMM
-------------------------------
MIT License
Copyright (c) 2025 DeepSeek
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
Notice for Dao-AILab/flash-attention
-------------------------------
BSD 3-Clause License
Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
* Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
python/sglang/kernels/aot/analyze_whl_kernel_sizes.py
+221
View File
@@ -0,0 +1,221 @@
import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
import zipfile
from pathlib import Path
def extract_whl(whl_file, extract_dir):
with zipfile.ZipFile(whl_file, "r") as zip_ref:
zip_ref.extractall(extract_dir)
def find_binary_files(extract_dir):
binary_files = []
extract_path = Path(extract_dir)
for so_file in extract_path.rglob("*.so"):
binary_files.append(str(so_file))
for cubin_file in extract_path.rglob("*.cubin"):
binary_files.append(str(cubin_file))
return sorted(binary_files)
def run_cubloaty(binary_file):
result = subprocess.run(
["cubloaty", binary_file, "--format", "json"],
capture_output=True,
text=True,
timeout=60,
)
if result.returncode != 0:
if (
"No CUDA binary sections found" in result.stderr
or "does not contain device code" in result.stderr
):
return {}
raise subprocess.CalledProcessError(
result.returncode, result.args, result.stdout, result.stderr
)
return json.loads(result.stdout)
def analyze_whl(whl_file):
temp_dir = tempfile.mkdtemp(prefix="sgl_kernel_analysis_")
try:
extract_whl(whl_file, temp_dir)
binary_files = find_binary_files(temp_dir)
if not binary_files:
print(f"No .so or .cubin files found in {whl_file}")
return []
all_kernels = []
for binary_file in binary_files:
file_name = os.path.basename(binary_file)
data = run_cubloaty(binary_file)
if not data or "kernels" not in data:
continue
for kernel in data["kernels"]:
all_kernels.append(
{
"file": file_name,
"name": kernel.get("name", "unknown"),
"size": kernel.get("size", 0),
"size_kb": kernel.get("size", 0) / 1024,
"size_mb": kernel.get("size", 0) / 1024 / 1024,
}
)
return all_kernels
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
def extract_kernel_prefix(kernel_name):
if "<" in kernel_name:
return kernel_name.split("<")[0]
return kernel_name
def generate_report(all_kernels, output_file):
if not all_kernels:
print("No kernels found")
return
sorted_kernels = sorted(all_kernels, key=lambda x: x["size"], reverse=True)
total_size = sum(k["size"] for k in all_kernels)
total_size_mb = total_size / 1024 / 1024
from collections import defaultdict
kernel_groups = defaultdict(lambda: {"size": 0, "count": 0})
for kernel in all_kernels:
prefix = extract_kernel_prefix(kernel["name"])
kernel_groups[prefix]["size"] += kernel["size"]
kernel_groups[prefix]["count"] += 1
sorted_groups = sorted(
kernel_groups.items(), key=lambda x: x[1]["size"], reverse=True
)
lines = []
lines.append("=" * 140)
lines.append("CUDA Kernel Size Analysis")
lines.append("=" * 140)
lines.append("")
lines.append(f"Total kernels: {len(all_kernels)}")
lines.append(f"Total size: {total_size_mb:.2f} MB ({total_size:,} bytes)")
lines.append(f"Average kernel size: {total_size / len(all_kernels) / 1024:.2f} KB")
lines.append("")
lines.append("=" * 140)
lines.append("Kernel Groups (by name prefix) - Top 20")
lines.append("=" * 140)
lines.append(
f"{'Rank':<6} {'Kernel Prefix':<80} {'Count':<8} {'Total (MB)':<12} {'%':<8}"
)
lines.append("-" * 140)
TOP_N = 20
for i, (prefix, stats) in enumerate(sorted_groups[:TOP_N], 1):
percentage = (stats["size"] / total_size * 100) if total_size > 0 else 0
size_mb = stats["size"] / 1024 / 1024
display_prefix = prefix
if len(display_prefix) > 77:
display_prefix = display_prefix[:74] + "..."
lines.append(
f"{i:<6} {display_prefix:<80} {stats['count']:<8} {size_mb:<12.2f} {percentage:<8.2f}"
)
if len(sorted_groups) > TOP_N:
other_size = sum(stats["size"] for _, stats in sorted_groups[TOP_N:])
other_count = sum(stats["count"] for _, stats in sorted_groups[TOP_N:])
other_percentage = (other_size / total_size * 100) if total_size > 0 else 0
other_size_mb = other_size / 1024 / 1024
lines.append(
f"{'Other':<6} {'(remaining ' + str(len(sorted_groups) - TOP_N) + ' kernel groups)':<80} "
f"{other_count:<8} {other_size_mb:<12.2f} {other_percentage:<8.2f}"
)
lines.append("")
lines.append("=" * 140)
lines.append("Individual Kernels (sorted by size) - Top 20")
lines.append("=" * 140)
lines.append(
f"{'Rank':<6} {'File':<40} {'Kernel Name':<70} {'Size (KB)':<12} {'Size (MB)':<12} {'%':<8}"
)
lines.append("-" * 140)
for i, kernel in enumerate(sorted_kernels[:TOP_N], 1):
percentage = (kernel["size"] / total_size * 100) if total_size > 0 else 0
kernel_name = kernel["name"]
if len(kernel_name) > 67:
kernel_name = kernel_name[:64] + "..."
file_name = kernel["file"]
if len(file_name) > 37:
file_name = file_name[:34] + "..."
lines.append(
f"{i:<6} {file_name:<40} {kernel_name:<70} "
f"{kernel['size_kb']:<12.2f} {kernel['size_mb']:<12.4f} {percentage:<8.2f}"
)
if len(sorted_kernels) > TOP_N:
other_size = sum(k["size"] for k in sorted_kernels[TOP_N:])
other_count = len(sorted_kernels) - TOP_N
other_percentage = (other_size / total_size * 100) if total_size > 0 else 0
other_size_kb = other_size / 1024
other_size_mb = other_size / 1024 / 1024
lines.append(
f"{'Other':<6} {'(remaining ' + str(other_count) + ' kernels)':<40} "
f"{'':<70} {other_size_kb:<12.2f} {other_size_mb:<12.4f} {other_percentage:<8.2f}"
)
report_text = "\n".join(lines)
with open(output_file, "w") as f:
f.write(report_text)
print(f"Report saved to: {output_file}")
def main():
parser = argparse.ArgumentParser(
description="Analyze CUDA kernel sizes in sglang-kernel wheel files"
)
parser.add_argument("whl", type=str, help="Path to whl file")
parser.add_argument(
"--output", type=str, default="kernel_analysis.txt", help="Output report file"
)
args = parser.parse_args()
if not os.path.exists(args.whl):
print(f"Error: {args.whl} not found")
sys.exit(1)
all_kernels = analyze_whl(args.whl)
if all_kernels:
generate_report(all_kernels, args.output)
else:
print("No kernel information extracted")
if __name__ == "__main__":
main()
python/sglang/kernels/aot/benchmark/bench_activation.py
+209
View File
@@ -0,0 +1,209 @@
# Benchmarks SGLang kernels versus vLLM across
# (kernel, dtype, batch_size, seq_len, dim) and prints speed-up.
import argparse
import itertools
import os
import re
from typing import List, Tuple
import sgl_kernel
import torch
import torch.nn.functional as F
import triton
import triton.testing
from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm import _custom_ops as vllm_ops
VLLM_AVAILABLE = True
except ImportError:
vllm_ops = None
VLLM_AVAILABLE = False
IS_CI = is_in_ci()
# gelu_quick is only available on HIP/ROCm platforms
try:
from sgl_kernel import gelu_quick
GELU_QUICK_AVAILABLE = True
except ImportError:
GELU_QUICK_AVAILABLE = False
gelu_quick = None
if VLLM_AVAILABLE and not hasattr(vllm_ops, "silu_and_mul"):
vllm_ops = torch.ops._C
def str2int_list(arg: str) -> List[int]:
if arg in ("", None):
return []
if re.fullmatch(r"\d+(,\d+)*", arg.strip()) is None:
raise argparse.ArgumentTypeError(f"Bad int list: {arg}")
return [int(x) for x in arg.split(",")]
def calculate_diff(
kernel: str, dtype: torch.dtype, batch_size: int, seq_len: int, dim: int
) -> bool:
"""Compare vLLM with SGLang for one shape."""
device = torch.device("cuda")
if not VLLM_AVAILABLE:
print(
f"[{kernel:14s} | {str(dtype):9s} | B={batch_size:3d} | "
f"L={seq_len:3d} | D={dim:5d}] ⚠️ vLLM not available, skipping comparison"
)
return True
# activation-only quick GELU
if kernel == "gelu_quick":
if not GELU_QUICK_AVAILABLE:
print(
f"[{kernel:14s} | {str(dtype):9s} | B={batch_size:3d} | "
f"L={seq_len:3d} | D={dim:5d}] ⚠️ not available on this platform"
)
return True
x = torch.randn(batch_size, seq_len, dim, dtype=dtype, device=device)
ref_out = torch.zeros_like(x)
getattr(vllm_ops, kernel)(ref_out, x)
test_out = getattr(sgl_kernel, kernel)(x)
# fused activation x mul kernels
else:
x = torch.randn(batch_size, seq_len, 2 * dim, dtype=dtype, device=device)
ref_out = torch.zeros(batch_size, seq_len, dim, dtype=dtype, device=device)
getattr(vllm_ops, kernel)(ref_out, x)
test_out = getattr(sgl_kernel, kernel)(x)
ok = torch.allclose(ref_out, test_out, rtol=1e-3, atol=1e-5)
tag = "✅ match" if ok else "❌ mismatch"
print(
f"[{kernel:14s} | {str(dtype):9s} | B={batch_size:3d} | "
f"L={seq_len:3d} | D={dim:5d}] {tag}"
)
return ok
# CI environment uses simplified parameters for kernels and dtypes too
if IS_CI:
kernels = ["silu_and_mul"] # Only test one kernel in CI
dtypes = [torch.float16] # Only test one dtype in CI
else:
kernels = ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul"]
if GELU_QUICK_AVAILABLE:
kernels.append("gelu_quick")
dtypes = [torch.float16, torch.bfloat16]
def make_configs(bsizes: List[int], slens: List[int], dims_: List[int]) -> List[Tuple]:
return list(itertools.product(kernels, dtypes, bsizes, slens, dims_))
# CI environment uses simplified parameters
if IS_CI:
default_batch_sizes = [1] # Single batch size for CI
default_seq_lens = [1] # Single sequence length for CI
default_dims = [1024] # Single dimension for CI
else:
default_batch_sizes = [2**i for i in range(0, 5, 2)] # 1,4,16
default_seq_lens = [2**i for i in range(0, 8, 2)] # 1,4,16,64
default_dims = [2**i for i in range(10, 15)] # 1024...16384
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["kernel", "dtype", "batch_size", "seq_len", "dim"],
x_vals=[],
line_arg="provider",
line_vals=["vllm", "sglang", "speedup"],
line_names=["vLLM", "SGL Kernel", "Speed-up (x)"],
styles=[("blue", "-"), ("green", "-"), ("red", "--")],
ylabel="µs (median) or × (speed-up)",
plot_name="activation-performance",
args={},
)
)
def benchmark(kernel, dtype, batch_size, seq_len, dim, provider):
device = torch.device("cuda")
in_mult = 1 if kernel == "gelu_quick" else 2
x = torch.randn(batch_size, seq_len, in_mult * dim, dtype=dtype, device=device)
y0 = torch.zeros(batch_size, seq_len, dim, dtype=dtype, device=device)
if not VLLM_AVAILABLE and provider in ["vllm", "speedup"]:
# Skip vLLM-related benchmarks if vLLM is not available
return (0, 0, 0)
if VLLM_AVAILABLE:
vllm_kernel = getattr(vllm_ops, kernel)
if kernel == "gelu_quick" and not GELU_QUICK_AVAILABLE:
# Skip benchmark for gelu_quick if not available
return (0, 0, 0)
sglang_kernel = getattr(sgl_kernel, kernel)
def baseline():
if VLLM_AVAILABLE:
tmp = y0.clone()
vllm_kernel(tmp, x)
return tmp
else:
return torch.zeros_like(y0)
def sglang():
return sglang_kernel(x)
# timing helper
def timed(fn):
for _ in range(5):
fn()
torch.cuda.synchronize()
ms, qmin, qmax = triton.testing.do_bench_cudagraph(
fn, quantiles=[0.5, 0.2, 0.8]
)
return 1000 * ms, 1000 * qmax, 1000 * qmin
if provider == "vllm":
return timed(baseline)
if provider == "sglang":
return timed(sglang)
# provider == "speedup"
t_ref, _, _ = timed(baseline)
t_sgl, _, _ = timed(sglang)
spd = t_ref / t_sgl if t_ref > 0 else 1.0
return (spd, spd, spd)
if __name__ == "__main__":
p = argparse.ArgumentParser("Activation kernel benchmark")
p.add_argument("--batch_sizes", type=str2int_list, default=default_batch_sizes)
p.add_argument("--seq_lens", type=str2int_list, default=default_seq_lens)
p.add_argument("--dims", type=str2int_list, default=default_dims)
p.add_argument("--verify_only", action="store_true")
args = p.parse_args()
# coerce lists
if isinstance(args.batch_sizes, str):
args.batch_sizes = str2int_list(args.batch_sizes)
if isinstance(args.seq_lens, str):
args.seq_lens = str2int_list(args.seq_lens)
if isinstance(args.dims, str):
args.dims = str2int_list(args.dims)
# patch perf_report grid
benchmark_grid = make_configs(args.batch_sizes, args.seq_lens, args.dims)
if hasattr(benchmark, "benchmarks"):
benchmark.benchmarks.x_vals = benchmark_grid
else:
benchmark.benchmark.x_vals = benchmark_grid
if args.verify_only:
# Test with the first available kernel
test_kernel = kernels[0]
ok = calculate_diff(test_kernel, torch.float16, 1, 1, args.dims[0])
print("✅ sanity pass" if ok else "❌ mismatch")
else:
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_amd_deterministic_allreduce.py
+689
View File
@@ -0,0 +1,689 @@
"""
Benchmark latency comparison between different all-reduce implementations.
Compares:
- NCCL all-reduce (may be non-deterministic)
- Reduce-scatter + all-gather (RS+AG, deterministic but slower)
- Deterministic 1-stage kernel (forces fixed accumulation order, deterministic)
Note: The "deterministic kernel" is NOT RS+AG. It uses the 1-stage kernel where
each GPU reads all data from all GPUs and reduces locally in a fixed order.
Usage:
python bench_amd_deterministic_allreduce.py
"""
import multiprocessing as mp
import os
import socket
import statistics
import sys
import time
import torch
import torch.distributed as dist
# Add python directory to path to import sglang modules
script_dir = os.path.dirname(os.path.abspath(__file__))
python_dir = os.path.join(script_dir, "python")
sys.path.insert(0, python_dir)
# Try to import custom all-reduce if available
from sglang.srt.environ import envs
try:
import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as custom_ar_ops
from sglang.srt.distributed.device_communicators.custom_all_reduce import (
CustomAllreduce,
)
CUSTOM_AR_AVAILABLE = custom_ar_ops.IS_CUSTOM_AR_AVAILABLE
except (ImportError, AttributeError):
CUSTOM_AR_AVAILABLE = False
CustomAllreduce = None
# Note: sglang's optimized all-reduce requires full runtime initialization
# and won't work in standalone benchmarks, so we skip it
SGLANG_AVAILABLE = False
def get_open_port():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def init_custom_ar_if_available(rank, world_size, device):
"""Check if custom all-reduce is available and applicable."""
if not CUSTOM_AR_AVAILABLE or CustomAllreduce is None:
return False
# Custom AR works best for single-node, even number of GPUs, world_size <= 8
if world_size <= 8 and world_size % 2 == 0:
return True
return False
def reduce_scatter_then_all_gather(tensor, rank, world_size, custom_ar=None):
"""
Deterministic all-reduce using reduce-scatter + all-gather.
This is deterministic because it uses fixed ordering (no atomics).
"""
total_size = tensor.numel()
if total_size % world_size != 0:
# Fallback to all-gather + local reduce if not divisible
gather_list = [torch.empty_like(tensor) for _ in range(world_size)]
dist.all_gather(gather_list, tensor)
stacked = torch.stack(gather_list, dim=0)
tensor.copy_(stacked.sum(dim=0))
return
chunk_size = total_size // world_size
# Flatten to 1D
tensor_flat = tensor.view(-1)
# Reduce-scatter: each rank gets its chunk of the reduced result
output_chunk = torch.empty(chunk_size, dtype=tensor.dtype, device=tensor.device)
# Split input into chunks for reduce-scatter
input_chunks = [
tensor_flat[i * chunk_size : (i + 1) * chunk_size].clone()
for i in range(world_size)
]
dist.reduce_scatter(output_chunk, input_chunks)
# All-gather: broadcast each rank's chunk to all ranks
output_chunks = [
torch.empty(chunk_size, dtype=tensor.dtype, device=tensor.device)
for _ in range(world_size)
]
dist.all_gather(output_chunks, output_chunk)
# Concatenate results back
result_flat = torch.cat(output_chunks, dim=0)
tensor.copy_(result_flat.view(tensor.shape))
def worker(world_size, rank, port, results_queue):
envs.SGLANG_USE_1STAGE_ALLREDUCE.set("1")
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
dist.init_process_group(
backend="nccl",
init_method=f"tcp://localhost:{port}",
rank=rank,
world_size=world_size,
)
# Try to initialize custom all-reduce if available
custom_ar = None
use_custom_ar = init_custom_ar_if_available(rank, world_size, device)
if use_custom_ar and CUSTOM_AR_AVAILABLE:
try:
# Create a gloo group for custom AR (it requires non-NCCL backend)
# All ranks must call new_group with the same parameters
from torch.distributed import new_group
dist.barrier() # Ensure all ranks are ready
ar_group = new_group(backend="gloo")
dist.barrier() # Ensure group creation is complete
custom_ar = CustomAllreduce(group=ar_group, device=device)
if rank == 0:
print(" Using custom all-reduce (deterministic)")
except Exception as e:
if rank == 0:
print(f" Custom AR init failed: {e}, using NCCL fallback")
custom_ar = None
dist.barrier() # Ensure all ranks continue even if one fails
# Test different batch sizes - similar to test_ar.py
batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256, 512]
hidden_dim = 16384 # Fixed hidden dimension
num_trials = 10 # Same as test_ar.py
# Different seed per rank - each GPU has DIFFERENT input (like test_ar.py)
torch.manual_seed(42 + rank)
results = {}
for bs in batch_sizes:
# Create fixed input for all trials (like test_ar.py)
base_input = torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device)
dist.barrier()
if rank == 0:
print(f"\nBatch size {bs:4d}:")
print(f" Testing determinism across {num_trials} trials...")
# Test all-reduce determinism
results_ar = []
latencies_ar = []
for trial in range(num_trials):
# Clone the same input for each trial
inp_ar = base_input.clone()
inp_flat_ar = inp_ar.view(-1)
# Measure latency
torch.cuda.synchronize()
start = time.perf_counter()
dist.all_reduce(inp_flat_ar, op=dist.ReduceOp.SUM)
torch.cuda.synchronize()
end = time.perf_counter()
latencies_ar.append(end - start)
# Store checksum and first values (like test_ar.py)
checksum = inp_flat_ar.sum().item()
first_vals = inp_flat_ar[:5].clone()
results_ar.append((checksum, first_vals))
# Test reduce-scatter + all-gather determinism
results_rs_ag = []
latencies_rs_ag = []
for trial in range(num_trials):
# Clone the same input for each trial
inp_rs_ag = base_input.clone()
inp_flat_rs_ag = inp_rs_ag.view(-1)
# Measure latency
torch.cuda.synchronize()
start = time.perf_counter()
reduce_scatter_then_all_gather(
inp_flat_rs_ag, rank, world_size, custom_ar=None
)
torch.cuda.synchronize()
end = time.perf_counter()
latencies_rs_ag.append(end - start)
# Store checksum and first values (like test_ar.py)
checksum = inp_flat_rs_ag.sum().item()
first_vals = inp_flat_rs_ag[:5].clone()
results_rs_ag.append((checksum, first_vals))
# Note: sglang's optimized all-reduce requires full runtime initialization
# and is not tested in this standalone benchmark
use_sglang_optimized = False
results_optimized_rs_ag = []
latencies_optimized_rs_ag = []
# Test custom all-reduce determinism (if available)
results_custom_ar = []
latencies_custom_ar = []
if custom_ar is not None:
for trial in range(num_trials):
# Clone the same input for each trial
inp_custom = base_input.clone()
inp_flat_custom = inp_custom.view(-1)
# Measure latency
torch.cuda.synchronize()
start = time.perf_counter()
reduce_scatter_then_all_gather(
inp_flat_custom, rank, world_size, custom_ar=custom_ar
)
torch.cuda.synchronize()
end = time.perf_counter()
latencies_custom_ar.append(end - start)
# Store checksum and first values (like test_ar.py)
checksum = inp_flat_custom.sum().item()
first_vals = inp_flat_custom[:5].clone()
results_custom_ar.append((checksum, first_vals))
# Test deterministic kernel (if available)
results_deterministic_kernel = []
latencies_deterministic_kernel = []
deterministic_kernel_available = False
if custom_ar is not None:
# Check if input size fits in buffer
input_size_bytes = base_input.numel() * base_input.element_size()
if input_size_bytes > custom_ar.max_size:
if rank == 0:
print(
f" Deterministic kernel skipped: input size ({input_size_bytes/(1024*1024):.1f} MB) > buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)"
)
deterministic_kernel_available = False
else:
try:
deterministic_kernel_available = True
for trial in range(num_trials):
# Clone the same input for each trial
inp_kernel = base_input.clone()
# Measure latency
torch.cuda.synchronize()
start = time.perf_counter()
result_kernel = custom_ar.custom_all_reduce(inp_kernel)
torch.cuda.synchronize()
end = time.perf_counter()
latencies_deterministic_kernel.append(end - start)
# Store checksum and first values
result_flat_kernel = result_kernel.view(-1)
checksum = result_flat_kernel.sum().item()
first_vals = result_flat_kernel[:5].clone()
results_deterministic_kernel.append((checksum, first_vals))
except Exception as e:
if rank == 0:
print(
f" Deterministic kernel test failed for batch size {bs}: {e}"
)
deterministic_kernel_available = False
dist.barrier()
if rank == 0:
# Check determinism for all-reduce
ar_deterministic = True
ar_ref_sum, ar_ref_vals = results_ar[0]
ar_variance = []
for i, (s, vals) in enumerate(results_ar[1:], 1):
if abs(ar_ref_sum - s) > 1e-3 or not torch.allclose(
ar_ref_vals, vals, rtol=1e-3
):
ar_deterministic = False
ar_variance.append(abs(ar_ref_sum - s))
# Check determinism for reduce-scatter + all-gather
rs_ag_deterministic = True
rs_ag_ref_sum, rs_ag_ref_vals = results_rs_ag[0]
rs_ag_variance = []
for i, (s, vals) in enumerate(results_rs_ag[1:], 1):
if abs(rs_ag_ref_sum - s) > 1e-3 or not torch.allclose(
rs_ag_ref_vals, vals, rtol=1e-3
):
rs_ag_deterministic = False
rs_ag_variance.append(abs(rs_ag_ref_sum - s))
# Check determinism for optimized RS+AG (if available)
optimized_rs_ag_deterministic = None
optimized_rs_ag_max_variance = None
lat_optimized_rs_ag_median = None
if use_sglang_optimized and results_optimized_rs_ag:
optimized_rs_ag_deterministic = True
opt_rs_ag_ref_sum, opt_rs_ag_ref_vals = results_optimized_rs_ag[0]
opt_rs_ag_variance = []
for i, (s, vals) in enumerate(results_optimized_rs_ag[1:], 1):
if abs(opt_rs_ag_ref_sum - s) > 1e-3 or not torch.allclose(
opt_rs_ag_ref_vals, vals, rtol=1e-3
):
optimized_rs_ag_deterministic = False
opt_rs_ag_variance.append(abs(opt_rs_ag_ref_sum - s))
optimized_rs_ag_max_variance = (
max(opt_rs_ag_variance) if opt_rs_ag_variance else 0.0
)
lat_optimized_rs_ag_median = statistics.median(
latencies_optimized_rs_ag
)
# Check determinism for custom all-reduce (if available)
custom_ar_deterministic = None
custom_ar_max_variance = None
lat_custom_ar_median = None
if custom_ar is not None and results_custom_ar:
custom_ar_deterministic = True
custom_ar_ref_sum, custom_ar_ref_vals = results_custom_ar[0]
custom_ar_variance = []
for i, (s, vals) in enumerate(results_custom_ar[1:], 1):
if abs(custom_ar_ref_sum - s) > 1e-3 or not torch.allclose(
custom_ar_ref_vals, vals, rtol=1e-3
):
custom_ar_deterministic = False
custom_ar_variance.append(abs(custom_ar_ref_sum - s))
custom_ar_max_variance = (
max(custom_ar_variance) if custom_ar_variance else 0.0
)
lat_custom_ar_median = statistics.median(latencies_custom_ar)
# Check determinism for deterministic kernel (if available)
deterministic_kernel_deterministic = None
deterministic_kernel_max_variance = None
lat_deterministic_kernel_median = None
if deterministic_kernel_available and results_deterministic_kernel:
deterministic_kernel_deterministic = True
kernel_ref_sum, kernel_ref_vals = results_deterministic_kernel[0]
kernel_variance = []
for i, (s, vals) in enumerate(results_deterministic_kernel[1:], 1):
if abs(kernel_ref_sum - s) > 1e-3 or not torch.allclose(
kernel_ref_vals, vals, rtol=1e-3
):
deterministic_kernel_deterministic = False
kernel_variance.append(abs(kernel_ref_sum - s))
deterministic_kernel_max_variance = (
max(kernel_variance) if kernel_variance else 0.0
)
lat_deterministic_kernel_median = statistics.median(
latencies_deterministic_kernel
)
# Calculate latency statistics
lat_ar_median = statistics.median(latencies_ar)
lat_rs_ag_median = statistics.median(latencies_rs_ag)
overhead_rs_ag = ((lat_rs_ag_median - lat_ar_median) / lat_ar_median) * 100
# Calculate variance statistics
ar_max_variance = max(ar_variance) if ar_variance else 0.0
rs_ag_max_variance = max(rs_ag_variance) if rs_ag_variance else 0.0
results[bs] = {
"all_reduce": {
"latency_median": lat_ar_median,
"deterministic": ar_deterministic,
"max_variance": ar_max_variance,
},
"rs_ag": {
"latency_median": lat_rs_ag_median,
"deterministic": rs_ag_deterministic,
"max_variance": rs_ag_max_variance,
},
"custom_ar": (
{
"latency_median": lat_custom_ar_median,
"deterministic": custom_ar_deterministic,
"max_variance": custom_ar_max_variance,
}
if custom_ar is not None
else None
),
"deterministic_kernel": (
{
"latency_median": lat_deterministic_kernel_median,
"deterministic": deterministic_kernel_deterministic,
"max_variance": deterministic_kernel_max_variance,
}
if lat_deterministic_kernel_median is not None
else None
),
"optimized_rs_ag": (
{
"latency_median": lat_optimized_rs_ag_median,
"deterministic": optimized_rs_ag_deterministic,
"max_variance": optimized_rs_ag_max_variance,
}
if lat_optimized_rs_ag_median is not None
else None
),
"overhead_rs_ag_pct": overhead_rs_ag,
}
print(
f" All-Reduce: {lat_ar_median*1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}"
)
print(
f" RS+All-Gather: {lat_rs_ag_median*1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}"
)
if custom_ar is not None and lat_custom_ar_median is not None:
overhead_custom = (
(lat_custom_ar_median - lat_ar_median) / lat_ar_median
) * 100
print(
f" Custom AR: {lat_custom_ar_median*1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%"
)
if lat_deterministic_kernel_median is not None:
overhead_kernel = (
(lat_deterministic_kernel_median - lat_ar_median) / lat_ar_median
) * 100
speedup_kernel_vs_rs_ag = (
(lat_rs_ag_median - lat_deterministic_kernel_median)
/ lat_rs_ag_median
) * 100
print(
f" Deterministic Kernel: {lat_deterministic_kernel_median*1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%"
)
if lat_optimized_rs_ag_median is not None:
overhead_opt = (
(lat_optimized_rs_ag_median - lat_ar_median) / lat_ar_median
) * 100
speedup_vs_rs_ag = (
(lat_rs_ag_median - lat_optimized_rs_ag_median) / lat_rs_ag_median
) * 100
print(
f" Optimized RS+AG: {lat_optimized_rs_ag_median*1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%"
)
print(f" RS+AG Overhead: {overhead_rs_ag:+.1f}%")
if rank == 0:
results_queue.put(results)
dist.destroy_process_group()
def main():
world_size = 8
available_gpus = torch.cuda.device_count()
print("=" * 80)
print("All-Reduce vs Reduce-Scatter + All-Gather Determinism & Latency Benchmark")
print("=" * 80)
print(f"Available GPUs: {available_gpus}")
print(f"Using world_size: {world_size}")
print(f"Hidden dimension: 16384")
print(f"Tensor dtype: bfloat16")
print(f"Trials per batch size: 10 (testing determinism)")
print(f"Testing batch sizes: [1, 4, 8, 16, 32, 64, 128, 256, 512]")
print("=" * 80)
if available_gpus < world_size:
print(
f"WARNING: Only {available_gpus} GPUs available, using {available_gpus} instead"
)
world_size = available_gpus
if world_size < 2:
print("ERROR: Need at least 2 GPUs for this benchmark")
return
mp.set_start_method("spawn", force=True)
port = get_open_port()
results_queue = mp.Queue()
procs = []
for rank in range(world_size):
p = mp.Process(target=worker, args=(world_size, rank, port, results_queue))
p.start()
procs.append(p)
for p in procs:
p.join()
# Collect results
if not results_queue.empty():
results = results_queue.get()
print("\n" + "=" * 80)
print("SUMMARY")
print("=" * 80)
header = f"{'Batch':<8} {'AR (ms)':<12} {'AR Det':<8} {'RS+AG (ms)':<15} {'RS+AG Det':<10} {'RS+AG Ovh':<12}"
if any(r.get("custom_ar") is not None for r in results.values()):
header += (
f" {'Custom AR (ms)':<18} {'Custom AR Det':<15} {'Custom AR Ovh':<15}"
)
if any(r.get("deterministic_kernel") is not None for r in results.values()):
header += f" {'Det Kernel (ms)':<18} {'Det Kernel Det':<15} {'Det Kernel Ovh':<15} {'Speedup':<10}"
if any(r.get("optimized_rs_ag") is not None for r in results.values()):
header += f" {'Opt RS+AG (ms)':<18} {'Opt RS+AG Det':<15} {'Opt RS+AG Ovh':<15} {'Speedup':<10}"
print(header)
print("-" * 150)
for bs in sorted(results.keys()):
r = results[bs]
ar_det_str = "✓" if r["all_reduce"]["deterministic"] else "✗"
rs_ag_det_str = "✓" if r["rs_ag"]["deterministic"] else "✗"
line = (
f"{bs:<8} {r['all_reduce']['latency_median']*1000:<12.3f} {ar_det_str:<8} "
f"{r['rs_ag']['latency_median']*1000:<15.3f} {rs_ag_det_str:<10} "
f"{r['overhead_rs_ag_pct']:<12.1f}"
)
if r.get("custom_ar") is not None:
custom_ar = r["custom_ar"]
custom_ar_det_str = "✓" if custom_ar["deterministic"] else "✗"
custom_ar_overhead = (
(custom_ar["latency_median"] - r["all_reduce"]["latency_median"])
/ r["all_reduce"]["latency_median"]
) * 100
line += f" {custom_ar['latency_median']*1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}"
if r.get("deterministic_kernel") is not None:
det_kernel = r["deterministic_kernel"]
det_kernel_det_str = "✓" if det_kernel["deterministic"] else "✗"
det_kernel_overhead = (
(det_kernel["latency_median"] - r["all_reduce"]["latency_median"])
/ r["all_reduce"]["latency_median"]
) * 100
speedup_kernel = (
(r["rs_ag"]["latency_median"] - det_kernel["latency_median"])
/ r["rs_ag"]["latency_median"]
) * 100
line += f" {det_kernel['latency_median']*1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}"
if r.get("optimized_rs_ag") is not None:
opt_rs_ag = r["optimized_rs_ag"]
opt_rs_ag_det_str = "✓" if opt_rs_ag["deterministic"] else "✗"
opt_rs_ag_overhead = (
(opt_rs_ag["latency_median"] - r["all_reduce"]["latency_median"])
/ r["all_reduce"]["latency_median"]
) * 100
speedup = (
(r["rs_ag"]["latency_median"] - opt_rs_ag["latency_median"])
/ r["rs_ag"]["latency_median"]
) * 100
line += f" {opt_rs_ag['latency_median']*1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}"
print(line)
print("=" * 80)
# Calculate statistics
overheads_rs_ag = [r["overhead_rs_ag_pct"] for r in results.values()]
ar_deterministic_count = sum(
1 for r in results.values() if r["all_reduce"]["deterministic"]
)
rs_ag_deterministic_count = sum(
1 for r in results.values() if r["rs_ag"]["deterministic"]
)
custom_ar_deterministic_count = sum(
1
for r in results.values()
if r.get("custom_ar") and r["custom_ar"]["deterministic"]
)
custom_ar_total_count = sum(
1 for r in results.values() if r.get("custom_ar") is not None
)
deterministic_kernel_deterministic_count = sum(
1
for r in results.values()
if r.get("deterministic_kernel")
and r["deterministic_kernel"]["deterministic"]
)
deterministic_kernel_total_count = sum(
1 for r in results.values() if r.get("deterministic_kernel") is not None
)
print(f"\nDeterminism Summary:")
print(
f" All-Reduce deterministic: {ar_deterministic_count}/{len(results)} batch sizes"
)
print(
f" RS+All-Gather deterministic: {rs_ag_deterministic_count}/{len(results)} batch sizes"
)
if custom_ar_total_count > 0:
print(
f" Custom AR deterministic: {custom_ar_deterministic_count}/{custom_ar_total_count} batch sizes"
)
if deterministic_kernel_total_count > 0:
print(
f" Deterministic Kernel deterministic: {deterministic_kernel_deterministic_count}/{deterministic_kernel_total_count} batch sizes"
)
print(f"\nLatency Overhead Statistics (RS+AG vs All-Reduce):")
avg_overhead = statistics.mean(overheads_rs_ag)
median_overhead = statistics.median(overheads_rs_ag)
min_overhead = min(overheads_rs_ag)
max_overhead = max(overheads_rs_ag)
print(f" Average: {avg_overhead:.1f}%")
print(f" Median: {median_overhead:.1f}%")
print(f" Min: {min_overhead:.1f}%")
print(f" Max: {max_overhead:.1f}%")
if custom_ar_total_count > 0:
overheads_custom = []
for r in results.values():
if r.get("custom_ar") is not None:
overhead = (
(
r["custom_ar"]["latency_median"]
- r["all_reduce"]["latency_median"]
)
/ r["all_reduce"]["latency_median"]
) * 100
overheads_custom.append(overhead)
print(f"\nLatency Overhead Statistics (Custom AR vs All-Reduce):")
print(f" Average: {statistics.mean(overheads_custom):.1f}%")
print(f" Median: {statistics.median(overheads_custom):.1f}%")
print(f" Min: {min(overheads_custom):.1f}%")
print(f" Max: {max(overheads_custom):.1f}%")
if deterministic_kernel_total_count > 0:
overheads_kernel = []
speedups_kernel = []
for r in results.values():
if r.get("deterministic_kernel") is not None:
overhead = (
(
r["deterministic_kernel"]["latency_median"]
- r["all_reduce"]["latency_median"]
)
/ r["all_reduce"]["latency_median"]
) * 100
overheads_kernel.append(overhead)
speedup = (
(
r["rs_ag"]["latency_median"]
- r["deterministic_kernel"]["latency_median"]
)
/ r["rs_ag"]["latency_median"]
) * 100
speedups_kernel.append(speedup)
print(
f"\nLatency Overhead Statistics (Deterministic Kernel vs All-Reduce):"
)
print(f" Average: {statistics.mean(overheads_kernel):.1f}%")
print(f" Median: {statistics.median(overheads_kernel):.1f}%")
print(f" Min: {min(overheads_kernel):.1f}%")
print(f" Max: {max(overheads_kernel):.1f}%")
print(f"\nSpeedup Statistics (Deterministic Kernel vs RS+AG):")
print(f" Average: {statistics.mean(speedups_kernel):.1f}%")
print(f" Median: {statistics.median(speedups_kernel):.1f}%")
print(f" Min: {min(speedups_kernel):.1f}%")
print(f" Max: {max(speedups_kernel):.1f}%")
# Show variance for non-deterministic cases
print(f"\nVariance Analysis (non-deterministic cases):")
for bs in sorted(results.keys()):
r = results[bs]
if not r["all_reduce"]["deterministic"]:
print(
f" Batch {bs}: All-Reduce max variance: {r['all_reduce']['max_variance']:.6f}"
)
if not r["rs_ag"]["deterministic"]:
print(
f" Batch {bs}: RS+All-Gather max variance: {r['rs_ag']['max_variance']:.6f}"
)
if r.get("custom_ar") is not None and not r["custom_ar"]["deterministic"]:
print(
f" Batch {bs}: Custom AR max variance: {r['custom_ar']['max_variance']:.6f}"
)
if (
r.get("deterministic_kernel") is not None
and not r["deterministic_kernel"]["deterministic"]
):
print(
f" Batch {bs}: Deterministic Kernel max variance: {r['deterministic_kernel']['max_variance']:.6f}"
)
if __name__ == "__main__":
main()
python/sglang/kernels/aot/benchmark/bench_awq_dequant.py
+151
View File
@@ -0,0 +1,151 @@
import itertools
import os
from typing import List, Tuple
import torch
import triton
import triton.testing
from sgl_kernel import awq_dequantize
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm import _custom_ops as ops
VLLM_AVAILABLE = True
except ImportError:
ops = None
VLLM_AVAILABLE = False
IS_CI = is_in_ci()
def vllm_awq_dequantize(
qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
if not VLLM_AVAILABLE:
# Fallback to SGLang implementation
return sglang_awq_dequantize(qweight, scales, qzeros)
return ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
def sglang_awq_dequantize(
qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
return awq_dequantize(qweight, scales, qzeros)
def calculate_diff(qweight_row: int, qweight_col: int):
"""Calculate difference between VLLM and SGLang implementations."""
device = torch.device("cuda")
qweight = torch.randint(
0,
torch.iinfo(torch.int32).max,
(qweight_row, qweight_col),
dtype=torch.int32,
device=device,
)
group_size = qweight_row
scales_row = qweight_row // group_size
scales_col = qweight_col * 8
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
qzeros = torch.randint(
0,
torch.iinfo(torch.int32).max,
(scales_row, qweight_col),
dtype=torch.int32,
device=device,
)
if not VLLM_AVAILABLE:
print("⚠️ vLLM not available, skipping comparison")
return
vllm_out = vllm_awq_dequantize(qweight, scales, qzeros)
sglang_out = sglang_awq_dequantize(qweight, scales, qzeros)
output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
if torch.allclose(
vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
):
print("✅ All implementations match")
else:
print("❌ Implementations differ")
# CI environment uses simplified parameters
if IS_CI:
qweight_row_range = [128] # Single row size for CI
qweight_cols_range = [16] # Single column size for CI
else:
qweight_row_range = [3584, 18944, 128, 256, 512, 1024]
qweight_cols_range = [448, 576, 4736, 16, 32, 64, 128]
configs = list(itertools.product(qweight_row_range, qweight_cols_range))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["qweight_row", "qweight_col"],
x_vals=configs,
line_arg="provider",
line_vals=["vllm", "sglang"] if VLLM_AVAILABLE else ["sglang"],
line_names=["VLLM", "SGL Kernel"] if VLLM_AVAILABLE else ["SGL Kernel"],
styles=[("blue", "-"), ("green", "-")] if VLLM_AVAILABLE else [("green", "-")],
ylabel="us",
plot_name="awq-dequantize-performance",
args={},
)
)
def benchmark(qweight_row, qweight_col, provider):
dtype = torch.float16
device = torch.device("cuda")
qweight = torch.randint(
0,
torch.iinfo(torch.int32).max,
(qweight_row, qweight_col),
dtype=torch.int32,
device=device,
)
group_size = qweight_row
scales_row = qweight_row // group_size
scales_col = qweight_col * 8
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
qzeros = torch.randint(
0,
torch.iinfo(torch.int32).max,
(scales_row, qweight_col),
dtype=torch.int32,
device=device,
)
quantiles = [0.5, 0.2, 0.8]
if provider == "vllm":
if not VLLM_AVAILABLE:
return (0, 0, 0)
fn = lambda: vllm_awq_dequantize(
qweight.clone(), scales.clone(), qzeros.clone()
)
elif provider == "sglang":
fn = lambda: sglang_awq_dequantize(
qweight.clone(), scales.clone(), qzeros.clone()
)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
# Simplify for CI environment
if IS_CI:
qweight_row, qweight_col = 128, 16 # Smaller values for CI
else:
qweight_row, qweight_col = 3584, 448
calculate_diff(qweight_row=qweight_row, qweight_col=qweight_col)
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_cutlass_mla.py
+175
View File
@@ -0,0 +1,175 @@
import argparse
import copy
import itertools
import os
import torch
import triton
from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size
from sglang.srt.utils import get_device_capability
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
# CI environment uses simplified parameters
if IS_CI:
bs_range = [1] # Single batch size for CI
qlen_range = [64] # Single sequence length for CI
else:
bs_range = [1, 8, 32, 64, 128, 256]
qlen_range = [1, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
configs = list(itertools.product(bs_range, qlen_range))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "seq_len"],
x_vals=configs,
x_log=False,
line_arg="provider",
line_vals=[
"128 heads",
"64 heads",
"32 heads",
"16 heads",
],
line_names=[
"128 heads",
"64 heads",
"32 heads",
"16 heads",
],
styles=[("green", "-"), ("green", "--"), ("blue", "-"), ("blue", "--")],
ylabel="GB/s",
plot_name="cutlass mla",
args={},
)
)
def benchmark(batch_size, seq_len, provider, block_size, num_kv_splits):
d = 576
dn = 64
dv = 512
h_q_map = {
"128": 128,
"64": 64,
"32": 32,
"16": 16,
}
parsed_h_q = next(
(value for key, value in h_q_map.items() if key in provider), None
)
if parsed_h_q is None:
raise ValueError(f"Unknown head configuration in provider: {provider}")
h_q = parsed_h_q
seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
max_seq_len = seq_lens.max().item()
block_num = (max_seq_len + block_size - 1) // block_size
# Pad block_num so that small blocks can be packed into full 128-sized CUTLASS tiles.
# One 128-wide tile can hold (128 // block_size) small blocks.
pack_factor = 128 // block_size
block_num = ((block_num + pack_factor - 1) // pack_factor) * pack_factor
qn = (
torch.randn(h_q, batch_size, d - dn, dtype=torch.bfloat16, device="cuda")
* 100.0
)
qr = torch.randn(batch_size, h_q, dn, dtype=torch.bfloat16, device="cuda") * 100.0
block_table = torch.randint(
0,
batch_size * block_num,
(batch_size, block_num),
dtype=torch.int32,
device="cuda",
)
kv_cache = torch.randn(
block_table.numel(), block_size, d, dtype=torch.bfloat16, device="cuda"
)
workspace_size = cutlass_mla_get_workspace_size(
block_num * block_size, batch_size, num_kv_splits=num_kv_splits
)
workspace = torch.empty(workspace_size, device="cuda", dtype=torch.uint8)
quantiles = [0.5, 0.2, 0.8]
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: cutlass_mla_decode(
qn.transpose(0, 1),
qr,
kv_cache,
seq_lens,
block_table,
workspace,
1.44,
num_kv_splits,
),
quantiles=quantiles,
)
q_size = qn.numel() * qn.element_size() + qr.numel() * qr.element_size()
gbps = (
lambda ms: (
q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size()
)
* 1e-9
/ (ms * 1e-3)
)
return gbps(ms), gbps(max_ms), gbps(min_ms)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--block-sizes",
nargs="+",
type=int,
default=[1, 32, 64, 128],
help="List of batch sizes",
)
parser.add_argument(
"--num-kv-splits",
nargs="+",
type=int,
default=[-1],
help="List of batch sizes",
)
args = parser.parse_args()
# Skip in CI environment or unsupported architectures
if IS_CI:
major, minor = get_device_capability()
if major is None or major < 10: # Requires compute capability 10.0+
print("Skipping Cutlass MLA benchmark in CI environment")
if major is not None:
print(
f"Cutlass MLA requires compute capability 10.0+, but found {major}.{minor}"
)
else:
print("Could not determine device capability")
else:
for block_size in args.block_sizes:
for kv_split in args.num_kv_splits:
print(f"block_size={block_size}, num_kv_splits={kv_split}: ")
benchmark.run(
print_data=True,
block_size=block_size,
num_kv_splits=kv_split,
)
print("Benchmark finished!")
else:
for block_size in args.block_sizes:
for kv_split in args.num_kv_splits:
print(f"block_size={block_size}, num_kv_splits={kv_split}: ")
benchmark.run(
print_data=True,
block_size=block_size,
num_kv_splits=kv_split,
)
print("Benchmark finished!")
python/sglang/kernels/aot/benchmark/bench_dsv4_norm_rope.py
+75
View File
@@ -0,0 +1,75 @@
"""Benchmark for DeepSeek-V4 fused norm + RoPE kernels."""
import itertools
import sgl_kernel
import torch
import triton
import triton.testing
try:
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
except ImportError:
IS_CI = False
batch_sizes = [1] if IS_CI else [1, 4, 16, 64, 256]
num_heads_list = [8] if IS_CI else [8, 16, 64]
head_dims = [192] if IS_CI else [128, 192]
configs = list(itertools.product(batch_sizes, num_heads_list, head_dims))
def torch_rmsnorm_rope(
q: torch.Tensor, freqs_cis: torch.Tensor, positions: torch.Tensor, eps: float
) -> torch.Tensor:
"""Naive PyTorch reference: RMSNorm + RoPE."""
rms = torch.sqrt(q.float().pow(2).mean(dim=-1, keepdim=True) + eps)
q_normed = (q.float() / rms).to(q.dtype)
return q_normed
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "num_heads", "head_dim"],
x_vals=configs,
line_arg="provider",
line_vals=["sglang", "torch"],
line_names=["SGL Kernel", "PyTorch"],
styles=[("green", "-"), ("red", "--")],
ylabel="µs (median)",
plot_name="dsv4-q-norm-rope-performance",
args={},
)
)
def benchmark_q_norm_rope(batch_size, num_heads, head_dim, provider):
torch.manual_seed(42)
eps = 1e-6
max_pos = 8192
rope_dim = 64
q_input = torch.randn(
batch_size, num_heads, head_dim, dtype=torch.bfloat16, device="cuda"
)
q_output = torch.empty_like(q_input)
freqs_cis = torch.randn(max_pos, rope_dim, dtype=torch.float32, device="cuda")
positions = torch.randint(
0, max_pos, (batch_size,), dtype=torch.int32, device="cuda"
)
if provider == "sglang":
fn = lambda: sgl_kernel.dsv4_fused_q_norm_rope(
q_input, freqs_cis, positions, eps, q_output
)
else:
fn = lambda: torch_rmsnorm_rope(q_input, freqs_cis, positions, eps)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
fn, quantiles=[0.5, 0.2, 0.8]
)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
benchmark_q_norm_rope.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_es_fp8_blockwise_grouped_gemm.py
+370
View File
@@ -0,0 +1,370 @@
import argparse
import random
from dataclasses import dataclass
from typing import List, Tuple
import numpy as np
import torch
from sgl_kernel import (
es_fp8_blockwise_scaled_grouped_mm,
fp8_blockwise_scaled_grouped_mm,
)
random.seed(28)
def ceil_div(x: int, y: int) -> int:
return (x + y - 1) // y
def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
pad_size = (128 - (n % 128)) % 128
x = torch.nn.functional.pad(x, (0, pad_size), value=0) if pad_size > 0 else x
x_view = x.view(m, -1, 128)
x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn)
return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
x_padded = torch.zeros(
(ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device
)
x_padded[:m, :n] = x
x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view(
x_view.size(0), x_view.size(2)
)
def create_unbalanced_expert_token_distribution(
batch_size: int, topk: int, num_experts: int
):
expert_ids = np.random.randint(0, num_experts, size=(batch_size * topk,)).tolist()
expert_to_count = dict()
for expert_id in range(num_experts):
expert_to_count[expert_id] = 0
for expert_id in expert_ids:
expert_to_count[expert_id] += 1
group_ms = []
for expert_id in range(num_experts):
group_ms.append(expert_to_count[expert_id])
return group_ms
def bench_es(
group_ms: List[int],
n: int,
k: int,
num_groups: int,
num_warmup: int,
num_run: int,
) -> Tuple[float, int]:
device = "cuda"
alignment = 128
n_g = ceil_div(n, alignment) * alignment
k_g = ceil_div(k, alignment) * alignment
out_dtype = torch.bfloat16
expert_offsets = torch.zeros((num_groups + 1), device=device, dtype=torch.int32)
problem_sizes = torch.zeros((num_groups, 3), device=device, dtype=torch.int32)
a_tensors = []
b_tensors = []
a_scales_tensors = []
b_scales_tensors = []
for g in range(num_groups):
m_g = group_ms[g]
expert_offsets[g + 1] = expert_offsets[g] + m_g
problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
if m_g != 0:
a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
a_tensors.append(a_g)
a_scales_tensors.append(a_scale)
b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
b_tensors.append(b_g)
b_scales_tensors.append(b_scale)
a_stack = torch.empty(
(expert_offsets[-1], k_g), device=device, dtype=torch.float8_e4m3fn
)
b_stack = torch.empty(
(num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
)
_aux_idx = 0
for g in range(num_groups):
if group_ms[g] != 0:
a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[_aux_idx]
_aux_idx += 1
b_stack[g] = b_tensors[g].t()
b_stack = b_stack.transpose(1, 2)
a_scale_stack = torch.empty(
(expert_offsets[-1], k_g // 128), device=device, dtype=torch.float32
)
b_scale_stack = torch.empty(
(num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32
)
_aux_idx = 0
for g in range(num_groups):
if group_ms[g] != 0:
a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[
_aux_idx
]
_aux_idx += 1
b_scale_stack[g] = b_scales_tensors[g].t()
b_scale_stack = b_scale_stack.transpose(1, 2)
workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8)
c_out = torch.empty((expert_offsets[-1], n_g), device=device, dtype=out_dtype)
a_strides = torch.full(
(num_groups,), a_stack.stride(0), device=device, dtype=torch.int64
)
d_strides = torch.full(
(num_groups,), c_out.stride(0), device=device, dtype=torch.int64
)
def run_cutlass():
es_fp8_blockwise_scaled_grouped_mm(
c_out,
a_stack,
b_stack,
a_scale_stack,
b_scale_stack,
a_strides,
a_strides,
d_strides,
problem_sizes,
expert_offsets[:-1],
workspace,
)
run_cutlass()
# warmup
for _ in range(num_warmup):
run_cutlass()
torch.cuda.synchronize()
# run
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
for _ in range(num_run):
run_cutlass()
end_event.record()
end_event.synchronize()
torch.cuda.synchronize()
avg = start_event.elapsed_time(end_event) / num_run * 1000 # us
return avg, expert_offsets[-1]
def bench_sgl(
group_ms: List[int],
n: int,
k: int,
num_groups: int,
num_warmup: int,
num_run: int,
) -> Tuple[float, int]:
device = "cuda"
alignment = 128
n_g = ceil_div(n, alignment) * alignment
k_g = ceil_div(k, alignment) * alignment
out_dtype = torch.bfloat16
expert_offsets = torch.zeros((num_groups + 1), device=device, dtype=torch.int32)
problem_sizes = torch.zeros((num_groups, 3), device=device, dtype=torch.int32)
layout_sfa = torch.zeros((num_groups, 5), device=device, dtype=torch.int32)
layout_sfb = torch.zeros((num_groups, 5), device=device, dtype=torch.int32)
a_tensors = []
b_tensors = []
a_scales_tensors = []
b_scales_tensors = []
for g in range(num_groups):
m_g = group_ms[g]
expert_offsets[g + 1] = expert_offsets[g] + m_g
problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
if m_g != 0:
a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
a_tensors.append(a_g)
a_scales_tensors.append(a_scale)
b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
b_tensors.append(b_g)
b_scales_tensors.append(b_scale)
a_stack = torch.empty(
(expert_offsets[-1], k_g), device=device, dtype=torch.float8_e4m3fn
)
b_stack = torch.empty(
(num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
)
_aux_idx = 0
for g in range(num_groups):
if group_ms[g] != 0:
a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[_aux_idx]
_aux_idx += 1
b_stack[g] = b_tensors[g].t()
b_stack = b_stack.transpose(1, 2)
a_scale_stack = torch.empty(
(expert_offsets[-1], k_g // 128), device=device, dtype=torch.float32
)
b_scale_stack = torch.empty(
(num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32
)
_aux_idx = 0
for g in range(num_groups):
if group_ms[g] != 0:
a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[
_aux_idx
]
_aux_idx += 1
b_scale_stack[g] = b_scales_tensors[g].t()
b_scale_stack = b_scale_stack.transpose(1, 2)
c_out = torch.empty((expert_offsets[-1], n_g), device=device, dtype=out_dtype)
a_strides = torch.full(
(num_groups,), a_stack.stride(0), device=device, dtype=torch.int64
)
c_strides = torch.full(
(num_groups,), c_out.stride(0), device=device, dtype=torch.int64
)
workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8)
a_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
b_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
out_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
a_scales_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
b_scales_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
def run_cutlass():
fp8_blockwise_scaled_grouped_mm(
c_out,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
a_stack,
b_stack,
a_scale_stack,
b_scale_stack,
a_strides,
a_strides,
c_strides,
layout_sfa,
layout_sfb,
problem_sizes,
expert_offsets[:-1],
workspace,
)
# warmup
for _ in range(num_warmup):
run_cutlass()
torch.cuda.synchronize()
# run
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
for _ in range(num_run):
run_cutlass()
end_event.record()
end_event.synchronize()
torch.cuda.synchronize()
avg = start_event.elapsed_time(end_event) / num_run * 1000 # us
return avg, expert_offsets[-1]
benchmark_kernels = {"es": bench_es, "sgl-kernel": bench_sgl}
@dataclass
class ShapeArg:
n: int
k: int
num_groups: int
def benchmark_one_shape(
shape_args: List[ShapeArg],
num_warmup: int,
num_run: int,
):
for shape in shape_args:
for batch_size in [
128,
256,
384,
512,
640,
768,
896,
1024,
1280,
1536,
2048,
3072,
]:
group_ms = create_unbalanced_expert_token_distribution(
batch_size, 8, shape.num_groups
)
print(
f"\nBenchmark: batch_size={batch_size}, n={shape.n}, k={shape.k}, num_groups={shape.num_groups}"
)
for kernel_name, kernel_func in benchmark_kernels.items():
average_time, m = kernel_func(
group_ms,
shape.n,
shape.k,
shape.num_groups,
num_warmup,
num_run,
)
print(f"{kernel_name}: {average_time} us")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--num-warmup", type=int, default=3)
parser.add_argument("--num-run", type=int, default=20)
shape_args = [
# DeepSeek-R1, gateup, TP = 8
ShapeArg(n=512, k=7168, num_groups=256),
# DeepSeek-R1, down, TP = 8
ShapeArg(n=7168, k=256, num_groups=256),
# DeepSeek-R1, gateup, TP = 4
ShapeArg(n=1024, k=7168, num_groups=256),
# DeepSeek-R1, down, TP = 4
ShapeArg(n=7168, k=512, num_groups=256),
# Qwen3-235B-A22B-FP8, gateup, TP = 4
ShapeArg(n=768, k=4096, num_groups=128),
# Qwen3-235B-A22B-FP8, down, TP = 4
ShapeArg(n=4096, k=384, num_groups=128),
# Qwen3-235B-A22B-FP8, gateup, TP = 2
ShapeArg(n=1536, k=4096, num_groups=128),
# Qwen3-235B-A22B-FP8, down, TP = 2
ShapeArg(n=4096, k=768, num_groups=128),
]
args = parser.parse_args()
benchmark_one_shape(shape_args, args.num_warmup, args.num_run)
if __name__ == "__main__":
main()
python/sglang/kernels/aot/benchmark/bench_fp4_gemm.py
Executable
+462
View File
@@ -0,0 +1,462 @@
import argparse
import csv
import logging
from functools import partial
from typing import List, Tuple
import torch
import triton
from flashinfer import fp4_quantize, mm_fp4
from flashinfer.autotuner import autotune
from flashinfer.jit.core import logger as flashinfer_logger
from flashinfer.testing import bench_gpu_time
flashinfer_logger.setLevel(logging.ERROR)
from sglang.srt.utils import (
get_device_capability,
is_sm100_supported,
is_sm120_supported,
)
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
DEEPSEEK_R1_MODEL = "deepseek-ai/DeepSeek-R1-0528-FP4"
# Weight shapes are in the format: ([K, N], TP_SPLIT_DIM)
# TP split dim 0 means split K by tp size; dim 1 means split N by tp size.
WEIGHT_SHAPES = {
"meta-llama/Llama-3.1-8B-Instruct": [
([4096, 6144], 1),
([4096, 4096], 0),
([4096, 28672], 1),
([14336, 4096], 0),
],
"meta-llama/Llama-3.3-70B-Instruct": [
([8192, 10240], 1),
([8192, 8192], 0),
([8192, 57344], 1),
([28672, 8192], 0),
],
"mistralai/Mistral-Large-Instruct-2407": [
([12288, 14336], 1),
([12288, 12288], 0),
([12288, 57344], 1),
([28672, 12288], 0),
],
"Qwen/Qwen2.5-7B-Instruct": [
([3584, 4608], 1),
([3584, 3584], 0),
([3584, 37888], 1),
([18944, 3584], 0),
],
"Qwen/Qwen2.5-32B-Instruct": [
([5120, 7168], 1),
([5120, 5120], 0),
([5120, 55296], 1),
([27648, 5120], 0),
],
"Qwen/Qwen2.5-72B-Instruct": [
([8192, 10240], 1),
([8192, 8192], 0),
([8192, 59136], 1),
([29568, 8192], 0),
],
"Qwen/Qwen3.5-27B": [
([5120, 8192], 1),
([6144, 5120], 0),
([5120, 34816], 1),
([17408, 5120], 0),
],
"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
([2048, 3072], 1),
([2048, 4096], 1),
([2048, 2048], 0),
([2048, 576], 0),
([2048, 21888], 1),
([10944, 2048], 0),
([2048, 2816], 1),
([1408, 2048], 0),
],
}
DEEPSEEK_R1_WEIGHT_SHAPES = {
4: [[1024, 3584], [7168, 256], [7168, 2304], [9216, 3584]],
8: [[512, 3584], [7168, 128], [7168, 1152], [4608, 3584]],
}
def get_weight_shapes(args) -> List[Tuple[int, int, str]]:
shapes: List[Tuple[int, int, str]] = []
for model in args.models:
if model == DEEPSEEK_R1_MODEL:
for tp_size in args.tp_sizes:
if tp_size in DEEPSEEK_R1_WEIGHT_SHAPES:
selected = DEEPSEEK_R1_WEIGHT_SHAPES[tp_size]
else:
selected = (
DEEPSEEK_R1_WEIGHT_SHAPES[4] + DEEPSEEK_R1_WEIGHT_SHAPES[8]
)
for n, packed_k in selected:
shapes.append((n, packed_k, model))
continue
if model not in WEIGHT_SHAPES:
raise ValueError(f"Unsupported model: {model}")
for tp_size in args.tp_sizes:
for k_n, tp_split_dim in WEIGHT_SHAPES[model]:
k, n = k_n
if tp_split_dim == 0:
k = k // tp_size
else:
n = n // tp_size
packed_k = k // 2
shapes.append((n, packed_k, model))
return shapes
if IS_CI:
batch_sizes = [1, 8]
else:
batch_sizes = [
1,
2,
4,
8,
16,
32,
64,
128,
256,
512,
1024,
2048,
3072,
4096,
8192,
16384,
]
def _run_mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend):
return mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend=backend)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size"],
x_vals=batch_sizes,
x_log=False,
line_arg="provider",
line_vals=(
["cutlass", "cudnn", "trtllm", "cute-dsl", "auto"]
if is_sm100_supported()
else ["cutlass", "cudnn", "cute-dsl", "auto"]
),
line_names=(
[
"flashinfer cutlass fp4",
"cudnn fp4",
"trtllm fp4",
"cute-dsl fp4",
"auto fp4 (cudnn/cutlass)",
]
if is_sm100_supported()
else [
"flashinfer cutlass fp4",
"cudnn fp4",
"cute-dsl fp4",
"auto fp4",
]
),
styles=(
[
("orange", "solid"),
("blue", "solid"),
("green", "solid"),
("brown", "solid"),
("purple", "solid"),
]
if is_sm100_supported()
else [
("orange", "solid"),
("blue", "solid"),
("brown", "solid"),
("purple", "solid"),
]
),
ylabel="bandwidth (GB/s)",
plot_name="fp4_gemm_benchmark",
args={},
)
)
def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
M = batch_size
packed_k = K
K = 2 * packed_k
a_dtype = torch.randn((M, K), dtype=dtype, device="cuda")
b_dtype = torch.randn((N, K), dtype=dtype, device="cuda")
a_global_scale = (
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
).to(torch.float32)
b_global_scale = (
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
).to(torch.float32)
alpha = 1.0 / (a_global_scale * b_global_scale)
a_fp4, a_scale_interleaved = fp4_quantize(a_dtype, a_global_scale)
b_fp4, b_scale_interleaved = fp4_quantize(b_dtype, b_global_scale)
# flashinfer.fp4_quantize returns scale factors as uint8 (e4m3fn bits stored
# in uint8 memory); the JIT cutlass kernel requires float8_e4m3fn dtype.
if a_scale_interleaved.dtype != torch.float8_e4m3fn:
a_scale_interleaved = a_scale_interleaved.view(torch.float8_e4m3fn)
if b_scale_interleaved.dtype != torch.float8_e4m3fn:
b_scale_interleaved = b_scale_interleaved.view(torch.float8_e4m3fn)
b_fp4_T = b_fp4.T
b_sf_T = b_scale_interleaved.T
res_fi = torch.empty((M, N), dtype=dtype, device="cuda")
if provider == "cutlass":
with autotune():
_run_mm_fp4(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
backend="cutlass",
)
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="cutlass"),
input_args=(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
),
use_cuda_graph=True,
)
elif provider == "cudnn":
with autotune():
_run_mm_fp4(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
backend="cudnn",
)
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="cudnn"),
input_args=(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
),
use_cuda_graph=True,
)
elif provider == "trtllm":
a_sf_u8 = a_scale_interleaved.to(torch.uint8)
b_sf_u8_T = b_sf_T.to(torch.uint8)
with autotune():
_run_mm_fp4(
a_fp4,
b_fp4_T,
a_sf_u8,
b_sf_u8_T,
alpha,
dtype,
res_fi,
backend="trtllm",
)
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="trtllm"),
input_args=(a_fp4, b_fp4_T, a_sf_u8, b_sf_u8_T, alpha, dtype, res_fi),
use_cuda_graph=True,
)
elif provider == "cute-dsl":
with autotune():
_run_mm_fp4(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
backend="cute-dsl",
)
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="cute-dsl"),
input_args=(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
),
use_cuda_graph=True,
)
elif provider == "auto":
with autotune():
_run_mm_fp4(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
backend="auto",
)
times_ms = bench_gpu_time(
fn=partial(_run_mm_fp4, backend="auto"),
input_args=(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
),
use_cuda_graph=True,
)
ms = torch.tensor(times_ms).median().item()
# A: M×packed_k bytes (fp4 packed), B: N×packed_k bytes, C: M×N×element_size bytes
element_size = torch.finfo(dtype).bits // 8
total_bytes = M * packed_k + N * packed_k + M * N * element_size
bandwidth_gbs = total_bytes / (ms * 1e-3) / 1e9
if correctness:
res_cutlass = torch.empty((M, N), dtype=dtype, device="cuda")
mm_fp4(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_cutlass,
backend="cutlass",
)
mm_fp4(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
backend="cudnn",
)
assert torch.allclose(
res_fi, res_cutlass, atol=1e-3, rtol=1e-3
), "cudnn fp4 doesn't match cutlass fp4"
mm_fp4(
a_fp4,
b_fp4_T,
a_scale_interleaved,
b_sf_T,
alpha,
dtype,
res_fi,
backend="trtllm",
)
assert torch.allclose(
res_fi, res_cutlass, atol=1e-3, rtol=1e-3
), "trtllm fp4 doesn't match cutlass fp4"
if csv_file:
with open(csv_file, "a", newline="") as f:
writer = csv.writer(f)
writer.writerow([provider, M, N, K, ms, bandwidth_gbs])
return bandwidth_gbs
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--models",
nargs="+",
type=str,
default=[DEEPSEEK_R1_MODEL],
help="List of models to benchmark. Supported: Llama 8B/70B, Qwen, Mistral, DeepSeek.",
)
parser.add_argument(
"--tp-sizes",
nargs="+",
type=int,
default=[1],
help="List of tensor parallel sizes",
)
parser.add_argument(
"--dtype",
type=torch.dtype,
default=torch.bfloat16,
help="Output data type",
)
parser.add_argument(
"--correctness",
action="store_true",
help="Check correctness",
)
parser.add_argument(
"--csv",
type=str,
default="results_cutlass_cudnn.csv",
help="CSV file to save results",
)
args = parser.parse_args()
if IS_CI:
args.tp_sizes = [args.tp_sizes[0]]
if args.csv:
with open(args.csv, "w", newline="") as f:
writer = csv.writer(f)
writer.writerow(["provider", "m", "n", "k", "time_ms", "bandwidth_gbs"])
major, minor = get_device_capability()
if not (is_sm100_supported() or is_sm120_supported()):
print("Skipping FP4 GEMM benchmark")
if major is not None:
print(f"FP4 operations require sm100+, but found sm{major}{minor}")
else:
print("Could not determine device capability")
else:
NKs = get_weight_shapes(args)
if IS_CI:
NKs = NKs[:2]
for N, K, model_name in NKs:
print(f"{model_name} N={N} packed_k={K}: ")
benchmark.run(
print_data=True,
N=N,
K=K,
dtype=args.dtype,
correctness=args.correctness,
csv_file=args.csv,
)
print("Benchmark finished!")
python/sglang/kernels/aot/benchmark/bench_fp8_blockwise_group_gemm.py
+340
View File
@@ -0,0 +1,340 @@
import argparse
import random
from dataclasses import dataclass
from typing import List, Tuple
import deep_gemm
import torch
from sgl_kernel import fp8_blockwise_scaled_grouped_mm
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
def get_m_alignment_for_contiguous_layout():
return 128
def ceil_div(x: int, y: int) -> int:
return (x + y - 1) // y
def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
pad_size = (128 - (n % 128)) % 128
x = torch.nn.functional.pad(x, (0, pad_size), value=0) if pad_size > 0 else x
x_view = x.view(m, -1, 128)
x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn)
return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
assert x.dim() == 2
m, n = x.shape
x_padded = torch.zeros(
(ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device
)
x_padded[:m, :n] = x
x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view(
x_view.size(0), x_view.size(2)
)
def construct_contiguous_grouped(
num_groups: int, expected_m_per_group: int, k: int, n: int
) -> Tuple[
int,
Tuple[torch.Tensor, torch.Tensor],
Tuple[torch.Tensor, torch.Tensor],
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
alignment = get_m_alignment_for_contiguous_layout()
group_ms = [int(expected_m_per_group) for _ in range(num_groups)]
m = sum([ceil_div(x, alignment) * alignment for x in group_ms])
x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
y = torch.randn((num_groups, n, k), device="cuda", dtype=torch.bfloat16)
m_indices = torch.empty(m, device="cuda", dtype=torch.int32)
out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
start = 0
for i, group_m in enumerate(group_ms):
actual_end = start + group_m
aligned_end = start + ceil_div(group_m, alignment) * alignment
m_indices[start:actual_end] = i
m_indices[actual_end:aligned_end] = -1
start = aligned_end
assert m % 4 == 0, f"TMA alignment error: {m}"
x_fp8 = per_token_cast_to_fp8(x)
y_fp8 = (
torch.empty_like(y, dtype=torch.float8_e4m3fn),
torch.empty(
(num_groups, ceil_div(n, 128), k // 128), device="cuda", dtype=torch.float
),
)
for i in range(num_groups):
y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i])
return m, x_fp8, y_fp8, m_indices, out
def bench_deepgemm(
expected_m_per_group: int,
n: int,
k: int,
num_groups: int,
num_warmup: int,
num_run: int,
) -> Tuple[float, int]:
# construct tensors
m, x_fp8, y_fp8, m_indices, out = construct_contiguous_grouped(
num_groups, expected_m_per_group, k, n
)
def run_deepgemm():
deep_gemm.m_grouped_fp8_gemm_nt_contiguous(x_fp8, y_fp8, out, m_indices)
# warmup
for _ in range(num_warmup):
run_deepgemm()
torch.cuda.synchronize()
# run
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
latencies: list[float] = []
start_event.record()
for _ in range(num_run):
run_deepgemm()
end_event.record()
end_event.synchronize()
torch.cuda.synchronize()
avg = start_event.elapsed_time(end_event) / num_run * 1000 # us
return avg, m
def bench_cutlass(
expected_m_per_group: int,
n: int,
k: int,
num_groups: int,
num_warmup: int,
num_run: int,
) -> Tuple[float, int]:
device = "cuda"
alignment = 16
n_g = ceil_div(n, alignment) * alignment
k_g = ceil_div(k, alignment) * alignment
out_dtype = torch.bfloat16
expert_offsets = torch.zeros((num_groups + 1), device=device, dtype=torch.int32)
problem_sizes = torch.zeros((num_groups, 3), device=device, dtype=torch.int32)
layout_sfa = torch.zeros((num_groups, 5), device=device, dtype=torch.int32)
layout_sfb = torch.zeros((num_groups, 5), device=device, dtype=torch.int32)
a_tensors = []
b_tensors = []
a_scales_tensors = []
b_scales_tensors = []
# TODO(@TianQiLin666666): Unique group_ms in all bench function
group_ms = [
alignment * ceil_div(int(expected_m_per_group), alignment)
for _ in range(num_groups)
]
for g in range(num_groups):
m_g = group_ms[g]
expert_offsets[g + 1] = expert_offsets[g] + m_g
problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device)
a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device))
b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t())
a_tensors.append(a_g)
b_tensors.append(b_g)
a_scales_tensors.append(a_scale)
b_scales_tensors.append(b_scale)
a_stack = torch.empty(
(expert_offsets[-1], k_g), device=device, dtype=torch.float8_e4m3fn
)
b_stack = torch.empty(
(num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
)
for g in range(num_groups):
a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[g]
b_stack[g] = b_tensors[g].t()
b_stack = b_stack.transpose(1, 2)
a_scale_stack = torch.empty(
(expert_offsets[-1], k_g // 128), device=device, dtype=torch.float32
)
b_scale_stack = torch.empty(
(num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32
)
for g in range(num_groups):
a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[g]
b_scale_stack[g] = b_scales_tensors[g].t()
b_scale_stack = b_scale_stack.transpose(1, 2)
c_out = torch.empty((expert_offsets[-1], n_g), device=device, dtype=out_dtype)
a_strides = torch.full(
(num_groups,), a_stack.stride(0), device=device, dtype=torch.int64
)
c_strides = torch.full(
(num_groups,), c_out.stride(0), device=device, dtype=torch.int64
)
workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8)
a_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
b_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
out_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
a_scales_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
b_scales_ptrs = torch.empty((num_groups,), device=device, dtype=torch.int64)
def run_cutlass():
fp8_blockwise_scaled_grouped_mm(
c_out,
a_ptrs,
b_ptrs,
out_ptrs,
a_scales_ptrs,
b_scales_ptrs,
a_stack,
b_stack,
a_scale_stack,
b_scale_stack,
a_strides,
a_strides,
c_strides,
layout_sfa,
layout_sfb,
problem_sizes,
expert_offsets[:-1],
workspace,
)
# warmup
for _ in range(num_warmup):
run_cutlass()
torch.cuda.synchronize()
# run
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
for _ in range(num_run):
run_cutlass()
end_event.record()
end_event.synchronize()
torch.cuda.synchronize()
avg = start_event.elapsed_time(end_event) / num_run * 1000 # us
return avg, expert_offsets[-1]
def bench_sglang_triton(
expected_m_per_group: int,
n: int,
k: int,
num_groups: int,
num_warmup: int,
num_run: int,
) -> Tuple[float, int]:
pass
benchmark_kernels = {
"deepgemm": bench_deepgemm,
"cutlass": bench_cutlass,
# "triton": bench_sglang_triton,
}
@dataclass
class ShapeArg:
expected_m_per_group: int
n: int
k: int
num_groups: int
def benchmark_one_shape(
shape_args: List[ShapeArg],
num_warmup: int,
num_run: int,
):
for shape in shape_args:
print(
f"\nBenchmark: expected_m_per_group={shape.expected_m_per_group}, n={shape.n}, k={shape.k}, num_groups={shape.num_groups}"
)
for kernel_name, kernel_func in benchmark_kernels.items():
average_time, m = kernel_func(
shape.expected_m_per_group,
shape.n,
shape.k,
shape.num_groups,
num_warmup,
num_run,
)
print(f"{kernel_name}: {average_time} us")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--num-warmup", type=int, default=3)
parser.add_argument("--num-run", type=int, default=10)
# CI environment uses simplified parameters
if IS_CI:
shape_args = [
# Only test one simple shape in CI
ShapeArg(expected_m_per_group=128, n=512, k=7168, num_groups=256),
]
else:
shape_args = [
# Prefill, DeepSeek-R1, gateup, chunk_size = 4096, TP = 8
ShapeArg(expected_m_per_group=128, n=512, k=7168, num_groups=256),
# Prefill, DeepSeek-R1, gateup, chunk_size = 8192, TP = 8
ShapeArg(expected_m_per_group=256, n=512, k=7168, num_groups=256),
# Prefill, DeepSeek-R1, gateup, chunk_size = 8192, TP = 16
ShapeArg(expected_m_per_group=256, n=256, k=7168, num_groups=256),
# Prefill, DeepSeek-R1, gateup, chunk_size = 16384, TP = 16
ShapeArg(expected_m_per_group=512, n=256, k=7168, num_groups=256),
# Decode, DeepSeek-R1, gateup, bs = 32, TP = 8
ShapeArg(expected_m_per_group=1, n=512, k=7168, num_groups=256),
# Decode, DeepSeek-R1, gateup, bs = 64, TP = 16
ShapeArg(expected_m_per_group=2, n=256, k=7168, num_groups=256),
# Prefill, DeepSeek-R1, gateup, chunk_size = 8192, EP = 8
ShapeArg(expected_m_per_group=256, n=4096, k=7168, num_groups=32),
# Prefill, DeepSeek-R1, gateup, chunk_size = 16384, EP = 16
ShapeArg(expected_m_per_group=512, n=4096, k=7168, num_groups=16),
# Decode, DeepSeek-R1, gateup, bs = 128, EP = 8
ShapeArg(expected_m_per_group=4, n=4096, k=7168, num_groups=32),
# Decode, DeepSeek-R1, gateup, bs = 256, EP = 16
ShapeArg(expected_m_per_group=8, n=4096, k=7168, num_groups=16),
# Prefill, Qwen3-235B-A22B-FP8, gateup, chunk_size = 16384, TP = 4
ShapeArg(expected_m_per_group=1024, n=768, k=4096, num_groups=128),
# Prefill, Qwen3-235B-A22B-FP8, down, chunk_size = 16384, TP = 4
ShapeArg(expected_m_per_group=1024, n=4096, k=384, num_groups=128),
# Decode, Qwen3-235B-A22B-FP8, gateup, bs = 256, TP = 4
ShapeArg(expected_m_per_group=16, n=768, k=4096, num_groups=128),
# Decode, Qwen3-235B-A22B-FP8, down, bs = 256, TP = 4
ShapeArg(expected_m_per_group=16, n=4096, k=384, num_groups=128),
]
args = parser.parse_args()
benchmark_one_shape(shape_args, args.num_warmup, args.num_run)
if __name__ == "__main__":
main()
python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py
+232
View File
@@ -0,0 +1,232 @@
import argparse
import copy
import itertools
import os
from typing import Optional, Tuple
import torch
import triton
from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm
from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import (
per_tensor_quant_fp8,
)
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
from vllm._custom_ops import scaled_fp8_quant as vllm_scaled_fp8_quant
VLLM_AVAILABLE = True
except ImportError:
vllm_scaled_mm = None
vllm_scaled_fp8_quant = None
VLLM_AVAILABLE = False
IS_CI = is_in_ci()
# Weight Shapes are in the format
# ([K, N], TP_SPLIT_DIM)
# Example:
# A shape of ([14336, 4096], 0) indicates the following GEMM shape,
# - TP1 : K = 14336, N = 4096
# - TP2 : K = 7168, N = 4096
# A shape of ([4096, 6144], 1) indicates the following GEMM shape,
# - TP1 : K = 4096, N = 6144
# - TP4 : K = 4096, N = 1536
# TP1 shapes
WEIGHT_SHAPES = {
"meta-llama/Llama-3.1-8B-Instruct": [
([4096, 6144], 1),
([4096, 4096], 0),
([4096, 28672], 1),
([14336, 4096], 0),
],
"meta-llama/Llama-3.3-70B-Instruct": [
([8192, 10240], 1),
([8192, 8192], 0),
([8192, 57344], 1),
([28672, 8192], 0),
],
"mistralai/Mistral-Large-Instruct-2407": [
([12288, 14336], 1),
([12288, 12288], 0),
([12288, 57344], 1),
([28672, 12288], 0),
],
"Qwen/Qwen2.5-7B-Instruct": [
([3584, 4608], 1),
([3584, 3584], 0),
([3584, 37888], 1),
([18944, 3584], 0),
],
"Qwen/Qwen2.5-32B-Instruct": [
([5120, 7168], 1),
([5120, 5120], 0),
([5120, 55296], 1),
([27648, 5120], 0),
],
"Qwen/Qwen2.5-72B-Instruct": [
([8192, 10240], 1),
([8192, 8192], 0),
([8192, 59136], 1),
([29568, 8192], 0),
],
"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
([2048, 3072], 1),
([2048, 4096], 1),
([2048, 2048], 0),
([2048, 576], 0),
([2048, 21888], 1),
([10944, 2048], 0),
([2048, 2816], 1),
([1408, 2048], 0),
],
}
def sglang_scaled_fp8_quant(
input: torch.Tensor,
scale: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
fp8_type_: torch.dtype = torch.float8_e4m3fn
output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
is_static = True
if scale is None:
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
is_static = False
per_tensor_quant_fp8(input, output, scale, is_static)
return output, scale
# CI environment uses simplified parameters
if IS_CI:
batch_sizes = [1] # Single batch size for CI
else:
batch_sizes = [1, 16, 64, 128, 256, 512, 1024, 2048]
# Filter line_vals based on vLLM availability
if VLLM_AVAILABLE:
line_vals = [
"vllm-fp8-fp16",
"vllm-fp8-bf16",
"sglang-fp8-fp16",
"sglang-fp8-bf16",
]
line_names = [
"vllm-fp8-fp16",
"vllm-fp8-bf16",
"sglang-fp8-fp16",
"sglang-fp8-bf16",
]
styles = [("green", "-"), ("green", "--"), ("blue", "-"), ("blue", "--")]
else:
line_vals = [
"sglang-fp8-fp16",
"sglang-fp8-bf16",
]
line_names = [
"sglang-fp8-fp16",
"sglang-fp8-bf16",
]
styles = [("blue", "-"), ("blue", "--")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size"],
x_vals=batch_sizes,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="GB/s",
plot_name="fp8 scaled matmul",
args={},
)
)
def benchmark(batch_size, provider, N, K):
# M, N, K = batch_size, 4096, 8192
M = batch_size
a = torch.ones((M, K), device="cuda") * 5.0
b = torch.ones((N, K), device="cuda") * 5.0
# vLLM expects scalar scales, while sglang can handle per-token scales
scale_a_scalar = torch.randn(1, device="cuda", dtype=torch.float32)
scale_b_scalar = torch.randn(1, device="cuda", dtype=torch.float32)
scale_a = torch.randn((M,), device="cuda", dtype=torch.float32)
scale_b = torch.randn((N,), device="cuda", dtype=torch.float32)
quantiles = [0.5, 0.2, 0.8]
dtype = torch.float16 if "fp16" in provider else torch.bfloat16
if "vllm-fp8" in provider:
if not VLLM_AVAILABLE:
# Return zero if vLLM is not available
return (0, 0, 0)
a_fp8, scale_a_fp8 = vllm_scaled_fp8_quant(a, scale_a_scalar)
b_fp8, scale_b_fp8 = vllm_scaled_fp8_quant(b, scale_b_scalar)
b_fp8 = b_fp8.t()
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype),
quantiles=quantiles,
)
elif "sglang-fp8" in provider:
a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, scale_a)
b_fp8, scale_b_fp8 = sglang_scaled_fp8_quant(b, scale_b)
b_fp8 = b_fp8.t()
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: sgl_scaled_mm(
a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype, bias=None
),
quantiles=quantiles,
)
gbps = lambda ms: (2 * M * N * K + M * N) * a.element_size() * 1e-9 / (ms * 1e-3)
return gbps(ms), gbps(max_ms), gbps(min_ms)
def prepare_shapes(args):
KN_model_names = []
models_tps = list(itertools.product(args.models, args.tp_sizes))
for model, tp_size in models_tps:
assert model in WEIGHT_SHAPES
for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
KN[tp_split_dim] = KN[tp_split_dim] // tp_size
KN.append(model)
KN_model_names.append(KN)
return KN_model_names
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--models",
nargs="+",
type=str,
default=["meta-llama/Llama-3.1-8B-Instruct"],
help="List of models to benchmark",
)
parser.add_argument(
"--tp-sizes",
nargs="+",
type=int,
default=[1],
help="List of tensor parallel sizes",
)
args = parser.parse_args()
# Simplify for CI environment
if IS_CI:
args.models = [args.models[0]] # Use only first model
args.tp_sizes = [args.tp_sizes[0]] # Use only first TP size
KN_model_names = prepare_shapes(args)
for K, N, model_name in KN_model_names:
print(f"{model_name} N={N} K={K}: ")
benchmark.run(print_data=True, N=N, K=K)
print("Benchmark finished!")
python/sglang/kernels/aot/benchmark/bench_fp8_gemm_swap_ab.py
+132
View File
@@ -0,0 +1,132 @@
"""Targeted benchmark for the SM90 FP8 swap-AB dispatch path.
Sweeps small batch sizes (M = 1..128) across N/K shapes that exercise each
dispatch bucket in `fp8_gemm_sm90_dispatch.cuh`. Output style matches
`bench_fp8_gemm.py`: `triton.testing.perf_report` + GB/s table per (N, K).
Compare against `main` by:
1. Run on `main`: `python bench_fp8_gemm_swap_ab.py > main.txt`
2. Run on feature branch:`python bench_fp8_gemm_swap_ab.py > swap_ab.txt`
3. Diff the two tables.
"""
import argparse
import os
from typing import Optional, Tuple
import torch
import triton
from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm
from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import (
per_tensor_quant_fp8,
)
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
# (N, K) shapes targeting each dispatch bucket boundary.
# Spans M16_smallN / M16_largeN / M32_largeN / M64_smallN / M64_largeN /
# M128_smallN / M128_largeN dispatch entries when crossed with batch sizes
# below.
NK_SHAPES = [
(1024, 4096),
(1024, 8192),
(1280, 4096), # n == kNThreshold boundary
(4096, 4096),
(4096, 8192), # n == kM128NThreshold boundary for M128 bucket
(8192, 4096),
(8192, 8192),
(14336, 4096),
(14336, 8192),
(28672, 4096), # Llama-3 70B MLP up_proj N
(28672, 8192),
]
# Batch sizes covering each M-bucket of the swap-AB dispatch.
# CI runs only M=1 to stay fast; full run probes the bucket transitions.
if IS_CI:
batch_sizes = [1]
else:
batch_sizes = [1, 8, 16, 17, 32, 48, 64, 96, 128]
line_vals = ["sglang-fp8-bf16", "sglang-fp8-fp16"]
line_names = line_vals
styles = [("blue", "-"), ("blue", "--")]
def sglang_scaled_fp8_quant(
input: torch.Tensor,
scale: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
fp8_type_ = torch.float8_e4m3fn
output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
is_static = True
if scale is None:
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
is_static = False
per_tensor_quant_fp8(input, output, scale, is_static)
return output, scale
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size"],
x_vals=batch_sizes,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="GB/s",
plot_name="fp8 swap-AB scaled matmul",
args={},
)
)
def benchmark(batch_size, provider, N, K):
M = batch_size
a = torch.ones((M, K), device="cuda") * 5.0
b = torch.ones((N, K), device="cuda") * 5.0
scale_a = torch.randn((M,), device="cuda", dtype=torch.float32)
scale_b = torch.randn((N,), device="cuda", dtype=torch.float32)
quantiles = [0.5, 0.2, 0.8]
dtype = torch.float16 if "fp16" in provider else torch.bfloat16
a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, scale_a)
b_fp8, scale_b_fp8 = sglang_scaled_fp8_quant(b, scale_b)
b_fp8 = b_fp8.t()
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: sgl_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype, bias=None),
quantiles=quantiles,
)
gbps = lambda ms: (2 * M * N * K + M * N) * a.element_size() * 1e-9 / (ms * 1e-3)
return gbps(ms), gbps(max_ms), gbps(min_ms)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--save-path",
type=str,
default=None,
help="Directory to save plots/CSVs (default: don't save)",
)
args = parser.parse_args()
if IS_CI:
# CI: probe a single (N, K) to stay quick.
N, K = NK_SHAPES[0]
print(f"N={N} K={K}: ")
benchmark.run(print_data=True, N=N, K=K)
else:
for N, K in NK_SHAPES:
print(f"N={N} K={K}: ")
kwargs = {"print_data": True, "N": N, "K": K}
if args.save_path:
os.makedirs(args.save_path, exist_ok=True)
kwargs["save_path"] = args.save_path
benchmark.run(**kwargs)
print("Benchmark finished!")
python/sglang/kernels/aot/benchmark/bench_int8_gemm.py
+183
View File
@@ -0,0 +1,183 @@
import argparse
import copy
import itertools
import os
import torch
import triton
from sgl_kernel import int8_scaled_mm
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
VLLM_AVAILABLE = True
except ImportError:
vllm_scaled_mm = None
VLLM_AVAILABLE = False
IS_CI = is_in_ci()
def to_int8(tensor: torch.Tensor) -> torch.Tensor:
return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)
WEIGHT_SHAPES = {
"meta-llama/Llama-3.1-8B-Instruct": [
([4096, 6144], 1),
([4096, 4096], 0),
([4096, 28672], 1),
([14336, 4096], 0),
],
"meta-llama/Llama-3.3-70B-Instruct": [
([8192, 10240], 1),
([8192, 8192], 0),
([8192, 57344], 1),
([28672, 8192], 0),
],
"mistralai/Mistral-Large-Instruct-2407": [
([12288, 14336], 1),
([12288, 12288], 0),
([12288, 57344], 1),
([28672, 12288], 0),
],
"Qwen/Qwen2.5-7B-Instruct": [
([3584, 4608], 1),
([3584, 3584], 0),
([3584, 37888], 1),
([18944, 3584], 0),
],
"Qwen/Qwen2.5-32B-Instruct": [
([5120, 7168], 1),
([5120, 5120], 0),
([5120, 55296], 1),
([27648, 5120], 0),
],
"Qwen/Qwen2.5-72B-Instruct": [
([8192, 10240], 1),
([8192, 8192], 0),
([8192, 59136], 1),
([29568, 8192], 0),
],
"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
([2048, 3072], 1),
([2048, 4096], 1),
([2048, 2048], 0),
([2048, 576], 0),
([2048, 21888], 1),
([10944, 2048], 0),
([2048, 2816], 1),
([1408, 2048], 0),
],
}
# CI environment uses simplified parameters
if IS_CI:
batch_sizes = [1] # Single batch size for CI
else:
batch_sizes = [1, 16, 32, 64, 128, 256, 512, 1024, 2048]
# Filter providers based on vLLM availability
if VLLM_AVAILABLE:
line_vals = ["vllm", "sgl-kernel"]
line_names = ["vllm int8 gemm", "sgl-kernel int8 gemm"]
styles = [("blue", "-"), ("orange", "-")]
else:
line_vals = ["sgl-kernel"]
line_names = ["sgl-kernel int8 gemm"]
styles = [("orange", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size"],
x_vals=batch_sizes,
x_log=False,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="GB/s",
plot_name="int8 scaled matmul",
args={},
)
)
def benchmark(batch_size, provider, N, K):
M = batch_size
a = to_int8(torch.randn((M, K), device="cuda") * 5)
b = to_int8(torch.randn((N, K), device="cuda").t() * 5)
scale_a = torch.randn((M,), device="cuda", dtype=torch.float32)
scale_b = torch.randn((N,), device="cuda", dtype=torch.float32)
bias = torch.randn((N,), device="cuda", dtype=torch.float16)
quantiles = [0.5, 0.2, 0.8]
if provider == "sgl-kernel":
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: int8_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias),
quantiles=quantiles,
)
elif provider == "vllm":
if not VLLM_AVAILABLE:
return (0, 0, 0)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: vllm_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias),
quantiles=quantiles,
)
gbps = (
lambda ms: (
(2 * M * N * K - M * N) * a.element_size()
+ (3 * M * N) * scale_a.element_size()
)
* 1e-9
/ (ms * 1e-3)
)
return gbps(ms), gbps(max_ms), gbps(min_ms)
def prepare_shapes(args):
KN_model_names = []
models_tps = list(itertools.product(args.models, args.tp_sizes))
for model, tp_size in models_tps:
assert model in WEIGHT_SHAPES
for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
KN[tp_split_dim] = KN[tp_split_dim] // tp_size
KN.append(model)
KN_model_names.append(KN)
return KN_model_names
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--models",
nargs="+",
type=str,
default=["meta-llama/Llama-3.1-8B-Instruct"],
help="List of models to benchmark",
)
parser.add_argument(
"--tp-sizes",
nargs="+",
type=int,
default=[1],
help="List of tensor parallel sizes",
)
args = parser.parse_args()
# Skip in CI environment due to architecture compatibility issues
if IS_CI:
print(
"Skipping INT8 GEMM benchmark in CI environment due to architecture compatibility issues"
)
print("INT8 operations may not be supported on all GPU architectures")
else:
KN_model_names = prepare_shapes(args)
for K, N, model_name in KN_model_names:
print(f"{model_name} N={N} K={K}: ")
benchmark.run(print_data=True, N=N, K=K)
print("Benchmark finished!")
python/sglang/kernels/aot/benchmark/bench_moe_align_block_size.py
+443
View File
@@ -0,0 +1,443 @@
import argparse
import itertools
import os
import torch
import triton
import triton.language as tl
from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size
from sglang.utils import is_in_ci
try:
from vllm import _custom_ops as ops
VLLM_AVAILABLE = True
except ImportError:
ops = None
VLLM_AVAILABLE = False
IS_CI = is_in_ci()
USE_RANDOM_PERM = False
def ceil_div(a, b):
return (a + b - 1) // b
@triton.jit
def moe_align_block_size_stage1(
topk_ids_ptr,
tokens_cnts_ptr,
num_experts: tl.constexpr,
numel: tl.constexpr,
tokens_per_thread: tl.constexpr,
):
pid = tl.program_id(0)
start_idx = pid * tokens_per_thread
off_c = (pid + 1) * num_experts
for i in range(tokens_per_thread):
if start_idx + i < numel:
idx = tl.load(topk_ids_ptr + start_idx + i)
token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)
@triton.jit
def moe_align_block_size_stage2(
tokens_cnts_ptr,
num_experts: tl.constexpr,
):
pid = tl.program_id(0)
last_cnt = 0
for i in range(1, num_experts + 1):
token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
last_cnt = last_cnt + token_cnt
tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)
@triton.jit
def moe_align_block_size_stage3(
total_tokens_post_pad_ptr,
tokens_cnts_ptr,
cumsum_ptr,
num_experts: tl.constexpr,
block_size: tl.constexpr,
):
last_cumsum = 0
off_cnt = num_experts * num_experts
for i in range(1, num_experts + 1):
token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
tl.store(cumsum_ptr + i, last_cumsum)
tl.store(total_tokens_post_pad_ptr, last_cumsum)
@triton.jit
def moe_align_block_size_stage4(
topk_ids_ptr,
sorted_token_ids_ptr,
expert_ids_ptr,
tokens_cnts_ptr,
cumsum_ptr,
num_experts: tl.constexpr,
block_size: tl.constexpr,
numel: tl.constexpr,
tokens_per_thread: tl.constexpr,
):
pid = tl.program_id(0)
start_idx = tl.load(cumsum_ptr + pid)
end_idx = tl.load(cumsum_ptr + pid + 1)
for i in range(start_idx, end_idx, block_size):
tl.store(expert_ids_ptr + i // block_size, pid)
start_idx = pid * tokens_per_thread
off_t = pid * num_experts
for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread, numel)):
expert_id = tl.load(topk_ids_ptr + i)
token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
tl.store(sorted_token_ids_ptr + rank_post_pad, i)
tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)
def moe_align_block_size_triton(
topk_ids: torch.Tensor,
num_experts: int,
block_size: int,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_post_pad: torch.Tensor,
) -> None:
numel = topk_ids.numel()
grid = (num_experts,)
tokens_cnts = torch.zeros(
(num_experts + 1, num_experts), dtype=torch.int32, device=topk_ids.device
)
cumsum = torch.zeros((num_experts + 1,), dtype=torch.int32, device=topk_ids.device)
tokens_per_thread = ceil_div(numel, num_experts)
moe_align_block_size_stage1[grid](
topk_ids,
tokens_cnts,
num_experts,
numel,
tokens_per_thread,
)
moe_align_block_size_stage2[grid](
tokens_cnts,
num_experts,
)
moe_align_block_size_stage3[(1,)](
num_tokens_post_pad,
tokens_cnts,
cumsum,
num_experts,
block_size,
)
moe_align_block_size_stage4[grid](
topk_ids,
sorted_token_ids,
expert_ids,
tokens_cnts,
cumsum,
num_experts,
block_size,
numel,
tokens_per_thread,
)
def calculate_diff(num_tokens, num_experts=256, block_size=128, topk=8):
topk_ids = torch.stack(
[
torch.randperm(num_experts, dtype=torch.int32, device="cuda")[:topk]
for _ in range(num_tokens)
]
)
# SGL kernel uses dynamic padding optimization
max_num_tokens_padded_sgl = topk_ids.numel() + num_experts * (block_size - 1)
if topk_ids.numel() < num_experts + 1:
max_num_tokens_padded_sgl = topk_ids.numel() * block_size
sorted_ids_cuda = torch.empty(
(max_num_tokens_padded_sgl,), dtype=torch.int32, device=topk_ids.device
)
sorted_ids_cuda.fill_(topk_ids.numel())
max_num_m_blocks_sgl = max_num_tokens_padded_sgl // block_size
expert_ids_cuda = torch.zeros(
(max_num_m_blocks_sgl,), dtype=torch.int32, device=topk_ids.device
)
num_tokens_post_pad_cuda = torch.empty(
(1), dtype=torch.int32, device=topk_ids.device
)
cumsum_buffer = torch.zeros(
num_experts + 1, dtype=torch.int32, device=topk_ids.device
)
# Triton and vLLM use original padding calculation
max_num_tokens_padded_triton = topk_ids.numel() + num_experts * (block_size - 1)
max_num_m_blocks_triton = max_num_tokens_padded_triton // block_size
sorted_ids_triton = torch.empty(
(max_num_tokens_padded_triton,), dtype=torch.int32, device=topk_ids.device
)
sorted_ids_triton.fill_(topk_ids.numel())
expert_ids_triton = torch.zeros(
(max_num_m_blocks_triton,), dtype=torch.int32, device=topk_ids.device
)
num_tokens_post_pad_triton = torch.empty_like(num_tokens_post_pad_cuda)
sorted_ids_vllm = torch.empty_like(sorted_ids_triton)
sorted_ids_vllm.fill_(topk_ids.numel())
expert_ids_vllm = torch.zeros_like(expert_ids_triton)
num_tokens_post_pad_vllm = torch.empty_like(num_tokens_post_pad_cuda)
# compare the performance of cuda, triton and vllm implementation
sgl_moe_align_block_size(
topk_ids,
num_experts,
block_size,
sorted_ids_cuda,
expert_ids_cuda,
num_tokens_post_pad_cuda,
cumsum_buffer,
)
moe_align_block_size_triton(
topk_ids,
num_experts,
block_size,
sorted_ids_triton,
expert_ids_triton,
num_tokens_post_pad_triton,
)
if VLLM_AVAILABLE:
try:
ops.moe_align_block_size(
topk_ids,
num_experts,
block_size,
sorted_ids_vllm,
expert_ids_vllm,
num_tokens_post_pad_vllm,
)
print(f"✅ VLLM implementation works with {num_experts} experts!")
vllm_works = True
except Exception as e:
print(f"❌ VLLM implementation failed with {num_experts} experts: {e}")
vllm_works = False
else:
print("⚠️ vLLM not available, skipping vLLM test")
vllm_works = False
if torch.allclose(expert_ids_cuda, expert_ids_triton) and torch.allclose(
num_tokens_post_pad_cuda, num_tokens_post_pad_triton
):
print("✅ SGL and Triton implementations match")
else:
print("❌ SGL and Triton implementations do not match")
print("SGL expert_ids:", expert_ids_cuda)
print("Triton expert_ids:", expert_ids_triton)
print("SGL num_tokens_post_pad:", num_tokens_post_pad_cuda)
print("Triton num_tokens_post_pad:", num_tokens_post_pad_triton)
if (
vllm_works
and torch.allclose(expert_ids_cuda, expert_ids_vllm)
and torch.allclose(num_tokens_post_pad_cuda, num_tokens_post_pad_vllm)
):
print("✅ SGL and VLLM implementations match")
else:
if not vllm_works:
print("⚠️ VLLM comparison skipped due to failure")
else:
print("❌ SGL and VLLM implementations do not match")
print("SGL expert_ids:", expert_ids_cuda)
print("VLLM expert_ids:", expert_ids_vllm)
print("SGL num_tokens_post_pad:", num_tokens_post_pad_cuda)
print("VLLM num_tokens_post_pad:", num_tokens_post_pad_vllm)
# Test range
num_tokens_range = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
num_experts_range = [8, 32, 64, 128, 256]
topk_range = [1, 2, 4, 8]
configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor:
topk_ids = torch.zeros((num_tokens, topk), dtype=torch.int32, device="cuda")
for i in range(num_tokens):
topk_ids[i, :] = torch.randperm(num_experts, dtype=torch.int32, device="cuda")[
:topk
]
return topk_ids
def sgl_moe_align_block_size_with_empty(
topk_ids,
num_experts,
block_size,
sorted_ids,
expert_ids,
num_tokens_post_pad,
pad_sorted_token_ids=False,
):
if not pad_sorted_token_ids:
sorted_ids.fill_(topk_ids.numel())
cumsum_buffer = torch.empty(
num_experts + 1, dtype=torch.int32, device=topk_ids.device
)
sgl_moe_align_block_size(
topk_ids,
num_experts,
block_size,
sorted_ids.clone(),
expert_ids.clone(),
num_tokens_post_pad.clone(),
cumsum_buffer,
pad_sorted_token_ids,
)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens", "num_experts", "topk"],
x_vals=configs,
line_arg="provider",
line_vals=["sgl", "sgl_fusion", "triton"],
line_names=["SGL", "SGL Fusion", "Triton"],
styles=[("blue", "-"), ("red", "-"), ("green", "-")],
ylabel="us",
plot_name="moe-align-block-size-performance",
args={},
)
)
def benchmark(num_tokens, num_experts, topk, provider):
block_size = 128
if USE_RANDOM_PERM:
topk_ids = get_topk_ids(num_tokens, num_experts, topk)
else:
topk_ids = torch.randint(
0,
num_experts,
(num_tokens, topk),
dtype=torch.int32,
device="cuda",
)
# Calculate max_num_tokens_padded based on provider
if provider == "sgl" or provider == "sgl_fusion":
# Apply dynamic padding optimization for SGL kernel
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
if topk_ids.numel() < num_experts:
max_num_tokens_padded = topk_ids.numel() * block_size
else: # triton
# Use original padding calculation for Triton
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
# Create tensors
sorted_ids = torch.empty(
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
)
max_num_m_blocks = max_num_tokens_padded // block_size
expert_ids = torch.empty(
(max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
)
num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)
quantiles = [0.5, 0.2, 0.8]
if provider == "sgl":
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: sgl_moe_align_block_size_with_empty(
topk_ids,
num_experts,
block_size,
sorted_ids,
expert_ids,
num_tokens_post_pad,
),
quantiles=quantiles,
)
elif provider == "sgl_fusion":
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: sgl_moe_align_block_size_with_empty(
topk_ids,
num_experts,
block_size,
sorted_ids,
expert_ids,
num_tokens_post_pad,
pad_sorted_token_ids=True,
),
quantiles=quantiles,
)
elif provider == "triton":
sorted_ids.fill_(topk_ids.numel())
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
lambda: moe_align_block_size_triton(
topk_ids,
num_experts,
block_size,
sorted_ids.clone(),
expert_ids.clone(),
num_tokens_post_pad.clone(),
),
quantiles=quantiles,
)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--save_path",
type=str,
default="./configs/benchmark_ops/moe_align_blocks/",
help="Path to save moe align benchmark results",
)
parser.add_argument(
"--num_experts",
type=int,
default=256,
choices=[8, 16, 32, 64, 128, 256],
help="Number of experts for benchmark",
)
parser.add_argument(
"--topk",
type=int,
default=8,
choices=[2, 4, 8],
help="Top-k value for benchmark",
)
parser.add_argument(
"--skip_full_benchmark",
action="store_true",
help="Only run the calculate_diff function, skip full benchmarking",
)
args = parser.parse_args()
# Simplify for CI environment
if IS_CI:
num_tokens = 256 # Smaller for CI
num_experts = 8 # Smaller for CI
topk = 2 # Smaller for CI
else:
num_tokens = 1024
num_experts = args.num_experts
topk = args.topk
calculate_diff(num_tokens=num_tokens, num_experts=num_experts, topk=topk)
if not args.skip_full_benchmark and not IS_CI: # Skip full benchmark in CI
print(f"\n📊 Running performance benchmark for {args.num_experts} experts...")
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_moe_ep_post_reorder.py
+85
View File
@@ -0,0 +1,85 @@
import torch
import triton
from sglang.kernels.ops.moe.ep_moe_kernels import post_reorder_triton_kernel
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
# CI environment uses simplified parameters
if IS_CI:
batch_sizes = [64, 128] # Only test 2 values in CI
else:
batch_sizes = [64, 128, 256, 512, 640, 768, 1024, 2048, 4096]
configs = [(bs,) for bs in batch_sizes]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size"],
x_vals=[list(_) for _ in configs],
line_arg="provider",
line_vals=["triton"],
line_names=["Triton Kernel"],
styles=[("orange", "-")],
ylabel="us",
plot_name="ep-moe-post-reorder-performance",
args={},
)
)
def benchmark(batch_size, provider):
dtype = torch.bfloat16
device = torch.device("cuda")
hidden_size, topk, start_expert_id, end_expert_id, block_size = 4096, 8, 0, 255, 512
def alloc_tensors():
down_output = torch.randn(
batch_size * topk, hidden_size, dtype=dtype, device=device
)
output = torch.zeros(batch_size, hidden_size, dtype=dtype, device=device)
src2dst = torch.randint(
0, batch_size * topk, (batch_size, topk), dtype=torch.int32, device=device
)
topk_ids = torch.randint(
start_expert_id,
end_expert_id + 1,
(batch_size, topk),
dtype=torch.int32,
device=device,
)
topk_weights = torch.rand(batch_size, topk, dtype=dtype, device=device)
return down_output, output, src2dst, topk_ids, topk_weights
quantiles = [0.5, 0.2, 0.8]
if provider == "triton":
d_out, out, s2d, tk_ids, tk_weights = alloc_tensors()
def run_triton():
post_reorder_triton_kernel[(batch_size,)](
d_out.view(-1),
out.view(-1),
s2d.view(-1),
tk_ids.view(-1),
tk_weights.view(-1),
start_expert_id,
end_expert_id,
topk,
hidden_size,
0,
block_size,
)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
run_triton, quantiles=quantiles
)
else:
raise ValueError(f"Unknown provider: {provider}")
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_moe_topk_sigmoid.py
+242
View File
@@ -0,0 +1,242 @@
import itertools
import os
import pytest
import torch
import triton
from sgl_kernel import topk_sigmoid
from sglang.utils import is_in_ci
# Optional MUSA import
try:
from sglang.srt.utils import is_musa
if is_musa():
from sglang.srt.hardware_backend.musa.kernels.topk import (
topk_sigmoid as musa_topk_sigmoid,
)
MUSA_AVAILABLE = True
else:
musa_topk_sigmoid = None
MUSA_AVAILABLE = False
except ImportError:
musa_topk_sigmoid = None
MUSA_AVAILABLE = False
IS_CI = is_in_ci()
def torch_topk_sigmoid_native(
gating_output: torch.Tensor,
topk: int,
renormalize: bool,
correction_bias: torch.Tensor = None,
):
scores = gating_output.sigmoid()
if correction_bias is not None:
n_routed_experts = gating_output.shape[-1]
scores_for_choice = scores.view(
-1, n_routed_experts
) + correction_bias.unsqueeze(0)
_, topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1)
topk_weights = scores.gather(1, topk_indices)
else:
topk_weights, topk_indices = torch.topk(scores, k=topk, dim=-1)
if renormalize:
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
return topk_weights, topk_indices
def sglang_topk_sigmoid(
gating_output: torch.Tensor,
topk: int,
renormalize: bool,
correction_bias: torch.Tensor = None,
):
num_tokens, num_experts = gating_output.shape
topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
topk_sigmoid(
topk_weights,
topk_indices,
gating_output,
renormalize=renormalize,
correction_bias=correction_bias,
)
return topk_weights, topk_indices
def musa_topk_sigmoid_fn(
gating_output: torch.Tensor,
topk: int,
renormalize: bool,
correction_bias: torch.Tensor = None,
):
num_tokens, num_experts = gating_output.shape
topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
musa_topk_sigmoid(
topk_weights,
topk_indices,
gating_output,
renormalize=renormalize,
correction_bias=correction_bias,
)
return topk_weights, topk_indices
def get_topk_sigmoid_input(num_tokens, num_experts):
gating_output = torch.randn(
(num_tokens, num_experts), dtype=torch.float32, device="cuda"
)
correction_bias = torch.randn((num_experts), dtype=torch.float32, device="cuda")
return gating_output, correction_bias
def calculate_diff(num_tokens, num_experts, topk):
gating_output, correction_bias = get_topk_sigmoid_input(num_tokens, num_experts)
weights_torch, indices_torch = torch_topk_sigmoid_native(
gating_output.clone(),
topk,
True,
correction_bias.clone(),
)
weights_sglang, indices_sglang = sglang_topk_sigmoid(
gating_output.clone(),
topk,
True,
correction_bias.clone(),
)
weights_diff = torch.abs(weights_torch - weights_sglang).mean().item()
indices_match = torch.equal(indices_torch, indices_sglang)
if (
torch.allclose(weights_torch, weights_sglang, atol=1e-3, rtol=1e-3)
and indices_match
):
print("✅ Torch and SGLang topk_sigmoid implementations match")
else:
print(
f"❌ Implementations differ: Weights diff={weights_diff}, Indices match={indices_match}"
)
if MUSA_AVAILABLE:
weights_musa, indices_musa = musa_topk_sigmoid_fn(
gating_output.clone(),
topk,
True,
correction_bias.clone(),
)
weights_diff_musa = torch.abs(weights_sglang - weights_musa).mean().item()
indices_match_musa = torch.equal(indices_sglang, indices_musa)
if (
torch.allclose(weights_sglang, weights_musa, atol=1e-3, rtol=1e-3)
and indices_match_musa
):
print("✅ SGLang and MUSA topk_sigmoid implementations match")
else:
print(
f"❌ MUSA vs SGLang differ: Weights diff={weights_diff_musa}, Indices match={indices_match_musa}"
)
else:
print("⚠️ MUSA not available, skipping MUSA comparison")
# CI environment uses simplified parameters
if IS_CI:
num_tokens_range = [128] # Single value for CI
num_experts_range = [32] # Single value for CI
topk_range = [2] # Single value for CI
else:
num_tokens_range = [128, 512, 1024, 2048, 4096, 8192, 16384, 32768]
num_experts_range = [32, 64, 128, 256, 12, 512]
topk_range = [1, 2, 4, 8]
configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
# Filter providers based on availability
line_vals = ["sglang", "torch"]
line_names = ["SGLang", "Torch"]
styles = [("blue", "-"), ("green", "-")]
if MUSA_AVAILABLE:
line_vals.append("musa")
line_names.append("MUSA")
styles.append(("red", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens", "num_experts", "topk"],
x_vals=configs,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="Latency (us)",
plot_name="topk-sigmoid-performance",
args={},
)
)
def benchmark(num_tokens, num_experts, topk, provider):
gating_output, correction_bias = get_topk_sigmoid_input(num_tokens, num_experts)
if provider == "torch" or provider == "torch1":
def fn():
return torch_topk_sigmoid_native(
gating_output,
topk,
True,
correction_bias,
)
elif provider == "sglang" or provider == "sglang1":
def fn():
return sglang_topk_sigmoid(gating_output, topk, True, correction_bias)
elif provider == "musa" or provider == "musa1":
if not MUSA_AVAILABLE:
return (0, 0, 0)
def fn():
return musa_topk_sigmoid_fn(gating_output, topk, True, correction_bias)
quantiles = [0.5, 0.2, 0.8]
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
# Simplify configs for CI environment
if IS_CI:
test_configs = [(20, 32, 2)] # Single config for CI
else:
test_configs = [
(20, 256, 4),
(20, 256, 8),
(20, 12, 4),
(20, 12, 1),
(20, 512, 4),
(20, 512, 1),
]
for num_tokens, num_experts, topk in test_configs:
calculate_diff(num_tokens, num_experts, topk)
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_moe_topk_softmax.py
+220
View File
@@ -0,0 +1,220 @@
import itertools
import os
import pytest
import torch
import triton
from sgl_kernel import topk_softmax
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm import _custom_ops as vllm_custom_ops
VLLM_AVAILABLE = True
except ImportError:
vllm_custom_ops = None
VLLM_AVAILABLE = False
# Optional MUSA import
try:
from sglang.srt.utils import is_musa
if is_musa():
from sglang.srt.hardware_backend.musa.kernels.topk import (
topk_softmax as musa_topk_softmax,
)
MUSA_AVAILABLE = True
else:
musa_topk_softmax = None
MUSA_AVAILABLE = False
except ImportError:
musa_topk_softmax = None
MUSA_AVAILABLE = False
IS_CI = is_in_ci()
def vllm_topk_softmax(gating_output, topk):
if not VLLM_AVAILABLE:
# Fallback to SGLang implementation if vLLM is not available
return sglang_topk_softmax(gating_output, topk)
num_tokens, num_experts = gating_output.shape
topk_weights = torch.empty(
(num_tokens, topk), device=gating_output.device, dtype=torch.float32
)
topk_indices = torch.empty(
(num_tokens, topk), dtype=torch.int32, device=gating_output.device
)
token_expert_indices = torch.empty(
(num_tokens, topk), dtype=torch.int32, device=gating_output.device
)
torch.ops._moe_C.topk_softmax(
topk_weights, topk_indices, token_expert_indices, gating_output
)
return topk_weights, topk_indices
def sglang_topk_softmax(gating_output, topk):
num_tokens, num_experts = gating_output.shape
topk_weights = torch.empty(
(num_tokens, topk), device=gating_output.device, dtype=torch.float32
)
topk_indices = torch.empty(
(num_tokens, topk), dtype=torch.int32, device=gating_output.device
)
topk_softmax(
topk_weights=topk_weights,
topk_ids=topk_indices,
gating_output=gating_output,
)
return topk_weights, topk_indices
def musa_topk_softmax_fn(gating_output, topk):
num_tokens, num_experts = gating_output.shape
topk_weights = torch.empty(
(num_tokens, topk), device=gating_output.device, dtype=torch.float32
)
topk_indices = torch.empty(
(num_tokens, topk), dtype=torch.int32, device=gating_output.device
)
musa_topk_softmax(
topk_weights,
topk_indices,
gating_output,
)
return topk_weights, topk_indices
def calculate_diff(num_tokens, num_experts, topk):
gating_output = torch.randn(
(num_tokens, num_experts), device="cuda", dtype=torch.float32
)
weights_sglang, indices_sglang = sglang_topk_softmax(gating_output.clone(), topk)
if MUSA_AVAILABLE:
weights_musa, indices_musa = musa_topk_softmax_fn(gating_output.clone(), topk)
weights_diff = torch.abs(weights_sglang - weights_musa).mean().item()
indices_match = torch.equal(indices_sglang, indices_musa)
if (
torch.allclose(weights_sglang, weights_musa, atol=1e-3, rtol=1e-3)
and indices_match
):
print("✅ SGLang and MUSA topk_softmax implementations match")
else:
print(
f"❌ Implementations differ: Weights diff={weights_diff}, Indices match={indices_match}"
)
else:
print("⚠️ MUSA not available, skipping MUSA comparison")
if VLLM_AVAILABLE:
weights_vllm, indices_vllm = vllm_topk_softmax(gating_output.clone(), topk)
weights_diff_vllm = torch.abs(weights_vllm - weights_sglang).mean().item()
indices_match_vllm = torch.equal(indices_vllm, indices_sglang)
if (
torch.allclose(weights_vllm, weights_sglang, atol=1e-3, rtol=1e-3)
and indices_match_vllm
):
print("✅ VLLM and SGLang topk_softmax implementations match")
else:
print(
f"❌ VLLM vs SGLang differ: Weights diff={weights_diff_vllm}, Indices match={indices_match_vllm}"
)
# CI environment uses simplified parameters
if IS_CI:
num_tokens_range = [128] # Single value for CI
num_experts_range = [32] # Single value for CI
topk_range = [2] # Single value for CI
else:
num_tokens_range = [128, 512, 1024, 2048, 4096, 8192, 16384, 32768]
num_experts_range = [32, 64, 128, 256, 12, 512]
topk_range = [1, 2, 4, 8, 10]
configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
# Filter providers based on availability
line_vals = ["sglang"]
line_names = ["SGLang"]
styles = [("blue", "-")]
if VLLM_AVAILABLE:
line_vals.append("vllm")
line_names.append("VLLM")
styles.append(("green", "-"))
if MUSA_AVAILABLE:
line_vals.append("musa")
line_names.append("MUSA")
styles.append(("red", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens", "num_experts", "topk"],
x_vals=configs,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="Latency (us)",
plot_name="topk-softmax-performance",
args={},
)
)
def benchmark(num_tokens, num_experts, topk, provider):
gating_output = torch.randn(
(num_tokens, num_experts), device="cuda", dtype=torch.float32
)
if provider == "vllm" or provider == "vllm1":
if not VLLM_AVAILABLE:
return (0, 0, 0)
fn = lambda: vllm_topk_softmax(gating_output, topk)
elif provider == "sglang" or provider == "sglang1":
fn = lambda: sglang_topk_softmax(gating_output, topk)
elif provider == "musa" or provider == "musa1":
if not MUSA_AVAILABLE:
return (0, 0, 0)
fn = lambda: musa_topk_softmax_fn(gating_output, topk)
quantiles = [0.5, 0.2, 0.8]
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
# Simplify configs for CI environment
if IS_CI:
test_configs = [(20, 32, 2)] # Single config for CI
else:
test_configs = [
(20, 256, 4),
(20, 256, 8),
(20, 12, 4),
(20, 12, 1),
(20, 512, 4),
(20, 512, 1),
]
for num_tokens, num_experts, topk in test_configs:
calculate_diff(num_tokens, num_experts, topk)
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_mrope.py
+250
View File
@@ -0,0 +1,250 @@
# Adapted from vLLM benchmark_mrope.py
# This script benchmarks the mrope kernel (mainly for Qwen2VL and Qwen2.5VL models).
# It generates test data, runs benchmarks, and saves results to a CSV file.
#
# The CSV file (named with current date/time) contains these columns:
# model_name, tp_size, num_tokens, num_heads, num_kv_heads, head_dim, max_position,
# rope_theta, is_neox_style, rope_scaling, dtype, torch_mean, torch_median, torch_p99,
# torch_min, torch_max, triton_mean, triton_median, triton_p99, triton_min, triton_max,
# speedup
#
# == Usage Examples ==
#
# Single model benchmark:
# python3 benchmark_mrope.py --model-name Qwen/Qwen2.5-VL-7B-Instruct --tp-size 8 \
# --warmup-iter 10 --benchmark-iter 100 --dtype bfloat16 --seed 0 --num-tokens 1024
import argparse
import time
from typing import Any
import numpy as np
import torch
from transformers import AutoConfig
from sglang.srt.layers.rotary_embedding import get_rope
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def get_model_config(model_name: str):
"""Get model configuration parameters"""
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
return config
def generate_test_data(
num_tokens: int,
num_q_heads: int,
num_kv_heads: int,
head_size: int,
max_position_embeddings: int,
dtype: torch.dtype,
device: torch.device,
):
"""Generate test data for given configuration."""
# Create 2D positions (3, num_tokens) for multimodal case
positions = torch.randint(
0, max_position_embeddings // 4, (3, num_tokens), device=device
)
# Create query and key tensors
query = torch.randn(num_tokens, num_q_heads * head_size, dtype=dtype, device=device)
key = torch.randn(num_tokens, num_kv_heads * head_size, dtype=dtype, device=device)
return positions, query, key
def calculate_stats(times: list[float]) -> dict[str, float]:
"""Calculate statistics from a list of times."""
times_array = np.array(times)
return {
"mean": np.mean(times_array),
"median": np.median(times_array),
"p99": np.percentile(times_array, 99),
"min": np.min(times_array),
"max": np.max(times_array),
}
def benchmark_mrope(
model_name: str,
num_tokens: int,
head_dim: int,
tp_size: int,
num_heads: int,
num_kv_heads: int,
max_position: int = 8192,
rope_theta: float = 10000,
is_neox_style: bool = True,
rope_scaling: dict[str, Any] = None,
dtype: torch.dtype = torch.bfloat16,
seed: int = 0,
warmup_iter: int = 10,
benchmark_iter: int = 100,
):
torch.manual_seed(seed)
torch.set_default_device(device)
# the parameters to compute the q k v size based on tp_size
mrope_helper_class = get_rope(
head_size=head_dim,
rotary_dim=head_dim,
max_position=max_position,
base=rope_theta,
is_neox_style=is_neox_style,
rope_scaling=rope_scaling,
dtype=dtype,
).to(device=device)
print(80 * "=")
print(
f"Evaluating model: {model_name} "
f"with tp_size: {tp_size} "
f"and num_tokens: {num_tokens}, "
f"dtype: {dtype}"
)
# create q k v input tensors
# create rotary pos emb input tensors
positions, query, key = generate_test_data(
num_tokens, num_heads, num_kv_heads, head_dim, max_position, dtype, device
)
# Warm up
for _ in range(warmup_iter):
mrope_helper_class.forward_native(
positions,
query.clone(),
key.clone(),
)
mrope_helper_class.forward(
positions,
query.clone(),
key.clone(),
)
torch.cuda.synchronize()
# Time reference implementation
torch_times = []
for _ in range(benchmark_iter):
query_clone = query.clone()
key_clone = key.clone()
torch.cuda.synchronize()
start_time = time.time()
mrope_helper_class.forward_native(
positions,
query_clone,
key_clone,
)
torch.cuda.synchronize()
torch_times.append(time.time() - start_time)
# Time triton kernel implementation
triton_times = []
for _ in range(benchmark_iter):
query_clone = query.clone()
key_clone = key.clone()
torch.cuda.synchronize()
start_time = time.time()
mrope_helper_class.forward(
positions,
query_clone,
key_clone,
)
torch.cuda.synchronize()
triton_times.append(time.time() - start_time)
# Calculate statistics
torch_stats = calculate_stats(torch_times)
triton_stats = calculate_stats(triton_times)
print(f"\nPerformance for config ({num_tokens}, {num_heads}, {num_kv_heads}):")
print(
f"Torch implementation: "
f"mean={torch_stats['mean']:.8f}s, "
f"median={torch_stats['median']:.8f}s, "
f"p99={torch_stats['p99']:.8f}s"
)
print(
f"Triton implementation: "
f"mean={triton_stats['mean']:.8f}s, "
f"median={triton_stats['median']:.8f}s, "
f"p99={triton_stats['p99']:.8f}s"
)
print(
f"Triton Speedup over Torch: {torch_stats['mean'] / triton_stats['mean']:.8f}x"
)
return torch_stats, triton_stats
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Benchmark the rotary embedding kernels."
)
parser.add_argument("--model-name", type=str, default="")
parser.add_argument("--tp-size", type=int, default=1)
parser.add_argument("--warmup-iter", type=int, default=10)
parser.add_argument("--benchmark-iter", type=int, default=100)
parser.add_argument("--dtype", type=str, choices=["bfloat16"], default="bfloat16")
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--num-tokens", type=int, nargs="+", required=False)
parser.add_argument("--trust-remote-code", action="store_true")
args = parser.parse_args()
print(args)
model_tp_dict = {}
if args.model_name == "":
model_tp_dict = {
"Qwen/Qwen2-VL-2B-Instruct": [1],
"Qwen/Qwen2-VL-7B-Instruct": [1],
"Qwen/Qwen2-VL-72B-Instruct": [2, 4, 8],
"Qwen/Qwen2.5-VL-3B-Instruct": [1, 2, 4, 8],
"Qwen/Qwen2.5-VL-7B-Instruct": [1, 2, 4, 8],
"Qwen/Qwen2.5-VL-72B-Instruct": [2, 4, 8],
}
else:
model_tp_dict[args.model_name] = [args.tp_size]
if args.num_tokens is None:
num_tokens_list = [2**i for i in range(0, 18)]
else:
num_tokens_list = args.num_tokens
for model_name, tp_list in model_tp_dict.items():
for tp_size in tp_list:
config = get_model_config(model_name)
# get the model config
total_num_kv_heads = config.num_key_value_heads
total_num_heads = config.num_attention_heads
num_heads = total_num_heads // tp_size
num_kv_heads = max(1, total_num_kv_heads // tp_size)
head_dim = config.hidden_size // total_num_heads
is_neox_style = True
rope_theta = config.rope_theta
max_position = config.max_position_embeddings
for num_tokens in num_tokens_list:
benchmark_mrope(
model_name=model_name,
num_tokens=num_tokens,
head_dim=head_dim,
tp_size=tp_size,
num_heads=num_heads,
num_kv_heads=num_kv_heads,
max_position=max_position,
rope_theta=rope_theta,
is_neox_style=is_neox_style,
rope_scaling=config.rope_scaling,
dtype=getattr(torch, args.dtype),
seed=args.seed,
warmup_iter=args.warmup_iter,
benchmark_iter=args.benchmark_iter,
)
python/sglang/kernels/aot/benchmark/bench_per_tensor_quant_fp8.py
+136
View File
@@ -0,0 +1,136 @@
import itertools
import math
import os
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import torch
import triton
import triton.testing
from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import (
per_tensor_quant_fp8,
)
from sglang.utils import is_in_ci
# Optional imports
try:
from vllm import _custom_ops as ops
VLLM_AVAILABLE = True
except ImportError:
ops = None
VLLM_AVAILABLE = False
from sglang.srt.utils import is_hip
_is_hip = is_hip()
IS_CI = is_in_ci()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
def vllm_scaled_fp8_quant(
input: torch.Tensor,
scale: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
if not VLLM_AVAILABLE:
# Fallback to SGLang implementation
return sglang_scaled_fp8_quant(input, scale)
return ops.scaled_fp8_quant(input, scale)
def sglang_scaled_fp8_quant(
input: torch.Tensor,
scale: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
fp8_type_: torch.dtype = torch.float8_e4m3fn
output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
is_static = True
if scale is None:
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
is_static = False
per_tensor_quant_fp8(input, output, scale, is_static)
return output, scale
def calculate_diff(batch_size: int, seq_len: int):
"""Calculate difference between VLLM and SGLang implementations."""
device = torch.device("cuda")
x = torch.rand((batch_size, seq_len), dtype=torch.float16, device=device)
if not VLLM_AVAILABLE:
print("⚠️ vLLM not available, skipping comparison")
return
vllm_out, vllm_scale = vllm_scaled_fp8_quant(x)
sglang_out, sglang_scale = sglang_scaled_fp8_quant(x)
scale_diff = torch.abs(vllm_scale - sglang_scale).item()
output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
if torch.allclose(
vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
) and torch.allclose(vllm_scale, sglang_scale, rtol=1e-3, atol=1e-5):
print("✅ All implementations match")
else:
print("❌ Implementations differ")
# CI environment uses simplified parameters
if IS_CI:
batch_size_range = [16] # Single batch size for CI
seq_len_range = [64] # Single sequence length for CI
else:
batch_size_range = [16, 32, 64, 128]
seq_len_range = [64, 128, 256, 512, 1024, 2048]
configs = list(itertools.product(batch_size_range, seq_len_range))
if VLLM_AVAILABLE:
line_vals = ["vllm", "sglang"]
line_names = ["VLLM", "SGL Kernel"]
styles = [("blue", "-"), ("green", "-")]
else:
line_vals = ["sglang"]
line_names = ["SGL Kernel"]
styles = [("green", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "seq_len"],
x_vals=configs,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="per-tensor-quant-fp8-performance",
args={},
)
)
def benchmark(batch_size, seq_len, provider):
dtype = torch.float16
device = torch.device("cuda")
x = torch.randn(batch_size * seq_len, 4096, device=device, dtype=dtype)
quantiles = [0.5, 0.2, 0.8]
if provider == "vllm":
fn = lambda: vllm_scaled_fp8_quant(x.clone())
elif provider == "sglang":
fn = lambda: sglang_scaled_fp8_quant(x.clone())
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
calculate_diff(batch_size=4, seq_len=4096)
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_per_token_group_quant_8bit.py
+241
View File
@@ -0,0 +1,241 @@
import itertools
import os
import torch
import triton
from sgl_kernel.test_utils import create_per_token_group_quant_test_data
from sglang.kernels.ops.quantization.fp8_kernel import (
per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
)
from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_8bit,
)
from sglang.srt.utils import is_hip
from sglang.srt.utils.bench_utils import bench_kineto
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
_is_hip = is_hip()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
mode_concentrated = IS_CI or (os.environ.get("SGLANG_BENCH_MODE", "") == "concentrated")
if int(os.environ.get("SGLANG_NSYS_PROFILING", "0")):
configs = [
[
768 * 8,
2048,
128,
48,
fp8_type_,
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
# masked_layout_mode=None,
masked_layout_mode="balanced",
# masked_layout_mode="extreme",
),
]
]
elif mode_concentrated:
configs = list(
itertools.product(
[768],
[1536, 7168, 16384],
[128],
[None],
[fp8_type_],
[
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
],
)
) + list(
itertools.product(
[768 * 8],
[2048],
[128],
[48],
[fp8_type_],
[
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="balanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="imbalanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="extreme",
),
],
)
)
else:
configs = list(
itertools.product(
[1, 4, 16, 64, 256, 768, 2048, 8192, 16384],
[1536, 7168, 16384],
[128],
[None],
[fp8_type_],
[
dict(
column_major_scales=False,
scale_tma_aligned=False,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=False,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=False,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=False,
masked_layout_mode=None,
),
],
)
) + list(
itertools.product(
[1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
[2048],
[128],
[8, 16, 32, 48],
[fp8_type_],
[
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode=None,
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="balanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="imbalanced",
),
dict(
column_major_scales=True,
scale_tma_aligned=True,
scale_ue8m0=True,
fuse_silu_and_mul=True,
masked_layout_mode="extreme",
),
],
)
)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=[
"num_tokens",
"hidden_dim",
"group_size",
"num_ranks",
"dst_dtype",
"flags",
],
x_vals=configs,
line_arg="provider",
line_vals=["triton", "sglang"],
# Triton has multi kernels and we only report the time for the core one
line_names=["Triton (Inaccurate)", "SGL Kernel"],
styles=[("blue", "-"), ("green", "-")],
ylabel="us",
plot_name="per-token-group-quant-8bit-performance",
args={},
)
)
def benchmark(
num_tokens, hidden_dim, group_size, num_ranks, dst_dtype, flags, provider
):
print(
f"Testing: {num_tokens=} {hidden_dim=} {group_size=} {num_ranks=} {dst_dtype=} {flags=} {provider=}"
)
x, masked_m = create_per_token_group_quant_test_data(
num_tokens=num_tokens, hidden_dim=hidden_dim, num_ranks=num_ranks, flags=flags
)
fn, kernel_names = {
"triton": (
triton_per_token_group_quant_8bit,
"_per_token_group_quant_8bit|_silu_and_mul_post_quant_kernel",
),
"sglang": (
sglang_per_token_group_quant_8bit,
"per_token_group_quant_8bit_kernel",
),
}[provider]
def bench_fn():
return fn(
x=x,
masked_m=masked_m,
group_size=group_size,
dst_dtype=dst_dtype,
**{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
)
time_s = bench_kineto(
bench_fn, kernel_names=kernel_names, num_tests=300 if mode_concentrated else 30
)
return time_s * 1e6
if __name__ == "__main__":
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_per_token_quant_fp8.py
+228
View File
@@ -0,0 +1,228 @@
import itertools
import os
from typing import Optional, Tuple
import torch
import triton
import triton.testing
from sgl_kernel import sgl_per_token_quant_fp8
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm import _custom_ops as ops
VLLM_AVAILABLE = True
except ImportError:
ops = None
VLLM_AVAILABLE = False
from sglang.srt.utils import is_hip
_is_hip = is_hip()
IS_CI = is_in_ci()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
# Get correct FP8 E4M3 maximum value
if _is_hip:
FP8_E4M3_MAX = 224.0 # ROCM uses 224.0
else:
# For CUDA, get the actual max value from the type
FP8_E4M3_MAX = float(torch.finfo(fp8_type_).max)
def torch_per_token_quant_fp8(
input: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Pure PyTorch reference implementation for per-token FP8 quantization."""
device = input.device
dtype = input.dtype
# Find max absolute value per token (row) - exactly like CUDA kernel
max_vals = torch.abs(input).max(dim=1)[0] # [num_tokens]
# Calculate scale per token - exactly like CUDA kernel: scale = max_value / FP8_E4M3_MAX
scales = max_vals / FP8_E4M3_MAX # [num_tokens]
# No special zero handling - directly compute 1.0 / scale like CUDA kernel
scale_inv = 1.0 / scales # [num_tokens]
# Quantize: input * scale_inv, then clamp to FP8 range
quantized_float = input * scale_inv.unsqueeze(1) # Broadcast scale_inv
quantized_float = torch.clamp(quantized_float, -FP8_E4M3_MAX, FP8_E4M3_MAX)
# Convert to FP8 - use more explicit conversion
quantized_fp8 = quantized_float.to(fp8_type_)
return quantized_fp8, scales
def vllm_per_token_quant_fp8(
input: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
if not VLLM_AVAILABLE:
# Fallback to SGLang implementation
return sglang_per_token_quant_fp8(input)
return ops.scaled_fp8_quant(input, use_per_token_if_dynamic=True)
def sglang_per_token_quant_fp8(
input: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
scale = torch.zeros(input.size(0), device=input.device, dtype=torch.float32)
output = torch.empty_like(input, device=input.device, dtype=fp8_type_)
sgl_per_token_quant_fp8(input, output, scale)
return output, scale
def calculate_diff(batch_size: int, seq_len: int, hidden_dim: int):
"""Compare Torch reference, VLLM, and SGLang implementations."""
device = torch.device("cuda")
x = torch.rand(
(batch_size * seq_len, hidden_dim), dtype=torch.float16, device=device
)
# Get all three implementations
torch_out, torch_scale = torch_per_token_quant_fp8(x)
vllm_out, vllm_scale = vllm_per_token_quant_fp8(x)
sglang_out, sglang_scale = sglang_per_token_quant_fp8(x)
if not VLLM_AVAILABLE:
print("⚠️ vLLM not available, skipping vLLM comparison")
# Only compare Torch vs SGLang
torch_sglang_scale_diff = torch.abs(torch_scale - sglang_scale).mean().item()
torch_sglang_out_diff = (
torch.abs(torch_out.float() - sglang_out.float()).mean().item()
)
print(f"Scale difference (Torch vs SGLang): {torch_sglang_scale_diff:.8f}")
print(f"Output difference (Torch vs SGLang): {torch_sglang_out_diff:.8f}")
return
print(f"\n=== Comparison for hidden_dim={hidden_dim} ===")
# Compare scales
torch_vllm_scale_diff = torch.abs(torch_scale - vllm_scale).mean().item()
torch_sglang_scale_diff = torch.abs(torch_scale - sglang_scale).mean().item()
vllm_sglang_scale_diff = torch.abs(vllm_scale - sglang_scale).mean().item()
print(f"Scale differences:")
print(f" Torch vs VLLM: {torch_vllm_scale_diff:.8f}")
print(f" Torch vs SGLang: {torch_sglang_scale_diff:.8f}")
print(f" VLLM vs SGLang: {vllm_sglang_scale_diff:.8f}")
# Compare outputs
torch_vllm_out_diff = torch.abs(torch_out.float() - vllm_out.float()).mean().item()
torch_sglang_out_diff = (
torch.abs(torch_out.float() - sglang_out.float()).mean().item()
)
vllm_sglang_out_diff = (
torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
)
print(f"Output differences:")
print(f" Torch vs VLLM: {torch_vllm_out_diff:.8f}")
print(f" Torch vs SGLang: {torch_sglang_out_diff:.8f}")
print(f" VLLM vs SGLang: {vllm_sglang_out_diff:.8f}")
# Check tolerances
rtol, atol = 1e-3, 1e-5
torch_vllm_match = torch.allclose(
torch_out.float(), vllm_out.float(), rtol=rtol, atol=atol
) and torch.allclose(torch_scale, vllm_scale, rtol=rtol, atol=atol)
torch_sglang_match = torch.allclose(
torch_out.float(), sglang_out.float(), rtol=rtol, atol=atol
) and torch.allclose(torch_scale, sglang_scale, rtol=rtol, atol=atol)
if hidden_dim == 1368:
rtol = 1e-2
# we found vllm sglang has diff when hidden dim is not dividable by 16
# and we believe SGLang is closer to Torch implementation
vllm_sglang_match = torch.allclose(
vllm_out.float(), sglang_out.float(), rtol=rtol, atol=atol
) and torch.allclose(vllm_scale, sglang_scale, rtol=rtol, atol=atol)
print(f"Matches (rtol={rtol}, atol={atol}):")
print(f" Torch vs VLLM: {'✅' if torch_vllm_match else '❌'}")
print(f" Torch vs SGLang: {'✅' if torch_sglang_match else '❌'}")
print(f" VLLM vs SGLang: {'✅' if vllm_sglang_match else '❌'}")
# CI environment uses simplified parameters
if IS_CI:
batch_size_range = [16] # Single batch size for CI
seq_len_range = [64] # Single sequence length for CI
hidden_dim_range = [2048] # Single hidden dimension for CI
else:
batch_size_range = [16, 32, 64, 128]
seq_len_range = [64, 128, 256, 512, 1024, 2048, 4096]
hidden_dim_range = [1368, 2048, 4096]
configs = list(itertools.product(batch_size_range, seq_len_range, hidden_dim_range))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "seq_len", "hidden_dim"],
x_vals=configs,
line_arg="provider",
line_vals=(
["torch", "vllm", "sglang"] if VLLM_AVAILABLE else ["torch", "sglang"]
),
line_names=(
["Torch Reference", "VLLM", "SGL Kernel"]
if VLLM_AVAILABLE
else ["Torch Reference", "SGL Kernel"]
),
styles=(
[("red", "-"), ("blue", "-"), ("green", "-")]
if VLLM_AVAILABLE
else [("red", "-"), ("green", "-")]
),
ylabel="us",
plot_name="per-token-dynamic-quant-fp8-performance",
args={},
)
)
def benchmark_quantization(batch_size, seq_len, hidden_dim, provider):
dtype = torch.float16
device = torch.device("cuda")
x = torch.randn(batch_size * seq_len, hidden_dim, device=device, dtype=dtype)
quantiles = [0.5, 0.2, 0.8]
if provider == "torch":
fn = lambda: torch_per_token_quant_fp8(x.clone())
elif provider == "vllm":
if not VLLM_AVAILABLE:
return (0, 0, 0)
fn = lambda: vllm_per_token_quant_fp8(x.clone())
elif provider == "sglang":
fn = lambda: sglang_per_token_quant_fp8(x.clone())
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
# Test various hidden dimensions for correctness - simplified for CI
if IS_CI:
test_dims = [2048] # Single dimension for CI
batch_size, seq_len = 4, 64 # Smaller values for CI
else:
test_dims = [1368, 2048, 4096]
batch_size, seq_len = 4, 4096
for dim in test_dims:
calculate_diff(batch_size=batch_size, seq_len=seq_len, hidden_dim=dim)
print("\n" + "=" * 60)
print("Starting performance benchmark...")
benchmark_quantization.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_rmsnorm.py
+395
View File
@@ -0,0 +1,395 @@
# Benchmarks SGLang RMSNorm kernels versus vLLM and FlashInfer across
# (batch_size, seq_len, hidden_size) and prints speed-up.
import argparse
import itertools
import os
import re
from typing import List, Optional, Tuple, Union
import sgl_kernel
import torch
import torch.nn as nn
import triton
import triton.testing
from sgl_kernel.utils import is_arch_support_pdl
from sglang.utils import is_in_ci
# Optional imports
try:
from flashinfer.norm import fused_add_rmsnorm, rmsnorm
FLASHINFER_AVAILABLE = True
except ImportError:
fused_add_rmsnorm = None
rmsnorm = None
FLASHINFER_AVAILABLE = False
try:
from vllm import _custom_ops as vllm_ops
VLLM_AVAILABLE = True
except ImportError:
vllm_ops = None
VLLM_AVAILABLE = False
IS_CI = is_in_ci()
def str2int_list(arg: str) -> List[int]:
if arg in ("", None):
return []
if re.fullmatch(r"\d+(,\d+)*", arg.strip()) is None:
raise argparse.ArgumentTypeError(f"Bad int list: {arg}")
return [int(x) for x in arg.split(",")]
class HuggingFaceRMSNorm(nn.Module):
def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
orig_dtype = x.dtype
x = x.to(torch.float32)
if residual is not None:
x = x + residual.to(torch.float32)
residual = x.to(orig_dtype)
variance = x.pow(2).mean(dim=-1, keepdim=True)
x = x * torch.rsqrt(variance + self.variance_epsilon)
x = x.to(orig_dtype) * self.weight
if residual is None:
return x
else:
return x, residual
def rmsnorm_naive(
x: torch.Tensor,
weight: torch.Tensor,
residual: Optional[torch.Tensor] = None,
eps: float = 1e-6,
):
naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps)
naive_norm.weight = nn.Parameter(weight)
naive_norm = naive_norm.to(x.device)
orig_shape = x.shape
x = x.view(-1, x.shape[-1])
if residual is not None:
residual = residual.view(-1, residual.shape[-1])
output = naive_norm(x, residual)
if isinstance(output, tuple):
output = (output[0].view(orig_shape), output[1].view(orig_shape))
else:
output = output.view(orig_shape)
return output
def rmsnorm_flashinfer(
x: torch.Tensor,
weight: torch.Tensor,
residual: Optional[torch.Tensor] = None,
eps: float = 1e-6,
):
if not FLASHINFER_AVAILABLE:
# Fallback to naive implementation if FlashInfer is not available
return rmsnorm_naive(x, weight, residual, eps)
orig_shape = x.shape
x = x.view(-1, x.shape[-1])
if residual is not None:
residual = residual.view(-1, residual.shape[-1])
if residual is not None:
fused_add_rmsnorm(x, residual, weight, eps)
output = (x, residual)
else:
output = rmsnorm(x, weight, eps)
if isinstance(output, tuple):
output = (output[0].view(orig_shape), output[1].view(orig_shape))
else:
output = output.view(orig_shape)
return output
def rmsnorm_vllm(
x: torch.Tensor,
weight: torch.Tensor,
residual: Optional[torch.Tensor] = None,
eps: float = 1e-6,
):
if not VLLM_AVAILABLE:
# Fallback to naive implementation if vLLM is not available
return rmsnorm_naive(x, weight, residual, eps)
orig_shape = x.shape
x = x.view(-1, x.shape[-1])
if residual is not None:
residual = residual.view(-1, residual.shape[-1])
if residual is not None:
vllm_ops.fused_add_rms_norm(x, residual, weight, eps)
output = (x, residual)
else:
out = torch.empty_like(x)
vllm_ops.rms_norm(out, x, weight, eps)
output = out
if isinstance(output, tuple):
output = (output[0].view(orig_shape), output[1].view(orig_shape))
else:
output = output.view(orig_shape)
return output
def rmsnorm_sglang(
x: torch.Tensor,
weight: torch.Tensor,
residual: Optional[torch.Tensor] = None,
eps: float = 1e-6,
enable_pdl: Optional[bool] = None,
):
orig_shape = x.shape
x = x.view(-1, x.shape[-1])
if residual is not None:
residual = residual.view(-1, residual.shape[-1])
if enable_pdl is None:
enable_pdl = is_arch_support_pdl()
if residual is not None:
sgl_kernel.fused_add_rmsnorm(x, residual, weight, eps, enable_pdl=enable_pdl)
output = (x, residual)
else:
out = torch.empty_like(x)
sgl_kernel.rmsnorm(x, weight, eps, out=out, enable_pdl=enable_pdl)
output = out
if isinstance(output, tuple):
output = (output[0].view(orig_shape), output[1].view(orig_shape))
else:
output = output.view(orig_shape)
return output
def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True):
dtype = torch.bfloat16
x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda")
weight = torch.ones(hidden_size, dtype=dtype, device="cuda")
residual = torch.randn_like(x) if use_residual else None
output_naive = rmsnorm_naive(
x.clone(), weight, residual.clone() if residual is not None else None
)
output_flashinfer = rmsnorm_flashinfer(
x.clone(), weight, residual.clone() if residual is not None else None
)
output_vllm = rmsnorm_vllm(
x.clone(), weight, residual.clone() if residual is not None else None
)
output_sglang = rmsnorm_sglang(
x.clone(), weight, residual.clone() if residual is not None else None
)
if use_residual:
output_naive = output_naive[0]
output_flashinfer = output_flashinfer[0]
output_vllm = output_vllm[0]
output_sglang = output_sglang[0]
print(f"Naive output={output_naive}")
if FLASHINFER_AVAILABLE:
print(f"FlashInfer output={output_flashinfer}")
else:
print("FlashInfer not available, skipped")
if VLLM_AVAILABLE:
print(f"VLLM output={output_vllm}")
else:
print("vLLM not available, skipped")
print(f"SGLang output={output_sglang}")
# Only compare available implementations
all_match = torch.allclose(output_naive, output_sglang, atol=1e-2, rtol=1e-2)
if FLASHINFER_AVAILABLE:
all_match = all_match and torch.allclose(
output_naive, output_flashinfer, atol=1e-2, rtol=1e-2
)
if VLLM_AVAILABLE:
all_match = all_match and torch.allclose(
output_naive, output_vllm, atol=1e-2, rtol=1e-2
)
if all_match:
print("✅ All available implementations match")
else:
print("❌ Implementations differ")
# CI environment uses simplified parameters
if IS_CI:
default_batch_sizes = [1] # Single batch size for CI
default_seq_lens = [64] # Single sequence length for CI
default_hidden_sizes = [4096] # Single hidden size for CI
else:
default_batch_sizes = [2**i for i in range(0, 7, 2)] # 1, 4, 16, 64
default_seq_lens = [2**i for i in range(6, 11, 1)] # 64, 128, 256, 512, 1024
default_hidden_sizes = [32 * 128, 48 * 128] # 4096, 6144
def make_configs(bsizes: List[int], slens: List[int], hsizes: List[int]) -> List[Tuple]:
return list(itertools.product(bsizes, slens, hsizes))
# Filter providers based on availability
available_providers = ["huggingface", "sglang"]
available_names = ["HuggingFace", "SGL Kernel"]
available_styles = [("blue", "-"), ("orange", "-")]
if FLASHINFER_AVAILABLE:
available_providers.insert(-1, "flashinfer")
available_names.insert(-1, "FlashInfer")
available_styles.insert(-1, ("green", "-"))
if VLLM_AVAILABLE:
available_providers.insert(-1, "vllm")
available_names.insert(-1, "vLLM")
available_styles.insert(-1, ("red", "-"))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "seq_len", "hidden_size"],
x_vals=[],
line_arg="provider",
line_vals=available_providers,
line_names=available_names,
styles=available_styles,
ylabel="µs (median) or × (speed-up)",
plot_name="rmsnorm-performance",
args={},
)
)
def benchmark(batch_size, seq_len, hidden_size, provider, use_residual):
device = torch.device("cuda")
dtype = torch.bfloat16
x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device)
weight = torch.ones(hidden_size, dtype=dtype, device=device)
residual = torch.randn_like(x) if use_residual else None
# timing helper
def timed(fn):
for _ in range(5):
fn()
torch.cuda.synchronize()
ms, qmin, qmax = triton.testing.do_bench_cudagraph(
fn, quantiles=[0.5, 0.2, 0.8]
)
return 1000 * ms, 1000 * qmax, 1000 * qmin
if provider == "huggingface":
return timed(
lambda: rmsnorm_naive(
x.clone(),
weight,
residual.clone() if residual is not None else None,
)
)
elif provider == "flashinfer":
if not FLASHINFER_AVAILABLE:
return (0, 0, 0)
return timed(
lambda: rmsnorm_flashinfer(
x.clone(),
weight,
residual.clone() if residual is not None else None,
)
)
elif provider == "vllm":
if not VLLM_AVAILABLE:
return (0, 0, 0)
return timed(
lambda: rmsnorm_vllm(
x.clone(),
weight,
residual.clone() if residual is not None else None,
)
)
elif provider == "sglang":
return timed(
lambda: rmsnorm_sglang(
x.clone(),
weight,
residual.clone() if residual is not None else None,
)
)
# provider == "speedup"
if VLLM_AVAILABLE:
t_ref, _, _ = timed(
lambda: rmsnorm_vllm(
x.clone(),
weight,
residual.clone() if residual is not None else None,
)
)
else:
t_ref, _, _ = timed(
lambda: rmsnorm_naive(
x.clone(),
weight,
residual.clone() if residual is not None else None,
)
)
t_sgl, _, _ = timed(
lambda: rmsnorm_sglang(
x.clone(),
weight,
residual.clone() if residual is not None else None,
)
)
spd = t_ref / t_sgl if t_ref > 0 else 1.0
return (spd, spd, spd)
if __name__ == "__main__":
p = argparse.ArgumentParser("RMSNorm kernel benchmark")
p.add_argument("--batch_sizes", type=str2int_list, default=default_batch_sizes)
p.add_argument("--seq_lens", type=str2int_list, default=default_seq_lens)
p.add_argument("--hidden_sizes", type=str2int_list, default=default_hidden_sizes)
p.add_argument(
"--use_residual", action="store_true", help="Whether to use residual connection"
)
p.add_argument("--verify_only", action="store_true")
args = p.parse_args()
# coerce lists
if isinstance(args.batch_sizes, str):
args.batch_sizes = str2int_list(args.batch_sizes)
if isinstance(args.seq_lens, str):
args.seq_lens = str2int_list(args.seq_lens)
if isinstance(args.hidden_sizes, str):
args.hidden_sizes = str2int_list(args.hidden_sizes)
# patch perf_report grid
benchmark_grid = make_configs(args.batch_sizes, args.seq_lens, args.hidden_sizes)
if hasattr(benchmark, "benchmarks"):
benchmark.benchmarks.x_vals = benchmark_grid
else:
benchmark.benchmark.x_vals = benchmark_grid
if args.verify_only:
ok = calculate_diff(4, 128, args.hidden_sizes[0], args.use_residual)
print("✅ sanity pass" if ok else "❌ mismatch")
else:
benchmark.run(print_data=True, use_residual=args.use_residual)
python/sglang/kernels/aot/benchmark/bench_rotary_embedding.py
+108
View File
@@ -0,0 +1,108 @@
import itertools
import os
import torch
import triton
from sgl_kernel.testing.rotary_embedding import (
FlashInferRotaryEmbedding,
FusedSetKVBufferArg,
MHATokenToKVPool,
RotaryEmbedding,
create_inputs,
)
from sglang.srt.utils.bench_utils import bench_kineto
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
# CI environment uses simplified parameters
if IS_CI:
batch_seq_configs = [(1, 1)] # Single config for CI
save_kv_configs = [False] # Single option for CI
else:
batch_seq_configs = [
(1, 1),
(32, 1),
(128, 1),
(512, 1),
(2, 512),
(4, 4096),
]
save_kv_configs = [False, True]
configs = [
(batch_size, seq_len, save_kv_cache)
for batch_size, seq_len in batch_seq_configs
for save_kv_cache in save_kv_configs
]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "seq_len", "save_kv_cache"],
x_vals=configs,
line_arg="provider",
line_vals=["sglang"],
line_names=["SGL Kernel"],
styles=[("green", "-")],
ylabel="us",
plot_name="bench_rotary_embedding",
args={},
)
)
def benchmark(batch_size, seq_len, save_kv_cache, provider):
device = torch.device("cuda")
num_q_heads = 32
num_kv_heads = 8
head_size = 64
dtype = torch.bfloat16
config = dict(
head_size=head_size,
rotary_dim=64,
max_position_embeddings=4096,
base=8000,
is_neox_style=True,
dtype=dtype,
)
rope_flashinfer = FlashInferRotaryEmbedding(**config).to(device)
pool_flashinfer = MHATokenToKVPool(head_num=num_kv_heads, head_dim=head_size)
inputs = create_inputs(
head_size=head_size,
batch_size=batch_size,
seq_len=seq_len,
device=device,
dtype=dtype,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
)
query_flashinfer, key_flashinfer = inputs["query"].clone(), inputs["key"].clone()
bench_fn = lambda: rope_flashinfer.forward_cuda(
inputs["pos_ids"],
query_flashinfer,
key_flashinfer,
fused_set_kv_buffer_arg=(
FusedSetKVBufferArg(
value=inputs["value"],
k_buffer=pool_flashinfer.k_buffer[0].view(-1, num_kv_heads * head_size),
v_buffer=pool_flashinfer.v_buffer[0].view(-1, num_kv_heads * head_size),
k_scale=None,
v_scale=None,
cache_loc=inputs["out_cache_loc"],
)
if save_kv_cache
else None
),
)
time_s = bench_kineto(bench_fn, kernel_names="BatchQKApplyRotaryPosIds")
return time_s * 1e6
if __name__ == "__main__":
benchmark.run(print_data=True)
python/sglang/kernels/aot/benchmark/bench_sum_scale.py
+249
View File
@@ -0,0 +1,249 @@
import os
import torch
import triton
import triton.language as tl
from sgl_kernel import moe_sum_reduce as moe_sum_reduce_cuda
from triton.testing import do_bench
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
@triton.jit
def _moe_sum_reduce_kernel(
input_ptr,
input_stride_0,
input_stride_1,
input_stride_2,
output_ptr,
output_stride_0,
output_stride_1,
token_num: int,
topk_num: int,
hidden_dim: int,
routed_scaling_factor: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_DIM: tl.constexpr,
NUM_STAGE: tl.constexpr,
):
input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64)
input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64)
output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64)
token_block_id = tl.program_id(0)
dim_block_id = tl.program_id(1)
offs_token = token_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
offs_dim = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM)
mask_token = offs_token < token_num
mask_dim = offs_dim < hidden_dim
base_ptrs = input_ptr + offs_token[:, None] * input_stride_0 + offs_dim[None, :]
accumulator = tl.zeros((BLOCK_M, BLOCK_DIM), dtype=tl.float32)
for i in tl.range(0, topk_num, num_stages=NUM_STAGE):
tile = tl.load(
base_ptrs + i * input_stride_1,
mask=mask_token[:, None] & mask_dim[None, :],
other=0.0,
)
accumulator += tile.to(tl.float32)
accumulator *= routed_scaling_factor
# -------- Write back --------
store_ptrs = output_ptr + offs_token[:, None] * output_stride_0 + offs_dim[None, :]
tl.store(
store_ptrs,
accumulator.to(input_ptr.dtype.element_ty),
mask=mask_token[:, None] & mask_dim[None, :],
)
# _moe_sum_reduce_kernel kernel modified from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/moe_sum_reduce.py
def moe_sum_reduce_triton(
input: torch.Tensor, output: torch.Tensor, routed_scaling_factor: float
):
assert input.is_contiguous()
assert output.is_contiguous()
token_num, topk_num, hidden_dim = input.shape
assert output.shape[0] == token_num and output.shape[1] == hidden_dim
BLOCK_M = 1
BLOCK_DIM = 2048
NUM_STAGE = 1
num_warps = 16
grid = (
triton.cdiv(token_num, BLOCK_M),
triton.cdiv(hidden_dim, BLOCK_DIM),
)
_moe_sum_reduce_kernel[grid](
input,
*input.stride(),
output,
*output.stride(),
token_num=token_num,
topk_num=topk_num,
hidden_dim=hidden_dim,
routed_scaling_factor=routed_scaling_factor,
BLOCK_M=BLOCK_M,
BLOCK_DIM=BLOCK_DIM,
NUM_STAGE=NUM_STAGE,
num_warps=num_warps,
)
return
def compute_sum_scaled_baseline(
x: torch.Tensor, out: torch.Tensor, routed_scaling_factor: float
) -> torch.Tensor:
torch.sum(x, dim=1, out=out)
out.mul_(routed_scaling_factor)
return out
@torch.compile
def compute_sum_scaled_compiled(
x: torch.Tensor, out: torch.Tensor, routed_scaling_factor: float
) -> torch.Tensor:
torch.sum(x * routed_scaling_factor, dim=1, out=out)
return out
def get_benchmark(dtype=torch.bfloat16):
num_tokens_range = [2**i for i in range(0, 13)]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_tokens"],
x_vals=num_tokens_range,
line_arg="version",
line_vals=["baseline", "compiled", "triton", "cuda"],
line_names=["Original", "TorchCompile", "TritonKernel", "CudaKernel"],
styles=[("blue", "-"), ("green", "-"), ("red", "-"), ("yellow", "-")],
ylabel="us",
plot_name=f"sum_scaled_performance_{str(dtype).split('.')[-1]}",
args={},
)
)
def benchmark(num_tokens, version):
topk = 9
hidden_size = 4096
dtype = torch.bfloat16
scaling_factor = 0.3
x = torch.randn(num_tokens, topk, hidden_size, dtype=dtype, device="cuda")
out = torch.empty(num_tokens, hidden_size, dtype=dtype, device="cuda")
# Warmup
for _ in range(3):
if version == "baseline":
compute_sum_scaled_baseline(x, out, scaling_factor)
elif version == "compiled":
compute_sum_scaled_compiled(x, out, scaling_factor)
elif version == "triton":
moe_sum_reduce_triton(x, out, scaling_factor)
else:
moe_sum_reduce_cuda(x, out, scaling_factor)
# Benchmark
quantiles = [0.5, 0.2, 0.8]
if version == "baseline":
ms, min_ms, max_ms = do_bench(
lambda: compute_sum_scaled_baseline(x, out, scaling_factor),
quantiles=quantiles,
)
elif version == "compiled":
ms, min_ms, max_ms = do_bench(
lambda: compute_sum_scaled_compiled(x, out, scaling_factor),
quantiles=quantiles,
)
elif version == "triton":
ms, min_ms, max_ms = do_bench(
lambda: moe_sum_reduce_triton(x, out, scaling_factor),
quantiles=quantiles,
)
else:
ms, min_ms, max_ms = do_bench(
lambda: moe_sum_reduce_cuda(x, out, scaling_factor),
quantiles=quantiles,
)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return benchmark
def verify_correctness(num_tokens=1024, dtype=torch.bfloat16):
x = torch.randn(num_tokens, 9, 4096, device="cuda", dtype=dtype)
scaling_factor = 0.3
out_baseline = torch.empty_like(x[:, 0])
compute_sum_scaled_baseline(x, out_baseline, scaling_factor)
out_compiled = torch.empty_like(out_baseline)
compute_sum_scaled_compiled(x, out_compiled, scaling_factor)
out_cuda = torch.empty_like(out_baseline)
moe_sum_reduce_cuda(x, out_cuda, scaling_factor)
triton_skipped = dtype == torch.float64
if not triton_skipped:
out_triton = torch.empty_like(out_baseline)
moe_sum_reduce_triton(x, out_triton, scaling_factor)
if dtype == torch.float64:
atol, rtol = 1e-12, 1e-12
elif dtype == torch.float32:
atol, rtol = 1e-6, 1e-6
else: # bfloat16 / float16
atol, rtol = 1e-2, 1e-2
ok_compiled = torch.allclose(out_baseline, out_compiled, atol=atol, rtol=rtol)
ok_cuda = torch.allclose(out_baseline, out_cuda, atol=atol, rtol=rtol)
ok_triton = (
True
if triton_skipped
else torch.allclose(out_baseline, out_triton, atol=atol, rtol=rtol)
)
if ok_compiled and ok_triton and ok_cuda:
msg = "✅ All implementations match"
if triton_skipped:
msg += " (Triton skipped for float64)"
print(msg)
else:
print("❌ Implementations differ")
print(
f"Baseline vs Compiled: {(out_baseline - out_compiled).abs().max().item()}"
)
if not triton_skipped:
print(
f"Baseline vs Triton: {(out_baseline - out_triton).abs().max().item()}"
)
print(f"Baseline vs Cuda: {(out_baseline - out_cuda).abs().max().item()}")
if __name__ == "__main__":
print("Running correctness verification for bfloat16...")
verify_correctness(dtype=torch.bfloat16)
# CI environment uses simplified parameters
if not IS_CI:
print("Running correctness verification for float64...")
verify_correctness(dtype=torch.float64)
print("Running correctness verification for float64...")
verify_correctness(dtype=torch.float64)
print("\nRunning performance benchmark for bfloat16...")
benchmark = get_benchmark(dtype=torch.bfloat16)
benchmark.run(
print_data=True,
# save_path="./configs/benchmark_ops/sum_scaled/"
)
python/sglang/kernels/aot/benchmark/bench_top_k_top_p_sampling.py
+146
View File
@@ -0,0 +1,146 @@
import itertools
import os
import flashinfer.sampling
import sgl_kernel
import torch
import triton
import triton.testing
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
def torch_top_k_top_p_joint_sampling_from_probs(
normalized_prob, top_k, top_p, eps=1e-4
):
"""Reference PyTorch implementation of joint top-k top-p sampling."""
batch_size, vocab_size = normalized_prob.shape
samples = torch.empty(batch_size, dtype=torch.int64, device=normalized_prob.device)
for i in range(batch_size):
p_val = top_p[i].item()
k_val = top_k[i].item()
# top-p mask
sorted_prob, indices = torch.sort(normalized_prob[i], descending=False)
cdf = torch.cumsum(sorted_prob, dim=-1)
mask_top_p = torch.zeros(
vocab_size, dtype=torch.int32, device=normalized_prob.device
)
mask_top_p.scatter_add_(0, indices, (cdf > (1 - p_val) - eps).int())
# top-k mask
sorted_prob_desc, _ = torch.sort(normalized_prob[i], descending=True)
pivot = sorted_prob_desc[k_val - 1]
mask_top_k = (normalized_prob[i] >= pivot).int()
# joint mask
mask = torch.minimum(mask_top_p, mask_top_k).bool()
# sample from masked probs
masked_probs = normalized_prob[i] * mask
masked_probs = masked_probs / masked_probs.sum()
idx = torch.multinomial(masked_probs, 1)
samples[i] = idx
return samples
def calculate_diff(batch_size, vocab_size, p):
"""Compare Torch reference and SGLang kernel for correctness."""
torch.manual_seed(42)
if p == 0.1:
k = int(vocab_size * 0.5)
elif p == 0.5:
k = int(vocab_size * 0.1)
else:
raise ValueError("p not recognized")
device = torch.device("cuda")
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
top_p_tensor = torch.full((batch_size,), p, device=device)
top_k_tensor = torch.full((batch_size,), k, device=device)
torch_samples = torch_top_k_top_p_joint_sampling_from_probs(
normalized_prob, top_k_tensor, top_p_tensor
)
sglang_samples = flashinfer.sampling.top_k_top_p_sampling_from_probs(
normalized_prob, top_k_tensor, top_p_tensor, filter_apply_order="joint"
)
# parameter space - simplified for CI
if IS_CI:
batch_size_range = [16] # Single batch size for CI
vocab_size_range = [111] # Single vocab size for CI
p_range = [0.1] # Single p value for CI
else:
batch_size_range = [16, 64, 128]
vocab_size_range = [111, 32000]
p_range = [0.1, 0.5]
configs = list(itertools.product(batch_size_range, vocab_size_range, p_range))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["batch_size", "vocab_size", "p"],
x_vals=configs,
line_arg="provider",
line_vals=["torch", "sglang"],
line_names=["Torch Reference", "SGL Kernel"],
styles=[("red", "-"), ("green", "-")],
ylabel="us",
plot_name="top-k-top-p-joint-sampling-performance",
args={},
)
)
def benchmark_sampling(batch_size, vocab_size, p, provider):
torch.manual_seed(42)
if p == 0.1:
k = int(vocab_size * 0.5)
elif p == 0.5:
k = int(vocab_size * 0.1)
else:
raise ValueError("p not recognized")
device = torch.device("cuda")
pre_norm_prob = torch.rand(batch_size, vocab_size, device=device)
normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True)
top_p_tensor = torch.full((batch_size,), p, device=device)
top_k_tensor = torch.full((batch_size,), k, device=device)
if provider == "torch":
fn = lambda: torch_top_k_top_p_joint_sampling_from_probs(
normalized_prob.clone(), top_k_tensor, top_p_tensor
)
elif provider == "sglang":
fn = lambda: flashinfer.sampling.top_k_top_p_sampling_from_probs(
normalized_prob.clone(),
top_k_tensor,
top_p_tensor,
filter_apply_order="joint",
)
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
if __name__ == "__main__":
# Correctness check - simplified for CI
if IS_CI:
# Only test one configuration in CI
test_configs = [configs[0]] if configs else [(16, 111, 0.1)]
else:
test_configs = configs
for cfg in test_configs:
calculate_diff(*cfg)
print("\n" + "=" * 60)
print("Starting performance benchmark...")
benchmark_sampling.run(print_data=True)
python/sglang/kernels/aot/build.sh
Executable
+164
View File
@@ -0,0 +1,164 @@
#!/bin/bash
set -ex
if [ $# -lt 2 ]; then
echo "Usage: $0 <PYTHON_VERSION> <CUDA_VERSION> [ARCH]"
exit 1
fi
PYTHON_VERSION="$1" # e.g. 3.10
CUDA_VERSION="$2" # e.g. 12.9
ARCH="${3:-$(uname -i)}" # optional override
if [ "${ARCH}" = "aarch64" ]; then
BASE_IMG="pytorch/manylinuxaarch64-builder"
else
BASE_IMG="pytorch/manylinux2_28-builder"
fi
# Create cache directories for persistent build artifacts in home directory
# Using home directory to persist across workspace cleanups/checkouts
CACHE_DIR="${HOME}/.cache/sgl-kernel"
BUILDX_CACHE_DIR="${CACHE_DIR}/buildx"
CCACHE_HOST_DIR="${CACHE_DIR}/ccache"
mkdir -p "${BUILDX_CACHE_DIR}" "${CCACHE_HOST_DIR}"
# Ensure a buildx builder with docker-container driver (required for cache export)
BUILDER_NAME="sgl-kernel-builder"
# RESET_BUILDER=1 removes and recreates the builder to clear corrupted internal
# state (e.g. stale containerd snapshots from base image layer GC).
if [ "${RESET_BUILDER:-0}" = "1" ]; then
echo "Resetting buildx builder: ${BUILDER_NAME}"
docker buildx rm "${BUILDER_NAME}" 2>/dev/null || true
rm -rf "${BUILDX_CACHE_DIR}"
mkdir -p "${BUILDX_CACHE_DIR}"
fi
if ! docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use --bootstrap
else
docker buildx use "${BUILDER_NAME}"
fi
PY_TAG="cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}"
# Output directory for wheels
DIST_DIR="dist"
mkdir -p "${DIST_DIR}"
echo "----------------------------------------"
echo "Build configuration"
echo "PYTHON_VERSION: ${PYTHON_VERSION}"
echo "CUDA_VERSION: ${CUDA_VERSION}"
echo "ARCH: ${ARCH}"
echo "BASE_IMG: ${BASE_IMG}"
echo "PYTHON_TAG: ${PY_TAG}"
echo "Output: ${DIST_DIR}/"
echo "Buildx cache: ${BUILDX_CACHE_DIR}"
echo "ccache dir: ${CCACHE_HOST_DIR}"
echo "Builder: ${BUILDER_NAME}"
echo "BUILD_JOBS: ${BUILD_JOBS:-auto}"
echo "NVCC_THREADS: ${NVCC_THREADS:-32}"
echo "USE_CCACHE: ${USE_CCACHE:-1}"
echo "RESET_BUILDER: ${RESET_BUILDER:-0}"
echo "GITHUB_ARTIFACTORY: ${GITHUB_ARTIFACTORY:-github.com}"
echo "PYTORCH_INDEX_BASE: ${PYTORCH_INDEX_BASE:-https://download.pytorch.org/whl}"
echo "PIP_DEFAULT_INDEX: ${PIP_DEFAULT_INDEX:-https://pypi.python.org/simple}"
echo "YUM_MIRROR: ${YUM_MIRROR:-(upstream)}"
echo "----------------------------------------"
# Optional build-args (empty string disables)
BUILD_ARGS=()
[ -n "${ENABLE_CMAKE_PROFILE:-}" ] && BUILD_ARGS+=(--build-arg ENABLE_CMAKE_PROFILE="${ENABLE_CMAKE_PROFILE}")
[ -n "${ENABLE_BUILD_PROFILE:-}" ] && BUILD_ARGS+=(--build-arg ENABLE_BUILD_PROFILE="${ENABLE_BUILD_PROFILE}")
[ -n "${USE_CCACHE:-}" ] && BUILD_ARGS+=(--build-arg USE_CCACHE="${USE_CCACHE}")
[ -n "${BUILD_JOBS:-}" ] && BUILD_ARGS+=(--build-arg BUILD_JOBS="${BUILD_JOBS}")
[ -n "${NVCC_THREADS:-}" ] && BUILD_ARGS+=(--build-arg NVCC_THREADS="${NVCC_THREADS}")
[ -n "${GITHUB_ARTIFACTORY:-}" ] && BUILD_ARGS+=(--build-arg GITHUB_ARTIFACTORY="${GITHUB_ARTIFACTORY}")
[ -n "${PYTORCH_INDEX_BASE:-}" ] && BUILD_ARGS+=(--build-arg PYTORCH_INDEX_BASE="${PYTORCH_INDEX_BASE}")
[ -n "${PIP_DEFAULT_INDEX:-}" ] && BUILD_ARGS+=(--build-arg PIP_DEFAULT_INDEX="${PIP_DEFAULT_INDEX}")
[ -n "${YUM_MIRROR:-}" ] && BUILD_ARGS+=(--build-arg YUM_MIRROR="${YUM_MIRROR}")
# ---- Step 1: Build deps image (layer cached, fast on repeat) ----
DEPS_TAG="sgl-kernel-deps:cuda${CUDA_VERSION}-${PY_TAG}-${ARCH}"
docker buildx build \
--builder "${BUILDER_NAME}" \
-f Dockerfile . \
--build-arg BASE_IMG="${BASE_IMG}" \
--build-arg CUDA_VERSION="${CUDA_VERSION}" \
--build-arg ARCH="${ARCH}" \
--build-arg PYTHON_VERSION="${PYTHON_VERSION}" \
--build-arg PYTHON_TAG="${PY_TAG}" \
"${BUILD_ARGS[@]}" \
--cache-from "type=local,src=${BUILDX_CACHE_DIR}" \
--cache-to "type=local,dest=${BUILDX_CACHE_DIR},mode=max" \
--target deps \
--load \
-t "${DEPS_TAG}" \
--network=host
echo "Deps image ready: ${DEPS_TAG}"
# ---- Step 2: Build wheel with host-mounted ccache ----
# This allows ccache to persist on the host filesystem across builds.
CCACHE_FLAG="${USE_CCACHE:-1}"
BUILD_JOBS_FLAG="${BUILD_JOBS:-0}"
NVCC_THREADS_FLAG="${NVCC_THREADS:-32}"
GITHUB_ARTIFACTORY_FLAG="${GITHUB_ARTIFACTORY:-github.com}"
docker run --rm \
--network=host \
-v "$(pwd):/sgl-kernel" \
-v "${CCACHE_HOST_DIR}:/ccache" \
-w /sgl-kernel \
-e ARCH="${ARCH}" \
-e GITHUB_ARTIFACTORY="${GITHUB_ARTIFACTORY_FLAG}" \
"${DEPS_TAG}" \
bash -c '
set -eux
USE_CCACHE='"${CCACHE_FLAG}"'
BUILD_JOBS='"${BUILD_JOBS_FLAG}"'
NVCC_THREADS='"${NVCC_THREADS_FLAG}"'
if [ "${USE_CCACHE}" = "1" ]; then
export CCACHE_DIR=/ccache
export CCACHE_BASEDIR=/sgl-kernel
export CCACHE_MAXSIZE=10G
export CCACHE_COMPILERCHECK=content
export CCACHE_COMPRESS=true
export CCACHE_SLOPPINESS=file_macro,time_macros,include_file_mtime,include_file_ctime
export CMAKE_C_COMPILER_LAUNCHER=ccache
export CMAKE_CXX_COMPILER_LAUNCHER=ccache
export CMAKE_CUDA_COMPILER_LAUNCHER=ccache
echo "=== ccache stats (before) ==="
ccache -sV
fi
if [ "'"${ARCH}"'" = "aarch64" ]; then
export CUDA_NVCC_FLAGS="-Xcudafe --threads=8"
export MAKEFLAGS="-j8"
export CMAKE_BUILD_PARALLEL_LEVEL=2
export NINJAFLAGS="-j4"
echo "ARM detected: Using extra conservative settings (2 parallel jobs)"
elif [ "${BUILD_JOBS}" -gt 0 ] 2>/dev/null; then
export CMAKE_BUILD_PARALLEL_LEVEL=${BUILD_JOBS}
else
export CMAKE_BUILD_PARALLEL_LEVEL=$(echo "$(( $(nproc) * 2 / 3 )) 64" | awk "{print (\$1 < \$2) ? \$1 : \$2}")
fi
export CMAKE_ARGS="${CMAKE_ARGS:-} -DSGL_KERNEL_COMPILE_THREADS=${NVCC_THREADS} -DGITHUB_ARTIFACTORY=${GITHUB_ARTIFACTORY}"
echo "Build parallelism: CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL}, NVCC_THREADS=${NVCC_THREADS}"
echo "GitHub mirror: GITHUB_ARTIFACTORY=${GITHUB_ARTIFACTORY}"
${PYTHON_ROOT_PATH}/bin/python -m uv build --wheel -Cbuild-dir=build . --color=always --no-build-isolation
PYTHON=${PYTHON_ROOT_PATH}/bin/python ./rename_wheels.sh
if [ "${USE_CCACHE}" = "1" ]; then
echo "=== ccache stats (after) ==="
ccache -s
fi
'
echo "Done. Wheels are in ${DIST_DIR}/"
ls -lh "${DIST_DIR}"/*.whl 2>/dev/null || true
python/sglang/kernels/aot/cmake/flashmla.cmake
+184
View File
@@ -0,0 +1,184 @@
# flash_mla
# sm90 dense decode HEAD_DIM_K=512 support (sgl-project/FlashMLA#9, merged).
FetchContent_Declare(
repo-flashmla
URL https://${GITHUB_ARTIFACTORY}/sgl-project/FlashMLA/archive/05e26647fe840b8baedae486c2d86d5ce4efeb7c.tar.gz
URL_HASH SHA256=ce369489bbfc42cdfbba9aa949de0270e64469d530748dea9f4f60b3c69dea9b
)
FetchContent_Populate(repo-flashmla)
# flashmla submodule pin: NVIDIA/cutlass @ 147f5673d0c1c3dcf66f78d677fd647e4a020219
FetchContent_Declare(
repo-flashmla-cutlass
URL https://${GITHUB_ARTIFACTORY}/NVIDIA/cutlass/archive/147f5673d0c1c3dcf66f78d677fd647e4a020219.tar.gz
URL_HASH SHA256=9f6c53320a85b4a570975e557918cde65168cd311f081920446c238437347dc6
SOURCE_DIR ${repo-flashmla_SOURCE_DIR}/csrc/cutlass
)
FetchContent_Populate(repo-flashmla-cutlass)
set(FLASHMLA_CUDA_FLAGS
"--expt-relaxed-constexpr"
"--expt-extended-lambda"
"--use_fast_math"
"-Xcudafe=--diag_suppress=177" # variable was declared but never referenced
)
set(FLASHMLA_ENABLE_SM100 OFF)
# The FlashMLA kernels only work on hopper and require CUDA 12.4 or later.
# Only build FlashMLA kernels if we are building for something compatible with
# sm90a
if(${CUDA_VERSION} VERSION_GREATER 12.4)
list(APPEND FLASHMLA_CUDA_FLAGS
"-gencode=arch=compute_90a,code=sm_90a"
)
endif()
if(${CUDA_VERSION} VERSION_GREATER 12.8)
list(APPEND FLASHMLA_CUDA_FLAGS
"-gencode=arch=compute_100a,code=sm_100a"
)
set(FLASHMLA_ENABLE_SM100 ON)
endif()
if(${CUDA_VERSION} VERSION_GREATER_EQUAL "13.0")
# Patch FlashMLA sources for SM103a support.
# These patches are only needed (and only valid) with CUDA 13+.
# Patch utils.h: widen IS_SM100 to cover the full SM100 family.
# Newer FlashMLA versions use csrc/utils.h.
set(FLASHMLA_UTILS_FILE "${repo-flashmla_SOURCE_DIR}/csrc/utils.h")
file(READ "${FLASHMLA_UTILS_FILE}" FLASHMLA_UTILS_CONTENT)
string(REPLACE
"#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000)
#define IS_SM100 1"
"#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100)
#define IS_SM100 1"
FLASHMLA_UTILS_CONTENT "${FLASHMLA_UTILS_CONTENT}")
file(WRITE "${FLASHMLA_UTILS_FILE}" "${FLASHMLA_UTILS_CONTENT}")
message(STATUS "Patched utils.h for SM103a support")
# Patch cutlass/arch/config.h: add SM103 architecture defines.
# The new block is inserted right before the existing "// SM101 and SM101a"
# anchor in the upstream header.
set(CUTLASS_CONFIG_FILE "${repo-flashmla_SOURCE_DIR}/csrc/cutlass/include/cutlass/arch/config.h")
file(READ "${CUTLASS_CONFIG_FILE}" CUTLASS_CONFIG_CONTENT)
string(FIND "${CUTLASS_CONFIG_CONTENT}" "SM103" SM103_FOUND)
if(SM103_FOUND EQUAL -1)
string(REPLACE
"// SM101 and SM101a"
"// SM103 and SM103a
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ >= 13)
#define CUTLASS_ARCH_MMA_SM103_SUPPORTED 1
#if (!defined(CUTLASS_ARCH_MMA_SM103_ENABLED) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030)
#define CUTLASS_ARCH_MMA_SM103_ENABLED 1
#if !defined(CUTLASS_ARCH_MMA_SM100A_ENABLED)
#define CUTLASS_ARCH_MMA_SM100A_ENABLED 1
#endif
#if !defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)
#define CUTLASS_ARCH_MMA_SM100F_ENABLED 1
#endif
#endif
#endif
/////////////////////////////////////////////////////////////////////////////////////////////////
// SM101 and SM101a"
CUTLASS_CONFIG_CONTENT "${CUTLASS_CONFIG_CONTENT}")
file(WRITE "${CUTLASS_CONFIG_FILE}" "${CUTLASS_CONFIG_CONTENT}")
message(STATUS "Patched cutlass/arch/config.h for SM103a support")
else()
message(STATUS "cutlass/arch/config.h already patched for SM103a")
endif()
list(APPEND FLASHMLA_CUDA_FLAGS
"-gencode=arch=compute_103a,code=sm_103a"
)
endif()
set(FlashMLA_SOURCES
"csrc/flashmla_extension.cc"
# Compatibility shim for sgl-kernel torch.ops API.
${repo-flashmla_SOURCE_DIR}/csrc/python_api.cpp
# Decode metadata/combine kernels.
${repo-flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
${repo-flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu
# sm90 dense decode.
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu
# sm90 sparse decode.
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu
# sm90 sparse prefill.
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/dense_fp8_python_api.cpp
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_metadata.cu
)
if(FLASHMLA_ENABLE_SM100)
list(APPEND FlashMLA_SOURCES
# sm100 dense prefill/bwd.
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu
# sm100 sparse prefill.
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
# sm100 sparse decode.
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
)
endif()
Python_add_library(flashmla_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${FlashMLA_SOURCES})
target_compile_options(flashmla_ops PRIVATE
$<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
$<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>
$<$<COMPILE_LANGUAGE:CUDA>:${FLASHMLA_CUDA_FLAGS}>
)
if(FLASHMLA_ENABLE_SM100)
target_compile_definitions(flashmla_ops PRIVATE FLASHMLA_ENABLE_SM100)
endif()
# CUDA 13 moved cuda/std/* under cccl/cuda/std/*. The vendored cutlass routes
# <cuda/std/...> to <cccl/cuda/std/...> when __CUDACC_VER_MAJOR__ >= 13, so the
# host C++ TU (compiled by g++, where that macro is unset for the legacy path)
# needs the cccl include root on the search path.
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "13.0")
find_path(FLASHMLA_CCCL_INCLUDE NAMES cuda/std/utility
HINTS ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}/cccl)
endif()
target_include_directories(flashmla_ops PRIVATE
${repo-flashmla_SOURCE_DIR}/csrc
${repo-flashmla_SOURCE_DIR}/csrc/kerutils/include
${repo-flashmla_SOURCE_DIR}/csrc/sm90
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/
${repo-flashmla_SOURCE_DIR}/csrc/cutlass/include
${repo-flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
${FLASHMLA_CCCL_INCLUDE}
)
target_link_libraries(flashmla_ops PRIVATE ${TORCH_LIBRARIES} c10 cuda)
install(TARGETS flashmla_ops LIBRARY DESTINATION "sgl_kernel")
target_compile_definitions(flashmla_ops PRIVATE)
python/sglang/kernels/aot/cmake/utils.cmake
+19
View File
@@ -0,0 +1,19 @@
# Adapt from: https://github.com/neuralmagic/vllm-flash-attention/blob/main/cmake/utils.cmake
#
# Clear all `-gencode` flags from `CMAKE_CUDA_FLAGS` and store them in
# `CUDA_ARCH_FLAGS`.
#
# Example:
# CMAKE_CUDA_FLAGS="-Wall -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75"
# clear_cuda_arches(CUDA_ARCH_FLAGS)
# CUDA_ARCH_FLAGS="-gencode arch=compute_70,code=sm_70;-gencode arch=compute_75,code=sm_75"
# CMAKE_CUDA_FLAGS="-Wall"
#
macro(clear_cuda_arches CUDA_ARCH_FLAGS)
# Extract all `-gencode` flags from `CMAKE_CUDA_FLAGS`
string(REGEX MATCHALL "-gencode arch=[^ ]+" CUDA_ARCH_FLAGS "${CMAKE_CUDA_FLAGS}")
# Remove all `-gencode` flags from `CMAKE_CUDA_FLAGS` since they will be modified
# and passed back via the `CUDA_ARCHITECTURES` property.
string(REGEX REPLACE "-gencode arch=[^ ]+ *" "" CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
endmacro()
python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce.cu
+137
View File
@@ -0,0 +1,137 @@
// Adapted from: https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cu
#include <ATen/cuda/Exceptions.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/all.h>
#include "custom_all_reduce.cuh"
// Fake pointer type, must match fptr_t type in ops.h.
// We use this type alias to indicate when pointers are passed in as int64_t.
using fptr_t = int64_t;
static_assert(sizeof(void*) == sizeof(fptr_t));
fptr_t
init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs, torch::Tensor& rank_data, int64_t rank, bool full_nvlink) {
int world_size = fake_ipc_ptrs.size();
if (world_size > 8) throw std::invalid_argument("world size > 8 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");
sglang::Signal* ipc_ptrs[8];
for (int i = 0; i < world_size; i++) {
ipc_ptrs[i] = reinterpret_cast<sglang::Signal*>(fake_ipc_ptrs[i]);
}
return (fptr_t) new sglang::CustomAllreduce(
ipc_ptrs, rank_data.data_ptr(), rank_data.numel(), rank, world_size, full_nvlink);
}
/**
* Make sure tensor t's data lies completely within ((char)t.data_ptr()) +
* t.numel() * t.element_size(). This is slightly weaker than t.is_contiguous()
* because it allows transpose of contiguous slice (i.e. slicing the first
* dimension). Currently, we require this because stride information is not
* passed into the kernels and we treat input tensors as flat.
*
* Examples
* A = torch.zeros(3, 3, 3)
* 1. A: OK
* 2. A[1:]: OK
* 3. A.permute(2, 0, 1): OK
* 4. A[1:].permute(2, 0, 1): OK
* 5. A[None].expand(2, -1, -1, -1): Not OK
* 6. A[:, 1:, 1:]: Not OK
*/
bool _is_weak_contiguous(torch::Tensor& t) {
return t.is_contiguous() ||
(t.storage().nbytes() - t.storage_offset() * t.element_size() == t.numel() * t.element_size());
}
/**
* Performs an out-of-place allreduce and stores result in out.
*
* If _reg_buffer is null, assumes inp.data_ptr() is already IPC-registered.
* Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
* copied into _reg_buffer.
*/
void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
auto stream = c10::cuda::getCurrentCUDAStream().stream();
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
TORCH_CHECK_EQ(inp.numel(), out.numel());
TORCH_CHECK(_is_weak_contiguous(out));
TORCH_CHECK(_is_weak_contiguous(inp));
auto input_size = inp.numel() * inp.element_size();
auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
if (reg_buffer) {
TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes);
AT_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.data_ptr(), input_size, cudaMemcpyDeviceToDevice, stream));
} else {
reg_buffer = inp.data_ptr();
}
switch (out.scalar_type()) {
case at::ScalarType::Float: {
fa->allreduce<float>(
stream, reinterpret_cast<float*>(reg_buffer), reinterpret_cast<float*>(out.data_ptr()), out.numel());
break;
}
case at::ScalarType::Half: {
fa->allreduce<half>(
stream, reinterpret_cast<half*>(reg_buffer), reinterpret_cast<half*>(out.data_ptr()), out.numel());
break;
}
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
case at::ScalarType::BFloat16: {
fa->allreduce<nv_bfloat16>(
stream,
reinterpret_cast<nv_bfloat16*>(reg_buffer),
reinterpret_cast<nv_bfloat16*>(out.data_ptr()),
out.numel());
break;
}
#endif
default:
throw std::runtime_error("custom allreduce only supports float32, float16 and bfloat16");
}
}
void dispose(fptr_t _fa) {
delete reinterpret_cast<sglang::CustomAllreduce*>(_fa);
}
int64_t meta_size() {
return sizeof(sglang::Signal);
}
void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
void* ipc_ptrs[8];
for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
}
fa->register_buffer(ipc_ptrs);
}
// Use vector<int64_t> to represent byte data for python binding compatibility.
std::tuple<std::vector<int64_t>, std::vector<int64_t>> get_graph_buffer_ipc_meta(fptr_t _fa) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
auto [handle, offsets] = fa->get_graph_buffer_ipc_meta();
std::vector<int64_t> bytes(handle.begin(), handle.end());
return std::make_tuple(bytes, offsets);
}
// Use vector<int64_t> to represent byte data for python binding compatibility.
void register_graph_buffers(
fptr_t _fa, const std::vector<std::vector<int64_t>>& handles, const std::vector<std::vector<int64_t>>& offsets) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
std::vector<std::string> bytes;
bytes.reserve(handles.size());
for (int i = 0; i < handles.size(); i++) {
bytes.emplace_back(handles[i].begin(), handles[i].end());
}
bytes.reserve(handles.size());
fa->register_graph_buffers(bytes, offsets);
}
python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce.cuh
+693
View File
@@ -0,0 +1,693 @@
// Adapted from https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cuh
#pragma once
#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <array>
#include <iostream>
#include <limits>
#include <map>
#include <unordered_map>
#include <vector>
#include "utils.h"
namespace sglang {
#ifndef USE_MUSA
constexpr int kMaxBlocks = 36;
constexpr int kDefaultThreads = 512;
constexpr int kDefaultBlockLimit = 36;
constexpr int kMaxThreadsPerBlock = 512;
#else
constexpr int kMaxBlocks = 60;
constexpr int kDefaultThreads = 1024;
constexpr int kDefaultBlockLimit = 60;
constexpr int kMaxThreadsPerBlock = 1024;
#endif
// Allreduce algorithm selection thresholds
constexpr int kAllReduceGPUSmall = 4;
constexpr int kAllReduceGPULarge = 8;
constexpr size_t kAllReduceSmallThreshold = 512 * 1024; // 512KB
constexpr size_t kAllReduceLargeThreshold = 256 * 1024; // 256KB
// Counter may overflow, but it's fine since unsigned int overflow is
// well-defined behavior.
using FlagType = uint32_t;
struct Signal {
alignas(128) FlagType self_counter[kMaxBlocks][8];
// Two sets of peer counters are needed for two syncs. The reason is that
// it's possible for peer GPU block to arrive at the second sync point while
// the current GPU block haven't passed the first sync point. Thus, peer GPU
// may write counter+1 while current GPU is busy waiting for counter. We use
// alternating counter array to avoid this possibility.
alignas(128) FlagType peer_counter[2][kMaxBlocks][8];
};
struct __align__(16) RankData {
const void* __restrict__ ptrs[8];
};
struct __align__(16) RankSignals {
Signal* signals[8];
};
// like std::array, but aligned
template <typename T, int sz>
struct __align__(alignof(T) * sz) array_t {
T data[sz];
using type = T;
static constexpr int size = sz;
};
// use packed type to maximize memory efficiency
// goal: generate ld.128 and st.128 instructions
template <typename T>
struct packed_t {
// the (P)acked type for load/store
using P = array_t<T, 16 / sizeof(T)>;
// the (A)ccumulator type for reduction
using A = array_t<float, 16 / sizeof(T)>;
};
#define DINLINE __device__ __forceinline__
// scalar cast functions
DINLINE float upcast_s(half val) {
return __half2float(val);
}
template <typename T>
DINLINE T downcast_s(float val);
template <>
DINLINE half downcast_s(float val) {
return __float2half(val);
}
// scalar add functions
// for some reason when compiling with Pytorch, the + operator for half and
// bfloat is disabled so we call the intrinsics directly
DINLINE half& assign_add(half& a, half b) {
a = __hadd(a, b);
return a;
}
DINLINE float& assign_add(float& a, float b) {
return a += b;
}
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
DINLINE float upcast_s(nv_bfloat16 val) {
return __bfloat162float(val);
}
template <>
DINLINE nv_bfloat16 downcast_s(float val) {
return __float2bfloat16(val);
}
DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
a = __hadd(a, b);
return a;
}
#endif
template <typename T, int N>
DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
#pragma unroll
for (int i = 0; i < N; i++) {
assign_add(a.data[i], b.data[i]);
}
return a;
}
template <typename T, int N>
DINLINE array_t<float, N> upcast(array_t<T, N> val) {
if constexpr (std::is_same<T, float>::value) {
return val;
} else {
array_t<float, N> out;
#pragma unroll
for (int i = 0; i < N; i++) {
out.data[i] = upcast_s(val.data[i]);
}
return out;
}
}
template <typename O>
DINLINE O downcast(array_t<float, O::size> val) {
if constexpr (std::is_same<typename O::type, float>::value) {
return val;
} else {
O out;
#pragma unroll
for (int i = 0; i < O::size; i++) {
out.data[i] = downcast_s<typename O::type>(val.data[i]);
}
return out;
}
}
static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
#ifdef USE_MUSA
volatile_store((uint32_t)flag, (uint32_t*)flag_addr);
#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
#else
asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
#endif
}
static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
#ifdef USE_MUSA
flushInv_byp();
return (uint32_t)volatile_load((uint32_t*)flag_addr);
#endif
FlagType flag;
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(flag) : "l"(flag_addr));
#else
asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" : "=r"(flag) : "l"(flag_addr));
#endif
return flag;
}
static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
}
static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
FlagType flag;
asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(flag) : "l"(flag_addr));
return flag;
}
// is_start: whether this is the very first synchronization barrier.
// need_fence: whether a memory fence is needed. If true, a release-acquire
// semantic is used to enforce memory access order before and after this
// barrier.
template <int ngpus, bool is_start, bool need_fence = false>
DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg, int rank) {
if constexpr (!is_start)
#ifndef USE_MUSA
__syncthreads();
#else
__syncthreads_lm();
#endif
static_assert(!(is_start && need_fence)); // Start barrier shouldn't need fence.
if (threadIdx.x < ngpus) {
// Increment the counter. Technically we only need one counter, but we use
// multiple per block to eliminate the need to share the counter via smem.
auto val = self_sg->self_counter[blockIdx.x][threadIdx.x] += 1;
// Write the expected counter value to peer and wait for correct value from
// peer.
auto peer_counter_ptr = &sg.signals[threadIdx.x]->peer_counter[val % 2][blockIdx.x][rank];
auto self_counter_ptr = &self_sg->peer_counter[val % 2][blockIdx.x][threadIdx.x];
if constexpr (need_fence) {
st_flag_release(peer_counter_ptr, val);
while (ld_flag_acquire(self_counter_ptr) != val)
;
} else {
st_flag_volatile(peer_counter_ptr, val);
while (ld_flag_volatile(self_counter_ptr) != val)
;
}
}
if constexpr (is_start || need_fence)
#ifndef USE_MUSA
__syncthreads();
#else
__syncthreads_lm();
#endif
}
template <typename P, int ngpus, typename A>
DINLINE P packed_reduce(const P* ptrs[], int idx) {
A tmp = upcast(ptrs[0][idx]);
#pragma unroll
for (int i = 1; i < ngpus; i++) {
packed_assign_add(tmp, upcast(ptrs[i][idx]));
}
return downcast<P>(tmp);
}
template <typename T, int ngpus>
__global__ void __launch_bounds__(kMaxThreadsPerBlock, 1) cross_device_reduce_1stage(
RankData* _dp, RankSignals sg, Signal* self_sg, T* __restrict__ result, int rank, int size) {
using P = typename packed_t<T>::P;
using A = typename packed_t<T>::A;
// note: we don't reorder the address so the accumulation order is the same
// for all ranks, ensuring bitwise identical results
auto dp = *_dp;
multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
// do the actual reduction
for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; idx += gridDim.x * blockDim.x) {
((P*)result)[idx] = packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], idx);
}
multi_gpu_barrier<ngpus, false>(sg, self_sg, rank);
}
template <typename P>
DINLINE P* get_tmp_buf(Signal* sg) {
return (P*)(((Signal*)sg) + 1);
}
#ifdef USE_MUSA
template <typename T, int32_t nranks, int32_t vlen = 8>
DINLINE void shfl_reduce(float* res) {
if constexpr (nranks >= 4) {
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
res[i] += __shfl_xor_sync(0xffffffff, res[i], 16);
}
}
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
res[i] += __shfl_xor_sync(0xffffffff, res[i], 8);
}
}
template <typename T, int32_t nranks, int32_t vlen = 8>
__global__ void __launch_bounds__(kMaxThreadsPerBlock, 1) custom_all_reduce_2shot(
RankData* _dp, RankSignals sg, Signal* self_sg, T* __restrict__ result, int32_t local_rank, int32_t size) {
constexpr int32_t nranks_sft = (nranks >> 1) - (nranks >> 3); // 8->3, 4->2, 2->1
constexpr int32_t coalesce_num = 8;
constexpr int32_t coalesce_sft = 3; // 8 threads per rank in group
constexpr int32_t group_size = nranks << coalesce_sft; // tp 8 -> 64 threads, tp 4 -> 32 threads, tp 2 -> 16 threads
constexpr int32_t group_stride_sft = nranks_sft + coalesce_sft;
const int32_t tidx = threadIdx.x;
const int32_t bidx = blockIdx.x;
const int32_t thread_num = blockDim.x;
const int32_t lane_idx = tidx & 31;
const int32_t warp_idx = tidx >> 5;
const int32_t group_num = thread_num >> group_stride_sft;
const int32_t target_rank = (tidx >> coalesce_sft) & (nranks - 1);
const int32_t group_id = tidx >> group_stride_sft;
const int32_t coalesce_tid = tidx & (coalesce_num - 1);
typedef int16_t Vec __attribute__((vector_size(16)));
const int32_t stride = gridDim.x * thread_num;
int32_t idx_base = bidx * thread_num;
int32_t idx_in_blk = coalesce_tid + (local_rank << coalesce_sft) + (group_id << group_stride_sft);
// first sync barrier
FlagType* target_barrier = nullptr;
FlagType* local_barrier = nullptr;
FlagType flag;
if (tidx < nranks) {
flag = atomicAdd(&(self_sg->self_counter[bidx][tidx]), 1);
target_barrier = &sg.signals[tidx]->peer_counter[flag & 1][bidx][local_rank];
local_barrier = &self_sg->peer_counter[flag & 1][bidx][tidx];
atomicExch(target_barrier, flag);
while (atomicAdd(local_barrier, 0) != flag) {
}
}
__syncthreads_lm();
// reduce scatter
Vec* target_ptr = (Vec*)_dp->ptrs[target_rank];
Vec* buffer_ptr = get_tmp_buf<Vec>(sg.signals[local_rank]);
do {
int32_t idx = idx_in_blk + idx_base;
float temp_res[vlen] = {0};
if (idx < size) {
T* data = reinterpret_cast<T*>(&(target_ptr[idx]));
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
temp_res[i] = upcast_s(data[i]);
}
}
shfl_reduce<T, nranks, vlen>(temp_res);
// reduce cross warp, only trigger when tp 8
if constexpr (nranks == 8) {
__shared__ float smem[kMaxThreadsPerBlock << 1];
if (lane_idx < coalesce_num) {
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
smem[warp_idx * vlen * coalesce_num + coalesce_tid * vlen + i] = temp_res[i];
}
}
__syncthreads_lm();
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
temp_res[i] += smem[(warp_idx ^ 1) * vlen * coalesce_num + coalesce_tid * vlen + i];
}
}
if (local_rank == target_rank && idx < size) {
Vec res;
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
reinterpret_cast<T*>(&res)[i] = downcast_s<T>(temp_res[i]);
}
buffer_ptr[idx] = res;
}
idx_base += stride;
} while (idx_base < size);
// make sure buffer_ptr data ready
__musa_barrier_slc();
__syncthreads_lm();
if (tidx == 0) {
__threadfence_system_noflush();
}
buffer_ptr = get_tmp_buf<Vec>(sg.signals[target_rank]);
// second sync barrier
if (tidx < nranks) {
flag = atomicAdd(&(self_sg->self_counter[bidx][tidx]), 1);
target_barrier = &sg.signals[tidx]->peer_counter[flag & 1][bidx][local_rank];
local_barrier = &self_sg->peer_counter[flag & 1][bidx][tidx];
atomicExch(target_barrier, flag);
while (atomicAdd(local_barrier, 0) != flag) {
}
}
__syncthreads_lm();
// all gather
idx_in_blk = coalesce_tid + (target_rank << coalesce_sft) + (group_id << group_stride_sft);
idx_base = bidx * thread_num;
do {
int32_t idx = idx_in_blk + idx_base;
if (idx < size) {
reinterpret_cast<Vec*>(result)[idx] = buffer_ptr[idx];
}
idx_base += stride;
} while (idx_base < size);
}
#endif // USE_MUSA
template <typename T, int ngpus>
__global__ void __launch_bounds__(kMaxThreadsPerBlock, 1) cross_device_reduce_2stage(
RankData* _dp, RankSignals sg, Signal* self_sg, T* __restrict__ result, int rank, int size) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int stride = gridDim.x * blockDim.x;
using P = typename packed_t<T>::P;
using A = typename packed_t<T>::A;
int part = size / ngpus;
int start = rank * part;
int end = rank == ngpus - 1 ? size : start + part;
int largest_part = part + size % ngpus;
const P* ptrs[ngpus];
P* tmps[ngpus];
#pragma unroll
for (int i = 0; i < ngpus; i++) {
int target = (rank + i) % ngpus;
ptrs[i] = (const P*)_dp->ptrs[target];
tmps[i] = get_tmp_buf<P>(sg.signals[target]);
}
auto tmp_out = tmps[0];
multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
// stage 1: reduce scatter
for (int idx = start + tid; idx < end; idx += stride) {
tmp_out[idx - start] = packed_reduce<P, ngpus, A>(ptrs, idx);
}
multi_gpu_barrier<ngpus, false, true>(sg, self_sg, rank);
// stage 2: allgather. Note: it's important to match the tid between
// the two stages, because visibility across devices is only guaranteed
// between threads that have the same tid. If thread i computes the sum of
// start + i in the first stage, then thread i also gathers start + i from all
// ranks.
for (int idx = tid; idx < largest_part; idx += stride) {
#pragma unroll
for (int i = 0; i < ngpus; i++) {
int gather_from_rank = ((rank + i) % ngpus);
if (gather_from_rank == ngpus - 1 || idx < part) {
int dst_idx = gather_from_rank * part + idx;
((P*)result)[dst_idx] = tmps[i][idx];
}
}
}
}
using IPC_KEY = std::array<uint8_t, sizeof(cudaIpcMemHandle_t)>;
static_assert(sizeof(IPC_KEY) == sizeof(cudaIpcMemHandle_t));
static_assert(alignof(IPC_KEY) == alignof(cudaIpcMemHandle_t));
class CustomAllreduce {
public:
int rank_;
int world_size_;
bool full_nvlink_;
RankSignals sg_;
// Stores an map from a pointer to its peer pointters from all ranks.
std::unordered_map<void*, RankData*> buffers_;
Signal* self_sg_;
// Stores rank data from all ranks. This is mainly for cuda graph purposes.
// For cuda graph to work, all kernel arguments must be fixed during graph
// capture time. However, the peer pointers are not known during graph capture
// time. Therefore, during capture, we increment the rank data pointer and use
// that as the argument to the kernel. The kernel arguments are stored in
// graph_unreg_buffers_. The actual peer pointers will be filled in at the
// memory pointed to by the pointers in graph_unreg_buffers_ when
// the IPC handles are exchanged between ranks.
//
// The overall process looks like this:
// 1. Graph capture.
// 2. Each rank obtains the IPC handles for each addresses used during cuda
// graph capture using get_graph_buffer_ipc_meta.
// 3. (In Python) all gather the IPC handles.
// 4. Obtain the peer pointers by opening the IPC handles, and store them in
// the rank data array at corresponding positions.
RankData *d_rank_data_base_, *d_rank_data_end_;
std::vector<void*> graph_unreg_buffers_;
// a map from IPC handles to opened IPC pointers
std::map<IPC_KEY, char*> ipc_handles_;
/**
* Signals are an array of ipc-enabled buffers from all ranks.
* For each of the buffer, the layout is as follows:
* | -- sizeof(Signal) -- | ------ a few MB ----- |
* The first section is for allreduce synchronization, and the second section
* is for storing the intermediate results required by some allreduce algos.
*
* Note: this class does not own any device memory. Any required buffers
* are passed in from the constructor.
*/
CustomAllreduce(
Signal** signals, void* rank_data, size_t rank_data_sz, int rank, int world_size, bool full_nvlink = true)
: rank_(rank),
world_size_(world_size),
full_nvlink_(full_nvlink),
self_sg_(signals[rank]),
d_rank_data_base_(reinterpret_cast<RankData*>(rank_data)),
d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) {
for (int i = 0; i < world_size_; i++) {
sg_.signals[i] = signals[i];
}
}
char* open_ipc_handle(const void* ipc_handle) {
auto [it, new_handle] = ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr});
if (new_handle) {
char* ipc_ptr;
CHECK_CUDA_SUCCESS(cudaIpcOpenMemHandle(
(void**)&ipc_ptr, *((const cudaIpcMemHandle_t*)ipc_handle), cudaIpcMemLazyEnablePeerAccess));
it->second = ipc_ptr;
}
return it->second;
}
std::pair<std::string, std::vector<int64_t>> get_graph_buffer_ipc_meta() {
auto num_buffers = graph_unreg_buffers_.size();
auto handle_sz = sizeof(cudaIpcMemHandle_t);
std::string handles(handle_sz * num_buffers, static_cast<char>(0));
std::vector<int64_t> offsets(num_buffers);
for (int i = 0; i < num_buffers; i++) {
auto ptr = graph_unreg_buffers_[i];
void* base_ptr;
// note: must share the base address of each allocation, or we get wrong
// address
if (cuPointerGetAttribute(&base_ptr, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, (CUdeviceptr)ptr) != CUDA_SUCCESS)
throw std::runtime_error("failed to get pointer attr");
CHECK_CUDA_SUCCESS(cudaIpcGetMemHandle((cudaIpcMemHandle_t*)&handles[i * handle_sz], base_ptr));
offsets[i] = ((char*)ptr) - ((char*)base_ptr);
}
return std::make_pair(handles, offsets);
}
void check_rank_data_capacity(size_t num = 1) {
if (d_rank_data_base_ + num > d_rank_data_end_)
throw std::runtime_error(
"Rank data buffer is overflowed by " + std::to_string(d_rank_data_base_ + num - d_rank_data_end_));
}
/**
* Register already-shared IPC pointers.
*/
void register_buffer(void** ptrs) {
check_rank_data_capacity();
RankData data;
for (int i = 0; i < world_size_; i++) {
data.ptrs[i] = ptrs[i];
}
auto d_data = d_rank_data_base_++;
CHECK_CUDA_SUCCESS(cudaMemcpy(d_data, &data, sizeof(RankData), cudaMemcpyHostToDevice));
buffers_[ptrs[rank_]] = d_data;
}
// Note: when registering graph buffers, we intentionally choose to not
// deduplicate the addresses. That means if the allocator reuses some
// addresses, they will be registered again. This is to account for the remote
// possibility of different allocation patterns between ranks. For example,
// rank 1 may get the same input address for the second allreduce, but rank 2
// got a different address. IPC handles have internal reference counting
// mechanism so overhead should be small.
void
register_graph_buffers(const std::vector<std::string>& handles, const std::vector<std::vector<int64_t>>& offsets) {
auto num_buffers = graph_unreg_buffers_.size();
check_rank_data_capacity(num_buffers);
std::vector<RankData> rank_data(num_buffers);
for (int i = 0; i < num_buffers; i++) {
auto self_ptr = graph_unreg_buffers_[i];
auto& rd = rank_data[i];
for (int j = 0; j < world_size_; j++) {
if (j != rank_) {
char* handle = open_ipc_handle(&handles[j][i * sizeof(cudaIpcMemHandle_t)]);
handle += offsets[j][i];
rd.ptrs[j] = handle;
} else {
rd.ptrs[j] = self_ptr;
}
}
}
CHECK_CUDA_SUCCESS(
cudaMemcpy(d_rank_data_base_, rank_data.data(), sizeof(RankData) * num_buffers, cudaMemcpyHostToDevice));
d_rank_data_base_ += num_buffers;
graph_unreg_buffers_.clear();
}
/**
* Performs allreduce, assuming input has already been registered.
*
* Block and grid default configs are results after careful grid search. Using
* 36 blocks give the best or close to the best runtime on the devices I
* tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also only
* take a small amount of SMs. Not quite sure the underlying reason, but my
* guess is that too many SMs will cause contention on NVLink bus.
*/
template <typename T>
void allreduce(
cudaStream_t stream,
T* input,
T* output,
int size,
int threads = kDefaultThreads,
int block_limit = kDefaultBlockLimit) {
auto d = packed_t<T>::P::size;
if (size % d != 0)
throw std::runtime_error(
"custom allreduce currently requires input length to be multiple "
"of " +
std::to_string(d));
if (block_limit > kMaxBlocks)
throw std::runtime_error(
"max supported block limit is " + std::to_string(kMaxBlocks) + ". Got " + std::to_string(block_limit));
RankData* ptrs;
cudaStreamCaptureStatus status;
CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &status));
if (status == cudaStreamCaptureStatusActive) {
ptrs = d_rank_data_base_ + graph_unreg_buffers_.size();
graph_unreg_buffers_.push_back(input);
} else {
auto it = buffers_.find(input);
if (it == buffers_.end())
throw std::runtime_error(
"buffer address " + std::to_string(reinterpret_cast<uint64_t>(input)) + " is not registered!");
ptrs = it->second;
}
size /= d;
auto bytes = size * sizeof(typename packed_t<T>::P);
int blocks = std::min(block_limit, (size + threads - 1) / threads);
// Check environment variable once
const char* env_algo = std::getenv("SGLANG_CUSTOM_ALLREDUCE_ALGO");
bool force_1stage = false;
bool force_2stage = false;
if (env_algo != nullptr) {
if (std::strcmp(env_algo, "1stage") == 0 || std::strcmp(env_algo, "oneshot") == 0) {
force_1stage = true;
} else if (std::strcmp(env_algo, "2stage") == 0 || std::strcmp(env_algo, "twoshot") == 0) {
force_2stage = true;
} else {
throw std::runtime_error(
"Invalid SGLANG_CUSTOM_ALLREDUCE_ALGO: " + std::string(env_algo) +
". Valid values: 1stage, oneshot, 2stage, twoshot");
}
}
#define KL(ngpus, name) name<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, rank_, size);
// TODO(hanzhi713): Threshold is different for A100 and H100.
// Add per device threshold.
#ifndef USE_MUSA
#define REDUCE_CASE(ngpus) \
case ngpus: { \
if (force_1stage) { \
KL(ngpus, cross_device_reduce_1stage); \
} else if (force_2stage) { \
KL(ngpus, cross_device_reduce_2stage); \
} else { \
if (world_size_ == 2) { \
KL(ngpus, cross_device_reduce_1stage); \
} else if (full_nvlink_) { \
if ((world_size_ <= kAllReduceGPUSmall && bytes < kAllReduceSmallThreshold) || \
(world_size_ <= kAllReduceGPULarge && bytes < kAllReduceLargeThreshold)) { \
KL(ngpus, cross_device_reduce_1stage); \
} else { \
KL(ngpus, cross_device_reduce_2stage); \
} \
} \
} \
break; \
}
#else
#define REDUCE_CASE(ngpus) \
case ngpus: { \
if constexpr (!std::is_same<T, float>::value) { \
custom_all_reduce_2shot<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, rank_, size); \
} else { \
if ((world_size_ <= kAllReduceGPUSmall && bytes < kAllReduceSmallThreshold) || \
(world_size_ <= kAllReduceGPULarge && bytes < kAllReduceLargeThreshold)) { \
KL(ngpus, cross_device_reduce_1stage); \
} else { \
KL(ngpus, cross_device_reduce_2stage); \
} \
} \
break; \
}
#endif
switch (world_size_) {
REDUCE_CASE(2)
REDUCE_CASE(4)
REDUCE_CASE(6)
REDUCE_CASE(8)
default:
throw std::runtime_error(
"custom allreduce only supports num gpus in (2,4,6,8). Actual num "
"gpus = " +
std::to_string(world_size_));
}
#undef REDUCE_CASE
#undef KL
}
~CustomAllreduce() {
for (auto [_, ptr] : ipc_handles_) {
CHECK_CUDA_SUCCESS(cudaIpcCloseMemHandle(ptr));
}
}
};
/**
* To inspect PTX/SASS, copy paste this header file to compiler explorer and add
a template instantiation:
* template void sglang::CustomAllreduce::allreduce<half>(cudaStream_t, half *,
half *, int, int, int);
*/
} // namespace sglang
python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce.hip
+180
View File
@@ -0,0 +1,180 @@
// !!! This is a file automatically generated by hipify!!!
#include <ATen/hip/Exceptions.h>
#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
#include <torch/all.h>
#include "custom_all_reduce_hip.cuh"
// fake pointer type, must match fptr_t type in ops.h
using fptr_t = int64_t;
static_assert(sizeof(void*) == sizeof(fptr_t));
fptr_t init_custom_ar(torch::Tensor& meta, torch::Tensor& rank_data,
const std::vector<std::string>& handles,
const std::vector<int64_t>& offsets, int64_t rank,
bool full_nvlink) {
int world_size = offsets.size();
if (world_size > 8)
throw std::invalid_argument("world size > 8 is not supported");
if (world_size % 2 != 0)
throw std::invalid_argument("Odd num gpus is not supported for now");
if (world_size != handles.size())
throw std::invalid_argument(
"handles length should equal to offsets length");
if (rank < 0 || rank >= world_size)
throw std::invalid_argument("invalid rank passed in");
hipIpcMemHandle_t ipc_handles[8];
for (int i = 0; i < world_size; i++) {
std::memcpy(&ipc_handles[i], handles[i].data(), sizeof(hipIpcMemHandle_t));
}
return (fptr_t) new sglang::CustomAllreduce(
reinterpret_cast<sglang::Signal*>(meta.data_ptr()), rank_data.data_ptr(),
rank_data.numel(), ipc_handles, offsets, rank, full_nvlink);
}
/**
* Make sure tensor t's data lies completely within ((char)t.data_ptr()) +
* t.numel() * t.element_size(). This is slightly weaker than t.is_contiguous()
* because it allows transpose of contiguous slice (i.e. slicing the first
* dimension). Currently, we require this because stride information is not
* passed into the kernels and we treat input tensors as flat.
*
* Examples
* A = torch.zeros(3, 3, 3)
* 1. A: OK
* 2. A[1:]: OK
* 3. A.permute(2, 0, 1): OK
* 4. A[1:].permute(2, 0, 1): OK
* 5. A[None].expand(2, -1, -1, -1): Not OK
* 6. A[:, 1:, 1:]: Not OK
*/
bool _is_weak_contiguous(torch::Tensor& t) {
return t.is_contiguous() ||
(t.storage().nbytes() - t.storage_offset() * t.element_size() ==
t.numel() * t.element_size());
}
void _all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
hipStream_t stream) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
TORCH_CHECK(_is_weak_contiguous(out));
switch (out.scalar_type()) {
case at::ScalarType::Float: {
fa->allreduce<float>(stream, reinterpret_cast<float*>(inp.data_ptr()),
reinterpret_cast<float*>(out.data_ptr()),
out.numel());
break;
}
case at::ScalarType::Half: {
fa->allreduce<half>(stream, reinterpret_cast<half*>(inp.data_ptr()),
reinterpret_cast<half*>(out.data_ptr()), out.numel());
break;
}
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
case at::ScalarType::BFloat16: {
fa->allreduce<nv_bfloat16>(
stream, reinterpret_cast<nv_bfloat16*>(inp.data_ptr()),
reinterpret_cast<nv_bfloat16*>(out.data_ptr()), out.numel());
break;
}
#endif
default:
throw std::runtime_error(
"custom allreduce only supports float32, float16 and bfloat16");
}
}
void all_reduce_reg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out) {
const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp));
auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
TORCH_CHECK_EQ(inp.numel(), out.numel());
_all_reduce(_fa, inp, out, stream);
}
void all_reduce_unreg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer,
torch::Tensor& out) {
const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp));
auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
auto input_size = inp.numel() * inp.element_size();
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
TORCH_CHECK_EQ(inp.numel(), out.numel());
TORCH_CHECK(input_size <= reg_buffer.numel() * reg_buffer.element_size(),
"registered buffer is too small to contain the input");
AT_CUDA_CHECK(hipMemcpyAsync(reg_buffer.data_ptr(), inp.data_ptr(),
input_size, hipMemcpyDeviceToDevice, stream));
_all_reduce(_fa, reg_buffer, out, stream);
}
void dispose(fptr_t _fa) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
delete fa;
}
int64_t meta_size() { return sizeof(sglang::Signal); }
void register_buffer(fptr_t _fa, torch::Tensor& t,
const std::vector<std::string>& handles,
const std::vector<int64_t>& offsets) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
fa->register_buffer(handles, offsets, t.data_ptr());
}
std::tuple<torch::Tensor, std::vector<int64_t>> get_graph_buffer_ipc_meta(
fptr_t _fa) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
auto [handle_bytes, offsets] = fa->get_graph_buffer_ipc_meta();
auto options =
torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
auto handles =
torch::empty({static_cast<int64_t>(handle_bytes.size())}, options);
std::memcpy(handles.data_ptr(), handle_bytes.data(), handle_bytes.size());
return {handles, std::move(offsets)};
}
void register_graph_buffers(fptr_t _fa, const std::vector<std::string>& handles,
const std::vector<std::vector<int64_t>>& offsets) {
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
fa->register_graph_buffers(handles, offsets);
}
void free_meta_buffer(void* buffer) { CUDACHECK(hipFree(buffer)); }
torch::Tensor get_meta_buffer_ipc_handle(torch::Tensor& inp) {
auto options =
torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
auto data_handle =
torch::empty({static_cast<int64_t>(sizeof(hipIpcMemHandle_t))}, options);
CUDACHECK(hipIpcGetMemHandle((hipIpcMemHandle_t*)data_handle.data_ptr(),
inp.data_ptr()));
return data_handle;
}
torch::Tensor allocate_meta_buffer(int64_t size) {
auto device_index = c10::hip::current_device();
at::DeviceGuard device_guard(at::Device(at::DeviceType::CUDA, device_index));
void* buffer;
hipStreamCaptureMode mode = hipStreamCaptureModeRelaxed;
auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
AT_CUDA_CHECK(hipThreadExchangeStreamCaptureMode(&mode));
AT_CUDA_CHECK(
hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached));
AT_CUDA_CHECK(hipMemsetAsync(buffer, 0, size, stream));
AT_CUDA_CHECK(hipStreamSynchronize(stream));
AT_CUDA_CHECK(hipThreadExchangeStreamCaptureMode(&mode));
auto options = torch::TensorOptions()
.dtype(torch::kI8)
.device(torch::kCUDA, device_index);
return torch::from_blob(buffer, {size}, free_meta_buffer, options);
}
std::vector<uint8_t> get_device_bdf(int dev) {
char busIdStr[] = "0000:00:00.0";
std::vector<uint8_t> bdf(sizeof(busIdStr), 0);
CUDACHECK(hipDeviceGetPCIBusId((char*)bdf.data(), sizeof(busIdStr), dev));
bdf.resize(bdf.size() - 1); // remove trailing NULL
return bdf;
}
python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce_hip.cuh
+582
View File
@@ -0,0 +1,582 @@
// !!! This is a file automatically generated by hipify!!!
#pragma once
#include <hip/hip_runtime.h>
#ifdef USE_ROCM
#include <hip/hip_bf16.h>
typedef __hip_bfloat16 nv_bfloat16;
#else
#include <hip/hip_bf16.h>
#endif
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
#include <iostream>
#include <limits>
#include <map>
#include <unordered_map>
#include <vector>
#define CUDACHECK(cmd) \
do { \
hipError_t e = cmd; \
if (e != hipSuccess) { \
printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, hipGetErrorString(e)); \
exit(EXIT_FAILURE); \
} \
} while (0)
namespace sglang {
constexpr int kMaxBlocks = 64;
// note: we don't want to use atomics for signals because peer atomics are no
// supported on PCIe links
struct Signal {
alignas(128) uint32_t start[kMaxBlocks][8];
alignas(128) uint32_t end[kMaxBlocks][8];
alignas(128) uint32_t _flag[kMaxBlocks]; // incremental flags for each rank
};
#ifdef USE_ROCM
struct __align__(16) RankData {
const void* ptrs[8];
};
#else
struct __align__(16) RankData {
const void* __restrict__ ptrs[8];
};
#endif
struct __align__(16) RankSignals {
#ifndef USE_ROCM
volatile
#endif
Signal* signals[8];
};
// like std::array, but aligned
template <typename T, int sz>
struct __align__(alignof(T) * sz) array_t {
T data[sz];
using type = T;
static constexpr int size = sz;
};
// use packed type to maximize memory efficiency
// goal: generate ld.128 and st.128 instructions
template <typename T>
struct packed_t {
// the (P)acked type for load/store
using P = array_t<T, 16 / sizeof(T)>;
// the (A)ccumulator type for reduction
using A = array_t<float, 16 / sizeof(T)>;
};
#define DINLINE __device__ __forceinline__
// scalar cast functions
DINLINE float upcast_s(half val) {
return __half2float(val);
}
template <typename T>
DINLINE T downcast_s(float val);
template <>
DINLINE half downcast_s(float val) {
return __float2half(val);
}
// scalar add functions
// for some reason when compiling with Pytorch, the + operator for half and
// bfloat is disabled so we call the intrinsics directly
DINLINE half& assign_add(half& a, half b) {
a = __hadd(a, b);
return a;
}
DINLINE float& assign_add(float& a, float b) {
return a += b;
}
#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
DINLINE float upcast_s(nv_bfloat16 val) {
return __bfloat162float(val);
}
template <>
DINLINE nv_bfloat16 downcast_s(float val) {
return __float2bfloat16(val);
}
DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) {
a = __hadd(a, b);
return a;
}
#endif
template <typename T, int N>
DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
#pragma unroll
for (int i = 0; i < N; i++) {
assign_add(a.data[i], b.data[i]);
}
return a;
}
template <typename T, int N>
DINLINE array_t<float, N> upcast(array_t<T, N> val) {
if constexpr (std::is_same<T, float>::value) {
return val;
} else {
array_t<float, N> out;
#pragma unroll
for (int i = 0; i < N; i++) {
out.data[i] = upcast_s(val.data[i]);
}
return out;
}
}
template <typename O>
DINLINE O downcast(array_t<float, O::size> val) {
if constexpr (std::is_same<typename O::type, float>::value) {
return val;
} else {
O out;
#pragma unroll
for (int i = 0; i < O::size; i++) {
out.data[i] = downcast_s<typename O::type>(val.data[i]);
}
return out;
}
}
// This function is meant to be used as the first synchronization in the all
// reduce kernel. Thus, it doesn't need to make any visibility guarantees for
// prior memory accesses. Note: volatile writes will not be reordered against
// other volatile writes.
template <int ngpus>
DINLINE void start_sync(
const RankSignals& sg,
#ifndef USE_ROCM
volatile
#endif
Signal* self_sg,
int rank) {
#ifdef USE_ROCM
uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
if (threadIdx.x < ngpus) {
// simultaneously write to the corresponding flag of all ranks.
// Latency = 1 p2p write
__scoped_atomic_store_n(
&sg.signals[threadIdx.x]->start[blockIdx.x][rank], flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
// wait until we got true from all ranks
while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], __ATOMIC_RELAXED, __MEMORY_SCOPE_DEVICE) <
flag)
;
}
__syncthreads();
// use one thread to update flag
if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
#else
if (threadIdx.x < ngpus) {
// reset flag for next time
self_sg->end[blockIdx.x][threadIdx.x] = 0;
// simultaneously write to the corresponding flag of all ranks.
// Latency = 1 p2p write
sg.signals[threadIdx.x]->start[blockIdx.x][rank] = 1;
// wait until we got true from all ranks
while (!self_sg->start[blockIdx.x][threadIdx.x])
;
}
__syncthreads();
#endif
}
// This function is meant to be used as the second or the final synchronization
// barrier in the all reduce kernel. If it's the final synchronization barrier,
// we don't need to make any visibility guarantees for prior memory accesses.
template <int ngpus, bool final_sync = false>
DINLINE void end_sync(
const RankSignals& sg,
#ifndef USE_ROCM
volatile
#endif
Signal* self_sg,
int rank) {
#ifdef USE_ROCM
__syncthreads();
// eliminate the case that prior writes are not visible after signals become
// visible. Note that I did not managed to make this happen through a lot of
// testing. Might be the case that hardware provides stronger guarantee than
// the memory model.
uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
if (threadIdx.x < ngpus) {
// simultaneously write to the corresponding flag of all ranks.
// Latency = 1 p2p write
__scoped_atomic_store_n(
&sg.signals[threadIdx.x]->end[blockIdx.x][rank],
flag,
final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE,
__MEMORY_SCOPE_SYSTEM);
// wait until we got true from all ranks
while (__scoped_atomic_load_n(
&self_sg->end[blockIdx.x][threadIdx.x],
final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
__MEMORY_SCOPE_DEVICE) < flag)
;
}
__syncthreads();
// use one thread to update flag
if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
#else
__syncthreads();
// eliminate the case that prior writes are not visible after signals become
// visible. Note that I did not managed to make this happen through a lot of
// testing. Might be the case that hardware provides stronger guarantee than
// the memory model.
if constexpr (!final_sync) __threadfence_system();
if (threadIdx.x < ngpus) {
// reset flag for next time
self_sg->start[blockIdx.x][threadIdx.x] = 0;
// simultaneously write to the corresponding flag of all ranks.
// Latency = 1 p2p write
sg.signals[threadIdx.x]->end[blockIdx.x][rank] = 1;
// wait until we got true from all ranks
while (!self_sg->end[blockIdx.x][threadIdx.x])
;
}
if constexpr (!final_sync) __syncthreads();
#endif
}
template <typename P, int ngpus, typename A>
DINLINE P packed_reduce(const P* ptrs[], int idx) {
A tmp = upcast(ptrs[0][idx]);
#pragma unroll
for (int i = 1; i < ngpus; i++) {
packed_assign_add(tmp, upcast(ptrs[i][idx]));
}
return downcast<P>(tmp);
}
template <typename T, int ngpus>
__global__ void __launch_bounds__(512, 1) cross_device_reduce_1stage(
RankData* _dp,
RankSignals sg,
#ifndef USE_ROCM
volatile
#endif
Signal* self_sg,
T* __restrict__ result,
int rank,
int size) {
using P = typename packed_t<T>::P;
using A = typename packed_t<T>::A;
// note: we don't reorder the address so the accumulation order is the same
// for all ranks, ensuring bitwise identical results
auto dp = *_dp;
start_sync<ngpus>(sg, self_sg, rank);
// do the actual reduction
for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; idx += gridDim.x * blockDim.x) {
((P*)result)[idx] = packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], idx);
}
end_sync<ngpus, true>(sg, self_sg, rank);
}
template <typename P>
#ifdef USE_ROCM
DINLINE P* get_tmp_buf(Signal* sg) {
#else
DINLINE P* get_tmp_buf(volatile Signal* sg) {
#endif
return (P*)(((Signal*)sg) + 1);
}
template <typename T, int ngpus>
__global__ void __launch_bounds__(512, 1) cross_device_reduce_2stage(
RankData* _dp,
RankSignals sg,
#ifndef USE_ROCM
volatile
#endif
Signal* self_sg,
T* __restrict__ result,
int rank,
int size) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int stride = gridDim.x * blockDim.x;
using P = typename packed_t<T>::P;
using A = typename packed_t<T>::A;
int part = size / ngpus;
int start = rank * part;
int end = rank == ngpus - 1 ? size : start + part;
int largest_part = part + size % ngpus;
const P* ptrs[ngpus];
P* tmps[ngpus];
#pragma unroll
for (int i = 0; i < ngpus; i++) {
int target = (rank + i) % ngpus;
ptrs[i] = (const P*)_dp->ptrs[target];
tmps[i] = get_tmp_buf<P>(sg.signals[target]);
}
auto tmp_out = tmps[0];
start_sync<ngpus>(sg, self_sg, rank);
// stage 1: reduce scatter
for (int idx = start + tid; idx < end; idx += stride) {
tmp_out[idx - start] = packed_reduce<P, ngpus, A>(ptrs, idx);
}
end_sync<ngpus>(sg, self_sg, rank);
// stage 2: allgather. Note: it's important to match the tid between
// the two stages, because visibility across devices is only guaranteed
// between threads that have the same tid. If thread i computes the sum of
// start + i in the first stage, then thread i also gathers start + i from all
// ranks.
for (int idx = tid; idx < largest_part; idx += stride) {
#pragma unroll
for (int i = 0; i < ngpus; i++) {
int gather_from_rank = ((rank + i) % ngpus);
if (gather_from_rank == ngpus - 1 || idx < part) {
int dst_idx = gather_from_rank * part + idx;
((P*)result)[dst_idx] = tmps[i][idx];
}
}
}
}
using IPC_KEY = std::array<uint8_t, sizeof(hipIpcMemHandle_t)>;
static_assert(sizeof(IPC_KEY) == sizeof(hipIpcMemHandle_t));
static_assert(alignof(IPC_KEY) == alignof(hipIpcMemHandle_t));
class CustomAllreduce {
public:
int rank_;
int world_size_;
bool full_nvlink_;
// below are device pointers
RankSignals sg_;
std::unordered_map<void*, RankData*> buffers_;
Signal* self_sg_;
// stores the registered device pointers from all ranks
RankData *d_rank_data_base_, *d_rank_data_end_;
std::vector<void*> graph_unreg_buffers_;
// a map from IPC handles to opened IPC pointers
std::map<IPC_KEY, char*> ipc_handles_;
/**
* meta is a pointer to device metadata and temporary buffer for allreduce.
*
* There's a total of sizeof(Signal) of prefix before the actual data,
* so meta + 1 points to actual temporary buffer.
*
* note: this class does not own any device memory. Any required buffers
* are passed in from the constructor
*/
CustomAllreduce(
Signal* meta,
void* rank_data,
size_t rank_data_sz,
const hipIpcMemHandle_t* handles,
const std::vector<int64_t>& offsets,
int rank,
bool full_nvlink = true)
: rank_(rank),
world_size_(offsets.size()),
full_nvlink_(full_nvlink),
self_sg_(meta),
d_rank_data_base_(reinterpret_cast<RankData*>(rank_data)),
d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) {
for (int i = 0; i < world_size_; i++) {
Signal* rank_sg;
if (i != rank_) {
char* handle = open_ipc_handle(&handles[i]);
handle += offsets[i];
rank_sg = (Signal*)handle;
} else {
rank_sg = self_sg_;
}
sg_.signals[i] = rank_sg;
}
}
char* open_ipc_handle(const void* ipc_handle) {
auto [it, new_handle] = ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr});
if (new_handle) {
char* ipc_ptr;
CUDACHECK(hipIpcOpenMemHandle(
(void**)&ipc_ptr, *((const hipIpcMemHandle_t*)ipc_handle), hipIpcMemLazyEnablePeerAccess));
it->second = ipc_ptr;
}
return it->second;
}
std::pair<std::vector<uint8_t>, std::vector<int64_t>> get_graph_buffer_ipc_meta() {
auto num_buffers = graph_unreg_buffers_.size();
auto handle_sz = sizeof(hipIpcMemHandle_t);
std::vector<uint8_t> handles(handle_sz * num_buffers, 0);
std::vector<int64_t> offsets(num_buffers);
for (int i = 0; i < num_buffers; i++) {
auto ptr = graph_unreg_buffers_[i];
void* base_ptr;
// note: must share the base address of each allocation, or we get wrong
// address
if (hipPointerGetAttribute(
&base_ptr,
#ifdef USE_ROCM
HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR,
#else
CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
#endif
(hipDeviceptr_t)ptr) != hipSuccess)
throw std::runtime_error("failed to get pointer attr");
CUDACHECK(hipIpcGetMemHandle((hipIpcMemHandle_t*)&handles[i * handle_sz], base_ptr));
offsets[i] = ((char*)ptr) - ((char*)base_ptr);
}
return std::make_pair(handles, offsets);
}
void check_rank_data_capacity(size_t num = 1) {
if (d_rank_data_base_ + num > d_rank_data_end_)
throw std::runtime_error(
"Rank data buffer is overflowed by " + std::to_string(d_rank_data_base_ + num - d_rank_data_end_));
}
void register_buffer(const std::vector<std::string>& handles, const std::vector<int64_t>& offsets, void* self) {
check_rank_data_capacity();
RankData data;
for (int i = 0; i < world_size_; i++) {
if (i != rank_) {
char* handle = open_ipc_handle(handles[i].data());
handle += offsets[i];
data.ptrs[i] = handle;
} else {
data.ptrs[i] = self;
}
}
auto d_data = d_rank_data_base_++;
CUDACHECK(hipMemcpy(d_data, &data, sizeof(RankData), hipMemcpyHostToDevice));
buffers_[self] = d_data;
}
// note: when registering graph buffers, we intentionally choose to not
// deduplicate the addresses. That means if the allocator reuses some
// addresses, they will be registered again. This is to account for the remote
// possibility of different allocation patterns between ranks. For example,
// rank 1 may get the same input address for the second allreduce, but rank 2
// got a different address. IPC handles have internal reference counting
// mechanism so overhead should be small.
void
register_graph_buffers(const std::vector<std::string>& handles, const std::vector<std::vector<int64_t>>& offsets) {
auto num_buffers = graph_unreg_buffers_.size();
check_rank_data_capacity(num_buffers);
std::vector<RankData> rank_data(num_buffers);
for (int i = 0; i < num_buffers; i++) {
auto self_ptr = graph_unreg_buffers_[i];
auto& rd = rank_data[i];
for (int j = 0; j < world_size_; j++) {
if (j != rank_) {
char* handle = open_ipc_handle(&handles[j][i * sizeof(hipIpcMemHandle_t)]);
handle += offsets[j][i];
rd.ptrs[j] = handle;
} else {
rd.ptrs[j] = self_ptr;
}
}
}
CUDACHECK(hipMemcpy(d_rank_data_base_, rank_data.data(), sizeof(RankData) * num_buffers, hipMemcpyHostToDevice));
d_rank_data_base_ += num_buffers;
graph_unreg_buffers_.clear();
}
/**
* This is the result after careful grid search. Using 36 blocks give the best
* or close to the best runtime on the devices I tried: A100, A10, A30, T4,
* V100. You'll notice that NCCL kernels also only take a small amount of SMs.
* Not quite sure the underlying reason, but my guess is that too many SMs
* will cause contention on NVLink bus.
*/
template <typename T>
void allreduce(
hipStream_t stream,
T* input,
T* output,
int size,
#ifndef USE_ROCM
int threads = 512,
int block_limit = 36){
#else
int threads = 512,
int block_limit = 16) {
#endif
auto d = packed_t<T>::P::size;
if (size % d != 0)
throw std::runtime_error(
"custom allreduce currently requires input length to be multiple "
"of " +
std::to_string(d));
if (block_limit > kMaxBlocks)
throw std::runtime_error(
"max supported block limit is " + std::to_string(kMaxBlocks) + ". Got " + std::to_string(block_limit));
RankData* ptrs;
hipStreamCaptureStatus status;
CUDACHECK(hipStreamIsCapturing(stream, &status));
if (status == hipStreamCaptureStatusActive) {
ptrs = d_rank_data_base_ + graph_unreg_buffers_.size();
graph_unreg_buffers_.push_back(input);
} else {
auto it = buffers_.find(input);
if (it == buffers_.end())
throw std::runtime_error(
"buffer address " + std::to_string(reinterpret_cast<uint64_t>(input)) + " is not registered!");
ptrs = it->second;
}
size /= d;
auto bytes = size * sizeof(typename packed_t<T>::P);
int blocks = ::min(block_limit, (size + threads - 1) / threads);
#define KL(ngpus, name) \
hipLaunchKernelGGL( \
(name<T, ngpus>), dim3(blocks), dim3(threads), 0, stream, ptrs, sg_, self_sg_, output, rank_, size);
#define REDUCE_CASE(ngpus) \
case ngpus: { \
if (world_size_ == 2) { \
KL(ngpus, cross_device_reduce_1stage); \
} else if (full_nvlink_) { \
if ((world_size_ <= 4 && bytes < 512 * 1024) || (world_size_ <= 8 && bytes < 256 * 1024)) { \
KL(ngpus, cross_device_reduce_1stage); \
} else { \
KL(ngpus, cross_device_reduce_2stage); \
} \
} \
break; \
}
switch (world_size_) {
REDUCE_CASE(2)
REDUCE_CASE(4)
REDUCE_CASE(6)
REDUCE_CASE(8)
default:
throw std::runtime_error(
"custom allreduce only supports num gpus in (2,4,6,8). Actual num "
"gpus = " +
std::to_string(world_size_));
}
#undef REDUCE_CASE
#undef KL
}
~CustomAllreduce() {
for (auto [_, ptr] : ipc_handles_) {
CUDACHECK(hipIpcCloseMemHandle(ptr));
}
}
}; // namespace sglang
/**
* To inspect PTX/SASS, copy paste this header file to compiler explorer and add
a template instantiation:
* template void sglang::CustomAllreduce::allreduce<half>(hipStream_t, half *,
half *, int, int, int);
*/
} // namespace sglang
python/sglang/kernels/aot/csrc/allreduce/deterministic_all_reduce.hip
+178
View File
@@ -0,0 +1,178 @@
// Deterministic All-Reduce for ROCm/HIP
//
// This is a wrapper that forces the use of the existing 1-stage all-reduce kernel
// (cross_device_reduce_1stage) which is inherently deterministic due to fixed
// accumulation ordering (no atomics, no race conditions).
//
// How the 1-stage kernel works:
// - Each GPU reads ALL data from ALL other GPUs via direct memory access
// - Each GPU reduces the data locally in a fixed order
// - Result: every GPU has the complete reduced output
//
// This is NOT a reduce-scatter + all-gather (RS+AG) approach.
// The 2-stage kernel (cross_device_reduce_2stage) implements RS+AG but may have
// non-deterministic behavior, so we explicitly avoid it here.
#include <ATen/hip/Exceptions.h>
#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
#include <torch/all.h>
#include "custom_all_reduce_hip.cuh"
using fptr_t = int64_t;
static_assert(sizeof(void*) == sizeof(fptr_t));
// Helper function for weak contiguity check
bool _is_weak_contiguous_det(torch::Tensor& t) {
return t.is_contiguous() ||
(t.storage().nbytes() - t.storage_offset() * t.element_size() == t.numel() * t.element_size());
}
// Deterministic all-reduce for registered buffers (ROCm)
// Uses the 1-stage kernel which is deterministic (fixed ordering)
void deterministic_all_reduce_reg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out) {
const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp));
auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
TORCH_CHECK_EQ(inp.numel(), out.numel());
TORCH_CHECK(_is_weak_contiguous_det(out));
TORCH_CHECK(_is_weak_contiguous_det(inp));
auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
// For ROCm, manually call the 1-stage kernel to ensure deterministic ordering
// Get rank data pointer
sglang::RankData* ptrs;
hipStreamCaptureStatus status;
AT_CUDA_CHECK(hipStreamIsCapturing(stream, &status));
if (status == hipStreamCaptureStatusActive) {
ptrs = fa->d_rank_data_base_ + fa->graph_unreg_buffers_.size();
fa->graph_unreg_buffers_.push_back(inp.data_ptr());
} else {
auto it = fa->buffers_.find(inp.data_ptr());
if (it == fa->buffers_.end()) {
throw std::runtime_error("buffer not registered!");
}
ptrs = it->second;
}
int size = out.numel();
int threads = 512;
switch (out.scalar_type()) {
case at::ScalarType::Float: {
using T = float;
using P = typename sglang::packed_t<T>::P;
auto d = P::size;
if (size % d != 0) {
throw std::runtime_error("size must be multiple of " + std::to_string(d));
}
size /= d;
int blocks = std::min(16, (size + threads - 1) / threads);
// Always use 1-stage kernel for determinism
switch (fa->world_size_) {
case 2:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 4:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 6:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 8:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
default:
throw std::runtime_error("world_size must be in (2,4,6,8)");
}
break;
}
case at::ScalarType::Half: {
using T = half;
using P = typename sglang::packed_t<T>::P;
auto d = P::size;
if (size % d != 0) {
throw std::runtime_error("size must be multiple of " + std::to_string(d));
}
size /= d;
int blocks = std::min(16, (size + threads - 1) / threads);
switch (fa->world_size_) {
case 2:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 4:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 6:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 8:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
default:
throw std::runtime_error("world_size must be in (2,4,6,8)");
}
break;
}
#if (__HIP_ARCH__ >= 800 || !defined(__HIP_ARCH__))
case at::ScalarType::BFloat16: {
using T = nv_bfloat16;
using P = typename sglang::packed_t<T>::P;
auto d = P::size;
if (size % d != 0) {
throw std::runtime_error("size must be multiple of " + std::to_string(d));
}
size /= d;
int blocks = std::min(16, (size + threads - 1) / threads);
switch (fa->world_size_) {
case 2:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 4:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 6:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
case 8:
hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
break;
default:
throw std::runtime_error("world_size must be in (2,4,6,8)");
}
break;
}
#endif
default:
throw std::runtime_error("deterministic allreduce only supports float32, float16 and bfloat16");
}
}
// Deterministic all-reduce for unregistered buffers (ROCm)
void deterministic_all_reduce_unreg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer, torch::Tensor& out) {
const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp));
auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
auto input_size = inp.numel() * inp.element_size();
TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
TORCH_CHECK_EQ(inp.numel(), out.numel());
TORCH_CHECK(input_size <= reg_buffer.numel() * reg_buffer.element_size(),
"registered buffer is too small to contain the input");
AT_CUDA_CHECK(hipMemcpyAsync(reg_buffer.data_ptr(), inp.data_ptr(),
input_size, hipMemcpyDeviceToDevice, stream));
deterministic_all_reduce_reg(_fa, reg_buffer, out);
}
python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce.cu
+111
View File
@@ -0,0 +1,111 @@
#include <ATen/cuda/Exceptions.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/all.h>
#ifdef USE_ROCM
#include "quick_all_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>;
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(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(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)
#endif // USE_ROCM
python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce.cuh
+634
View File
@@ -0,0 +1,634 @@
#pragma once
#include <hip/hip_runtime.h>
#include "quick_all_reduce_base.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;
}
}
};
// 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);
}
// That's pretty much it...
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;
int grid_size = gridDim.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
python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce.h
+259
View File
@@ -0,0 +1,259 @@
#pragma once
#include <hip/hip_runtime.h>
#include <vector>
#include "quick_all_reduce.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;
// Read this block's counter from device memory and bump it here in the
// kernel. Keeping the value in device memory (instead of a host scalar
// baked into the launch) lets every CUDA-graph replay see a fresh color.
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++;
}
// The whole block ends up with the same value, so a single writer suffices.
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); \
}
enum QuickReduceQuantLevel {
F16 = 0,
INT8 = 1,
INT6 = 2,
INT4 = 3,
};
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));
// A per-block color counter that the kernel advances itself. Seed it with
// 1 rather than 0 so it never matches the freshly zeroed flags buffer.
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() {
// This buffer is created before `initialized` becomes true, so release it
// on its own check to keep a half-finished init from leaking it.
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;
default:
TWOSHOT_DISPATCH(CodecFP)
break;
}
HIP_CHECK(cudaGetLastError());
// The color now advances on-device inside the kernel; no host-side bump.
}
};
} // namespace quickreduce
python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce_base.h
+318
View File
@@ -0,0 +1,318 @@
#pragma once
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
#include <cstdint>
#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
};
};
__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");
__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;
}
// 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
python/sglang/kernels/aot/csrc/attention/cutlass_mla_kernel.cu
+274
View File
@@ -0,0 +1,274 @@
/*
Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cutlass/cutlass.h>
#include <cutlass/kernel_hardware_info.h>
#include <torch/all.h>
#include <cute/tensor.hpp>
#include <iostream>
#include "cutlass_sm100_mla/device/sm100_mla.hpp"
#include "cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp"
#include "utils.h"
// clang-format off
#if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
void cutlass_mla_decode(
torch::Tensor const& out,
torch::Tensor const& q_nope,
torch::Tensor const& q_pe,
torch::Tensor const& kv_c_and_k_pe_cache,
torch::Tensor const& seq_lens,
torch::Tensor const& page_table,
torch::Tensor const& workspace,
int64_t num_kv_splits) {
TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
}
int64_t cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
}
#else
#define CUTLASS_CHECK(status) \
{ \
cutlass::Status error = status; \
TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
}
using namespace cute;
using namespace cutlass::fmha::kernel;
template <bool v>
struct IsPersistent {
static const bool value = v;
};
template <typename T, bool IsPaged128, typename PersistenceOption = IsPersistent<true>>
struct MlaSm100 {
using Element = T;
using ElementAcc = float;
using ElementOut = T;
using TileShape = Shape<_128, _128, Shape<_512, _64>>;
using TileShapeH = cute::tuple_element_t<0, TileShape>;
using TileShapeD = cute::tuple_element_t<2, TileShape>;
// H K (D_latent D_rope) B
using ProblemShape = cute::tuple<TileShapeH, int, TileShapeD, int>;
using StrideQ = cute::tuple<int64_t, _1, int64_t>; // H D B
using StrideK = cute::tuple<int64_t, _1, int64_t>; // K D B
using StrideO = StrideK; // H D B
using StrideLSE = cute::tuple<_1, int>; // H B
using TileScheduler =
std::conditional_t<PersistenceOption::value, Sm100MlaPersistentTileScheduler, Sm100MlaIndividualTileScheduler>;
using FmhaKernel = cutlass::fmha::kernel::Sm100FmhaMlaKernelTmaWarpspecialized<
TileShape,
Element,
ElementAcc,
ElementOut,
ElementAcc,
TileScheduler,
/*kIsCpAsync=*/!IsPaged128>;
using Fmha = cutlass::fmha::device::MLA<FmhaKernel>;
};
template <typename T>
typename T::Fmha::Arguments args_from_options(
at::Tensor const& out,
at::Tensor const& q_nope,
at::Tensor const& q_pe,
at::Tensor const& kv_c_and_k_pe_cache,
at::Tensor const& seq_lens,
at::Tensor const& page_table,
double sm_scale,
int64_t num_kv_splits) {
cutlass::KernelHardwareInfo hw_info;
hw_info.device_id = q_nope.device().index();
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
int batches = q_nope.size(0);
int page_count_per_seq = page_table.size(1);
int page_count_total = kv_c_and_k_pe_cache.size(0);
int page_size = kv_c_and_k_pe_cache.size(1);
int max_seq_len = page_size * page_count_per_seq;
using TileShapeH = typename T::TileShapeH;
using TileShapeD = typename T::TileShapeD;
auto problem_shape = cute::make_tuple(TileShapeH{}, max_seq_len, TileShapeD{}, batches);
auto [H, K, D, B] = problem_shape;
auto [D_latent, D_rope] = D;
float scale = float(sm_scale);
using StrideQ = typename T::StrideQ;
using StrideK = typename T::StrideK;
using StrideO = typename T::StrideO;
using StrideLSE = typename T::StrideLSE;
StrideQ stride_Q_nope = cute::make_tuple(
static_cast<int64_t>(q_nope.stride(1)), _1{}, static_cast<int64_t>(q_nope.stride(0)));
StrideQ stride_Q_pe = cute::make_tuple(
static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));
StrideK stride_C = cute::make_tuple(
static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));
using Element = typename T::Element;
using ElementOut = typename T::ElementOut;
using ElementAcc = typename T::ElementAcc;
auto Q_nope_ptr = static_cast<Element*>(q_nope.data_ptr());
auto Q_pe_ptr = static_cast<Element*>(q_pe.data_ptr());
auto C_ptr = static_cast<Element*>(kv_c_and_k_pe_cache.data_ptr());
typename T::Fmha::Arguments arguments{
problem_shape,
{scale,
Q_nope_ptr,
stride_Q_nope,
Q_pe_ptr,
stride_Q_pe,
C_ptr,
stride_C,
C_ptr + D_latent,
stride_C,
static_cast<int*>(seq_lens.data_ptr()),
static_cast<int*>(page_table.data_ptr()),
stride_PT,
page_count_total,
page_size},
{static_cast<ElementOut*>(out.data_ptr()), stride_O, static_cast<ElementAcc*>(nullptr), stride_LSE},
hw_info,
// TODO(trevor-m): Change split_kv back to -1 when
// https://github.com/NVIDIA/cutlass/issues/2274 is fixed. Split_kv=1 will
// perform worse with larger context length and smaller batch sizes.
static_cast<int>(num_kv_splits), // split_kv
nullptr, // is_var_split_kv
};
// TODO(kaixih@nvidia): When split_kv=-1 and is_var_split_kv=false, we compute
// split_kv automatically based on batch size and sequence length to balance
// workload across available SMs. Consider using var_split_kv for manual
// control if needed.
T::Fmha::set_split_kv(arguments);
return arguments;
}
template <typename Element, bool IsPaged128, typename PersistenceOption>
void runMla(
at::Tensor const& out,
at::Tensor const& q_nope,
at::Tensor const& q_pe,
at::Tensor const& kv_c_and_k_pe_cache,
at::Tensor const& seq_lens,
at::Tensor const& page_table,
at::Tensor const& workspace,
double sm_scale,
int64_t num_kv_splits,
cudaStream_t stream) {
using MlaSm100Type = MlaSm100<Element, IsPaged128, PersistenceOption>;
typename MlaSm100Type::Fmha fmha;
auto arguments = args_from_options<MlaSm100Type>(out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, sm_scale, num_kv_splits);
CUTLASS_CHECK(fmha.can_implement(arguments));
CUTLASS_CHECK(fmha.initialize(arguments, workspace.data_ptr(), stream));
CUTLASS_CHECK(fmha.run(arguments, workspace.data_ptr(), stream));
}
#define DISPATCH_BOOL(expr, const_expr, ...) \
[&]() -> bool { \
if (expr) { \
constexpr bool const_expr = true; \
return __VA_ARGS__(); \
} else { \
constexpr bool const_expr = false; \
return __VA_ARGS__(); \
} \
}()
void cutlass_mla_decode(
torch::Tensor const& out,
torch::Tensor const& q_nope,
torch::Tensor const& q_pe,
torch::Tensor const& kv_c_and_k_pe_cache,
torch::Tensor const& seq_lens,
torch::Tensor const& page_table,
torch::Tensor const& workspace,
double sm_scale,
int64_t num_kv_splits) {
auto sm_version = getSMVersion();
// On SM103a, half of the accuracy tests are failing.
TORCH_CHECK(sm_version == 100, "cutlass_mla_decode is only supported on compute capability 10.0, but found sm version ", sm_version);
auto in_dtype = q_nope.dtype();
at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device());
const int page_size = kv_c_and_k_pe_cache.size(1);
// NOTE(alcanderian): IsPersistent has bug with manual split_kv.
// Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8)
// Maybe per batch split kv will fix this.
DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
if (in_dtype == at::ScalarType::Half) {
runMla<cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
} else if (in_dtype == at::ScalarType::BFloat16) {
runMla<cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
} else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
runMla<cutlass::float_e4m3_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
} else {
TORCH_CHECK(false, "Unsupported input data type of MLA");
}
return true;
});
return true;
});
}
int64_t cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
// Workspace size depends on ElementAcc and ElementLSE (same as ElementAcc)
// which are float, so Element type here doesn't matter.
using MlaSm100Type = MlaSm100<cutlass::half_t, true>;
// Get split kv. Requires problem shape and sm_count only.
typename MlaSm100Type::Fmha::Arguments arguments;
using TileShapeH = typename MlaSm100Type::TileShapeH;
using TileShapeD = typename MlaSm100Type::TileShapeD;
arguments.problem_shape =
cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
// Assumes device 0 when getting sm_count.
arguments.hw_info.sm_count =
sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
arguments.split_kv = static_cast<int>(num_kv_splits);
MlaSm100Type::Fmha::set_split_kv(arguments);
return MlaSm100Type::Fmha::get_workspace_size(arguments);
}
#endif
// clang-format on
python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/device/sm100_mla.hpp
+358
View File
@@ -0,0 +1,358 @@
/***************************************************************************************************
* Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*!
\file
\brief An universal device layer for cutlass 3.x-style kernels.
*/
// clang-format off
#pragma once
// common
#include "cutlass/cutlass.h"
#include "cutlass/device_kernel.h"
#if !defined(__CUDACC_RTC__)
#include "cutlass/cluster_launch.hpp"
#include "cutlass/trace.h"
#endif // !defined(__CUDACC_RTC__)
#include "../kernel/sm100_fmha_mla_tma_warpspecialized.hpp"
#include "../kernel/sm100_fmha_mla_reduction.hpp"
////////////////////////////////////////////////////////////////////////////////
namespace cutlass::fmha::device {
using namespace cute;
using namespace cutlass::fmha::kernel;
////////////////////////////////////////////////////////////////////////////////
////////////////////////////// CUTLASS 3.x API /////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
template<
class Kernel_
>
class MLA {
public:
using Kernel = Kernel_;
using ReductionKernel = cutlass::fmha::kernel::Sm100FmhaMlaReductionKernel<
typename Kernel::ElementOut,
typename Kernel::ElementAcc,
typename Kernel::ElementAcc,
Kernel::TileShapeH::value,
Kernel::TileShapeL::value,
256 /*Max split*/
>;
/// Argument structure: User API
using KernelArguments = typename Kernel::Arguments;
using ReductionArguments = typename ReductionKernel::Arguments;
using Arguments = KernelArguments;
/// Argument structure: Kernel API
using KernelParams = typename Kernel::Params;
using ReductionParams = typename ReductionKernel::Params;
struct Params {
KernelParams fmha_params;
ReductionParams reduction_params;
};
private:
/// Kernel API parameters object
Params params_;
bool is_initialized(bool set = false) {
static bool initialized = false;
if (set) initialized = true;
return initialized;
}
static ReductionArguments to_reduction_args(Arguments const& args) {
auto [H, K, D, B] = args.problem_shape;
return ReductionArguments{
nullptr, args.epilogue.ptr_o, nullptr, args.epilogue.ptr_lse,
args.mainloop.softmax_scale, B, args.split_kv, K, args.mainloop.ptr_seq,
args.ptr_split_kv, Kernel::TileShapeS::value
};
}
public:
/// Access the Params structure
Params const& params() const {
return params_;
}
static void set_split_kv (KernelArguments& args) {
if (args.split_kv >= 1) return;
auto [H, K, D, B] = args.problem_shape;
int sm_count = args.hw_info.sm_count;
int max_splits = ceil_div(K, 128);
int sms_per_batch = max(1, sm_count / B);
int split_heur = min(max_splits, sms_per_batch);
int waves = ceil_div(B * split_heur, sm_count);
int k_waves = ceil_div(max_splits, split_heur);
int split_wave_aware = ceil_div(max_splits, k_waves);
args.split_kv = split_wave_aware;
}
/// Determines whether the GEMM can execute the given problem.
static Status
can_implement(Arguments const& args) {
if (! Kernel::can_implement(args)) {
return Status::kInvalid;
}
if (! ReductionKernel::can_implement(to_reduction_args(args))) {
return Status::kInvalid;
}
return Status::kSuccess;
}
/// Gets the workspace size
static size_t
get_workspace_size(Arguments const& args) {
size_t workspace_bytes = 0;
workspace_bytes += Kernel::get_workspace_size(args);
workspace_bytes += ReductionKernel::get_workspace_size(to_reduction_args(args));
return workspace_bytes;
}
/// Computes the maximum number of active blocks per multiprocessor
static int maximum_active_blocks(int /* smem_capacity */ = -1) {
CUTLASS_TRACE_HOST("MLA::maximum_active_blocks()");
int max_active_blocks = -1;
int smem_size = Kernel::SharedStorageSize;
// first, account for dynamic smem capacity if needed
cudaError_t result;
if (smem_size >= (48 << 10)) {
CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
result = cudaFuncSetAttribute(
device_kernel<Kernel>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(
" cudaFuncSetAttribute() returned error: "
<< cudaGetErrorString(result));
return -1;
}
}
// query occupancy after setting smem size
result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&max_active_blocks,
device_kernel<Kernel>,
Kernel::MaxThreadsPerBlock,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(
" cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: "
<< cudaGetErrorString(result));
return -1;
}
CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks);
return max_active_blocks;
}
/// Initializes GEMM state from arguments.
Status
initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
CUTLASS_TRACE_HOST("MLA::initialize() - workspace "
<< workspace << ", stream: " << (stream ? "non-null" : "null"));
// Initialize the workspace
Status status = Kernel::initialize_workspace(args, workspace, stream);
if (status != Status::kSuccess) {
return status;
}
status = ReductionKernel::initialize_workspace(to_reduction_args(args), workspace, stream);
if (status != Status::kSuccess) {
return status;
}
KernelParams kernel_params = Kernel::to_underlying_arguments(args, workspace);
ReductionArguments reduction_args = to_reduction_args(args);
if (reduction_args.split_kv > 1) {
reduction_args.ptr_oaccum = kernel_params.epilogue.ptr_o_acc;
reduction_args.ptr_lseaccum = kernel_params.epilogue.ptr_lse_acc;
}
ReductionParams reduction_params = ReductionKernel::to_underlying_arguments(reduction_args, workspace);
// Initialize the Params structure
params_ = Params {kernel_params, reduction_params};
if (is_initialized()) return Status::kSuccess;
// account for dynamic smem capacity if needed
// no dynamic smem is needed for reduction kernel
int smem_size = Kernel::SharedStorageSize;
if (smem_size >= (48 << 10)) {
CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
cudaError_t result = cudaFuncSetAttribute(
device_kernel<Kernel>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result));
return Status::kErrorInternal;
}
}
is_initialized(true);
return Status::kSuccess;
}
/// Update API is preserved in 3.0, but does not guarantee a lightweight update of params.
Status
update(Arguments const& args, void* workspace = nullptr) {
CUTLASS_TRACE_HOST("MLA()::update() - workspace: " << workspace);
size_t workspace_bytes = get_workspace_size(args);
if (workspace_bytes > 0 && nullptr == workspace) {
return Status::kErrorWorkspaceNull;
}
auto fmha_params = Kernel::to_underlying_arguments(args, workspace);
ReductionArguments reduction_args = to_reduction_args(args);
if (reduction_args.split_kv > 1) {
reduction_args.ptr_oaccum = fmha_params.epilogue.ptr_o_acc;
reduction_args.ptr_lseaccum = fmha_params.epilogue.ptr_lse_acc;
}
ReductionParams reduction_params = ReductionKernel::to_underlying_arguments(reduction_args, workspace);
// Initialize the Params structure
params_ = Params {fmha_params, reduction_params};
return Status::kSuccess;
}
/// Primary run() entry point API that is static allowing users to create and manage their own params.
/// Supplied params struct must be construct by calling Kernel::to_underling_arguments()
static Status
run(Params& params, cudaStream_t stream = nullptr) {
CUTLASS_TRACE_HOST("MLA::run()");
dim3 const block = Kernel::get_block_shape();
dim3 const grid = Kernel::get_grid_shape(params.fmha_params);
// configure smem size and carveout
int smem_size = Kernel::SharedStorageSize;
Status launch_result;
// Use extended launch API only for mainloops that use it
if constexpr(Kernel::ArchTag::kMinComputeCapability >= 90) {
dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}),
cute::size<1>(typename Kernel::ClusterShape{}),
cute::size<2>(typename Kernel::ClusterShape{}));
void const* kernel = (void const*) device_kernel<Kernel>;
void* kernel_params[] = {&params.fmha_params};
launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params);
}
else {
launch_result = Status::kSuccess;
device_kernel<Kernel><<<grid, block, smem_size, stream>>>(params.fmha_params);
}
cudaError_t result = cudaGetLastError();
if (cudaSuccess != result or Status::kSuccess != launch_result) {
//return Status::kSuccess;
CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result);
return Status::kErrorInternal;
}
if (params.reduction_params.split_kv > 1) {
// launch reduction kernel
dim3 const block = ReductionKernel::get_block_shape();
dim3 const grid = ReductionKernel::get_grid_shape(params.reduction_params);
device_kernel<ReductionKernel><<<grid, block, 0, stream>>>(params.reduction_params);
cudaError_t result = cudaGetLastError();
if (cudaSuccess == result) {
return Status::kSuccess;
}
else {
CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result);
return Status::kErrorInternal;
}
}
else {
return Status::kSuccess;
}
}
//
// Non-static launch overloads that first create and set the internal params struct of this kernel handle.
//
/// Launches the kernel after first constructing Params internal state from supplied arguments.
Status
run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
Status status = initialize(args, workspace, stream);
if (Status::kSuccess == status) {
status = run(params_, stream);
}
return status;
}
/// Launches the kernel after first constructing Params internal state from supplied arguments.
Status
operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
return run(args, workspace, stream);
}
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
Status
run(cudaStream_t stream = nullptr) {
return run(params_, stream);
}
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
Status
operator()(cudaStream_t stream = nullptr) {
return run(params_, stream);
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::fmha::device
////////////////////////////////////////////////////////////////////////////////
python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp
+198
View File
@@ -0,0 +1,198 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
// clang-format off
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/arch/arch.h"
#include "cute/tensor.hpp"
namespace cutlass::fmha::kernel {
using namespace cute;
template<
class ElementOut,
class ElementAcc,
class ElementScale,
size_t kNumHeads,
size_t kHeadDimLatent,
int kMaxSplits
>
struct Sm100FmhaMlaReductionKernel {
static const int SharedStorageSize = 0;
static const int MaxThreadsPerBlock = 128;
static const int MinBlocksPerMultiprocessor = 1;
using ArchTag = cutlass::arch::Sm100;
static_assert(kHeadDimLatent % MaxThreadsPerBlock == 0);
struct Arguments {
ElementAcc* ptr_oaccum = nullptr;
ElementOut* ptr_o = nullptr;
ElementAcc* ptr_lseaccum = nullptr;
ElementAcc* ptr_lse = nullptr;
ElementScale scale = 1.f;
int num_batches = 0;
int split_kv = -1;
int dim_k = -1;
int* ptr_seq = nullptr;
int* ptr_split_kv = nullptr;
int tile_shape_s = 128;
};
using Params = Arguments;
static Params to_underlying_arguments(Arguments const& args, void* workspace) {
return {args.ptr_oaccum, args.ptr_o, args.ptr_lseaccum, args.ptr_lse,
args.scale, args.num_batches, args.split_kv, args.dim_k, args.ptr_seq,
args.ptr_split_kv, args.tile_shape_s};
}
static size_t get_workspace_size(Arguments const& /*args*/) {
return 0;
}
static Status initialize_workspace(
Arguments const& /*args*/, void* /*ws*/, cudaStream_t /*stream*/) {
return Status::kSuccess;
}
static dim3 get_grid_shape(Params const& params) {
return dim3(kNumHeads, 1, params.num_batches);
}
static dim3 get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
static bool can_implement(Arguments const& args) {
if (args.num_batches <= 0) return false;
if (args.split_kv <= 0) return false;
return true;
}
CUTLASS_DEVICE void operator() (Params const& params, char* smem_raw) {
if (params.split_kv <= 1) return;
auto blk_coord = make_coord(blockIdx.x, _0{}, blockIdx.z);
__shared__ ElementAcc sLseScale[kMaxSplits];
const size_t offset_lseaccum = get<0>(blk_coord) + kNumHeads * params.split_kv * get<2>(blk_coord);
const size_t offset_lse = get<0>(blk_coord) + kNumHeads * get<2>(blk_coord);
Tensor gLSEaccum = make_tensor(make_gmem_ptr(params.ptr_lseaccum + offset_lseaccum),
make_shape(params.split_kv), Stride<Int<kNumHeads>>{});
Tensor gLSE = make_tensor(make_gmem_ptr(params.ptr_lse + offset_lse),
Shape<_1>{}, Stride<_1>{});
auto dim_k = params.ptr_seq == nullptr ? params.dim_k : params.ptr_seq[get<2>(blk_coord)];
auto local_split_kv = params.ptr_split_kv == nullptr ? params.split_kv : params.ptr_split_kv[get<2>(blk_coord)];
auto k_tile_total = ceil_div(dim_k, params.tile_shape_s);
auto k_tile_per_cta = ceil_div(k_tile_total, local_split_kv);
local_split_kv = ceil_div(k_tile_total, k_tile_per_cta);
int warp_idx = cutlass::canonical_warp_idx_sync();
if (warp_idx == 0) {
constexpr int kNLsePerThread = cute::ceil_div(kMaxSplits, 32);
ElementAcc local_lse[kNLsePerThread];
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kNLsePerThread; ++i) {
const int split = i * 32 + threadIdx.x;
local_lse[i] = split < local_split_kv ? gLSEaccum(split) : -std::numeric_limits<ElementAcc>::infinity();
}
ElementAcc lse_max = -std::numeric_limits<ElementAcc>::infinity();
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kNLsePerThread; ++i) {
lse_max = max(lse_max, local_lse[i]);
}
CUTLASS_PRAGMA_UNROLL
for (int offset = 16; offset >= 1; offset /= 2) {
lse_max = max(lse_max, __shfl_xor_sync(0xffffffff, lse_max, offset));
}
lse_max = lse_max == -std::numeric_limits<ElementAcc>::infinity() ? 0.0f : lse_max; // In case all local LSEs are -inf
lse_max = __shfl_sync(0xffffffff, lse_max, 0);
ElementAcc sum_lse = 0;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kNLsePerThread; ++i) {
sum_lse = sum_lse + expf(local_lse[i] - lse_max);
}
CUTLASS_PRAGMA_UNROLL
for (int offset = 16; offset >= 1; offset /= 2) {
sum_lse = sum_lse + __shfl_xor_sync(0xffffffff, sum_lse, offset);
}
sum_lse = __shfl_sync(0xffffffff, sum_lse, 0);
ElementAcc global_lse = (sum_lse == 0.f || sum_lse != sum_lse) ? std::numeric_limits<ElementAcc>::infinity() : logf(sum_lse) + lse_max;
if (threadIdx.x == 0 and params.ptr_lse != nullptr) {
gLSE(0) = global_lse;
}
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kNLsePerThread; ++i) {
const int split = i * 32 + threadIdx.x;
if (split < local_split_kv) {
sLseScale[split] = expf(local_lse[i] - global_lse);
}
}
}
__syncthreads();
constexpr int Elements = kHeadDimLatent / MaxThreadsPerBlock;
const size_t offset_oaccum = kHeadDimLatent * params.split_kv * (get<0>(blk_coord) + kNumHeads * get<2>(blk_coord));
Tensor gOaccum = make_tensor(make_gmem_ptr(params.ptr_oaccum + offset_oaccum),
Shape<Int<kHeadDimLatent>>{}, Stride<_1>{});
ElementAcc local_val[Elements] = {0};
for (int split = 0; split < local_split_kv; ++split) {
ElementAcc lse_scale = sLseScale[split];
CUTLASS_PRAGMA_UNROLL
for(int i = 0; i < Elements; ++i) {
local_val[i] += lse_scale * gOaccum(threadIdx.x + MaxThreadsPerBlock * i);
}
gOaccum.data() = gOaccum.data() + kHeadDimLatent;
}
auto ptr_o_local = params.ptr_o + (get<0>(blk_coord) + get<2>(blk_coord) * kNumHeads) * kHeadDimLatent;
Tensor gO = make_tensor(make_gmem_ptr(ptr_o_local), Shape<Int<kHeadDimLatent>>{}, Stride<_1>{});
CUTLASS_PRAGMA_UNROLL
for(int i = 0; i < Elements; ++i) {
gO(threadIdx.x + MaxThreadsPerBlock * i) = static_cast<ElementOut>(local_val[i]);
}
}
};
} // namespace cutlass::fmha::kernel
python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp
+2018
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/aot/csrc/attention/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp
+160
View File
@@ -0,0 +1,160 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
// clang-format off
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h"
#include "cutlass/kernel_hardware_info.h"
namespace cutlass::fmha::kernel {
////////////////////////////////////////////////////////////////////////////////
struct Sm100MlaIndividualTileScheduler {
struct Params {
dim3 grid;
};
bool valid_ = true;
CUTLASS_DEVICE
Sm100MlaIndividualTileScheduler(Params const&) {}
template<class ProblemShape, class ClusterShape>
static Params to_underlying_arguments(
ProblemShape const& problem_shape, KernelHardwareInfo hw_info,
ClusterShape const& cluster_shape, int const& split_kv) {
using namespace cute;
dim3 grid(get<0>(cluster_shape), get<3>(problem_shape) /* Batch */, split_kv /*Maximum Split KV*/);
return Params{ grid };
}
static dim3 get_grid_shape(Params const& params) {
return params.grid;
}
CUTLASS_DEVICE
bool is_valid() {
return valid_;
}
CUTLASS_DEVICE
auto get_block_coord() {
using namespace cute;
return make_coord(blockIdx.x, _0{}, blockIdx.y, blockIdx.z);
}
CUTLASS_DEVICE
Sm100MlaIndividualTileScheduler& operator++() {
valid_ = false;
return *this;
}
};
////////////////////////////////////////////////////////////////////////////////
struct Sm100MlaPersistentTileScheduler {
struct Params {
int num_blocks;
FastDivmod divmod_m_block;
FastDivmod divmod_b;
FastDivmod divmod_split_kv;
KernelHardwareInfo hw_info;
};
int block_idx = 0;
Params params;
CUTLASS_DEVICE
Sm100MlaPersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {}
template<class ProblemShape, class ClusterShape>
static Params to_underlying_arguments(
ProblemShape const& problem_shape, KernelHardwareInfo hw_info,
ClusterShape const& cluster_shape, int const& split_kv) {
using namespace cute;
// Get SM count if needed, otherwise use user supplied SM count
int sm_count = hw_info.sm_count;
if (sm_count <= 1 || sm_count % size<0>(cluster_shape) != 0) {
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
}
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
hw_info.sm_count = sm_count;
int num_m_blocks = size<0>(cluster_shape);
int num_blocks = num_m_blocks * get<3>(problem_shape) /* Batch */;
num_blocks *= split_kv; /* Maximum Split KV*/
return Params {
num_blocks,
{ num_m_blocks}, { get<3>(problem_shape) }, {split_kv},
hw_info
};
}
static dim3 get_grid_shape(Params const& params) {
dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1);
return grid;
}
CUTLASS_DEVICE
bool is_valid() {
return block_idx < params.num_blocks;
}
CUTLASS_DEVICE
auto get_block_coord() {
using namespace cute;
int block_decode = block_idx;
int m_block, bidb, n_split_kv;
params.divmod_m_block(block_decode, m_block, block_decode);
params.divmod_b(block_decode, bidb, block_decode);
params.divmod_split_kv(block_decode, n_split_kv, block_decode);
return make_coord(m_block, _0{}, bidb, n_split_kv);
}
CUTLASS_DEVICE
Sm100MlaPersistentTileScheduler& operator++() {
block_idx += gridDim.x;
return *this;
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::fmha::kernel
python/sglang/kernels/aot/csrc/attention/merge_attn_states.cu
+206
View File
@@ -0,0 +1,206 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <algorithm>
#include <optional>
#include "pytorch_extension_utils.h"
// Helper functions to convert between different data types
// (float, half, bfloat16) for the merge attention states kernel.
inline __device__ float to_float(float u) {
return u;
}
inline __device__ float to_float(half u) {
return __half2float(u);
}
inline __device__ float to_float(__nv_bfloat16 u) {
return __bfloat162float(u);
}
inline __device__ void from_float(float& d, float s) {
d = s;
}
inline __device__ void from_float(half& d, float s) {
d = __float2half(s);
}
inline __device__ void from_float(__nv_bfloat16& d, float s) {
d = __float2bfloat16(s);
}
// Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
template <typename scalar_t, const uint NUM_THREADS>
__global__ void merge_attn_states_kernel(
scalar_t* output,
float* output_lse,
const scalar_t* prefix_output,
const float* prefix_lse,
const scalar_t* suffix_output,
const float* suffix_lse,
const uint num_tokens,
const uint num_heads,
const uint head_size) {
using pack_128b_t = uint4;
const uint pack_size = 16 / sizeof(scalar_t);
const uint threads_per_head = head_size / pack_size;
const uint global_idx = blockIdx.x * NUM_THREADS + threadIdx.x;
const uint token_head_threads = num_tokens * num_heads * threads_per_head;
if (global_idx >= token_head_threads) return;
// global_idx -> token_idx + head_idx + pack_idx
const uint token_head_idx = global_idx / threads_per_head;
const uint pack_idx = global_idx % threads_per_head;
const uint token_idx = token_head_idx / num_heads;
const uint head_idx = token_head_idx % num_heads;
const uint pack_offset = pack_idx * pack_size; // (0~15)*8, etc.
const uint head_offset = token_idx * num_heads * head_size + head_idx * head_size;
const scalar_t* prefix_head_ptr = prefix_output + head_offset;
const scalar_t* suffix_head_ptr = suffix_output + head_offset;
scalar_t* output_head_ptr = output + head_offset;
// float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
// float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
float p_lse = prefix_lse[token_idx * num_heads + head_idx];
float s_lse = suffix_lse[token_idx * num_heads + head_idx];
p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;
s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;
const float max_lse = fmaxf(p_lse, s_lse);
p_lse = p_lse - max_lse;
s_lse = s_lse - max_lse;
const float p_se = expf(p_lse);
const float s_se = expf(s_lse);
const float out_se = p_se + s_se;
const float p_scale = p_se / out_se;
const float s_scale = s_se / out_se;
if (pack_offset < head_size) {
// Pack 128b load
pack_128b_t p_out_pack = reinterpret_cast<const pack_128b_t*>(prefix_head_ptr)[pack_offset / pack_size];
pack_128b_t s_out_pack = reinterpret_cast<const pack_128b_t*>(suffix_head_ptr)[pack_offset / pack_size];
pack_128b_t o_out_pack;
#pragma unroll
for (uint i = 0; i < pack_size; ++i) {
// Always use float for FMA to keep high precision.
// half(uint16_t), bfloat16, float -> float.
const float p_out_f = to_float(reinterpret_cast<const scalar_t*>(&p_out_pack)[i]);
const float s_out_f = to_float(reinterpret_cast<const scalar_t*>(&s_out_pack)[i]);
// fma: a * b + c = p_out_f * p_scale + (s_out_f * s_scale)
const float o_out_f = p_out_f * p_scale + (s_out_f * s_scale);
// float -> half(uint16_t), bfloat16, float.
from_float(reinterpret_cast<scalar_t*>(&o_out_pack)[i], o_out_f);
}
// Pack 128b storage
reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] = o_out_pack;
}
// We only need to write to output_lse once per head.
if (output_lse != nullptr && pack_idx == 0) {
float out_lse = logf(out_se) + max_lse;
output_lse[token_idx * num_heads + head_idx] = out_lse;
}
}
// The following macro is used to dispatch the conversion function based on
// the output data type. The FN is a macro that calls a function with
// template<typename scalar_t>.
#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn) \
{ \
if (scalar_dtype == at::ScalarType::Float) { \
fn(float); \
} else if (scalar_dtype == at::ScalarType::Half) { \
fn(half); \
} else if (scalar_dtype == at::ScalarType::BFloat16) { \
fn(__nv_bfloat16); \
} else { \
TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
} \
}
#define LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS) \
{ \
merge_attn_states_kernel<scalar_t, NUM_THREADS><<<grid, block, 0, stream>>>( \
reinterpret_cast<scalar_t*>(output.data_ptr()), \
reinterpret_cast<float*>(output_lse.data_ptr()), \
reinterpret_cast<scalar_t*>(prefix_output.data_ptr()), \
reinterpret_cast<float*>(prefix_lse.data_ptr()), \
reinterpret_cast<scalar_t*>(suffix_output.data_ptr()), \
reinterpret_cast<float*>(suffix_lse.data_ptr()), \
num_tokens, \
num_heads, \
head_size); \
}
/*@brief Merges the attention states from prefix and suffix
* into the output tensor. NUM_TOKENS: n, NUM_HEADS: h, HEAD_SIZE: d
*
* @param output [n,h,d] The output tensor to store the merged attention states.
* @param output_lse [h,d] Optional tensor to store the log-sum-exp values.
* @param prefix_output [n,h,d] The prefix attention states.
* @param prefix_lse [n,h] The log-sum-exp values for the prefix attention
* states.
* @param suffix_output [n,h,d] The suffix attention states.
* @param suffix_lse [n,h] The log-sum-exp values for the suffix attention
* states.
*/
template <typename scalar_t>
void merge_attn_states_launcher(
const at::Tensor& prefix_output, // [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
const at::Tensor& prefix_lse, // [NUM_TOKENS, NUM_HEADS]
const at::Tensor& suffix_output, // [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
const at::Tensor& suffix_lse, // [NUM_TOKENS, NUM_HEADS]
at::Tensor& output, // [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
at::Tensor& output_lse // [NUM_TOKENS, NUM_HEADS]
) {
constexpr uint NUM_THREADS = 128;
const uint num_tokens = output.size(0);
const uint num_heads = output.size(1);
const uint head_size = output.size(2);
const uint pack_size = 16 / sizeof(scalar_t);
TORCH_CHECK(head_size % pack_size == 0, "headsize must be multiple of pack_size:", pack_size);
// Process one pack elements per thread. for float, the
// pack_size is 4 for half/bf16, the pack_size is 8.
const uint threads_per_head = head_size / pack_size;
const uint total_threads = num_tokens * num_heads * threads_per_head;
dim3 block(NUM_THREADS);
dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);
const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
auto stream = at::cuda::getCurrentCUDAStream();
LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS);
}
#define CALL_MERGE_ATTN_STATES_LAUNCHER(scalar_t) \
{ \
merge_attn_states_launcher<scalar_t>(v_a, s_a, v_b, s_b, v_merged, s_merged); \
}
void merge_state_v2(
at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, at::Tensor v_merged, at::Tensor s_merged) {
// Input tensors must be contiguous
CHECK_INPUT(v_a); // v_a prefix_output (seq_len, num_heads, head_dim)
CHECK_INPUT(s_a); // s_a prefix_lse (seq_len, num_heads)
CHECK_INPUT(v_b); // v_b suffix_output (seq_len, num_heads, head_dim)
CHECK_INPUT(s_b); // s_b suffix_lse (seq_len, num_heads)
// v_merged output (seq_len, num_heads, head_dim)
// s_merged output_lse (seq_len, num_heads)
auto device = v_a.device();
CHECK_EQ(s_a.device(), device);
CHECK_EQ(v_b.device(), device);
CHECK_EQ(s_b.device(), device);
CHECK_DIM(3, v_a);
CHECK_DIM(2, s_a);
CHECK_DIM(3, v_b);
CHECK_DIM(2, s_b);
CHECK_SHAPE(v_a, v_b);
CHECK_SHAPE(s_a, s_b);
CHECK_EQ(v_a.size(0), s_a.size(0));
CHECK_EQ(v_a.size(1), s_b.size(1));
DISPATCH_BY_SCALAR_DTYPE(v_merged.dtype(), CALL_MERGE_ATTN_STATES_LAUNCHER);
}
python/sglang/kernels/aot/csrc/attention/vertical_slash_index.cu
+462
View File
@@ -0,0 +1,462 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.
// This file is for blocksparse attention utils cuda kernel.
#include <assert.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda.h>
#include <torch/all.h>
// Save the start index of each block in the given range into block_offset.
// Returns the updated block count.
__device__ int64_t save_blocks(
int* block_offset,
int64_t range_start,
int64_t range_end,
int64_t block_size,
int64_t input_block_count,
int64_t kv_seqlen) {
if (range_start >= kv_seqlen) {
return input_block_count;
}
if (range_end > kv_seqlen) {
range_end = kv_seqlen;
}
int64_t current_block_count = input_block_count;
for (int idx = range_start; idx < range_end; idx += block_size) {
block_offset[current_block_count++] = idx;
}
return current_block_count;
}
// CUDA kernel: convert sparse vertical/slash indices to block/column offsets.
__global__ void convert_vertical_slash_indexes_kernel(
const int* q_seqlens, // [BATCH, ]
const int* kv_seqlens, // [BATCH, ]
const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
int64_t N_HEADS,
int64_t N_ROWS,
int64_t BLOCK_SIZE_M,
int64_t BLOCK_SIZE_N,
int64_t NNZ_V,
int64_t NNZ_S,
bool causal // True for intra, False for succ
) {
const int batch_idx = blockIdx.y;
const int head_idx = blockIdx.x;
const int group_idx = blockIdx.z;
int64_t q_seqlen = q_seqlens[batch_idx];
int64_t kv_seqlen = kv_seqlens[batch_idx];
int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
int64_t start_m = block_idx_m * BLOCK_SIZE_M;
if (start_m >= q_seqlen) {
return;
}
int64_t end_m = start_m + BLOCK_SIZE_M;
vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
block_count += row_offset;
block_offset += row_offset * NNZ_S;
column_count += row_offset;
column_index += row_offset * NNZ_V;
bool has_slash = true;
int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
int64_t s = 0, v = 0;
int64_t v_idx = vertical_indexes[v++];
int64_t s_idx = slash_indexes[s++];
if (causal) {
while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
s_idx = slash_indexes[s++];
}
if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
} else {
while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
s_idx = slash_indexes[s++];
}
if (s_idx > end_m + kv_seqlen) has_slash = false;
s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
}
int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
if (!has_slash) {
if (causal) {
range_start = (kv_seqlen - q_seqlen) + end_m;
range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
} else {
range_start = kv_seqlen;
range_end = kv_seqlen + BLOCK_SIZE_N;
}
}
bool slash_finished = false;
while (1) {
if (v_idx < range_end) {
if (v_idx < range_start) {
column_index[tmp_col_cnt++] = v_idx;
}
if (v < NNZ_V) {
v_idx = vertical_indexes[v++];
} else {
if (causal)
v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
else
v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
}
} else {
if ((s < NNZ_S && causal) || (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
if (causal)
s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++], BLOCK_SIZE_M);
else
s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
} else {
if (v == NNZ_V || (v_idx > range_start && causal)) {
// add the last vertical if no more slash
if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
column_index[tmp_col_cnt++] = v_idx;
}
tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
break;
} else {
if (causal) {
range_start = (kv_seqlen - q_seqlen) + end_m;
range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
} else {
// if slash_finished but there are vertical left, save current
// blocks
tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
range_start = kv_seqlen;
range_end = kv_seqlen + BLOCK_SIZE_N;
}
slash_finished = true;
}
}
if (!slash_finished) {
if (s_idx > range_end + BLOCK_SIZE_M) {
tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
range_start = s_idx - BLOCK_SIZE_M;
range_end = s_idx;
} else if (s_idx > range_end) {
range_end += BLOCK_SIZE_M;
}
}
}
}
block_count[0] = tmp_blk_cnt;
column_count[0] = tmp_col_cnt;
}
// Host function: launches the kernel with 64 threads per block.
void convert_vertical_slash_indexes_64x64(
const int* q_seqlens, // [BATCH, ]
const int* kv_seqlens, // [BATCH, ]
const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
int64_t BATCH_SIZE,
int64_t N_HEADS,
int64_t N_ROWS,
int64_t BLOCK_SIZE_M,
int64_t BLOCK_SIZE_N,
int64_t NNZ_V,
int64_t NNZ_S,
bool causal) {
const int N_THREADS = 64;
const dim3 dimBlock((int32_t)N_THREADS);
const dim3 dimGrid(
(int32_t)N_HEADS, (int32_t)BATCH_SIZE, ((int32_t)N_ROWS + (int32_t)N_THREADS - 1) / (int32_t)N_THREADS);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
convert_vertical_slash_indexes_kernel<<<dimGrid, dimBlock, 0, stream>>>(
q_seqlens,
kv_seqlens,
vertical_indexes,
slash_indexes,
block_count,
block_offset,
column_count,
column_index,
N_HEADS,
N_ROWS,
BLOCK_SIZE_M,
BLOCK_SIZE_N,
NNZ_V,
NNZ_S,
causal);
}
// Host function: prepares tensor pointers and launches the CUDA kernel.
void convert_vertical_slash_indexes(
torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
torch::Tensor q_seqlens, // [BATCH, ]
torch::Tensor kv_seqlens, // [BATCH, ]
torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
int64_t context_size,
int64_t block_size_M,
int64_t block_size_N,
bool causal) {
cudaSetDevice(q_seqlens.get_device());
int64_t batch_size = slash_indexes.size(0);
int64_t num_heads = slash_indexes.size(1);
int64_t nnz_slash = slash_indexes.size(2);
int64_t nnz_vertical = vertical_indexes.size(2);
int64_t num_rows = (context_size + block_size_M - 1) / block_size_M;
convert_vertical_slash_indexes_64x64(
q_seqlens.data_ptr<int>(),
kv_seqlens.data_ptr<int>(),
vertical_indexes.data_ptr<int>(),
slash_indexes.data_ptr<int>(),
block_count.data_ptr<int>(),
block_offset.data_ptr<int>(),
column_count.data_ptr<int>(),
column_index.data_ptr<int>(),
batch_size,
num_heads,
num_rows,
block_size_M,
block_size_N,
nnz_vertical,
nnz_slash,
causal);
}
// --- mergehead kernels --- //
// Kernel: like above, but supports per-head variable NNZ_V/NNZ_S.
__global__ void convert_vertical_slash_indexes_kernel_mergehead(
const int* q_seqlens, // [BATCH, ]
const int* kv_seqlens, // [BATCH, ]
const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
const int* per_head_vertical_topkv,
const int* per_head_slash_topkv,
int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
int64_t N_HEADS,
int64_t N_ROWS,
int64_t BLOCK_SIZE_M,
int64_t BLOCK_SIZE_N,
int64_t NNZ_V,
int64_t NNZ_S,
bool causal // True for intra, False for succ
) {
const int batch_idx = blockIdx.y;
const int head_idx = blockIdx.x;
const int group_idx = blockIdx.z;
int64_t q_seqlen = q_seqlens[batch_idx];
int64_t kv_seqlen = kv_seqlens[batch_idx];
int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
int64_t start_m = block_idx_m * BLOCK_SIZE_M;
if (start_m >= q_seqlen) {
return;
}
int64_t end_m = start_m + BLOCK_SIZE_M;
vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
block_count += row_offset;
block_offset += row_offset * NNZ_S;
column_count += row_offset;
column_index += row_offset * NNZ_V;
// MergeHead: each head has it's unique max topk NNZ_V，NNZ_S. (NNZ_V，NNZ_S
// above is buffer size, use to compute offset)
NNZ_S = per_head_slash_topkv[head_idx];
NNZ_V = per_head_vertical_topkv[head_idx];
bool has_slash = true;
int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
int64_t s = 0, v = 0;
int64_t v_idx = vertical_indexes[v++];
int64_t s_idx = slash_indexes[s++];
if (causal) {
while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
s_idx = slash_indexes[s++];
}
if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
} else {
while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
s_idx = slash_indexes[s++];
}
if (s_idx > end_m + kv_seqlen) has_slash = false;
s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
}
int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
if (!has_slash) {
if (causal) {
range_start = (kv_seqlen - q_seqlen) + end_m;
range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
} else {
range_start = kv_seqlen;
range_end = kv_seqlen + BLOCK_SIZE_N;
}
}
bool slash_finished = false;
while (1) {
if (v_idx < range_end) {
if (v_idx < range_start) {
column_index[tmp_col_cnt++] = v_idx;
}
if (v < NNZ_V) {
v_idx = vertical_indexes[v++];
} else {
if (causal)
v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
else
v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
}
} else {
if ((s < NNZ_S && causal) || (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
if (causal)
s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++], BLOCK_SIZE_M);
else
s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
} else {
if (v == NNZ_V || (v_idx > range_start && causal)) {
// add the last vertical if no more slash
if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
column_index[tmp_col_cnt++] = v_idx;
}
tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
break;
} else {
if (causal) {
range_start = (kv_seqlen - q_seqlen) + end_m;
range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
} else {
// if slash_finished but there are vertical left, save current
// blocks
tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
range_start = kv_seqlen;
range_end = kv_seqlen + BLOCK_SIZE_N;
}
slash_finished = true;
}
}
if (!slash_finished) {
if (s_idx > range_end + BLOCK_SIZE_M) {
tmp_blk_cnt = save_blocks(block_offset, range_start, range_end, BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
range_start = s_idx - BLOCK_SIZE_M;
range_end = s_idx;
} else if (s_idx > range_end) {
range_end += BLOCK_SIZE_M;
}
}
}
}
block_count[0] = tmp_blk_cnt;
column_count[0] = tmp_col_cnt;
}
// Launch the mergehead kernel with 64 threads per block.
void convert_vertical_slash_indexes_64x64_mergehead(
const int* q_seqlens, // [BATCH, ]
const int* kv_seqlens, // [BATCH, ]
const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
int* per_head_vertical_topkv,
int* per_head_slash_topkv,
int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
int64_t BATCH_SIZE,
int64_t N_HEADS,
int64_t N_ROWS,
int64_t BLOCK_SIZE_M,
int64_t BLOCK_SIZE_N,
int64_t NNZ_V,
int64_t NNZ_S,
bool causal) {
const int N_THREADS = 64;
const dim3 dimBlock(N_THREADS);
const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
convert_vertical_slash_indexes_kernel_mergehead<<<dimGrid, dimBlock, 0, stream>>>(
q_seqlens,
kv_seqlens,
vertical_indexes,
slash_indexes,
per_head_vertical_topkv,
per_head_slash_topkv,
block_count,
block_offset,
column_count,
column_index,
N_HEADS,
N_ROWS,
BLOCK_SIZE_M,
BLOCK_SIZE_N,
NNZ_V,
NNZ_S,
causal);
}
// Host wrapper for mergehead kernel.
void convert_vertical_slash_indexes_mergehead(
torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
torch::Tensor q_seqlens, // [BATCH, ]
torch::Tensor kv_seqlens, // [BATCH, ]
torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
torch::Tensor vertical_indices_count, // [N_HEADS, ]
torch::Tensor slash_indices_count,
int64_t context_size,
int64_t block_size_M,
int64_t block_size_N,
bool causal) {
cudaSetDevice(q_seqlens.get_device());
int batch_size = slash_indexes.size(0);
int num_heads = slash_indexes.size(1);
int nnz_slash = slash_indexes.size(2);
int nnz_vertical = vertical_indexes.size(2);
int num_rows = (context_size + block_size_M - 1) / block_size_M;
convert_vertical_slash_indexes_64x64_mergehead(
q_seqlens.data_ptr<int>(),
kv_seqlens.data_ptr<int>(),
vertical_indexes.data_ptr<int>(),
slash_indexes.data_ptr<int>(),
vertical_indices_count.data_ptr<int>(),
slash_indices_count.data_ptr<int>(),
block_count.data_ptr<int>(),
block_offset.data_ptr<int>(),
column_count.data_ptr<int>(),
column_index.data_ptr<int>(),
batch_size,
num_heads,
num_rows,
block_size_M,
block_size_N,
nnz_vertical,
nnz_slash,
causal);
}
python/sglang/kernels/aot/csrc/common_extension.cc
+475
View File
@@ -0,0 +1,475 @@
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/core/dispatch/Dispatcher.h>
#include <torch/all.h>
#include <torch/library.h>
#include "sgl_kernel_ops.h"
TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
/*
* From csrc/allreduce
*/
m.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta);
m.def("register_graph_buffers", &register_graph_buffers);
m.def("dispose", &dispose);
m.def("meta_size", &meta_size);
m.def("register_buffer", &register_buffer);
m.def(
"init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
"int rank, bool full_nvlink) -> int");
m.impl("init_custom_ar", torch::kCUDA, &init_custom_ar);
m.def(
"all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
"int reg_buffer_sz_bytes) -> ()");
m.impl("all_reduce", torch::kCUDA, &all_reduce);
/*
* From csrc/attention
*/
m.def("merge_state_v2(Tensor v_a, Tensor s_a, Tensor v_b, Tensor s_b, Tensor! v_merged, Tensor! s_merged) -> ()");
m.impl("merge_state_v2", torch::kCUDA, &merge_state_v2);
m.def(
"cutlass_mla_decode(Tensor! out, Tensor q_nope, Tensor q_pe, Tensor kv_c_and_k_pe_cache, Tensor seq_lens, Tensor "
"page_table, Tensor! workspace, float sm_scale, int num_kv_splits) -> ()");
m.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode);
m.def("cutlass_mla_get_workspace_size", &cutlass_mla_get_workspace_size);
/*
* From csrc/infllm_v2
*/
m.def(
"infllm_v2_max_pooling_1d_varlen(Tensor input, Tensor! output, Tensor cu_seqlens_q, Tensor cu_seqlens_k, "
"Tensor cache_lens, int max_seqlen_q, int max_seqlen_k, int kernel_size, int stride, int padding, "
"int block_size, int local_blocks, int init_blocks, int total_q) -> ()");
m.impl("infllm_v2_max_pooling_1d_varlen", torch::kCUDA, &infllm_v2_max_pooling_1d_varlen);
/*
* From csrc/elementwise
*/
m.def("rmsnorm(Tensor! output, Tensor input, Tensor weight, float eps, bool enable_pdl) -> ()");
m.impl("rmsnorm", torch::kCUDA, &rmsnorm);
m.def("fused_add_rmsnorm(Tensor! input, Tensor! residual, Tensor weight, float eps, bool enable_pdl) -> ()");
m.impl("fused_add_rmsnorm", torch::kCUDA, &sgl_fused_add_rmsnorm);
m.def("gemma_rmsnorm(Tensor! output, Tensor input, Tensor weight, float eps, bool enable_pdl) -> ()");
m.impl("gemma_rmsnorm", torch::kCUDA, &gemma_rmsnorm);
m.def("gemma_fused_add_rmsnorm(Tensor! input, Tensor! residual, Tensor weight, float eps, bool enable_pdl) -> ()");
m.impl("gemma_fused_add_rmsnorm", torch::kCUDA, &gemma_fused_add_rmsnorm);
m.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
m.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("gelu_tanh_and_mul", torch::kCUDA, &gelu_tanh_and_mul);
m.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);
m.def(
"rotary_embedding(Tensor positions, Tensor! query,"
" Tensor!? key, int head_size,"
" Tensor cos_sin_cache, bool is_neox) -> ()");
m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
m.def("copy_to_gpu_no_ce(Tensor input, Tensor! output) -> ()");
m.impl("copy_to_gpu_no_ce", torch::kCUDA, &copy_to_gpu_no_ce);
m.def("concat_mla_k(Tensor! k, Tensor k_nope, Tensor k_rope) -> ()");
m.impl("concat_mla_k", torch::kCUDA, &concat_mla_k);
m.def("concat_mla_absorb_q(Tensor a, Tensor b, Tensor! out) -> ()");
m.impl("concat_mla_absorb_q", torch::kCUDA, &concat_mla_absorb_q);
m.def("fast_topk(Tensor score, Tensor indices, Tensor lengths, Tensor? row_starts) -> ()");
m.impl("fast_topk", torch::kCUDA, &fast_topk_interface);
m.def(
"fast_topk_transform_fused(Tensor score, Tensor lengths, Tensor dst_page_table, Tensor src_page_table, Tensor "
"cu_seqlens_q, Tensor? row_starts) -> ()");
m.impl("fast_topk_transform_fused", torch::kCUDA, &fast_topk_transform_interface);
m.def(
"fast_topk_transform_ragged_fused(Tensor score, Tensor lengths, Tensor topk_indices_ragged, Tensor "
"topk_indices_offset, Tensor ? row_starts) -> ()");
m.impl("fast_topk_transform_ragged_fused", torch::kCUDA, &fast_topk_transform_ragged_interface);
/*
* From csrc/gemm
*/
m.def("awq_dequantize(Tensor qweight, Tensor scales, Tensor qzeros) -> Tensor");
m.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
m.def(
"int8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? "
"bias) -> Tensor");
m.impl("int8_scaled_mm", torch::kCUDA, &int8_scaled_mm);
m.def(
"fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? "
"bias) -> Tensor");
m.impl("fp8_scaled_mm", torch::kCUDA, &fp8_scaled_mm);
m.def(
"sgl_per_token_group_quant_8bit(Tensor input, Tensor! output_q, Tensor! output_s, int group_size,"
" float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()");
m.impl("sgl_per_token_group_quant_8bit", torch::kCUDA, &sgl_per_token_group_quant_8bit);
m.def(
"sgl_per_token_group_quant_8bit_v2(Tensor input, Tensor! output_q, Tensor! output_s, int group_size,"
" float eps, float fp8_min, float fp8_max, bool scale_ue8m0, bool fuse_silu_and_mul, Tensor? masked_m) -> ()");
m.impl("sgl_per_token_group_quant_8bit_v2", torch::kCUDA, &sgl_per_token_group_quant_8bit_v2);
m.def("sgl_per_token_quant_fp8(Tensor input, Tensor! output_q, Tensor! output_s) -> ()");
m.impl("sgl_per_token_quant_fp8", torch::kCUDA, &sgl_per_token_quant_fp8);
/*
* From csrc/gemm/gptq
*/
m.def(
"gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, Tensor b_gptq_scales, Tensor b_g_idx, bool "
"use_shuffle, int bit) -> Tensor");
m.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);
m.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()");
m.impl("gptq_shuffle", torch::kCUDA, &gptq_shuffle);
/*
* From csrc/moe
*/
m.def(
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! cumsum_buffer, bool "
"pad_sorted_token_ids, bool ignore_invalid_expert) -> ()");
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, float "
"moe_softcapping, Tensor? correction_bias) -> ()");
m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
m.def(
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, Tensor? "
"correction_bias) -> ()");
m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
m.def("moe_sum_reduce(Tensor input, Tensor output, float routed_scaling_factor) -> ()");
m.impl("moe_sum_reduce", torch::kCUDA, &moe_sum_reduce);
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
m.impl("moe_sum", torch::kCUDA, &moe_sum);
// moe_fused_gate / kimi_k2_moe_fused_gate (AOT) retired: the CUDA gate/topk path
// now routes through the unified Triton router
// (python/sglang/kernels/ops/moe/moe_fused_gate.py).
m.def(
"fp8_blockwise_scaled_grouped_mm(Tensor output, Tensor a_ptrs, Tensor b_ptrs, Tensor out_ptrs, Tensor "
"a_scales_ptrs, Tensor b_scales_ptrs, Tensor a, Tensor b, Tensor scales_a, Tensor scales_b, Tensor "
"stride_a, Tensor stride_b, Tensor stride_c, Tensor layout_sfa, Tensor layout_sfb, Tensor problem_sizes, Tensor "
"expert_offsets, Tensor workspace) -> ()");
m.impl("fp8_blockwise_scaled_grouped_mm", torch::kCUDA, &fp8_blockwise_scaled_grouped_mm);
m.def(
"prepare_moe_input(Tensor topk_ids, Tensor expert_offsets, Tensor? blockscale_offsets, Tensor problem_sizes1,"
" Tensor problem_sizes2, Tensor input_permutation, Tensor output_permutation, int num_experts, int n, int k) -> "
"()");
m.impl("prepare_moe_input", torch::kCUDA, &prepare_moe_input);
m.def("shuffle_rows(Tensor input, Tensor dst2src_map, Tensor output) -> ()");
m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows);
m.def("apply_shuffle_mul_sum(Tensor input, Tensor output, Tensor permutation, Tensor? factors) -> ()");
m.impl("apply_shuffle_mul_sum", torch::kCUDA, &apply_shuffle_mul_sum);
// DeepSeek-V4 fused norm + rope
m.def(
"dsv4_fused_q_norm_rope(Tensor q_input, Tensor! q_output, Tensor freqs_cis, Tensor positions, float eps) -> ()");
m.impl("dsv4_fused_q_norm_rope", torch::kCUDA, &dsv4_fused_q_norm_rope);
m.def(
"dsv4_fused_k_norm_rope_flashmla(Tensor kv, Tensor kv_weight, Tensor freqs_cis, Tensor positions, "
"Tensor out_loc, Tensor! kvcache, float eps, int page_size) -> ()");
m.impl("dsv4_fused_k_norm_rope_flashmla", torch::kCUDA, &dsv4_fused_k_norm_rope_flashmla);
m.def(
"dsv4_fused_q_indexer_rope_hadamard_quant(Tensor q_input, Tensor! q_fp8, Tensor weight, "
"Tensor! weights_out, float weight_scale, Tensor freqs_cis, Tensor positions) -> ()");
m.impl("dsv4_fused_q_indexer_rope_hadamard_quant", torch::kCUDA, &dsv4_fused_q_indexer_rope_hadamard_quant);
m.def(
"fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
"int num_heads_k, int num_heads_v, int head_dim, float eps, "
"Tensor q_weight, Tensor k_weight, float base, "
"bool is_neox, Tensor position_ids, float factor, float low, float high, float attention_factor, int rotary_dim) "
"-> ()");
m.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
/*
* From csrc/moe/cutlass_moe/w4a8
*/
m.def(
"get_cutlass_w4a8_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, "
" Tensor! problem_sizes1, Tensor! problem_sizes2, "
" Tensor! input_permutation, "
" Tensor! output_permutation, int num_experts, "
" int n, int k) -> ()");
m.impl("get_cutlass_w4a8_moe_mm_data", torch::kCUDA, &get_cutlass_w4a8_moe_mm_data);
m.def(
"cutlass_w4a8_moe_mm(Tensor! d, Tensor a, Tensor b, "
" Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
" Tensor problem_sizes, Tensor a_strides, "
" Tensor b_strides, Tensor d_strides, Tensor s_strides,"
" int chunk_size, int topk) -> ()");
m.impl("cutlass_w4a8_moe_mm", torch::kCUDA, &cutlass_w4a8_moe_mm);
/*
* From csrc/speculative
*/
m.def(
"tree_speculative_sampling_target_only(Tensor! predicts, Tensor! accept_index, Tensor! accept_token_num, "
"Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, "
"Tensor uniform_samples, Tensor uniform_samples_for_final_sampling, Tensor target_probs, Tensor draft_probs, "
"float threshold_single, float threshold_acc, "
"bool deterministic) -> ()");
m.impl("tree_speculative_sampling_target_only", torch::kCUDA, &tree_speculative_sampling_target_only);
m.def(
"verify_tree_greedy(Tensor! predicts, Tensor! accept_index, Tensor! accept_token_num, "
"Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, "
"Tensor target_predict) -> ()");
m.impl("verify_tree_greedy", torch::kCUDA, &verify_tree_greedy);
m.def(
"reconstruct_indices_from_tree_mask(Tensor tree_mask, Tensor verified_seq_len, Tensor positions, "
"Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, "
"int batch_size, int draft_token_num) -> ()");
m.impl("reconstruct_indices_from_tree_mask", torch::kCUDA, &reconstruct_indices_from_tree_mask);
m.def(
"build_tree_kernel_efficient(Tensor parent_list, Tensor selected_index, Tensor verified_seq_len, "
"Tensor! tree_mask, Tensor! positions, Tensor! retrive_index, Tensor! retrive_next_token, "
"Tensor! retrive_next_sibling, int topk, int depth, int draft_token_num, int tree_mask_mode) -> "
"()");
m.impl("build_tree_kernel_efficient", torch::kCUDA, &build_tree_kernel_efficient);
m.def(
"segment_packbits(Tensor x, Tensor input_indptr, Tensor output_indptr, Tensor! y, int batch_size, "
"int cuda_stream) -> ()");
m.impl("segment_packbits", torch::kCUDA, &segment_packbits);
/*
* From csrc/kvcacheio
*/
m.def(
"transfer_kv_per_layer(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int item_size, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer", torch::kCUDA, &transfer_kv_per_layer);
m.def(
"transfer_kv_per_layer_pf_lf(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int layer_id, int item_size, int src_layout_dim, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_pf_lf", torch::kCUDA, &transfer_kv_per_layer_pf_lf);
m.def(
"transfer_kv_per_layer_ph_lf(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int layer_id, int item_size, int src_layout_dim, int page_size, int head_num, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_ph_lf", torch::kCUDA, &transfer_kv_per_layer_ph_lf);
m.def(
"transfer_kv_all_layer(Tensor src_k_layers, Tensor dst_k_layers, Tensor src_v_layers, Tensor dst_v_layers, "
"Tensor src_indices, Tensor dst_indices, int item_size, int num_layers, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer", torch::kCUDA, &transfer_kv_all_layer);
m.def(
"transfer_kv_all_layer_lf_pf(Tensor src_k_layers, Tensor dst_k, Tensor src_v_layers, Tensor dst_v, "
"Tensor src_indices, Tensor dst_indices, int item_size, int dst_layout_dim, int num_layers, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_lf_pf", torch::kCUDA, &transfer_kv_all_layer_lf_pf);
m.def(
"transfer_kv_all_layer_lf_ph(Tensor src_k_layers, Tensor dst_k, Tensor src_v_layers, Tensor dst_v, "
"Tensor src_indices, Tensor dst_indices, int item_size, int dst_layout_dim, int num_layers, int page_size, int "
"head_num, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_lf_ph", torch::kCUDA, &transfer_kv_all_layer_lf_ph);
m.def(
"transfer_kv_per_layer_mla(Tensor src, Tensor dst, Tensor src_indices, Tensor dst_indices, int item_size, int "
"block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_mla", torch::kCUDA, &transfer_kv_per_layer_mla);
m.def(
"transfer_kv_per_layer_mla_pf_lf(Tensor src, Tensor dst, Tensor src_indices, Tensor dst_indices, int layer_id, "
"int item_size, int src_layout_dim, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_mla_pf_lf", torch::kCUDA, &transfer_kv_per_layer_mla_pf_lf);
m.def(
"transfer_kv_all_layer_mla(Tensor src_layers, Tensor dst_layers, Tensor src_indices, Tensor dst_indices, int "
"item_size, int num_layers, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_mla", torch::kCUDA, &transfer_kv_all_layer_mla);
m.def(
"transfer_kv_all_layer_mla_lf_pf(Tensor src_layers, Tensor dst, Tensor src_indices, Tensor dst_indices, "
"int item_size, int dst_layout_dim, int num_layers, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_mla_lf_pf", torch::kCUDA, &transfer_kv_all_layer_mla_lf_pf);
m.def(
"transfer_kv_direct(Tensor[] src_layers, Tensor[] dst_layers, Tensor src_indices, Tensor dst_indices, int "
"page_size) -> ()");
m.impl("transfer_kv_direct", torch::kCUDA, &transfer_kv_direct);
m.def(
"transfer_kv_per_layer_direct_pf_lf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int layer_id, int page_size)->() ");
m.impl("transfer_kv_per_layer_direct_pf_lf", torch::kCUDA, &transfer_kv_per_layer_direct_pf_lf);
m.def(
"transfer_kv_all_layer_direct_lf_pf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int page_size) ->() ");
m.impl("transfer_kv_all_layer_direct_lf_pf", torch::kCUDA, &transfer_kv_all_layer_direct_lf_pf);
/*
* From csrc/memory
*/
m.def("weak_ref_tensor(Tensor tensor) -> Tensor");
m.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
/*
* From FlashInfer
*/
m.def("top_k_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_k_arr, int top_k_val) -> ()");
m.impl("top_k_renorm_probs", torch::kCUDA, &top_k_renorm_probs);
m.def("top_p_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_p_arr, float top_p_val) -> ()");
m.impl("top_p_renorm_probs", torch::kCUDA, &top_p_renorm_probs);
/*
* From Sparse Flash Attention
*/
m.def(
"fwd_sparse(Tensor! q, Tensor k, Tensor v, "
"Tensor block_count, Tensor block_offset, Tensor column_count, Tensor column_index, "
"Tensor!? out, Tensor? alibi_slopes, "
"float p_dropout, float softmax_scale, bool is_causal, "
"float softcap, bool return_softmax, Generator? gen)"
"-> Tensor[]");
m.impl("fwd_sparse", torch::kCUDA, &flash::mha_fwd_sparse);
m.def(
"varlen_fwd_sparse(Tensor! q, Tensor k, Tensor v, "
"Tensor block_count, Tensor block_offset, Tensor column_count, Tensor column_index, "
"Tensor!? out, Tensor cu_seqlens_q, "
"Tensor cu_seqlens_k, Tensor? seqused_k, Tensor? alibi_slopes, "
"int max_seqlen_q, int max_seqlen_k, float p_dropout, float softmax_scale, bool zero_tensors, "
"bool is_causal, float softcap, bool return_softmax, "
"Generator? gen) -> Tensor[]");
m.impl("varlen_fwd_sparse", torch::kCUDA, &flash::mha_varlen_fwd_sparse);
// Sparse Attention utils
m.def(
"convert_vertical_slash_indexes("
" Tensor! block_count, Tensor! block_offset, "
" Tensor! column_count, Tensor! column_index, "
" Tensor q_seqlens, Tensor q_seqlens, "
" Tensor vertical_indexes, Tensor slash_indexes, "
" int context_size, int block_size_M, int block_size_N, "
" bool causal) -> ()");
m.impl("convert_vertical_slash_indexes", torch::kCUDA, &convert_vertical_slash_indexes);
m.def(
"convert_vertical_slash_indexes_mergehead("
" Tensor! block_count, Tensor! block_offset, "
" Tensor! column_count, Tensor! column_index, "
" Tensor q_seqlens, Tensor q_seqlens, "
" Tensor vertical_indexes, Tensor slash_indexes, "
" Tensor vertical_indices_count, Tensor slash_indices_count, "
" int context_size, int block_size_M, int block_size_N, "
" bool causal) -> ()");
m.impl("convert_vertical_slash_indexes_mergehead", torch::kCUDA, &convert_vertical_slash_indexes_mergehead);
/*
* From csrc/grammar
*/
m.def("apply_token_bitmask_inplace_cuda(Tensor logits, Tensor bitmask, Tensor? indices=None) -> ()");
m.impl("apply_token_bitmask_inplace_cuda", &ApplyTokenBitmaskInplace);
/*
* From csrc/quantization/gguf
*/
m.def(
"ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? "
"dtype) -> Tensor");
m.impl("ggml_dequantize", torch::kCUDA, &ggml_dequantize);
m.def(
"ggml_mul_mat_vec_a8(Tensor W, Tensor X, int type, SymInt row) "
"-> Tensor");
m.impl("ggml_mul_mat_vec_a8", torch::kCUDA, &ggml_mul_mat_vec_a8);
m.def("ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor");
m.impl("ggml_mul_mat_a8", torch::kCUDA, &ggml_mul_mat_a8);
m.def(
"ggml_moe_a8(Tensor X, Tensor W, "
"Tensor sorted_token_ids, Tensor expert_ids, Tensor "
"num_tokens_post_padded, "
"int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
m.impl("ggml_moe_a8", torch::kCUDA, &ggml_moe_a8);
m.def(
"ggml_moe_a8_vec(Tensor X, Tensor W, "
"Tensor topk_ids, int top_k, "
"int type, SymInt row, SymInt tokens) -> Tensor");
m.impl("ggml_moe_a8_vec", torch::kCUDA, &ggml_moe_a8_vec);
m.def("ggml_moe_get_block_size(int type) -> int");
m.impl("ggml_moe_get_block_size", torch::kCUDA, &ggml_moe_get_block_size);
/*
* From csrc/mamba
*/
m.def(
"causal_conv1d_update(Tensor! x,"
"Tensor! conv_state,"
"Tensor! weight,"
"Tensor? bias_,"
"bool silu_activation,"
"Tensor? cache_seqlens_,"
"Tensor? conv_state_indices,"
"int pad_slot_id) -> ()");
m.impl("causal_conv1d_update", torch::kCUDA, &causal_conv1d_update);
m.def(
"causal_conv1d_fwd(Tensor! x, Tensor! weight,"
"Tensor? bias_,"
"Tensor!? conv_states,"
"Tensor? query_start_loc,"
"Tensor? cache_indices,"
"Tensor? has_initial_state,"
"bool silu_activation,"
"int pad_slot_id) -> ()");
m.impl("causal_conv1d_fwd", torch::kCUDA, &causal_conv1d_fwd);
/*
* From csrc/expert_sepcialization
*/
m.def(
"es_fp8_blockwise_scaled_grouped_mm(Tensor output, Tensor a, Tensor b, Tensor scales_a, Tensor scales_b, Tensor "
"stride_a, Tensor stride_b, Tensor stride_d, Tensor problem_sizes, Tensor expert_offsets, Tensor workspace) -> "
"()");
m.impl("es_fp8_blockwise_scaled_grouped_mm", &es_fp8_blockwise_scaled_grouped_mm);
m.def(
"es_sm100_mxfp8_blockscaled_grouped_mm(Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor d, Tensor "
"problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets) -> ()");
m.impl("es_sm100_mxfp8_blockscaled_grouped_mm", &es_sm100_mxfp8_blockscaled_grouped_mm);
m.def(
"es_sm100_mxfp8_blockscaled_grouped_quant(Tensor input, Tensor problem_sizes, Tensor expert_offsets, Tensor "
"blockscale_offsets, Tensor quant_output, Tensor scale_factor) -> () ");
m.impl("es_sm100_mxfp8_blockscaled_grouped_quant", &es_sm100_mxfp8_blockscaled_grouped_quant);
}
REGISTER_EXTENSION(common_ops)
python/sglang/kernels/aot/csrc/common_extension_musa.cc
+315
View File
@@ -0,0 +1,315 @@
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/core/dispatch/Dispatcher.h>
#include <torch/library.h>
#include "sgl_kernel_ops.h"
#include "torch_musa/csrc/aten/musa/MUSAContext.h"
TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
/*
* From csrc/allreduce
*/
m.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta);
m.def("register_graph_buffers", &register_graph_buffers);
m.def("dispose", &dispose);
m.def("meta_size", &meta_size);
m.def("register_buffer", &register_buffer);
m.def(
"init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
"int rank, bool full_nvlink) -> int");
m.impl("init_custom_ar", torch::kMUSA, &init_custom_ar);
m.def(
"all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
"int reg_buffer_sz_bytes) -> ()");
m.impl("all_reduce", torch::kMUSA, &all_reduce);
/*
* From csrc/attention
*/
m.def("merge_state_v2(Tensor v_a, Tensor s_a, Tensor v_b, Tensor s_b, Tensor! v_merged, Tensor! s_merged) -> ()");
m.impl("merge_state_v2", torch::kMUSA, &merge_state_v2);
/*
* From csrc/elementwise
*/
m.def("rmsnorm(Tensor! output, Tensor input, Tensor weight, float eps, bool enable_pdl) -> ()");
m.impl("rmsnorm", torch::kMUSA, &rmsnorm);
m.def("fused_add_rmsnorm(Tensor! input, Tensor! residual, Tensor weight, float eps, bool enable_pdl) -> ()");
m.impl("fused_add_rmsnorm", torch::kMUSA, &musa_fused_add_rms_norm);
m.def("gemma_rmsnorm(Tensor! output, Tensor input, Tensor weight, float eps, bool enable_pdl) -> ()");
m.impl("gemma_rmsnorm", torch::kMUSA, &gemma_rmsnorm);
m.def("gemma_fused_add_rmsnorm(Tensor! input, Tensor! residual, Tensor weight, float eps, bool enable_pdl) -> ()");
m.impl("gemma_fused_add_rmsnorm", torch::kMUSA, &gemma_fused_add_rmsnorm);
m.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("silu_and_mul", torch::kMUSA, &silu_and_mul);
m.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("gelu_tanh_and_mul", torch::kMUSA, &gelu_tanh_and_mul);
m.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("gelu_and_mul", torch::kMUSA, &gelu_and_mul);
m.def("concat_mla_k(Tensor! k, Tensor k_nope, Tensor k_rope) -> ()");
m.impl("concat_mla_k", torch::kMUSA, &concat_mla_k);
m.def(
"rotary_embedding(Tensor positions, Tensor! query,"
" Tensor!? key, int head_size,"
" Tensor cos_sin_cache, bool is_neox) -> ()");
m.impl("rotary_embedding", torch::kMUSA, &rotary_embedding);
/*
* From csrc/gemm
*/
m.def("awq_dequantize(Tensor qweight, Tensor scales, Tensor qzeros) -> Tensor");
m.impl("awq_dequantize", torch::kMUSA, &awq_dequantize);
m.def(
"sgl_per_token_group_quant_8bit(Tensor input, Tensor output_q, Tensor output_s, int group_size,"
" float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()");
m.impl("sgl_per_token_group_quant_8bit", torch::kMUSA, &sgl_per_token_group_quant_8bit);
m.def(
"sgl_per_token_group_quant_8bit_v2(Tensor input, Tensor output_q, Tensor output_s, int group_size,"
" float eps, float fp8_min, float fp8_max, bool scale_ue8m0, bool fuse_silu_and_mul, Tensor? masked_m) -> ()");
m.impl("sgl_per_token_group_quant_8bit_v2", torch::kMUSA, &sgl_per_token_group_quant_8bit_v2);
m.def("sgl_per_token_quant_fp8(Tensor input, Tensor output_q, Tensor output_s) -> ()");
m.impl("sgl_per_token_quant_fp8", torch::kMUSA, &sgl_per_token_quant_fp8);
/*
* From csrc/moe
*/
m.def(
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! cumsum_buffer, bool "
"pad_sorted_token_ids, bool ignore_invalid_expert) -> ()");
m.impl("moe_align_block_size", torch::kMUSA, &moe_align_block_size);
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, float "
"moe_softcapping, Tensor? correction_bias) -> ()");
m.impl("topk_softmax", torch::kMUSA, &topk_softmax);
m.def("moe_sum_reduce(Tensor input, Tensor output, float routed_scaling_factor) -> ()");
m.impl("moe_sum_reduce", torch::kMUSA, &moe_sum_reduce);
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
m.impl("moe_sum", torch::kMUSA, &moe_sum);
// moe_fused_gate / kimi_k2_moe_fused_gate (AOT gate kernels) retired: gate/topk
// is consolidated onto the unified Triton router (sglang issue #26771). sglang's
// MUSA path uses `mate.moe_fused_gate`, so dropping the sgl_kernel MUSA op here
// has no runtime impact.
/*
* From csrc/speculative
*/
m.def(
"tree_speculative_sampling_target_only(Tensor! predicts, Tensor! accept_index, Tensor! accept_token_num, "
"Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, "
"Tensor uniform_samples, Tensor uniform_samples_for_final_sampling, Tensor target_probs, Tensor draft_probs, "
"float threshold_single, float threshold_acc, "
"bool deterministic) -> ()");
m.impl("tree_speculative_sampling_target_only", torch::kMUSA, &tree_speculative_sampling_target_only);
m.def(
"verify_tree_greedy(Tensor! predicts, Tensor! accept_index, Tensor! accept_token_num, "
"Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, "
"Tensor target_predict) -> ()");
m.impl("verify_tree_greedy", torch::kMUSA, &verify_tree_greedy);
m.def(
"reconstruct_indices_from_tree_mask(Tensor tree_mask, Tensor verified_seq_len, Tensor positions, "
"Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, "
"int batch_size, int draft_token_num) -> ()");
m.impl("reconstruct_indices_from_tree_mask", torch::kMUSA, &reconstruct_indices_from_tree_mask);
m.def(
"build_tree_kernel_efficient(Tensor parent_list, Tensor selected_index, Tensor verified_seq_len, "
"Tensor! tree_mask, Tensor! positions, Tensor! retrive_index, Tensor! retrive_next_token, "
"Tensor! retrive_next_sibling, int topk, int depth, int draft_token_num, int tree_mask_mode) -> "
"()");
m.impl("build_tree_kernel_efficient", torch::kMUSA, &build_tree_kernel_efficient);
/*
* From csrc/grammar
*/
m.def("apply_token_bitmask_inplace_cuda(Tensor logits, Tensor bitmask, Tensor? indices=None) -> ()");
m.impl("apply_token_bitmask_inplace_cuda", &ApplyTokenBitmaskInplace);
/*
* From csrc/quantization/gguf
*/
m.def(
"ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? "
"dtype) -> Tensor");
m.impl("ggml_dequantize", torch::kMUSA, &ggml_dequantize);
m.def(
"ggml_mul_mat_vec_a8(Tensor W, Tensor X, int type, SymInt row) "
"-> Tensor");
m.impl("ggml_mul_mat_vec_a8", torch::kMUSA, &ggml_mul_mat_vec_a8);
m.def("ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor");
m.impl("ggml_mul_mat_a8", torch::kMUSA, &ggml_mul_mat_a8);
m.def(
"ggml_moe_a8(Tensor X, Tensor W, "
"Tensor sorted_token_ids, Tensor expert_ids, Tensor "
"num_tokens_post_padded, "
"int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
m.impl("ggml_moe_a8", torch::kMUSA, &ggml_moe_a8);
m.def(
"ggml_moe_a8_vec(Tensor X, Tensor W, "
"Tensor topk_ids, int top_k, "
"int type, SymInt row, SymInt tokens) -> Tensor");
m.impl("ggml_moe_a8_vec", torch::kMUSA, &ggml_moe_a8_vec);
m.def("ggml_moe_get_block_size(int type) -> int");
m.impl("ggml_moe_get_block_size", torch::kMUSA, &ggml_moe_get_block_size);
/*
* From csrc/kvcacheio
*/
m.def(
"transfer_kv_per_layer(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int item_size, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer", torch::kMUSA, &transfer_kv_per_layer);
m.def(
"transfer_kv_per_layer_pf_lf(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int layer_id, int item_size, int src_layout_dim, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_pf_lf", torch::kMUSA, &transfer_kv_per_layer_pf_lf);
m.def(
"transfer_kv_per_layer_ph_lf(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int layer_id, int item_size, int src_layout_dim, int page_size, int head_num, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_ph_lf", torch::kMUSA, &transfer_kv_per_layer_ph_lf);
m.def(
"transfer_kv_all_layer(Tensor src_k_layers, Tensor dst_k_layers, Tensor src_v_layers, Tensor dst_v_layers, "
"Tensor src_indices, Tensor dst_indices, int item_size, int num_layers, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer", torch::kMUSA, &transfer_kv_all_layer);
m.def(
"transfer_kv_all_layer_lf_pf(Tensor src_k_layers, Tensor dst_k, Tensor src_v_layers, Tensor dst_v, "
"Tensor src_indices, Tensor dst_indices, int item_size, int dst_layout_dim, int num_layers, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_lf_pf", torch::kMUSA, &transfer_kv_all_layer_lf_pf);
m.def(
"transfer_kv_all_layer_lf_ph(Tensor src_k_layers, Tensor dst_k, Tensor src_v_layers, Tensor dst_v, "
"Tensor src_indices, Tensor dst_indices, int item_size, int dst_layout_dim, int num_layers, int page_size, int "
"head_num, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_lf_ph", torch::kMUSA, &transfer_kv_all_layer_lf_ph);
m.def(
"transfer_kv_per_layer_mla(Tensor src, Tensor dst, Tensor src_indices, Tensor dst_indices, int item_size, int "
"block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_mla", torch::kMUSA, &transfer_kv_per_layer_mla);
m.def(
"transfer_kv_per_layer_mla_pf_lf(Tensor src, Tensor dst, Tensor src_indices, Tensor dst_indices, int layer_id, "
"int item_size, int src_layout_dim, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_mla_pf_lf", torch::kMUSA, &transfer_kv_per_layer_mla_pf_lf);
m.def(
"transfer_kv_all_layer_mla(Tensor src_layers, Tensor dst_layers, Tensor src_indices, Tensor dst_indices, int "
"item_size, int num_layers, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_mla", torch::kMUSA, &transfer_kv_all_layer_mla);
m.def(
"transfer_kv_all_layer_mla_lf_pf(Tensor src_layers, Tensor dst, Tensor src_indices, Tensor dst_indices, "
"int item_size, int dst_layout_dim, int num_layers, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_mla_lf_pf", torch::kMUSA, &transfer_kv_all_layer_mla_lf_pf);
m.def(
"transfer_kv_direct(Tensor[] src_layers, Tensor[] dst_layers, Tensor src_indices, Tensor dst_indices, int "
"page_size) -> ()");
m.impl("transfer_kv_direct", torch::kMUSA, &transfer_kv_direct);
m.def(
"transfer_kv_per_layer_direct_pf_lf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int layer_id, int page_size)->() ");
m.impl("transfer_kv_per_layer_direct_pf_lf", torch::kMUSA, &transfer_kv_per_layer_direct_pf_lf);
m.def(
"transfer_kv_all_layer_direct_lf_pf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int page_size) ->() ");
m.impl("transfer_kv_all_layer_direct_lf_pf", torch::kMUSA, &transfer_kv_all_layer_direct_lf_pf);
/*
* From FlashInfer
*/
m.def("top_k_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_k_arr, int top_k_val) -> ()");
m.impl("top_k_renorm_probs", torch::kMUSA, &top_k_renorm_probs);
m.def("top_p_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_p_arr, float top_p_val) -> ()");
m.impl("top_p_renorm_probs", torch::kMUSA, &top_p_renorm_probs);
m.def(
"min_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_min_p_arr, float "
"min_p_val, bool deterministic, Generator? gen) -> ()");
m.impl("min_p_sampling_from_probs", torch::kMUSA, &min_p_sampling_from_probs);
m.def(
"top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_top_p_arr, "
"float top_p_val, bool deterministic, Generator? gen) -> ()");
m.impl("top_p_sampling_from_probs", torch::kMUSA, &top_p_sampling_from_probs);
/*
* From csrc/musa
*/
m.def(
"musa_batched_rotary_embedding_contiguous(Tensor! positions, Tensor! query, Tensor! key, "
"int head_size, Tensor! cos_sin_cache, bool is_neox, int rot_dim, Tensor! cos_sin_cache_offsets) -> ()");
m.impl("musa_batched_rotary_embedding_contiguous", torch::kMUSA, &batched_rotary_embedding_contiguous);
m.def(
"musa_rotary_embedding_contiguous(Tensor! positions, Tensor! query, Tensor! key, "
"int head_size, Tensor! cos_sin_cache, bool is_neox) -> ()");
m.impl("musa_rotary_embedding_contiguous", torch::kMUSA, &rotary_embedding_contiguous);
m.def(
"musa_fused_moe_gemv(Tensor! A, Tensor! B, Tensor! C, Tensor? A_scale, Tensor? B_scale,"
"Tensor! topk_weights, Tensor! topk_ids, bool mul_routed_weight, int topk, bool use_int4_w4a16,"
"bool use_swigelu) -> ()");
m.impl("fused_moe_gemv", torch::kMUSA, &fused_moe_gemv);
m.def(
"musa_fused_gemv(Tensor! A, Tensor! B, Tensor! C, Tensor? A_scale, Tensor? B_scale,"
"bool use_int4_w4a16, bool use_swigelu, bool use_rms_norm, Tensor? gamma,"
"float eps) -> ()");
m.impl("musa_fused_gemv", torch::kMUSA, &musa_fused_gemv);
m.def(
"musa_fused_mul_add(Tensor! output, Tensor! self, Tensor! bias,"
"float scale) -> ()");
m.impl("musa_fused_mul_add", torch::kMUSA, &fused_mul_add);
m.def(
"musa_top_k_top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? "
"maybe_top_k_arr, "
"float top_k_val, Tensor? maybe_top_p_arr, float top_p_val, bool deterministic, Generator? gen) -> ()");
m.impl("musa_top_k_top_p_sampling_from_probs", torch::kMUSA, &musa_top_k_top_p_sampling_from_probs);
/*
* From csrc/memory
*/
m.def("weak_ref_tensor(Tensor tensor) -> Tensor");
m.impl("weak_ref_tensor", torch::kMUSA, &weak_ref_tensor);
}
REGISTER_EXTENSION(common_ops)
python/sglang/kernels/aot/csrc/common_extension_rocm.cc
+249
View File
@@ -0,0 +1,249 @@
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/core/dispatch/Dispatcher.h>
#include <torch/library.h>
#include "sgl_kernel_ops.h"
TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
/*
* From csrc/elementwise
*/
m.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
m.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("gelu_tanh_and_mul", torch::kCUDA, &gelu_tanh_and_mul);
m.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);
m.def("gelu_quick(Tensor! out, Tensor input) -> ()");
m.impl("gelu_quick", torch::kCUDA, &gelu_quick);
m.def("fast_topk(Tensor score, Tensor indices, Tensor lengths, Tensor? row_starts) -> ()");
m.impl("fast_topk", torch::kCUDA, &fast_topk_interface);
m.def(
"fast_topk_transform_fused(Tensor score, Tensor lengths, Tensor dst_page_table, Tensor src_page_table, Tensor "
"cu_seqlens_q, Tensor? row_starts) -> ()");
m.impl("fast_topk_transform_fused", torch::kCUDA, &fast_topk_transform_interface);
m.def(
"fast_topk_transform_ragged_fused(Tensor score, Tensor lengths, Tensor topk_indices_ragged, Tensor "
"topk_indices_offset, Tensor ? row_starts) -> ()");
m.impl("fast_topk_transform_ragged_fused", torch::kCUDA, &fast_topk_transform_ragged_interface);
m.def(
"deepseek_v4_topk_transform_512(Tensor scores, Tensor seq_lens, Tensor page_table, Tensor! "
"page_indices, int page_size, Tensor!? raw_indices) -> ()");
m.impl("deepseek_v4_topk_transform_512", torch::kCUDA, &deepseek_v4_topk_transform_512);
m.def(
"dsv4_fused_q_norm_rope(Tensor q_input, Tensor! q_output, Tensor freqs_cis, Tensor positions, float eps) -> ()");
m.impl("dsv4_fused_q_norm_rope", torch::kCUDA, &dsv4_fused_q_norm_rope);
m.def(
"dsv4_fused_k_norm_rope_flashmla(Tensor kv, Tensor kv_weight, Tensor freqs_cis, Tensor positions, "
"Tensor out_loc, Tensor! kvcache, float eps, int page_size) -> ()");
m.impl("dsv4_fused_k_norm_rope_flashmla", torch::kCUDA, &dsv4_fused_k_norm_rope_flashmla);
m.def(
"dsv4_fused_q_indexer_rope_hadamard_quant(Tensor q_input, Tensor! q_fp8, Tensor weight, "
"Tensor! weights_out, float weight_scale, Tensor freqs_cis, Tensor positions) -> ()");
m.impl("dsv4_fused_q_indexer_rope_hadamard_quant", torch::kCUDA, &dsv4_fused_q_indexer_rope_hadamard_quant);
/*
* From csrc/allreduce
*/
m.def(
"init_custom_ar(Tensor meta, Tensor rank_data, "
"str[] handles, int[] offsets, int rank, "
"bool full_nvlink) -> int");
m.impl("init_custom_ar", torch::kCUDA, &init_custom_ar);
m.def("all_reduce_reg(int fa, Tensor inp, Tensor! out) -> ()");
m.impl("all_reduce_reg", torch::kCUDA, &all_reduce_reg);
m.def(
"all_reduce_unreg(int fa, Tensor inp, Tensor reg_buffer, Tensor! out) -> "
"()");
m.impl("all_reduce_unreg", torch::kCUDA, &all_reduce_unreg);
// Deterministic all-reduce for ROCm
extern void deterministic_all_reduce_reg(int64_t _fa, torch::Tensor& inp, torch::Tensor& out);
extern void deterministic_all_reduce_unreg(
int64_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer, torch::Tensor& out);
m.def("deterministic_all_reduce_reg(int fa, Tensor inp, Tensor! out) -> ()");
m.impl("deterministic_all_reduce_reg", torch::kCUDA, &deterministic_all_reduce_reg);
m.def("deterministic_all_reduce_unreg(int fa, Tensor inp, Tensor reg_buffer, Tensor! out) -> ()");
m.impl("deterministic_all_reduce_unreg", torch::kCUDA, &deterministic_all_reduce_unreg);
m.def("dispose", &dispose);
m.def("meta_size", &meta_size);
m.def(
"register_buffer(int fa, Tensor t, str[] handles, "
"int[] offsets) -> ()");
m.impl("register_buffer", torch::kCUDA, &register_buffer);
m.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta);
m.def("register_graph_buffers", &register_graph_buffers);
m.def("allocate_meta_buffer", &allocate_meta_buffer);
m.impl("allocate_meta_buffer", torch::kCUDA, &allocate_meta_buffer);
m.def("get_meta_buffer_ipc_handle", &get_meta_buffer_ipc_handle);
m.impl("get_meta_buffer_ipc_handle", torch::kCPU, &get_meta_buffer_ipc_handle);
// quick allreduce
m.def(
"qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
"cast_bf2half) -> ()");
m.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
m.def("init_custom_qr", &init_custom_qr);
m.def("qr_destroy", &qr_destroy);
m.def("qr_get_handle", &qr_get_handle);
m.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
m.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
// Max input size in bytes
m.def("qr_max_size", &qr_max_size);
/*
* From csrc/moe
*/
m.def(
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! cumsum_buffer, bool "
"pad_sorted_token_ids, bool ignore_invalid_expert) -> ()");
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, float "
"moe_softcapping, Tensor? correction_bias) -> ()");
m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
m.def(
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, Tensor? "
"correction_bias) -> ()");
m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
/*
* From csrc/speculative
*/
m.def(
"verify_tree_greedy(Tensor! predicts, Tensor! accept_index, Tensor! accept_token_num, "
"Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, "
"Tensor target_predict) -> ()");
m.impl("verify_tree_greedy", torch::kCUDA, &verify_tree_greedy);
m.def(
"build_tree_kernel_efficient(Tensor parent_list, Tensor selected_index, Tensor verified_seq_len, "
"Tensor! tree_mask, Tensor! positions, Tensor! retrive_index, Tensor! retrive_next_token, "
"Tensor! retrive_next_sibling, int topk, int depth, int draft_token_num, int tree_mask_mode) -> "
"()");
m.impl("build_tree_kernel_efficient", torch::kCUDA, &build_tree_kernel_efficient);
/*
* From csrc/kvcacheio
*/
m.def(
"transfer_kv_per_layer(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int item_size, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer", torch::kCUDA, &transfer_kv_per_layer);
m.def(
"transfer_kv_per_layer_pf_lf(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int layer_id, int item_size, int src_layout_dim, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_pf_lf", torch::kCUDA, &transfer_kv_per_layer_pf_lf);
m.def(
"transfer_kv_all_layer(Tensor src_k_layers, Tensor dst_k_layers, Tensor src_v_layers, Tensor dst_v_layers, "
"Tensor src_indices, Tensor dst_indices, int item_size, int num_layers, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer", torch::kCUDA, &transfer_kv_all_layer);
m.def(
"transfer_kv_all_layer_lf_pf(Tensor src_k_layers, Tensor dst_k, Tensor src_v_layers, Tensor dst_v, "
"Tensor src_indices, Tensor dst_indices, int item_size, int dst_layout_dim, int num_layers, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_lf_pf", torch::kCUDA, &transfer_kv_all_layer_lf_pf);
m.def(
"transfer_kv_per_layer_mla(Tensor src, Tensor dst, Tensor src_indices, Tensor dst_indices, int item_size, int "
"block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_mla", torch::kCUDA, &transfer_kv_per_layer_mla);
m.def(
"transfer_kv_per_layer_mla_pf_lf(Tensor src, Tensor dst, Tensor src_indices, Tensor dst_indices, int layer_id, "
"int item_size, int src_layout_dim, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_mla_pf_lf", torch::kCUDA, &transfer_kv_per_layer_mla_pf_lf);
m.def(
"transfer_kv_all_layer_mla(Tensor src_layers, Tensor dst_layers, Tensor src_indices, Tensor dst_indices, int "
"item_size, int num_layers, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_mla", torch::kCUDA, &transfer_kv_all_layer_mla);
m.def(
"transfer_kv_all_layer_mla_lf_pf(Tensor src_layers, Tensor dst, Tensor src_indices, Tensor dst_indices, "
"int item_size, int dst_layout_dim, int num_layers, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_mla_lf_pf", torch::kCUDA, &transfer_kv_all_layer_mla_lf_pf);
m.def(
"transfer_kv_direct(Tensor[] src_layers, Tensor[] dst_layers, Tensor src_indices, Tensor dst_indices, int "
"page_size) -> ()");
m.impl("transfer_kv_direct", torch::kCUDA, &transfer_kv_direct);
m.def(
"transfer_kv_per_layer_direct_pf_lf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int layer_id, int page_size)->() ");
m.impl("transfer_kv_per_layer_direct_pf_lf", torch::kCUDA, &transfer_kv_per_layer_direct_pf_lf);
m.def(
"transfer_kv_all_layer_direct_lf_pf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int page_size) ->() ");
m.impl("transfer_kv_all_layer_direct_lf_pf", torch::kCUDA, &transfer_kv_all_layer_direct_lf_pf);
m.def(
"transfer_kv_all_layer_lf_ph(Tensor src_k_layers, Tensor dst_k, Tensor src_v_layers, Tensor dst_v, "
"Tensor src_indices, Tensor dst_indices, int item_size, int dst_layout_dim, int num_layers, int page_size, int "
"head_num, int block_quota, int num_warps_per_block) -> ()");
m.impl("transfer_kv_all_layer_lf_ph", torch::kCUDA, &transfer_kv_all_layer_lf_ph);
m.def(
"transfer_kv_per_layer_ph_lf(Tensor src_k, Tensor dst_k, Tensor src_v, Tensor dst_v, Tensor src_indices, Tensor "
"dst_indices, int layer_id, int item_size, int src_layout_dim, int page_size, int head_num, int block_quota, int "
"num_warps_per_block) -> ()");
m.impl("transfer_kv_per_layer_ph_lf", torch::kCUDA, &transfer_kv_per_layer_ph_lf);
/*
* From csrc/grammar
*/
m.def("apply_token_bitmask_inplace_cuda(Tensor logits, Tensor bitmask, Tensor? indices=None) -> ()");
m.impl("apply_token_bitmask_inplace_cuda", &ApplyTokenBitmaskInplace);
/*
* From csrc/elementwise
*/
m.def(
"rotary_embedding(Tensor positions, Tensor! query,"
" Tensor!? key, int head_size,"
" Tensor cos_sin_cache, bool is_neox) -> ()");
m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
/*
* From csrc/memory
*/
m.def("weak_ref_tensor(Tensor tensor) -> Tensor");
m.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
}
REGISTER_EXTENSION(common_ops)
python/sglang/kernels/aot/csrc/cpu/CMakeLists.txt
+146
View File
@@ -0,0 +1,146 @@
cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
project(sgl_kernel)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)
find_package(Python COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT} REQUIRED)
execute_process(
COMMAND ${Python_EXECUTABLE}
-c "import torch; print(torch.utils.cmake_prefix_path)"
OUTPUT_VARIABLE TORCH_PY_PREFIX
OUTPUT_STRIP_TRAILING_WHITESPACE
)
message(STATUS ${TORCH_PY_PREFIX})
list(APPEND CMAKE_PREFIX_PATH ${TORCH_PY_PREFIX}/Torch)
find_package(Torch REQUIRED)
include_directories(
${TORCH_INCLUDE_DIRS}
${TORCH_INSTALL_PREFIX}/include
${Python_INCLUDE_DIRS}
${CMAKE_CURRENT_SOURCE_DIR}/../../csrc
${CMAKE_CURRENT_SOURCE_DIR}/../../include
${CMAKE_CURRENT_SOURCE_DIR}
)
# Build all cpp files in current and sub dirs
file(GLOB_RECURSE SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
# Exclude all arch dependent files, then add back only files for this arch
set(ARCH_DIRS "x86_64|aarch64|ppc64")
list(FILTER SOURCES EXCLUDE REGEX "^${CMAKE_CURRENT_SOURCE_DIR}/(${ARCH_DIRS})/")
# Platform-specific source and library directory
set(MY_ARCH_DIR "")
if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64")
set(PLAT_LIB_DIR "/usr/lib/x86_64-linux-gnu")
set(MY_ARCH_DIR "x86_64")
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64")
set(PLAT_LIB_DIR "/usr/lib/aarch64-linux-gnu")
set(MY_ARCH_DIR "aarch64")
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "ppc64le|ppc64")
set(PLAT_LIB_DIR "/usr/lib/powerpc64le-linux-gnu")
set(MY_ARCH_DIR "ppc64")
else()
set(PLAT_LIB_DIR "/usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu")
endif()
link_directories(${PLAT_LIB_DIR})
if(MY_ARCH_DIR)
file(GLOB_RECURSE ARCH_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/${MY_ARCH_DIR}/*.cpp")
list(APPEND SOURCES ${ARCH_SOURCES})
endif()
# Conda library path support
if(DEFINED ENV{CONDA_PREFIX})
set(CONDA_LIB_DIR "$ENV{CONDA_PREFIX}/lib")
message(STATUS "Using Conda lib dir: ${CONDA_LIB_DIR}")
link_directories(${CONDA_LIB_DIR})
set(CONDA_INCLUDE_DIR "$ENV{CONDA_PREFIX}/include")
include_directories(${CONDA_INCLUDE_DIR})
# Look for libnuma in Conda's lib directory
find_library(NUMA_LIB numa HINTS "${CONDA_LIB_DIR}")
if(NUMA_LIB)
message(STATUS "Found libnuma: ${NUMA_LIB}")
else()
message(FATAL_ERROR "libnuma not found in Conda environment at ${CONDA_LIB_DIR}\n"
"Please install it using: conda install libnuma numactl\n")
endif()
else()
if(DEFINED ENV{VIRTUAL_ENV})
set(VENV_LIB_DIR "$ENV{VIRTUAL_ENV}/lib")
message(STATUS "Using venv lib dir: ${VENV_LIB_DIR}")
link_directories(${VENV_LIB_DIR})
set(VENV_INCLUDE_DIR "$ENV{VIRTUAL_ENV}/include")
include_directories(${VENV_INCLUDE_DIR})
endif()
# Look for libnuma in system env paths
find_library(NUMA_LIB numa)
if(NUMA_LIB)
message(STATUS "Found libnuma: ${NUMA_LIB}")
else()
message(FATAL_ERROR "libnuma not found in system environment\n"
"Please install it using: apt-get install libnuma numactl\n")
endif()
endif()
# These kernels still rely on x86-specific AMX/AVX512 implementations.
# Keep them out of Arm64 bootstrap builds until native Arm paths land.
set(SGLANG_CPU_X86_ONLY_SOURCES
${CMAKE_CURRENT_SOURCE_DIR}/gemm_int4.cpp
${CMAKE_CURRENT_SOURCE_DIR}/gemm_int8.cpp
${CMAKE_CURRENT_SOURCE_DIR}/moe.cpp
${CMAKE_CURRENT_SOURCE_DIR}/moe_fp8.cpp
${CMAKE_CURRENT_SOURCE_DIR}/moe_int4.cpp
${CMAKE_CURRENT_SOURCE_DIR}/moe_int8.cpp
${CMAKE_CURRENT_SOURCE_DIR}/qkv_proj.cpp
${CMAKE_CURRENT_SOURCE_DIR}/mamba/conv.cpp
)
if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64")
add_compile_definitions(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
list(REMOVE_ITEM SOURCES ${SGLANG_CPU_X86_ONLY_SOURCES})
endif()
if(NOT DEFINED ENV{SGLANG_CPU_FP8_CVT_FTZ})
set(ENV{SGLANG_CPU_FP8_CVT_FTZ} "1")
endif()
if("$ENV{SGLANG_CPU_FP8_CVT_FTZ}" STREQUAL "1")
message(STATUS "Enabling macro: SGLANG_CPU_FP8_CVT_FTZ")
add_compile_definitions(SGLANG_CPU_FP8_CVT_FTZ)
endif()
if(MY_ARCH_DIR STREQUAL "x86_64")
add_compile_options(
-O3
-Wno-unknown-pragmas
-march=x86-64-v4
-mavx512bf16
-mavx512vnni
-mamx-tile
-mamx-bf16
-mamx-int8
-fopenmp
)
else()
add_compile_options(
-O3
-Wno-unknown-pragmas
-march=native
-fopenmp
)
endif()
Python_add_library(common_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES})
target_link_libraries(common_ops PRIVATE ${TORCH_LIBRARIES} ${NUMA_LIB})
target_include_directories(common_ops PRIVATE ${TORCH_INCLUDE_DIRS})
install(TARGETS common_ops
LIBRARY DESTINATION sgl_kernel
)
python/sglang/kernels/aot/csrc/cpu/aarch64/gemm_int8.cpp
+149
View File
@@ -0,0 +1,149 @@
#include "../common.h"
#include "op.h"
namespace {
// out = mat1 @ mat2 + bias
template <typename scalar_t>
void int8_scaled_mm_impl(
scalar_t* __restrict__ out, // [M, N], row major
const int8_t* __restrict__ mat1, // [M, K], row major
const int8_t* __restrict__ mat2, // [K, N], column major
const float* __restrict__ scales1, // [M, 1], mat1 scales
const float* __restrict__ scales2, // [1, N], mat2 scales
const float* __restrict__ bias, // [1, N]
int64_t M,
int64_t N,
int64_t K) {
TORCH_CHECK(false, "not supported yet");
}
template <>
void int8_scaled_mm_impl<at::BFloat16>(
at::BFloat16* __restrict__ out,
const int8_t* __restrict__ mat1,
const int8_t* __restrict__ mat2,
const float* __restrict__ scales1,
const float* __restrict__ scales2,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K) {
const int slice_size = (M * K * sizeof(int8_t)) > kL2Size ? 64 : 8;
const int num_slices = (N + slice_size - 1) / slice_size;
auto mm = [mat1, mat2, out, M, N, K, scales1, scales2, bias, slice_size](int64_t begin, int64_t end) {
for (int64_t slice_idx = begin; slice_idx < end; ++slice_idx) {
const int64_t n_start = slice_idx * slice_size;
const int64_t n_end = std::min(n_start + slice_size, N);
const int slice_width = static_cast<int>(n_end - n_start);
const int8_t* a_ptr = mat1;
const int8_t* b_ptr = mat2 + n_start * K;
bfloat16_t* c_ptr = reinterpret_cast<bfloat16_t*>(out) + n_start;
op::i8mm_matmul(a_ptr, b_ptr, c_ptr, M, K, N, slice_width, scales1, scales2 + n_start);
// NOTE: matmul reduces matrix values to BF16, may influence precision
if (bias) {
op::add_bias(c_ptr, bias + n_start, M, N, slice_width);
}
}
};
at::parallel_for(0, num_slices, 0, mm);
}
} // anonymous namespace
std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& /*A*/) {
TORCH_CHECK(false, "not implemented yet");
return {at::Tensor(), at::Tensor()};
}
at::Tensor int8_scaled_mm_cpu(
at::Tensor& /*mat1*/,
at::Tensor& /*mat2*/,
at::Tensor& /*scales1*/,
at::Tensor& /*scales2*/,
const std::optional<at::Tensor>& /*bias*/,
at::ScalarType /*out_dtype*/,
bool /*is_vnni*/) {
TORCH_CHECK(false, "not implemented yet");
return at::Tensor();
}
// weight : static, per-channel, symmetric
// activation : dynamic, per-token, symmetric
//
// mat1 : [M, K]
// mat2 : [N, K]
// scales1 : [M]
// scales2 : [N]
// bias : [N]
// out : [M, N]
//
// fused activation quantization and matmul
at::Tensor int8_scaled_mm_with_quant(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales2,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool /*is_vnni*/) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
CHECK_INPUT(mat2);
CHECK_INPUT(scales2);
CHECK_DIM(2, mat1);
CHECK_DIM(2, mat2);
int64_t M = mat1.size(0);
int64_t N = mat2.size(0);
int64_t K = mat1.size(1);
int64_t lda = mat1.stride(0);
CHECK_EQ(mat2.size(1), K);
CHECK_EQ(scales2.numel(), N);
const auto st = mat1.scalar_type();
TORCH_CHECK(st == at::kBFloat16, "int8_scaled_mm_with_quant: expect A to be bfloat16.");
TORCH_CHECK(st == out_dtype, "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm_with_quant: expect mat2 to be int8.");
TORCH_CHECK(scales2.scalar_type() == at::kFloat, "int8_scaled_mm_with_quant: expect scales to be float32.");
const int64_t buffer_size = M * K + M * sizeof(float);
auto buffer = at::empty({buffer_size}, mat1.options().dtype(at::kChar));
auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
const bool has_bias = bias.has_value();
const float* bias_data = nullptr;
if (has_bias) {
CHECK_EQ(bias.value().size(0), N);
bias_data = bias.value().data_ptr<float>();
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_with_quant_kernel_impl", [&] {
int8_t* __restrict__ Aq_data = buffer.data_ptr<int8_t>();
float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K));
const scalar_t* __restrict__ A_data = mat1.data_ptr<scalar_t>();
const int64_t grain = kL1Size / (K * sizeof(scalar_t));
at::parallel_for(0, M, grain, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
op::quantize_row_int8(Aq_data + m * K, As_data + m, A_data + m * lda, K);
}
});
int8_scaled_mm_impl<scalar_t>(
out.data_ptr<scalar_t>(),
Aq_data,
mat2.data_ptr<int8_t>(),
As_data,
scales2.data_ptr<float>(),
bias_data,
M,
N,
K);
});
return out;
}
python/sglang/kernels/aot/csrc/cpu/aarch64/moe.cpp
+327
View File
@@ -0,0 +1,327 @@
#include "../common.h"
#include "op.h"
namespace {
// key: expert id, value: input rows and weights for this expert
using expert_to_rows_t = std::map<int, std::vector<std::tuple<int, float>>>;
// for expert_id, row_weight_list in x_per_expert.items():
// rows, weights = zip(*row_weight_list)
// x_rows = x[rows]
// w1, w3 = torch.chunk(w13[expert_id], chunks=2)
// gate = x_rows @ w1
// up = x_rows @ w3
// up *= silu(gate)
// down = up @ w2[expert_id]
// down *= weights
// y.index_add_(0, rows, down)
template <typename scalar_t>
void fused_experts_int8_kernel_impl(
scalar_t* __restrict__ y, // [M, K], row major
const int8_t* __restrict__ x, // [M, K], row major
const int8_t* __restrict__ w13, // [E, K, 2N], per expert [K, N], column major, w1 before w3
const int8_t* __restrict__ w2, // [E, N, K], per expert [N, K], column major
const float* __restrict__ x_scale, // [M, 1]
const float* __restrict__ w13_scale, // [E, 1, 2N], per expert [1, N], w1 before w3
const float* __restrict__ w2_scale, // [E, 1, K], per expert [1, K]
const expert_to_rows_t& x_per_expert, // expert id -> related x rows and weights
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk) {
TORCH_CHECK(false, "not implemented yet");
}
template <>
void fused_experts_int8_kernel_impl<at::BFloat16>(
at::BFloat16* __restrict__ y,
const int8_t* __restrict__ x,
const int8_t* __restrict__ w13,
const int8_t* __restrict__ w2,
const float* __restrict__ x_scale,
const float* __restrict__ w13_scale,
const float* __restrict__ w2_scale,
const expert_to_rows_t& x_per_expert,
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk) {
// x dispatch buffer to aggregate all rows per expert
int64_t max_agg_rows = 0;
for (const auto& [eid, rows] : x_per_expert) {
max_agg_rows = std::max<int64_t>(max_agg_rows, rows.size());
}
// x_scale_agg[max_agg_rows] + up_scale[max_agg_rows] +
// gate[max_agg_rows,N] + up[max_agg_rows,N] + down[max_agg_rows,K]
auto f32_buffer = at::empty({max_agg_rows, 1 + 1 + N + N + K}, at::kFloat);
float* x_scale_agg = f32_buffer.data_ptr<float>();
float* up_scale = x_scale_agg + max_agg_rows;
float* gate = up_scale + max_agg_rows;
float* up = gate + max_agg_rows * N;
float* down = up + max_agg_rows * N;
// x_agg[max_agg_rows,K] + up_q8[max_agg_rows,N]
auto int8_buffer = at::empty({max_agg_rows, K + N}, at::kChar);
int8_t* x_agg = int8_buffer.data_ptr<int8_t>();
int8_t* up_q8 = x_agg + max_agg_rows * K;
// out[M,K]: accumulated output
auto out_buffer = at::zeros({M, K}, at::kFloat);
float* out = out_buffer.data_ptr<float>();
// iterate used experts
for (const auto& [eid, rows] : x_per_expert) {
const int64_t n_agg = rows.size();
// copy input rows using this expert to contiguous buffer
{
int8_t* x_agg_ptr = x_agg;
float* x_scale_agg_ptr = x_scale_agg;
for (const auto [row, weight] : rows) {
// int row; float weight;
std::memcpy(x_agg_ptr, x + row * K, K * sizeof(int8_t));
*x_scale_agg_ptr = x_scale[row];
x_agg_ptr += K;
++x_scale_agg_ptr;
}
}
// gate = x_agg @ w1
// up = x_agg @ w3
// up *= silu(gate)
{
// expert specific tensors
const int8_t* w1e = w13 + eid * 2 * N * K;
const int8_t* w3e = w1e + N * K;
const float* w1e_scale = w13_scale + eid * 2 * N;
const float* w3e_scale = w1e_scale + N;
// tensor shapes
// - x_agg: [n_agg, K], int8, row major
// - x_scale_agg: [n_agg, 1], float
// - w{1,3}e: [K, N], int8, col major
// - w{1,3}e_scale: [1, N], float
// - gate: [n_agg, N], float, row major
// - up: [n_agg, N], float, row major
// - up_q8: [n_agg, N], int8, row major
// - up_scale: [n_agg, 1], float
const int slice_size = (n_agg * K * sizeof(int8_t)) > kL2Size ? 64 : 8;
const int num_slices = (N + slice_size - 1) / slice_size;
auto mm = [&](int64_t begin, int64_t end) {
for (int64_t slice_idx = begin; slice_idx < end; ++slice_idx) {
const int64_t n_start = slice_idx * slice_size;
const int64_t n_end = std::min(n_start + slice_size, N);
const int slice_width = static_cast<int>(n_end - n_start);
const int8_t* w1e_ptr = w1e + n_start * K;
const int8_t* w3e_ptr = w3e + n_start * K;
const float* w1e_scale_ptr = w1e_scale + n_start;
const float* w3e_scale_ptr = w3e_scale + n_start;
float* gate_ptr = gate + n_start;
float* up_ptr = up + n_start;
op::i8mm_matmul(x_agg, w1e_ptr, gate_ptr, n_agg, K, N, slice_width, x_scale_agg, w1e_scale_ptr);
op::i8mm_matmul(x_agg, w3e_ptr, up_ptr, n_agg, K, N, slice_width, x_scale_agg, w3e_scale_ptr);
for (int i = 0; i < n_agg; ++i) {
const float* __restrict__ gate_ptr = gate + n_start + i * N;
float* __restrict__ up_ptr = up + n_start + i * N;
// TODO: vectorize
for (int j = 0; j < slice_width; ++j) {
up_ptr[j] *= gate_ptr[j] / (1 + std::exp(-gate_ptr[j]));
}
}
}
};
at::parallel_for(0, num_slices, 0, mm);
}
// quantize
{
const int64_t grain = kL1Size / (K * sizeof(float));
at::parallel_for(0, n_agg, grain, [&](int64_t begin, int64_t end) {
for (int64_t i = begin; i < end; ++i) {
op::quantize_row_int8(up_q8 + i * N, up_scale + i, up + i * N, N);
}
});
}
// down = up @ w2
{
// expert specific tensors
const int8_t* w2e = w2 + eid * K * N;
const float* w2e_scale = w2_scale + eid * K;
// tensor shapes
// - up_q8: [n_agg, N], int8, row major
// - up_scale: [n_agg, 1], float
// - w2e: [N, K], int8, col major
// - w2e_scale: [1, K], float
// - down: [n_agg, K], float, row major
// - out: [M, K], float, row major
const int slice_size = (n_agg * N * sizeof(int8_t)) > kL2Size ? 64 : 8;
const int num_slices = (K + slice_size - 1) / slice_size;
auto mm = [&](int64_t begin, int64_t end) {
for (int64_t slice_idx = begin; slice_idx < end; ++slice_idx) {
const int64_t n_start = slice_idx * slice_size;
const int64_t n_end = std::min(n_start + slice_size, K);
const int slice_width = static_cast<int>(n_end - n_start);
{
const int8_t* w2e_ptr = w2e + n_start * N;
const float* w2e_scale_ptr = w2e_scale + n_start;
float* down_ptr = down + n_start;
op::i8mm_matmul(up_q8, w2e_ptr, down_ptr, n_agg, N, K, slice_width, up_scale, w2e_scale_ptr);
}
// accumulate to out buffer
{
const float* __restrict__ down_ptr = down + n_start;
for (const auto [row, weight] : rows) {
// int row; float weight;
float* __restrict__ out_ptr = out + n_start + row * K;
// auto vectorizable
for (int i = 0; i < slice_width; ++i) {
out_ptr[i] += down_ptr[i] * weight;
}
down_ptr += K;
}
}
}
};
at::parallel_for(0, num_slices, 0, mm);
}
}
// copy output: float -> bf16
{
// tensor shapes
// - out: [M, K], float, row major
// - y: [M, K], bf16, row major
const int64_t grain = kL1Size / (K * sizeof(float));
at::parallel_for(0, M, grain, [&](int64_t begin, int64_t end) {
const float* out_ptr = out + begin * K;
bfloat16_t* y_ptr = reinterpret_cast<bfloat16_t*>(y) + begin * K;
op::f32_to_bf16(out_ptr, y_ptr, (end - begin) * K);
});
}
}
} // anonymous namespace
// hidden_states: [M, K]
// w13: [E, 2N, K]
// w2: [E, K, N]
// topk_weights: [M, topk]
// topk_ids: [M, topk] (int32_t)
// w13_scale: [E, 2N]
// w2_scale: [E, K]
at::Tensor fused_experts_cpu(
at::Tensor& hidden_states,
at::Tensor& w13,
at::Tensor& w2,
at::Tensor& topk_weights,
at::Tensor& topk_ids,
bool inplace,
int64_t moe_comp_method,
const std::optional<at::Tensor>& w13_scale,
const std::optional<at::Tensor>& w2_scale,
const std::optional<at::Tensor>& /*w13_zero*/,
const std::optional<at::Tensor>& /*w2_zero*/,
const std::optional<std::vector<int64_t>> block_size,
const std::optional<at::Tensor>& /*w1_bias*/,
const std::optional<at::Tensor>& /*w2_bias*/,
const std::optional<double>& /*alpha*/,
const std::optional<double>& /*limit*/,
bool /*is_vnni*/) {
const auto st = hidden_states.scalar_type();
CHECK_INPUT(hidden_states);
CHECK_INPUT(w13);
CHECK_INPUT(w2);
CHECK_EQ(topk_weights.sizes(), topk_ids.sizes());
CHECK_DIM(2, hidden_states);
CHECK_DIM(3, w13);
CHECK_DIM(3, w2);
CHECK_DIM(2, topk_weights);
CHECK_DIM(2, topk_ids);
CHECK_EQ(topk_ids.scalar_type(), at::kInt);
// TODO: support topk_weights to be bf16 or fp16 in the kernel
auto topk_weights_ = topk_weights.to(at::kFloat);
int64_t M = hidden_states.size(0);
int64_t K = hidden_states.size(1);
int64_t N = w13.size(1) / 2;
int64_t E = w13.size(0);
int64_t topk = topk_weights_.size(1);
// check weight shapes
CHECK_EQ(w2.size(0), E);
CHECK_EQ(w2.size(1), K);
CHECK_EQ(w13.size(2), K);
CHECK_EQ(w2.size(2), N);
CHECK_EQ(inplace, false);
at::Tensor out = at::empty_like(hidden_states);
// expert id -> related input rows and weights
expert_to_rows_t x_per_expert; // std::map<int, std::vector<std::tuple<int, float>>>
{
const int* ids = topk_ids.data_ptr<int>();
const float* weights = topk_weights_.data_ptr<float>();
for (int i = 0; i < M; ++i) {
for (int j = 0; j < topk; ++j) {
x_per_expert[*ids].emplace_back(i, *weights);
++ids;
++weights;
}
}
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_experts_kernel_impl", [&] {
auto& w13s = w13_scale.value();
auto& w2s = w2_scale.value();
TORCH_CHECK(w13s.numel() == E * 2 * N);
TORCH_CHECK(w2s.numel() == E * K);
// quantize hidden_states
auto x_buffer = at::empty({M * K}, hidden_states.options().dtype(at::kChar));
auto x_scale_buffer = at::empty({M}, at::kFloat);
int8_t* x = x_buffer.data_ptr<int8_t>();
float* x_scale = x_scale_buffer.data_ptr<float>();
scalar_t* in = hidden_states.data_ptr<scalar_t>();
const int64_t grain = kL1Size / (K * sizeof(scalar_t));
at::parallel_for(0, M, grain, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
op::quantize_row_int8(x + m * K, x_scale + m, in + m * K, K);
}
});
fused_experts_int8_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
x,
w13.data_ptr<int8_t>(),
w2.data_ptr<int8_t>(),
x_scale,
w13s.data_ptr<float>(),
w2s.data_ptr<float>(),
x_per_expert,
M,
N,
K,
E,
topk);
});
return out;
}
python/sglang/kernels/aot/csrc/cpu/aarch64/op.h
+343
View File
@@ -0,0 +1,343 @@
#pragma once
#include <arm_neon.h>
constexpr int64_t kL1Size = 64 * 1024;
constexpr int64_t kL2Size = 1 * 1024 * 1024;
// simd optimized operators
namespace op {
// do matmul in "R rows x C cols" tile with sdot
// - a is the full [M, K] matrix, row major
// - b is one slice of a [K, N] matrix, column major ([N, K] row major)
// - c is one slice of a [M, N] matrix, row major
//
// slice_width slice_width
// a b / c /
// | |----| |----|
// v | | v v
// / ------ \ v v / ------ \
// | ------ | / |||||| \ | ------ |
// M | ------ | @ | |||||| | = | ------ |
// | ------ | | |||||| | | ------ |
// | ------ | \ |||||| / | ------ |
// \ ------ / \ ------ /
// K N
//
template <int R = 4, int C = 8, typename T>
__attribute__((target("+dotprod+bf16"))) void sdot_matmul(
const int8_t* __restrict__ a,
const int8_t* __restrict__ b,
T* c,
int64_t M,
int64_t K,
int64_t N,
int slice_width,
const float* __restrict__ scales1,
const float* __restrict__ scales2) {
static_assert(std::is_same_v<T, float> || std::is_same_v<T, bfloat16_t>);
int row = 0;
for (; row + R <= M; row += R) {
const int8_t* a_rows[R];
T* c_rows[R];
for (int i = 0; i < R; ++i) {
a_rows[i] = a + (row + i) * K;
c_rows[i] = c + (row + i) * N;
}
int col = 0;
for (; col + C <= slice_width; col += C) {
const int8_t* b_cols[C];
for (int i = 0; i < C; ++i) {
b_cols[i] = b + (col + i) * K;
}
int32x4_t vsums[R][C]{};
// TODO: accumulated integer sum may overflow when K >= 65536
int k = 0;
for (; k + 16 <= K; k += 16) {
int8x16_t va[R];
int8x16_t vb[C];
for (int i = 0; i < R; ++i) {
va[i] = vld1q_s8(a_rows[i] + k);
}
for (int i = 0; i < C; ++i) {
vb[i] = vld1q_s8(b_cols[i] + k);
}
for (int i = 0; i < R; ++i) {
for (int j = 0; j < C; ++j) {
vsums[i][j] = vdotq_s32(vsums[i][j], va[i], vb[j]);
}
}
}
if (k < K) {
int8_t abuf[16]{};
int8_t bbuf[16]{};
for (int i = 0; i < R; ++i) {
memcpy(abuf, a_rows[i] + k, K - k);
const int8x16_t va = vld1q_s8(abuf);
for (int j = 0; j < C; ++j) {
memcpy(bbuf, b_cols[j] + k, K - k);
const int8x16_t vb = vld1q_s8(bbuf);
vsums[i][j] = vdotq_s32(vsums[i][j], va, vb);
}
}
}
for (int i = 0; i < R; ++i) {
for (int j = 0; j < C; ++j) {
const float sum = vaddvq_s32(vsums[i][j]);
const float sum_scaled = sum * scales1[row + i] * scales2[col + j];
if constexpr (std::is_same_v<T, bfloat16_t>) {
c_rows[i][col + j] = vcvth_bf16_f32(sum_scaled);
} else {
c_rows[i][col + j] = sum_scaled;
}
}
}
}
if (col < slice_width) {
sdot_matmul<R, 1>(a, b + col * K, c + col, M, K, N, slice_width - col, scales1, scales2 + col);
}
}
if (row < M) {
sdot_matmul<1, C>(a + row * K, b, c + row * N, M - row, K, N, slice_width, scales1 + row, scales2);
}
}
// do matmul in "R rows x C cols" tile with i8mm
template <int R = 4, int C = 8, typename T>
__attribute__((target("+i8mm+bf16"))) void i8mm_matmul(
const int8_t* __restrict__ a,
const int8_t* __restrict__ b,
T* c,
int64_t M,
int64_t K,
int64_t N,
int slice_width,
const float* __restrict__ scales1,
const float* __restrict__ scales2) {
static_assert(std::is_same_v<T, float> || std::is_same_v<T, bfloat16_t>);
static_assert(R % 2 == 0 && C % 2 == 0);
int row = 0;
for (; row + R <= M; row += R) {
const int8_t* a_rows[R];
T* c_rows[R];
for (int i = 0; i < R; ++i) {
a_rows[i] = a + (row + i) * K;
c_rows[i] = c + (row + i) * N;
}
int col = 0;
for (; col + C <= slice_width; col += C) {
const int8_t* b_cols[C];
for (int i = 0; i < C; ++i) {
b_cols[i] = b + (col + i) * K;
}
int8x16_t va[R], vb[C];
int32x4_t vsums[R / 2][C / 2]{};
// TODO: accumulated integer sum may overflow when K >= 65536
int k = 0;
for (; k + 16 <= K; k += 16) {
for (int i = 0; i < R; i += 2) {
const int64x2_t va0_s64 = vreinterpretq_s64_s8(vld1q_s8(a_rows[i + 0] + k));
const int64x2_t va1_s64 = vreinterpretq_s64_s8(vld1q_s8(a_rows[i + 1] + k));
va[i + 0] = vreinterpretq_s8_s64(vzip1q_s64(va0_s64, va1_s64));
va[i + 1] = vreinterpretq_s8_s64(vzip2q_s64(va0_s64, va1_s64));
}
for (int i = 0; i < C; i += 2) {
const int64x2_t vb0_s64 = vreinterpretq_s64_s8(vld1q_s8(b_cols[i + 0] + k));
const int64x2_t vb1_s64 = vreinterpretq_s64_s8(vld1q_s8(b_cols[i + 1] + k));
vb[i + 0] = vreinterpretq_s8_s64(vzip1q_s64(vb0_s64, vb1_s64));
vb[i + 1] = vreinterpretq_s8_s64(vzip2q_s64(vb0_s64, vb1_s64));
}
for (int i = 0; i < R / 2; ++i) {
for (int j = 0; j < C / 2; ++j) {
vsums[i][j] = vmmlaq_s32(vsums[i][j], va[i * 2 + 0], vb[j * 2 + 0]);
vsums[i][j] = vmmlaq_s32(vsums[i][j], va[i * 2 + 1], vb[j * 2 + 1]);
}
}
}
if (k < K) {
int8_t buf0[16]{}, buf1[16]{};
for (int i = 0; i < R; i += 2) {
memcpy(buf0, a_rows[i + 0] + k, (K - k) * sizeof(int8_t));
memcpy(buf1, a_rows[i + 1] + k, (K - k) * sizeof(int8_t));
const int64x2_t va0_s64 = vreinterpretq_s64_s8(vld1q_s8(buf0));
const int64x2_t va1_s64 = vreinterpretq_s64_s8(vld1q_s8(buf1));
va[i + 0] = vreinterpretq_s8_s64(vzip1q_s64(va0_s64, va1_s64));
va[i + 1] = vreinterpretq_s8_s64(vzip2q_s64(va0_s64, va1_s64));
}
for (int i = 0; i < C; i += 2) {
memcpy(buf0, b_cols[i + 0] + k, (K - k) * sizeof(int8_t));
memcpy(buf1, b_cols[i + 1] + k, (K - k) * sizeof(int8_t));
const int64x2_t vb0_s64 = vreinterpretq_s64_s8(vld1q_s8(buf0));
const int64x2_t vb1_s64 = vreinterpretq_s64_s8(vld1q_s8(buf1));
vb[i + 0] = vreinterpretq_s8_s64(vzip1q_s64(vb0_s64, vb1_s64));
vb[i + 1] = vreinterpretq_s8_s64(vzip2q_s64(vb0_s64, vb1_s64));
}
for (int i = 0; i < R / 2; ++i) {
for (int j = 0; j < C / 2; ++j) {
vsums[i][j] = vmmlaq_s32(vsums[i][j], va[i * 2 + 0], vb[j * 2 + 0]);
vsums[i][j] = vmmlaq_s32(vsums[i][j], va[i * 2 + 1], vb[j * 2 + 1]);
}
}
}
for (int i = 0; i < R; i += 2) {
for (int j = 0; j < C; j += 2) {
float32x4_t vsum_f32 = vcvtq_f32_s32(vsums[i / 2][j / 2]);
const float32x4_t scales = {
scales1[row + i + 0] * scales2[col + j + 0],
scales1[row + i + 0] * scales2[col + j + 1],
scales1[row + i + 1] * scales2[col + j + 0],
scales1[row + i + 1] * scales2[col + j + 1],
};
vsum_f32 = vmulq_f32(vsum_f32, scales);
if constexpr (std::is_same_v<T, bfloat16_t>) {
const bfloat16x4_t vsum_bf16 = vcvt_bf16_f32(vsum_f32);
c_rows[i + 0][col + j + 0] = vget_lane_bf16(vsum_bf16, 0);
c_rows[i + 0][col + j + 1] = vget_lane_bf16(vsum_bf16, 1);
c_rows[i + 1][col + j + 0] = vget_lane_bf16(vsum_bf16, 2);
c_rows[i + 1][col + j + 1] = vget_lane_bf16(vsum_bf16, 3);
} else {
c_rows[i + 0][col + j + 0] = vgetq_lane_f32(vsum_f32, 0);
c_rows[i + 0][col + j + 1] = vgetq_lane_f32(vsum_f32, 1);
c_rows[i + 1][col + j + 0] = vgetq_lane_f32(vsum_f32, 2);
c_rows[i + 1][col + j + 1] = vgetq_lane_f32(vsum_f32, 3);
}
}
}
}
if (col < slice_width) {
sdot_matmul<R, 1>(a, b + col * K, c + col, M, K, N, slice_width - col, scales1, scales2 + col);
}
}
if (row < M) {
sdot_matmul<1, C>(a + row * K, b, c + row * N, M - row, K, N, slice_width, scales1 + row, scales2);
}
}
__attribute__((target("+bf16"))) inline void
add_bias(bfloat16_t* __restrict__ out, const float* __restrict__ bias, int64_t M, int64_t N, int width) {
int col = 0;
for (; col + 4 <= width; col += 4) {
const float32x4_t vbias32 = vld1q_f32(bias + col);
bfloat16_t* out_ptr = out + col;
for (int64_t i = 0; i < M; ++i) {
bfloat16x4_t vout16 = vld1_bf16(out_ptr);
float32x4_t vout32 = vcvt_f32_bf16(vout16);
vout32 = vaddq_f32(vout32, vbias32);
vout16 = vcvt_bf16_f32(vout32);
vst1_bf16(out_ptr, vout16);
out_ptr += N;
}
}
for (; col < width; ++col) {
const float vbias32 = bias[col];
bfloat16_t* out_ptr = out + col;
for (int64_t i = 0; i < M; ++i) {
bfloat16_t vout16 = *out_ptr;
float vout32 = vcvtah_f32_bf16(vout16);
vout32 += vbias32;
vout16 = vcvth_bf16_f32(vout32);
*out_ptr = vout16;
out_ptr += N;
}
}
}
constexpr float eps = 1e-7;
template <typename scalar_t>
inline void quantize_row_int8(int8_t* __restrict__ q, float* scale, const scalar_t* __restrict__ x, int64_t n) {
float max_abs_val = eps;
for (int64_t i = 0; i < n; ++i) {
max_abs_val = std::max(std::abs(static_cast<float>(x[i])), max_abs_val);
}
*scale = max_abs_val / 127.0f;
const float inv_scale = 127.0f / max_abs_val;
for (int64_t i = 0; i < n; ++i) {
q[i] = static_cast<int8_t>(std::round(static_cast<float>(x[i]) * inv_scale));
}
}
// manually optimize for bf16
template <>
__attribute__((target("+bf16"))) inline void
quantize_row_int8<bfloat16_t>(int8_t* __restrict__ q, float* scale, const bfloat16_t* __restrict__ x, int64_t n) {
float max_abs_val = eps;
for (int64_t i = 0; i < n; ++i) {
max_abs_val = std::max(std::abs(vcvtah_f32_bf16(x[i])), max_abs_val);
}
*scale = max_abs_val / 127.0f;
const float inv_scale = 127.0f / max_abs_val;
int64_t i = 0;
for (; i + 16 <= n; i += 16) {
int32x4_t qv_s32[4];
{
const bfloat16x8x2_t xv_bf16 = vld1q_bf16_x2(x + i);
const float32x4_t xv_f32[4] = {
vcvtq_low_f32_bf16(xv_bf16.val[0]),
vcvtq_high_f32_bf16(xv_bf16.val[0]),
vcvtq_low_f32_bf16(xv_bf16.val[1]),
vcvtq_high_f32_bf16(xv_bf16.val[1]),
};
for (int j = 0; j < 4; ++j) {
float32x4_t qv_f32 = vmulq_n_f32(xv_f32[j], inv_scale);
qv_f32 = vrndaq_f32(qv_f32);
qv_s32[j] = vcvtq_s32_f32(qv_f32);
}
}
const int16x8_t qv_s16[2] = {
vuzp1q_s16(vreinterpretq_s16_s32(qv_s32[0]), vreinterpretq_s16_s32(qv_s32[1])),
vuzp1q_s16(vreinterpretq_s16_s32(qv_s32[2]), vreinterpretq_s16_s32(qv_s32[3])),
};
const int8x16_t qv_s8 = vuzp1q_s8(vreinterpretq_s8_s16(qv_s16[0]), vreinterpretq_s8_s16(qv_s16[1]));
vst1q_s8(q + i, qv_s8);
}
for (; i < n; ++i) {
q[i] = static_cast<int8_t>(std::round(vcvtah_f32_bf16(x[i]) * inv_scale));
}
}
template <>
inline void
quantize_row_int8<at::BFloat16>(int8_t* __restrict__ q, float* scale, const at::BFloat16* __restrict__ x, int64_t n) {
quantize_row_int8(q, scale, reinterpret_cast<const bfloat16_t*>(x), n);
}
__attribute__((target("+bf16"))) inline void f32_to_bf16(const float* f32, bfloat16_t* bf16, int64_t n) {
int64_t i = 0;
for (; i + 4 <= n; i += 4) {
const float32x4_t vf32 = vld1q_f32(f32 + i);
const bfloat16x4_t vbf16 = vcvt_bf16_f32(vf32);
vst1_bf16(bf16 + i, vbf16);
}
for (; i < n; ++i) {
bf16[i] = vcvth_bf16_f32(f32[i]);
}
}
} // namespace op
python/sglang/kernels/aot/csrc/cpu/aarch64/shm.h
+127
View File
@@ -0,0 +1,127 @@
#pragma once
#include <arm_neon.h>
#define VECTOR_LENGTH_IN_BYTES 16
__attribute__((target("+bf16"))) inline float32x4x2_t cvt_bf16_to_fp32(const bfloat16x8_t src) {
float32x4x2_t y;
y.val[0] = vcvtq_low_f32_bf16(src);
y.val[1] = vcvtq_high_f32_bf16(src);
return y;
}
__attribute__((target("+bf16"))) inline bfloat16x8_t cvt_fp32_to_bf16(const float32x4x2_t src) {
return vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(src.val[0]), src.val[1]);
}
__attribute__((target("+bf16"))) inline void
reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) {
const int element_size = 2;
const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size;
int main_elements = num_elements - (num_elements % vector_length);
int remain_elements = num_elements % vector_length;
// process aligned part
#pragma omp parallel for
for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size;
i += VECTOR_LENGTH_IN_BYTES) {
float32x4x2_t inout_val = cvt_bf16_to_fp32(vld1q_bf16((const bfloat16_t*)(buffers[0] + i)));
for (int j = 1; j < world_size; j++) {
const float32x4x2_t in_val = cvt_bf16_to_fp32(vld1q_bf16((const bfloat16_t*)(buffers[j] + i)));
inout_val.val[0] = vaddq_f32(inout_val.val[0], in_val.val[0]);
inout_val.val[1] = vaddq_f32(inout_val.val[1], in_val.val[1]);
}
vst1q_bf16((bfloat16_t*)(to_buffer + i), cvt_fp32_to_bf16(inout_val));
}
// process remaining part
int i = (start_elements + main_elements) * element_size;
while (remain_elements > 0) {
float val = 0.0f;
for (int j = 0; j < world_size; j++) {
val += vcvtah_f32_bf16(*(bfloat16_t*)(buffers[j] + i));
}
*(bfloat16_t*)(to_buffer + i) = vcvth_bf16_f32(val);
remain_elements--;
i += element_size;
}
}
inline void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) {
const int element_size = 2;
const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size;
int main_elements = num_elements - (num_elements % vector_length);
int remain_elements = num_elements % vector_length;
// process aligned part
#pragma omp parallel for
for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size;
i += VECTOR_LENGTH_IN_BYTES) {
float16x8_t inout_val = vld1q_f16((const float16_t*)(buffers[0] + i));
for (int j = 1; j < world_size; j++) {
const float16x8_t in_val = vld1q_f16((const float16_t*)(buffers[j] + i));
inout_val = vaddq_f16(inout_val, in_val);
}
vst1q_f16((float16_t*)(to_buffer + i), inout_val);
}
// process remaining part
int i = (start_elements + main_elements) * element_size;
while (remain_elements > 0) {
float16_t val = 0.0f;
for (int j = 0; j < world_size; j++) {
val = vaddh_f16(val, *(float16_t*)(buffers[j] + i));
}
*(float16_t*)(to_buffer + i) = val;
remain_elements--;
i += element_size;
}
}
inline void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) {
const int element_size = 4;
const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size;
int main_elements = num_elements - (num_elements % vector_length);
int remain_elements = num_elements % vector_length;
// process aligned part
#pragma omp parallel for
for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size;
i += VECTOR_LENGTH_IN_BYTES) {
float32x4_t inout_val = vld1q_f32((const float*)(buffers[0] + i));
for (int j = 1; j < world_size; j++) {
const float32x4_t in_val = vld1q_f32((const float*)(buffers[j] + i));
inout_val = vaddq_f32(inout_val, in_val);
}
vst1q_f32((float32_t*)(to_buffer + i), inout_val);
}
// process remaining part
int i = (start_elements + main_elements) * element_size;
while (remain_elements > 0) {
float val = 0.0f;
for (int j = 0; j < world_size; j++) {
val += *(float*)(buffers[j] + i);
}
*(float*)(to_buffer + i) = val;
remain_elements--;
i += element_size;
}
}
inline void parallel_memcpy(void* to, void* from, size_t n_bytes) {
auto aligned_bytes = n_bytes - (n_bytes % VECTOR_LENGTH_IN_BYTES);
// process aligned part
#pragma omp parallel for
for (size_t i = 0; i < aligned_bytes; i += VECTOR_LENGTH_IN_BYTES) {
const uint8x16_t val = vld1q_u8((uint8_t*)from + i);
vst1q_u8((uint8_t*)to + i, val);
}
// process remaining part
for (size_t i = aligned_bytes; i < n_bytes; i++) {
*((uint8_t*)to + i) = *((uint8_t*)from + i);
}
}
#undef VECTOR_LENGTH_IN_BYTES
python/sglang/kernels/aot/csrc/cpu/activation.cpp
+212
View File
@@ -0,0 +1,212 @@
#include "common.h"
#include "vec.h"
namespace {
template <typename scalar_t, typename func_t, typename vec_func_t>
void act_and_mul_kernel_impl(
scalar_t* __restrict__ output,
const scalar_t* __restrict__ input,
int64_t num_tokens,
int64_t dim,
const func_t& f,
const vec_func_t& vf) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int64_t kVecSize = bVec::size();
at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
for (int64_t i = begin; i < end; ++i) {
// local ptrs
const scalar_t* __restrict__ input_ptr = input + i * 2 * dim;
const scalar_t* __restrict__ input_other_ptr = input_ptr + dim;
scalar_t* __restrict__ output_ptr = output + i * dim;
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= dim - kVecSize; d += kVecSize) {
auto [x_fvec0, x_fvec1] = load_float_vec2(input_ptr + d);
auto [y_fvec0, y_fvec1] = load_float_vec2(input_other_ptr + d);
x_fvec0 = vf(x_fvec0);
x_fvec1 = vf(x_fvec1);
x_fvec0 = x_fvec0 * y_fvec0;
x_fvec1 = x_fvec1 * y_fvec1;
convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1).store(output_ptr + d);
}
#pragma GCC unroll 4
for (; d < dim; ++d) {
float x_val = static_cast<float>(input_ptr[d]);
float y_val = static_cast<float>(input_other_ptr[d]);
output_ptr[d] = f(x_val) * y_val;
}
}
});
}
// input : [num_tokens, dim] contiguous
// gate : [num_tokens, num_heads, head_dim] 2d or 3d, maybe strided
template <typename scalar_t>
void fused_sigmoid_mul_kernel_impl(
scalar_t* __restrict__ output,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ gate,
int64_t num_tokens,
int64_t dim,
int64_t num_heads,
int64_t head_dim,
int64_t g_strideT,
int64_t g_strideH) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int64_t kVecSize = bVec::size();
at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
for (int64_t i = begin; i < end; ++i) {
const scalar_t* __restrict__ i_ptr = input + i * dim;
const scalar_t* __restrict__ g_ptr = gate + i * g_strideT;
scalar_t* __restrict__ o_ptr = output + i * dim;
for (int64_t h = 0; h < num_heads; ++h) {
const scalar_t* __restrict__ attn_ptr = i_ptr + h * head_dim;
const scalar_t* __restrict__ gate_ptr = g_ptr + h * g_strideH;
scalar_t* __restrict__ out_ptr = o_ptr + h * head_dim;
int64_t d = 0;
#pragma GCC unroll 4
for (; d <= head_dim - kVecSize; d += kVecSize) {
auto [x_fvec0, x_fvec1] = load_float_vec2(attn_ptr + d);
auto [g_fvec0, g_fvec1] = load_float_vec2(gate_ptr + d);
x_fvec0 = x_fvec0 * fast_sigmoid(g_fvec0);
x_fvec1 = x_fvec1 * fast_sigmoid(g_fvec1);
convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1).store(out_ptr + d);
}
#pragma GCC unroll 4
for (; d < head_dim; ++d) {
float x_val = static_cast<float>(attn_ptr[d]);
float g_val = static_cast<float>(gate_ptr[d]);
out_ptr[d] = static_cast<scalar_t>(x_val / (1.f + std::exp(-g_val)));
}
}
}
});
}
} // anonymous namespace
// input : {num_tokens, 2 * d}
// output : {num_tokens, d}
at::Tensor silu_and_mul_cpu(at::Tensor& input) {
auto sizes = input.sizes().vec();
int64_t last_dim = input.ndimension() - 1;
int64_t d = sizes[last_dim] / 2;
sizes[last_dim] = d;
int64_t num_tokens = input.numel() / input.size(-1);
at::Tensor out = at::empty(sizes, input.options());
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "silu_and_mul", [&] {
using Vec = at::vec::Vectorized<float>;
act_and_mul_kernel_impl(
out.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
num_tokens,
d,
[](float x) { return x / (1.f + std::exp(-x)); },
[](Vec x) { return fast_silu(x); });
});
return out;
}
at::Tensor gelu_tanh_and_mul_cpu(const at::Tensor& input) {
auto sizes = input.sizes().vec();
int64_t last_dim = input.ndimension() - 1;
int64_t d = sizes[last_dim] / 2;
sizes[last_dim] = d;
int64_t num_tokens = input.numel() / input.size(-1);
at::Tensor out = at::empty(sizes, input.options());
const float sqrt_2_div_pi = std::sqrt(2.f / M_PI);
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_and_mul", [&] {
using Vec = at::vec::Vectorized<float>;
act_and_mul_kernel_impl(
out.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
num_tokens,
d,
[sqrt_2_div_pi](float x) {
float x3 = x * x * x;
float tanh_arg = sqrt_2_div_pi * (x + 0.044715f * x3);
return 0.5f * x * (1.f + std::tanh(tanh_arg));
},
[sqrt_2_div_pi](Vec x) {
Vec x3 = x * x * x;
Vec tanh_arg = Vec(sqrt_2_div_pi) * (x + Vec(0.044715f) * x3);
return Vec(0.5f) * x * (Vec(1.f) + tanh_arg.tanh());
});
});
return out;
}
at::Tensor gelu_and_mul_cpu(const at::Tensor& input) {
auto sizes = input.sizes().vec();
int64_t last_dim = input.ndimension() - 1;
int64_t d = sizes[last_dim] / 2;
sizes[last_dim] = d;
int64_t num_tokens = input.numel() / input.size(-1);
at::Tensor out = at::empty(sizes, input.options());
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gelu_and_mul", [&] {
using Vec = at::vec::Vectorized<float>;
const float inv_sqrt2 = 1.0f / std::sqrt(2.0f);
act_and_mul_kernel_impl(
out.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
num_tokens,
d,
[inv_sqrt2](float x) { return 0.5f * x * (1.f + std::erf(x * inv_sqrt2)); },
[inv_sqrt2](Vec x) { return Vec(0.5f) * x * (Vec(1.f) + (x * Vec(inv_sqrt2)).erf()); });
});
return out;
}
at::Tensor fused_sigmoid_mul_cpu(at::Tensor& input, const at::Tensor& gate, bool inplace) {
CHECK_DIM(2, input);
const int64_t gate_dim = gate.dim();
TORCH_CHECK(gate_dim == 2 || gate_dim == 3, "gate must be a 2D or 3D tensor");
CHECK_CONTIGUOUS(input);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(gate);
const auto st = input.scalar_type();
CHECK_EQ(gate.scalar_type(), st);
int64_t num_tokens = input.size(0);
int64_t d = input.size(1);
const bool is_gate_3d = gate_dim == 3;
int64_t num_heads = is_gate_3d ? gate.size(1) : 1;
int64_t head_dim = gate.size(-1);
CHECK_EQ(gate.size(0), num_tokens);
CHECK_EQ(d, num_heads * head_dim);
int64_t g_strideT = gate.stride(0);
int64_t g_strideH = is_gate_3d ? gate.stride(1) : 0;
at::Tensor out = inplace ? input : at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_sigmoid_mul", [&] {
fused_sigmoid_mul_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
gate.data_ptr<scalar_t>(),
num_tokens,
d,
num_heads,
head_dim,
g_strideT,
g_strideH);
});
return out;
}
python/sglang/kernels/aot/csrc/cpu/bmm.cpp
+203
View File
@@ -0,0 +1,203 @@
#include "common.h"
#include "gemm.h"
#include "vec.h"
namespace {
template <typename scalar_t, typename packed_t>
void bmm_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ mat1,
const packed_t* __restrict__ mat2,
int64_t B,
int64_t M,
int64_t N,
int64_t K,
int64_t mat1_strideB,
int64_t mat1_strideM,
int64_t out_strideB,
int64_t out_strideM,
float scale = 0.f) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
// mat2 contiguous in [B, N, K]
int64_t mat2_strideB = N * K;
int64_t mat2_strideN = K;
const bool use_brgemm = can_use_brgemm<scalar_t>(M);
// parallel on [B, MB, NB]
at::parallel_for(0, B * MB * NB, 0, [&](int64_t begin, int64_t end) {
int64_t bs{0}, mb{0}, nb{0};
data_index_init(begin, bs, B, mb, MB, nb, NB);
// for brgemm, use float32 for accumulate
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
for (int i = begin; i < end; ++i) {
UNUSED(i);
int mb_start = mb * BLOCK_M;
int mb_size = std::min(M - mb_start, BLOCK_M);
int nb_start = nb * BLOCK_N;
int nb_size = std::min(N - nb_start, BLOCK_N);
tinygemm_kernel<scalar_t>(
/* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM,
/* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */,
/* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start,
/* Ctmp*/ Ctmp,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ mat1_strideM,
/* ldb */ nb_size,
/* ldc */ out_strideM,
/* brg */ use_brgemm);
// move to the next index
data_index_step(bs, B, mb, MB, nb, NB);
}
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
template <>
void bmm_kernel_impl(
at::BFloat16* __restrict__ out,
const at::BFloat16* __restrict__ mat1,
const at::Float8_e4m3fn* __restrict__ mat2,
int64_t B,
int64_t M,
int64_t N,
int64_t K,
int64_t mat1_strideB,
int64_t mat1_strideM,
int64_t out_strideB,
int64_t out_strideM,
float scale) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
// mat2 contiguous in [B, N, K]
int64_t mat2_strideB = N * K;
int64_t mat2_strideN = K;
const bool use_brgemm = can_use_brgemm<at::BFloat16>(M);
// parallel on [B, MB, NB]
parallel_2d(B * MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// for brgemm, use float32 for accumulate
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
// for brgemm when mat2 is float8_e4m3
alignas(64) at::BFloat16 Btmp[BLOCK_N * BLOCK_K];
loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t bs = mb / MB;
int64_t mb_start = (mb % MB) * BLOCK_M;
int64_t mb_size = std::min(M - mb_start, BLOCK_M);
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(N - nb_start, BLOCK_N);
tinygemm_kernel(
/* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM,
/* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */,
/* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start,
/* Btmp*/ Btmp,
/* Ctmp*/ Ctmp,
/*scale*/ scale,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ mat1_strideM,
/* ldb */ nb_size,
/* ldc */ out_strideM,
/* brg */ use_brgemm);
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
} // anonymous namespace
// mat1 : [B, M, K]
// mat2 : [B, N, K] or [B, OC, IC]
// out : [B, M, N]
// scale: [] 0-dim tensor for per tensor quant
//
void bmm_cpu(
at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale) {
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
// input and out could be non-contiguous
// weight needs to be contiguous in [OC, IC] order
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(out);
CHECK_INPUT(mat2);
CHECK_DIM(3, out);
CHECK_DIM(3, mat1);
CHECK_DIM(3, mat2);
int64_t B = mat1.size(0);
int64_t M = mat1.size(1);
int64_t N = mat2.size(1);
int64_t K = mat1.size(2);
const bool use_fp8_w8a16 = scale.has_value();
TORCH_CHECK(N % 32 == 0, "tinygemm requires N to be 32x.");
int64_t mat1_strideB = mat1.stride(0);
int64_t mat1_strideM = mat1.stride(1);
int64_t out_strideB = out.stride(0);
int64_t out_strideM = out.stride(1);
// check shapes
TORCH_CHECK(mat2.size(0) == B && mat2.size(2) == K, "bmm: mat2 shape mismatch!");
TORCH_CHECK(out.size(0) == B && out.size(1) == M, "bmm: out shape mismatch!");
if (!use_fp8_w8a16) {
AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] {
bmm_kernel_impl<scalar_t, scalar_t>(
out.data_ptr<scalar_t>(),
mat1.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
B,
M,
N,
K,
mat1_strideB,
mat1_strideM,
out_strideB,
out_strideM);
});
} else { // fp8 bmm
float scale_val = 0.f;
auto scale_tensor = scale.value();
TORCH_CHECK(scale_tensor.ndimension() == 0, "bmm: expect scale to be 0-dim tensor.");
scale_val = scale_tensor.item<float>();
bmm_kernel_impl<at::BFloat16, at::Float8_e4m3fn>(
out.data_ptr<at::BFloat16>(),
mat1.data_ptr<at::BFloat16>(),
packed_w.data_ptr<at::Float8_e4m3fn>(),
B,
M,
N,
K,
mat1_strideB,
mat1_strideM,
out_strideB,
out_strideM,
scale_val);
}
}
python/sglang/kernels/aot/csrc/cpu/common.h
+436
View File
@@ -0,0 +1,436 @@
#pragma once
#include <ATen/ATen.h>
#include <ATen/Dispatch.h>
#include <ATen/Parallel.h>
#if defined(_OPENMP)
#include <omp.h>
#endif
namespace {
// dispatch bool
#define AT_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...) \
[&] { \
if (BOOL_V) { \
constexpr bool BOOL_NAME = true; \
return __VA_ARGS__(); \
} else { \
constexpr bool BOOL_NAME = false; \
return __VA_ARGS__(); \
} \
}()
#define AT_DISPATCH_BOOL2(BOOL_V1, BOOL_NAME1, BOOL_V2, BOOL_NAME2, ...) \
[&] { \
if (BOOL_V1) { \
constexpr bool BOOL_NAME1 = true; \
if (BOOL_V2) { \
constexpr bool BOOL_NAME2 = true; \
return __VA_ARGS__(); \
} else { \
constexpr bool BOOL_NAME2 = false; \
return __VA_ARGS__(); \
} \
} else { \
constexpr bool BOOL_NAME1 = false; \
if (BOOL_V2) { \
constexpr bool BOOL_NAME2 = true; \
return __VA_ARGS__(); \
} else { \
constexpr bool BOOL_NAME2 = false; \
return __VA_ARGS__(); \
} \
} \
}()
// Half + BFloat16, plus one extra scalar type
#define AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, ...) \
AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES(__VA_ARGS__) \
AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
#define AT_DISPATCH_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, __VA_ARGS__))
// dispatch: bfloat16, float16, int8_t, fp8_e4m3, uint8_t(mxfp4/int4)
#define CPU_DISPATCH_PACKED_TYPES(TYPE, ...) \
[&] { \
switch (TYPE) { \
case at::ScalarType::BFloat16: { \
using packed_t = at::BFloat16; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Half: { \
using packed_t = at::Half; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Char: { \
using packed_t = int8_t; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Float8_e4m3fn: { \
using packed_t = at::Float8_e4m3fn; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Byte: { \
using packed_t = uint8_t; \
return __VA_ARGS__(); \
} \
default: \
TORCH_CHECK(false, "Unsupported floating data type.\n"); \
} \
}()
// Helper MICRO for CPU_DISPATCH_FLOATING_TYPES_EXT:
// TYPE1: the primary dtype (input, output, weight);
// TYPE2: defined as PARAM_T input
#define CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, PARAM_T, ...) \
switch (TYPE1) { \
case at::ScalarType::BFloat16: { \
using scalar_t = at::BFloat16; \
using param_t = PARAM_T; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Half: { \
using scalar_t = at::Half; \
using param_t = PARAM_T; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Float: { \
using scalar_t = float; \
using param_t = PARAM_T; \
return __VA_ARGS__(); \
} \
default: \
TORCH_CHECK(false, "Unsupported floating data type."); \
}
// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
// TYPE1: the primary dtype (input, output, weight);
// TYPE2: defined as PARAM_T input
#define CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, PARAM_T, ...) \
switch (TYPE1) { \
case at::ScalarType::BFloat16: { \
using scalar_t = at::BFloat16; \
using param_t = PARAM_T; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Half: { \
using scalar_t = at::Half; \
using param_t = PARAM_T; \
return __VA_ARGS__(); \
} \
default: \
TORCH_CHECK(false, "Unsupported floating data type."); \
}
// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
// TYPE1: the dtype both for scalar_t and param_t
#define CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, ...) \
switch (TYPE1) { \
case at::ScalarType::BFloat16: { \
using scalar_t = at::BFloat16; \
using param_t = at::BFloat16; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Half: { \
using scalar_t = at::Half; \
using param_t = at::Half; \
return __VA_ARGS__(); \
} \
default: \
TORCH_CHECK(false, "Unsupported reduced floating data type."); \
}
// dispatch with mixed dtypes (TYPE1, TYPE2):
// TYPE1: the primary dtype (input, output, weight);
// TYPE2: the secondary dtype (bias, etc.).
#define CPU_DISPATCH_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...) \
[&] { \
if (TYPE2 == at::kFloat) { \
CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, float, __VA_ARGS__) \
} else if (TYPE2 == at::ScalarType::BFloat16) { \
CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, at::BFloat16, __VA_ARGS__) \
} else if (TYPE2 == at::ScalarType::Half) { \
CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, at::Half, __VA_ARGS__) \
} else { \
TORCH_CHECK(false, "Unsupported floating data type."); \
} \
}()
// dispatch with mixed dtypes (reduced one, no float for TYPE1) (TYPE1, TYPE2):
// TYPE1: the primary dtype (input, output, weight);
// TYPE2: the secondary dtype (bias, etc.).
#define CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...) \
[&] { \
if (TYPE2 == at::kFloat) { \
CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, float, __VA_ARGS__) \
} else { \
TORCH_CHECK(TYPE1 == TYPE2); \
CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, __VA_ARGS__) \
} \
}()
#define UNUSED(x) (void)(x)
#define CHECK_CPU(x) TORCH_CHECK(x.device().type() == at::kCPU, #x " must be a CPU tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_LAST_DIM_CONTIGUOUS(x) \
TORCH_CHECK(x.strides()[x.strides().size() - 1] == 1, #x "must be contiguous at last dimension")
#define CHECK_INPUT(x) \
CHECK_CPU(x); \
CHECK_CONTIGUOUS(x)
#define CHECK_LAST_DIM_CONTIGUOUS_INPUT(x) \
CHECK_CPU(x); \
CHECK_LAST_DIM_CONTIGUOUS(x)
#define CHECK_DIM(d, x) TORCH_CHECK(x.dim() == d, #x " must be a " #d "D tensor")
#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b)
#define CHECK_GT(a, b) TORCH_CHECK((a) > (b), "CHECK_GT(" #a ", " #b ") failed. ", a, " vs ", b)
#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b), "CHECK_GE(" #a ", " #b ") failed. ", a, " vs ", b)
template <bool is_only_lastdim_contiguous>
static inline void CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const at::IntArrayRef sizes, at::ScalarType st) {
TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes());
TORCH_CHECK(tensor.scalar_type() == st, "Input tensor dtype mismatch");
if constexpr (is_only_lastdim_contiguous) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor);
} else {
CHECK_INPUT(tensor);
}
}
// [NB] Parallel Routines
//
// * at::parallel_for - applies for most of generic use cases, this will be compiled
// against openmp in default torch release.
//
// * parallel_for - same function as above, can choose payload partition scheme in
// balance211.
//
// * parallel_2d - parallel for 2 dimensions, used in GEMM, etc.
// this one will do payload balance across 2 dimensions.
//
// grain size for each thread
constexpr int GRAIN_SIZE = 1024;
template <typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0>
inline T div_up(T x, T y) {
return (x + y - 1) / y;
}
// you can only use at::get_thread_num() with at::parallel_for()
// as it is lazy initialized, otherwise it will always return 0.
inline int get_thread_num() {
#if defined(_OPENMP)
return omp_get_thread_num();
#else
return 0;
#endif
}
// balance payload across each thread
template <typename T>
inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
#if 0
// onednn partition pattern
T& n_my = n_end;
if (nth <= 1 || n == 0) {
n_start = 0;
n_my = n;
} else {
T n1 = div_up(n, nth);
T n2 = n1 - 1;
T T1 = n - n2 * nth;
n_my = ith < T1 ? n1 : n2;
n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2;
}
n_end += n_start;
#else
// pytorch aten partition pattern
T n_my = div_up(n, nth);
n_start = ith * n_my;
n_end = std::min(n_start + n_my, n);
#endif
}
template <typename func_t>
inline void parallel_for(int n, const func_t& f) {
#if defined(_OPENMP)
#pragma omp parallel
{
int nth = omp_get_num_threads();
int ith = omp_get_thread_num();
int tbegin, tend;
balance211(n, nth, ith, tbegin, tend);
f(tbegin, tend);
}
#else
f(0, n);
#endif
}
// for 1d parallel, use `actual_nth`
// for 2d parallel, use even nths, e.g. 43->42
int inline adjust_num_threads(int m) {
int actual_nth = at::get_num_threads();
if (m == 1) {
return actual_nth;
}
return std::max(1, (actual_nth >> 1) * 2);
}
template <typename func_t>
inline void parallel_2d(int m, int n, const func_t& f) {
// make sure we have even num_threads
int nth = adjust_num_threads(m);
// [NOTE] thread blocking:
//
// 1) prefer square block per thread
// 2) use even number of CPU cores
// 3) use all `num_threads` cores
//
// we have:
// TM * TN = T
// BM / TM = BN / TN
// then:
// TM = ((BM / BN) * T) ^ 0.5
//
float r = float(m) / n;
int nth_m = std::ceil(std::sqrt(r * nth));
int nth_n = 1;
for (; nth_m > 0; --nth_m) {
nth_n = nth / nth_m;
if (nth_m * nth_n == nth) {
break;
}
}
#if defined(_OPENMP)
#pragma omp parallel num_threads(nth)
{
int ith = omp_get_thread_num();
int ith_m = ith / nth_n;
int ith_n = ith % nth_n;
int thread_block_m = div_up(m, nth_m);
int thread_block_n = div_up(n, nth_n);
int begin_m = ith_m * thread_block_m;
int end_m = std::min(m, begin_m + thread_block_m);
int begin_n = ith_n * thread_block_n;
int end_n = std::min(n, begin_n + thread_block_n);
f(begin_m, end_m, begin_n, end_n);
}
#else
f(0, m, 0, n);
#endif
}
// limit max cache blocks
// when we need to do pre-unpack for weights, e.g. fp8
#define MAX_CACHE_BLOCK_SIZE 4
template <typename T>
inline int get_cache_blocks(int chunk_size) {
// L2 2MB and ratio of 50%
const int L2_size = 2048 * 1024 >> 1;
return std::max(1, int(L2_size / (chunk_size * sizeof(T))));
}
template <>
inline int get_cache_blocks<at::Float8_e4m3fn>(int chunk_size) {
// fp8 uses bf16 as accumulate type
int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
}
template <>
inline int get_cache_blocks<uint8_t>(int chunk_size) {
// mxfp4 uses bf16 as accumulate type
int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
}
// 2d sequential loop in range : [mb0, mb1), [nb0, nb1)
template <typename T, typename func_t>
inline void loop_2d(int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1, int64_t chunk_size, const func_t& f) {
// get number of blocks for L2 in most inner loop
int64_t cache_blocks_nb = get_cache_blocks<T>(chunk_size);
// loop order: [NB / cache_blocks_nb, MB, cache_blocks_nb]
// TODO: implement reverse order of [MB / cache_blocks_mb, NB, cache_blocks_mb]
for (int64_t nbb = nb0; nbb < nb1; nbb += cache_blocks_nb) {
for (int64_t mb = mb0; mb < mb1; ++mb) {
for (int64_t nb = nbb; nb < std::min(nbb + cache_blocks_nb, nb1); ++nb) {
f(mb, nb, nb - nbb);
}
}
}
}
// data indexing for dimension collapse
template <typename T>
inline T data_index_init(T offset) {
return offset;
}
template <typename T, typename... Args>
inline T data_index_init(T offset, T& x, const T& X, Args&&... args) {
offset = data_index_init(offset, std::forward<Args>(args)...);
x = offset % X;
return offset / X;
}
inline bool data_index_step() {
return true;
}
template <typename T, typename... Args>
inline bool data_index_step(T& x, const T& X, Args&&... args) {
if (data_index_step(std::forward<Args>(args)...)) {
x = ((x + 1) == X) ? 0 : (x + 1);
return x == 0;
}
return false;
}
// forced unroll for perf critical path
#if __has_attribute(always_inline)
#define ALWAYS_INLINE __attribute__((__always_inline__)) inline
#else
#define ALWAYS_INLINE inline
#endif
template <int n>
struct Unroll {
template <typename Func, typename... Args>
ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
Unroll<n - 1>{}(f, args...);
f(std::integral_constant<int, n - 1>{}, args...);
}
};
template <>
struct Unroll<1> {
template <typename Func, typename... Args>
ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
f(std::integral_constant<int, 0>{}, args...);
}
};
// conditional data ptr for optional tensor
template <typename T>
inline T* conditional_data_ptr(const std::optional<at::Tensor>& opt) {
return opt.has_value() ? opt.value().data_ptr<T>() : nullptr;
}
} // anonymous namespace
python/sglang/kernels/aot/csrc/cpu/conv3d.cpp
+232
View File
@@ -0,0 +1,232 @@
#include "common.h"
#include "gemm.h"
#include "vec.h"
namespace {
// convert to vnni format
// from [N, K] to [K/2, N, 2] for bfloat16 and float16
template <typename scalar_t>
inline void
pack_vnni(scalar_t* __restrict__ packed, const scalar_t* __restrict__ weight, int64_t N, int64_t K, int64_t lda) {
const int64_t VNNI_BLK = 2;
for (int64_t n = 0; n < N; ++n) {
for (int64_t k = 0; k < K / VNNI_BLK; ++k) {
for (int64_t d = 0; d < VNNI_BLK; ++d) {
packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * lda + k * VNNI_BLK + d];
}
}
}
}
#if defined(CPU_CAPABILITY_AVX512)
template <>
inline void pack_vnni(
at::BFloat16* __restrict__ packed, const at::BFloat16* __restrict__ weight, int64_t N, int64_t K, int64_t lda) {
const float* src = reinterpret_cast<const float*>(weight);
float* dst = reinterpret_cast<float*>(packed);
int64_t K2 = K >> 1;
int64_t lda2 = lda >> 1;
int64_t ldb2 = N * 2 >> 1;
__m512i vinputs[16];
for (int64_t n = 0; n < N; n += 16) {
for (int64_t k2 = 0; k2 < K2; k2 += 16) {
for (int64_t d = 0; d < 16; ++d) {
vinputs[d] = _mm512_loadu_si512(src + (n + d) * lda2 + k2);
}
transpose_16x16_32bit(vinputs);
for (int64_t d = 0; d < 16; ++d) {
_mm512_storeu_si512(dst + (k2 + d) * ldb2 + n, vinputs[d]);
}
}
}
}
#endif
// apply bias: C [M, N] ldc, Ctmp: [M, N]
template <typename scalar_t>
inline void copy_add_stub(
scalar_t* __restrict__ C,
const float* __restrict__ Ctmp,
const scalar_t* __restrict__ bias,
int64_t M,
int64_t N,
int64_t ldc) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
for (int64_t d = 0; d < N; d += kVecSize) {
auto [bias0, bias1] = load_float_vec2(bias + d);
for (int64_t m = 0; m < M; ++m) {
auto [data0, data1] = load_float_vec2(Ctmp + m * N + d);
data0 = data0 + bias0;
data1 = data1 + bias1;
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
out_vec.store(C + m * ldc + d);
}
}
}
template <typename scalar_t>
void conv3d_embed_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ weight,
const scalar_t* __restrict__ bias,
int64_t N,
int64_t IC,
int64_t OC,
int64_t D,
int64_t H,
int64_t W) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(N, BLOCK_M);
const int64_t NB = div_up(OC, BLOCK_N);
// K in gemm
const int64_t K = IC * D * H * W;
// input : [ N/BLOCK_M, BLOCK_M, IC, D, H, W]
// weight: [OC/BLOCK_N, IC, D, H*W/2, BLOCK_N, 2]
// out : [N/BLOCK_M, BLOCK_M, OC/BLOCK_N, BLOCK_N]
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(N - mb_start, BLOCK_M);
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(OC - nb_start, BLOCK_N);
const scalar_t* __restrict__ A = input + mb_start * K;
const scalar_t* __restrict__ B = weight + nb_start * K;
#if 0
// only access 1st index of D dimension
for (int64_t ic = 0; ic < IC; ++ic) {
for (int64_t d = 0; d < D; ++d) {
at::native::cpublas::brgemm(
mb_size,
nb_size,
H * W,
K,
BLOCK_N,
BLOCK_N,
/* add_C */ ic > 0 || d > 0,
A + ic * (D * H * W) + /* d */ 0 * (H * W), // dimension D for input is repeated
B + ic * (D * BLOCK_N * H * W) + d * (BLOCK_N * H * W),
Ctmp);
}
#else
// accumulates K normally, this is still marginally faster than above
at::native::cpublas::brgemm(mb_size, nb_size, K, K, BLOCK_N, BLOCK_N, false, A, B, Ctmp);
#endif
// update bias
copy_add_stub(out + mb_start * OC + nb_start, Ctmp, bias + nb_start, mb_size, nb_size, OC);
});
at::native::cpublas::brgemm_release();
});
}
} // anonymous namespace
// [NB]: use blocked format for weight of OIDHW
//
// from [OC, Cin, D, H, W]
// view [OC / BLOCK_N, BLOCK_N, Cin, D, H * W]
// view [OC / BLOCK_N, IC, D, BLOCK_N, H * W]
// to [OC / BLOCK_N][IC, D][H * W / 2, BLOCK_N, 2]
// +- parallel -+- seq -+------ mma ----------+
//
at::Tensor conv3d_embed_weight_pack(const at::Tensor& weight) {
CHECK_INPUT(weight);
int64_t OC = weight.size(0);
int64_t IC = weight.size(1);
int64_t D = weight.size(2);
int64_t H = weight.size(3);
int64_t W = weight.size(4);
constexpr int64_t BLOCK_N = block_size_n();
TORCH_CHECK(OC % BLOCK_N == 0, "conv3d_embed_weight_pack: expect OC dividable by ", BLOCK_N);
TORCH_CHECK((H * W) % TILE_K == 0, "conv3d_embed_weight_pack: expect IC dividable by ", TILE_K);
// strides
int64_t stride_nb = BLOCK_N * IC * D * H * W;
int64_t stride_ic = D * H * W;
int64_t stride_d = H * W;
const int64_t NB = div_up(OC, BLOCK_N);
at::Tensor packed_weight = at::empty_like(weight);
AT_DISPATCH_REDUCED_FLOATING_TYPES(weight.scalar_type(), "conv3d_embed_weight_pack", [&] {
// parallel {NB, IC, D}
at::parallel_for(0, NB * IC * D, 0, [&](int64_t begin, int64_t end) {
int64_t nb{0}, ic{0}, d{0};
data_index_init(begin, nb, NB, ic, IC, d, D);
const scalar_t* w_data = weight.data_ptr<scalar_t>();
scalar_t* packed_data = packed_weight.data_ptr<scalar_t>();
for (int64_t i = begin; i < end; ++i) {
int64_t n = nb * BLOCK_N;
int64_t n_size = std::min(BLOCK_N, OC - n); // BLOCK_N
pack_vnni<scalar_t>(
packed_data + i * (BLOCK_N * H * W),
w_data + nb * stride_nb + ic * stride_ic + d * stride_d,
n_size,
H * W,
IC * D * H * W);
// move to the next index
data_index_step(nb, NB, ic, IC, d, D);
}
});
});
return packed_weight;
}
// conv3d mapped to gemm in embedding
at::Tensor conv3d_embed_cpu(const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias, bool is_vnni) {
auto packed_w = is_vnni ? weight : conv3d_embed_weight_pack(weight);
CHECK_CONTIGUOUS(input);
CHECK_CONTIGUOUS(weight);
CHECK_DIM(5, input);
CHECK_DIM(5, weight);
const int64_t N = input.size(0);
const int64_t IC = input.size(1);
const int64_t OC = weight.size(0);
const int64_t D = input.size(2);
const int64_t H = input.size(3);
const int64_t W = input.size(4);
const auto st = input.scalar_type();
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {OC, IC, D, H, W}, st);
CHECK_INPUT_SHAPE_DTYPE<false>(bias, {OC}, st);
// allocate {D, H, W} for out is 1
at::Tensor out = at::empty({N, OC}, input.options());
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "conv3d_embed_kernel_impl", [&] {
conv3d_embed_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
bias.data_ptr<scalar_t>(),
N,
IC,
OC,
D,
H,
W);
});
return out;
}
python/sglang/kernels/aot/csrc/cpu/decode.cpp
+1814
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/aot/csrc/cpu/extend.cpp
+563
View File
@@ -0,0 +1,563 @@
#include "common.h"
#include "flash_attn.h"
#include "gemm.h"
namespace {
// [NOTE]: extend attention for CPU
// 1. BLOCK_M and BLOCK_N tuned for various seq lengths
// 2. can handle non-contiguous k_extend and v_extend
// 3. computes attention for prefix and extend separately
// 4. TODO: apply head dimension blocking to optimize GQA
// 5. optional tree mask for speculative decoding TARGET_VERIFY (EAGLE topk > 1):
// `tree_mask` is a flat [batches * qlen * qlen] bool tensor in
// TreeMaskMode::QLEN_ONLY layout, where qlen == extend_seq_lens[bs] ==
// max_len_extend (uniform across the batch, equal to draft_token_num).
// Row i = query draft token, column j = key draft token; true means query i
// may attend key j (each row marks self + ancestors + root). The committed
// prefix (stage 1) is implicitly fully visible to every draft token, which
// is why the mask only covers the qlen x qlen new-token block; the GPU
// FULL_MASK layout carries the prefix columns explicitly but they are
// all-true for EAGLE. When tree_mask is absent, stage 2 falls back to the
// plain causal mask (correct for non-spec extend and topk == 1 chains).
//
template <typename scalar_t, typename index_t, int BLOCK_M, int BLOCK_N>
void extend_attention_kernel_impl(
scalar_t* __restrict__ o_extend,
const scalar_t* __restrict__ q_extend,
const scalar_t* __restrict__ k_extend,
const scalar_t* __restrict__ v_extend,
const scalar_t* __restrict__ k_buffer,
const scalar_t* __restrict__ v_buffer,
const index_t* __restrict__ req_to_token,
const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
const index_t* __restrict__ extend_seq_lens,
const index_t* __restrict__ extend_start_loc,
const void* __restrict__ buffer,
const scalar_t* __restrict__ sinks,
const bool* __restrict__ tree_mask,
int batches,
int num_heads,
int num_heads_kv,
int head_size,
int head_size_v,
int q_strideM,
int q_strideH,
int ke_strideN,
int ke_strideH,
int ve_strideN,
int ve_strideH,
int k_strideN,
int k_strideH,
int v_strideN,
int v_strideH,
float sm_scale,
int max_num_reqs,
int max_context_len,
int max_total_num_tokens,
int max_len_extend,
int buffer_size_per_thread,
int64_t sliding_window_size,
bool is_prefix_skipped,
bool is_cross_attn,
bool has_encoder_lens,
bool has_sink) {
// strides
const int o_strideM = num_heads * head_size_v;
const int o_strideH = head_size_v;
// we use same buffer for packed key and value
const int ldb_tmp = std::max(head_size, head_size_v);
const int num_groups = num_heads / num_heads_kv;
TORCH_CHECK(num_groups * num_heads_kv == num_heads);
// number of blocks along M
int MB = div_up(max_len_extend, BLOCK_M);
// parallel on [batches, num_heads, BM]
at::parallel_for(0, batches * num_heads * MB, 0, [&](int begin, int end) {
int bs{0}, head_id{0}, mb{0};
data_index_init(begin, bs, batches, head_id, num_heads, mb, MB);
int tid = at::get_thread_num();
// s_i: [BLOCK_M, BLOCK_N]
float* __restrict__ s_i = reinterpret_cast<float*>((char*)(buffer) + tid * buffer_size_per_thread);
// v_prime: [BLOCK_M, head_size_v]
float* __restrict__ v_prime = s_i + BLOCK_M * BLOCK_N;
// s_delta: [BLOCK_M, BLOCK_N]
scalar_t* __restrict__ s_delta = reinterpret_cast<scalar_t*>(v_prime + BLOCK_M * head_size_v);
// Btmp: [BLOCK_N, max(head_size, head_size_v)]
scalar_t* __restrict__ Btmp = reinterpret_cast<scalar_t*>(s_delta + BLOCK_M * BLOCK_N);
// init Btmp just once for each thread to prevent NaN
fill_stub(Btmp, 0.f, BLOCK_N * ldb_tmp);
fill_stub(s_delta, 0.f, BLOCK_M * BLOCK_N);
alignas(64) float s_prime[BLOCK_M];
alignas(64) float m_prime[BLOCK_M];
for (int i = begin; i < end; ++i) {
// seq_len = prefix + extend
int head_kv_id = head_id / num_groups;
int seq_len = seq_lens[bs];
int seq_len_extend = extend_seq_lens[bs];
int seq_len_prefix = seq_len - seq_len_extend;
int seq_extend_start_loc = extend_start_loc[bs];
int req_pool_id = req_pool_indices[bs];
int kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
TORCH_CHECK(seq_len_prefix >= 0, "prefix len < 0!");
TORCH_CHECK(seq_len <= max_context_len, "seq_len out of scope!");
TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
if (is_prefix_skipped) {
TORCH_CHECK(seq_len_prefix == 0, "extend attention: expect seq_len_prefix to be 0, got ", seq_len_prefix);
}
if (tree_mask != nullptr) {
// QLEN_ONLY layout assumes a uniform qlen across the batch (TARGET_VERIFY)
TORCH_CHECK(
seq_len_extend == max_len_extend,
"extend attention: tree_mask requires uniform extend_seq_lens, got ",
seq_len_extend,
" vs ",
max_len_extend);
}
// offset and size in MB
int m = mb * BLOCK_M;
int m_size = std::min(BLOCK_M, seq_len_extend - m);
if (m_size <= 0) {
data_index_step(bs, batches, head_id, num_heads, mb, MB);
continue;
}
// get query
const scalar_t* __restrict__ q_ptr = q_extend + (seq_extend_start_loc + m) * q_strideM + head_id * q_strideH;
// init v', s' and m'
fill_stub(v_prime, 0.f, m_size * head_size_v);
fill_stub(s_prime, 0.f, m_size);
fill_stub(m_prime, -std::numeric_limits<scalar_t>::infinity(), m_size);
// stage 1: compute scores with prefix
int kv_start = 0;
int kv_end = is_cross_attn ? encoder_lens[bs] : seq_len_prefix;
for (int n = kv_start; n < kv_end; n += BLOCK_N) {
int n_size = std::min(BLOCK_N, kv_end - n);
// `n_size` is K in 2nd gemm, pad to TILE_K;
const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
// get key and pack
pack_vnni<scalar_t, index_t>(
/* dst */ Btmp,
/* src */ k_buffer + head_kv_id * k_strideH,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
/* N */ n_size,
/* K */ head_size,
/* ld_src */ k_strideN,
/* ld_dst */ BLOCK_N);
// calculate s_i <- Q @ K
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ head_size,
/* lda */ q_strideM,
/* ldb */ BLOCK_N,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ q_ptr,
/* B */ Btmp,
/* C */ s_i);
for (int row = 0; row < m_size; ++row) {
if (sliding_window_size > 0) {
int last_col = seq_len_prefix + row + m - sliding_window_size + 1;
if (last_col >= n + n_size) {
continue;
}
fill_stub(s_i + row * BLOCK_N, -std::numeric_limits<float>::infinity(), last_col - n);
}
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row);
}
// get value and pack
pack_vnni2<scalar_t>(
/* dst */ Btmp,
/* src */ v_buffer + head_kv_id * v_strideH,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
/* K */ n_size,
/* N */ head_size_v,
/* ld_src */ v_strideN,
/* ld_dst */ head_size_v);
// calculate V' <- s_delta @ V + V'
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ head_size_v,
/* K */ padded_n_size, // n_size
/* lda */ BLOCK_N,
/* ldb */ head_size_v,
/* ldc */ head_size_v,
/* add_C */ true,
/* A */ s_delta,
/* B */ Btmp,
/* C */ v_prime);
} // loop with seq_len_prefix
if (!is_cross_attn) {
// stage 2: compute the triangle part
int num_keys = std::min(seq_len_extend, m + BLOCK_M);
for (int n = 0; n < num_keys; n += BLOCK_N) {
int n_size = std::min(BLOCK_N, num_keys - n);
// `n_size` is K in 2nd gemm, pad to TILE_K;
const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
// get key and pack
pack_vnni<scalar_t>(
/* dst */ Btmp,
/* src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH,
/* N */ n_size,
/* K */ head_size,
/* ld_src */ ke_strideN,
/* ld_dst */ BLOCK_N);
// calculate s_i <- Q @ K
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ head_size,
/* lda */ q_strideM,
/* ldb */ BLOCK_N,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ q_ptr,
/* B */ Btmp,
/* C */ s_i);
// apply tree mask (speculative TARGET_VERIFY) or causal mask
if (tree_mask != nullptr) {
// [Note] tree mask for EAGLE topk > 1 (TreeMaskMode::QLEN_ONLY).
// mask[bs][m + row][n + col] == false -> query draft token (m + row)
// may not attend key draft token (n + col); set the score to -inf
// before softmax. The tree mask subsumes the causal constraint:
// ancestors always precede descendants in the draft token ordering,
// so permitted keys satisfy j <= i and the causal `num_keys` bound
// above remains valid.
const bool* __restrict__ mask_base =
tree_mask + (static_cast<int64_t>(bs) * seq_len_extend + m) * seq_len_extend + n;
for (int row = 0; row < m_size; ++row) {
float* __restrict__ row_ptr = s_i + row * BLOCK_N;
const bool* __restrict__ mask_ptr = mask_base + static_cast<int64_t>(row) * seq_len_extend;
for (int col = 0; col < n_size; ++col) {
if (!mask_ptr[col]) {
row_ptr[col] = -std::numeric_limits<float>::infinity();
}
}
}
} else if (n + n_size - 1 > m) {
// apply causal mask
// [Note] condition to apply causal mask.
// Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
// n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was
// correct when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys.
// With BLOCK_M=512, BLOCK_N=768:
// BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
// Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
// Query row=0 is at position 512, so keys 513..767 are future and must be
// masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
// producing wrong (non-causal) attention for rows 0..254 of this m-block.
for (int row = 0; row < m_size; ++row) {
int last_col = m + row - n;
// [Note] mask the entire row if last_col < 0.
// Clamp to -1: when n > m + row every key in this block is a future
// key, so the entire row should be masked. Without this clamp,
// last_col+1 <= 0 and fill_stub would write before row_ptr.
last_col = std::max(last_col, -1);
// fill [last_col + 1, n_size) to -inf
float* row_ptr = s_i + row * BLOCK_N;
fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
}
}
for (int row = 0; row < m_size; ++row) {
if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size - 1 &&
row + m + 1 < n + sliding_window_size + n_size) {
fill_stub(
s_i + row * BLOCK_N, -std::numeric_limits<float>::infinity(), row + m - n - sliding_window_size + 1);
} else if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size) {
continue;
}
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row);
}
// get value and pack
pack_vnni2<scalar_t>(
/* dst */ Btmp,
/* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH,
/* K */ n_size,
/* N */ head_size_v,
/* ld_src */ ve_strideN,
/* ld_dst */ head_size_v);
// calculate V' <- s_delta @ V + V'
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ head_size_v,
/* K */ padded_n_size, // n_size
/* lda */ BLOCK_N,
/* ldb */ head_size_v,
/* ldc */ head_size_v,
/* add_C */ true,
/* A */ s_delta,
/* B */ Btmp,
/* C */ v_prime);
} // loop with seq_len_extend
}
scalar_t* __restrict__ out_ptr = o_extend + (seq_extend_start_loc + m) * o_strideM + head_id * o_strideH;
for (int row = 0; row < m_size; ++row) {
if (has_sink) {
s_prime[row] += std::exp(sinks[head_id] - m_prime[row]);
}
float s = 1 / s_prime[row];
copy_stub<scalar_t>(out_ptr + row * o_strideM, v_prime + row * head_size_v, s, head_size_v);
}
// move to the next index
data_index_step(bs, batches, head_id, num_heads, mb, MB);
}
at::native::cpublas::brgemm_release();
});
}
} // anonymous namespace
template <int BLOCK_M, int BLOCK_N>
inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int head_size_v) {
static_assert(BLOCK_M <= BLOCK_N, "Make sure BLOCK_M <= BLOCK_N to prevent buffer overflows during causal masking");
const int size_per_thread =
/* s_i */ BLOCK_M * BLOCK_N * sizeof(float) +
/* v_prime */ BLOCK_M * head_size_v * sizeof(float) +
/* s_delta */ BLOCK_M * BLOCK_N * sizeof(uint16_t) +
/* Btmp */ BLOCK_N * std::max(head_size, head_size_v) * sizeof(uint16_t);
buffer.resize_({num_threads, size_per_thread});
return size_per_thread;
}
#define LAUNCH_EXTEND_ATTENTION_KERNEL(BLOCK_M, BLOCK_N) \
do { \
int sz = resize_buffer<BLOCK_M, BLOCK_N>(buffer, num_threads, head_size, head_size_v); \
\
extend_attention_kernel_impl<scalar_t, index_t, BLOCK_M, BLOCK_N>( \
o_extend.data_ptr<scalar_t>(), \
q_extend.data_ptr<scalar_t>(), \
k_extend.data_ptr<scalar_t>(), \
v_extend.data_ptr<scalar_t>(), \
k_buffer.data_ptr<scalar_t>(), \
v_buffer.data_ptr<scalar_t>(), \
req_to_token.data_ptr<index_t>(), \
req_pool_indices.data_ptr<int64_t>(), \
seq_lens.data_ptr<int64_t>(), \
encoder_lens_t.data_ptr<int64_t>(), \
extend_seq_lens.data_ptr<index_t>(), \
extend_start_loc.data_ptr<index_t>(), \
buffer.data_ptr(), \
sinks_tensor.data_ptr<scalar_t>(), \
tree_mask_ptr, \
num_seqs, \
num_heads, \
num_heads_kv, \
head_size, \
head_size_v, \
q_strideM, \
q_strideH, \
ke_strideN, \
ke_strideH, \
ve_strideN, \
ve_strideH, \
k_strideN, \
k_strideH, \
v_strideN, \
v_strideH, \
sm_scale, \
max_num_reqs, \
max_context_len, \
max_total_num_tokens, \
max_len_extend, \
sz, \
sliding_window_size, \
is_prefix_skipped, \
is_cross_attn, \
has_encoder_lens, \
has_sink); \
} while (0)
// q_extend, k_extend, v_extend, o_extend: contiguous tensors
// k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
//
// q_extend: [num_tokens, num_heads, head_size]
// k_extend: [num_extend_tokens, num_heads, head_size]
// v_extend: [num_extend_tokens, num_heads, head_size]
// o_extend: [num_tokens, num_heads, head_size]
// k_buffer: [max_total_num_tokens, num_heads, head_size]
// v_buffer: [max_total_num_tokens, num_heads, head_size]
// req_to_token: [max_num_reqs, max_context_len] int32 or int64
// req_pool_indices: [num_seqs] int64
// seq_lens: [num_seqs] int64
// extend_seq_lens: [num_seqs]
// extend_start_loc: [num_seqs]
// encoder_lens: [num_seqs] int64 or None
// sinks: [num_heads] or None
// tree_mask: [num_seqs * max_len_extend * max_len_extend] bool or None
// TreeMaskMode::QLEN_ONLY tree mask for speculative TARGET_VERIFY; see [NOTE] 5 above.
void extend_attention_cpu(
at::Tensor& q_extend,
const std::optional<at::Tensor>& k_extend_opt,
const std::optional<at::Tensor>& v_extend_opt,
at::Tensor& o_extend,
at::Tensor& k_buffer,
at::Tensor& v_buffer,
at::Tensor& req_to_token,
at::Tensor& req_pool_indices,
at::Tensor& seq_lens,
at::Tensor& extend_seq_lens,
at::Tensor& extend_start_loc,
int64_t max_len_extend,
double sm_scale,
double logit_cap,
bool is_cross_attn,
int64_t sliding_window_size,
std::optional<at::Tensor> encoder_lens,
std::optional<at::Tensor> sinks,
std::optional<at::Tensor> tree_mask) {
if (!is_cross_attn) {
TORCH_CHECK(
k_extend_opt.has_value() && v_extend_opt.has_value(),
"k_extend and v_extend are required for non-cross attention");
}
// Since k_extend and v_extend are not used for cross attention, they can be initialized as k_buffer and v_buffer
// here.
auto k_extend = k_extend_opt.has_value() ? k_extend_opt.value() : k_buffer;
auto v_extend = v_extend_opt.has_value() ? v_extend_opt.value() : v_buffer;
CHECK_LAST_DIM_CONTIGUOUS_INPUT(q_extend);
CHECK_INPUT(o_extend);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_extend);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer);
int num_seqs = seq_lens.size(0);
int max_num_reqs = req_to_token.size(0);
int max_context_len = req_to_token.size(1);
int max_total_num_tokens = k_buffer.size(0);
int num_heads = q_extend.size(1);
int num_heads_kv = k_extend.size(1);
int head_size = q_extend.size(2);
int head_size_v = v_extend.size(2);
// strides for q_extend, k_extend and v_extend
int q_strideM = q_extend.stride(0);
int q_strideH = q_extend.stride(1);
int ke_strideN = k_extend.stride(0);
int ke_strideH = k_extend.stride(1);
int ve_strideN = v_extend.stride(0);
int ve_strideH = v_extend.stride(1);
// strides for k_buffer and v_buffer
int k_strideN = k_buffer.stride(0);
int k_strideH = k_buffer.stride(1);
int v_strideN = v_buffer.stride(0);
int v_strideH = v_buffer.stride(1);
// check sizes
CHECK_EQ(req_pool_indices.size(0), num_seqs);
CHECK_EQ(extend_seq_lens.size(0), num_seqs);
CHECK_EQ(extend_start_loc.size(0), num_seqs);
CHECK_EQ(v_extend.size(1), num_heads_kv);
CHECK_EQ(k_buffer.size(1), v_buffer.size(1));
// MLA will skip prefix part
const bool is_prefix_skipped = k_buffer.size(1) != num_heads_kv;
// check index data types
const auto index_dtype = req_to_token.scalar_type();
TORCH_CHECK(
index_dtype == at::kInt || index_dtype == at::kLong,
"extend: expect req_to_token to be int32 or int64, got ",
index_dtype);
TORCH_CHECK(seq_lens.scalar_type() == at::kLong, "extend: expect req_lens to be int64, got ", seq_lens.scalar_type());
TORCH_CHECK(
req_pool_indices.scalar_type() == at::kLong,
"extend: expect req_pool_indices to be int64, got ",
req_pool_indices.scalar_type());
TORCH_CHECK(
extend_seq_lens.scalar_type() == index_dtype && extend_start_loc.scalar_type() == index_dtype,
"extend: expect extend_seq_lens and extend_start_loc to have same dtype as req_to_token.");
// D and DV need to be 32x as we transpose by 512-bit
TORCH_CHECK(head_size % 32 == 0, "invalid head_size ", head_size);
TORCH_CHECK(head_size_v % 32 == 0, "invalid head_size_v ", head_size_v);
int num_threads = at::get_num_threads();
auto buffer = at::empty({}, q_extend.options().dtype(at::kChar));
bool has_encoder_lens = encoder_lens.has_value();
// Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype.
at::Tensor encoder_lens_t = seq_lens;
if (has_encoder_lens) {
encoder_lens_t = encoder_lens.value();
CHECK_EQ(encoder_lens_t.size(0), num_seqs);
}
bool has_sink = sinks.has_value();
at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, q_extend.options());
CHECK_DIM(1, sinks_tensor);
CHECK_EQ(sinks_tensor.size(0), num_heads);
const bool* tree_mask_ptr = nullptr;
if (tree_mask.has_value()) {
const at::Tensor& tree_mask_t = tree_mask.value();
CHECK_INPUT(tree_mask_t);
TORCH_CHECK(
tree_mask_t.scalar_type() == at::kBool, "extend: expect tree_mask to be bool, got ", tree_mask_t.scalar_type());
TORCH_CHECK(
tree_mask_t.numel() == static_cast<int64_t>(num_seqs) * max_len_extend * max_len_extend,
"extend: expect tree_mask numel to be num_seqs * max_len_extend^2 = ",
static_cast<int64_t>(num_seqs) * max_len_extend * max_len_extend,
", got ",
tree_mask_t.numel());
TORCH_CHECK(!is_cross_attn, "extend: tree_mask is not supported for cross attention");
// The window mask derives query positions from the row index
// (seq_len_prefix + m + row), but tree-mask rows sit at their tree depth,
// which is <= the row index; combining the two would over-mask the prefix.
TORCH_CHECK(sliding_window_size <= 0, "extend: tree_mask is not supported with sliding window attention");
tree_mask_ptr = tree_mask_t.data_ptr<bool>();
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] {
AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] {
if (max_len_extend <= 256) {
LAUNCH_EXTEND_ATTENTION_KERNEL(32, 64);
} else if (max_len_extend <= 1024) {
LAUNCH_EXTEND_ATTENTION_KERNEL(128, 256);
} else if (max_len_extend <= 4096) {
LAUNCH_EXTEND_ATTENTION_KERNEL(256, 768);
} else { // max_len_extend > 4096
LAUNCH_EXTEND_ATTENTION_KERNEL(512, 768);
}
});
});
}
python/sglang/kernels/aot/csrc/cpu/flash_attn.cpp
+555
View File
@@ -0,0 +1,555 @@
/*****************************************************************************************
* Copyright (c) 2025 - 2025 Codeplay Software Ltd. All rights reserved.
* Copyright (C) 2025 Intel Corporation, All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
****************************************************************************************/
#include "flash_attn.h"
#include "common.h"
#include "gemm.h"
// [NOTE]: flash attention interface for CPU
namespace {
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
void flash_attn_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ q,
const scalar_t* __restrict__ k,
const scalar_t* __restrict__ v,
void* __restrict__ buffer,
int seqlen_q,
int seqlen_k,
int batches,
int num_heads,
int num_heads_kv,
int head_size,
int head_size_v,
int q_strideM,
int q_strideH,
int k_strideN,
int k_strideH,
int v_strideN,
int v_strideH,
float sm_scale,
int buffer_size_per_thread,
bool causal) {
// strides
const int o_strideM = num_heads * head_size_v;
const int o_strideH = head_size_v;
// we use same buffer for packed key and value
const int ldb_tmp = std::max(head_size, head_size_v);
const int num_groups = num_heads / num_heads_kv;
TORCH_CHECK(num_groups * num_heads_kv == num_heads);
// number of super locks along M
int MB = div_up(seqlen_q, BLOCK_M);
// parallel on [batches, num_heads, MB]
parallel_for(batches * num_heads * MB, [&](int begin, int end) {
int bs{0}, head_id{0}, mb{0};
data_index_init(begin, bs, batches, head_id, num_heads, mb, MB);
int tid = get_thread_num();
// s_i and s_delta: [BLOCK_M, BLOCK_N]
float* __restrict__ s_i = reinterpret_cast<float*>((char*)(buffer) + tid * buffer_size_per_thread);
scalar_t* __restrict__ s_delta = reinterpret_cast<scalar_t*>(s_i);
// v_prime: [BLOCK_M, head_size_v]
float* __restrict__ v_prime = s_i + BLOCK_M * BLOCK_N;
// Btmp: [BLOCK_N, max(head_size, head_size_v)]
scalar_t* __restrict__ Btmp = reinterpret_cast<scalar_t*>(v_prime + BLOCK_M * head_size_v);
// init Btmp and Btmp2 just once for each thread to prevent NaN
fill_stub(Btmp, 0.f, BLOCK_N * ldb_tmp);
alignas(64) float s_prime[BLOCK_M];
alignas(64) float m_prime[BLOCK_M];
for (int i = begin; i < end; ++i) {
// [Note] use int64_t to avoid overflow
// For large inputs, for example bs = 4096, seqlen_q = 4097, m = 0, q_strideM = 128:
// The index calculated below: (seq_q_start_loc + m) * q_strideM = 4096 * 4097 * 128 will overflow int
int64_t seq_q_start_loc = bs * seqlen_q;
int64_t seq_k_start_loc = bs * seqlen_k;
// offset and size in MB
int m = mb * BLOCK_M;
int m_size = std::min(BLOCK_M, seqlen_q - m);
assert(m_size > 0);
int head_kv_id = head_id / num_groups;
// get query
const scalar_t* __restrict__ q_ptr = q + (seq_q_start_loc + m) * q_strideM + head_id * q_strideH;
// init v', s' and m'
fill_stub(v_prime, 0.f, m_size * head_size_v);
fill_stub(s_prime, 0.f, m_size);
fill_stub(m_prime, -std::numeric_limits<scalar_t>::infinity(), m_size);
int num_keys = causal ? std::min(m + m_size, seqlen_k) : seqlen_k;
for (int n = 0; n < num_keys; n += BLOCK_N) {
int n_size = std::min(BLOCK_N, num_keys - n);
// `n_size` is K in 2nd gemm, pad to TILE_K;
const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
// get key and pack
pack_vnni<scalar_t>(
/* dst */ Btmp,
/* src */ k + (seq_k_start_loc + n) * k_strideN + head_kv_id * k_strideH,
/* N */ n_size,
/* K */ head_size,
/* ld_src */ k_strideN,
/* ld_dst */ BLOCK_N);
// calculate s_i <- Q @ K
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ head_size,
/* lda */ q_strideM,
/* ldb */ BLOCK_N,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ q_ptr,
/* B */ Btmp,
/* C */ s_i);
// apply causal mask
// See [Note] condition to apply causal mask.
if (causal && n + n_size - 1 > m) {
for (int row = 0; row < m_size; ++row) {
int last_col = m + row - n;
// See [Note] mask the entire row if last_col < 0.
last_col = std::max(last_col, -1);
// fill [last_col + 1, n_size) to -inf
float* row_ptr = s_i + row * BLOCK_N;
fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
}
}
for (int row = 0; row < m_size; ++row) {
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row);
}
// get value and pack
pack_vnni2<scalar_t>(
/* dst */ Btmp,
/* src */ v + (seq_k_start_loc + n) * v_strideN + head_kv_id * v_strideH,
/* K */ n_size,
/* N */ head_size_v,
/* ld_src */ v_strideN,
/* ld_dst */ head_size_v);
// calculate V' <- s_delta @ V + V'
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ head_size_v,
/* K */ padded_n_size, // n_size
/* lda */ BLOCK_N,
/* ldb */ head_size_v,
/* ldc */ head_size_v,
/* add_C */ true,
/* A */ s_delta,
/* B */ Btmp,
/* C */ v_prime);
} // loop with seqlen_k
scalar_t* __restrict__ out_ptr = out + (seq_q_start_loc + m) * o_strideM + head_id * o_strideH;
for (int row = 0; row < m_size; ++row) {
float s = 1 / s_prime[row];
copy_stub<scalar_t>(out_ptr + row * o_strideM, v_prime + row * head_size_v, s, head_size_v);
}
// move to the next index
data_index_step(bs, batches, head_id, num_heads, mb, MB);
}
at::native::cpublas::brgemm_release();
});
}
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
void flash_attn_varlen_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ q,
const scalar_t* __restrict__ k,
const scalar_t* __restrict__ v,
const int32_t* __restrict__ cu_seqlens_q,
const int32_t* __restrict__ cu_seqlens_k,
void* __restrict__ buffer,
int32_t* __restrict__ indices,
int max_seqlen_q,
int max_seqlen_k,
int batches,
int num_heads,
int num_heads_kv,
int head_size,
int head_size_v,
int q_strideM,
int q_strideH,
int k_strideN,
int k_strideH,
int v_strideN,
int v_strideH,
float sm_scale,
int buffer_size_per_thread,
bool causal) {
// strides
const int o_strideM = num_heads * head_size_v;
const int o_strideH = head_size_v;
// compute index (bs, mb_offset) for Query blocks
// do this sequentially as usually problem size won't be big
int idx = 0;
for (int32_t bs = 0; bs < batches; ++bs) {
int32_t seqlen_q = cu_seqlens_q[bs + 1] - cu_seqlens_q[bs];
int32_t seqlen_k = cu_seqlens_k[bs + 1] - cu_seqlens_k[bs];
TORCH_CHECK(seqlen_q <= max_seqlen_q && seqlen_k <= max_seqlen_k);
int32_t blocks = div_up(seqlen_q, BLOCK_M);
for (int32_t offset = 0; offset < blocks; ++offset) {
indices[idx * 2 + 0] = bs;
indices[idx * 2 + 1] = offset;
idx++;
}
}
// number of query blocks
int MB = idx;
// we use same buffer for packed key and value
const int ldb_tmp = std::max(head_size, head_size_v);
const int num_groups = num_heads / num_heads_kv;
TORCH_CHECK(num_groups * num_heads_kv == num_heads);
// parallel on [MB, num_heads]
parallel_for(num_heads * MB, [&](int begin, int end) {
int head_id{0}, mb{0};
data_index_init(begin, head_id, num_heads, mb, MB);
int tid = get_thread_num();
// s_i and s_delta: [BLOCK_M, BLOCK_N]
float* __restrict__ s_i = reinterpret_cast<float*>((char*)(buffer) + tid * buffer_size_per_thread);
scalar_t* __restrict__ s_delta = reinterpret_cast<scalar_t*>(s_i);
// v_prime: [BLOCK_M, head_size_v]
float* __restrict__ v_prime = s_i + BLOCK_M * BLOCK_N;
// Btmp: [BLOCK_N, max(head_size, head_size_v)]
scalar_t* __restrict__ Btmp = reinterpret_cast<scalar_t*>(v_prime + BLOCK_M * head_size_v);
// init Btmp just once for each thread to prevent NaN
fill_stub(Btmp, 0.f, BLOCK_N * ldb_tmp);
alignas(64) float s_prime[BLOCK_M];
alignas(64) float m_prime[BLOCK_M];
for (int i = begin; i < end; ++i) {
int32_t bs = indices[mb * 2 + 0];
// See [Note] use int64_t to avoid overflow
int64_t seq_q_start_loc = cu_seqlens_q[bs];
int64_t seq_k_start_loc = cu_seqlens_k[bs];
int32_t seqlen_q = cu_seqlens_q[bs + 1] - cu_seqlens_q[bs];
// offset and size in MB
int m = indices[mb * 2 + 1] * BLOCK_M;
int m_size = std::min(BLOCK_M, seqlen_q - m);
assert(m_size > 0);
int head_kv_id = head_id / num_groups;
// get query
const scalar_t* __restrict__ q_ptr = q + (seq_q_start_loc + m) * q_strideM + head_id * q_strideH;
// init v', s' and m'
fill_stub(v_prime, 0.f, m_size * head_size_v);
fill_stub(s_prime, 0.f, m_size);
fill_stub(m_prime, -std::numeric_limits<scalar_t>::infinity(), m_size);
int seqlen_k = cu_seqlens_k[bs + 1] - cu_seqlens_k[bs];
int num_keys = causal ? std::min(m + m_size, seqlen_k) : seqlen_k;
for (int n = 0; n < num_keys; n += BLOCK_N) {
int n_size = std::min(BLOCK_N, num_keys - n);
// `n_size` is K in 2nd gemm, pad to TILE_K;
const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
// get key and pack
pack_vnni<scalar_t>(
/* dst */ Btmp,
/* src */ k + (seq_k_start_loc + n) * k_strideN + head_kv_id * k_strideH,
/* N */ n_size,
/* K */ head_size,
/* ld_src */ k_strideN,
/* ld_dst */ BLOCK_N);
// calculate s_i <- Q @ K
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ head_size,
/* lda */ q_strideM,
/* ldb */ BLOCK_N,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ q_ptr,
/* B */ Btmp,
/* C */ s_i);
// apply causal mask
// See [Note] condition to apply causal mask.
if (causal && n + n_size - 1 > m) {
for (int row = 0; row < m_size; ++row) {
int last_col = m + row - n;
// See [Note] mask the entire row if last_col < 0.
last_col = std::max(last_col, -1);
// fill [last_col + 1, n_size) to -inf
float* row_ptr = s_i + row * BLOCK_N;
fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
}
}
for (int row = 0; row < m_size; ++row) {
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row);
}
// get value and pack
pack_vnni2<scalar_t>(
/* dst */ Btmp,
/* src */ v + (seq_k_start_loc + n) * v_strideN + head_kv_id * v_strideH,
/* K */ n_size,
/* N */ head_size_v,
/* ld_src */ v_strideN,
/* ld_dst */ head_size_v);
// calculate V' <- s_delta @ V + V'
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ head_size_v,
/* K */ padded_n_size, // n_size
/* lda */ BLOCK_N,
/* ldb */ head_size_v,
/* ldc */ head_size_v,
/* add_C */ true,
/* A */ s_delta,
/* B */ Btmp,
/* C */ v_prime);
} // loop with seqlen_k
scalar_t* __restrict__ out_ptr = out + (seq_q_start_loc + m) * o_strideM + head_id * o_strideH;
for (int row = 0; row < m_size; ++row) {
float s = 1 / s_prime[row];
copy_stub<scalar_t>(out_ptr + row * o_strideM, v_prime + row * head_size_v, s, head_size_v);
}
// move to the next index
data_index_step(head_id, num_heads, mb, MB);
}
at::native::cpublas::brgemm_release();
});
}
} // anonymous namespace
template <typename index_t>
inline bool has_varlen_sequences(
const at::Tensor& cu_seqlens_q,
const at::Tensor& cu_seqlens_k,
int batches,
index_t max_seqlen_q,
index_t max_seqlen_k) {
const index_t* cu_seqlens_q_data = cu_seqlens_q.data_ptr<index_t>();
const index_t* cu_seqlens_k_data = cu_seqlens_k.data_ptr<index_t>();
for (int bs = 0; bs < batches; ++bs) {
index_t seqlen_q = cu_seqlens_q_data[bs + 1] - cu_seqlens_q_data[bs];
index_t seqlen_k = cu_seqlens_k_data[bs + 1] - cu_seqlens_k_data[bs];
if (seqlen_q != max_seqlen_q || seqlen_k != max_seqlen_k) {
return true;
}
}
return false;
}
template <int BLOCK_M, int BLOCK_N>
inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int head_size_v) {
static_assert(BLOCK_M <= BLOCK_N, "Make sure BLOCK_M <= BLOCK_N to prevent buffer overflows during causal masking");
const int size_per_thread =
/* s_i */ BLOCK_M * BLOCK_N * sizeof(float) +
/* v_prime */ BLOCK_M * head_size_v * sizeof(float) +
/* Btmp */ BLOCK_N * std::max(head_size, head_size_v) * sizeof(uint16_t);
buffer.resize_({num_threads, size_per_thread});
return size_per_thread;
}
template <int BLOCK_M>
inline void resize_indices(at::Tensor& indices, int num_seqs, int max_seqlen_q) {
// we allocate memory based on max seqlen
indices.resize_({num_seqs, div_up(max_seqlen_q, BLOCK_M), 2});
}
// [NOTE]: `flash_attn_varlen_func` AMX kernel
//
// q: [num_tokens, num_heads, head_size]
// k: [num_tokens, num_heads_kv, head_size]
// v: [num_tokens, num_heads_kv, head_size_v]
// cu_seqlens_q: [num_seqs + 1]
// cu_seqlens_k: [num_seqs + 1]
// out: [num_tokens, num_heads, head_size_v]
//
at::Tensor flash_attn_varlen_func(
const at::Tensor& q,
const at::Tensor& k,
const at::Tensor& v,
const at::Tensor& cu_seqlens_q,
const at::Tensor& cu_seqlens_k,
int64_t max_seqlen_q,
int64_t max_seqlen_k,
bool causal) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v);
CHECK_DIM(3, q);
CHECK_DIM(3, k);
CHECK_DIM(3, v);
CHECK_INPUT(cu_seqlens_q);
CHECK_INPUT(cu_seqlens_k);
CHECK_EQ(cu_seqlens_q.scalar_type(), at::kInt);
CHECK_EQ(cu_seqlens_k.scalar_type(), at::kInt);
int num_seqs = cu_seqlens_q.size(0) - 1;
int num_tokens = q.size(0);
int num_heads = q.size(1);
int num_heads_kv = k.size(1);
int head_size = q.size(2);
int head_size_v = v.size(2);
// strides for q, k and v
int q_strideM = q.stride(0);
int q_strideH = q.stride(1);
int k_strideN = k.stride(0);
int k_strideH = k.stride(1);
int v_strideN = v.stride(0);
int v_strideH = v.stride(1);
// check sizes
CHECK_EQ(k.size(2), head_size);
CHECK_EQ(v.size(1), num_heads_kv);
CHECK_EQ(cu_seqlens_k.size(0), num_seqs + 1);
// D and DV need to be even as we transpose by 512-bit
TORCH_CHECK(head_size % 2 == 0, "invalid head_size ", head_size);
TORCH_CHECK(head_size_v % 2 == 0, "invalid head_size_v ", head_size_v);
// softmax scale
double sm_scale = 1.0 / std::sqrt(static_cast<double>(head_size));
// check whether the batch has variant lengths
const bool is_varlen =
has_varlen_sequences<int32_t>(cu_seqlens_q, cu_seqlens_k, num_seqs, max_seqlen_q, max_seqlen_k);
int num_threads = at::get_num_threads();
at::Tensor buffer = at::empty({}, q.options().dtype(at::kChar));
at::Tensor indices = at::empty({}, q.options().dtype(at::kInt));
at::Tensor out = at::empty({num_tokens, num_heads, head_size_v}, q.options());
// TODO: tune the block size
constexpr int BLOCK_M = 512;
constexpr int BLOCK_N = 768;
AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "flash_attn_varlen_func", [&] {
int sz = resize_buffer<BLOCK_M, BLOCK_N>(buffer, num_threads, head_size, head_size_v);
if (is_varlen) {
resize_indices<BLOCK_M>(indices, num_seqs, max_seqlen_q);
flash_attn_varlen_kernel_impl<scalar_t, BLOCK_M, BLOCK_N>(
out.data_ptr<scalar_t>(),
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
cu_seqlens_q.data_ptr<int32_t>(),
cu_seqlens_k.data_ptr<int32_t>(),
buffer.data_ptr(),
indices.data_ptr<int32_t>(),
max_seqlen_q,
max_seqlen_k,
num_seqs,
num_heads,
num_heads_kv,
head_size,
head_size_v,
q_strideM,
q_strideH,
k_strideN,
k_strideH,
v_strideN,
v_strideH,
sm_scale,
sz,
causal);
} else {
flash_attn_kernel_impl<scalar_t, BLOCK_M, BLOCK_N>(
out.data_ptr<scalar_t>(),
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
buffer.data_ptr(),
max_seqlen_q,
max_seqlen_k,
num_seqs,
num_heads,
num_heads_kv,
head_size,
head_size_v,
q_strideM,
q_strideH,
k_strideN,
k_strideH,
v_strideN,
v_strideH,
sm_scale,
sz,
causal);
}
});
return out;
}
python/sglang/kernels/aot/csrc/cpu/flash_attn.h
+245
View File
@@ -0,0 +1,245 @@
#pragma once
#include "common.h"
#include "vec.h"
#include "vec_pack.h"
template <typename scalar_t>
inline void fill_stub(scalar_t* __restrict__ out, float val, int size) {
using Vec = at::vec::Vectorized<scalar_t>;
constexpr int kVecSize = Vec::size();
const Vec data_vec = Vec(static_cast<scalar_t>(val));
int d = 0;
#pragma GCC unroll 4
for (; d <= size - kVecSize; d += kVecSize) {
data_vec.store(out + d);
}
if (size - d > 0) {
data_vec.store(out + d, size - d);
}
}
template <typename scalar_t, int BLOCK_N>
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) {
static_assert(BLOCK_N % 32 == 0);
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int COLS = BLOCK_N / 16;
auto store = [&](auto i) {
constexpr int col = i % COLS;
// for COLS = 2, 4 use 512bit store
if constexpr (col % 2 == 0) {
auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16);
bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
out_bvec.store(out + col * 16);
}
};
Unroll<COLS>{}(store);
}
template <typename scalar_t>
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
const fVec s_fvec = fVec(s);
int d = 0;
#pragma GCC unroll 4
for (; d <= size - kVecSize; d += kVecSize) {
auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d);
a_fvec0 = a_fvec0 * s_fvec;
a_fvec1 = a_fvec1 * s_fvec;
bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
out_bvec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(acc[d] * s);
}
}
#if defined(CPU_CAPABILITY_AVX512)
template <>
inline void copy_stub<at::BFloat16>(at::BFloat16* __restrict__ out, const float* __restrict__ acc, float s, int size) {
const __m512 vscale = _mm512_set1_ps(s);
int d = 0;
#pragma GCC unroll 4
for (; d <= size - 32; d += 32) {
__m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale);
__m512 va1 = _mm512_mul_ps(_mm512_loadu_ps(acc + d + 16), vscale);
__m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0));
_mm512_storeu_si512(out + d, vb);
}
int remainder = size - d;
if (remainder > 0) {
if (remainder <= 16) {
const __mmask16 vmask = (1ULL << remainder) - 1;
__m512 va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d), vscale);
__m256i vb = (__m256i)(_mm512_cvtneps_pbh(va));
_mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(out + d), vmask, vb);
} else { // remainder > 16
const __mmask16 vmask = (1ULL << (remainder - 16)) - 1;
__m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale);
__m512 va1 = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d + 16), vscale);
__m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0));
const __mmask32 vmask2 = (1ULL << remainder) - 1;
_mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(out + d), vmask2, vb);
}
}
}
#endif
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
struct flash_attn_softmax {
static inline void apply(
float* __restrict__ s_i,
scalar_t* __restrict__ s_delta2,
float* __restrict__ v_prime,
float* __restrict__ s_prime,
float* __restrict__ m_prime,
int m_size,
int n_size,
int padded_n_size,
int head_size_v,
const float sm_scale,
int row) {
using Vec = at::vec::Vectorized<float>;
const Vec scale_vec = Vec(sm_scale);
float* s_delta = s_i;
// s_i <- s_i * scale
at::vec::map<float>([scale_vec](Vec x) { return x * scale_vec; }, s_i + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
// m_i: max value per row
float m_i =
at::vec::reduce_all<float>([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + row * BLOCK_N, n_size);
m_i = std::max(m_i, m_prime[row]);
// m_delta <- exp(m' - m_i)
float m_delta = std::exp(m_prime[row] - m_i);
// s_delta <- exp(s_i - m_i)
at::vec::map<float>(
[m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
// s' <- s' * m_delta + sum(s_delta)
s_prime[row] *= m_delta;
s_prime[row] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + row * BLOCK_N, n_size);
m_prime[row] = m_i;
// v' <- v' * m_delta
at::vec::map<float>(
[m_delta](Vec x) { return x * Vec(m_delta); },
v_prime + row * head_size_v,
v_prime + row * head_size_v,
head_size_v);
// Keep s_delta row-major for the following brgemm(P @ V), and only
// convert the columns that brgemm will consume.
fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size);
copy_stub<scalar_t>(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N, 1.f, padded_n_size);
}
};
#if defined(CPU_CAPABILITY_AVX512)
template <int BLOCK_M, int BLOCK_N>
struct flash_attn_softmax<at::BFloat16, BLOCK_M, BLOCK_N> {
static inline void apply(
float* __restrict__ s_i,
at::BFloat16* __restrict__ s_delta2,
float* __restrict__ v_prime,
float* __restrict__ s_prime,
float* __restrict__ m_prime,
int m_size,
int n_size,
int padded_n_size,
int head_size_v,
const float sm_scale,
int row) {
float* s_delta = s_i;
const __m512 vscale = _mm512_set1_ps(sm_scale);
int n_remainder = n_size & 15; // 0xF
const __mmask16 vmask = (1ULL << n_remainder) - 1;
int v_remainder = head_size_v & 15; // 0xF
const __mmask16 vmask1 = (1ULL << v_remainder) - 1;
constexpr float NEG_INF = -std::numeric_limits<float>::infinity();
__m512 va;
__m256i vb;
__m512 vmax;
__m512 vsum;
__m512 vmdelta;
const __m512 vneg_inf = _mm512_set1_ps(NEG_INF);
int m = row;
vmax = vneg_inf;
// s_i <- s_i * scale
int n = 0;
for (; n <= n_size - 16; n += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
vmax = _mm512_max_ps(va, vmax);
}
if (n_remainder > 0) {
va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
vmax = _mm512_max_ps(va, vmax);
}
// m_i: max value per row
float m_i = _mm512_reduce_max_ps(vmax);
m_i = std::max(m_i, m_prime[m]);
vmax = _mm512_set1_ps(m_i);
// m_delta <- exp(m' - m_i)
float m_delta = std::exp(m_prime[m] - m_i);
// s_delta <- exp(s_i - m_i)
vsum = _mm512_setzero_ps();
for (n = 0; n <= n_size - 16; n += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
vsum = _mm512_add_ps(vsum, va);
vb = (__m256i)(_mm512_cvtneps_pbh(va));
_mm256_storeu_si256(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vb);
}
if (n_remainder > 0) {
va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
vsum = _mm512_add_ps(vsum, va);
vb = (__m256i)(_mm512_cvtneps_pbh(va));
_mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vmask, vb);
}
// s' <- s' * m_delta + sum(s_delta)
s_prime[m] *= m_delta;
s_prime[m] += _mm512_reduce_add_ps(vsum);
m_prime[m] = m_i;
// pad s_delta with 0, pad_size range from [0, 32)
int pad_size = padded_n_size - n_size;
if (pad_size > 0) {
const __m512i vzero = _mm512_setzero_si512();
__mmask32 vmask2 = (1ULL << pad_size) - 1;
_mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(s_delta2 + m * BLOCK_N + n_size), vmask2, vzero);
}
// v' <- v' * m_delta
vmdelta = _mm512_set1_ps(m_delta);
int k = 0;
for (; k <= head_size_v - 16; k += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(v_prime + m * head_size_v + k), vmdelta);
_mm512_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), va);
}
if (v_remainder > 0) {
va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask1, v_prime + m * head_size_v + k), vmdelta);
_mm512_mask_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), vmask1, va);
}
}
};
#endif
python/sglang/kernels/aot/csrc/cpu/gemm.cpp
+842
View File
@@ -0,0 +1,842 @@
#include "gemm.h"
#include "common.h"
#include "vec.h"
namespace {
// packed layout:
// quants {N, K} int8_t
// comp {N} int32_t
template <int BLOCK_N>
inline void s8s8_compensation(int8_t* __restrict__ packed, int K) {
#if defined(CPU_CAPABILITY_AVX512)
constexpr int COLS = BLOCK_N / 16;
__m512i vcomp[COLS];
for (int col = 0; col < COLS; ++col) {
vcomp[col] = _mm512_setzero_si512();
}
const int64_t offset = BLOCK_N * K;
const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
for (int k = 0; k < K / 4; ++k) {
for (int col = 0; col < COLS; ++col) {
__m512i vb = _mm512_loadu_si512((const __m512i*)(packed + k * BLOCK_N * 4 + col * 64));
vcomp[col] = _mm512_dpbusd_epi32(vcomp[col], off, vb);
}
}
for (int col = 0; col < COLS; ++col) {
_mm512_storeu_si512((__m512i*)(packed + offset + col * 64), vcomp[col]);
}
#else
TORCH_CHECK(false, "s8s8_compensation not implemented!");
#endif
}
// convert to vnni format
// from [N, K] to [K/2, N, 2] for bfloat16 and float16
template <typename packed_t>
inline void pack_vnni(packed_t* __restrict__ packed, const packed_t* __restrict__ weight, int N, int K) {
const int VNNI_BLK = 2;
for (int n = 0; n < N; ++n) {
for (int k = 0; k < K / VNNI_BLK; ++k) {
for (int d = 0; d < VNNI_BLK; ++d) {
packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * K + k * VNNI_BLK + d];
}
}
}
}
template <>
inline void pack_vnni<int8_t>(int8_t* __restrict__ packed, const int8_t* __restrict__ weight, int N, int K) {
constexpr int BLOCK_N = block_size_n();
TORCH_CHECK(N == BLOCK_N);
const int VNNI_BLK = 4;
for (int n = 0; n < N; ++n) {
for (int k = 0; k < K / VNNI_BLK; ++k) {
for (int d = 0; d < VNNI_BLK; ++d) {
packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * K + k * VNNI_BLK + d];
}
}
}
s8s8_compensation<BLOCK_N>(packed, K);
}
// uint8_t: mxfp4 or int4
// pack to vnni2 format as they are computed with bfloat16
//
// from [N, K'/2, 2] to [K'/2, N, 2], view 2x int4 as unit8:
// from [N, K ] to [K, N ] where K = K'/2
//
template <>
inline void pack_vnni<uint8_t>(uint8_t* __restrict__ packed, const uint8_t* __restrict__ weight, int N, int K) {
constexpr int BLOCK_N = block_size_n();
uint8_t unpacked[2 * BLOCK_N];
// 32-way pack (align with BLOCK_N), faster for avx512 unpacking
//
// for a range of (64):
// {0, 1, 2, ..., 63}
//
// original format:
// { 1|0, 3|2, ..., 63|62}
//
// packed format:
// {32|0, 31|1, ..., 63|31}
//
for (int k = 0; k < K; ++k) {
// unpack first
for (int n = 0; n < N; ++n) {
uint8_t value = weight[n * K + k];
unpacked[n * 2 + 0] = value & 0xF; // lower 4 bits
unpacked[n * 2 + 1] = value >> 4; // higher 4 bits
}
// re-pack to 32-way
for (int n = 0; n < N; ++n) {
packed[k * N + n] = (unpacked[n + BLOCK_N] << 4) | unpacked[n];
}
}
}
template <typename scalar_t>
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [data0, data1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d]);
}
}
template <typename scalar_t>
inline void copy_stub(float* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [data0, data1] = load_float_vec2(input + d);
data0.store(out + d);
data1.store(out + d + fVec::size());
}
for (; d < size; ++d) {
out[d] = static_cast<float>(input[d]);
}
}
template <typename scalar_t>
inline void copy_add_stub(
scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [data0, data1] = load_float_vec2(input + d);
auto [bias0, bias1] = load_float_vec2(bias + d);
bVec out_vec = convert_from_float_ext<scalar_t>(data0 + bias0, data1 + bias1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] + bias[d]);
}
}
template <typename scalar_t, bool has_bias>
inline void scalar_sigmoid_and_mul(
scalar_t* __restrict__ out,
const float* __restrict__ input,
const float* __restrict__ bias,
const scalar_t* __restrict__ mul,
int SIZE) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
// scalar sigmoid
fVec X;
if constexpr (has_bias) {
assert(bias != nullptr);
X = fVec(input[0] + bias[0]);
} else {
X = fVec(input[0]);
}
X = fast_sigmoid(X);
// vec mul
constexpr int kVecSize = bVec::size();
for (int d = 0; d < SIZE; d += kVecSize) {
auto [m_fvec0, m_fvec1] = load_float_vec2(mul + d);
bVec out_vec = convert_from_float_ext<scalar_t>(m_fvec0 * X, m_fvec1 * X);
out_vec.store(out + d);
}
}
template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn {
static inline void apply(
const scalar_t* __restrict__ A,
const scalar_t* __restrict__ B,
scalar_t* __restrict__ C,
const float* __restrict__ bias,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
}
};
#if defined(CPU_CAPABILITY_AVX512)
template <bool has_bias, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
static inline void apply(
const at::BFloat16* __restrict__ A,
const at::BFloat16* __restrict__ B,
at::BFloat16* __restrict__ C,
const float* __restrict__ bias,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
constexpr int ROWS = BLOCK_M;
constexpr int COLS = BLOCK_N / 16;
// prefetch distance
constexpr int PREFETCH_SIZE_K = 0;
__m512bh va;
__m512bh vb[COLS];
__m512 vc[ROWS * COLS];
auto loadc = [&](auto i) {
constexpr int col = i % COLS;
if constexpr (has_bias) {
vc[i] = _mm512_loadu_ps(bias + col * 16);
} else {
vc[i] = _mm512_set1_ps(0.f);
}
};
Unroll<ROWS * COLS>{}(loadc);
const int64_t K2 = K >> 1;
const int64_t lda2 = lda >> 1;
const int64_t ldb2 = ldb; // ldb * 2 >> 1;
const float* a_ptr = reinterpret_cast<const float*>(A);
const float* b_ptr = reinterpret_cast<const float*>(B);
auto compute = [&](auto i, int64_t k) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
if constexpr (col == 0) {
va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
}
if constexpr (row == 0) {
vb[col] = (__m512bh)(_mm512_loadu_si512(b_ptr + k * ldb2 + col * 16));
if constexpr (PREFETCH_SIZE_K > 0) {
_mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
}
}
vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
};
for (int64_t k = 0; k < K2; ++k) {
Unroll<ROWS * COLS>{}(compute, k);
}
auto storec = [&](auto i) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
// for COLS = 2, 4 use 512bit store
// for COLS = 1, 3 use 256bit store
if constexpr (COLS % 2 == 0) {
if constexpr (col % 2 == 0) {
_mm512_storeu_si512(
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
(__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
}
} else {
_mm256_storeu_si256(reinterpret_cast<__m256i*>(C + row * ldc + col * 16), (__m256i)(_mm512_cvtneps_pbh(vc[i])));
}
};
Unroll<ROWS * COLS>{}(storec);
}
};
#endif
#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
A + mb_start * lda, \
B + nb_start * 2, \
C + mb_start * ldc + nb_start, \
has_bias ? bias + nb_start : nullptr, \
K, \
lda, \
ldb, \
ldc);
template <typename scalar_t, bool has_bias>
struct brgemm {
static inline void apply(
const scalar_t* __restrict__ A,
const scalar_t* __restrict__ B,
scalar_t* __restrict__ C,
float* __restrict__ Ctmp,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
constexpr int BLOCK_N = block_size_n();
at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
// copy from Ctmp to C
for (int64_t m = 0; m < M; ++m) {
if constexpr (has_bias) {
copy_add_stub(C + m * ldc, Ctmp + m * BLOCK_N, bias, N);
} else {
copy_stub(C + m * ldc, Ctmp + m * BLOCK_N, N);
}
}
}
static inline void apply(
const float* __restrict__ A,
const float* __restrict__ B,
scalar_t* __restrict__ C,
float* __restrict__ Ctmp,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
constexpr int BLOCK_N = block_size_n();
at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
}
};
template <typename scalar_t, bool has_bias>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const scalar_t* __restrict__ B,
scalar_t* __restrict__ C,
float* __restrict__ Ctmp,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg) {
if (brg) {
brgemm<scalar_t, has_bias>::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc);
return;
}
// pattern: 1-4-16, N = 16, 32, 48, 64
constexpr int64_t BLOCK_M = 4;
constexpr int64_t BLOCK_N = 64;
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
for (int mb = 0; mb < MB; ++mb) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
for (int64_t nb = 0; nb < NB; ++nb) {
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
switch (mb_size << 4 | nb_size >> 4) {
// mb_size = 1
case 0x11:
LAUNCH_TINYGEMM_KERNEL_NN(1, 16);
break;
case 0x12:
LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
break;
case 0x13:
LAUNCH_TINYGEMM_KERNEL_NN(1, 48);
break;
case 0x14:
LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
break;
// mb_size = 2
case 0x21:
LAUNCH_TINYGEMM_KERNEL_NN(2, 16);
break;
case 0x22:
LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
break;
case 0x23:
LAUNCH_TINYGEMM_KERNEL_NN(2, 48);
break;
case 0x24:
LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
break;
// mb_size = 3
case 0x31:
LAUNCH_TINYGEMM_KERNEL_NN(3, 16);
break;
case 0x32:
LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
break;
case 0x33:
LAUNCH_TINYGEMM_KERNEL_NN(3, 48);
break;
case 0x34:
LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
break;
// mb_size = 4
case 0x41:
LAUNCH_TINYGEMM_KERNEL_NN(4, 16);
break;
case 0x42:
LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
break;
case 0x43:
LAUNCH_TINYGEMM_KERNEL_NN(4, 48);
break;
case 0x44:
LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", mb_size, " x ", nb_size);
}
}
}
}
template <typename scalar_t, bool has_bias>
void tinygemm_kernel(
const float* __restrict__ A,
const float* __restrict__ B,
scalar_t* __restrict__ C,
float* __restrict__ Ctmp,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg) {
TORCH_CHECK(brg, "Expected to use fp32 brgemm for small N GEMM");
if (brg) {
brgemm<scalar_t, has_bias>::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc);
return;
}
// TODO : add intrinsic path
}
template <typename scalar_t>
void weight_packed_linear_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ mat1,
const scalar_t* __restrict__ mat2,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K,
int64_t mat1_strideM,
int64_t out_strideM) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
const bool use_brgemm = can_use_brgemm<scalar_t>(M);
// parallel on [MB, NB]
AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// for brgemm, use float32 for accumulate
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(M - mb_start, BLOCK_M);
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(N - nb_start, BLOCK_N);
tinygemm_kernel<scalar_t, has_bias>(
/* A */ mat1 + mb_start * mat1_strideM,
/* B */ mat2 + nb_start * K /* nb * BLOCK_N * K */,
/* C */ out + mb_start * out_strideM + nb_start,
/* Ctmp*/ Ctmp,
/* bias*/ bias + nb_start,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ mat1_strideM,
/* ldb */ nb_size,
/* ldc */ out_strideM,
/* brg */ use_brgemm);
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
});
}
template <typename scalar_t>
void weight_packed_linear_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ mat1,
const float* __restrict__ mat2,
const float* __restrict__ bias,
const scalar_t* __restrict__ post_mul_mat,
int64_t M,
int64_t N,
int64_t K,
int64_t mat1_strideM,
int64_t out_strideM) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
const bool use_brgemm = true; // TODO: add intrinsic path
// parallel on [MB, NB]
AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// for brgemm, use float32 for accumulate
alignas(64) float Atmp[BLOCK_M * K];
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
loop_2d<float>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(M - mb_start, BLOCK_M);
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(N - nb_start, BLOCK_N);
for (int64_t m = 0; m < mb_size; ++m) {
copy_stub<scalar_t>(Atmp + m * K, mat1 + mb_start * mat1_strideM + m * K, K);
}
tinygemm_kernel<scalar_t, has_bias>(
/* A */ Atmp,
/* B */ mat2 + nb_start * K /* nb * BLOCK_N * K */,
/* C */ out + mb_start * out_strideM + nb_start,
/* Ctmp*/ Ctmp,
/* bias*/ bias + nb_start,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ mat1_strideM,
/* ldb */ nb_size,
/* ldc */ out_strideM,
/* brg */ use_brgemm);
if (post_mul_mat != nullptr) {
for (int64_t m = 0; m < mb_size; ++m) {
scalar_sigmoid_and_mul<scalar_t, has_bias>(
out + mb_start * out_strideM + nb_start + m * out_strideM,
Ctmp + m * BLOCK_N,
bias + nb_start,
post_mul_mat + mb_start * out_strideM + m * out_strideM,
out_strideM);
}
} else {
for (int64_t m = 0; m < mb_size; ++m) {
if constexpr (has_bias) {
copy_add_stub(
out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, bias + nb_start, N);
} else {
copy_stub(out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, N);
}
}
}
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
});
}
} // anonymous namespace
// tinygemm interface
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const scalar_t* __restrict__ B,
scalar_t* __restrict__ C,
float* __restrict__ Ctmp,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg) {
tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, nullptr, M, N, K, lda, ldb, ldc, brg);
}
#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
template void tinygemm_kernel<TYPE>( \
const TYPE* __restrict__ A, \
const TYPE* __restrict__ B, \
TYPE* __restrict__ C, \
float* __restrict__ Ctmp, \
int64_t M, \
int64_t N, \
int64_t K, \
int64_t lda, \
int64_t ldb, \
int64_t ldc, \
bool brg)
INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
at::Tensor convert_weight_packed(at::Tensor& weight) {
// for 3d moe weights
// weight : [E, OC, IC]
// w1 : [E, 2N, K]
// w2 : [E, K, N]
CHECK_INPUT(weight);
const int64_t ndim = weight.ndimension();
TORCH_CHECK(ndim == 2 || ndim == 3, "expect weight to be 2d or 3d, got ", ndim, "d tensor.");
if (ndim == 2 && weight.size(0) < TILE_N) {
// for 2D weight and small OC shape, we use fma linear path, which needs transpose not pack
return weight.to(at::kFloat).t().contiguous();
}
const auto st = weight.scalar_type();
const int64_t E = ndim == 3 ? weight.size(0) : 1;
const int64_t OC = ndim == 3 ? weight.size(1) : weight.size(0);
const int64_t IC = ndim == 3 ? weight.size(2) : weight.size(1);
// mxfp4 or int4 are packed with uint8
const int64_t actual_IC = st == at::kByte ? IC * 2 : IC;
// we handle 2 TILE_N at a time.
TORCH_CHECK(OC % TILE_N == 0, "invalid weight out features ", OC);
TORCH_CHECK(actual_IC % TILE_K == 0, "invalid weight input features ", actual_IC);
constexpr int64_t BLOCK_N = block_size_n();
const int64_t NB = div_up(OC, BLOCK_N);
// use phony sizes here [E, OC, IC], for each [E], [OC, IC] -> [IC / 2, OC, 2]
auto packed_weight = at::empty({}, weight.options());
const int64_t stride = OC * IC;
// Note: for `kByte` (uint8), it represents either `mxfp4` or `int4`.
TORCH_CHECK(
st == at::kBFloat16 || st == at::kHalf || st == at::kChar || st == at::kFloat8_e4m3fn || st == at::kByte,
"expect weight to be bfloat16, float16, int8, fp8_e4m3 or uint8(mxfp4 or int4).");
CPU_DISPATCH_PACKED_TYPES(st, [&] {
// adjust most inner dimension size
const int packed_row_size = get_row_size<packed_t>(actual_IC);
auto sizes = weight.sizes().vec();
sizes[ndim - 1] = packed_row_size;
packed_weight.resize_(sizes);
const packed_t* w_data = weight.data_ptr<packed_t>();
packed_t* packed_data = packed_weight.data_ptr<packed_t>();
// parallel on {E, NB}
at::parallel_for(0, E * NB, 0, [&](int64_t begin, int64_t end) {
int64_t e{0}, nb{0};
data_index_init(begin, e, E, nb, NB);
for (int64_t i = begin; i < end; ++i) {
UNUSED(i);
int64_t n = nb * BLOCK_N;
int64_t n_size = std::min(BLOCK_N, OC - n);
pack_vnni<packed_t>(
packed_data + e * OC * packed_row_size + n * packed_row_size, w_data + e * stride + n * IC, n_size, IC);
// move to the next index
data_index_step(e, E, nb, NB);
}
});
});
return packed_weight;
}
at::Tensor convert_scale_packed(at::Tensor& scale) {
CHECK_INPUT(scale);
const int64_t ndim = scale.ndimension();
TORCH_CHECK(ndim == 2 || ndim == 3, "expect scale to be 2d or 3d, got ", ndim, "d tensor.");
const auto st = scale.scalar_type();
const int64_t E = ndim == 3 ? scale.size(0) : 1;
const int64_t N = ndim == 3 ? scale.size(1) : scale.size(0);
// number of groups, e.g. K/32
const int64_t G = ndim == 3 ? scale.size(2) : scale.size(1);
constexpr int64_t BLOCK_N = block_size_n();
TORCH_CHECK(N % BLOCK_N == 0, "invalid weight out features ", N);
const int64_t NB = N / BLOCK_N;
auto packed_scale = at::empty_like(scale);
TORCH_CHECK(st == at::kByte, "expect scale to be uint8.");
const uint8_t* s_data = scale.data_ptr<uint8_t>();
uint8_t* packed_data = packed_scale.data_ptr<uint8_t>();
// parallel on src {E, NB, BLOCK_N, G}, dst {E, NB, G, BLOCK_N}
at::parallel_for(0, E * NB * BLOCK_N * G, 0, [&](int64_t begin, int64_t end) {
int64_t e{0}, nb{0}, n{0}, g{0};
data_index_init(begin, e, E, nb, NB, n, BLOCK_N, g, G);
for (int64_t i = begin; i < end; ++i) {
packed_data[e * N * G + nb * G * BLOCK_N + g * BLOCK_N + n] = s_data[i];
// move to the next index
data_index_step(e, E, nb, NB, n, BLOCK_N, g, G);
}
});
return packed_scale;
}
// mat1 : [*, K]
// mat2 : [N, K] ([K, N] if use_fma_gemm)
// bias : [N]
// out : [*, N]
//
at::Tensor
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
bool use_fma_gemm = false;
if (packed_w.scalar_type() == at::kFloat) {
use_fma_gemm = true;
}
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
CHECK_INPUT(mat2);
const int64_t ndim = mat1.ndimension();
auto input_sizes = mat1.sizes().vec();
int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
int64_t K = use_fma_gemm ? mat1.size(1) : mat2.size(1);
int64_t M = use_fma_gemm ? mat1.size(0) : mat1.numel() / K;
CHECK_DIM(2, mat2);
if (use_fma_gemm) {
CHECK_DIM(2, mat1);
} else {
CHECK_EQ(mat1.size(ndim - 1), K);
}
auto dispatch_type = mat1.scalar_type();
auto out = at::empty({M, N}, mat1.options());
// strides
int64_t out_strideM = out.stride(0);
int64_t mat1_strideM = mat1.stride(-2);
const bool has_bias = bias.has_value();
const float* bias_data = nullptr;
if (has_bias) {
CHECK_EQ(bias.value().size(0), N);
bias_data = bias.value().data_ptr<float>();
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "weight_packed_linear_kernel_impl", [&] {
if (use_fma_gemm) {
weight_packed_linear_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
mat1.data_ptr<scalar_t>(),
packed_w.data_ptr<float>(),
bias_data,
nullptr,
M,
N,
K,
mat1_strideM,
out_strideM);
} else {
weight_packed_linear_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
mat1.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
bias_data,
M,
N,
K,
mat1_strideM,
out_strideM);
}
});
input_sizes[ndim - 1] = N;
return out.view(input_sizes);
}
// mat1 : [M, K]
// mat2 : [K, 1]
// post_mul_mat : [M, K]
// bias : [N]
// out : [M, N]
//
at::Tensor fused_linear_sigmoid_mul(
at::Tensor& mat1,
at::Tensor& mat2,
const std::optional<at::Tensor>& bias,
bool is_vnni,
const at::Tensor& post_mul_mat) {
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
TORCH_CHECK(packed_w.scalar_type() == at::kFloat, "fused_linear_sigmoid_mul requires packed float weight")
int64_t M = mat1.size(0);
int64_t K = mat1.size(1);
int64_t N = mat2.size(1);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
CHECK_INPUT(mat2);
CHECK_DIM(2, mat1);
CHECK_DIM(2, mat2);
int64_t out_strideM = post_mul_mat.size(1);
int64_t mat1_strideM = mat1.stride(0);
auto dispatch_type = mat1.scalar_type();
auto out = at::empty({M, out_strideM}, mat1.options());
TORCH_CHECK(
N == 1 && out_strideM % 32 == 0,
"post_mul_mat tensor size(1) should be 32 dividable, and the mat2 OC=1 (Mx1 as linear output shape)")
const bool has_bias = bias.has_value();
const float* bias_data = nullptr;
if (has_bias) {
CHECK_EQ(bias.value().size(0), N);
bias_data = bias.value().data_ptr<float>();
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "fused_linear_sigmoid_mul", [&] {
weight_packed_linear_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
mat1.data_ptr<scalar_t>(),
packed_w.data_ptr<float>(),
bias_data,
post_mul_mat.data_ptr<scalar_t>(),
M,
N,
K,
mat1_strideM,
out_strideM);
});
return out;
}
python/sglang/kernels/aot/csrc/cpu/gemm.h
+373
View File
@@ -0,0 +1,373 @@
#pragma once
#include <ATen/native/CPUBlas.h>
#include "common.h"
// amx-bf16
#define TILE_M 16
#define TILE_N 16
#define TILE_K 32
// block size for AMX gemm
constexpr int block_size_m() {
return 2 * TILE_M;
}
constexpr int block_size_n() {
return 2 * TILE_N;
}
// define threshold using brgemm (intel AMX)
template <typename T>
inline bool can_use_brgemm(int M);
template <>
inline bool can_use_brgemm<at::BFloat16>(int M) {
return M > 4;
}
template <>
inline bool can_use_brgemm<at::Half>(int M) {
return true;
}
// this requires PyTorch 2.7 or above
template <>
inline bool can_use_brgemm<int8_t>(int M) {
return M > 4;
}
template <>
inline bool can_use_brgemm<uint8_t>(int M) {
return M > 4;
}
template <>
inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
return M > 4;
}
// work around compiler internal error
#define BLOCK_K 128 // 4 * TILE_K
// adjust leading dimension size for K
template <typename T>
inline int64_t get_row_size(int64_t K) {
return K;
}
template <>
inline int64_t get_row_size<int8_t>(int64_t K) {
return K + sizeof(int32_t);
}
// uint8: mxfp4 or int4
template <>
inline int64_t get_row_size<uint8_t>(int64_t K) {
return K >> 1;
}
inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
return use_int8_w8a8 ? K + sizeof(int32_t) : K;
}
enum class CPUActMethod : int {
silu_and_mul = 0,
swiglu = 1,
gelu_and_mul = 2,
};
enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };
constexpr bool operator==(CPUQuantMethod a, int64_t b) {
return static_cast<int64_t>(a) == b;
}
constexpr bool operator==(int64_t a, CPUQuantMethod b) {
return a == static_cast<int64_t>(b);
}
enum class CPUQuantAlgo : int64_t { AWQ = 0, GPTQ = 1 };
constexpr bool operator==(CPUQuantAlgo a, int64_t b) {
return static_cast<int64_t>(a) == b;
}
constexpr bool operator==(int64_t a, CPUQuantAlgo b) {
return a == static_cast<int64_t>(b);
}
inline int64_t get_row_size(CPUQuantMethod quant, int64_t K) {
switch (quant) {
case CPUQuantMethod::INT8_W8A8:
return K + sizeof(int32_t);
case CPUQuantMethod::MXFP4:
return K >> 1;
default:
return K;
}
}
inline int64_t get_4bit_block_k_size(int64_t group_size) {
return group_size > 128 ? 128 : group_size;
}
// pack weight to vnni format
at::Tensor convert_weight_packed(at::Tensor& weight);
// pack weight to vnni format for int4
std::tuple<at::Tensor, at::Tensor, at::Tensor>
convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
// moe implementations for int8 w8a8
template <typename scalar_t>
void fused_experts_int8_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic1,
scalar_t* __restrict__ ic2,
uint8_t* __restrict__ A_tmp,
float* __restrict__ C_tmp,
uint8_t* __restrict__ Aq_tmp,
float* __restrict__ As_tmp,
const scalar_t* __restrict__ input,
const int8_t* __restrict__ packed_w1,
const int8_t* __restrict__ packed_w2,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
const float* __restrict__ topk_weights,
const int32_t* __restrict__ sorted_ids,
const int32_t* __restrict__ expert_ids,
const int32_t* __restrict__ offsets,
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad);
// moe implementations for fp8 w8a16 and mxfp4
template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
void fused_experts_fp_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic0,
scalar_t* __restrict__ ic1,
scalar_t* __restrict__ ic2,
scalar_t* __restrict__ A_tmp,
scalar_t* __restrict__ B_tmp,
float* __restrict__ C_tmp,
const scalar_t* __restrict__ input,
const packed_t* __restrict__ packed_w1,
const packed_t* __restrict__ packed_w2,
const float* __restrict__ w1_bias,
const float* __restrict__ w2_bias,
const param_t* __restrict__ w1s,
const param_t* __restrict__ w2s,
int64_t block_size_N,
int64_t block_size_K,
const float* __restrict__ topk_weights,
const int32_t* __restrict__ sorted_ids,
const int32_t* __restrict__ expert_ids,
const int32_t* __restrict__ offsets,
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad,
float alpha,
float limit,
CPUActMethod act_func,
bool with_bias);
// shared expert implementation for int8 w8a8
template <typename scalar_t>
void shared_expert_int8_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic1,
float* __restrict__ C_tmp,
uint8_t* __restrict__ Aq_tmp,
float* __restrict__ As_tmp,
const scalar_t* __restrict__ input,
const int8_t* __restrict__ packed_w1,
const int8_t* __restrict__ packed_w2,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
const scalar_t* __restrict__ fused_experts_out,
float routed_scaling_factor,
int64_t M,
int64_t N,
int64_t K);
template <typename scalar_t>
void fused_experts_int4_w4a8_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic0,
scalar_t* __restrict__ ic1,
scalar_t* __restrict__ ic2,
uint8_t* __restrict__ A_tmp,
uint8_t* __restrict__ Aq_tmp,
float* __restrict__ As_tmp,
int32_t* __restrict__ Azp_tmp,
float* __restrict__ C_tmp,
int8_t* __restrict__ dqB_tmp,
const scalar_t* __restrict__ input,
const uint8_t* __restrict__ packed_w1,
const uint8_t* __restrict__ packed_w2,
const int8_t* __restrict__ w1z,
const int8_t* __restrict__ w2z,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
int group_size,
const float* __restrict__ topk_weights,
const int32_t* __restrict__ sorted_ids,
const int32_t* __restrict__ expert_ids,
const int32_t* __restrict__ offsets,
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad);
template <typename scalar_t>
void shared_expert_fp8_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic0,
scalar_t* __restrict__ ic1,
scalar_t* __restrict__ B_tmp,
float* __restrict__ C_tmp,
const scalar_t* __restrict__ input,
const at::Float8_e4m3fn* __restrict__ packed_w1,
const at::Float8_e4m3fn* __restrict__ packed_w2,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
int64_t block_size_N,
int64_t block_size_K,
const scalar_t* __restrict__ fused_experts_out,
float routed_scaling_factor,
int64_t M,
int64_t N,
int64_t K);
// tinygemm interface
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const scalar_t* __restrict__ B,
scalar_t* __restrict__ C,
float* __restrict__ Ctmp,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg);
template <typename scalar_t>
void tinygemm_kernel(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B,
scalar_t* __restrict__ C,
int32_t* __restrict__ Ctmp,
const float* __restrict__ As,
const float* __restrict__ Bs,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg);
// block quantization
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const at::Float8_e4m3fn* __restrict__ B,
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const float* __restrict__ Bbias,
const float* __restrict__ scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg,
int64_t block_size_K,
bool do_unpack = true);
// per tensor quantization
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const at::Float8_e4m3fn* __restrict__ B,
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
float scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg);
// mxfp4
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const uint8_t* __restrict__ B,
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const float* __restrict__ Bbias,
const uint8_t* __restrict__ scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg,
int64_t block_size_K,
bool do_unpack = true);
template <typename scalar_t>
void tinygemm_kernel(
scalar_t* C,
float* C_temp,
const uint8_t* A,
const float* scales_a,
const int32_t* qzeros_a,
const uint8_t* B,
const float* scales_b,
const int8_t* qzeros_b,
const int32_t* compensation,
int8_t* dqB_tmp,
int64_t M,
int64_t K,
int64_t lda,
int64_t ldc_f,
int64_t ldc_s,
bool store_out,
bool use_brgemm);
// mxfp4
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const uint8_t* __restrict__ B,
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const uint8_t* __restrict__ scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg,
int64_t block_size_K,
bool do_unpack = true);
python/sglang/kernels/aot/csrc/cpu/gemm_fp8.cpp
+1236
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/aot/csrc/cpu/gemm_int4.cpp
+889
View File
@@ -0,0 +1,889 @@
#include <torch/all.h>
#include "gemm.h"
#include "vec.h"
namespace {
#define BLOCK_N block_size_n()
#define BLOCK_M 128
template <bool sym_quant_act>
struct ActDtype;
template <>
struct ActDtype<true> {
using type = int8_t;
};
template <>
struct ActDtype<false> {
using type = uint8_t;
};
#if defined(CPU_CAPABILITY_AVX512)
struct alignas(32) m256i_wrapper {
__m256i data;
};
inline std::array<m256i_wrapper, 2> load_zps_4vnni(const int8_t* __restrict__ zps) {
// broadcast 01234567 to
// 01234567012345670123456701234567
__m256i vzps_low = _mm256_set1_epi64x(*reinterpret_cast<const long*>(zps));
__m256i vzps_high = _mm256_set1_epi64x(*reinterpret_cast<const long*>(zps + 8));
// shuffle from
// 01234567012345670123456701234567
// to
// 00001111222233334444555566667777
__m256i shuffle_mask =
_mm256_set_epi8(7, 7, 7, 7, 6, 6, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0);
vzps_low = _mm256_shuffle_epi8(vzps_low, shuffle_mask);
vzps_high = _mm256_shuffle_epi8(vzps_high, shuffle_mask);
m256i_wrapper vzps_low_wp, vzps_high_wp;
vzps_low_wp.data = vzps_low;
vzps_high_wp.data = vzps_high;
return {vzps_low_wp, vzps_high_wp};
}
inline std::array<m256i_wrapper, 2> load_uint4_as_int8(const uint8_t* __restrict__ qB) {
__m256i packed = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(qB));
const __m256i low_mask = _mm256_set1_epi8(0x0f);
__m256i high = _mm256_srli_epi16(packed, 4);
high = _mm256_and_si256(high, low_mask);
__m256i low = _mm256_and_si256(packed, low_mask);
m256i_wrapper low_wp, high_wp;
low_wp.data = low;
high_wp.data = high;
return {low_wp, high_wp};
}
template <int64_t N, int64_t ldb>
void _dequant_weight_zp_only(const uint8_t* __restrict__ B, int8_t* dqB, const int8_t* __restrict__ qzeros, int64_t K) {
// unpack weight int8 -> two int4
// subtract zero point
// B shape = [K, ldb] = [K, N / 2], actual shape = [K / 4, N / 2, 4]
// dqB shape = [K, N], actual shape = [K / 4, N, 4]
#pragma GCC unroll 2
for (int n = 0; n < N; n += 16) {
auto [zps_low_wp, zps_high_wp] = load_zps_4vnni(&qzeros[n]);
auto zps_low = zps_low_wp.data;
auto zps_high = zps_high_wp.data;
for (int k = 0; k < K; k += 4) {
auto [vb_low_wp, vb_high_wp] = load_uint4_as_int8(B + ldb * k + n / 2 * 4);
auto vb_low = vb_low_wp.data;
auto vb_high = vb_high_wp.data;
vb_high = _mm256_sub_epi8(vb_high, zps_high);
vb_low = _mm256_sub_epi8(vb_low, zps_low);
// store vb to B
_mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + n * 4), vb_low);
_mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + (n + 8) * 4), vb_high);
}
}
}
template <bool accum, int64_t N, bool sym_quant_act>
void _dequant_and_store(
float* __restrict__ output,
const int32_t* __restrict__ input,
const float* __restrict__ scale_a,
const int32_t* __restrict__ zp_a,
const float* __restrict__ scale_b,
const int32_t* __restrict__ comp_b,
int M,
int ldi,
int ldo,
int ldsa = 1) {
for (int m = 0; m < M; ++m) {
float a_scale = *(scale_a + m * ldsa);
__m512 va_scale = _mm512_set1_ps(a_scale);
int32_t a_zp;
__m512i va_zp;
if constexpr (!sym_quant_act) {
a_zp = *(zp_a + m * ldsa);
va_zp = _mm512_set1_epi32(a_zp);
}
int n = 0;
#pragma GCC unroll 2
for (; n < N; n += 16) {
__m512i vc = _mm512_loadu_si512(input + m * ldi + n);
if constexpr (!sym_quant_act) {
__m512i vb_comp = _mm512_loadu_si512(comp_b + n);
vc = _mm512_sub_epi32(vc, _mm512_mullo_epi32(vb_comp, va_zp));
}
__m512 vc_f = _mm512_cvtepi32_ps(vc);
__m512 vc_f_mul = _mm512_mul_ps(vc_f, va_scale);
__m512 vb_s = _mm512_loadu_ps(scale_b + n);
vc_f_mul = _mm512_mul_ps(vc_f_mul, vb_s);
if constexpr (accum) {
__m512 vo = _mm512_loadu_ps(output + m * ldo + n);
_mm512_storeu_ps(output + m * ldo + n, _mm512_add_ps(vo, vc_f_mul));
} else {
_mm512_storeu_ps(output + m * ldo + n, vc_f_mul);
}
}
for (; n < N; ++n) {
float dq_val;
if constexpr (sym_quant_act) {
dq_val = (float)input[m * ldi + n] * a_scale * scale_b[n];
} else {
dq_val = (float)(input[m * ldi + n] - a_zp * comp_b[n]) * a_scale * scale_b[n];
}
if constexpr (accum) {
output[m * ldo + n] += dq_val;
} else {
output[m * ldo + n] = dq_val;
}
}
}
}
#else
template <int64_t N, int64_t ldb>
void _dequant_weight_zp_only(const uint8_t* B, int8_t* dqB, const int8_t* qzeros, int64_t K) {
// B shape = [K, N / 2]
// dqB shape = [K, N]
for (int k = 0; k < K; ++k) {
for (int n = 0; n < N / 2; ++n) {
int32_t b = (int32_t)B[k * ldb + n];
dqB[k * N + n * 2] = (b & 0xf) - qzeros[n];
dqB[k * N + n * 2 + 1] = (b >> 4) - qzeros[n];
}
}
}
#endif
#if defined(CPU_CAPABILITY_AVX512)
inline __m512i combine_m256i(__m256i a, __m256i b) {
__m512i c = _mm512_castsi256_si512(a);
return _mm512_inserti64x4(c, b, 1);
}
inline __m512i combine_m256i(std::array<m256i_wrapper, 2> two_256) {
return combine_m256i(two_256[0].data, two_256[1].data);
}
// negate elements in a according to b's sign
static inline __m512i _mm512_sign_epi8(__m512i a, __m512i b) {
__m512i zero = _mm512_setzero_si512();
__mmask64 blt0 = _mm512_movepi8_mask(b);
return _mm512_mask_sub_epi8(a, blt0, zero, a);
}
template <int64_t M, int64_t N, int64_t ldb, bool sym_quant_act>
void _dequant_gemm_accum_small_M(
float* __restrict__ C,
const uint8_t* A,
const float* scales_a,
const int32_t* qzeros_a,
const uint8_t* B,
const float* scales_b,
const int8_t* qzeros_b,
int64_t K,
int64_t lda,
int64_t ldc) {
// if sym_quant_act is true, A pointer type is passed in as uint8_t* but actually int8_t*.
constexpr int COLS = N / 16;
// Computing compensation is faster than loading it for small M
// because it's memory bound.
__m512i ones = _mm512_set1_epi8(1); // used for computing compensation
__m512i va;
__m512i vb[COLS];
__m512i vc[M * COLS];
__m512 vscales[COLS];
__m512i vzps[COLS];
__m512i vcompensate[COLS];
// Load scales and zps
Unroll<COLS>{}([&](auto i) {
vscales[i] = _mm512_loadu_ps(scales_b + i * 16);
vzps[i] = combine_m256i(load_zps_4vnni(qzeros_b + i * 16));
if constexpr (!sym_quant_act) {
vcompensate[i] = _mm512_setzero_epi32();
}
});
Unroll<M * COLS>{}([&](auto i) { vc[i] = _mm512_setzero_epi32(); });
auto compute = [&](auto i, int k) {
constexpr const int row = i / COLS;
constexpr const int col = i % COLS;
if constexpr (col == 0) {
va = _mm512_set1_epi32(*(int32_t*)(A + row * lda + k));
}
if constexpr (row == 0) {
int B_offset = k * ldb + col * 16 * 2;
vb[col] = combine_m256i(load_uint4_as_int8(B + B_offset));
vb[col] = _mm512_sub_epi8(vb[col], vzps[col]);
if constexpr (!sym_quant_act) {
vcompensate[col] = _mm512_dpbusd_epi32(vcompensate[col], ones, vb[col]);
}
_mm_prefetch(B + B_offset + 128 * ldb, _MM_HINT_T0);
}
if constexpr (sym_quant_act) {
auto vsb = _mm512_sign_epi8(vb[col], va);
auto vabsa = _mm512_sign_epi8(va, va);
vc[i] = _mm512_dpbusds_epi32(vc[i], vabsa, vsb);
} else {
vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
}
};
// Accumulate along k
constexpr const int unroll = 4;
int k = 0;
for (; k < K / 4 / unroll; k++) {
Unroll<unroll>{}([&](auto i) { Unroll<M * COLS>{}(compute, 4 * (k * unroll + i)); });
}
k *= 4 * unroll;
for (; k < K; k += 4) {
Unroll<M * COLS>{}(compute, k);
}
// Store to C
auto store = [&](auto i) {
constexpr const int row = i / COLS;
constexpr const int col = i % COLS;
// compute (qC - compensate * zp_a) * scale_a * scale_b
__m512 vc_float;
if constexpr (!sym_quant_act) {
vc[i] = _mm512_sub_epi32(vc[i], _mm512_mullo_epi32(vcompensate[col], _mm512_set1_epi32(*(qzeros_a + row))));
}
vc_float = _mm512_cvtepi32_ps(vc[i]);
vc_float = _mm512_mul_ps(vc_float, _mm512_set1_ps(*(scales_a + row)));
vc_float = _mm512_mul_ps(vc_float, vscales[col]);
auto vc_old = _mm512_loadu_ps(C + row * ldc + col * 16);
vc_float = _mm512_add_ps(vc_float, vc_old);
_mm512_storeu_ps(C + row * ldc + col * 16, vc_float);
};
Unroll<M * COLS>{}(store);
}
#define CALL_DEQUANT_GEMM_ACCUM_SMALL_M(M) \
_dequant_gemm_accum_small_M<M, N, ldb, sym_quant_act>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
#endif
template <int64_t N, int64_t ldb, bool sym_quant_act>
void _dequant_gemm_accum(
float* C,
const uint8_t* A,
const float* scales_a,
const int32_t* qzeros_a,
const uint8_t* B,
const float* scales_b,
const int8_t* qzeros_b,
const int32_t* compensation,
int8_t* dqB,
int64_t M,
int64_t K,
int64_t lda,
int64_t ldc,
bool use_brgemm) {
// Compute GEMM int8 * int8 -> int32
// dequant result to float by applying scales/qzeros
#if defined(CPU_CAPABILITY_AVX512)
if (!use_brgemm) {
switch (M) {
case 1:
CALL_DEQUANT_GEMM_ACCUM_SMALL_M(1);
break;
case 2:
CALL_DEQUANT_GEMM_ACCUM_SMALL_M(2);
break;
case 3:
CALL_DEQUANT_GEMM_ACCUM_SMALL_M(3);
break;
case 4:
CALL_DEQUANT_GEMM_ACCUM_SMALL_M(4);
break;
default:
TORCH_CHECK(false, "tinygemm_kernel: unexpected M for AVX path!");
}
return;
}
_dequant_weight_zp_only<N, ldb>(B, dqB, qzeros_b, K);
using Tin = typename ActDtype<sym_quant_act>::type;
Tin* A_ptr = (Tin*)A;
if (use_brgemm) {
int32_t C_i32[M * N];
at::native::cpublas::brgemm(
M, N, K, lda, N /*ldb*/, N /*ldc*/, false /* add_C */, A_ptr, dqB, C_i32, true /* is_vnni */);
_mm_prefetch(B + N * K / 2, _MM_HINT_T0);
_mm_prefetch(A + K, _MM_HINT_T0);
_dequant_and_store<true, N, sym_quant_act>(
C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/);
} else
#endif
{
TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!");
}
}
template <int64_t N>
inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
if (bias_ptr) {
for (int i = 0; i < m; ++i) {
int j = 0;
#if defined(CPU_CAPABILITY_AVX512)
#pragma GCC unroll 2
for (; j < N; j += 16) {
__m512 bias_vec = _mm512_loadu_ps(bias_ptr + j);
_mm512_storeu_ps(y_buf + i * N + j, bias_vec);
}
#endif
for (; j < N; ++j) {
y_buf[i * N + j] = bias_ptr[j];
}
}
} else { // initialize to zero
for (int i = 0; i < m; ++i) {
int j = 0;
#if defined(CPU_CAPABILITY_AVX512)
#pragma GCC unroll 2
for (; j < N; j += 16) {
__m512 zero_vec = _mm512_setzero_ps();
_mm512_storeu_ps(y_buf + i * N + j, zero_vec);
}
#endif
for (; j < N; ++j) {
y_buf[i * N + j] = 0;
}
}
}
}
template <typename out_dtype, int64_t N>
inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m, /* int64_t n, */ int64_t lda) {
for (int i = 0; i < m; ++i) {
int j = 0;
if constexpr (std::is_same<out_dtype, float>::value) {
#if defined(CPU_CAPABILITY_AVX512)
#pragma GCC unroll 2
for (; j < N; j += 16) {
__m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
_mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
}
#endif
for (; j < N; ++j) {
c_ptr[i * lda + j] = y_buf[i * N + j];
}
} else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) {
#if defined(CPU_CAPABILITY_AVX512)
#pragma GCC unroll 2
for (; j < N; j += 16) {
__m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
__m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec);
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec);
}
#endif
for (; j < N; ++j) {
c_ptr[i * lda + j] = at::BFloat16(y_buf[i * N + j]);
}
} else if constexpr (std::is_same<out_dtype, at::Half>::value) {
#if defined(CPU_CAPABILITY_AVX512)
#pragma GCC unroll 2
for (; j < N; j += 16) {
__m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
__m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec);
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec);
}
#endif
for (; j < N; ++j) {
c_ptr[i * lda + j] = at::Half(y_buf[i * N + j]);
}
} else {
TORCH_CHECK(false, "Unsupported output dtype");
}
}
}
void fill_val_stub(int32_t* __restrict__ output, int32_t value, int64_t size) {
using iVec = at::vec::Vectorized<int32_t>;
constexpr int VecSize = iVec::size();
const iVec fill_val_vec = iVec(value);
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - VecSize; d += VecSize) {
fill_val_vec.store(output + d);
}
for (; d < size; ++d) {
output[d] = value;
}
}
template <typename act_dtype, typename out_dtype, bool sym_quant_act>
void _da8w4_linear_impl(
act_dtype* __restrict__ input,
const float* __restrict__ input_scales,
const int32_t* __restrict__ input_qzeros,
const uint8_t* __restrict__ weight,
const float* __restrict__ weight_scales,
const int8_t* __restrict__ weight_qzeros,
const float* __restrict__ bias,
out_dtype* __restrict__ output,
float* __restrict__ output_temp,
int8_t* __restrict__ dequant_weight_temp,
int64_t M,
int64_t N,
int64_t K,
int64_t num_groups) {
// weight + compensation shape = [Nc, Kc, BLOCK_N * _block_k / 2 + BLOCK_N*sizeof(int32_t)]
// scales/qzeros shape = [Nc, G, BLOCK_N]
const bool use_brgemm = can_use_brgemm<int8_t>(M);
int64_t block_m = [&]() -> long {
if (M <= 48) {
return M;
} else if (M < 64) {
return 32;
} else if (M < 96) {
return 64;
} else {
return 128;
}
}();
int64_t Mc = div_up(M, block_m);
bool parallel_on_M = M > 128;
int64_t Nc = N / BLOCK_N;
int64_t num_blocks = parallel_on_M ? Mc * Nc : Nc;
int64_t group_size = div_up(K, num_groups);
int64_t _block_k = get_4bit_block_k_size(group_size);
int64_t Kc = K / _block_k;
int64_t block_per_group = group_size / _block_k;
at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) {
int tid = get_thread_num();
float* C_tmp = output_temp + tid * block_m * BLOCK_N;
int8_t* dqB_tmp = dequant_weight_temp + tid * _block_k * BLOCK_N;
for (const auto i : c10::irange(begin, end)) {
int64_t mc = parallel_on_M ? i / Nc : 0;
int64_t nc = parallel_on_M ? i % Nc : i;
int64_t mc_end = parallel_on_M ? mc + 1 : Mc;
for (int mci = mc; mci < mc_end; ++mci) {
int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m;
// copy bias to y_buf if bias is not None
auto bias_data = bias ? bias + nc * BLOCK_N : nullptr;
copy_bias<BLOCK_N>(bias_data, C_tmp, m_size);
for (int kci = 0; kci < Kc; ++kci) {
int32_t* compensation_ptr =
sym_quant_act
? nullptr
: (int32_t*)(void*)(weight + (nc * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
_block_k * BLOCK_N / 2) /*Bcomp*/;
_dequant_gemm_accum<BLOCK_N, BLOCK_N / 2, sym_quant_act>(
/*C*/ C_tmp,
/*A*/ (uint8_t*)input + mci * block_m * K + kci * _block_k,
/*scales_a*/ input_scales + mci * block_m,
/*qzeros_a*/ input_qzeros + mci * block_m,
/*B*/ weight + (nc * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
/*scales_b*/ weight_scales + nc * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N,
/*qzeros_b*/ weight_qzeros + nc * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N,
/*Bcomp*/ compensation_ptr,
/*dqB_tmp*/ dqB_tmp,
/*M*/ m_size,
/*K*/ _block_k,
/*lda*/ K,
/*ldc*/ BLOCK_N,
/*use_brgemm*/ use_brgemm);
}
// store y_buf to output with dtype conversion
store_out<out_dtype, BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/);
}
}
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
} // anonymous namespace
/*
return: packed_weight, packed_scales, packed_qzeros
*/
std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_int4_weight_packed_with_compensation(
const at::Tensor& weight, const at::Tensor& scales, const at::Tensor& qzeros) {
// weight shape = [N, K]
// scales shape = [N, G]
// qzeros shape = [N, G]
TORCH_CHECK(weight.dim() == 2, "DA8W4 CPU: Weight should be a 2D tensor for packing");
TORCH_CHECK(weight.size(1) % 2 == 0, "DA8W4 CPU: Weight should have even number of columns for packing");
auto new_scales = scales;
auto new_qzeros = qzeros;
if (new_scales.dim() == 1) {
new_scales.unsqueeze_(1);
}
new_scales = new_scales.to(at::kFloat);
if (new_qzeros.dim() == 1) {
new_qzeros.unsqueeze_(1);
}
new_qzeros = new_qzeros.to(at::kChar);
int64_t N = weight.size(0);
int64_t K = weight.size(1);
int64_t G = scales.size(1);
int64_t group_size = K / G;
int64_t _block_k = get_4bit_block_k_size(group_size);
constexpr int block_n = block_size_n();
int64_t Nc = N / block_n;
int64_t Kc = K / _block_k;
// Reorder weight to [N/block_n, K/_block_k, _block_k, block_n]
// Reorder scales/qzeros to [N/block_n, G, block_n]
// weight + compensation shape = [Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)]
// scales/qzeros shape = [Nc, G, block_n]
auto weight_view = weight.view({Nc, block_n, Kc, _block_k});
at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous();
at::Tensor blocked_weight;
at::Tensor blocked_scales = new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
at::Tensor blocked_qzeros = new_qzeros.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
// Compensation = Σ(k)(W[k][n] - ZP[n]) for each block.
auto weight_sub_qzero = weight.view({Nc, block_n, G, -1}).to(at::kInt) - new_qzeros.view({Nc, block_n, G, -1});
weight_sub_qzero = weight_sub_qzero.view({Nc, block_n, Kc, _block_k});
at::Tensor compensation = weight_sub_qzero.sum(-1);
compensation = compensation.permute({0, 2, 1}).contiguous().to(at::kInt);
int64_t buffer_size_nbytes = _block_k * block_n / 2 + block_n * sizeof(int32_t);
blocked_weight = at::empty({Nc, Kc, buffer_size_nbytes}, weight.options());
auto weight_ptr = weight_reordered.data_ptr<uint8_t>();
auto compensation_ptr = compensation.data_ptr<int32_t>();
auto blocked_weight_ptr = blocked_weight.data_ptr<uint8_t>();
int64_t num_blocks = Nc * Kc;
at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) {
for (const auto i : c10::irange(begin, end)) {
auto in_ptr = weight_ptr + i * _block_k * block_n;
auto out_ptr = blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t));
int32_t* comp_in_prt = compensation_ptr + i * block_n;
int32_t* comp_out_prt = (int32_t*)(void*)(blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t)) +
_block_k * block_n / 2);
// Reorder weight block to VNNI4 and pack two lanes along N
// N=16 viewed as two lanes: a0, ...a7, b0, ...b7
// pack two lanes: [a0, b0], ..., [a7, b7]
// plain shape = [_block_k, block_n]
// packed shape = [_block_k / 4, block_n / 2, 4] viewed as [_block_k, block_n / 2]
constexpr int n_group_size = 8;
constexpr int vnni_size = 4;
constexpr int n_group = block_n / n_group_size; // 4
for (int nb = 0; nb < n_group; nb += 2) {
for (int k = 0; k < _block_k; k += vnni_size) {
for (int ni = 0; ni < n_group_size; ++ni) {
for (int ki = 0; ki < vnni_size; ++ki) {
int src_idx_1 = nb * n_group_size + ni + (k + ki) * block_n;
int src_idx_2 = (nb + 1) * n_group_size + ni + (k + ki) * block_n;
int dst_idx = (nb / 2 * n_group_size + ni) * vnni_size + k * block_n / 2 + ki;
uint8_t src_1 = *(in_ptr + src_idx_1);
uint8_t src_2 = *(in_ptr + src_idx_2);
uint8_t dst = (src_1 & 0x0f) | ((src_2 & 0x0f) << 4);
*(out_ptr + dst_idx) = dst;
}
}
}
}
// compensation [block_n]
for (int nb = 0; nb < block_n; nb++) {
*(comp_out_prt + nb) = *(comp_in_prt + nb);
}
}
});
return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales), std::move(blocked_qzeros));
}
std::tuple<at::Tensor, at::Tensor> unpack_4bit_to_32bit_signed(const at::Tensor& qweight, const at::Tensor& qzeros) {
TORCH_CHECK(qweight.scalar_type() == at::kInt, "qweight must be int32");
TORCH_CHECK(qzeros.scalar_type() == at::kInt, "qzeros must be int32");
const auto W0 = qweight.size(0);
const auto W1 = qweight.size(1);
const auto Z0 = qzeros.size(0);
const auto Z1 = qzeros.size(1);
// unpacked_weights: (W0 * 8, W1), int8
auto unpacked_weights = at::zeros({W0 * 8, W1}, at::TensorOptions().dtype(at::kChar));
// unpacked_zeros: (Z0, Z1 * 8), int8
auto unpacked_zeros = at::zeros({Z0, Z1 * 8}, at::TensorOptions().dtype(at::kChar));
const int32_t* qw_ptr = qweight.data_ptr<int32_t>();
const int32_t* qz_ptr = qzeros.data_ptr<int32_t>();
int8_t* uw_ptr = unpacked_weights.data_ptr<int8_t>();
int8_t* uz_ptr = unpacked_zeros.data_ptr<int8_t>();
// ---- unpack qweight ----
for (int64_t row = 0; row < W0 * 8; ++row) {
const int i = row & 7; // row % 8
const int src_row = row >> 3; // row // 8
const int shift = 4 * i;
for (int64_t col = 0; col < W1; ++col) {
int32_t v = qw_ptr[src_row * W1 + col];
uw_ptr[row * W1 + col] = static_cast<int8_t>((v >> shift) & 0xF);
}
}
// ---- unpack qzeros ----
for (int64_t col = 0; col < Z1 * 8; ++col) {
const int i = col & 7;
const int src_col = col >> 3;
const int shift = 4 * i;
for (int64_t row = 0; row < Z0; ++row) {
int32_t v = qz_ptr[row * Z1 + src_col];
uz_ptr[row * (Z1 * 8) + col] = static_cast<int8_t>((v >> shift) & 0xF);
}
}
return std::make_tuple(unpacked_weights, unpacked_zeros + 1);
}
std::tuple<at::Tensor, at::Tensor>
autogptq_to_int4pack(const at::Tensor& qweight_tensor, const at::Tensor& qzeros_tensor) {
TORCH_CHECK(qweight_tensor.scalar_type() == at::kInt, "qweight_tensor must be int32");
TORCH_CHECK(qzeros_tensor.scalar_type() == at::kInt, "qzeros_tensor must be int32");
TORCH_CHECK(qweight_tensor.is_cpu(), "CPU only implementation");
if (qweight_tensor.dim() == 3) {
const int64_t B = qweight_tensor.size(0);
std::vector<at::Tensor> qweight_list;
std::vector<at::Tensor> qzeros_list;
qweight_list.reserve(B);
qzeros_list.reserve(B);
for (int64_t i = 0; i < B; ++i) {
auto outputs = unpack_4bit_to_32bit_signed(qweight_tensor[i], qzeros_tensor[i]);
at::Tensor unpacked_qweight = std::get<0>(outputs);
at::Tensor unpacked_qzeros = std::get<1>(outputs);
qweight_list.push_back(unpacked_qweight.transpose(0, 1).contiguous().to(at::kByte));
qzeros_list.push_back(unpacked_qzeros.contiguous().to(at::kByte));
}
return std::make_tuple(at::stack(qweight_list).detach(), at::stack(qzeros_list).detach());
}
auto outputs = unpack_4bit_to_32bit_signed(qweight_tensor, qzeros_tensor);
at::Tensor unpacked_qweight = std::get<0>(outputs);
at::Tensor unpacked_qzeros = std::get<1>(outputs);
at::Tensor return_qweight = unpacked_qweight.transpose(0, 1).contiguous().to(at::kByte);
at::Tensor return_qzeros = unpacked_qzeros.contiguous().to(at::kByte);
return std::make_tuple(return_qweight, return_qzeros);
}
std::tuple<at::Tensor, at::Tensor> int4pack(at::Tensor qweight, at::Tensor qzeros, int64_t quant_method_4bit) {
if (quant_method_4bit == CPUQuantAlgo::AWQ) {
// autoawq unpacking
qweight = qweight.contiguous();
qzeros = qzeros.contiguous();
// bitshifts: [0, 4, 1, 5, 2, 6, 3, 7] * 4
auto bitshifts = at::tensor({0, 4, 1, 5, 2, 6, 3, 7}, at::kInt) * 4;
auto qweight_unsq = qweight.unsqueeze(-1); // [..., K, N/8, 1]
auto unpacked = (at::bitwise_right_shift(qweight_unsq, bitshifts) & 0xF).contiguous();
auto qweight_final = unpacked.flatten(-2).transpose(-1, -2).to(at::kByte).clone();
auto qzeros_unsq = qzeros.unsqueeze(-1);
auto qzeros_unpacked = (at::bitwise_right_shift(qzeros_unsq, bitshifts) & 0xF).contiguous();
auto qzeros_final = qzeros_unpacked.flatten(-2).to(at::kByte).clone();
return std::make_tuple(qweight_final, qzeros_final);
} else if (quant_method_4bit == CPUQuantAlgo::GPTQ) {
// autogptq unpacking
auto outputs = autogptq_to_int4pack(qweight, qzeros);
at::Tensor unpacked_qweight = std::get<0>(outputs);
at::Tensor unpacked_qzeros = std::get<1>(outputs);
return std::make_tuple(unpacked_qweight, unpacked_qzeros);
} else {
TORCH_CHECK(false, "CPU int4 pack only support AWQ or GPTQ...");
}
}
std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32
at::Tensor qzeros, // awq: (*, K / group_size, N / 8) || gptq: (*, K / group_size, N / 8) , int32
at::Tensor scales, // awq: (*, K / group_size, N) || gptq: (*, K / group_size, N) , bfloat16
int64_t quant_method_4bit) {
at::Tensor _qweight;
at::Tensor _qzeros;
auto res = int4pack(qweight, qzeros, quant_method_4bit);
_qweight = std::get<0>(res);
_qzeros = std::get<1>(res);
auto _scales = scales;
_qzeros = _qzeros.transpose(-2, -1).contiguous(); // .T
_scales = _scales.transpose(-2, -1).contiguous();
if (_qweight.dim() == 3) { // Dim=3 for MOE packing, TODO: refine a unified loop
int64_t E = _qweight.size(0);
int64_t K = _qweight.size(2);
int64_t G = _scales.size(2);
int64_t group_size = K / G;
int64_t _block_k = get_4bit_block_k_size(group_size);
int64_t block_n = block_size_n();
int64_t Nc = _qweight.size(1) / block_n;
int64_t Kc = K / _block_k;
int64_t buffer_size_nbytes = _block_k * block_n / 2 + block_n * sizeof(int32_t);
auto blocked_weight = at::empty({E, Nc, Kc, buffer_size_nbytes}, _qweight.options());
auto blocked_scales = at::empty({E, Nc, G, block_n}, _scales.options()).to(at::kFloat);
auto blocked_qzeros = at::empty({E, Nc, G, block_n}, _qzeros.options()).to(at::kChar);
for (int i = 0; i < _qweight.size(0); i++) {
auto res_ = convert_int4_weight_packed_with_compensation(_qweight[i], _scales[i], _qzeros[i]);
blocked_weight[i] = std::get<0>(res_);
blocked_scales[i] = std::get<1>(res_);
blocked_qzeros[i] = std::get<2>(res_);
}
_qweight = blocked_weight;
_scales = blocked_scales;
_qzeros = blocked_qzeros;
} else {
auto res_ = convert_int4_weight_packed_with_compensation(_qweight, _scales, _qzeros);
_qweight = std::get<0>(res_);
_scales = std::get<1>(res_);
_qzeros = std::get<2>(res_);
}
return std::make_tuple(_qweight, _qzeros, _scales);
}
at::Tensor int4_scaled_mm_cpu_with_quant(
const at::Tensor& input,
const at::Tensor& weight,
const at::Tensor& weight_scales,
const at::Tensor& weight_qzeros,
const std::optional<at::Tensor>& bias,
at::ScalarType output_dtype) {
int64_t M_a = input.size(0);
int64_t K_a = input.size(1);
int64_t lda = input.stride(0);
const auto st = input.scalar_type();
TORCH_CHECK(
st == at::kBFloat16 || st == at::kHalf, "int4_scaled_mm_cpu_with_quant: expect A to be bfloat16 or half.");
constexpr bool sym_quant_act = false; // TODO: add sym quant path
using Tin = typename ActDtype<sym_quant_act>::type;
int64_t act_buffer_size = /* act quant */ M_a * K_a +
/* act scale */ M_a * sizeof(float) +
/* act zp */ M_a * sizeof(int32_t);
auto act_buffer = at::empty({act_buffer_size}, input.options().dtype(at::kByte));
// asym path, activation quants into uint8_t
auto Aq_data = act_buffer.data_ptr<uint8_t>();
auto As_data = reinterpret_cast<float*>(Aq_data + M_a * K_a);
auto Azp_data = reinterpret_cast<int32_t*>(As_data + M_a);
fill_val_stub(Azp_data, 128, M_a); // sym_a s8s8 is unified to u8s8 with compensation (128)
auto out_sizes = input.sizes().vec();
int64_t N = weight_scales.size(0) * weight_scales.size(-1);
out_sizes.back() = N;
auto output = at::empty(out_sizes, input.options());
// weight + compensation shape = [Nc, Kc, BLOCK_N * _block_k / 2 + BLOCK_N*sizeof(int32_t)]
// scales/qzeros shape = [Nc, G, BLOCK_N]
int64_t Nc = weight.size(0);
int64_t Kc = weight.size(1);
int64_t _block_k = K_a / Kc;
TORCH_CHECK(N == Nc * BLOCK_N, "DA8W4: weight and input shapes mismatch");
// scales/qzeros shape = [Nc, G, BLOCK_N]
int64_t num_groups = weight_scales.size(1);
const uint8_t* b_ptr = weight.data_ptr<uint8_t>();
const float* b_scales_ptr = weight_scales.data_ptr<float>();
const int8_t* b_qzeros_ptr = weight_qzeros.data_ptr<int8_t>();
const float* bias_ptr = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
int num_threads = at::get_num_threads();
int64_t temp_buffer_size = /* output temp */ num_threads * BLOCK_M * BLOCK_N * sizeof(float) +
/* weight dequant temp */ num_threads * _block_k * BLOCK_N;
auto c_temp_buffer = at::empty({temp_buffer_size}, input.options().dtype(at::kChar));
float* c_temp_ptr = (float*)((void*)(c_temp_buffer.data_ptr<int8_t>()));
int8_t* dqB_temp_ptr = (int8_t*)((void*)(c_temp_ptr + num_threads * BLOCK_M * BLOCK_N));
#define LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act) \
AT_DISPATCH_FLOATING_TYPES_AND2( \
at::ScalarType::BFloat16, at::ScalarType::Half, output_dtype, "int4_scaled_mm_cpu_with_quant", [&] { \
const scalar_t* __restrict__ A_data = input.data_ptr<scalar_t>(); \
scalar_t* __restrict__ c_ptr = output.data_ptr<scalar_t>(); \
at::parallel_for(0, M_a, 0, [&](int64_t begin, int64_t end) { \
for (int64_t m = begin; m < end; ++m) { \
quantize_row_int8<scalar_t>(Aq_data + m * K_a, As_data[m], A_data + m * lda, K_a); \
} \
}); \
_da8w4_linear_impl<Tin, scalar_t, sym_quant_act>( \
Aq_data, \
As_data, \
Azp_data, \
b_ptr, \
b_scales_ptr, \
b_qzeros_ptr, \
bias_ptr, \
c_ptr, \
c_temp_ptr, \
dqB_temp_ptr, \
M_a, \
N, \
K_a, \
num_groups); \
});
LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act);
return output;
}
template <typename scalar_t>
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
using Vec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += Vec::size()) {
fVec x0 = fVec::loadu(input + d);
fVec x1 = fVec::loadu(input + d + fVec::size());
Vec res = convert_from_float_ext<scalar_t>(x0, x1);
res.store(out + d);
}
}
template <typename scalar_t>
void tinygemm_kernel(
scalar_t* C,
float* C_temp,
const uint8_t* A,
const float* scales_a,
const int32_t* qzeros_a,
const uint8_t* B,
const float* scales_b,
const int8_t* qzeros_b,
const int32_t* compensation,
int8_t* dqB_tmp,
int64_t M,
int64_t K,
int64_t lda,
int64_t ldc_f,
int64_t ldc_s,
bool store_out,
bool use_brgemm) {
// TODO: add sym quant act, now only asym
_dequant_gemm_accum<BLOCK_N, BLOCK_N / 2, false>(
C_temp, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, dqB_tmp, M, K, lda, ldc_f, use_brgemm);
if (store_out) {
// copy from Ctmp to C
for (int64_t m = 0; m < M; ++m) {
copy_stub<scalar_t>(C + m * ldc_s, C_temp + m * ldc_f, BLOCK_N);
}
}
}
#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
template void tinygemm_kernel<TYPE>( \
TYPE * C, \
float* C_temp, \
const uint8_t* A, \
const float* scales_a, \
const int32_t* qzeros_a, \
const uint8_t* B, \
const float* scales_b, \
const int8_t* qzeros_b, \
const int32_t* compensation, \
int8_t* dqB_tmp, \
int64_t M, \
int64_t K, \
int64_t lda, \
int64_t ldc_f, \
int64_t ldc_s, \
bool store_out, \
bool use_brgemm)
INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
// int4 gemm dispatch api register
at::Tensor int4_scaled_mm_cpu(
at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias) {
return int4_scaled_mm_cpu_with_quant(x, w, w_scales, w_zeros, bias, x.scalar_type());
}
python/sglang/kernels/aot/csrc/cpu/gemm_int8.cpp
+541
View File
@@ -0,0 +1,541 @@
#include "common.h"
#include "gemm.h"
#include "vec.h"
namespace {
template <typename scalar_t, bool has_bias, int BLOCK_N>
struct scale_C {
static inline void apply(
scalar_t* __restrict__ C,
const int32_t* __restrict__ Ctmp,
const int32_t* __restrict__ Bcomp,
const float* __restrict__ bias,
float As,
const float* __restrict__ Bs) {
TORCH_CHECK(false, "scale_C: scalar path not implemented!");
}
};
#if defined(CPU_CAPABILITY_AVX512)
template <bool has_bias, int BLOCK_N>
struct scale_C<at::BFloat16, has_bias, BLOCK_N> {
static inline void apply(
at::BFloat16* __restrict__ C,
const int32_t* __restrict__ Ctmp,
const int32_t* __restrict__ Bcomp,
const float* __restrict__ bias,
float As,
const float* __restrict__ Bs) {
constexpr int COLS = BLOCK_N / 16;
static_assert(COLS % 2 == 0);
__m512 vc[COLS];
__m512 vd0 = _mm512_set1_ps(As);
auto compute = [&](auto col) {
__m512 vd1 = _mm512_loadu_ps(Bs + col * 16);
__m512i vcomp = _mm512_loadu_si512(Bcomp + col * 16);
__m512i vc32 = _mm512_loadu_si512(Ctmp + col * 16);
vc[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32, vcomp));
if constexpr (has_bias) {
__m512 vbias = _mm512_loadu_ps(bias + col * 16);
vc[col] = _mm512_fmadd_ps(_mm512_mul_ps(vc[col], vd0), vd1, vbias);
} else {
vc[col] = _mm512_mul_ps(_mm512_mul_ps(vc[col], vd0), vd1);
}
};
Unroll<COLS>{}(compute);
auto storec = [&](auto col) {
// for COLS = 2, 4 use 512bit store
if constexpr (col % 2 == 0) {
_mm512_storeu_si512(
reinterpret_cast<__m512i*>((C + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc[col + 1], vc[col + 0])));
}
};
Unroll<COLS>{}(storec);
}
};
#endif
template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn {
static inline void apply(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B,
scalar_t* __restrict__ C,
const float* __restrict__ As,
const float* __restrict__ Bs,
const int32_t* __restrict__ Bcomp,
const float* __restrict__ bias,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
}
};
#if defined(CPU_CAPABILITY_AVX512)
template <bool has_bias, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
static inline void apply(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B,
at::BFloat16* __restrict__ C,
const float* __restrict__ As,
const float* __restrict__ Bs,
const int32_t* __restrict__ Bcomp,
const float* __restrict__ bias,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
constexpr int ROWS = BLOCK_M;
constexpr int COLS = BLOCK_N / 16;
static_assert(COLS % 2 == 0);
// prefetch distance
constexpr int PREFETCH_SIZE_K = 0;
__m512i va;
__m512i vb[COLS];
__m512i vc[ROWS * COLS];
__m512i vcomp[COLS];
__m512 vd0;
__m512 vd1[COLS];
// oops! 4x4 spills but we use 4x2
__m512 vbias[COLS];
// [NOTE]: s8s8 igemm compensation in avx512-vnni
//
// avx512-vnni has no s8s8, so we need to change s8s8 to u8s8 with compensate:
//
// a * b = (a + 128) * b - 128 * b
// s s u s u s
//
// 1) 128 * b is pre-computed when packing B to vnni formats
// 2) a + 128 is fused when dynamically quantize A
//
auto loadc = [&](auto i) { vc[i] = _mm512_set1_epi32(0); };
Unroll<ROWS * COLS>{}(loadc);
const int64_t K4 = K >> 2;
const int64_t lda4 = lda >> 2;
const int64_t ldb4 = ldb; // ldb * 4 >> 2;
const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
auto compute = [&](auto i, int64_t k) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
if constexpr (col == 0) {
va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
}
if constexpr (row == 0) {
vb[col] = _mm512_loadu_si512(b_ptr + k * ldb4 + col * 16);
if constexpr (PREFETCH_SIZE_K > 0) {
_mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb4 + col * 16, _MM_HINT_T0);
}
}
vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
};
for (int64_t k = 0; k < K4; ++k) {
Unroll<ROWS * COLS>{}(compute, k);
}
auto storec = [&](auto i) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
// load a scale
if constexpr (col == 0) {
vd0 = _mm512_set1_ps(As[row]);
}
// load b scale and vcomp per 2 vectors
// also load bias if any
if constexpr (row == 0) {
if constexpr (col % 2 == 0) {
vd1[col + 0] = _mm512_loadu_ps(Bs + col * 16);
vd1[col + 1] = _mm512_loadu_ps(Bs + col * 16 + 16);
vcomp[col + 0] = _mm512_loadu_si512(Bcomp + col * 16);
vcomp[col + 1] = _mm512_loadu_si512(Bcomp + col * 16 + 16);
if constexpr (has_bias) {
vbias[col + 0] = _mm512_loadu_ps(bias + col * 16);
vbias[col + 1] = _mm512_loadu_ps(bias + col * 16 + 16);
}
}
}
// for COLS = 2, 4 use 512bit store
if constexpr (col % 2 == 0) {
__m512 vc0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[row * COLS + col + 0], vcomp[col + 0]));
__m512 vc1 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[row * COLS + col + 1], vcomp[col + 1]));
if constexpr (has_bias) {
vc0 = _mm512_fmadd_ps(_mm512_mul_ps(vc0, vd0), vd1[col + 0], vbias[col + 0]);
vc1 = _mm512_fmadd_ps(_mm512_mul_ps(vc1, vd0), vd1[col + 1], vbias[col + 1]);
} else {
vc0 = _mm512_mul_ps(_mm512_mul_ps(vc0, vd0), vd1[col + 0]);
vc1 = _mm512_mul_ps(_mm512_mul_ps(vc1, vd0), vd1[col + 1]);
}
_mm512_storeu_si512(
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
}
};
Unroll<ROWS * COLS>{}(storec);
}
};
#endif
#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
A + mb_start * lda, \
B + nb_start * 4, \
C + mb_start * ldc + nb_start, \
As + mb_start, \
Bs + nb_start, \
Bcomp + nb_start, \
has_bias ? bias + nb_start : nullptr, \
K, \
lda, \
ldb, \
ldc);
template <typename scalar_t, bool has_bias>
void tinygemm_kernel(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B,
scalar_t* __restrict__ C,
int32_t* __restrict__ Ctmp,
const float* __restrict__ As,
const float* __restrict__ Bs,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg) {
// B compensation
const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
if (brg) {
constexpr int BLOCK_N = block_size_n();
at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
// apply compensation and scale
for (int64_t m = 0; m < M; ++m) {
scale_C<scalar_t, has_bias, BLOCK_N>::apply(C + m * ldc, Ctmp + m * BLOCK_N, Bcomp, bias, As[m], Bs);
}
return;
}
// pattern: 1-4-16
constexpr int64_t BLOCK_M = 4;
constexpr int64_t BLOCK_N = 64;
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
for (int64_t mb = 0; mb < MB; ++mb) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
for (int64_t nb = 0; nb < NB; ++nb) {
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
switch (mb_size << 4 | nb_size >> 4) {
// mb_size = 1
case 0x12:
LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
break;
case 0x14:
LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
break;
// mb_size = 2
case 0x22:
LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
break;
case 0x24:
LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
break;
// mb_size = 3
case 0x32:
LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
break;
case 0x34:
LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
break;
// mb_size = 4
case 0x42:
LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
break;
case 0x44:
LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
}
}
}
}
template <typename scalar_t>
void int8_scaled_mm_kernel_impl(
scalar_t* __restrict__ out,
const uint8_t* __restrict__ mat1,
const int8_t* __restrict__ mat2,
const float* __restrict__ scales1,
const float* __restrict__ scales2,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
const bool use_brgemm = can_use_brgemm<int8_t>(M);
// K + 4 after compensation
const int64_t packed_row_size = get_row_size<int8_t>(K);
AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// for brgemm, use int32_t for accumulate
alignas(64) int32_t Ctmp[BLOCK_M * BLOCK_N];
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int mb_start = mb * BLOCK_M;
int mb_size = std::min(M - mb_start, BLOCK_M);
int nb_start = nb * BLOCK_N;
int nb_size = std::min(N - nb_start, BLOCK_N);
tinygemm_kernel<scalar_t, has_bias>(
/* A */ mat1 + mb_start * K,
/* B */ mat2 + nb_start * packed_row_size /* nb * BLOCK_N * (K + 4) */,
/* C */ out + mb_start * N + nb_start,
/* Ctmp*/ Ctmp,
/* As */ scales1 + mb_start,
/* Bs */ scales2 + nb_start,
/* bias*/ bias + nb_start,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ K,
/* ldb */ nb_size,
/* ldc */ N,
/* brg */ use_brgemm);
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
});
}
} // anonymous namespace
// tinygemm interface
template <typename scalar_t>
void tinygemm_kernel(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B,
scalar_t* __restrict__ C,
int32_t* __restrict__ Ctmp,
const float* __restrict__ As,
const float* __restrict__ Bs,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg) {
tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, As, Bs, nullptr, M, N, K, lda, ldb, ldc, brg);
}
#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
template void tinygemm_kernel<TYPE>( \
const uint8_t* __restrict__ A, \
const int8_t* __restrict__ B, \
TYPE* __restrict__ C, \
int32_t* __restrict__ Ctmp, \
const float* __restrict__ As, \
const float* __restrict__ Bs, \
int64_t M, \
int64_t N, \
int64_t K, \
int64_t lda, \
int64_t ldb, \
int64_t ldc, \
bool brg)
INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(A);
CHECK_DIM(2, A);
int64_t M = A.size(0);
int64_t K = A.size(1);
int64_t lda = A.stride(0);
const auto st = A.scalar_type();
TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "per_token_quant_int8: expect A to be bfloat16 or half.");
auto Aq = at::empty({M, K}, A.options().dtype(at::kByte));
auto As = at::empty({M}, A.options().dtype(at::kFloat));
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "per_token_quant_int8", [&] {
uint8_t* __restrict__ Aq_data = Aq.data_ptr<uint8_t>();
float* __restrict__ As_data = As.data_ptr<float>();
const scalar_t* __restrict__ A_data = A.data_ptr<scalar_t>();
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_data + m * K, As_data[m], A_data + m * lda, K);
}
});
});
return std::make_tuple(Aq, As);
}
// weight : static, per-channel, symmetric
// activation : dynamic, per-token, symmetric
//
// mat1 : [M, K]
// mat2 : [N, K]
// scales1 : [M]
// scales2 : [N]
// bias : [N]
// out : [M, N]
//
at::Tensor int8_scaled_mm_cpu(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales1,
at::Tensor& scales2,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni) {
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
CHECK_INPUT(mat1);
CHECK_INPUT(mat2);
CHECK_INPUT(scales1);
CHECK_INPUT(scales2);
CHECK_DIM(2, mat1);
CHECK_DIM(2, mat2);
int64_t M = mat1.size(0);
int64_t N = mat2.size(0);
int64_t K = mat1.size(1);
// see [NOTE]: s8s8 igemm compensation in avx512-vnni
CHECK_EQ(mat2.size(1), (int64_t)(is_vnni ? K + sizeof(int32_t) : K));
CHECK_EQ(scales1.numel(), M);
CHECK_EQ(scales2.numel(), N);
TORCH_CHECK(mat1.scalar_type() == at::kByte, "int8_scaled_mm: expect mat1 to be uint8.");
TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm: expect mat2 to be int8.");
TORCH_CHECK(
scales1.scalar_type() == at::kFloat && scales2.scalar_type() == at::kFloat,
"int8_scaled_mm: expect scales to be float32.");
auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
const bool has_bias = bias.has_value();
const float* bias_data = nullptr;
if (has_bias) {
CHECK_EQ(bias.value().size(0), N);
bias_data = bias.value().data_ptr<float>();
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_kernel_impl", [&] {
int8_scaled_mm_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
mat1.data_ptr<uint8_t>(),
packed_w.data_ptr<int8_t>(),
scales1.data_ptr<float>(),
scales2.data_ptr<float>(),
bias_data,
M,
N,
K);
});
return out;
}
// fused `per_token_quant_int8_cpu` and `int8_scaled_mm_cpu`
at::Tensor int8_scaled_mm_with_quant(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales2,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni) {
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
CHECK_INPUT(mat2);
CHECK_INPUT(scales2);
CHECK_DIM(2, mat1);
CHECK_DIM(2, mat2);
int64_t M = mat1.size(0);
int64_t N = mat2.size(0);
int64_t K = mat1.size(1);
int64_t lda = mat1.stride(0);
// see [NOTE]: s8s8 igemm compensation in avx512-vnni
CHECK_EQ(mat2.size(1), (int64_t)(is_vnni ? K + sizeof(int32_t) : K));
CHECK_EQ(scales2.numel(), N);
const auto st = mat1.scalar_type();
TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "int8_scaled_mm_with_quant: expect A to be bfloat16 or half.");
TORCH_CHECK(st == out_dtype, "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm_with_quant: expect mat2 to be int8.");
TORCH_CHECK(scales2.scalar_type() == at::kFloat, "int8_scaled_mm_with_quant: expect scales to be float32.");
const int64_t buffer_size = M * K + M * sizeof(float);
auto buffer = at::empty({buffer_size}, mat1.options().dtype(at::kByte));
auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
const bool has_bias = bias.has_value();
const float* bias_data = nullptr;
if (has_bias) {
CHECK_EQ(bias.value().size(0), N);
bias_data = bias.value().data_ptr<float>();
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_with_quant_kernel_impl", [&] {
uint8_t* __restrict__ Aq_data = buffer.data_ptr<uint8_t>();
float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K));
const scalar_t* __restrict__ A_data = mat1.data_ptr<scalar_t>();
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_data + m * K, As_data[m], A_data + m * lda, K);
}
});
int8_scaled_mm_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
Aq_data,
packed_w.data_ptr<int8_t>(),
As_data,
scales2.data_ptr<float>(),
bias_data,
M,
N,
K);
});
return out;
}
python/sglang/kernels/aot/csrc/cpu/interface.cpp
+92
View File
@@ -0,0 +1,92 @@
#include <ATen/record_function.h>
#include <torch/all.h>
#include "shm.h"
// Communication settings
static int world_rank = -1;
static int world_size = -1;
static bool is_initialized = false;
static bool all_ranks_local_p = false;
void initialize(int64_t size, int64_t rank) {
if (is_initialized) {
return;
}
// Check whether all ranks is on the same physical machine.
// If true, we will use an SHM based low latency allreduce
auto ls_string = std::getenv("LOCAL_SIZE");
int ls = 0;
if (ls_string != NULL) {
ls = std::stoi(std::getenv("LOCAL_SIZE"));
}
if (size >= 1 && size == ls) {
all_ranks_local_p = true;
}
world_size = size;
world_rank = rank;
is_initialized = true;
const char* addr_string = std::getenv("MASTER_ADDR");
if (addr_string == NULL) {
addr_string = "";
}
const char* port_string = std::getenv("MASTER_PORT");
if (port_string == NULL) {
port_string = "";
}
if (all_ranks_local_p) {
shm_initialize(size, rank, addr_string, port_string);
}
}
void shm_allreduce(torch::Tensor& data, int64_t op) {
TORCH_CHECK(op == c10d::ReduceOp::SUM, "Only torch.distributed.ReduceOp.SUM is supported");
auto numel = data.numel();
int data_size = numel * data.element_size();
all_reduce_outer_loop(data, numel, data_size);
return;
}
torch::Tensor shm_allgather(torch::Tensor& data, int64_t dim) {
auto numel = data.numel();
int data_size = numel * data.element_size();
if (dim < 0) {
dim += data.dim();
}
std::vector<int64_t> result_shape = data.sizes().vec();
result_shape[dim] *= world_size;
torch::Tensor result_tensor = torch::empty(result_shape, data.options());
return all_gather<STATE_GROUP_ALL_GATHER>(result_tensor, data, dim, numel, data_size);
}
void shm_allgather_into_tensor(torch::Tensor& output_tensor, torch::Tensor& data) {
RECORD_FUNCTION("sgl-kernel::shm_allgather_into_tensor", std::vector<c10::IValue>({data}));
auto numel = data.numel();
int data_size = numel * data.element_size();
int64_t dim = 0;
all_gather<STATE_GROUP_ALL_GATHER_INTO_TENSOR>(output_tensor, data, dim, numel, data_size);
}
void shm_reduce_scatter_tensor(at::Tensor& output_tensor, at::Tensor& data, int64_t op) {
RECORD_FUNCTION("sgl-kernel::shm_reduce_scatter_tensor", std::vector<c10::IValue>({data}));
TORCH_CHECK(op == c10d::ReduceOp::SUM, "Only torch.distributed.ReduceOp.SUM is supported");
auto numel = data.numel();
int data_size = numel * data.element_size();
reduce_scatter_outer_loop(output_tensor, data, numel, data_size);
return;
}
python/sglang/kernels/aot/csrc/cpu/kvcache.cpp
+184
View File
@@ -0,0 +1,184 @@
#include "common.h"
#include "vec.h"
namespace {
template <typename scalar_t>
inline void copy_stub(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int size) {
int d = 0;
#if defined(CPU_CAPABILITY_AVX512)
using Vec = at::vec::Vectorized<scalar_t>;
constexpr int kVecSize = Vec::size();
for (; d <= size - kVecSize; d += kVecSize) {
Vec data = Vec::loadu(src + d);
data.store(dst + d);
}
#endif
for (; d < size; ++d) {
dst[d] = src[d];
}
}
template <typename scalar_t, typename index_t>
void store_cache_kernel_impl(
const scalar_t* __restrict__ k,
const scalar_t* __restrict__ v,
scalar_t* __restrict__ k_cache,
scalar_t* __restrict__ v_cache,
const index_t* __restrict__ indices,
int64_t batch_size,
int64_t num_pages,
int64_t row_dim,
int64_t k_stride,
int64_t v_stride,
int64_t kc_stride,
int64_t vc_stride) {
at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
for (int64_t bs = begin; bs < end; ++bs) {
const int64_t idx = static_cast<int64_t>(indices[bs]);
const scalar_t* k_ptr = k + bs * k_stride;
const scalar_t* v_ptr = v + bs * v_stride;
scalar_t* kc_ptr = k_cache + idx * kc_stride;
scalar_t* vc_ptr = v_cache + idx * vc_stride;
copy_stub(kc_ptr, k_ptr, row_dim);
copy_stub(vc_ptr, v_ptr, row_dim);
}
});
}
} // anonymous namespace
// check tensor last two dimensions are contiguous
#define CHECK_LAST2_DIM_CONTIGUOUS(x, ndim) \
do { \
const auto& _x = (x); \
const auto _ndim = _x.dim(); \
const auto _strides = _x.strides(); \
const auto _sizes = _x.sizes(); \
TORCH_CHECK(_ndim == ndim, #x " must have " #ndim " dimensions"); \
TORCH_CHECK( \
_ndim >= 2 && _strides[_ndim - 1] == 1 && _strides[_ndim - 2] == _sizes[_ndim - 1], \
#x " must be contiguous at the last two dimensions"); \
} while (0)
// [NB]: store_cache takes 3 dimension tensors,
// This is to avoid the overhead of creating a new TensorImpl
// from .view(-1, row_dim)
//
// k : [batch_size, num_heads, head_size] -> [batch_size, row_dim]
// v : [batch_size, num_heads, head_size] -> [batch_size, row_dim]
// k_cache : [num_pages, num_heads, head_size] -> [num_pages, row_dim]
// v_cache : [num_pages, num_heads, head_size] -> [num_pages, row_dim]
// indices : [batch_size]
//
void store_cache_cpu(
const at::Tensor& k,
const at::Tensor& v,
const at::Tensor& k_cache,
const at::Tensor& v_cache,
const at::Tensor& indices,
std::optional<int64_t> row_dim) {
CHECK_LAST2_DIM_CONTIGUOUS(k, 3);
CHECK_LAST2_DIM_CONTIGUOUS(v, 3);
CHECK_LAST2_DIM_CONTIGUOUS(k_cache, 3);
CHECK_LAST2_DIM_CONTIGUOUS(v_cache, 3);
CHECK_INPUT(indices);
int64_t batch_size = k.size(0);
int64_t num_heads = k.size(1);
int64_t head_size = k.size(2);
int64_t num_pages = k_cache.size(0);
int64_t row_dim_value = num_heads * head_size;
if (row_dim.has_value()) {
CHECK_EQ(row_dim.value(), row_dim_value);
}
CHECK_EQ(indices.size(0), batch_size);
// strides: batch dimension (dim 0) stride in elements
int64_t k_stride = k.stride(0);
int64_t v_stride = v.stride(0);
int64_t kc_stride = k_cache.stride(0);
int64_t vc_stride = v_cache.stride(0);
const auto dtype = k.scalar_type();
TORCH_CHECK(
dtype == v.scalar_type() && dtype == k_cache.scalar_type() && dtype == v_cache.scalar_type(),
"store_cache_cpu: input tensors must have the same dtype");
const auto index_dtype = indices.scalar_type();
TORCH_CHECK(index_dtype == at::kLong || index_dtype == at::kInt, "indices must be int64 or int32");
// dtype : [bfloat16, float16, uint8] for fp8 KV stored as uint8
// index_dtype : [int64, int32]
AT_DISPATCH_REDUCED_FLOATING_TYPES_AND(at::ScalarType::Byte, dtype, "store_cache_cpu", [&] {
AT_DISPATCH_INDEX_TYPES(index_dtype, "store_cache_cpu_index", [&] {
store_cache_kernel_impl<scalar_t, index_t>(
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
k_cache.data_ptr<scalar_t>(),
v_cache.data_ptr<scalar_t>(),
indices.data_ptr<index_t>(),
batch_size,
num_pages,
row_dim_value,
k_stride,
v_stride,
kc_stride,
vc_stride);
});
});
}
// CPU counterpart of the Triton kernel `copy_all_layer_kv_cache_tiled`:
// for every K/V buffer b, copy the slot rows `src_loc` to `tgt_loc`:
// buf_b[tgt_loc[i], :] = buf_b[src_loc[i], :] for i in [0, num_locs)
//
// data_ptrs : [2 * layer_num] uint64; base address of each K/V buffer
// strides : [2 * layer_num] int64; bytes per slot row of each buffer
// tgt_loc : [num_locs] int64/int32 slot indices
// src_loc : [num_locs] int64/int32 slot indices
//
// Like the Triton kernel, the copy is safe when tgt_loc and src_loc overlap
// arbitrarily: all source rows of a buffer are staged before any target row
// of that buffer is written (gather then scatter).
void copy_all_layer_kv_cache_cpu(
const at::Tensor& data_ptrs, const at::Tensor& strides, const at::Tensor& tgt_loc, const at::Tensor& src_loc) {
CHECK_INPUT(data_ptrs);
CHECK_INPUT(strides);
CHECK_INPUT(tgt_loc);
CHECK_INPUT(src_loc);
CHECK_EQ(data_ptrs.scalar_type(), at::kUInt64);
CHECK_EQ(strides.scalar_type(), at::kLong);
CHECK_EQ(tgt_loc.scalar_type(), src_loc.scalar_type());
int64_t num_bufs = data_ptrs.numel();
CHECK_EQ(strides.numel(), num_bufs);
int64_t num_locs = tgt_loc.numel();
CHECK_EQ(src_loc.numel(), num_locs);
if (num_bufs == 0 || num_locs == 0) {
return;
}
const uint64_t* __restrict__ ptrs = reinterpret_cast<const uint64_t*>(data_ptrs.data_ptr());
const int64_t* __restrict__ stride_ptr = strides.data_ptr<int64_t>();
AT_DISPATCH_INDEX_TYPES(tgt_loc.scalar_type(), "copy_all_layer_kv_cache_cpu", [&] {
const index_t* __restrict__ tgt_ptr = tgt_loc.data_ptr<index_t>();
const index_t* __restrict__ src_ptr = src_loc.data_ptr<index_t>();
at::parallel_for(0, num_bufs, 0, [&](int64_t begin, int64_t end) {
std::vector<uint8_t> staging;
for (int64_t b = begin; b < end; ++b) {
uint8_t* base = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(ptrs[b]));
const int64_t stride = stride_ptr[b];
staging.resize(num_locs * stride);
for (int64_t i = 0; i < num_locs; ++i) {
std::memcpy(staging.data() + i * stride, base + src_ptr[i] * stride, stride);
}
for (int64_t i = 0; i < num_locs; ++i) {
std::memcpy(base + tgt_ptr[i] * stride, staging.data() + i * stride, stride);
}
}
});
});
}
python/sglang/kernels/aot/csrc/cpu/mamba/conv.cpp
+703
View File
@@ -0,0 +1,703 @@
#include "common.h"
#include "gemm.h"
#include "vec.h"
namespace {
template <typename scalar_t>
inline void copy_stub(scalar_t* __restrict__ y, const scalar_t* __restrict__ x, int64_t size) {
using Vec = at::vec::Vectorized<scalar_t>;
const bool is_padding = (x == nullptr);
for (int64_t d = 0; d < size; d += Vec::size()) {
Vec data_vec = is_padding ? Vec(0.f) : Vec::loadu(x + d);
data_vec.store(y + d);
}
}
// no remainder
template <typename scalar_t>
void inline update_conv_state(
scalar_t* __restrict__ conv_states,
const scalar_t* __restrict__ input,
int64_t width,
int64_t dim,
int64_t seqlen,
bool has_initial_states) {
// width for `conv_states`
int64_t width1 = width - 1;
int64_t w = 0;
for (; w < width1 - seqlen; ++w) {
scalar_t* y = conv_states + w * dim;
const scalar_t* x = has_initial_states ? conv_states + (w + seqlen) * dim : nullptr;
copy_stub(y, x, dim);
}
for (; w < width1; ++w) {
scalar_t* y = conv_states + w * dim;
const scalar_t* x = input + (w + seqlen - width1) * dim;
copy_stub(y, x, dim);
}
}
// A : [M, BLOCK_N]
// B : [BLOCK_N, K], prepacked as [K/2, BLOCK_N, 2]
// C : [M, BLOCK_N]
// bias : [BLOCK_N]
//
// lda : leading dimension of `input` and `out`
//
template <typename scalar_t, int K, int BLOCK_N, bool has_bias, bool has_silu>
struct tinygemm_kernel {
static inline void apply(
const scalar_t* __restrict__ A,
const scalar_t* __restrict__ B,
scalar_t* __restrict__ C,
const scalar_t* __restrict__ bias,
const scalar_t* __restrict__ conv_states,
bool has_initial_state,
int64_t M,
int64_t lda,
bool is_first_token) {
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
}
};
#if defined(CPU_CAPABILITY_AVX512)
template <int K, int BLOCK_N, bool has_bias, bool has_silu>
struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
static inline void apply(
const at::BFloat16* __restrict__ A,
const at::BFloat16* __restrict__ B,
at::BFloat16* __restrict__ C,
const at::BFloat16* __restrict__ bias,
const at::BFloat16* __restrict__ conv_states,
bool has_initial_state,
int64_t M,
int64_t lda,
bool is_first_token) {
assert(K == 4);
constexpr int ROWS = K;
constexpr int COLS = BLOCK_N / block_size_n();
// leading dimension size for b for next block [K/2, 32, 2]
constexpr int ldb = block_size_n() * K;
__m512bh va[ROWS * COLS];
__m512bh vb[ROWS * COLS];
__m512 vc[COLS * 2];
// k: {-3, -2, -1} -> {0, 1, 2}
auto set_conv_states = [&](int k, int col) -> __m512i {
return has_initial_state ? _mm512_loadu_si512(conv_states + (k + K - 1) * lda + col * 32)
: _mm512_setzero_si512();
};
#define MM512_LOAD_A(idx) \
((idx) < 0 && is_first_token) ? (__m512bh)(set_conv_states((idx), col)) \
: (__m512bh)(_mm512_loadu_si512(A + (idx) * lda + col * 32))
#define MM512_PACK_A(ap, bp, a, b) \
do { \
__m512i r0 = (__m512i)(a); \
__m512i r1 = (__m512i)(b); \
__m512i d0 = _mm512_unpacklo_epi16(r0, r1); \
__m512i d1 = _mm512_unpackhi_epi16(r0, r1); \
r0 = _mm512_shuffle_i32x4(d0, d1, 0x88); \
r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd); \
(ap) = (__m512bh)_mm512_shuffle_i32x4(r0, r1, 0x88); \
(bp) = (__m512bh)_mm512_shuffle_i32x4(r0, r1, 0xdd); \
} while (0)
// step 0 : preload a at time step [-3][-2][-1]
auto preloada = [&](auto i) {
constexpr int col = i;
int64_t m = 0;
va[1 * COLS + col] = MM512_LOAD_A(m - 3);
va[2 * COLS + col] = MM512_LOAD_A(m - 2);
va[3 * COLS + col] = MM512_LOAD_A(m - 1);
};
Unroll<COLS>{}(preloada);
auto loada = [&](auto i, int64_t m) {
constexpr int col = i;
// update previous time step
va[0 * COLS + col] = va[1 * COLS + col];
va[1 * COLS + col] = va[2 * COLS + col];
va[2 * COLS + col] = va[3 * COLS + col];
// load current time step
va[3 * COLS + col] = MM512_LOAD_A(m);
};
// step 1 : load weight for just once
auto loadb = [&](auto i) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
vb[row * COLS + col] = (__m512bh)(_mm512_loadu_si512(B + col * ldb + row * 32));
};
Unroll<ROWS * COLS>{}(loadb);
// [NB] accumulates 4x32 bfloat16 blocks
//
// +------------+------------+
// | col0 | col1 |
// +------------+------------+
// | va0 va1 | va0 va1 |
// | va2 va3 | va2 va3 |
// +------------+------------+
// | vc0 vc1 | vc0 vc1 |
// +------------+------------+
//
// * va and vb shares the same memory layout
// * block_n 32 with 4 rows equals to 4 registers
// * 37 uops with avx512bf16 v.s. 57 uops with avx512f
//
auto compute = [&](auto i) {
constexpr int col = i;
// init accumulators
if constexpr (has_bias) {
__m512i b16 = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(bias + col * 32));
vc[col * 2 + 0] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
vc[col * 2 + 1] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
} else {
vc[col * 2 + 0] = _mm512_set1_ps(0.f);
vc[col * 2 + 1] = _mm512_set1_ps(0.f);
}
// convert to vnni2 format
__m512bh va0, va1, va2, va3;
MM512_PACK_A(va0, va1, va[0 * COLS + col], va[1 * COLS + col]);
MM512_PACK_A(va2, va3, va[2 * COLS + col], va[3 * COLS + col]);
// accumulate
vc[col * 2 + 0] = _mm512_dpbf16_ps(vc[col * 2 + 0], va0, vb[0 * COLS + col]);
vc[col * 2 + 0] = _mm512_dpbf16_ps(vc[col * 2 + 0], va2, vb[2 * COLS + col]);
vc[col * 2 + 1] = _mm512_dpbf16_ps(vc[col * 2 + 1], va1, vb[1 * COLS + col]);
vc[col * 2 + 1] = _mm512_dpbf16_ps(vc[col * 2 + 1], va3, vb[3 * COLS + col]);
};
using fVec = at::vec::Vectorized<float>;
using bVec = at::vec::Vectorized<at::BFloat16>;
auto storec = [&](auto i, int64_t m) {
constexpr int col = i;
fVec x0 = fVec(vc[col * 2 + 0]);
fVec x1 = fVec(vc[col * 2 + 1]);
if constexpr (has_silu) {
x0 = fast_silu(x0);
x1 = fast_silu(x1);
}
bVec out_vec = convert_from_float_ext<at::BFloat16>(x0, x1);
out_vec.store(C + m * lda + col * 32);
};
for (int64_t m = 0; m < M; ++m) {
// step 3.a : load a at current time step
Unroll<COLS>{}(loada, m);
// step 3.b : accumulate for window size (4)
Unroll<COLS>{}(compute);
// step 3.c : store c at current time step
Unroll<COLS>{}(storec, m);
}
}
};
#endif
#define LAUNCH_TINYGEMM_KERNEL(K, NB_SIZE) \
tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
input + bs * seqlen * dim + mb_start * dim + nb_start, \
weight + nb_start * width, \
out + bs * seqlen * dim + mb_start * dim + nb_start, \
has_bias ? bias + nb_start : nullptr, \
has_conv_states ? conv_states + conv_state_index * (K - 1) * dim + nb_start : nullptr, \
has_initial_states_value, \
mb_size, \
dim, \
mb_start == 0);
template <typename scalar_t>
void causal_conv1d_fwd_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ weight,
const scalar_t* __restrict__ bias,
scalar_t* __restrict__ conv_states,
const int32_t* __restrict__ conv_indices,
const bool* __restrict__ has_initial_state,
bool silu_activation,
int64_t batch,
int64_t dim,
int64_t seqlen,
int64_t width,
int64_t num_seq_blocks) {
// handle 32 x 64 per block
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n() * 2;
const int64_t NB = div_up(dim, BLOCK_N);
const int64_t num_blocks_per_seq = div_up(seqlen, BLOCK_M);
const bool has_conv_states = conv_states != nullptr;
const bool has_conv_indices = conv_indices != nullptr;
// parallel on [batch, seq, NB]
AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
at::parallel_for(0, num_seq_blocks * NB, 0, [&](int64_t begin, int64_t end) {
int64_t mb{0}, nb{0};
data_index_init(begin, mb, num_seq_blocks, nb, NB);
for (int64_t i = begin; i < end; ++i) {
int64_t bs = mb / num_blocks_per_seq;
int64_t mb_start = (mb % num_blocks_per_seq) * BLOCK_M;
int64_t mb_size = std::min(seqlen - mb_start, BLOCK_M);
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false;
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
switch (width << 4 | nb_size >> 4) {
case 0x42:
LAUNCH_TINYGEMM_KERNEL(4, 32);
break;
case 0x44:
LAUNCH_TINYGEMM_KERNEL(4, 64);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
}
// move to the next index
data_index_step(mb, num_seq_blocks, nb, NB);
}
});
});
// update conv_states if necessary
if (has_conv_states) {
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t bs = begin; bs < end; ++bs) {
update_conv_state(
conv_states + bs * (width - 1) * dim, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);
}
});
}
}
#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \
tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
input + batch_offset * dim + mb_start * dim + nb_start, \
weight + nb_start * width, \
out + batch_offset * dim + mb_start * dim + nb_start, \
has_bias ? bias + nb_start : nullptr, \
nullptr, \
false, \
mb_size, \
dim, \
mb_start == 0);
// TODO: add `has_initial_state` support for varlen kernel
template <typename scalar_t>
void causal_conv1d_fwd_varlen_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ weight,
const scalar_t* __restrict__ bias,
scalar_t* __restrict__ conv_states,
const int32_t* __restrict__ query_start_loc,
const int32_t* __restrict__ conv_indices,
const bool* __restrict__ has_initial_state,
const int32_t* __restrict__ block_indices,
bool silu_activation,
int64_t batch,
int64_t dim,
int64_t width,
int64_t num_seq_blocks) {
// handle 32 x 64 per block
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n() * 2;
const int64_t NB = div_up(dim, BLOCK_N);
const bool has_conv_states = conv_states != nullptr;
const bool has_conv_indices = conv_indices != nullptr;
// parallel on [batch, seq, NB]
AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
at::parallel_for(0, num_seq_blocks * NB, 0, [&](int64_t begin, int64_t end) {
int64_t mb{0}, nb{0};
data_index_init(begin, mb, num_seq_blocks, nb, NB);
for (int64_t i = begin; i < end; ++i) {
int32_t bs = block_indices[mb * 2 + 0];
int32_t batch_offset = query_start_loc[bs];
int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
int64_t mb_start = block_indices[mb * 2 + 1] * BLOCK_M;
int64_t mb_size = std::min(seqlen - mb_start, BLOCK_M);
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
switch (width << 4 | nb_size >> 4) {
case 0x42:
LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32);
break;
case 0x44:
LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 64);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
}
// move to the next index
data_index_step(mb, num_seq_blocks, nb, NB);
}
});
});
// update conv_states if necessary
if (has_conv_states) {
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t bs = begin; bs < end; ++bs) {
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
int32_t batch_offset = query_start_loc[bs];
update_conv_state(
conv_states + conv_state_index * (width - 1) * dim,
input + batch_offset * dim,
width,
dim,
seqlen,
/* has_initial_state */ false);
}
});
}
}
template <typename scalar_t>
void causal_conv1d_update_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
scalar_t* __restrict__ conv_states,
const scalar_t* __restrict__ weight,
const scalar_t* __restrict__ bias,
const int32_t* __restrict__ conv_indices,
bool silu_activation,
int64_t batch,
int64_t dim,
int64_t seqlen,
int64_t width) {
// handle 32 x 64 per block
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n() * 2;
const int64_t NB = div_up(dim, BLOCK_N);
const bool has_conv_states = conv_states != nullptr;
const bool has_conv_indices = conv_indices != nullptr;
// parallel on [batch, NB]
AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) {
int64_t bs{0}, nb{0};
data_index_init(begin, bs, batch, nb, NB);
for (int64_t i = begin; i < end; ++i) {
int64_t mb_start = 0;
int64_t mb_size = 1;
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
const bool has_initial_states_value = true;
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
switch (width << 4 | nb_size >> 4) {
case 0x42:
LAUNCH_TINYGEMM_KERNEL(4, 32);
break;
case 0x44:
LAUNCH_TINYGEMM_KERNEL(4, 64);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
}
// move to the next index
data_index_step(bs, batch, nb, NB);
}
});
});
#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*(width - 1) * dim + (w) * dim
// update conv_states
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t bs = begin; bs < end; ++bs) {
// update old states, range [1, width - 1)
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
for (int64_t w = 1; w < width - 1; ++w) {
std::memcpy(CONV_STATE_INDEXR(w - 1), CONV_STATE_INDEXR(w), dim * sizeof(scalar_t));
}
// copy new states
std::memcpy(CONV_STATE_INDEXR(width - 2), input + bs * dim, dim * sizeof(scalar_t));
}
});
}
} // anonymous namespace
// from [dim, width] or [N, K]
// to [N/BLOCK_N, K/2, BLOCK_N, 2]
at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight) {
CHECK_INPUT(weight);
int64_t dim = weight.size(0);
int64_t width = weight.size(1);
constexpr int64_t BLOCK_N = block_size_n();
TORCH_CHECK(width == 4, "causal_conv1d_weight_pack: support only width of 4");
TORCH_CHECK(dim % BLOCK_N == 0, "causal_conv1d_weight_pack: invalid dim size ", dim);
const int64_t N = dim, K2 = width >> 1;
const int64_t NB = div_up(N, BLOCK_N);
auto packed_weight = at::empty_like(weight);
AT_DISPATCH_REDUCED_FLOATING_TYPES(weight.scalar_type(), "causal_conv1d_fwd_kernel_impl", [&] {
// cast to float32 as vnni size is 2
const float* w_data = reinterpret_cast<float*>(weight.data_ptr<scalar_t>());
float* packed_data = reinterpret_cast<float*>(packed_weight.data_ptr<scalar_t>());
at::parallel_for(0, NB * K2 * BLOCK_N, 0, [&](int64_t begin, int64_t end) {
int64_t nb{0}, k2{0}, n{0};
data_index_init(begin, nb, NB, k2, K2, n, BLOCK_N);
// TODO: optimize this if we need to online prepacking.
for (int64_t i = begin; i < end; ++i) {
packed_data[i] = w_data[nb * BLOCK_N * K2 + n * K2 + k2];
// move to the next index
data_index_step(nb, NB, k2, K2, n, BLOCK_N);
}
});
});
return packed_weight;
}
#define CHECK_OPTIONAL_SHAPE_DTYPE(OPT, SIZE, DTYPE) \
if (OPT.has_value()) { \
const auto tensor = OPT.value(); \
CHECK_CONTIGUOUS(tensor); \
CHECK_EQ(tensor.size(0), SIZE); \
CHECK_EQ(tensor.scalar_type(), DTYPE); \
}
template <int BLOCK_M>
int64_t get_block_count(const std::optional<at::Tensor>& offsets, int64_t batch, int64_t seqlen) {
if (offsets.has_value()) {
const int32_t* offsets_data = offsets.value().data_ptr<int32_t>();
int32_t num_seq_blocks = 0;
for (int64_t row = 0; row < batch; ++row) {
num_seq_blocks += div_up(offsets_data[row + 1] - offsets_data[row], BLOCK_M);
}
return num_seq_blocks;
}
return batch * div_up(seqlen, int64_t(BLOCK_M));
}
template <int BLOCK_M>
at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t num_seq_blocks) {
if (!offsets.has_value()) {
return at::Tensor();
}
const at::Tensor& offsets_ = offsets.value();
at::Tensor indices = at::empty({num_seq_blocks, 2}, offsets_.options());
int64_t batch = offsets_.size(0) - 1;
const int32_t* offsets_data = offsets_.data_ptr<int32_t>();
int32_t* indices_data = indices.data_ptr<int32_t>();
int64_t idx = 0;
for (int32_t row = 0; row < batch; ++row) {
int32_t blocks = div_up(offsets_data[row + 1] - offsets_data[row], BLOCK_M);
for (int32_t col = 0; col < blocks; ++col) {
indices_data[idx * 2 + 0] = row;
indices_data[idx * 2 + 1] = col;
idx++;
}
}
return indices;
}
// API aligned with GPUs
//
// x: (batch, dim, seqlen) or (dim, cu_seq_len) for varlen
// weight: (dim, width)
// bias: (dim,)
// query_start_loc: (batch + 1) int32
// cache_indices: (batch) int32
// has_initial_state: (batch) bool
// conv_states: (..., dim, width - 1) itype
// activation: either None or "silu" or "swish"
// pad_slot_id: int
//
at::Tensor causal_conv1d_fwd_cpu(
const at::Tensor& x,
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
const std::optional<at::Tensor>& conv_states,
const std::optional<at::Tensor>& query_start_loc,
const std::optional<at::Tensor>& conv_state_indices,
const std::optional<at::Tensor>& has_initial_state,
bool silu_activation,
int64_t pad_slot_id,
bool is_vnni) {
CHECK_CONTIGUOUS(weight);
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
const bool is_var_seqlen = query_start_loc.has_value();
const int64_t input_ndim = is_var_seqlen ? 2 : 3;
TORCH_CHECK(x.dim() == input_ndim, "causal_conv1d_fwd_cpu: expect x to be ", input_ndim, "D tensor.");
TORCH_CHECK(x.stride(-2) == 1 && x.stride(-1) == x.size(-2), "causal_conv1d_fwd_cpu: expect x to be transposed.");
const int64_t batch = is_var_seqlen ? query_start_loc.value().size(0) - 1 : x.size(0);
const int64_t dim = x.size(-2);
const int64_t seqlen = x.size(-1);
const int64_t width = weight.size(-1);
const auto scalar_type = x.scalar_type();
CHECK_EQ(weight.scalar_type(), scalar_type);
CHECK_OPTIONAL_SHAPE_DTYPE(bias, dim, scalar_type);
CHECK_OPTIONAL_SHAPE_DTYPE(query_start_loc, batch + 1, at::kInt);
CHECK_OPTIONAL_SHAPE_DTYPE(conv_state_indices, batch, at::kInt);
CHECK_OPTIONAL_SHAPE_DTYPE(has_initial_state, batch, at::kBool);
if (conv_states.has_value()) {
auto& conv_states_val = conv_states.value();
int64_t padded_batch = conv_states_val.size(0);
CHECK_EQ(conv_states_val.scalar_type(), scalar_type);
CHECK_GE(padded_batch, batch);
CHECK_EQ(conv_states_val.size(1), dim);
CHECK_EQ(conv_states_val.size(2), width - 1);
// adjust `conv_states` to be contiguous on `dim`
// should happen only once
if (conv_states_val.stride(-2) != 1) {
auto conv_states_copy = conv_states_val.clone();
conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim});
conv_states_val.copy_(conv_states_copy);
}
}
// block size for sequence blocks, 32
constexpr int64_t BLOCK_M = block_size_m();
// total number of sequence blocks
int64_t num_seq_blocks = get_block_count<BLOCK_M>(query_start_loc, batch, seqlen);
at::Tensor out = at::empty_like(x);
AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_fwd_kernel_impl", [&] {
if (is_var_seqlen) {
// record seq blocks in Coordinate format, aka [num_seq_blocks, 2]
at::Tensor block_indices = get_block_indices<BLOCK_M>(query_start_loc, num_seq_blocks);
causal_conv1d_fwd_varlen_kernel_impl(
out.data_ptr<scalar_t>(),
x.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
conditional_data_ptr<scalar_t>(bias),
conditional_data_ptr<scalar_t>(conv_states),
conditional_data_ptr<int32_t>(query_start_loc),
conditional_data_ptr<int32_t>(conv_state_indices),
conditional_data_ptr<bool>(has_initial_state),
block_indices.data_ptr<int32_t>(),
silu_activation,
batch,
dim,
width,
num_seq_blocks);
} else {
causal_conv1d_fwd_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
x.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
conditional_data_ptr<scalar_t>(bias),
conditional_data_ptr<scalar_t>(conv_states),
conditional_data_ptr<int32_t>(conv_state_indices),
conditional_data_ptr<bool>(has_initial_state),
silu_activation,
batch,
dim,
seqlen,
width,
num_seq_blocks);
}
});
return out;
}
// API aligned with GPUs
//
// x: (batch, dim) or (batch, dim, seqlen)
// conv_state: (..., dim, state_len), where state_len >= width - 1
// weight: (dim, width)
// bias: (dim,)
// cache_seqlens: (batch,), dtype int32.
// conv_state_indices: (batch,), dtype int32
// pad_slot_id: int
// out: (batch, dim) or (batch, dim, seqlen)
//
at::Tensor causal_conv1d_update_cpu(
const at::Tensor& x,
const at::Tensor& conv_states,
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
bool silu_activation,
const std::optional<at::Tensor>& cache_seqlens,
const std::optional<at::Tensor>& conv_state_indices,
int64_t pad_slot_id,
bool is_vnni) {
CHECK_CONTIGUOUS(x);
CHECK_CONTIGUOUS(weight);
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
// TODO: add multi-token prediction support
TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor.");
TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens.");
int64_t batch = x.size(0);
int64_t dim = x.size(1);
int64_t seqlen = 1;
int64_t width = weight.size(-1);
const auto scalar_type = x.scalar_type();
CHECK_EQ(weight.scalar_type(), scalar_type);
CHECK_OPTIONAL_SHAPE_DTYPE(bias, dim, scalar_type);
CHECK_OPTIONAL_SHAPE_DTYPE(conv_state_indices, batch, at::kInt);
CHECK_EQ(conv_states.scalar_type(), scalar_type);
CHECK_EQ(conv_states.size(1), dim);
CHECK_EQ(conv_states.size(2), width - 1);
// adjust `conv_states` to be contiguous on `dim`
if (conv_states.stride(-2) != 1) {
int64_t num_cache_lines = conv_states.size(0);
auto conv_states_copy = conv_states.clone();
conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim});
conv_states.copy_(conv_states_copy);
}
at::Tensor out = at::empty_like(x);
AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_update_kernel_impl", [&] {
causal_conv1d_update_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
x.data_ptr<scalar_t>(),
conv_states.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
conditional_data_ptr<scalar_t>(bias),
conditional_data_ptr<int32_t>(conv_state_indices),
silu_activation,
batch,
dim,
seqlen,
width);
});
return out;
}
python/sglang/kernels/aot/csrc/cpu/mamba/fla.cpp
+1730
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/aot/csrc/cpu/model/qwen3.cpp
+287
View File
@@ -0,0 +1,287 @@
#include "common.h"
#include "gemm.h"
#include "vec.h"
namespace {
template <typename scalar_t>
inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ src, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
constexpr int kVecSize = bVec::size();
int64_t d = 0;
#pragma GCC unroll 4
for (; d <= size - kVecSize; d += kVecSize) {
bVec out_bvec = bVec::loadu(src + d);
out_bvec.store(out + d);
}
for (; d < size; ++d) {
out[d] = src[d];
}
}
template <typename scalar_t>
void fused_qkvzba_split_reshape_cat_impl(
const scalar_t* __restrict__ mixed_qkvz,
const scalar_t* __restrict__ mixed_ba,
scalar_t* __restrict__ mixed_qkv,
scalar_t* __restrict__ z,
scalar_t* __restrict__ b,
scalar_t* __restrict__ a,
int64_t batch,
int64_t num_heads_qk,
int64_t num_heads_v,
int64_t head_qk,
int64_t group,
int64_t head_v,
int64_t qkv_strideB,
int64_t qkvz_strideB,
int64_t ba_strideB) {
int64_t qkvz_stride_per_head = head_qk * 2 + head_v * 2 * group;
at::parallel_for(0, batch * num_heads_qk, 0, [&](int64_t begin, int64_t end) {
int64_t bi{0}, hi{0};
data_index_init(begin, bi, batch, hi, num_heads_qk);
for (int64_t i = begin; i < end; ++i) {
scalar_t* __restrict__ q_out_ptr = mixed_qkv + bi * qkv_strideB + hi * head_qk;
const scalar_t* __restrict__ q_in_ptr = mixed_qkvz + bi * qkvz_strideB + hi * qkvz_stride_per_head;
scalar_t* __restrict__ k_out_ptr = q_out_ptr + num_heads_qk * head_qk;
const scalar_t* __restrict__ k_in_ptr = q_in_ptr + head_qk;
scalar_t* __restrict__ v_out_ptr = k_out_ptr + num_heads_qk * head_qk + hi * head_qk * (group - 1);
const scalar_t* __restrict__ v_in_ptr = k_in_ptr + head_qk;
scalar_t* __restrict__ z_out_ptr = z + bi * num_heads_v * head_v + hi * group * head_v;
const scalar_t* __restrict__ z_in_ptr = v_in_ptr + head_qk * group;
copy_stub(q_out_ptr, q_in_ptr, head_qk);
copy_stub(k_out_ptr, k_in_ptr, head_qk);
copy_stub(v_out_ptr, v_in_ptr, head_qk * group);
copy_stub(z_out_ptr, z_in_ptr, head_qk * group);
scalar_t* __restrict__ b_out_ptr = b + bi * num_heads_v + hi * group;
const scalar_t* __restrict__ b_in_ptr = mixed_ba + bi * ba_strideB + hi * group * 2;
scalar_t* __restrict__ a_out_ptr = a + bi * num_heads_v + hi * group;
const scalar_t* __restrict__ a_in_ptr = b_in_ptr + group;
copy_stub(b_out_ptr, b_in_ptr, group);
copy_stub(a_out_ptr, a_in_ptr, group);
data_index_step(bi, batch, hi, num_heads_qk);
}
});
}
template <typename scalar_t>
void fused_qkvzba_split_reshape_cat_contiguous_impl(
const scalar_t* __restrict__ mixed_qkvz,
const scalar_t* __restrict__ mixed_ba,
scalar_t* __restrict__ mixed_qkv,
scalar_t* __restrict__ z,
scalar_t* __restrict__ b,
scalar_t* __restrict__ a,
int64_t batch,
int64_t v_tp,
int64_t num_heads_v,
int64_t qkv_dim,
int64_t qkv_strideB,
int64_t qkvz_strideB,
int64_t ba_strideB) {
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t bi = begin; bi < end; ++bi) {
scalar_t* __restrict__ qkv_out_ptr = mixed_qkv + bi * qkv_strideB;
const scalar_t* __restrict__ qkv_in_ptr = mixed_qkvz + bi * qkvz_strideB;
scalar_t* __restrict__ z_out_ptr = z + bi * v_tp;
const scalar_t* __restrict__ z_in_ptr = qkv_in_ptr + qkv_dim;
copy_stub(qkv_out_ptr, qkv_in_ptr, qkv_dim);
copy_stub(z_out_ptr, z_in_ptr, v_tp);
scalar_t* __restrict__ b_out_ptr = b + bi * num_heads_v;
const scalar_t* __restrict__ b_in_ptr = mixed_ba + bi * ba_strideB;
scalar_t* __restrict__ a_out_ptr = a + bi * num_heads_v;
const scalar_t* __restrict__ a_in_ptr = b_in_ptr + num_heads_v;
copy_stub(b_out_ptr, b_in_ptr, num_heads_v);
copy_stub(a_out_ptr, a_in_ptr, num_heads_v);
}
});
}
template <typename scalar_t>
void fused_input_proj_kernel_impl(
scalar_t* __restrict__ out,
scalar_t* __restrict__ out2,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ weight,
const scalar_t* __restrict__ weight2,
int64_t M,
int64_t N,
int64_t N2,
int64_t K) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N + N2, BLOCK_N);
const bool use_brgemm = can_use_brgemm<scalar_t>(M);
// parallel on [MB, NB]
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// for brgemm, use float32 for accumulate
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(M - mb_start, BLOCK_M);
int64_t nb_start = nb * BLOCK_N;
const bool is_first = nb_start < N;
int64_t local_nb_start = is_first ? nb_start : nb_start - N;
int64_t nb_size = std::min((is_first ? N : N2) - local_nb_start, BLOCK_N);
scalar_t* __restrict__ curr_out = is_first ? out : out2;
const scalar_t* __restrict__ curr_weight = is_first ? weight : weight2;
int64_t local_out_strideM = is_first ? N : N2;
tinygemm_kernel<scalar_t>(
/* A */ input + mb_start * K,
/* B */ curr_weight + local_nb_start * K,
/* C */ curr_out + mb_start * local_out_strideM + local_nb_start,
/* Ctmp*/ Ctmp,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ K,
/* ldb */ nb_size,
/* ldc */ local_out_strideM,
/* brg */ use_brgemm);
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
} // anonymous namespace
// mixed_qkvz: [batch, num_heads_qk * head_qk * 2 + num_heads_v * head_v * 2]
// mixed_ba: [batch, num_heads_v * 2]
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_reshape_cat_cpu(
const at::Tensor& mixed_qkvz,
const at::Tensor& mixed_ba,
int64_t num_heads_qk,
int64_t num_heads_v,
int64_t head_qk,
int64_t head_v) {
int64_t batch = mixed_qkvz.size(0);
int64_t qkv_dim = num_heads_qk * head_qk * 2 + num_heads_v * head_v;
int64_t ba_dim = num_heads_v * 2;
int64_t expected_dim = qkv_dim + num_heads_v * head_v;
CHECK_INPUT_SHAPE_DTYPE<false>(mixed_qkvz, {batch, expected_dim}, mixed_qkvz.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<false>(mixed_ba, {batch, ba_dim}, mixed_qkvz.scalar_type());
CHECK_EQ(num_heads_v % num_heads_qk, 0);
at::Tensor mixed_qkv = at::empty({batch, qkv_dim}, mixed_qkvz.options());
at::Tensor z = at::empty({batch, num_heads_v, head_v}, mixed_qkvz.options());
at::Tensor b = at::empty({batch, num_heads_v}, mixed_ba.options());
at::Tensor a = at::empty({batch, num_heads_v}, mixed_ba.options());
int64_t group = num_heads_v / num_heads_qk;
int64_t qkvz_strideB = mixed_qkvz.size(1);
int64_t qkv_strideB = mixed_qkv.size(1);
int64_t ba_strideB = mixed_ba.size(1);
AT_DISPATCH_REDUCED_FLOATING_TYPES(mixed_qkvz.scalar_type(), "fused_qkvzba_split_reshape_cat_impl", [&] {
fused_qkvzba_split_reshape_cat_impl<scalar_t>(
mixed_qkvz.data_ptr<scalar_t>(),
mixed_ba.data_ptr<scalar_t>(),
mixed_qkv.data_ptr<scalar_t>(),
z.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
a.data_ptr<scalar_t>(),
batch,
num_heads_qk,
num_heads_v,
head_qk,
group,
head_v,
qkv_strideB,
qkvz_strideB,
ba_strideB);
});
return std::make_tuple(mixed_qkv, z, b, a);
}
// mixed_qkvz: [batch, num_heads_qk * head_qk * 2 + num_heads_v * head_v * 2]
// mixed_ba: [batch, num_heads_v * 2]
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_reshape_cat_contiguous_cpu(
const at::Tensor& mixed_qkvz,
const at::Tensor& mixed_ba,
int64_t num_heads_qk,
int64_t num_heads_v,
int64_t head_qk,
int64_t head_v) {
int64_t batch = mixed_qkvz.size(0);
int64_t k_tp = num_heads_qk * head_qk;
int64_t v_tp = num_heads_v * head_v;
int64_t qkv_dim = k_tp * 2 + v_tp;
int64_t ba_dim = num_heads_v * 2;
int64_t expected_dim = qkv_dim + v_tp;
CHECK_INPUT_SHAPE_DTYPE<false>(mixed_qkvz, {batch, expected_dim}, mixed_qkvz.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<false>(mixed_ba, {batch, ba_dim}, mixed_qkvz.scalar_type());
at::Tensor mixed_qkv = at::empty({batch, qkv_dim}, mixed_qkvz.options());
at::Tensor z = at::empty({batch, num_heads_v, head_v}, mixed_qkvz.options());
at::Tensor b = at::empty({batch, num_heads_v}, mixed_ba.options());
at::Tensor a = at::empty({batch, num_heads_v}, mixed_ba.options());
int64_t qkvz_strideB = mixed_qkvz.size(1);
int64_t qkv_strideB = mixed_qkv.size(1);
int64_t ba_strideB = mixed_ba.size(1);
AT_DISPATCH_REDUCED_FLOATING_TYPES(mixed_qkvz.scalar_type(), "fused_qkvzba_split_reshape_cat_contiguous_impl", [&] {
fused_qkvzba_split_reshape_cat_contiguous_impl<scalar_t>(
mixed_qkvz.data_ptr<scalar_t>(),
mixed_ba.data_ptr<scalar_t>(),
mixed_qkv.data_ptr<scalar_t>(),
z.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
a.data_ptr<scalar_t>(),
batch,
v_tp,
num_heads_v,
qkv_dim,
qkv_strideB,
qkvz_strideB,
ba_strideB);
});
return std::make_tuple(mixed_qkv, z, b, a);
}
// [projected_states_qkvz |projected_states_ba]
// = hidden_states @ [qkvz_weight.T | ba_weight.T]
//
// hidden_states : [batch, hidden_size]
// qkvz_weight : [qkvz_dim, hidden_size]
// ba_weight : [ba_dim, hidden_size]
// projected_states_qkvz : [batch, qkvz_dim]
// projected_states_ba : [batch, ba_dim]
//
std::tuple<at::Tensor, at::Tensor>
fused_input_proj_cpu(at::Tensor& hidden_states, at::Tensor& qkvz_weight, at::Tensor& ba_weight, bool is_vnni) {
const auto st = hidden_states.scalar_type();
TORCH_CHECK(st == at::ScalarType::BFloat16, "fused_input_proj_cpu only supports BFloat16");
int64_t batch = hidden_states.size(0);
int64_t hidden_size = hidden_states.size(1);
int64_t qkvz_dim = qkvz_weight.size(0);
int64_t ba_dim = ba_weight.size(0);
CHECK_INPUT(hidden_states);
CHECK_INPUT_SHAPE_DTYPE<false>(qkvz_weight, {qkvz_dim, hidden_size}, st);
CHECK_INPUT_SHAPE_DTYPE<false>(ba_weight, {ba_dim, hidden_size}, st);
TORCH_CHECK(qkvz_dim % block_size_n() == 0, "qkvz_weight out features must be divisible by ", block_size_n());
TORCH_CHECK(ba_dim % block_size_n() == 0, "ba_weight out features must be divisible by ", block_size_n());
TORCH_CHECK(hidden_size % TILE_K == 0, "hidden_size must be divisible by ", TILE_K);
// weight prepacking if necessary
at::Tensor packed_w = is_vnni ? qkvz_weight : convert_weight_packed(qkvz_weight);
at::Tensor packed_w2 = is_vnni ? ba_weight : convert_weight_packed(ba_weight);
at::Tensor projected_states_qkvz = at::empty({batch, qkvz_dim}, hidden_states.options());
at::Tensor projected_states_ba = at::empty({batch, ba_dim}, hidden_states.options());
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_input_proj_cpu", [&] {
fused_input_proj_kernel_impl<scalar_t>(
projected_states_qkvz.data_ptr<scalar_t>(),
projected_states_ba.data_ptr<scalar_t>(),
hidden_states.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
packed_w2.data_ptr<scalar_t>(),
batch,
qkvz_dim,
ba_dim,
hidden_size);
});
return std::make_tuple(projected_states_qkvz, projected_states_ba);
}
python/sglang/kernels/aot/csrc/cpu/moe.cpp
+1371
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/aot/csrc/cpu/moe.h
+238
View File
@@ -0,0 +1,238 @@
#pragma once
#include "vec.h"
template <typename scalar_t>
inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) {
using Vec = at::vec::Vectorized<scalar_t>;
const Vec data_vec(val);
at::vec::map<scalar_t>([data_vec](Vec out) { return out = data_vec; }, out, out, size);
}
template <typename scalar_t>
inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
using Vec = at::vec::Vectorized<scalar_t>;
constexpr int kVecSize = Vec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
Vec data = Vec::loadu(input + d);
data.store(out + d);
}
for (; d < size; ++d) {
out[d] = input[d];
}
}
template <typename scalar_t>
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d]);
}
}
template <>
inline void copy_stub<uint8_t>(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) {
// size might be 64x + 32
std::memcpy(out, input, size * sizeof(uint8_t));
}
template <typename scalar_t, typename input_t>
inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict__ input, float weight, int64_t size) {
static_assert(
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
"copy_mul_stub only supports input_t == float or input_t == scalar_t");
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
const fVec weight_vec = fVec(weight);
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] * weight);
}
}
// acc from [topk, K] to [K]
template <typename scalar_t>
inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
if (topk == 1) {
// do copy for topk = 1
copy_stub(out, input, K);
} else {
// do sum for topk != 1
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= K - kVecSize; d += kVecSize) {
fVec sum_fvec0 = fVec(0.f);
fVec sum_fvec1 = fVec(0.f);
for (int t = 0; t < topk; ++t) {
auto [x_fvec0, x_fvec1] = load_float_vec2(input + t * K + d);
sum_fvec0 += x_fvec0;
sum_fvec1 += x_fvec1;
}
bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
out_bvec.store(out + d);
}
for (; d < K; ++d) {
float sum_val = 0.f;
for (int t = 0; t < topk; ++t) {
sum_val += static_cast<float>(input[t * K + d]);
}
out[d] = static_cast<scalar_t>(sum_val);
}
}
}
// out = input + input2 * scale
template <typename scalar_t, typename input_t>
inline void add_mul_stub(
scalar_t* __restrict__ out,
const input_t* __restrict__ input,
const scalar_t* __restrict__ input2,
float scale,
int64_t size) {
static_assert(
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
"add_mul_stub only supports input_t == float or input_t == scalar_t");
// out = input (without scale factor)
if (input2 == nullptr) {
copy_stub(out, input, size);
return;
}
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
const fVec s_vec = fVec(scale);
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
auto [y0, y1] = load_float_vec2(input2 + d);
x0 = x0 + y0 * s_vec;
x1 = x1 + y1 * s_vec;
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
}
}
template <typename scalar_t, typename input_t>
inline void silu_and_mul_stub(
scalar_t* __restrict__ out, const input_t* __restrict__ input, const input_t* __restrict__ input2, int64_t size) {
static_assert(
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
"silu_and_mul_stub only supports input_t == float or input_t == scalar_t");
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
auto [x0, x1] = load_float_vec2(input + d);
auto [y0, y1] = load_float_vec2(input2 + d);
x0 = fast_silu(x0) * y0;
x1 = fast_silu(x1) * y1;
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(out + d);
}
}
template <typename scalar_t, typename input_t>
inline void clamp_sigmoid_and_mul_stub(
scalar_t* __restrict__ out, const input_t* __restrict__ input, int64_t size, const float alpha, const float limit) {
static_assert(
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
"clamp_sigmoid_and_mul_stub only supports input_t == float or input_t == scalar_t");
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
const fVec limit_v = fVec(limit);
const fVec nlimit_v = fVec(-limit);
const fVec alpha_v = fVec(alpha);
#pragma GCC unroll 4
for (int64_t d = 0; d < 2 * size; d += bVec::size()) {
auto [x0_, y0_] = load_float_vec2(input + d);
auto [x0, y0] = at::vec::deinterleave2<float>(x0_, y0_);
x0 = at::vec::minimum(x0, limit_v);
y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
x0 = fast_sigmoid_glu(x0, alpha_v) * (y0 + one);
store_from_float_ext(out + d / 2, x0);
}
}
template <typename scalar_t>
inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
const fVec weight_vec = fVec(weight);
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] * weight);
}
}
// input = input + input2
inline void add_bias_stub(float* __restrict__ input, const float* __restrict__ input2, int64_t size) {
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = fVec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec x_fvec = fVec::loadu(input + d);
fVec y_fvec = fVec::loadu(input2 + d);
x_fvec = x_fvec + y_fvec;
x_fvec.store(input + d);
}
for (; d < size; ++d) {
input[d] = input[d] + input2[d];
}
}
template <typename scalar_t>
inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
const fVec weight_vec = fVec(weight);
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] * weight);
}
}
python/sglang/kernels/aot/csrc/cpu/moe_fp8.cpp
+400
View File
@@ -0,0 +1,400 @@
#include "common.h"
#include "gemm.h"
#include "moe.h"
template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
void fused_experts_fp_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic0,
scalar_t* __restrict__ ic1,
scalar_t* __restrict__ ic2,
scalar_t* __restrict__ A_tmp,
scalar_t* __restrict__ B_tmp,
float* __restrict__ C_tmp,
const scalar_t* __restrict__ input,
const packed_t* __restrict__ packed_w1,
const packed_t* __restrict__ packed_w2,
const float* __restrict__ w1_bias,
const float* __restrict__ w2_bias,
const param_t* __restrict__ w1s,
const param_t* __restrict__ w2s,
int64_t block_size_N,
int64_t block_size_K,
const float* __restrict__ topk_weights,
const int32_t* __restrict__ sorted_ids,
const int32_t* __restrict__ expert_ids,
const int32_t* __restrict__ offsets,
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad,
float alpha,
float limit,
CPUActMethod act_func,
bool with_bias) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
// stage 1: intermediate_cache0 = hidden_states @ w1
const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
const int64_t NB = div_up(2 * N, BLOCK_N);
int64_t scale_size_N = div_up(2 * N, block_size_N);
int64_t scale_size_K = div_up(K, block_size_K);
int64_t blocks_n_per_group = block_size_N / BLOCK_N;
std::function<int64_t(int64_t)> scale_offset_per_block;
if constexpr (is_mxfp4) {
scale_offset_per_block = [&](int64_t a) { return a * BLOCK_N; };
} else {
scale_offset_per_block = [&](int64_t a) { return a / blocks_n_per_group; };
}
const int64_t packed_K = get_row_size<packed_t>(K);
const int64_t stride_e = 2 * N * packed_K;
const int64_t stride_n = packed_K;
int64_t avg_M = std::max(int64_t(1), M * topk / E);
const bool use_brgemm = can_use_brgemm<packed_t>(avg_M);
int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// get local pointers
int tid = get_thread_num();
scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
// B shape [K, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const packed_t* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
const param_t* __restrict__ Bs =
w1s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
const float* __restrict__ B_bias = with_bias ? w1_bias + expert_id * 2 * N + nb * BLOCK_N : nullptr;
// do unpacking for the first row or a new expert
int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
int64_t m_size = offsets[mb + 1] - offsets[mb];
if (nb_offset == 0) {
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, input + index * K, K);
}
}
const int64_t offset = offsets[mb];
tinygemm_kernel<scalar_t>(
/* A */ A,
/* B */ B,
/* C */ ic0 + offset * 2 * N + nb * BLOCK_N,
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
/* Bbias */ B_bias,
/* scale */ Bs,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ 2 * N,
/* brg */ use_brgemm,
/* block_size_K */ block_size_K,
/* do_unpack */ do_unpack);
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
if (act_func == CPUActMethod::silu_and_mul) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
} else if (act_func == CPUActMethod::swiglu) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N / 2, alpha, limit);
clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N / 2, alpha, limit);
}
});
}
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
// w2 : [E, K, N] as [E, OC, IC]
const int64_t OC = K; // rename K as OC
const int64_t IC = N; // rename N as IC
const int64_t MB2 = MB;
const int64_t NB2 = div_up(OC, BLOCK_N);
scale_size_N = div_up(K, block_size_N);
scale_size_K = div_up(N, block_size_K);
const int64_t packed_IC = get_row_size<packed_t>(IC);
const int64_t stride_e2 = OC * packed_IC;
const int64_t stride_oc = packed_IC;
// parallel on [MB2, NB2]
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
int tid = get_thread_num();
alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t m_size = offsets[mb + 1] - offsets[mb];
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
// A ptr from ic1 of [M * topk, N] in sorted order
// so as to avoid copy A to tmp buffer again
const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
// B shape [IC, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const packed_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
const param_t* __restrict__ Bs =
w2s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
const float* __restrict__ B_bias = with_bias ? w2_bias + expert_id * OC + nb * BLOCK_N : nullptr;
// do unpacking for the first row or a new expert
int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
tinygemm_kernel<scalar_t>(
/* A */ A,
/* B */ B,
/* C */ C,
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
/* Bbias */ B_bias,
/* scale */ Bs,
/* M */ m_size,
/* N */ n_size,
/* K */ IC,
/* lda */ IC,
/* ldb */ n_size,
/* ldc */ BLOCK_N,
/* brg */ use_brgemm,
/* block_size_K */ block_size_K,
/* do_unpack */ do_unpack);
// 2.b copy from C to ic2 in original order
// and also mul topk_weights in float32
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m];
float weight = topk_weights[index];
copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
}
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
// stage 3: out = intermediate_cache2.sum(dim=1)
// from [M, topk, K] to [M, K]
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
}
});
}
#define INSTANTIATE_MOE_FP_TEMPLATE(TYPE1, TYPE2, TYPE3, IS_MXFP4) \
template void fused_experts_fp_kernel_impl<TYPE1, TYPE2, TYPE3, IS_MXFP4>( \
TYPE1* __restrict__ output, \
TYPE1* __restrict__ ic0, \
TYPE1* __restrict__ ic1, \
TYPE1* __restrict__ ic2, \
TYPE1* __restrict__ A_tmp, \
TYPE1* __restrict__ B_tmp, \
float* __restrict__ C_tmp, \
const TYPE1* __restrict__ input, \
const TYPE2* __restrict__ packed_w1, \
const TYPE2* __restrict__ packed_w2, \
const float* __restrict__ w1_bias, \
const float* __restrict__ w2_bias, \
const TYPE3* __restrict__ w1s, \
const TYPE3* __restrict__ w2s, \
int64_t block_size_N, \
int64_t block_size_K, \
const float* __restrict__ topk_weights, \
const int32_t* __restrict__ sorted_ids, \
const int32_t* __restrict__ expert_ids, \
const int32_t* __restrict__ offsets, \
int64_t M, \
int64_t N, \
int64_t K, \
int64_t E, \
int64_t topk, \
int64_t num_tokens_post_pad, \
float alpha, \
float limit, \
CPUActMethod act_func, \
bool with_bias)
INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false);
INSTANTIATE_MOE_FP_TEMPLATE(at::Half, at::Float8_e4m3fn, float, false);
INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, uint8_t, uint8_t, true);
INSTANTIATE_MOE_FP_TEMPLATE(at::Half, uint8_t, uint8_t, true);
template <typename scalar_t>
void shared_expert_fp8_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic0,
scalar_t* __restrict__ ic1,
scalar_t* __restrict__ B_tmp,
float* __restrict__ C_tmp,
const scalar_t* __restrict__ input,
const at::Float8_e4m3fn* __restrict__ packed_w1,
const at::Float8_e4m3fn* __restrict__ packed_w2,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
int64_t block_size_N,
int64_t block_size_K,
const scalar_t* __restrict__ fused_experts_out,
float routed_scaling_factor,
int64_t M,
int64_t N,
int64_t K) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
// stage 1: intermediate_cache0 = hidden_states @ w1
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(2 * N, BLOCK_N);
int64_t scale_size_K = div_up(K, block_size_K);
int64_t blocks_n_per_group = block_size_N / BLOCK_N;
const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
const bool apply_scaling_factor = fused_experts_out != nullptr;
int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
int tid = get_thread_num();
loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
// do unpacking for the first row
bool do_unpack = (mb == mb0);
tinygemm_kernel<scalar_t>(
/* A */ input + mb * BLOCK_M * K,
/* B */ packed_w1 + nb * BLOCK_N * K,
/* C */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
/* Bbias */ nullptr,
/* scale */ w1s + (nb / blocks_n_per_group) * scale_size_K,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ 2 * N,
/* brg */ use_brgemm,
/* block_size_K */ block_size_K,
/* do_unpack */ do_unpack);
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
// w2 : [K, N] as [OC, IC]
const int64_t OC = K; // rename K as OC
const int64_t IC = N; // rename N as IC
const int64_t MB2 = MB;
const int64_t NB2 = div_up(K, BLOCK_N);
scale_size_K = div_up(N, block_size_K);
// parallel on [MB2, NB2]
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
int tid = get_thread_num();
alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
// do unpacking for the first row
bool do_unpack = (mb == mb0);
// 2.a gemm: C = A @ B
tinygemm_kernel<scalar_t>(
/* A */ ic1 + mb * BLOCK_M * N,
/* B */ packed_w2 + nb * BLOCK_N * N,
/* C */ C,
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
/* Bbias */ nullptr,
/* scale */ w2s + (nb / blocks_n_per_group) * scale_size_K,
/* M */ m_size,
/* N */ n_size,
/* K */ IC,
/* lda */ IC,
/* ldb */ n_size,
/* ldc */ BLOCK_N,
/* brg */ use_brgemm,
/* block_size_K */ block_size_K,
/* do_unpack */ do_unpack);
// 2.b copy from C to output and add fused_experts_out
scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
const scalar_t* __restrict__ fused_out =
apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
for (int64_t m = 0; m < m_size; ++m) {
const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
}
});
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
}
#define INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(TYPE) \
template void shared_expert_fp8_kernel_impl<TYPE>( \
TYPE* __restrict__ output, \
TYPE* __restrict__ ic0, \
TYPE* __restrict__ ic1, \
TYPE* __restrict__ B_tmp, \
float* __restrict__ C_tmp, \
const TYPE* __restrict__ input, \
const at::Float8_e4m3fn* __restrict__ packed_w1, \
const at::Float8_e4m3fn* __restrict__ packed_w2, \
const float* __restrict__ w1s, \
const float* __restrict__ w2s, \
int64_t block_size_N, \
int64_t block_size_K, \
const TYPE* __restrict__ fused_experts_out, \
float routed_scaling_factor, \
int64_t M, \
int64_t N, \
int64_t K)
INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::BFloat16);
INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::Half);
python/sglang/kernels/aot/csrc/cpu/moe_int4.cpp
+318
View File
@@ -0,0 +1,318 @@
#include "common.h"
#include "gemm.h"
#include "moe.h"
template <int64_t N>
inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m, int64_t ldn) {
using Vec = at::vec::Vectorized<float>;
constexpr int kVecSize = Vec::size();
static_assert(N % kVecSize == 0, "copy_bias requires N to be a multiple of Vectorized<float>::size()");
const bool has_bias = bias_ptr != nullptr;
const Vec zero_vec(0.f);
for (int i = 0; i < m; ++i) {
#pragma GCC unroll 2
for (int j = 0; j < N; j += kVecSize) {
Vec vec = has_bias ? Vec::loadu(bias_ptr + j) : zero_vec;
vec.store(y_buf + i * ldn + j);
}
}
}
template <typename scalar_t>
void fused_experts_int4_w4a8_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic0,
scalar_t* __restrict__ ic1,
scalar_t* __restrict__ ic2,
uint8_t* __restrict__ A_tmp,
uint8_t* __restrict__ Aq_tmp,
float* __restrict__ As_tmp,
int32_t* __restrict__ Azp_tmp,
float* __restrict__ C_tmp,
int8_t* __restrict__ dqB_tmp,
const scalar_t* __restrict__ input,
const uint8_t* __restrict__ packed_w1,
const uint8_t* __restrict__ packed_w2,
const int8_t* __restrict__ w1z,
const int8_t* __restrict__ w2z,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
int group_size,
const float* __restrict__ topk_weights,
const int32_t* __restrict__ sorted_ids,
const int32_t* __restrict__ expert_ids,
const int32_t* __restrict__ offsets,
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
int num_threads = at::get_num_threads();
// int64_t buffer_size_nbytes = M * topk * N * 2
// M * topk * K * 2 +
// num_threads * BLOCK_M * K +
// num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float) +
// M * topk * 2 * N * 2 +
// max(M * K, M * topk * N) +
// M * topk * sizeof(float);
// intermediate_cache1 (scalar_t): START + M * topk * N
// intermediate_cache2 (scalar_t): + M * topk * K
// A_tmp (uint8_t): + num_threads * BLOCK_M * K
// C_tmp (float): + num_threads * 2 * BLOCK_M * BLOCK_N
// intermediate_cache0 (scalar_t): + M * topk * 2 * N
// Aq_tmp (uint8_t): + max(M * K, M * topk * N)
// As_tmp (float): + M * topk
// dqB_tmp (int8_t) + num_threads * _block_k * BlOCK_N
// stage 0: quantize input to uint8, [M, K]
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
}
});
int64_t _block_k = get_4bit_block_k_size(group_size);
auto Azp = at::ones({M * topk}).to(at::kInt).mul(128);
auto Azp_ptr = Azp.data_ptr<int32_t>();
// stage 1: intermediate_cache0 = hidden_states @ w1
const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
int64_t block_per_group = group_size / _block_k;
int64_t Kc = K / _block_k;
int64_t num_groups = K / group_size;
const int64_t stride_e = 2 * NB * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t)));
const bool sym_quant_act = false;
// weight + compensation shape = [E, Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)]
// scales/qzeros shape = [E, Nc, G, block_n]
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
// get local pointers
int tid = at::get_thread_num();
int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N;
int8_t* dqB_tmp2 = dqB_tmp1 + _block_k * BLOCK_N;
alignas(64) float As[BLOCK_M];
uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
bool is_brgemm_used = false;
for (int64_t i = begin; i < end; ++i) {
int64_t mb = i / NB;
int64_t nb = i % NB;
int64_t nb1 = nb + NB;
int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
// B shape [K, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const uint8_t* __restrict__ B = packed_w1 + expert_id * stride_e;
// Bz and Bs: [E, K/gs, 2N]
const int8_t* __restrict__ Bz = w1z + expert_id * (num_groups) * (2 * N);
const float* __restrict__ Bs = w1s + expert_id * (num_groups) * (2 * N);
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
int64_t m_size = offsets[mb + 1] - offsets[mb];
const bool use_brgemm = can_use_brgemm<int8_t>(m_size);
is_brgemm_used = is_brgemm_used || use_brgemm;
// copy to A [BLOCK_M, K]
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, Aq_tmp + index * K, K);
As[m] = As_tmp[index];
}
const int64_t offset = offsets[mb];
copy_bias<BLOCK_N>(nullptr, C0, m_size, BLOCK_N);
copy_bias<BLOCK_N>(nullptr, C1, m_size, BLOCK_N);
for (int kci = 0; kci < Kc; ++kci) {
int32_t* compensation_ptr =
sym_quant_act ? nullptr
: (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
_block_k * BLOCK_N / 2) /*Bcomp*/;
tinygemm_kernel<scalar_t>(
ic0 + offset * 2 * N + nb * BLOCK_N,
C0,
A + kci * _block_k,
As,
Azp_ptr,
B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/,
Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/,
compensation_ptr,
dqB_tmp1,
m_size,
_block_k,
K,
BLOCK_N,
2 * N,
kci == Kc - 1,
use_brgemm);
}
for (int kci = 0; kci < Kc; ++kci) {
int32_t* compensation_ptr =
sym_quant_act ? nullptr
: (int32_t*)(void*)(B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
_block_k * BLOCK_N / 2) /*Bcomp*/;
tinygemm_kernel<scalar_t>(
ic0 + offset * 2 * N + nb1 * BLOCK_N,
C1,
A + kci * _block_k,
As,
Azp_ptr,
B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/,
Bs + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
Bz + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/,
compensation_ptr,
dqB_tmp2,
m_size,
_block_k,
K,
BLOCK_N,
2 * N,
kci == Kc - 1,
use_brgemm);
}
}
if (is_brgemm_used) {
at::native::cpublas::brgemm_release();
}
});
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
// stage 1.5: quantize ic1 to uint8, [M * topk, N]
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
}
});
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
// w2 : [E, K, N] as [E, OC, IC]
const int64_t OC = K; // rename K as OC
const int64_t IC = N; // rename N as IC
const int64_t MB2 = MB;
const int64_t NB2 = div_up(OC, BLOCK_N);
const int64_t stride_oc = IC;
num_groups = IC / group_size;
Kc = IC / _block_k;
const int64_t stride_e2 = NB2 * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t)));
// parallel on [MB2, NB2]
at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
int tid = at::get_thread_num();
int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N;
float* __restrict__ C2 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
bool is_brgemm_used = false;
for (int64_t i = begin; i < end; ++i) {
int64_t mb = i / NB2;
int64_t nb = i % NB2;
int64_t m_size = offsets[mb + 1] - offsets[mb];
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
const bool use_brgemm = can_use_brgemm<int8_t>(m_size);
is_brgemm_used = is_brgemm_used || use_brgemm;
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
// B shape [IC, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const uint8_t* __restrict__ B = packed_w2 + expert_id * stride_e2;
// Bz and Bs: [E, IC/gs, OC]
const int8_t* __restrict__ Bz = w2z + expert_id * (num_groups)*OC;
const float* __restrict__ Bs = w2s + expert_id * (num_groups)*OC;
// A ptr from ic1 of [M * topk, N] in sorted order
// so as to avoid copy A to tmp buffer again
const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * IC;
const float* __restrict__ As = As_tmp + offsets[mb];
copy_bias<BLOCK_N>(nullptr, C2, m_size, BLOCK_N);
for (int kci = 0; kci < Kc; ++kci) {
int32_t* compensation_ptr =
sym_quant_act ? nullptr
: (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
_block_k * BLOCK_N / 2) /*Bcomp*/;
tinygemm_kernel<scalar_t>(
nullptr, /*store_out is false*/
C2,
A + kci * _block_k,
As,
Azp_ptr,
B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*zeros_b*/,
compensation_ptr,
dqB_tmp1,
m_size,
_block_k,
IC,
BLOCK_N,
BLOCK_N,
false,
use_brgemm);
}
// 2.b copy from C to ic2 in original order
// and also mul topk_weights in float32
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m];
float weight = topk_weights[index];
copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C2 + m * BLOCK_N, weight, n_size);
}
}
if (is_brgemm_used) {
at::native::cpublas::brgemm_release();
}
});
// stage 3: out = intermediate_cache2.sum(dim=1)
// from [M, topk, K] to [M, K]
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
}
});
}
#define INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(TYPE) \
template void fused_experts_int4_w4a8_kernel_impl<TYPE>( \
TYPE* __restrict__ output, \
TYPE* __restrict__ ic0, \
TYPE* __restrict__ ic1, \
TYPE* __restrict__ ic2, \
uint8_t* __restrict__ A_tmp, \
uint8_t* __restrict__ Aq_tmp, \
float* __restrict__ As_tmp, \
int32_t* __restrict__ Azp_tmp, \
float* __restrict__ C_tmp, \
int8_t* __restrict__ dqB_tmp, \
const TYPE* __restrict__ input, \
const uint8_t* __restrict__ packed_w1, \
const uint8_t* __restrict__ packed_w2, \
const int8_t* __restrict__ w1z, \
const int8_t* __restrict__ w2z, \
const float* __restrict__ w1s, \
const float* __restrict__ w2s, \
int group_size, \
const float* __restrict__ topk_weights, \
const int32_t* __restrict__ sorted_ids, \
const int32_t* __restrict__ expert_ids, \
const int32_t* __restrict__ offsets, \
int64_t M, \
int64_t N, \
int64_t K, \
int64_t E, \
int64_t topk, \
int64_t num_tokens_post_pad)
INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::BFloat16);
INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::Half);
python/sglang/kernels/aot/csrc/cpu/moe_int8.cpp
+960
View File
@@ -0,0 +1,960 @@
#include "common.h"
#include "gemm.h"
#include "moe.h"
namespace {
template <typename scalar_t, int BLOCK_N>
inline void silu_and_mul(
scalar_t* __restrict__ C,
const int32_t* __restrict__ C0, // x: x0, x1
const int32_t* __restrict__ C1, // y: y0, y1
const float* __restrict__ As,
const float* __restrict__ Bs0,
const float* __restrict__ Bs1,
const int32_t* __restrict__ Bcomp0,
const int32_t* __restrict__ Bcomp1,
int64_t m_size,
int64_t N) {
#if defined(CPU_CAPABILITY_AVX512)
constexpr int COLS = BLOCK_N / 16;
static_assert(COLS % 2 == 0);
__m512 vc0[COLS];
__m512 vc1[COLS];
__m512i vcomp0[COLS];
__m512i vcomp1[COLS];
__m512 vas;
__m512 vbs0[COLS];
__m512 vbs1[COLS];
auto load_scale_and_comp = [&](auto col) {
vcomp0[col] = _mm512_loadu_si512(Bcomp0 + col * 16);
vcomp1[col] = _mm512_loadu_si512(Bcomp1 + col * 16);
vbs0[col] = _mm512_loadu_ps(Bs0 + col * 16);
vbs1[col] = _mm512_loadu_ps(Bs1 + col * 16);
};
Unroll<COLS>{}(load_scale_and_comp);
auto scalec = [&](auto col, int64_t m) {
// update As
vas = _mm512_set1_ps(As[m]);
// C = As * (C - Bcomp) * Bs
__m512i vc32_0 = _mm512_loadu_si512(C0 + m * BLOCK_N + col * 16);
__m512i vc32_1 = _mm512_loadu_si512(C1 + m * BLOCK_N + col * 16);
vc0[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32_0, vcomp0[col]));
vc1[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32_1, vcomp1[col]));
vc0[col] = _mm512_mul_ps(_mm512_mul_ps(vc0[col], vas), vbs0[col]);
vc1[col] = _mm512_mul_ps(_mm512_mul_ps(vc1[col], vas), vbs1[col]);
};
auto silu_and_mul = [&](auto col) {
__m512 x = vc0[col];
__m512 y = vc1[col];
vc0[col] = _mm512_mul_ps(_mm512_rcp14_silu_ps(x), y);
};
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
auto storec = [&](auto col, int64_t m) {
if constexpr (col % 2 == 0) {
fVec x0 = fVec(vc0[col + 0]);
fVec x1 = fVec(vc0[col + 1]);
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(C + m * N + col * 16);
}
};
for (int64_t m = 0; m < m_size; ++m) {
Unroll<COLS>{}(scalec, m);
Unroll<COLS>{}(silu_and_mul);
Unroll<COLS>{}(storec, m);
}
#else
TORCH_CHECK(false, "silu_and_mul: scalar path not implemented!");
#endif
}
template <int BLOCK_N>
inline void scale_C(
float* __restrict__ C,
const int32_t* __restrict__ Ctmp,
const float* __restrict__ As,
const float* __restrict__ Bs,
const int32_t* __restrict__ Bcomp,
int64_t m_size) {
#if defined(CPU_CAPABILITY_AVX512)
constexpr int COLS = BLOCK_N / 16;
static_assert(COLS % 2 == 0);
__m512 vc[COLS];
__m512i vcomp[COLS];
__m512 vas;
__m512 vbs[COLS];
auto load_scale_and_comp = [&](auto col) {
vcomp[col] = _mm512_loadu_si512(Bcomp + col * 16);
vbs[col] = _mm512_loadu_ps(Bs + col * 16);
};
Unroll<COLS>{}(load_scale_and_comp);
auto scalec = [&](auto col, int64_t m) {
// update As
vas = _mm512_set1_ps(As[m]);
// C = As * (C - Bcomp) * Bs
__m512i vc32 = _mm512_loadu_si512(Ctmp + m * BLOCK_N + col * 16);
vc[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32, vcomp[col]));
vc[col] = _mm512_mul_ps(_mm512_mul_ps(vc[col], vas), vbs[col]);
_mm512_storeu_ps(C + m * BLOCK_N + col * 16, vc[col]);
};
for (int64_t m = 0; m < m_size; ++m) {
Unroll<COLS>{}(scalec, m);
}
#else
TORCH_CHECK(false, "scale_C: scalar path not implemented!");
#endif
}
/// gemm for w13
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_vnni {
static inline void apply(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B0,
const int8_t* __restrict__ B1,
scalar_t* __restrict__ C,
const float* __restrict__ As,
const float* __restrict__ Bs0,
const float* __restrict__ Bs1,
const int32_t* __restrict__ Bcomp0,
const int32_t* __restrict__ Bcomp1,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
}
};
#if defined(CPU_CAPABILITY_AVX512)
template <int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
static inline void apply(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B0,
const int8_t* __restrict__ B1,
at::BFloat16* __restrict__ C,
const float* __restrict__ As,
const float* __restrict__ Bs0,
const float* __restrict__ Bs1,
const int32_t* __restrict__ Bcomp0,
const int32_t* __restrict__ Bcomp1,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
constexpr int ROWS = BLOCK_M;
constexpr int COLS = BLOCK_N / 16;
static_assert(COLS % 2 == 0);
__m512i va;
__m512i vb0[COLS];
__m512i vb1[COLS];
__m512i vc0[ROWS * COLS];
__m512i vc1[ROWS * COLS];
__m512i vcomp0[COLS];
__m512i vcomp1[COLS];
__m512 vas;
__m512 vbs0[COLS];
__m512 vbs1[COLS];
auto loadc = [&](auto i) {
vc0[i] = _mm512_set1_epi32(0);
vc1[i] = _mm512_set1_epi32(0);
};
Unroll<ROWS * COLS>{}(loadc);
const int64_t K4 = K >> 2;
const int64_t lda4 = lda >> 2;
const int64_t ldb4 = ldb; // ldb * 4 >> 2;
const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
const int32_t* b0_ptr = reinterpret_cast<const int32_t*>(B0);
const int32_t* b1_ptr = reinterpret_cast<const int32_t*>(B1);
auto compute = [&](auto i, int64_t k) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
if constexpr (col == 0) {
va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
}
if constexpr (row == 0) {
vb0[col] = _mm512_loadu_si512(b0_ptr + k * ldb4 + col * 16);
vb1[col] = _mm512_loadu_si512(b1_ptr + k * ldb4 + col * 16);
}
vc0[i] = _mm512_dpbusd_epi32(vc0[i], va, vb0[col]);
vc1[i] = _mm512_dpbusd_epi32(vc1[i], va, vb1[col]);
};
for (int64_t k = 0; k < K4; ++k) {
Unroll<ROWS * COLS>{}(compute, k);
}
auto scalec = [&](auto i) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
// load a scale
if constexpr (col == 0) {
vas = _mm512_set1_ps(As[row]);
}
// load b scale and vcomp
if constexpr (row == 0) {
vbs0[col] = _mm512_loadu_ps(Bs0 + col * 16);
vbs1[col] = _mm512_loadu_ps(Bs1 + col * 16);
vcomp0[col] = _mm512_loadu_si512(Bcomp0 + col * 16);
vcomp1[col] = _mm512_loadu_si512(Bcomp1 + col * 16);
}
__m512 c0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc0[i], vcomp0[col]));
__m512 c1 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc1[i], vcomp1[col]));
vc0[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c0, vas), vbs0[col]));
vc1[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c1, vas), vbs1[col]));
};
Unroll<ROWS * COLS>{}(scalec);
auto storec = [&](auto i) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
// for COLS = 2, 4 use 512bit store
if constexpr (col % 2 == 0) {
__m512 x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]);
__m512 x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
__m512 y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
__m512 y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0);
x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1);
_mm512_storeu_si512(
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
(__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0))));
}
};
Unroll<ROWS * COLS>{}(storec);
}
};
#endif
#define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE) \
tinygemm_kernel_vnni<scalar_t, MB_SIZE, NB_SIZE>::apply( \
A + mb_start * lda, \
B0 + nb_start * 4, \
B1 + nb_start * 4, \
C + mb_start * ldc + nb_start, \
As + mb_start, \
Bs0 + nb_start, \
Bs1 + nb_start, \
Bcomp0 + nb_start, \
Bcomp1 + nb_start, \
K, \
lda, \
ldb, \
ldc);
template <typename scalar_t>
void tinygemm_kernel(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B0,
const int8_t* __restrict__ B1,
scalar_t* __restrict__ C,
const float* __restrict__ As,
const float* __restrict__ Bs0,
const float* __restrict__ Bs1,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
// pattern: 1-(2+2)-(8+8)
constexpr int64_t BLOCK_M = 4;
constexpr int64_t BLOCK_N = 32;
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
for (int mb = 0; mb < MB; ++mb) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
for (int64_t nb = 0; nb < NB; ++nb) {
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
switch (mb_size << 4 | nb_size >> 4) {
case 0x12:
LAUNCH_TINYGEMM_KERNEL_VNNI(1, 32);
break;
case 0x22:
LAUNCH_TINYGEMM_KERNEL_VNNI(2, 32);
break;
case 0x32:
LAUNCH_TINYGEMM_KERNEL_VNNI(3, 32);
break;
case 0x42:
LAUNCH_TINYGEMM_KERNEL_VNNI(4, 32);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
}
}
}
}
/// gemm for w2
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_vnni2 {
static inline void apply(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B,
float* __restrict__ C,
const float* __restrict__ As,
const float* __restrict__ Bs,
const int32_t* __restrict__ Bcomp,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
}
};
#if defined(CPU_CAPABILITY_AVX512)
template <int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
static inline void apply(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B,
float* __restrict__ C,
const float* __restrict__ As,
const float* __restrict__ Bs,
const int32_t* __restrict__ Bcomp,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
constexpr int ROWS = BLOCK_M;
constexpr int COLS = BLOCK_N / 16;
static_assert(COLS % 2 == 0);
__m512i va;
__m512i vb[COLS];
__m512i vc[ROWS * COLS];
__m512i vcomp[COLS];
__m512 vas;
__m512 vbs[COLS];
auto loadc = [&](auto i) { vc[i] = _mm512_set1_epi32(0); };
Unroll<ROWS * COLS>{}(loadc);
const int64_t K4 = K >> 2;
const int64_t lda4 = lda >> 2;
const int64_t ldb4 = ldb; // ldb * 4 >> 2;
const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
auto compute = [&](auto i, int64_t k) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
if constexpr (col == 0) {
va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
}
if constexpr (row == 0) {
vb[col] = _mm512_loadu_si512(b_ptr + k * ldb4 + col * 16);
}
vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
};
for (int64_t k = 0; k < K4; ++k) {
Unroll<ROWS * COLS>{}(compute, k);
}
auto storec = [&](auto i) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
// load a scale
if constexpr (col == 0) {
vas = _mm512_set1_ps(As[row]);
}
// load b scale and vcomp per 2 vectors
// also load bias if any
if constexpr (row == 0) {
if constexpr (col % 2 == 0) {
vbs[col + 0] = _mm512_loadu_ps(Bs + col * 16);
vbs[col + 1] = _mm512_loadu_ps(Bs + col * 16 + 16);
vcomp[col + 0] = _mm512_loadu_si512(Bcomp + col * 16);
vcomp[col + 1] = _mm512_loadu_si512(Bcomp + col * 16 + 16);
}
}
__m512 x = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[i], vcomp[col]));
x = _mm512_mul_ps(_mm512_mul_ps(x, vas), vbs[col]);
_mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), x);
};
Unroll<ROWS * COLS>{}(storec);
}
};
#endif
#define LAUNCH_TINYGEMM_KERNEL_VNNI2(MB_SIZE, NB_SIZE) \
tinygemm_kernel_vnni2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
A + mb_start * lda, \
B + nb_start * 4, \
C + mb_start * ldc + nb_start, \
As + mb_start, \
Bs + nb_start, \
Bcomp + nb_start, \
K, \
lda, \
ldb, \
ldc);
template <typename scalar_t>
void tinygemm_kernel(
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B,
float* __restrict__ C,
const float* __restrict__ As,
const float* __restrict__ Bs,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc) {
// B compensation
const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
// pattern: 1-4-16
constexpr int64_t BLOCK_M = 4;
constexpr int64_t BLOCK_N = 64;
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
for (int64_t mb = 0; mb < MB; ++mb) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
for (int64_t nb = 0; nb < NB; ++nb) {
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
switch (mb_size << 4 | nb_size >> 4) {
case 0x12:
LAUNCH_TINYGEMM_KERNEL_VNNI2(1, 32);
break;
case 0x22:
LAUNCH_TINYGEMM_KERNEL_VNNI2(2, 32);
break;
case 0x32:
LAUNCH_TINYGEMM_KERNEL_VNNI2(3, 32);
break;
case 0x42:
LAUNCH_TINYGEMM_KERNEL_VNNI2(4, 32);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
}
}
}
}
} // anonymous namespace
template <typename scalar_t>
void fused_experts_int8_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic1,
scalar_t* __restrict__ ic2,
uint8_t* __restrict__ A_tmp,
float* __restrict__ C_tmp,
uint8_t* __restrict__ Aq_tmp,
float* __restrict__ As_tmp,
const scalar_t* __restrict__ input,
const int8_t* __restrict__ packed_w1,
const int8_t* __restrict__ packed_w2,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
const float* __restrict__ topk_weights,
const int32_t* __restrict__ sorted_ids,
const int32_t* __restrict__ expert_ids,
const int32_t* __restrict__ offsets,
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad) {
// handle 2 tiles per block
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
// stage 0: quantize input to uint8, [M, K]
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
}
});
// stage 1: intermediate_cache1 = silu(hidden_states @ w1)
const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
// strides for w1: [E, 2N, K]
TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
// K and N are packed for int8
const int64_t packed_K = get_row_size<int8_t>(K);
const int64_t packed_N = get_row_size<int8_t>(N);
const int64_t stride_e = 2 * N * packed_K;
const int64_t stride_n = packed_K;
int64_t avg_M = std::max(int64_t(1), M * topk / E);
const bool use_brgemm = can_use_brgemm<int8_t>(avg_M);
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// get local pointers
int tid = get_thread_num();
uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
int32_t* __restrict__ C0 = reinterpret_cast<int32_t*>(C_tmp) + tid * 2 * BLOCK_M * BLOCK_N;
int32_t* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
alignas(64) float As[BLOCK_M];
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
// nb_upper from top half and nb_lower from bottom half
int64_t nb_upper = nb, nb_lower = nb + NB;
int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
// B shape [K, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const int8_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n;
const int8_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;
const float* __restrict__ Bs0 = w1s + expert_id * 2 * N + nb_upper * BLOCK_N;
const float* __restrict__ Bs1 = w1s + expert_id * 2 * N + nb_lower * BLOCK_N;
int64_t m_size = offsets[mb + 1] - offsets[mb];
if (nb_offset == 0) {
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, Aq_tmp + index * K, K);
As[m] = As_tmp[index];
}
}
if (use_brgemm) {
// 1.b gemm: C0 = A @ B0
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ A,
/* B */ B0,
/* C */ C0);
// 1.c gemm: C1 = A @ B1
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ A,
/* B */ B1,
/* C */ C1);
const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
// 1.d silu and mul
const int64_t offset = offsets[mb];
silu_and_mul<scalar_t, BLOCK_N>(
ic1 + offset * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
} else {
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
const int64_t offset = offsets[mb];
tinygemm_kernel(
/* A */ A,
/* B0 */ B0,
/* B1 */ B1,
/* C */ ic1 + offset * N + nb * BLOCK_N,
/* As */ As,
/* Bs0 */ Bs0,
/* Bs1 */ Bs1,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ N);
}
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
// stage 1.5: quantize ic1 to uint8, [M * topk, N]
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
}
});
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
// w2 : [E, K, N] as [E, OC, IC]
const int64_t OC = K; // rename K as OC
const int64_t IC = N; // rename N as IC
const int64_t MB2 = MB;
const int64_t NB2 = div_up(OC, BLOCK_N);
const int64_t stride_e2 = OC * packed_N;
const int64_t stride_oc = packed_N;
// parallel on [MB2, NB2]
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// get local pointers
int tid = get_thread_num();
float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
int32_t* __restrict__ C32 = reinterpret_cast<int32_t*>(C + BLOCK_M * BLOCK_N);
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t m_size = offsets[mb + 1] - offsets[mb];
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
// A ptr from ic1 of [M * topk, N] in sorted order
// so as to avoid copy A to tmp buffer again
const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * N;
const float* __restrict__ As = As_tmp + offsets[mb];
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
// B shape [IC, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const int8_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
const float* __restrict__ Bs = w2s + expert_id * K + nb * BLOCK_N;
// 2.a gemm: C = A @ B
if (use_brgemm) {
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ IC,
/* lda */ IC,
/* ldb */ n_size,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ A,
/* B */ B,
/* C */ C32);
// apply scales
const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * IC);
scale_C<BLOCK_N>(C, C32, As, Bs, Bcomp, m_size);
} else {
tinygemm_kernel<scalar_t>(
/* A */ A,
/* B */ B,
/* C */ C,
/* As */ As,
/* Bs */ Bs,
/* M */ m_size,
/* N */ n_size,
/* K */ IC,
/* lda */ IC,
/* ldb */ n_size,
/* ldc */ BLOCK_N);
}
// 2.b copy from C to ic2 in original order
// and also mul topk_weights in float32
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m];
float weight = topk_weights[index];
copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
}
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
// stage 3: out = intermediate_cache2.sum(dim=1)
// from [M, topk, K] to [M, K]
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
}
});
}
#define INSTANTIATE_MOE_INT8_TEMPLATE(TYPE) \
template void fused_experts_int8_kernel_impl<TYPE>( \
TYPE* __restrict__ output, \
TYPE* __restrict__ ic1, \
TYPE* __restrict__ ic2, \
uint8_t* __restrict__ A_tmp, \
float* __restrict__ C_tmp, \
uint8_t* __restrict__ Aq_tmp, \
float* __restrict__ As_tmp, \
const TYPE* __restrict__ input, \
const int8_t* __restrict__ packed_w1, \
const int8_t* __restrict__ packed_w2, \
const float* __restrict__ w1s, \
const float* __restrict__ w2s, \
const float* __restrict__ topk_weights, \
const int32_t* __restrict__ sorted_ids, \
const int32_t* __restrict__ expert_ids, \
const int32_t* __restrict__ offsets, \
int64_t M, \
int64_t N, \
int64_t K, \
int64_t E, \
int64_t topk, \
int64_t num_tokens_post_pad)
INSTANTIATE_MOE_INT8_TEMPLATE(at::BFloat16);
INSTANTIATE_MOE_INT8_TEMPLATE(at::Half);
template <typename scalar_t>
void shared_expert_int8_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic1,
float* __restrict__ C_tmp,
uint8_t* __restrict__ Aq_tmp,
float* __restrict__ As_tmp,
const scalar_t* __restrict__ input,
const int8_t* __restrict__ packed_w1,
const int8_t* __restrict__ packed_w2,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
const scalar_t* __restrict__ fused_experts_out,
float routed_scaling_factor,
int64_t M,
int64_t N,
int64_t K) {
// handle 2 tiles per block
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
// stage 0: quantize input to uint8, [M, K]
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
}
});
// stage 1: intermediate_cache1 = silu(hidden_states @ w1)
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
// K and N are packed for int8
const int64_t packed_K = get_row_size<int8_t>(K);
const int64_t packed_N = get_row_size<int8_t>(N);
const int64_t stride_n = packed_K;
const bool use_brgemm = can_use_brgemm<int8_t>(M);
const bool apply_scaling_factor = fused_experts_out != nullptr;
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// get local pointers
int tid = get_thread_num();
int32_t* __restrict__ C0 = reinterpret_cast<int32_t*>(C_tmp) + tid * 2 * BLOCK_M * BLOCK_N;
int32_t* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
// nb_upper from top half and nb_lower from bottom half
int64_t nb_upper = nb, nb_lower = nb + NB;
int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
// A shape [m_size, K]
const uint8_t* A = Aq_tmp + mb * BLOCK_M * K;
const float* As = As_tmp + mb * BLOCK_M;
// B shape [K, n_size] in vnni format
const int8_t* __restrict__ B0 = packed_w1 + nb_upper * BLOCK_N * stride_n;
const int8_t* __restrict__ B1 = packed_w1 + nb_lower * BLOCK_N * stride_n;
const float* __restrict__ Bs0 = w1s + nb_upper * BLOCK_N;
const float* __restrict__ Bs1 = w1s + nb_lower * BLOCK_N;
if (use_brgemm) {
// 1.b gemm: C0 = A @ B0
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ A,
/* B */ B0,
/* C */ C0);
// 1.c gemm: C1 = A @ B1
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ A,
/* B */ B1,
/* C */ C1);
const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
// 1.d silu and mul
silu_and_mul<scalar_t, BLOCK_N>(
ic1 + mb * BLOCK_M * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
} else {
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
tinygemm_kernel(
/* A */ A,
/* B0 */ B0,
/* B1 */ B1,
/* C */ ic1 + mb * BLOCK_M * N + nb * BLOCK_N,
/* As */ As,
/* Bs0 */ Bs0,
/* Bs1 */ Bs1,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ N);
}
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
// stage 1.5: quantize ic1 to uint8, [M * topk, N]
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
}
});
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
// w2 : [K, N] as [OC, IC]
const int64_t OC = K; // rename K as OC
const int64_t IC = N; // rename N as IC
const int64_t MB2 = MB;
const int64_t NB2 = div_up(OC, BLOCK_N);
const int64_t stride_oc = packed_N;
// parallel on [MB2, NB2]
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// get local pointers
int tid = get_thread_num();
float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
int32_t* __restrict__ C32 = reinterpret_cast<int32_t*>(C + BLOCK_M * BLOCK_N);
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
// A shape [m_size, IC]
const uint8_t* __restrict__ A = Aq_tmp + mb * BLOCK_M * N;
const float* __restrict__ As = As_tmp + mb * BLOCK_M;
// B shape [IC, n_size] in vnni format
const int8_t* __restrict__ B = packed_w2 + nb * BLOCK_N * stride_oc;
const float* __restrict__ Bs = w2s + nb * BLOCK_N;
if (use_brgemm) {
at::native::cpublas::brgemm(
/* M */ m_size,
/* N */ n_size,
/* K */ IC,
/* lda */ IC,
/* ldb */ n_size,
/* ldc */ BLOCK_N,
/* add_C */ false,
/* A */ A,
/* B */ B,
/* C */ C32);
// apply scales
const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * IC);
scale_C<BLOCK_N>(C, C32, As, Bs, Bcomp, m_size);
} else {
// 2.a gemm: C = A @ B
tinygemm_kernel<scalar_t>(
/* A */ A,
/* B */ B,
/* C */ C,
/* As */ As,
/* Bs */ Bs,
/* M */ m_size,
/* N */ n_size,
/* K */ IC,
/* lda */ IC,
/* ldb */ n_size,
/* ldc */ BLOCK_N);
}
// 2.b copy from C to output and add fused_experts_out
scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
const scalar_t* __restrict__ fused_out =
apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
for (int64_t m = 0; m < m_size; ++m) {
const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
}
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
#define INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(TYPE) \
template void shared_expert_int8_kernel_impl<TYPE>( \
TYPE* __restrict__ output, \
TYPE* __restrict__ ic1, \
float* __restrict__ C_tmp, \
uint8_t* __restrict__ Aq_tmp, \
float* __restrict__ As_tmp, \
const TYPE* __restrict__ input, \
const int8_t* __restrict__ packed_w1, \
const int8_t* __restrict__ packed_w2, \
const float* __restrict__ w1s, \
const float* __restrict__ w2s, \
const TYPE* __restrict__ fused_experts_out, \
float routed_scaling_factor, \
int64_t M, \
int64_t N, \
int64_t K)
INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::BFloat16);
INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::Half);
python/sglang/kernels/aot/csrc/cpu/norm.cpp
+797
View File
@@ -0,0 +1,797 @@
#include "common.h"
#include "vec.h"
namespace {
struct NormParams {
// Treat all tensors as [B, H, T, D]:
// 2D -> [B, 1, 1, D]
// 3D -> [B, 1, T, D]
// 4D -> [B, H, T, D]
//
// Input: last dimension contiguous.
// Output: contiguous.
int ndim{0};
int64_t B{1}, H{1}, T{1}, D{1};
int64_t i_strideB{0}, i_strideH{0}, i_strideT{0};
float eps{1e-5f};
float shift{0.f};
const void* weight{nullptr};
const void* bias{nullptr};
explicit NormParams(const at::Tensor& input, float eps_) : ndim(input.dim()), eps(eps_) {
TORCH_CHECK(ndim >= 2 && ndim <= 4, "Expected a 2D/3D/4D tensor, got ", ndim, "D.");
B = input.size(0);
D = input.size(ndim - 1);
i_strideB = input.stride(0);
switch (ndim) {
case 2:
break;
case 3:
T = input.size(1);
i_strideT = input.stride(1);
break;
case 4:
H = input.size(1);
T = input.size(2);
i_strideH = input.stride(1);
i_strideT = input.stride(2);
break;
default:
TORCH_INTERNAL_ASSERT(false);
}
}
inline int64_t rows() const {
return B * H * T;
}
inline int64_t input_offset(int64_t b, int64_t h, int64_t t) const {
return b * i_strideB + h * i_strideH + t * i_strideT;
}
inline int64_t output_offset(int64_t b, int64_t h, int64_t t) const {
return ((b * H + h) * T + t) * D;
}
};
enum class NormMode {
L2Norm, // y = x / sqrt(mean(x^2) + eps)
RMSNorm, // y = x * weight / sqrt(mean(x^2) + eps)
GemmaNorm, // y = x * (weight + scale_shift) / sqrt(mean(x^2) + eps)
LayerNorm, // y = (x - mean(x)) * weight / sqrt(var(x) + eps) + bias
RMSNormGated, // y = x * weight / sqrt(mean(x^2) + eps) * SiLU(gate)
};
struct NormTraitsBase {
static constexpr bool has_weight = false;
static constexpr bool has_bias = false;
static constexpr bool has_shift = false;
static constexpr bool has_mean = false;
static constexpr bool has_gate = false;
template <typename VT>
static inline VT apply_weight(VT x, VT w) {
return x * w;
}
#if defined(CPU_CAPABILITY_AVX512)
static inline __m512 apply_weight(__m512 x, __m512 w) {
return _mm512_mul_ps(x, w);
}
#endif
};
template <NormMode M>
struct NormTraits : NormTraitsBase {};
template <>
struct NormTraits<NormMode::RMSNorm> : NormTraitsBase {
static constexpr bool has_weight = true;
};
template <>
struct NormTraits<NormMode::GemmaNorm> : NormTraitsBase {
static constexpr bool has_weight = true;
static constexpr bool has_shift = true;
template <typename VT>
static inline VT apply_shift(VT w, VT shift) {
return w + shift;
}
#if defined(CPU_CAPABILITY_AVX512)
static inline __m512 apply_shift(__m512 w, __m512 shift) {
return _mm512_add_ps(w, shift);
}
#endif
};
// LayerNorm: Var(X) = E(X^2) - (E(X))^2, refer to FlashInfer impl:
// https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/norm.cuh#L552
template <>
struct NormTraits<NormMode::LayerNorm> : NormTraitsBase {
static constexpr bool has_weight = true;
static constexpr bool has_bias = true;
static constexpr bool has_mean = true;
template <typename VT>
static inline VT apply_bias(VT x, VT bias) {
return x + bias;
}
#if defined(CPU_CAPABILITY_AVX512)
static inline __m512 apply_bias(__m512 x, __m512 bias) {
return _mm512_add_ps(x, bias);
}
#endif
};
template <>
struct NormTraits<NormMode::RMSNormGated> : NormTraitsBase {
static constexpr bool has_weight = true;
static constexpr bool has_gate = true;
static inline float apply_gate(float x, float gate) {
return x * (gate / (1.f + std::exp(-gate)));
}
static inline at::vec::Vectorized<float> apply_gate(at::vec::Vectorized<float> x, at::vec::Vectorized<float> gate) {
const auto one = at::vec::Vectorized<float>(1.f);
return x * (gate / (one + gate.neg().exp_u20()));
}
#if defined(CPU_CAPABILITY_AVX512)
static inline __m512 apply_gate(__m512 x, __m512 gate) {
__m512 minus_gate = _mm512_xor_ps(_mm512_set1_ps(-0.f), gate);
__m512 denom = _mm512_add_ps(_mm512_exp_u20_ps(minus_gate), _mm512_set1_ps(1.0f));
// NOTE: avoid vdivps -> use reciprocal
__m512 sigmoid = _mm512_mul_ps(gate, _mm512_rcp14_ps(denom));
return _mm512_mul_ps(x, sigmoid);
}
#endif
};
template <NormMode M, typename scalar_t, int D>
struct NormReduce;
#if defined(CPU_CAPABILITY_AVX512)
template <NormMode M, int D>
struct NormReduce<M, at::BFloat16, D> {
static inline void apply(
at::BFloat16* __restrict__ out,
const at::BFloat16* __restrict__ input,
const at::BFloat16* __restrict__ gate,
const NormParams& params) {
static_assert(D % 32 == 0);
constexpr int COLS = D / 32;
const bool use_bias = params.bias != nullptr;
__m512bh va[COLS];
__m512 vmean, vrscale;
const __m512 vshift = _mm512_set1_ps(params.shift);
// step 1: load input and do reduce with avx512-bf16
__m512 vsum = _mm512_set1_ps(0.f);
__m512 vsum2 = _mm512_set1_ps(0.f);
Unroll<COLS>{}([&](auto col) {
va[col] = (__m512bh)(_mm512_loadu_si512(input + col * 32));
if constexpr (NormTraits<M>::has_mean) {
vsum = _mm512_add_ps(vsum, CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)va[col], 0)));
vsum = _mm512_add_ps(vsum, CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)va[col], 1)));
}
vsum2 = _mm512_dpbf16_ps(vsum2, va[col], va[col]);
});
// compute mean (if has_mean) and rscale
float sum2 = _mm512_reduce_add_ps(vsum2);
float variance = sum2 / D;
if constexpr (NormTraits<M>::has_mean) {
float sum = _mm512_reduce_add_ps(vsum);
float mean = sum / D;
variance -= mean * mean;
vmean = _mm512_set1_ps(mean);
}
float rscale = 1.f / std::sqrt(variance + params.eps);
vrscale = _mm512_set1_ps(rscale);
// step 2: apply scale to output
Unroll<COLS>{}([&](auto col) {
__m512i a16 = (__m512i)va[col];
__m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0));
__m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1));
if constexpr (NormTraits<M>::has_mean) {
va0 = _mm512_sub_ps(va0, vmean);
va1 = _mm512_sub_ps(va1, vmean);
}
va0 = _mm512_mul_ps(va0, vrscale);
va1 = _mm512_mul_ps(va1, vrscale);
if constexpr (NormTraits<M>::has_weight) {
// TODO: need to block B to hide weight reload
const at::BFloat16* weight = static_cast<const at::BFloat16*>(params.weight);
__m512i w16 = (__m512i)(_mm512_loadu_si512(weight + col * 32));
__m512 w0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(w16, 0));
__m512 w1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(w16, 1));
if constexpr (NormTraits<M>::has_shift) {
w0 = NormTraits<M>::apply_shift(w0, vshift);
w1 = NormTraits<M>::apply_shift(w1, vshift);
}
va0 = NormTraits<M>::apply_weight(va0, w0);
va1 = NormTraits<M>::apply_weight(va1, w1);
}
if constexpr (NormTraits<M>::has_bias) {
if (use_bias) {
const at::BFloat16* bias = static_cast<const at::BFloat16*>(params.bias);
__m512i b16 = (__m512i)(_mm512_loadu_si512(bias + col * 32));
__m512 vbias0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
__m512 vbias1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
va0 = NormTraits<M>::apply_bias(va0, vbias0);
va1 = NormTraits<M>::apply_bias(va1, vbias1);
}
}
if constexpr (NormTraits<M>::has_gate) {
__m512i g16 = (__m512i)(_mm512_loadu_si512(gate + col * 32));
__m512 vgate0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(g16, 0));
__m512 vgate1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(g16, 1));
va0 = NormTraits<M>::apply_gate(va0, vgate0);
va1 = NormTraits<M>::apply_gate(va1, vgate1);
}
_mm512_storeu_si512(out + col * 32, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)));
});
}
};
#endif
template <NormMode M, typename scalar_t, bool has_residual>
struct NormReduceGeneric {
static inline void apply(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ gate,
scalar_t* __restrict__ residual,
const NormParams& params,
int D) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
const bool use_bias = params.bias != nullptr;
fVec sum_fvec{0.f}, sum2_fvec{0.f};
float sum_val{0.f}, sum2_val{0.f};
int d;
#pragma GCC unroll 4
for (d = 0; d <= D - kVecSize; d += kVecSize) {
auto [x_fvec0, x_fvec1] = load_float_vec2(input + d);
if constexpr (has_residual) {
auto [r_fvec0, r_fvec1] = load_float_vec2(residual + d);
x_fvec0 += r_fvec0;
x_fvec1 += r_fvec1;
}
sum2_fvec += x_fvec0 * x_fvec0;
sum2_fvec += x_fvec1 * x_fvec1;
if constexpr (NormTraits<M>::has_mean) {
sum_fvec += x_fvec0;
sum_fvec += x_fvec1;
}
}
#pragma GCC unroll 4
for (; d < D; ++d) {
float x_val = static_cast<float>(input[d]);
if constexpr (has_residual) {
x_val += static_cast<float>(residual[d]);
}
sum2_val += x_val * x_val;
if constexpr (NormTraits<M>::has_mean) {
sum_val += x_val;
}
}
float mean = 0.f;
float variance = sum2_val + vec_reduce_sum(sum2_fvec);
variance /= D;
if constexpr (NormTraits<M>::has_mean) {
sum_val += vec_reduce_sum(sum_fvec);
mean = sum_val / D;
variance -= mean * mean;
}
float rsqrt_var = float(1) / std::sqrt(variance + params.eps);
const fVec mean_fvec = fVec(mean);
const fVec scale_fvec = fVec(rsqrt_var);
const fVec shift_fvec = fVec(params.shift);
#pragma GCC unroll 4
for (d = 0; d <= D - kVecSize; d += kVecSize) {
auto [x_fvec0, x_fvec1] = load_float_vec2(input + d);
if constexpr (has_residual) {
auto [r_fvec0, r_fvec1] = load_float_vec2(residual + d);
x_fvec0 += r_fvec0;
x_fvec1 += r_fvec1;
convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1).store(residual + d);
}
if constexpr (NormTraits<M>::has_mean) {
x_fvec0 = x_fvec0 - mean_fvec;
x_fvec1 = x_fvec1 - mean_fvec;
}
x_fvec0 = x_fvec0 * scale_fvec;
x_fvec1 = x_fvec1 * scale_fvec;
if constexpr (NormTraits<M>::has_weight) {
auto [w_fvec0, w_fvec1] = load_float_vec2(static_cast<const scalar_t*>(params.weight) + d);
if constexpr (NormTraits<M>::has_shift) {
w_fvec0 = NormTraits<M>::apply_shift(w_fvec0, shift_fvec);
w_fvec1 = NormTraits<M>::apply_shift(w_fvec1, shift_fvec);
}
x_fvec0 = NormTraits<M>::apply_weight(x_fvec0, w_fvec0);
x_fvec1 = NormTraits<M>::apply_weight(x_fvec1, w_fvec1);
}
if constexpr (NormTraits<M>::has_bias) {
if (use_bias) {
auto [b_fvec0, b_fvec1] = load_float_vec2(static_cast<const scalar_t*>(params.bias) + d);
x_fvec0 = NormTraits<M>::apply_bias(x_fvec0, b_fvec0);
x_fvec1 = NormTraits<M>::apply_bias(x_fvec1, b_fvec1);
}
}
if constexpr (NormTraits<M>::has_gate) {
auto [g_fvec0, g_fvec1] = load_float_vec2(static_cast<const scalar_t*>(gate) + d);
x_fvec0 = NormTraits<M>::apply_gate(x_fvec0, g_fvec0);
x_fvec1 = NormTraits<M>::apply_gate(x_fvec1, g_fvec1);
}
bVec out_bvec = convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1);
out_bvec.store(out + d);
}
#pragma GCC unroll 4
for (; d < D; ++d) {
float x_val = static_cast<float>(input[d]);
if constexpr (has_residual) {
x_val += static_cast<float>(residual[d]);
residual[d] = static_cast<scalar_t>(x_val);
}
if constexpr (NormTraits<M>::has_mean) {
x_val -= mean;
}
x_val *= rsqrt_var;
if constexpr (NormTraits<M>::has_weight) {
float w_val = static_cast<float>(static_cast<const scalar_t*>(params.weight)[d]);
if constexpr (NormTraits<M>::has_shift) {
w_val = NormTraits<M>::apply_shift(w_val, params.shift);
}
x_val = NormTraits<M>::apply_weight(x_val, w_val);
}
if constexpr (NormTraits<M>::has_bias) {
if (use_bias) {
float b_val = static_cast<float>(static_cast<const scalar_t*>(params.bias)[d]);
x_val = NormTraits<M>::apply_bias(x_val, b_val);
}
}
if constexpr (NormTraits<M>::has_gate) {
float g_val = static_cast<float>(static_cast<const scalar_t*>(gate)[d]);
x_val = NormTraits<M>::apply_gate(x_val, g_val);
}
out[d] = static_cast<scalar_t>(x_val);
}
}
};
// TODO: add generic avx512-bf16 path here
#define LAUNCH_PARALLEL_LOOP(...) \
at::parallel_for(0, p.rows(), 0, [&](int64_t begin, int64_t end) { \
int64_t b{0}, h{0}, t{0}; \
data_index_init(begin, b, p.B, h, p.H, t, p.T); \
for (int64_t i = begin; i < end; ++i) { \
__VA_ARGS__; \
data_index_step(b, p.B, h, p.H, t, p.T); \
} \
})
#define LAUNCH_PARALLEL_LOOP_HD(DIM) \
case DIM: \
LAUNCH_PARALLEL_LOOP( \
const scalar_t* __restrict__ gate_ptr{nullptr}; if constexpr (NormTraits<M>::has_gate) { \
gate_ptr = gate + p.output_offset(b, h, t); \
} NormReduce<M, scalar_t, DIM>:: \
apply(out + p.output_offset(b, h, t), input + p.input_offset(b, h, t), gate_ptr, p)); \
return
template <NormMode M, typename scalar_t>
void norm4d_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
const NormParams& p,
const scalar_t* __restrict__ gate = nullptr) {
#if defined(CPU_CAPABILITY_AVX512)
// fast path only applies to bfloat16 when D in {32, 64, 128, 256, 512}
if constexpr (std::is_same_v<scalar_t, at::BFloat16>) {
switch (p.D) {
LAUNCH_PARALLEL_LOOP_HD(32);
LAUNCH_PARALLEL_LOOP_HD(64);
LAUNCH_PARALLEL_LOOP_HD(128);
LAUNCH_PARALLEL_LOOP_HD(256);
LAUNCH_PARALLEL_LOOP_HD(512);
default:
break;
}
}
#endif
// generic path
LAUNCH_PARALLEL_LOOP(
const scalar_t* __restrict__ gate_ptr{nullptr}; if constexpr (NormTraits<M>::has_gate) {
gate_ptr = gate + p.output_offset(b, h, t);
} NormReduceGeneric<M, scalar_t, false>::
apply(out + p.output_offset(b, h, t), input + p.input_offset(b, h, t), gate_ptr, nullptr, p, p.D));
}
template <NormMode M, typename scalar_t>
void fused_add_norm4d_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
scalar_t* __restrict__ residual,
const NormParams& p,
bool output_uses_input_stride = false) {
LAUNCH_PARALLEL_LOOP(
const int64_t out_offset = output_uses_input_stride ? p.input_offset(b, h, t) : p.output_offset(b, h, t);
scalar_t* __restrict__ residual_ptr = residual + p.output_offset(b, h, t);
NormReduceGeneric<M, scalar_t, true>::apply(
out + out_offset, input + p.input_offset(b, h, t), nullptr, residual_ptr, p, p.D));
}
template <NormMode M, typename scalar_t, bool copy_gate = false>
void fused_qk_norm4d_kernel_impl(
scalar_t* __restrict__ q_out,
scalar_t* __restrict__ k_out,
scalar_t* __restrict__ gate_out,
const scalar_t* __restrict__ q,
const scalar_t* __restrict__ k,
const NormParams& params_q,
const NormParams& params_k) {
at::parallel_for(0, params_q.B, 0, [&](int64_t begin, int64_t end) {
for (int64_t b = begin; b < end; ++b) {
for (int64_t h = 0; h < params_q.H /*num_head*/; ++h) {
const int64_t q_offset = params_q.input_offset(b, h, /*t*/ 0);
const int64_t out_offset = params_q.output_offset(b, h, /*t*/ 0);
NormReduceGeneric<M, scalar_t, false>::apply(
q_out + out_offset, q + q_offset, nullptr, nullptr, params_q, params_q.D);
if constexpr (copy_gate) {
std::memcpy(gate_out + out_offset, q + q_offset + params_q.D, params_q.D * sizeof(scalar_t));
}
}
for (int64_t h = 0; h < params_k.H /*num_head_kv*/; ++h) {
NormReduceGeneric<M, scalar_t, false>::apply(
k_out + params_k.output_offset(b, h, /*t*/ 0),
k + params_k.input_offset(b, h, /*t*/ 0),
nullptr,
nullptr,
params_k,
params_k.D);
}
}
});
}
#undef LAUNCH_PARALLEL_LOOP
#undef LAUNCH_PARALLEL_LOOP_HD
} // anonymous namespace
template <int... Dims>
inline void CHECK_INPUT_ND(const at::Tensor& tensor) {
static_assert(sizeof...(Dims) > 0);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor);
const int64_t dim = tensor.dim();
const bool dim_ok = ((dim == Dims) || ...);
TORCH_CHECK(dim_ok, "Expected input dim to match template constraints, got ", dim);
}
// input : {batch_size, hidden_size}
at::Tensor l2norm_cpu(at::Tensor& input, double eps) {
const auto st = input.scalar_type();
CHECK_INPUT_ND<2>(input);
NormParams p{input, static_cast<float>(eps)};
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "l2norm_kernel", [&] {
norm4d_kernel_impl<NormMode::L2Norm, scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), p);
});
return output;
}
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight: {hidden_size}
at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
const auto st = input.scalar_type();
CHECK_INPUT_ND<2, 3>(input);
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {input.size(-1)}, st);
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "rmsnorm_kernel", [&] {
norm4d_kernel_impl<NormMode::RMSNorm, scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), p);
});
return output;
}
// input : {batch_size, hidden_size}
// weight: {hidden_size}
at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
CHECK_INPUT_ND<2>(input);
const auto st = input.scalar_type();
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {input.size(-1)}, st);
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
p.shift = 1.f;
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "gemma_rmsnorm_kernel", [&] {
norm4d_kernel_impl<NormMode::GemmaNorm, scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), p);
});
return output;
}
// input : {batch_size, hidden_size} or {batch_size, num_head, seq_len, head_dim}
// weight: {hidden_size}
at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
const auto st = input.scalar_type();
CHECK_INPUT_ND<2, 4>(input);
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {input.size(-1)}, st);
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
p.shift = 1.f;
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "gemma3_rmsnorm_kernel", [&] {
norm4d_kernel_impl<NormMode::GemmaNorm, scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), p);
});
return output;
}
// Gemma4RMSNorm: with_scale ? norm(x) * (weight + scale_shift) : norm(x)
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight: {hidden_size}
at::Tensor gemma4_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps, double scale_shift, bool with_scale) {
const auto st = input.scalar_type();
CHECK_INPUT_ND<2, 3>(input);
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {input.size(-1)}, st);
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
p.shift = static_cast<float>(scale_shift);
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "gemma4_rmsnorm_kernel", [&] {
if (with_scale) {
norm4d_kernel_impl<NormMode::GemmaNorm, scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), p);
} else {
norm4d_kernel_impl<NormMode::L2Norm, scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), p);
}
});
return output;
}
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight: {hidden_size}
// bias : {hidden_size}
at::Tensor
layernorm_cpu(const at::Tensor& input, const at::Tensor& weight, const std::optional<at::Tensor>& bias, double eps) {
const auto st = input.scalar_type();
const int64_t hidden_size = input.size(-1);
CHECK_INPUT_ND<2, 3>(input);
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {hidden_size}, st);
if (bias.has_value()) {
CHECK_INPUT_SHAPE_DTYPE<false>(bias.value(), {hidden_size}, st);
}
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
p.bias = bias.has_value() ? bias.value().data_ptr() : nullptr;
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "layernorm_kernel", [&] {
norm4d_kernel_impl<NormMode::LayerNorm, scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), p);
});
return output;
}
// input : {batch_size, hidden_size}
// weight: {hidden_size}
// gate: {batch_size, hidden_size}
at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Tensor& gate, double eps) {
const auto st = input.scalar_type();
const int64_t batch_size = input.size(0);
const int64_t hidden_size = input.size(-1);
CHECK_INPUT_ND<2>(input);
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {hidden_size}, st);
CHECK_INPUT_SHAPE_DTYPE<false>(gate, {batch_size, hidden_size}, st);
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_rmsnorm_gated_kernel", [&] {
norm4d_kernel_impl<NormMode::RMSNormGated, scalar_t>(
output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), p, gate.data_ptr<scalar_t>());
});
return output;
}
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// residual: {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight : {hidden_size}
void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) {
const auto st = input.scalar_type();
CHECK_INPUT_ND<2, 3>(input);
CHECK_EQ(input.sizes(), residual.sizes());
CHECK_EQ(st, residual.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {input.size(-1)}, st);
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_add_rmsnorm_kernel", [&] {
fused_add_norm4d_kernel_impl<NormMode::RMSNorm, scalar_t>(
input.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
residual.data_ptr<scalar_t>(),
p,
/*output_uses_input_stride=*/true);
});
}
// input : {batch_size, hidden_size}
// residual: {batch_size, hidden_size}
// weight : {hidden_size}
void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) {
const auto st = input.scalar_type();
CHECK_INPUT_ND<2>(input);
CHECK_EQ(input.sizes(), residual.sizes());
CHECK_EQ(st, residual.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {input.size(-1)}, st);
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
p.shift = 1.f;
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "gemma_fused_add_rmsnorm_kernel", [&] {
fused_add_norm4d_kernel_impl<NormMode::GemmaNorm, scalar_t>(
input.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
residual.data_ptr<scalar_t>(),
p,
/*output_uses_input_stride=*/true);
});
}
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// residual: {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight : {hidden_size}
// bias : {hidden_size}
at::Tensor fused_add_layernorm_cpu(
const at::Tensor& input,
at::Tensor& residual,
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
double eps) {
const auto st = input.scalar_type();
const int64_t hidden_size = input.size(-1);
CHECK_INPUT_ND<2, 3>(input);
CHECK_EQ(input.sizes(), residual.sizes());
CHECK_EQ(st, residual.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<false>(weight, {hidden_size}, st);
if (bias.has_value()) {
CHECK_INPUT_SHAPE_DTYPE<false>(bias.value(), {hidden_size}, st);
}
NormParams p{input, static_cast<float>(eps)};
p.weight = weight.data_ptr();
p.bias = bias.has_value() ? bias.value().data_ptr() : nullptr;
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_add_layernorm_kernel", [&] {
fused_add_norm4d_kernel_impl<NormMode::LayerNorm, scalar_t>(
output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), residual.data_ptr<scalar_t>(), p);
});
return output;
}
// q : {batch_size, num_head * head_dim} 2D
// k : {batch_size, num_head_kv * head_dim} 2D
std::tuple<at::Tensor, at::Tensor> fused_qk_gemma_rmsnorm_cpu(
const at::Tensor& q,
const at::Tensor& k,
const at::Tensor& q_weight,
const at::Tensor& k_weight,
double eps,
int64_t head_dim) {
const auto st = q.scalar_type();
CHECK_INPUT_ND<2>(q);
CHECK_INPUT_ND<2>(k);
int64_t batch_size = q.size(0);
int64_t num_head = q.size(1) / head_dim;
int64_t num_head_kv = k.size(1) / head_dim;
CHECK_EQ(k.size(0), batch_size);
CHECK_EQ(k.scalar_type(), st);
CHECK_INPUT_SHAPE_DTYPE<false>(q_weight, {head_dim}, st);
CHECK_INPUT_SHAPE_DTYPE<false>(k_weight, {head_dim}, st);
NormParams q_params{q, static_cast<float>(eps)};
q_params.H = num_head;
q_params.D = head_dim;
q_params.i_strideH = head_dim;
q_params.weight = q_weight.data_ptr();
q_params.shift = 1.f;
NormParams k_params{k, static_cast<float>(eps)};
k_params.H = num_head_kv;
k_params.D = head_dim;
k_params.i_strideH = head_dim;
k_params.weight = k_weight.data_ptr();
k_params.shift = 1.f;
at::Tensor q_out = at::empty_like(q);
at::Tensor k_out = at::empty_like(k);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_qk_gemma_rmsnorm_kernel", [&] {
fused_qk_norm4d_kernel_impl<NormMode::GemmaNorm, scalar_t, false>(
q_out.data_ptr<scalar_t>(),
k_out.data_ptr<scalar_t>(),
nullptr,
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
q_params,
k_params);
});
return std::make_tuple(q_out, k_out);
}
// q_gate : {batch_size, num_head * head_dim * 2} 2D, interleaved per head as [q_h, gate_h]
// k : {batch_size, num_head_kv * head_dim} 2D
std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qk_gemma_rmsnorm_with_gate_cpu(
const at::Tensor& q_gate,
const at::Tensor& k,
const at::Tensor& q_weight,
const at::Tensor& k_weight,
double eps,
int64_t head_dim,
int64_t num_head) {
const auto st = q_gate.scalar_type();
CHECK_INPUT_ND<2>(q_gate);
CHECK_INPUT_ND<2>(k);
int64_t batch_size = q_gate.size(0);
int64_t num_head_kv = k.size(1) / head_dim;
CHECK_EQ(q_gate.size(1), num_head * head_dim * 2);
CHECK_EQ(k.size(0), batch_size);
CHECK_EQ(k.scalar_type(), st);
CHECK_INPUT_SHAPE_DTYPE<false>(q_weight, {head_dim}, st);
CHECK_INPUT_SHAPE_DTYPE<false>(k_weight, {head_dim}, st);
NormParams q_params{q_gate, static_cast<float>(eps)};
q_params.H = num_head;
q_params.D = head_dim;
q_params.i_strideH = head_dim * 2;
q_params.weight = q_weight.data_ptr();
q_params.shift = 1.f;
NormParams k_params{k, static_cast<float>(eps)};
k_params.H = num_head_kv;
k_params.D = head_dim;
k_params.i_strideH = head_dim;
k_params.weight = k_weight.data_ptr();
k_params.shift = 1.f;
at::Tensor q_out = at::empty({batch_size * num_head, head_dim}, q_gate.options());
at::Tensor k_out = at::empty({batch_size * num_head_kv, head_dim}, k.options());
at::Tensor gate_out = at::empty_like(q_out);
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_qk_gemma_rmsnorm_with_gate_kernel", [&] {
fused_qk_norm4d_kernel_impl<NormMode::GemmaNorm, scalar_t, true>(
q_out.data_ptr<scalar_t>(),
k_out.data_ptr<scalar_t>(),
gate_out.data_ptr<scalar_t>(),
q_gate.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
q_params,
k_params);
});
return std::make_tuple(q_out, k_out, gate_out);
}
python/sglang/kernels/aot/csrc/cpu/numa_utils.cpp
+98
View File
@@ -0,0 +1,98 @@
#include <numa.h>
#include <sched.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <unistd.h>
#include <string>
#include "common.h"
std::string init_cpu_threads_env(const std::string& cpu_ids) {
bitmask* omp_cpu_mask = numa_parse_cpustring(cpu_ids.c_str());
TORCH_CHECK(omp_cpu_mask->size > 0);
std::vector<int> omp_cpu_ids;
omp_cpu_ids.reserve(omp_cpu_mask->size);
constexpr int group_size = 8 * sizeof(*omp_cpu_mask->maskp);
for (int offset = 0; offset < omp_cpu_mask->size; offset += group_size) {
unsigned long group_mask = omp_cpu_mask->maskp[offset / group_size];
int i = 0;
while (group_mask) {
if (group_mask & 1) {
omp_cpu_ids.emplace_back(offset + i);
}
++i;
group_mask >>= 1;
}
}
// Memory node binding
if (numa_available() != -1) {
TORCH_CHECK(!omp_cpu_ids.empty(), "Cannot bind memory, no CPUs specified.");
int mem_node_id_st = numa_node_of_cpu(omp_cpu_ids.front());
int mem_node_id_ed = numa_node_of_cpu(omp_cpu_ids.back());
if (mem_node_id_st > mem_node_id_ed) {
std::swap(mem_node_id_st, mem_node_id_ed);
}
bitmask* mask =
numa_parse_nodestring((std::to_string(mem_node_id_st) + "-" + std::to_string(mem_node_id_ed)).c_str());
bitmask* src_mask = numa_get_membind();
int pid = getpid();
// move all existing pages to the specified numa node.
*(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
int page_num = numa_migrate_pages(pid, src_mask, mask);
if (page_num == -1) {
TORCH_WARN(false, "numa_migrate_pages failed. errno: " + std::to_string(errno));
}
// restrict memory allocation node.
numa_set_membind(mask);
numa_set_strict(1);
}
// OMP threads binding
omp_set_num_threads((int)omp_cpu_ids.size());
at::set_num_threads((int)omp_cpu_ids.size());
TORCH_CHECK_EQ(omp_cpu_ids.size(), at::get_num_threads());
TORCH_CHECK_EQ(omp_cpu_ids.size(), omp_get_max_threads());
std::vector<std::pair<int, int>> thread_core_mapping;
thread_core_mapping.reserve(omp_cpu_ids.size());
omp_lock_t writelock;
omp_init_lock(&writelock);
#pragma omp parallel for schedule(static, 1)
for (size_t i = 0; i < omp_cpu_ids.size(); ++i) {
cpu_set_t mask;
CPU_ZERO(&mask);
CPU_SET(omp_cpu_ids[i], &mask);
int ret = sched_setaffinity(0, sizeof(cpu_set_t), &mask);
if (ret == -1) {
TORCH_CHECK(false, "sched_setaffinity failed. errno: " + std::to_string(errno));
}
omp_set_lock(&writelock);
thread_core_mapping.emplace_back(syscall(SYS_gettid), omp_cpu_ids[i]);
omp_unset_lock(&writelock);
}
omp_destroy_lock(&writelock);
numa_free_nodemask(omp_cpu_mask);
std::stringstream ss;
ss << "OMP threads binding of Process " << getpid() << ":\n";
std::sort(
thread_core_mapping.begin(), thread_core_mapping.end(), [](auto&& a, auto&& b) { return a.second < b.second; });
for (auto&& item : thread_core_mapping) {
ss << "\t"
<< "OMP tid: " << item.first << ", core " << item.second << "\n";
}
return ss.str();
}
python/sglang/kernels/aot/csrc/cpu/preprocessor.cpp
+371
View File
@@ -0,0 +1,371 @@
/*****************************************************************************************
* Copyright (c) 2025 - 2025 Codeplay Software Ltd. All rights reserved.
* Copyright (C) 2025 Intel Corporation, All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
****************************************************************************************/
#include "common.h"
#include "vec.h"
// [NOTE] Preprocessor Optimization
// 1. this file is apple-to-apple to `Qwen2VLImageProcessorFast`.
// 2. `out_dtype` set to torch.bfloat16 skips outplace dtype conversion.
// 3. skip all redundant memory copy and dtype conversion.
// 4. TODO: rewrite `_upsample_bicubic2d_aa`.
//
// ref: https://github.com/huggingface/transformers/blob/main/src/transformers
// /models/qwen2_vl/image_processing_qwen2_vl_fast.py
//
namespace {
template <typename scalar_t>
inline void normalize(
scalar_t* __restrict__ out,
const uint8_t* __restrict__ input,
const std::vector<float>& image_mean,
const std::vector<float>& image_std,
int64_t channel,
int64_t temporal_patch_size,
int64_t patch_size,
int64_t stride_ch,
int64_t stride_pt,
int64_t stride_ph) {
TORCH_CHECK(false, "normalize: scalar path not implemented.");
}
#if defined(CPU_CAPABILITY_AVX512)
template <>
inline void normalize<float>(
float* __restrict__ out,
const uint8_t* __restrict__ input,
const std::vector<float>& image_mean,
const std::vector<float>& image_std,
int64_t channel,
int64_t temporal_patch_size,
int64_t patch_size,
int64_t stride_ch,
int64_t stride_pt,
int64_t stride_ph) {
// we do vectorization on patch_size dim
assert(patch_size == 16);
// loop last 4 dimensions:
// {channel, patch_t(repeated), patch_h, patch_w}
for (int64_t c = 0; c < channel; ++c) {
__m512 vmean = _mm512_set1_ps(image_mean[c]);
__m512 vrstd = _mm512_set1_ps(1.f / image_std[c]);
float* __restrict__ out_ptr = out + c * temporal_patch_size * patch_size * patch_size;
#pragma GCC unroll 4
for (int64_t ph = 0; ph < patch_size; ++ph) {
__m128i u8 = _mm_loadu_si128((const __m128i*)(input + c * stride_ch + /* pt */ 0 * stride_pt + ph * stride_ph));
__m512 x = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(u8));
x = _mm512_mul_ps(_mm512_sub_ps(x, vmean), vrstd);
#pragma GCC unroll 2
for (int64_t pt = 0; pt < temporal_patch_size; ++pt) {
_mm512_storeu_ps(out_ptr + pt * patch_size * patch_size + ph * patch_size, x);
}
}
}
}
template <>
inline void normalize<at::BFloat16>(
at::BFloat16* __restrict__ out,
const uint8_t* __restrict__ input,
const std::vector<float>& image_mean,
const std::vector<float>& image_std,
int64_t channel,
int64_t temporal_patch_size,
int64_t patch_size,
int64_t stride_ch,
int64_t stride_pt,
int64_t stride_ph) {
// we do vectorization on patch_size dim
assert(patch_size == 16);
// loop last 4 dimensions:
// {channel, patch_t(repeated), patch_h, patch_w}
for (int64_t c = 0; c < channel; ++c) {
__m512 vmean = _mm512_set1_ps(image_mean[c]);
__m512 vrstd = _mm512_set1_ps(1.f / image_std[c]);
at::BFloat16* __restrict__ out_ptr = out + c * temporal_patch_size * patch_size * patch_size;
#pragma GCC unroll 4
for (int64_t ph = 0; ph < patch_size; ++ph) {
__m128i u8 = _mm_loadu_si128((const __m128i*)(input + c * stride_ch + /* pt */ 0 * stride_pt + ph * stride_ph));
__m512 x = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(u8));
x = _mm512_mul_ps(_mm512_sub_ps(x, vmean), vrstd);
__m256i x16 = (__m256i)_mm512_cvtneps_pbh(x);
#pragma GCC unroll 2
for (int64_t pt = 0; pt < temporal_patch_size; ++pt) {
_mm256_storeu_si256(reinterpret_cast<__m256i*>(out_ptr + pt * patch_size * patch_size + ph * patch_size), x16);
}
}
}
}
#endif
template <typename scalar_t>
void rescale_and_normalize_kernel_impl(
scalar_t* __restrict__ out,
const uint8_t* __restrict__ input,
const std::vector<float>& image_mean,
const std::vector<float>& image_std,
int64_t grid_t,
int64_t grid_h,
int64_t grid_w,
int64_t merge_size,
int64_t channel,
int64_t temporal_patch_size,
int64_t patch_size) {
// [NOTE]: temporal patching uses repeat on last image
//
// input : {grid_t, patch_t, channel, grid_h, merge_h, patch_h, grid_w, merge_w, patch_w}
// out : {grid_t, grid_h, grid_w, merge_h, merge_w, channel, patch_t, patch_h, patch_w}
//
int64_t height = grid_h * merge_size * patch_size;
int64_t width = grid_w * merge_size * patch_size;
int64_t stride_gt = /* temporal_patch_size */ 1 * channel * height * width;
int64_t stride_gh = merge_size * patch_size * width;
int64_t stride_gw = merge_size * patch_size;
int64_t stride_mh = patch_size * width;
int64_t stride_mw = patch_size;
int64_t stride_ch = height * width;
int64_t stride_pt = channel * height * width;
int64_t stride_ph = width;
int64_t stride_grid = channel * temporal_patch_size * patch_size * patch_size;
// parallel on first 5 dims, aka, grids
at::parallel_for(0, grid_t * grid_h * grid_w * merge_size * merge_size, 0, [&](int64_t begin, int64_t end) {
int64_t gt{0}, gh{0}, gw{0}, mh{0}, mw{0};
data_index_init(begin, gt, grid_t, gh, grid_h, gw, grid_w, mh, merge_size, mw, merge_size);
for (int64_t i = begin; i < end; ++i) {
normalize<scalar_t>(
out + i * stride_grid,
input + gt * stride_gt + gh * stride_gh + gw * stride_gw + mh * stride_mh + mw * stride_mw,
image_mean,
image_std,
channel,
temporal_patch_size,
patch_size,
stride_ch,
stride_pt,
stride_ph);
// move to the next index
data_index_step(gt, grid_t, gh, grid_h, gw, grid_w, mh, merge_size, mw, merge_size);
}
});
}
} // anonymous namespace
void check_input_image(const at::Tensor& image) {
TORCH_CHECK(image.scalar_type() == at::kByte, "expect image to be uint8.");
TORCH_CHECK(image.dim() == 3, "expect image to be CHW.");
}
// https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py
std::pair<int64_t, int64_t>
smart_resize(int64_t height, int64_t width, int64_t factor, int64_t min_pixels, int64_t max_pixels) {
// aspect ratio check
int64_t mx = std::max(height, width);
int64_t mn = std::min(height, width);
TORCH_CHECK(static_cast<double>(mx) / mn <= 200.0, "absolute aspect ratio must be smaller than 200");
// round to nearest multiple of factor
auto round_to_factor = [&](int64_t x) {
return static_cast<int64_t>(std::round(static_cast<double>(x) / factor)) * factor;
};
int64_t h_bar = round_to_factor(height);
int64_t w_bar = round_to_factor(width);
int64_t area = h_bar * w_bar;
if (area > max_pixels) {
double beta = std::sqrt((1.0 * height * width) / max_pixels);
h_bar = std::max(factor, (static_cast<int64_t>(std::floor(height / beta / factor)) * factor));
w_bar = std::max(factor, (static_cast<int64_t>(std::floor(width / beta / factor)) * factor));
} else if (area < min_pixels) {
double beta = std::sqrt((double)min_pixels / (height * width));
h_bar = static_cast<int64_t>(std::ceil(height * beta / factor)) * factor;
w_bar = static_cast<int64_t>(std::ceil(width * beta / factor)) * factor;
}
return {h_bar, w_bar};
}
// do rescale and normalize
// from `resized_image` to `pixel_values`
void rescale_and_normalize_image(
at::Tensor& pixel_values,
const at::Tensor& image,
double rescale_factor,
c10::ArrayRef<double> image_mean,
c10::ArrayRef<double> image_std,
int64_t grid_t,
int64_t grid_h,
int64_t grid_w,
int64_t merge_size,
int64_t channel,
int64_t temporal_patch_size,
int64_t patch_size,
int64_t grid_offset,
int64_t grid_stride) {
// update mean and std
std::vector<float> mean_vec(channel), std_vec(channel);
for (int64_t c = 0; c < channel; ++c) {
mean_vec[c] = static_cast<float>(image_mean[c] * (1 / rescale_factor));
std_vec[c] = static_cast<float>(image_std[c] * (1 / rescale_factor));
}
AT_DISPATCH_FLOATING_TYPES_AND(at::kBFloat16, pixel_values.scalar_type(), "rescale_and_normalize_image", [&] {
rescale_and_normalize_kernel_impl<scalar_t>(
pixel_values.data_ptr<scalar_t>() + grid_offset * grid_stride,
image.data_ptr<uint8_t>(),
mean_vec,
std_vec,
grid_t,
grid_h / merge_size,
grid_w / merge_size,
merge_size,
channel,
temporal_patch_size,
patch_size);
});
}
std::tuple<at::Tensor, at::Tensor> image_preprocess_cpu(
at::TensorList images,
bool do_convert_rgb,
bool do_resize,
int64_t shortest_edge,
int64_t longest_edge,
const std::string& interpolation,
bool do_rescale,
double rescale_factor,
bool do_normalize,
c10::ArrayRef<double> image_mean,
c10::ArrayRef<double> image_std,
int64_t patch_size,
int64_t temporal_patch_size,
int64_t merge_size,
bool disable_grouping,
at::ScalarType out_dtype) {
// TODO: lift C++ kernel limitations
TORCH_CHECK(interpolation == "bicubic", "image_preprocess_cpu: support only bicubic mode.");
TORCH_CHECK(do_rescale && do_normalize, "image_preprocess_cpu: support only do_rescale and do_normalize.");
TORCH_CHECK(disable_grouping, "image_preprocess_cpu: support only disable_grouping.");
// support only float32 or bfloat16 as output
TORCH_CHECK(
out_dtype == at::kFloat || out_dtype == at::kBFloat16,
"image_preprocess_cpu: support only float32 and bfloat16 as pixel_values dtype.");
int64_t batch_size = images.size();
int64_t channel = image_mean.size();
CHECK_GT(batch_size, 0);
CHECK_EQ(channel, image_std.size());
CHECK_EQ(channel, 3);
const at::Tensor& first_image = images[0];
const auto options = first_image.options();
at::Tensor pixel_values = at::empty({}, options.dtype(out_dtype));
at::Tensor image_grid_thw = at::empty({batch_size, channel}, options.dtype(at::kLong));
// index type use int64_t
int64_t* image_grid_thw_data = image_grid_thw.data_ptr<int64_t>();
// resized image sizes and global grid offset
std::vector<std::pair<int64_t, int64_t>> image_sizes(batch_size);
std::vector<int64_t> grid_offsets(batch_size + 1, 0);
// Stage 1: compute resized shapes and fill in `image_grid_thw`
for (int64_t idx = 0; idx < batch_size; ++idx) {
const auto& image = images[idx];
check_input_image(image);
auto [resized_h, resized_w] =
smart_resize(image.size(-2), image.size(-1), patch_size * merge_size, shortest_edge, longest_edge);
image_sizes[idx] = {resized_h, resized_w};
// temporal dimension for image is 1
int64_t grid_t = div_up((int64_t)1, temporal_patch_size);
int64_t grid_h = div_up(resized_h, patch_size);
int64_t grid_w = div_up(resized_w, patch_size);
// fill in image_grid_thw
image_grid_thw_data[idx * 3 + 0] = grid_t;
image_grid_thw_data[idx * 3 + 1] = grid_h;
image_grid_thw_data[idx * 3 + 2] = grid_w;
// fill in global grid offset
grid_offsets[idx + 1] = grid_offsets[idx] + grid_t * grid_h * grid_w;
}
// last element holds the total sum of grids
int64_t grid_size = grid_offsets[batch_size];
int64_t grid_stride = channel * temporal_patch_size * patch_size * patch_size;
// allocate memory
pixel_values.resize_({grid_size, grid_stride});
// Stage 2: compute `pixel_values`
for (int64_t idx = 0; idx < batch_size; ++idx) {
const auto& image = images[idx];
int64_t resized_h = image_sizes[idx].first;
int64_t resized_w = image_sizes[idx].second;
auto resized_image = at::_upsample_bicubic2d_aa(
image.unsqueeze(0),
{resized_h, resized_w},
/* align_corners */ false);
rescale_and_normalize_image(
pixel_values,
resized_image,
rescale_factor,
image_mean,
image_std,
/* grid_t */ image_grid_thw_data[idx * 3 + 0],
/* grid_h */ image_grid_thw_data[idx * 3 + 1],
/* grid_w */ image_grid_thw_data[idx * 3 + 2],
merge_size,
channel,
temporal_patch_size,
patch_size,
grid_offsets[idx],
grid_stride);
}
return std::make_tuple(pixel_values, image_grid_thw);
}
python/sglang/kernels/aot/csrc/cpu/qkv_proj.cpp
+690
View File
@@ -0,0 +1,690 @@
#include "common.h"
#include "gemm.h"
#include "vec.h"
namespace {
// [NOTE]: Fused kernel for QKV projection with weight absorption and RoPE
//
// 1. `q_a_proj` and `kv_a_proj_with_mqa` fused into one gemm,
// otherwise we need to split IC for the 2nd gemm.
// 2. `q_a_layernorm` and `kv_a_layernorm` fused into one parallel loop.
// 3. k_input and v_input share the same storage, the torch API did
// this in `set_kv_buffer`. No additional memory movement.
//
// [C0, C1] = A @ [B0, B1]
template <typename scalar_t>
void segment_gemm_kernel_impl(
scalar_t* __restrict__ C0,
scalar_t* __restrict__ C1,
const scalar_t* __restrict__ A,
const scalar_t* __restrict__ B0,
const scalar_t* __restrict__ B1,
int64_t M,
int64_t N0,
int64_t N1,
int64_t K) {
// convert_weight_packed make sure N0 and N1 are 32x
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB0 = div_up(N0, BLOCK_N);
const int64_t NB1 = div_up(N1, BLOCK_N);
const int64_t NB = NB0 + NB1;
const bool use_brgemm = can_use_brgemm<scalar_t>(M);
// parallel on [MB, NB0 + NB1]
at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
int64_t mb{0}, nb{0};
data_index_init(begin, mb, MB, nb, NB);
// for brgemm, use float32 for accumulate
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
for (int64_t i = begin; i < end; ++i) {
UNUSED(i);
int mb_start = mb * BLOCK_M;
int mb_size = std::min(M - mb_start, BLOCK_M);
int nb_start = nb * BLOCK_N;
int nb_size = BLOCK_N;
const scalar_t* __restrict__ B = nb < NB0 ? B0 : B1;
scalar_t* __restrict__ C = nb < NB0 ? C0 : C1;
int64_t ldc = nb < NB0 ? N0 : N1;
int64_t local_nb_start = nb < NB0 ? nb_start : nb_start - N0;
tinygemm_kernel<scalar_t>(
/* A */ A + mb_start * K,
/* B */ B + local_nb_start * K /* nb * BLOCK_N * K */,
/* C */ C + mb_start * ldc + local_nb_start,
/* Ctmp*/ Ctmp,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ K,
/* ldb */ nb_size,
/* ldc */ ldc,
/* brg */ use_brgemm);
// move to the next index
data_index_step(mb, MB, nb, NB);
}
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
// [C0, C1] = A @ [B0, B1]
template <typename scalar_t>
void segment_gemm_kernel_impl(
scalar_t* __restrict__ C0,
scalar_t* __restrict__ C1,
const uint8_t* __restrict__ A,
const int8_t* __restrict__ B0,
const int8_t* __restrict__ B1,
const float* __restrict__ As,
const float* __restrict__ Bs0,
const float* __restrict__ Bs1,
int64_t M,
int64_t N0,
int64_t N1,
int64_t K) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB0 = div_up(N0, BLOCK_N);
const int64_t NB1 = div_up(N1, BLOCK_N);
const int64_t NB = NB0 + NB1;
const bool use_brgemm = can_use_brgemm<int8_t>(M);
// K + 4 after compensation
const int64_t packed_row_size = get_row_size<int8_t>(K);
// parallel on [MB, NB0 + NB1]
at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
int64_t mb{0}, nb{0};
data_index_init(begin, mb, MB, nb, NB);
// for brgemm, use float32 for accumulate
alignas(64) int32_t Ctmp[BLOCK_M * BLOCK_N];
for (int64_t i = begin; i < end; ++i) {
UNUSED(i);
int mb_start = mb * BLOCK_M;
int mb_size = std::min(M - mb_start, BLOCK_M);
int nb_start = nb * BLOCK_N;
int nb_size = BLOCK_N;
const int8_t* __restrict__ B = nb < NB0 ? B0 : B1;
const float* __restrict__ Bs = nb < NB0 ? Bs0 : Bs1;
scalar_t* __restrict__ C = nb < NB0 ? C0 : C1;
int64_t ldc = nb < NB0 ? N0 : N1;
int64_t local_nb_start = nb < NB0 ? nb_start : nb_start - N0;
tinygemm_kernel<scalar_t>(
/* A */ A + mb_start * K,
/* B */ B + local_nb_start * packed_row_size /* nb * BLOCK_N * (K + 4) */,
/* C */ C + mb_start * ldc + local_nb_start,
/* Ctmp*/ Ctmp,
/* As */ As + mb_start,
/* Bs */ Bs + local_nb_start,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ K,
/* ldb */ nb_size,
/* ldc */ ldc,
/* brg */ use_brgemm);
// move to the next index
data_index_step(mb, MB, nb, NB);
}
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
// [C0, C1] = A @ [B0, B1]
template <typename scalar_t>
void segment_gemm_kernel_impl(
scalar_t* __restrict__ C0,
scalar_t* __restrict__ C1,
const scalar_t* __restrict__ A,
const at::Float8_e4m3fn* __restrict__ B0,
const at::Float8_e4m3fn* __restrict__ B1,
const float* __restrict__ Bs0,
const float* __restrict__ Bs1,
scalar_t* __restrict__ Btmp,
int64_t M,
int64_t N0,
int64_t N1,
int64_t K,
int64_t block_size_N,
int64_t block_size_K) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB0 = div_up(N0, BLOCK_N);
const int64_t NB1 = div_up(N1, BLOCK_N);
const int64_t NB = NB0 + NB1;
const int64_t scale_size_K = div_up(K, block_size_K);
const int64_t blocks_n_per_group = block_size_N / BLOCK_N;
const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
// parallel on [MB, NB0 + NB1]
at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
int64_t mb{0}, nb{0};
data_index_init(begin, mb, MB, nb, NB);
int tid = at::get_thread_num();
// for brgemm, use float32 for accumulate
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
for (int64_t i = begin; i < end; ++i) {
UNUSED(i);
int mb_start = mb * BLOCK_M;
int mb_size = std::min(M - mb_start, BLOCK_M);
int nb_start = nb * BLOCK_N;
int nb_size = BLOCK_N;
const at::Float8_e4m3fn* __restrict__ B = nb < NB0 ? B0 : B1;
const float* __restrict__ Bs = nb < NB0 ? Bs0 : Bs1;
scalar_t* __restrict__ C = nb < NB0 ? C0 : C1;
int64_t ldc = nb < NB0 ? N0 : N1;
int64_t local_nb_start = nb < NB0 ? nb_start : nb_start - N0;
int64_t new_nb = nb < NB0 ? nb : nb - NB0;
tinygemm_kernel<scalar_t>(
/* A */ A + mb_start * K,
/* B */ B + local_nb_start * K /* nb * BLOCK_N * K */,
/* C */ C + mb_start * ldc + local_nb_start,
/* Btmp*/ Btmp + tid * BLOCK_N * K,
/* Ctmp*/ Ctmp,
/*Bbias*/ nullptr,
/* Bs */ Bs + (new_nb / blocks_n_per_group) * scale_size_K,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ K,
/* ldb */ nb_size,
/* ldc */ ldc,
/* brg */ use_brgemm,
/* block_size_K */ block_size_K);
// move to the next index
data_index_step(mb, MB, nb, NB);
}
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
template <typename scalar_t>
inline float reduce(const scalar_t* __restrict__ x, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
fVec sum_fvec = fVec(float(0));
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
auto [x_fvec0, x_fvec1] = load_float_vec2(x + d);
sum_fvec += x_fvec0 * x_fvec0;
sum_fvec += x_fvec1 * x_fvec1;
}
return vec_reduce_sum(sum_fvec);
}
// map2 from aten functional doesn't have fast bf16->fp32 conversion
template <typename scalar_t>
inline void map2(scalar_t* y, const scalar_t* x, const scalar_t* __restrict__ w, float scale, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
fVec scale_fvec = fVec(scale);
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
auto [x_fvec0, x_fvec1] = load_float_vec2(x + d);
auto [w_fvec0, w_fvec1] = load_float_vec2(w + d);
x_fvec0 = x_fvec0 * scale_fvec * w_fvec0;
x_fvec1 = x_fvec1 * scale_fvec * w_fvec1;
bVec out_bvec = convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1);
out_bvec.store(y + d);
}
}
template <typename scalar_t>
void rms_norm_kernel_impl(
scalar_t* __restrict__ input0,
scalar_t* __restrict__ input1,
const scalar_t* __restrict__ weight0,
const scalar_t* __restrict__ weight1,
int64_t M,
int64_t N0,
int64_t N1,
int64_t stride1,
float eps = 1e-5) {
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
scalar_t* x0 = input0 + m * N0;
scalar_t* x1 = input1 + m * stride1;
float scale0 = reduce(x0, N0);
float scale1 = reduce(x1, N1);
scale0 = float(1) / std::sqrt(scale0 / N0 + eps);
scale1 = float(1) / std::sqrt(scale1 / N1 + eps);
map2(x0, x0, weight0, scale0, N0);
map2(x1, x1, weight1, scale1, N1);
}
});
}
template <typename scalar_t>
inline void rotary(const scalar_t* input, scalar_t* out, const scalar_t* cos, const scalar_t* sin, int64_t size) {
TORCH_CHECK(false, "rotary scalar path not implemented.");
}
#if defined(CPU_CAPABILITY_AVX512)
template <>
inline void rotary<at::BFloat16>(
const at::BFloat16* input, at::BFloat16* out, const at::BFloat16* cos, const at::BFloat16* sin, int64_t size) {
// permute indices
const __m512i idx1 = _mm512_set_epi32(30, 28, 26, 24, 22, 20, 18, 16, 14, 12, 10, 8, 6, 4, 2, 0);
const __m512i idx2 = _mm512_set_epi32(31, 29, 27, 25, 23, 21, 19, 17, 15, 13, 11, 9, 7, 5, 3, 1);
const __m512i idy1 = _mm512_set_epi32(23, 7, 22, 6, 21, 5, 20, 4, 19, 3, 18, 2, 17, 1, 16, 0);
const __m512i idy2 = _mm512_set_epi32(31, 15, 30, 14, 29, 13, 28, 12, 27, 11, 26, 10, 25, 9, 24, 8);
// rotary dim is 64, just 2 iters
#pragma GCC unroll 2
for (int64_t d = 0; d < size; d += 32) {
int64_t d2 = d >> 1;
// load coefs
__m512 vcos = CVT_BF16_TO_FP32(_mm256_loadu_si256(reinterpret_cast<const __m256i*>(cos + d2)));
__m512 vsin = CVT_BF16_TO_FP32(_mm256_loadu_si256(reinterpret_cast<const __m256i*>(sin + d2)));
// load input
__m512i a16 = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(input + d));
__m512 a = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0));
__m512 b = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1));
// from [16, 2] to [2, 16]
__m512 in1 = _mm512_mask_permutex2var_ps(a, 0xffff, idx1, b);
__m512 in2 = _mm512_mask_permutex2var_ps(a, 0xffff, idx2, b);
// out1 = in1 * cos - in2 * sin;
// out2 = in2 * cos + in1 * sin
__m512 out1 = _mm512_sub_ps(_mm512_mul_ps(in1, vcos), _mm512_mul_ps(in2, vsin));
__m512 out2 = _mm512_add_ps(_mm512_mul_ps(in2, vcos), _mm512_mul_ps(in1, vsin));
// from [2, 16] to [16, 2]
a = _mm512_mask_permutex2var_ps(out1, 0xffff, idy1, out2);
b = _mm512_mask_permutex2var_ps(out1, 0xffff, idy2, out2);
_mm512_storeu_si512(reinterpret_cast<__m512i*>((out + d)), (__m512i)(_mm512_cvtne2ps_pbh(b, a)));
}
}
#endif
template <typename scalar_t>
void rotary_emb_kernel_impl(
scalar_t* q_pe_out,
scalar_t* k_pe_out,
const scalar_t* q_pe,
const scalar_t* k_pe,
const int64_t* pos,
const scalar_t* cos_sin,
int64_t num_seqs,
int64_t num_heads,
int64_t rotary_dim,
int64_t q_strideB,
int64_t q_strideH,
int64_t k_strideB,
int64_t oq_strideB,
int64_t oq_strideH,
int64_t ok_strideB) {
TORCH_CHECK(rotary_dim % 32 == 0, "rotary_dim is not 32x.");
const int64_t rotary_offset = rotary_dim / 2;
// parallel on [num_seqs, num_heads + 1]
// top [num_heads] handle q_pe and bottom [1] handle k_pe
at::parallel_for(0, num_seqs * (num_heads + 1), GRAIN_SIZE / rotary_dim, [&](int64_t begin, int64_t end) {
int64_t seq{0}, head_id{0};
data_index_init(begin, seq, num_seqs, head_id, num_heads + 1);
for (int64_t i = begin; i < end; ++i) {
UNUSED(i);
// get cos and sin cache ptr
int64_t index = pos[seq];
const scalar_t* cos = cos_sin + index * rotary_dim;
const scalar_t* sin = cos + rotary_offset;
const scalar_t* input =
(head_id < num_heads) ? q_pe + seq * q_strideB + head_id * q_strideH : k_pe + seq * k_strideB;
scalar_t* out =
(head_id < num_heads) ? q_pe_out + seq * oq_strideB + head_id * oq_strideH : k_pe_out + seq * ok_strideB;
rotary<scalar_t>(input, out, cos, sin, rotary_dim);
// move to the next index
data_index_step(seq, num_seqs, head_id, num_heads + 1);
}
});
}
} // anonymous namespace
extern at::Tensor
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni);
extern at::Tensor int8_scaled_mm_with_quant(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales2,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni);
extern void
bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale);
extern at::Tensor fp8_scaled_mm_cpu(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales2,
std::vector<int64_t> block_size,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni);
// NB: shapes in DeepDeek R1
//
// hidden_states : [num_seqs, hidden_size] [1, 7168]
// q_a_proj_weight : [q_lora_rank, hidden_size] [1536, 7168]
// q_b_proj_weight : [num_heads * qk_head_dim, q_lora_rank] [4224, 1536]
// kv_a_proj_weight : [kv_lora_rank + qk_rope_head_dim, hidden_size] [576, 7168]
// w_kc : [num_heads, kv_lora_rank, qk_nope_head_dim] [22, 512, 128]
// q_a_layernorm_weight : [q_lora_rank] [1536]
// kv_a_layernorm_weight : [kv_lora_rank] [512]
//
std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope(
at::Tensor& hidden_states,
at::Tensor& q_a_proj_weight,
at::Tensor& q_b_proj_weight,
at::Tensor& kv_a_proj_weight,
at::Tensor& w_kc,
at::Tensor& q_a_layernorm_weight,
at::Tensor& kv_a_layernorm_weight,
at::Tensor& positions,
at::Tensor& cos_sin_cache,
double eps,
bool use_int8_w8a8,
bool use_fp8_w8a16,
std::optional<at::Tensor> q_a_proj_scale,
std::optional<at::Tensor> q_b_proj_scale,
std::optional<at::Tensor> kv_a_proj_scale,
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size) {
const auto st = hidden_states.scalar_type();
CHECK_INPUT(hidden_states);
CHECK_INPUT(positions);
CHECK_INPUT(cos_sin_cache);
CHECK_EQ(q_a_layernorm_weight.scalar_type(), st);
CHECK_EQ(kv_a_layernorm_weight.scalar_type(), st);
CHECK_EQ(positions.scalar_type(), at::kLong);
CHECK_EQ(cos_sin_cache.scalar_type(), st);
CHECK_DIM(2, hidden_states);
CHECK_DIM(3, w_kc);
CHECK_DIM(1, q_a_layernorm_weight);
CHECK_DIM(1, kv_a_layernorm_weight);
CHECK_DIM(1, positions);
CHECK_DIM(2, cos_sin_cache);
// skip contiguous checks for weights, expect prepacked
TORCH_CHECK(is_vnni, "qkv_proj_with_rope: expect weights are prepacked!");
int64_t num_seqs = hidden_states.size(0);
int64_t hidden_size = hidden_states.size(1);
int64_t q_lora_rank = q_a_proj_weight.size(0);
int64_t num_heads = w_kc.size(0);
int64_t kv_lora_rank = w_kc.size(1);
int64_t qk_head_dim = q_b_proj_weight.size(0) / num_heads;
int64_t qk_nope_head_dim = w_kc.size(2);
int64_t qk_rope_head_dim = kv_a_proj_weight.size(0) - kv_lora_rank;
int64_t rotary_dim = cos_sin_cache.size(1);
CHECK_EQ(positions.numel(), num_seqs);
CHECK_EQ(rotary_dim, qk_rope_head_dim);
CHECK_EQ(q_a_layernorm_weight.numel(), q_lora_rank);
CHECK_EQ(kv_a_layernorm_weight.numel(), kv_lora_rank);
// check the packed dimension
CHECK_EQ(q_a_proj_weight.size(1), get_row_size(hidden_size, use_int8_w8a8));
CHECK_EQ(q_b_proj_weight.size(1), get_row_size(q_lora_rank, use_int8_w8a8));
CHECK_EQ(kv_a_proj_weight.size(1), get_row_size(hidden_size, use_int8_w8a8));
if (use_int8_w8a8) {
TORCH_CHECK(q_a_proj_scale.has_value(), "missing q_a_proj_scale for int8 w8a8.");
TORCH_CHECK(q_b_proj_scale.has_value(), "missing q_b_proj_scale for int8 w8a8.");
TORCH_CHECK(kv_a_proj_scale.has_value(), "missing kv_a_proj_scale for int8 w8a8.");
}
if (use_fp8_w8a16) {
TORCH_CHECK(q_a_proj_scale.has_value(), "missing q_a_proj_scale for fp8 w8a16.");
TORCH_CHECK(q_b_proj_scale.has_value(), "missing q_b_proj_scale for fp8 w8a16.");
TORCH_CHECK(kv_a_proj_scale.has_value(), "missing kv_a_proj_scale for fp8 w8a16.");
TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a16.");
TORCH_CHECK(block_size.value().size() == 2, "block_size should be 2D for fp8 w8a16.");
}
// outputs and temp buffer
const auto options = hidden_states.options();
auto q_input = at::empty({num_seqs, num_heads, kv_lora_rank + qk_rope_head_dim}, options);
auto k_input = at::empty({num_seqs, 1, kv_lora_rank + qk_rope_head_dim}, options);
auto v_input = k_input.narrow(-1, 0, kv_lora_rank);
// outputs of q_a_proj and q_b_proj
auto qa = at::empty({num_seqs, q_lora_rank}, options);
// stage 1: q_a_proj and kv_a_proj
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "qkv_proj_kernel_impl", [&] {
if (use_int8_w8a8) {
auto q_a_proj_s = q_a_proj_scale.value();
auto kv_a_proj_s = kv_a_proj_scale.value();
TORCH_CHECK(q_a_proj_s.numel() == q_lora_rank);
TORCH_CHECK(kv_a_proj_s.numel() == kv_lora_rank + qk_rope_head_dim);
auto buffer = at::empty({num_seqs * hidden_size + num_seqs * 4}, options.dtype(at::kByte));
uint8_t* __restrict__ Aq_data = buffer.data_ptr<uint8_t>();
float* __restrict__ As_data = (float*)((void*)(Aq_data + num_seqs * hidden_size));
const scalar_t* __restrict__ A_data = hidden_states.data_ptr<scalar_t>();
at::parallel_for(0, num_seqs, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
quantize_row_int8<scalar_t>(Aq_data + m * hidden_size, As_data[m], A_data + m * hidden_size, hidden_size);
}
});
segment_gemm_kernel_impl<scalar_t>(
qa.data_ptr<scalar_t>(),
k_input.data_ptr<scalar_t>(),
Aq_data,
q_a_proj_weight.data_ptr<int8_t>(),
kv_a_proj_weight.data_ptr<int8_t>(),
As_data,
q_a_proj_s.data_ptr<float>(),
kv_a_proj_s.data_ptr<float>(),
num_seqs,
q_lora_rank,
kv_lora_rank + qk_rope_head_dim,
hidden_size);
} else if (use_fp8_w8a16) {
int64_t block_size_N = block_size.value()[0];
int64_t block_size_K = block_size.value()[1];
auto q_a_proj_s = q_a_proj_scale.value();
auto kv_a_proj_s = kv_a_proj_scale.value();
CHECK_EQ(q_a_proj_s.size(0), div_up(q_lora_rank, block_size_N));
CHECK_EQ(q_a_proj_s.size(1), div_up(hidden_size, block_size_K));
CHECK_EQ(kv_a_proj_s.size(0), div_up(kv_lora_rank + qk_rope_head_dim, block_size_N));
CHECK_EQ(kv_a_proj_s.size(1), div_up(hidden_size, block_size_K));
const int BLOCK_N = block_size_n();
const int num_threads = at::get_num_threads();
auto buffer = at::empty({num_threads, BLOCK_N * hidden_size}, options);
segment_gemm_kernel_impl<scalar_t>(
qa.data_ptr<scalar_t>(),
k_input.data_ptr<scalar_t>(),
hidden_states.data_ptr<scalar_t>(),
q_a_proj_weight.data_ptr<at::Float8_e4m3fn>(),
kv_a_proj_weight.data_ptr<at::Float8_e4m3fn>(),
q_a_proj_s.data_ptr<float>(),
kv_a_proj_s.data_ptr<float>(),
buffer.data_ptr<scalar_t>(),
num_seqs,
q_lora_rank,
kv_lora_rank + qk_rope_head_dim,
hidden_size,
block_size_N,
block_size_K);
} else {
segment_gemm_kernel_impl<scalar_t>(
qa.data_ptr<scalar_t>(),
k_input.data_ptr<scalar_t>(),
hidden_states.data_ptr<scalar_t>(),
q_a_proj_weight.data_ptr<scalar_t>(),
kv_a_proj_weight.data_ptr<scalar_t>(),
num_seqs,
q_lora_rank,
kv_lora_rank + qk_rope_head_dim,
hidden_size);
}
});
// stage 2: apply rmsnorm inplace
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "rms_norm_kernel_impl", [&] {
rms_norm_kernel_impl<scalar_t>(
qa.data_ptr<scalar_t>(),
v_input.data_ptr<scalar_t>(),
q_a_layernorm_weight.data_ptr<scalar_t>(),
kv_a_layernorm_weight.data_ptr<scalar_t>(),
num_seqs,
q_lora_rank,
kv_lora_rank,
kv_lora_rank + qk_rope_head_dim,
eps);
});
// stage 3: q_b_proj
at::Tensor qb;
std::optional<at::Tensor> bias;
if (use_int8_w8a8) {
qb = int8_scaled_mm_with_quant(qa, q_b_proj_weight, q_b_proj_scale.value(), bias, at::kBFloat16, is_vnni);
} else if (use_fp8_w8a16) {
qb = fp8_scaled_mm_cpu(
qa, q_b_proj_weight, q_b_proj_scale.value(), block_size.value(), bias, at::kBFloat16, is_vnni);
} else {
qb = weight_packed_linear(qa, q_b_proj_weight, bias, is_vnni);
}
qb.as_strided_({num_seqs, num_heads, qk_head_dim}, {num_heads * qk_head_dim, qk_head_dim, 1});
// stage 4: bmm
auto q_nope = qb.narrow(2, 0, qk_nope_head_dim).transpose_(0, 1);
auto q_nope_out = q_input.narrow(2, 0, kv_lora_rank).transpose_(0, 1);
bmm_cpu(q_nope_out, q_nope, w_kc, is_vnni, w_scale);
// stage 5: rope
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "rotary_emb_kernel_impl", [&] {
rotary_emb_kernel_impl<scalar_t>(
q_input.data_ptr<scalar_t>() + kv_lora_rank,
k_input.data_ptr<scalar_t>() + kv_lora_rank,
qb.data_ptr<scalar_t>() + qk_nope_head_dim,
k_input.data_ptr<scalar_t>() + kv_lora_rank,
positions.data_ptr<int64_t>(),
cos_sin_cache.data_ptr<scalar_t>(),
num_seqs,
num_heads,
rotary_dim,
num_heads * qk_head_dim,
qk_head_dim,
kv_lora_rank + qk_rope_head_dim,
num_heads * (kv_lora_rank + qk_rope_head_dim),
kv_lora_rank + qk_rope_head_dim,
kv_lora_rank + qk_rope_head_dim);
});
return std::make_tuple(q_input, k_input, v_input);
}
std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope_fused_weight(
at::Tensor& hidden_states,
at::Tensor& qkv_a_proj_weight,
at::Tensor& q_b_proj_weight,
at::Tensor& w_kc,
at::Tensor& q_a_layernorm_weight,
at::Tensor& kv_a_layernorm_weight,
at::Tensor& positions,
at::Tensor& cos_sin_cache,
double eps,
bool use_int8_w8a8,
bool use_fp8_w8a16,
std::optional<at::Tensor> qkv_a_proj_scale,
std::optional<at::Tensor> q_b_proj_scale,
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size,
int64_t q_lora_rank,
int64_t kv_lora_rank,
int64_t qk_rope_head_dim) {
int64_t hidden_size = hidden_states.size(1);
CHECK_EQ(qkv_a_proj_weight.size(0), q_lora_rank + kv_lora_rank + qk_rope_head_dim);
CHECK_EQ(qkv_a_proj_weight.size(1), get_row_size(hidden_size, use_int8_w8a8));
std::vector<at::Tensor> weight_chunks =
at::split(qkv_a_proj_weight, {q_lora_rank, kv_lora_rank + qk_rope_head_dim}, 0);
at::Tensor q_a_proj_weight = weight_chunks[0];
at::Tensor kv_a_proj_weight = weight_chunks[1];
at::Tensor q_a_proj_s;
at::Tensor kv_a_proj_s;
if (use_int8_w8a8) {
TORCH_CHECK(qkv_a_proj_scale.has_value(), "missing qkv_a_proj_scale for int8 w8a8.");
std::vector<at::Tensor> scale_chunks =
at::split(qkv_a_proj_scale.value(), {q_lora_rank, kv_lora_rank + qk_rope_head_dim}, 0);
q_a_proj_s = scale_chunks[0];
kv_a_proj_s = scale_chunks[1];
}
if (use_fp8_w8a16) {
TORCH_CHECK(qkv_a_proj_scale.has_value(), "missing qkv_a_proj_scale for fp8 w8a16.");
int64_t block_size_N = block_size.value()[0];
int64_t q_a_proj_s_dim0 = div_up(q_lora_rank, block_size_N);
int64_t kv_a_proj_s_dim0 = div_up(kv_lora_rank + qk_rope_head_dim, block_size_N);
std::vector<at::Tensor> scale_chunks = at::split(qkv_a_proj_scale.value(), {q_a_proj_s_dim0, kv_a_proj_s_dim0}, 0);
q_a_proj_s = scale_chunks[0];
kv_a_proj_s = scale_chunks[1];
}
return qkv_proj_with_rope(
hidden_states,
q_a_proj_weight,
q_b_proj_weight,
kv_a_proj_weight,
w_kc,
q_a_layernorm_weight,
kv_a_layernorm_weight,
positions,
cos_sin_cache,
eps,
use_int8_w8a8,
use_fp8_w8a16,
q_a_proj_s,
q_b_proj_scale,
kv_a_proj_s,
w_scale,
is_vnni,
block_size);
}
python/sglang/kernels/aot/csrc/cpu/rope.cpp
+453
View File
@@ -0,0 +1,453 @@
#include "common.h"
#include "vec.h"
namespace {
struct RopeParams {
// Treat all tensors as [B, S, H, D]
// 2D [S, H * D] -> [1, S, H, D]
// 3D [S, H, D] -> [1, S, H, D]
// 4D [B, S, H, D]
int64_t rotary_dim{0};
int64_t head_size{0};
int64_t batches{1}, seqlen{1}, num_heads{1}, num_heads_kv{1};
int64_t q_strideB{0}, q_strideS{0}, q_strideH{0};
int64_t k_strideB{0}, k_strideS{0}, k_strideH{0};
RopeParams(const at::Tensor& query, const at::Tensor& key, int64_t head_size_, int64_t rotary_dim_)
: rotary_dim(rotary_dim_), head_size(head_size_) {
int64_t ndim = query.dim();
switch (ndim) {
case 2:
seqlen = query.size(0);
num_heads = query.size(1) / head_size;
num_heads_kv = key.size(1) / head_size;
q_strideS = query.stride(0);
k_strideS = key.stride(0);
q_strideH = head_size;
k_strideH = head_size;
break;
case 3:
seqlen = query.size(0);
num_heads = query.size(1);
num_heads_kv = key.size(1);
q_strideS = query.stride(0);
k_strideS = key.stride(0);
q_strideH = query.stride(1);
k_strideH = key.stride(1);
break;
case 4:
batches = query.size(0);
seqlen = query.size(1);
num_heads = query.size(2);
num_heads_kv = key.size(2);
q_strideB = query.stride(0);
k_strideB = key.stride(0);
q_strideS = query.stride(1);
k_strideS = key.stride(1);
q_strideH = query.stride(2);
k_strideH = key.stride(2);
break;
default:
TORCH_CHECK(false, "Expected a 2D/3D/4D tensor, got ", ndim, "D.");
}
}
inline int64_t rows() const {
return batches * seqlen;
}
inline int64_t q_offset(int64_t b, int64_t s, int64_t h) const {
return b * q_strideB + s * q_strideS + h * q_strideH;
}
inline int64_t k_offset(int64_t b, int64_t s, int64_t h) const {
return b * k_strideB + s * k_strideS + h * k_strideH;
}
inline int64_t q_out_offset(int64_t b, int64_t s, int64_t h) const {
return ((b * seqlen + s) * num_heads + h) * head_size;
}
inline int64_t k_out_offset(int64_t b, int64_t s, int64_t h) const {
return ((b * seqlen + s) * num_heads_kv + h) * head_size;
}
};
enum class RotaryMode {
Interleaved, // GPT-J / packed [cos|sin]
Neox, // packed [cos|sin]
NeoxFull, // split cos/sin each of length head_size (HF rotate_half)
};
// Already-indexed cos/sin rows for apply_rotary_pos_emb style.
template <typename param_t>
struct SplitCosSinRow {
const param_t* cos;
const param_t* sin;
};
// Already-indexed T/H/W cache rows for 2D mRoPE (no gathered buffer).
template <typename scalar_t>
struct MropeCosSinRow {
const scalar_t* cache_t;
const scalar_t* cache_h;
const scalar_t* cache_w;
int64_t section_t;
int64_t section_h;
int64_t section_w;
bool interleaved;
inline const scalar_t* ptr_at(int64_t j) const {
if (interleaved) {
if (j % 3 == 1 && j <= section_h * 3) return cache_h;
if (j % 3 == 2 && j <= section_w * 3) return cache_w;
return cache_t;
}
if (j < section_t) return cache_t;
if (j < section_t + section_h) return cache_h;
return cache_w;
}
};
template <typename scalar_t, RotaryMode rotary_mode>
struct RotaryEmbedInternal;
template <typename scalar_t>
struct RotaryEmbedInternal<scalar_t, RotaryMode::Interleaved> {
static inline void
apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ cache, int size) {
constexpr int kVecSize = at::vec::Vectorized<scalar_t>::size();
const int half_size = size / 2;
int d = 0;
for (; d <= size - kVecSize; d += kVecSize) {
auto [xy0, xy1] = load_float_vec2(input + d);
auto [x, y] = at::vec::deinterleave2(xy0, xy1);
auto cos = load_float_vec(cache + d / 2);
auto sin = load_float_vec(cache + half_size + d / 2);
auto out0 = x * cos - y * sin;
auto out1 = y * cos + x * sin;
std::tie(xy0, xy1) = at::vec::interleave2(out0, out1);
convert_from_float_ext<scalar_t>(xy0, xy1).store(out + d);
}
for (; d < size; d += 2) {
float x = input[d], y = input[d + 1];
float cos = cache[d >> 1], sin = cache[half_size + (d >> 1)];
out[d] = static_cast<scalar_t>(x * cos - y * sin);
out[d + 1] = static_cast<scalar_t>(y * cos + x * sin);
}
}
// mRoPE: cos/sin may come from different T/H/W rows per pair index.
static inline void
apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, MropeCosSinRow<scalar_t> cache, int size) {
const int half_size = size / 2;
for (int j = 0; j < half_size; ++j) {
const scalar_t* src = cache.ptr_at(j);
float cos = src[j], sin = src[j + half_size];
float x = input[2 * j], y = input[2 * j + 1];
out[2 * j] = static_cast<scalar_t>(x * cos - y * sin);
out[2 * j + 1] = static_cast<scalar_t>(y * cos + x * sin);
}
}
};
template <typename scalar_t>
struct RotaryEmbedInternal<scalar_t, RotaryMode::Neox> {
static inline void
apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ cache, int size) {
constexpr int kVecSize = at::vec::Vectorized<scalar_t>::size();
const int half_size = size / 2;
int d = 0;
for (; d <= half_size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
auto [y0, y1] = load_float_vec2(input + half_size + d);
auto [cos0, cos1] = load_float_vec2(cache + d);
auto [sin0, sin1] = load_float_vec2(cache + half_size + d);
auto out0 = x0 * cos0 - y0 * sin0;
auto out1 = x1 * cos1 - y1 * sin1;
auto out2 = y0 * cos0 + x0 * sin0;
auto out3 = y1 * cos1 + x1 * sin1;
convert_from_float_ext<scalar_t>(out0, out1).store(out + d);
convert_from_float_ext<scalar_t>(out2, out3).store(out + half_size + d);
}
for (; d < half_size; ++d) {
float x = input[d], y = input[d + half_size];
float cos = cache[d], sin = cache[d + half_size];
out[d] = static_cast<scalar_t>(x * cos - y * sin);
out[d + half_size] = static_cast<scalar_t>(y * cos + x * sin);
}
}
// mRoPE: cos/sin may come from different T/H/W rows per rotary index.
static inline void
apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, MropeCosSinRow<scalar_t> cache, int size) {
const int half_size = size / 2;
for (int j = 0; j < half_size; ++j) {
const scalar_t* src = cache.ptr_at(j);
float cos = src[j], sin = src[j + half_size];
float x = input[j], y = input[j + half_size];
out[j] = static_cast<scalar_t>(x * cos - y * sin);
out[j + half_size] = static_cast<scalar_t>(y * cos + x * sin);
}
}
};
template <typename scalar_t>
struct RotaryEmbedInternal<scalar_t, RotaryMode::NeoxFull> {
template <typename CosT>
static inline void
apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, SplitCosSinRow<CosT> cache, int size) {
constexpr int kVecSize = at::vec::Vectorized<scalar_t>::size();
const int half_size = size / 2;
int d = 0;
for (; d <= half_size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
auto [y0, y1] = load_float_vec2(input + half_size + d);
auto [cos_x0, cos_x1] = load_float_vec2(cache.cos + d);
auto [sin_x0, sin_x1] = load_float_vec2(cache.sin + d);
auto [cos_y0, cos_y1] = load_float_vec2(cache.cos + half_size + d);
auto [sin_y0, sin_y1] = load_float_vec2(cache.sin + half_size + d);
auto out0 = x0 * cos_x0 - y0 * sin_x0;
auto out1 = x1 * cos_x1 - y1 * sin_x1;
auto out2 = y0 * cos_y0 + x0 * sin_y0;
auto out3 = y1 * cos_y1 + x1 * sin_y1;
convert_from_float_ext<scalar_t>(out0, out1).store(out + d);
convert_from_float_ext<scalar_t>(out2, out3).store(out + half_size + d);
}
for (; d < half_size; ++d) {
float x = input[d], y = input[d + half_size];
float cos_x = static_cast<float>(cache.cos[d]);
float sin_x = static_cast<float>(cache.sin[d]);
float cos_y = static_cast<float>(cache.cos[d + half_size]);
float sin_y = static_cast<float>(cache.sin[d + half_size]);
out[d] = static_cast<scalar_t>(x * cos_x - y * sin_x);
out[d + half_size] = static_cast<scalar_t>(y * cos_y + x * sin_y);
}
}
};
template <typename scalar_t, RotaryMode mode, bool inplace, typename CachePos>
void rotary_embedding_kernel_impl(
scalar_t* __restrict__ query_out,
scalar_t* __restrict__ key_out,
scalar_t* __restrict__ query,
scalar_t* __restrict__ key,
const RopeParams& p,
const CachePos& cache_pos) {
at::parallel_for(0, p.rows(), 0, [&](int64_t begin, int64_t end) {
int64_t bs = 0, seq = 0;
data_index_init(begin, bs, p.batches, seq, p.seqlen);
for (int64_t i = begin; i < end; ++i) {
auto cache = cache_pos(bs * p.seqlen + seq);
for (int64_t h = 0; h < p.num_heads; ++h) {
scalar_t* q_in = query + p.q_offset(bs, seq, h);
scalar_t* q_out;
if constexpr (inplace) {
q_out = q_in;
} else {
q_out = query_out + p.q_out_offset(bs, seq, h);
}
RotaryEmbedInternal<scalar_t, mode>::apply(q_out, q_in, cache, p.rotary_dim);
}
for (int64_t h = 0; h < p.num_heads_kv; ++h) {
scalar_t* k_in = key + p.k_offset(bs, seq, h);
scalar_t* k_out;
if constexpr (inplace) {
k_out = k_in;
} else {
k_out = key_out + p.k_out_offset(bs, seq, h);
}
RotaryEmbedInternal<scalar_t, mode>::apply(k_out, k_in, cache, p.rotary_dim);
}
data_index_step(bs, p.batches, seq, p.seqlen);
}
});
}
} // namespace
// 2D : [num_tokens, num_heads*head_size] inplace
// 3D : [num_tokens, num_heads, head_size] outplace
// 4D : [batch_size, seq_len, num_heads, head_size] inplace
std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
at::Tensor& positions,
at::Tensor& query,
at::Tensor& key,
int64_t head_size,
at::Tensor& cos_sin_cache,
bool is_neox) {
CHECK_DIM(1, positions);
const auto input_dim = query.dim();
const auto input_dtype = query.scalar_type();
TORCH_CHECK(input_dim >= 2 && input_dim <= 4, "Query/Key must be 2D/3D/4D, got ", input_dim, "D.");
CHECK_DIM(2, cos_sin_cache);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(key);
TORCH_CHECK(positions.scalar_type() == at::kLong, "expect positions to be int64, got ", positions.scalar_type());
TORCH_CHECK(input_dtype == key.scalar_type(), "query and key must have the same data type");
TORCH_CHECK(input_dtype == cos_sin_cache.scalar_type(), "query and cos_sin_cache must have the same data type");
int64_t rotary_dim = cos_sin_cache.size(1);
const RopeParams p{query, key, head_size, rotary_dim};
TORCH_CHECK(p.rotary_dim <= p.head_size, "rotary_dim must be <= head_size");
TORCH_CHECK(p.rotary_dim % 2 == 0, "rotary_dim must be even");
TORCH_CHECK(positions.numel() == p.rows(), "positions.numel() must equal batch * seqlen");
if (input_dim <= 3) {
CHECK_EQ(key.size(0), query.size(0));
}
if (input_dim == 2) {
CHECK_EQ(query.size(1), p.num_heads * p.head_size);
CHECK_EQ(key.size(1), p.num_heads_kv * p.head_size);
}
if (input_dim == 3) {
// out-of-place path: align with legacy behavior, no partial rotary
CHECK_EQ(query.size(-1), rotary_dim);
CHECK_EQ(key.size(-1), rotary_dim);
CHECK_EQ(head_size, rotary_dim);
}
if (input_dim == 4) {
CHECK_EQ(query.size(0), key.size(0));
CHECK_EQ(query.size(1), key.size(1));
}
at::Tensor query_out = (input_dim != 3) ? query : at::empty(query.sizes(), query.options());
at::Tensor key_out = (input_dim != 3) ? key : at::empty(key.sizes(), key.options());
AT_DISPATCH_REDUCED_FLOATING_TYPES(input_dtype, "rotary_embedding_cpu", [&] {
AT_DISPATCH_BOOL(input_dim != 3, inplace, [&] {
const scalar_t* cache_base = cos_sin_cache.data_ptr<scalar_t>();
const int64_t* pos_ptr = positions.data_ptr<int64_t>();
auto cache_pos = [cache_base, pos_ptr, rotary_dim](int64_t token) -> const scalar_t* {
return cache_base + pos_ptr[token] * rotary_dim;
};
scalar_t* q_ptr = query.data_ptr<scalar_t>();
scalar_t* k_ptr = key.data_ptr<scalar_t>();
scalar_t* q_out_ptr = query_out.data_ptr<scalar_t>();
scalar_t* k_out_ptr = key_out.data_ptr<scalar_t>();
if (is_neox) {
rotary_embedding_kernel_impl<scalar_t, RotaryMode::Neox, inplace>(
q_out_ptr, k_out_ptr, q_ptr, k_ptr, p, cache_pos);
} else {
rotary_embedding_kernel_impl<scalar_t, RotaryMode::Interleaved, inplace>(
q_out_ptr, k_out_ptr, q_ptr, k_ptr, p, cache_pos);
}
});
});
return std::make_tuple(query_out, key_out);
}
// query: [num_tokens, num_heads, head_size]
// key: [num_tokens, num_heads, head_size]
// cos: [num_tokens, head_size]
// sin: [num_tokens, head_size]
std::tuple<at::Tensor, at::Tensor>
apply_rotary_pos_emb_cpu(at::Tensor& query, at::Tensor& key, at::Tensor& cos, at::Tensor& sin) {
CHECK_DIM(3, query);
const auto input_dtype = query.scalar_type();
int64_t num_tokens = query.size(0);
int64_t num_heads = query.size(1);
int64_t head_size = query.size(2);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
CHECK_INPUT_SHAPE_DTYPE<true>(key, {num_tokens, num_heads, head_size}, input_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(cos, {num_tokens, head_size}, cos.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<false>(sin, {num_tokens, head_size}, sin.scalar_type());
CHECK_EQ(cos.scalar_type(), sin.scalar_type());
TORCH_CHECK(head_size % 2 == 0, "head_size must be even");
const RopeParams p{query, key, head_size, head_size};
CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(input_dtype, cos.scalar_type(), [&] {
scalar_t* q_ptr = query.data_ptr<scalar_t>();
scalar_t* k_ptr = key.data_ptr<scalar_t>();
const param_t* cos_ptr = cos.data_ptr<param_t>();
const param_t* sin_ptr = sin.data_ptr<param_t>();
auto cache_pos = [cos_ptr, sin_ptr, head_size](int64_t token) -> SplitCosSinRow<param_t> {
return {cos_ptr + token * head_size, sin_ptr + token * head_size};
};
rotary_embedding_kernel_impl<scalar_t, RotaryMode::NeoxFull, true>(q_ptr, k_ptr, q_ptr, k_ptr, p, cache_pos);
});
return std::make_tuple(query, key);
}
// positions: [num_tokens] (text only) or [3, num_tokens] (T/H/W positions with multimodal inputs)
// query: [num_tokens, num_heads * head_size]
// key: [num_tokens, num_kv_heads * head_size]
// cos_sin_cache: [max_position_embeddings, rotary_dim]
// mrope_section: [t, h, w]
std::tuple<at::Tensor, at::Tensor> multimodal_rotary_embedding_cpu(
at::Tensor& positions,
at::Tensor& query,
at::Tensor& key,
int64_t head_size,
at::Tensor& cos_sin_cache,
const std::optional<std::vector<int64_t>>& mrope_section,
bool mrope_interleaved,
bool is_neox) {
TORCH_CHECK(positions.dim() == 1 || positions.dim() == 2, "positions must be a 1D or 2D tensor");
CHECK_EQ(positions.scalar_type(), at::kLong);
CHECK_DIM(2, query);
const auto input_dtype = query.scalar_type();
int64_t rotary_dim = cos_sin_cache.size(1);
int64_t num_tokens = positions.size(-1);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(key);
CHECK_EQ(query.size(0), num_tokens);
CHECK_EQ(key.size(0), num_tokens);
CHECK_EQ(query.size(-1) % head_size, 0);
CHECK_EQ(key.size(-1) % head_size, 0);
CHECK_EQ(input_dtype, key.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<false>(cos_sin_cache, {cos_sin_cache.size(0), rotary_dim}, input_dtype);
const RopeParams p{query, key, head_size, rotary_dim};
TORCH_CHECK(p.rotary_dim <= p.head_size, "rotary_dim must be <= head_size");
TORCH_CHECK(p.rotary_dim % 2 == 0, "rotary_dim must be even");
TORCH_CHECK(positions.size(-1) == p.rows(), "positions.size(-1) must equal batch * seqlen");
AT_DISPATCH_REDUCED_FLOATING_TYPES(input_dtype, "multimodal_rotary_embedding_cpu", [&] {
const scalar_t* cache_base = cos_sin_cache.data_ptr<scalar_t>();
const int64_t* pos_ptr = positions.data_ptr<int64_t>();
scalar_t* q_ptr = query.data_ptr<scalar_t>();
scalar_t* k_ptr = key.data_ptr<scalar_t>();
if (positions.dim() == 2) {
TORCH_CHECK(mrope_section.has_value(), "mrope_section must be provided when positions is 2D");
auto mrope_section_val = mrope_section.value();
CHECK_EQ(mrope_section_val.size(), 3);
CHECK_EQ(positions.size(0), 3);
const int64_t section_t = mrope_section_val[0];
const int64_t section_h = mrope_section_val[1];
const int64_t section_w = mrope_section_val[2];
const int64_t p_stride0 = positions.stride(0);
auto cache_pos = [=](int64_t token) -> MropeCosSinRow<scalar_t> {
return {
cache_base + pos_ptr[0 * p_stride0 + token] * rotary_dim,
cache_base + pos_ptr[1 * p_stride0 + token] * rotary_dim,
cache_base + pos_ptr[2 * p_stride0 + token] * rotary_dim,
section_t,
section_h,
section_w,
mrope_interleaved};
};
if (is_neox) {
rotary_embedding_kernel_impl<scalar_t, RotaryMode::Neox, true>(q_ptr, k_ptr, q_ptr, k_ptr, p, cache_pos);
} else {
rotary_embedding_kernel_impl<scalar_t, RotaryMode::Interleaved, true>(q_ptr, k_ptr, q_ptr, k_ptr, p, cache_pos);
}
} else { // positions.dim() == 1
auto cache_pos = [cache_base, pos_ptr, rotary_dim](int64_t token) -> const scalar_t* {
return cache_base + pos_ptr[token] * rotary_dim;
};
if (is_neox) {
rotary_embedding_kernel_impl<scalar_t, RotaryMode::Neox, true>(q_ptr, k_ptr, q_ptr, k_ptr, p, cache_pos);
} else {
rotary_embedding_kernel_impl<scalar_t, RotaryMode::Interleaved, true>(q_ptr, k_ptr, q_ptr, k_ptr, p, cache_pos);
}
}
});
return std::make_tuple(query, key);
}
python/sglang/kernels/aot/csrc/cpu/shm.cpp
+522
View File
@@ -0,0 +1,522 @@
#include "shm.h"
#if defined(__x86_64__)
#include "x86_64/shm.h"
#elif defined(__aarch64__)
#include "aarch64/shm.h"
#else
#error "unsupported architecture"
#endif
#include <ATen/ATen.h>
#include <errno.h>
#include <fcntl.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <unistd.h>
// states for collectives
enum coll_state {
coll_begin = 0,
coll_allreduce_naive__copy_in_done,
coll_allreduce_naive__reduce_done,
// alternative state when allreduce is working on alternative buffer
// of the double buffer.
coll_alt1_allreduce_naive__copy_in_done,
coll_alt2_allreduce_naive__copy_in_done,
coll_alt1_allreduce_naive__reduce_done,
coll_allgather_naive__copy_in_done,
coll_alt1_allgather_naive__copy_in_done,
coll_alt2_allgather_naive__copy_in_done,
coll_reduce_scatter_naive__copy_in_done,
coll_reduce_scatter_naive__reduce_done,
coll_alt1_reduce_scatter_naive__copy_in_done,
coll_alt2_reduce_scatter_naive__copy_in_done,
};
// SHM building blocks
struct SharedData {
const char* name;
int descriptor;
void* bytes;
size_t nbytes;
};
void shared_open(SharedData* data, const char* name, size_t nbytes) {
int d = shm_open(name, O_RDWR, S_IRUSR | S_IWUSR);
if (d != -1) {
void* bytes = mmap(NULL, nbytes, PROT_READ | PROT_WRITE, MAP_SHARED, d, 0);
data->name = name;
data->descriptor = d;
data->bytes = bytes;
data->nbytes = nbytes;
} else {
if (errno != ENOENT) {
// don't print if shm can not be found because we want to loop over from
// caller again until the other ranks created the shm
printf("shared_open %s failed, errno=%d\n", name, errno);
}
data->descriptor = -1;
}
}
void shared_create(SharedData* data, const char* name, void* bytes, size_t nbytes) {
int d = shm_open(name, O_CREAT | O_RDWR, S_IRUSR | S_IWUSR);
if (d != -1) {
nbytes = write(d, bytes, nbytes);
if (nbytes > 0) {
shared_open(data, name, nbytes);
}
} else {
printf("shared_create %s failed\n", name);
}
}
static int world_size;
// SHM based allreduce helper functions
// buffer that holds shm name
#define NAME_BUF_SIZE 1000
#define MAX_BUF_SIZE 1048576 * 32
#define NAIVE_ALLREDUCE_THRESHOLD 1048576
#define SHM_BUFFER_NAME "deepspeed_allreduce_buffer"
struct allreduce_workspace {
enum coll_state states[5]; // idx=0 -- state for symmetric_naive_all_reduce
// idx=1 -- state for distributed_naive_all_reduce
// idx=2 -- state for all_gather
// idx=3 -- state for all_gather_into_tensor
// idx=4 -- state for reduce_scatter
// double buffer to avoid syncing between rounds
// offset=0 -- 2*NAIVE_ALLREDUCE_THRESHOLD : buffer for
// symmetric_naive_all_reduce after that : buffer for
// distributed_naive_all_reduce
char buffer
[2 * NAIVE_ALLREDUCE_THRESHOLD + // symmetric allreduce
2 * MAX_BUF_SIZE + // distributed naive reduce
2 * MAX_BUF_SIZE + // allgather
2 * MAX_BUF_SIZE + // allgather_into_tensor
2 * MAX_BUF_SIZE // reduce_scatter
];
};
#define BUFFER0_OFFSET(current_buffer) current_buffer* NAIVE_ALLREDUCE_THRESHOLD
#define BUFFER1_OFFSET(current_buffer) 2 * NAIVE_ALLREDUCE_THRESHOLD + current_buffer* MAX_BUF_SIZE
#define BUFFER2_OFFSET(current_buffer) \
(2 * NAIVE_ALLREDUCE_THRESHOLD + 2 * MAX_BUF_SIZE + current_buffer * MAX_BUF_SIZE) // allgather
#define BUFFER3_OFFSET(current_buffer) \
(2 * NAIVE_ALLREDUCE_THRESHOLD + 4 * MAX_BUF_SIZE + current_buffer * MAX_BUF_SIZE) // allgather_into_tensor
#define BUFFER4_OFFSET(current_buffer) \
(2 * NAIVE_ALLREDUCE_THRESHOLD + 6 * MAX_BUF_SIZE + current_buffer * MAX_BUF_SIZE) // reduce_scatter
struct allreduce_workspace** workspace;
// buffer for small messages, double buffer
char** symmetric_buffer[2];
// buffer for large messages, double buffer
char** distributed_buffer[2];
char** allgather_buffer[2];
char** allgather_into_tensor_buffer[2];
char** reduce_scatter_buffer[2];
void wait_buffer_state_until_2(int index, enum coll_state state0, enum coll_state state1, int state_group) {
volatile enum coll_state* state_ptr = &(workspace[index]->states[state_group]);
while (1) {
volatile enum coll_state cur_state = *state_ptr;
if (cur_state == state0 || cur_state == state1) break;
}
}
void reduce_all_buffers(
int start_elements,
int num_elements,
c10::ScalarType scalar_type,
int to_buffer_idx,
char* to_buffer,
char** buffers) {
switch (scalar_type) {
case c10::ScalarType::BFloat16:
reduce_bf16_buffers(start_elements, num_elements, to_buffer, buffers, world_size);
break;
case c10::ScalarType::Half:
reduce_fp16_buffers(start_elements, num_elements, to_buffer, buffers, world_size);
break;
case c10::ScalarType::Float:
reduce_fp32_buffers(start_elements, num_elements, to_buffer, buffers, world_size);
break;
default:
assert(!"Should not get here");
}
}
static bool is_initialized = false;
static int world_rank;
void shm_initialize(int size, int rank, const char* addr_string, const char* port_string) {
if (is_initialized) {
return;
}
is_initialized = true;
world_size = size;
world_rank = rank;
char shm_name_prefix[NAME_BUF_SIZE];
char shm_name[NAME_BUF_SIZE];
snprintf(shm_name_prefix, NAME_BUF_SIZE, "%s_%d_%s_%s", SHM_BUFFER_NAME, getuid(), addr_string, port_string);
// create shared workspace for SHM based allreduce
SharedData allreduce_buffer;
// allocate workspace_buf for current rank
struct allreduce_workspace* workspace_buf;
struct allreduce_workspace* workspace_buf_other;
workspace_buf = (struct allreduce_workspace*)malloc(sizeof(struct allreduce_workspace));
snprintf(shm_name, NAME_BUF_SIZE, "%.900s_%d", shm_name_prefix, rank);
shared_create(&allreduce_buffer, shm_name, workspace_buf, sizeof(struct allreduce_workspace));
workspace_buf = (struct allreduce_workspace*)allreduce_buffer.bytes;
workspace_buf->states[STATE_GROUP_SYMMETRIC_ALLREDUCE] =
coll_alt2_allreduce_naive__copy_in_done; // symmetric_naive_all_reduce
workspace_buf->states[STATE_GROUP_DISTRIBUTED_ALLREDUCE] = coll_begin; // distributed_naive_reduce
workspace_buf->states[STATE_GROUP_ALL_GATHER] = coll_alt2_allgather_naive__copy_in_done; // all_gather
workspace_buf->states[STATE_GROUP_ALL_GATHER_INTO_TENSOR] =
coll_alt2_allgather_naive__copy_in_done; // all_gather_into_tensor
workspace_buf->states[STATE_GROUP_REDUCE_SCATTER] = coll_begin; // reduce_scatter
// create the workspace pointer list
workspace = (struct allreduce_workspace**)malloc(size * sizeof(struct allreduce_workspace*));
symmetric_buffer[0] = (char**)malloc(size * sizeof(char**));
symmetric_buffer[1] = (char**)malloc(size * sizeof(char**));
distributed_buffer[0] = (char**)malloc(size * sizeof(char**));
distributed_buffer[1] = (char**)malloc(size * sizeof(char**));
allgather_buffer[0] = (char**)malloc(size * sizeof(char*));
allgather_buffer[1] = (char**)malloc(size * sizeof(char*));
allgather_into_tensor_buffer[0] = (char**)malloc(size * sizeof(char*));
allgather_into_tensor_buffer[1] = (char**)malloc(size * sizeof(char*));
reduce_scatter_buffer[0] = (char**)malloc(size * sizeof(char*));
reduce_scatter_buffer[1] = (char**)malloc(size * sizeof(char*));
// map shm of all ranks
for (int i = 0; i < size; i++) {
if (i != rank) {
snprintf(shm_name, NAME_BUF_SIZE, "%.900s_%d", shm_name_prefix, i);
// printf("open %s, %d\n", shm_name, rank);
do {
shared_open(&allreduce_buffer, shm_name, sizeof(struct allreduce_workspace));
} while (allreduce_buffer.descriptor == -1 && errno == ENOENT);
workspace_buf_other = (struct allreduce_workspace*)allreduce_buffer.bytes;
workspace[i] = workspace_buf_other;
} else {
workspace[i] = workspace_buf;
}
symmetric_buffer[0][i] = workspace[i]->buffer + BUFFER0_OFFSET(0);
symmetric_buffer[1][i] = workspace[i]->buffer + BUFFER0_OFFSET(1);
distributed_buffer[0][i] = workspace[i]->buffer + BUFFER1_OFFSET(0);
distributed_buffer[1][i] = workspace[i]->buffer + BUFFER1_OFFSET(1);
allgather_buffer[0][i] = workspace[i]->buffer + BUFFER2_OFFSET(0);
allgather_buffer[1][i] = workspace[i]->buffer + BUFFER2_OFFSET(1);
allgather_into_tensor_buffer[0][i] = workspace[i]->buffer + BUFFER3_OFFSET(0);
allgather_into_tensor_buffer[1][i] = workspace[i]->buffer + BUFFER3_OFFSET(1);
reduce_scatter_buffer[0][i] = workspace[i]->buffer + BUFFER4_OFFSET(0);
reduce_scatter_buffer[1][i] = workspace[i]->buffer + BUFFER4_OFFSET(1);
}
}
#define positive_mod(num, mod) ((((num) % (mod)) + (mod)) % (mod))
#define rank_mod(rank) positive_mod(rank, world_size)
size_t slice_size(size_t chunk_el, int slice_idx) {
size_t slice_size = chunk_el / world_size;
return slice_idx == world_size - 1 ? slice_size + (chunk_el % world_size) : slice_size;
}
char* slice_data(char* data_ptr, size_t chunk_el, int el_size, int slice_idx) {
size_t slice_size = chunk_el / world_size;
size_t el_offset = slice_size * slice_idx;
return data_ptr + el_offset * el_size;
}
size_t slice_el_start(size_t chunk_el, int slice_idx) {
size_t slice_size = chunk_el / world_size;
return slice_size * slice_idx;
}
void symmetric_naive_all_reduce(char* data_ptr, c10::ScalarType scalar_type, size_t chunk_size, size_t chunk_el) {
const int state_group = STATE_GROUP_SYMMETRIC_ALLREDUCE;
static int current_buffer = 0;
static int state_idx = 0;
// init states to case 0 to get rid of "maybe-uninitialized" warning.
enum coll_state copy_current = coll_allreduce_naive__copy_in_done;
enum coll_state copy_next = coll_alt1_allreduce_naive__copy_in_done;
switch (state_idx) {
case 0:
copy_current = coll_allreduce_naive__copy_in_done;
copy_next = coll_alt1_allreduce_naive__copy_in_done;
break;
case 1:
copy_current = coll_alt1_allreduce_naive__copy_in_done;
copy_next = coll_alt2_allreduce_naive__copy_in_done;
break;
case 2:
copy_current = coll_alt2_allreduce_naive__copy_in_done;
copy_next = coll_allreduce_naive__copy_in_done;
break;
default:
assert(!"Should not get here.");
}
state_idx = (state_idx + 1) % 3;
parallel_memcpy(symmetric_buffer[current_buffer][world_rank], data_ptr, chunk_size);
std::atomic_thread_fence(std::memory_order_release);
workspace[world_rank]->states[state_group] = copy_current;
for (int i = 0; i < world_size; i++) {
// wait until the other rank copy the buffer
if (i != world_rank) {
wait_buffer_state_until_2(i, copy_current, copy_next, state_group);
}
}
// each rank reduce the buffer independently so therre is no need for
// synchronization afterward
reduce_all_buffers(0, chunk_el, scalar_type, world_rank, data_ptr, symmetric_buffer[current_buffer]);
// switch buffer
current_buffer = 1 - current_buffer;
}
// naive allreduce distributed, each rank do naive reduce on its slice
void distributed_naive_reduce(char* data_ptr, c10::ScalarType scalar_type, size_t chunk_size, size_t chunk_el) {
const int state_group = STATE_GROUP_DISTRIBUTED_ALLREDUCE;
static int current_buffer = 0;
static int state_idx = 0;
// init states to case 0 to get rid of "maybe-uninitialized" warning.
enum coll_state copy_current = coll_allreduce_naive__copy_in_done;
enum coll_state reduce_current = coll_allreduce_naive__reduce_done;
enum coll_state copy_next = coll_alt1_allreduce_naive__copy_in_done;
// similar to symmetric_naive_allreduce, but here we only need two sets of
// states, because distributed naive reduce has two barriers in the algorithm
switch (state_idx) {
case 0:
copy_current = coll_allreduce_naive__copy_in_done;
reduce_current = coll_allreduce_naive__reduce_done;
copy_next = coll_alt1_allreduce_naive__copy_in_done;
break;
case 1:
copy_current = coll_alt1_allreduce_naive__copy_in_done;
reduce_current = coll_alt1_allreduce_naive__reduce_done;
copy_next = coll_allreduce_naive__copy_in_done;
break;
default:
assert(!"Should not get here.");
}
state_idx = (state_idx + 1) % 2;
int data_size = chunk_size / chunk_el;
parallel_memcpy(distributed_buffer[current_buffer][world_rank], data_ptr, chunk_size);
std::atomic_thread_fence(std::memory_order_release);
workspace[world_rank]->states[state_group] = copy_current;
for (int i = 0; i < world_size; i++) {
// wait until all the other ranks copy the buffer
if (i != world_rank) wait_buffer_state_until_2(i, copy_current, reduce_current, state_group);
}
// reduce scatter
reduce_all_buffers(
slice_el_start(chunk_el, world_rank),
slice_size(chunk_el, world_rank),
scalar_type,
world_rank,
distributed_buffer[current_buffer][world_rank],
distributed_buffer[current_buffer]);
std::atomic_thread_fence(std::memory_order_release);
workspace[world_rank]->states[state_group] = reduce_current;
for (int i = 0; i < world_size; i++) {
// wait until all the other ranks reduce the buffer
if (i != world_rank) wait_buffer_state_until_2(i, reduce_current, copy_next, state_group);
}
for (int i = 0; i < world_size; i++) {
int rank = (i + world_rank) % world_size;
parallel_memcpy(
slice_data(data_ptr, chunk_el, data_size, rank),
slice_data(distributed_buffer[current_buffer][rank], chunk_el, chunk_size / chunk_el, rank),
slice_size(chunk_el, rank) * data_size);
}
current_buffer = 1 - current_buffer;
}
void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size) {
for (int offset = 0; offset < data_size; offset += MAX_BUF_SIZE) {
auto data_ptr = ((char*)(data.data_ptr()) + offset);
size_t chunk_size = data_size - offset > MAX_BUF_SIZE ? MAX_BUF_SIZE : data_size - offset;
size_t chunk_el = chunk_size / (data_size / numel);
if (chunk_size < NAIVE_ALLREDUCE_THRESHOLD) {
symmetric_naive_all_reduce(data_ptr, data.scalar_type(), chunk_size, chunk_el);
} else {
distributed_naive_reduce(data_ptr, data.scalar_type(), chunk_size, chunk_el);
}
}
}
template <int STATE_GROUP>
void naive_all_gather(char* result_ptr, char* data_ptr, size_t res_stride, size_t chunk_size, size_t chunk_el) {
static int current_buffer = 0;
static int state_idx = 0;
char*** buffer = nullptr;
if constexpr (STATE_GROUP == STATE_GROUP_ALL_GATHER) {
buffer = allgather_buffer;
} else if constexpr (STATE_GROUP == STATE_GROUP_ALL_GATHER_INTO_TENSOR) {
buffer = allgather_into_tensor_buffer;
} else {
static_assert(
STATE_GROUP == STATE_GROUP_ALL_GATHER || STATE_GROUP == STATE_GROUP_ALL_GATHER_INTO_TENSOR,
"Unsupported STATE_GROUP");
}
// init states to case 0 to get rid of "maybe-uninitialized" warning.
enum coll_state copy_current = coll_allgather_naive__copy_in_done;
enum coll_state copy_next = coll_alt1_allgather_naive__copy_in_done;
switch (state_idx) {
case 0:
copy_current = coll_allgather_naive__copy_in_done;
copy_next = coll_alt1_allgather_naive__copy_in_done;
break;
case 1:
copy_current = coll_alt1_allgather_naive__copy_in_done;
copy_next = coll_alt2_allgather_naive__copy_in_done;
break;
case 2:
copy_current = coll_alt2_allgather_naive__copy_in_done;
copy_next = coll_allgather_naive__copy_in_done;
break;
default:
assert(!"Should not get here.");
}
state_idx = (state_idx + 1) % 3;
parallel_memcpy(buffer[current_buffer][world_rank], data_ptr, chunk_size);
std::atomic_thread_fence(std::memory_order_release);
workspace[world_rank]->states[STATE_GROUP] = copy_current;
for (int i = 0; i < world_size; i++) {
// wait until all the other ranks copy the buffer
if (i != world_rank) wait_buffer_state_until_2(i, copy_current, copy_next, STATE_GROUP);
}
for (int i = 0; i < world_size; i++) {
parallel_memcpy(result_ptr + i * res_stride, buffer[current_buffer][i], chunk_size);
}
current_buffer = 1 - current_buffer;
}
template <int STATE_GROUP>
torch::Tensor& all_gather(torch::Tensor& result, torch::Tensor& data, int dim, size_t numel, int data_size) {
size_t dim_el = data.stride(dim) * data.size(dim);
int dtype_size = data_size / numel;
size_t dim_size = dim_el * dtype_size;
int dim_count = data_size / dim_size;
auto data_ptr = (char*)(data.data_ptr());
auto result_ptr = (char*)(result.data_ptr());
for (int i = 0; i < dim_count; i++) {
for (size_t offset = 0; offset < dim_size; offset += MAX_BUF_SIZE) {
size_t chunk_size = dim_size - offset > MAX_BUF_SIZE ? MAX_BUF_SIZE : dim_size - offset;
size_t chunk_el = chunk_size / dtype_size;
naive_all_gather<STATE_GROUP>(
result_ptr + i * dim_size * world_size + offset,
data_ptr + i * dim_size + offset,
dim_size,
chunk_size,
chunk_el);
}
}
return result;
}
template torch::Tensor& all_gather<STATE_GROUP_ALL_GATHER>(torch::Tensor&, torch::Tensor&, int, size_t, int);
template torch::Tensor&
all_gather<STATE_GROUP_ALL_GATHER_INTO_TENSOR>(torch::Tensor&, torch::Tensor&, int, size_t, int);
void naive_reduce_scatter(
char* output_ptr,
char* data_ptr,
c10::ScalarType scalar_type,
size_t chunk_size,
size_t chunk_el,
int element_size) {
const int state_group = STATE_GROUP_REDUCE_SCATTER;
static int current_buffer = 0;
static int state_idx = 0;
enum coll_state copy_current = coll_reduce_scatter_naive__copy_in_done;
enum coll_state copy_next = coll_alt1_reduce_scatter_naive__copy_in_done;
switch (state_idx) {
case 0:
copy_current = coll_reduce_scatter_naive__copy_in_done;
copy_next = coll_alt1_reduce_scatter_naive__copy_in_done;
break;
case 1:
copy_current = coll_alt1_reduce_scatter_naive__copy_in_done;
copy_next = coll_alt2_reduce_scatter_naive__copy_in_done;
break;
case 2:
copy_current = coll_alt2_reduce_scatter_naive__copy_in_done;
copy_next = coll_reduce_scatter_naive__copy_in_done;
break;
default:
assert(!"Should not get here.");
}
state_idx = (state_idx + 1) % 3;
// Step 1: copy local data to shared buffer
parallel_memcpy(reduce_scatter_buffer[current_buffer][world_rank], data_ptr, chunk_size);
std::atomic_thread_fence(std::memory_order_release);
workspace[world_rank]->states[state_group] = copy_current;
// Step 2: wait for all ranks to copy in
for (int i = 0; i < world_size; i++) {
if (i != world_rank) wait_buffer_state_until_2(i, copy_current, copy_next, state_group);
}
// // Step 3: do local reduce on this rank’s slice only
int start_el = slice_el_start(chunk_el, world_rank);
// each rank reduce its slice of buffer independently so therre is no need for
// synchronization afterward
reduce_all_buffers(
start_el,
slice_size(chunk_el, world_rank),
scalar_type,
world_rank,
output_ptr -
start_el * element_size, // in reduce_all_buffers, the output_ptr is the buffer for all ranks, but here
// output_ptr is already the local buffer for one rank. Adjust it here.
reduce_scatter_buffer[current_buffer]);
// done
current_buffer = 1 - current_buffer;
}
void reduce_scatter_outer_loop(torch::Tensor& output, torch::Tensor& data, size_t numel, int data_size) {
for (int offset = 0; offset < data_size; offset += MAX_BUF_SIZE) {
auto data_ptr = ((char*)(data.data_ptr()) + offset);
auto output_ptr = ((char*)(output.data_ptr()) + offset);
size_t chunk_size = std::min((size_t)MAX_BUF_SIZE, (size_t)(data_size - offset));
size_t chunk_el = chunk_size / (data_size / numel);
naive_reduce_scatter(output_ptr, data_ptr, data.scalar_type(), chunk_size, chunk_el, data.element_size());
}
}
python/sglang/kernels/aot/csrc/cpu/shm.h
+19
View File
@@ -0,0 +1,19 @@
#include <torch/all.h>
#include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
#ifndef __SHM_COLLECTIVES__
#define __SHM_COLLECTIVES__
constexpr int STATE_GROUP_SYMMETRIC_ALLREDUCE = 0;
constexpr int STATE_GROUP_DISTRIBUTED_ALLREDUCE = 1;
constexpr int STATE_GROUP_ALL_GATHER = 2;
constexpr int STATE_GROUP_ALL_GATHER_INTO_TENSOR = 3;
constexpr int STATE_GROUP_REDUCE_SCATTER = 4;
void shm_initialize(int size, int rank, const char* addr_string, const char* port_string);
void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size);
template <int STATE_GROUP>
torch::Tensor& all_gather(torch::Tensor& result, torch::Tensor& data, int dim, size_t numel, int data_size);
void reduce_scatter_outer_loop(torch::Tensor& output, torch::Tensor& data, size_t numel, int data_size);
#endif
python/sglang/kernels/aot/csrc/cpu/spec.cpp
+854
View File
@@ -0,0 +1,854 @@
#include "common.h"
namespace {
// Contract shared by every kernel in this file: all tensors are dense,
// contiguous CPU tensors (checked below), so strides are the canonical
// row-major ones; per-function comments list shapes and dtypes only.
// `index_t` params accept int32 or int64 via AT_DISPATCH_INDEX_TYPES so
// callers never pay a dtype-conversion copy.
template <typename rpi_t, typename off_t>
void assign_req_to_token_pool_kernel_impl(
const rpi_t* __restrict__ req_pool_indices,
int32_t* __restrict__ req_to_token,
const off_t* __restrict__ start_offset,
const off_t* __restrict__ end_offset,
const int64_t* __restrict__ out_cache_loc,
int64_t num_cache_locs,
int64_t batch_size,
int64_t pool_len) {
// Pre-compute exclusive prefix sum of (end - start) to avoid O(N^2) work.
std::vector<int64_t> prefix(batch_size + 1, 0);
for (int64_t i = 0; i < batch_size; ++i) {
prefix[i + 1] = prefix[i] + (end_offset[i] - start_offset[i]);
}
TORCH_CHECK(
prefix[batch_size] <= num_cache_locs,
"assign_req_to_token_pool: out_cache_loc has ",
num_cache_locs,
" entries but offsets require ",
prefix[batch_size]);
at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
for (int64_t pid = begin; pid < end; ++pid) {
int64_t kv_start = start_offset[pid];
int64_t kv_end = end_offset[pid];
int32_t* token_pool = req_to_token + req_pool_indices[pid] * pool_len;
int64_t out_offset = prefix[pid];
for (int64_t j = kv_start; j < kv_end; ++j) {
token_pool[j] = static_cast<int32_t>(out_cache_loc[out_offset + (j - kv_start)]);
}
}
});
}
template <typename index_t>
void verify_tree_greedy_kernel_impl(
int32_t* __restrict__ predicts,
int32_t* __restrict__ accept_index,
int32_t* __restrict__ accept_token_num,
const index_t* __restrict__ candidates,
const index_t* __restrict__ retrive_index,
const index_t* __restrict__ retrive_next_token,
const index_t* __restrict__ retrive_next_sibling,
const index_t* __restrict__ target_predict,
int64_t batch_size,
int64_t num_spec_step,
int64_t num_draft_tokens) {
at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
for (int64_t bx = begin; bx < end; ++bx) {
int64_t off = bx * num_draft_tokens;
int64_t ai_off = bx * num_spec_step;
int64_t last_accept_index = retrive_index[off]; // retrive_index[bx, 0]
accept_index[ai_off] = static_cast<int32_t>(last_accept_index);
int32_t num_correct_drafts = 0;
int64_t cur = 0;
for (int64_t j = 1; j < num_spec_step; ++j) {
cur = retrive_next_token[off + cur]; // move to next token
while (cur != -1) {
int64_t draft_idx = retrive_index[off + cur];
int64_t draft_tok = candidates[off + cur];
int64_t target_tok = target_predict[last_accept_index];
if (draft_tok == target_tok) {
predicts[last_accept_index] = static_cast<int32_t>(target_tok);
++num_correct_drafts;
accept_index[ai_off + num_correct_drafts] = static_cast<int32_t>(draft_idx);
last_accept_index = draft_idx;
break;
}
cur = retrive_next_sibling[off + cur]; // try sibling
}
if (cur == -1) break;
}
accept_token_num[bx] = num_correct_drafts;
predicts[last_accept_index] = static_cast<int32_t>(target_predict[last_accept_index]);
}
});
}
// Find the node index in `selected_index[bid]` holding `token_idx`; -1 when the
// tree is malformed and the parent is absent (callers warn and stop the walk,
// mirroring the CUDA kernel's "invalid eagle tree" printf).
template <typename index_t>
int64_t
find_parent_node(const index_t* __restrict__ selected_index, int64_t row_off, int64_t sel_stride, int64_t token_idx) {
for (int64_t i = 0; i < sel_stride; ++i) {
if (selected_index[row_off + i] == token_idx) {
return i;
}
}
return -1;
}
template <typename index_t>
void build_tree_kernel_efficient_impl(
const index_t* __restrict__ parent_list,
const index_t* __restrict__ selected_index,
const index_t* __restrict__ verified_seq_len,
bool* __restrict__ tree_mask,
index_t* __restrict__ positions,
index_t* __restrict__ retrive_index,
index_t* __restrict__ retrive_next_token,
index_t* __restrict__ retrive_next_sibling,
int64_t bs,
int64_t topk,
int64_t depth,
int64_t draft_token_num,
int64_t tree_mask_mode) {
int64_t parent_stride = topk * (depth - 1) + 1;
int64_t sel_stride = draft_token_num - 1;
// FULL_MASK row offsets depend on a prefix sum over verified_seq_len;
// precompute it so the batch loop can run in parallel.
std::vector<int64_t> mask_offsets(bs, 0);
if (tree_mask_mode == 0) { // FULL_MASK
int64_t acc = 0;
for (int64_t i = 0; i < bs; ++i) {
mask_offsets[i] = i * draft_token_num * draft_token_num + acc;
acc += static_cast<int64_t>(verified_seq_len[i]) * draft_token_num;
}
}
at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
for (int64_t bid = begin; bid < end; ++bid) {
int64_t off = bid * draft_token_num;
int64_t sel_off = bid * sel_stride;
int64_t seq_len = verified_seq_len[bid];
// tid == 0 logic: build retrive_index, retrive_next_token, retrive_next_sibling
positions[off] = seq_len;
retrive_index[off] = off; // retrive_index[bid, 0] = bid * draft_token_num
for (int64_t i = draft_token_num - 1; i > 0; --i) {
retrive_index[off + i] = off + i;
int64_t parent_tb_idx = selected_index[sel_off + i - 1] / topk;
int64_t parent_position = 0;
if (parent_tb_idx > 0) {
int64_t parent_token_idx = parent_list[bid * parent_stride + parent_tb_idx];
int64_t found = find_parent_node(selected_index, sel_off, sel_stride, parent_token_idx);
if (found < 0) {
TORCH_WARN("build_tree_kernel_efficient_cpu: invalid eagle tree, parent of node ", i, " not found");
continue; // skip invalid
}
parent_position = found + 1;
}
if (retrive_next_token[off + parent_position] == -1) {
retrive_next_token[off + parent_position] = i;
} else {
int64_t origin = retrive_next_token[off + parent_position];
retrive_next_token[off + parent_position] = i;
retrive_next_sibling[off + i] = origin;
}
}
// Build tree_mask and positions for tid > 0
if (tree_mask_mode == 1) { // QLEN_ONLY
int64_t mask_stride = draft_token_num;
for (int64_t tid = 0; tid < draft_token_num; ++tid) {
int64_t row_start = (off + tid) * mask_stride;
tree_mask[row_start] = true; // attend to the root token (column 0)
for (int64_t j = 1; j < draft_token_num; ++j) {
tree_mask[row_start + j] = false;
}
if (tid == 0) {
continue;
}
int64_t position = 0;
int64_t cur = tid - 1;
// A valid root-ward walk has at most `depth` steps; the bound turns a
// malformed (cyclic) tree into a warning instead of a scheduler hang.
while (position < depth) {
position++;
tree_mask[row_start + cur + 1] = true;
int64_t ptb = selected_index[sel_off + cur] / topk;
if (ptb == 0) break;
int64_t tok_idx = parent_list[bid * parent_stride + ptb];
cur = find_parent_node(selected_index, sel_off, sel_stride, tok_idx);
if (cur < 0) {
TORCH_WARN("build_tree_kernel_efficient_cpu: invalid eagle tree, ancestor of node ", tid, " not found");
break; // stop the walk on a malformed tree
}
}
positions[off + tid] = position + seq_len;
}
} else { // FULL_MASK (mode 0)
// Full mask includes the seq_len prefix
int64_t seq_tree_idx = mask_offsets[bid];
for (int64_t tid = 0; tid < draft_token_num; ++tid) {
int64_t row_start = seq_tree_idx + (seq_len + draft_token_num) * tid + seq_len;
tree_mask[row_start] = true; // attend to the root token (column 0)
for (int64_t j = 1; j < draft_token_num; ++j) {
tree_mask[row_start + j] = false;
}
if (tid == 0) {
continue;
}
int64_t position = 0;
int64_t cur = tid - 1;
// Same depth bound as the QLEN_ONLY branch above.
while (position < depth) {
position++;
tree_mask[row_start + cur + 1] = true;
int64_t ptb = selected_index[sel_off + cur] / topk;
if (ptb == 0) {
break;
}
int64_t tok_idx = parent_list[bid * parent_stride + ptb];
cur = find_parent_node(selected_index, sel_off, sel_stride, tok_idx);
if (cur < 0) {
TORCH_WARN("build_tree_kernel_efficient_cpu: invalid eagle tree, ancestor of node ", tid, " not found");
break; // stop the walk on a malformed tree
}
}
positions[off + tid] = position + seq_len;
}
}
}
});
}
} // anonymous namespace
// Greedy tree verification: walk each request's draft tree, accepting the
// longest root path whose draft tokens match the target model's argmax.
//
// predicts: [bs * num_draft_tokens] int32; out, verified tokens by flat draft index
// accept_index: [bs, num_spec_step] int32; out, flat indices of accepted
// tokens; caller pre-fills with -1 (rejected slots keep it)
// accept_token_num: [bs] int32; out, accepted drafts per request (bonus excluded)
// candidates: [bs, num_draft_tokens] int32 or int64; draft tokens
// retrive_index: [bs, num_draft_tokens] int32 or int64; flat index of each tree node
// retrive_next_token: [bs, num_draft_tokens] int32 or int64; first child, -1 = none
// retrive_next_sibling:[bs, num_draft_tokens] int32 or int64; next sibling, -1 = none
// target_predict: [bs, num_draft_tokens] int32 or int64; target argmax per draft slot
void verify_tree_greedy_cpu(
at::Tensor predicts,
at::Tensor accept_index,
at::Tensor accept_token_num,
const at::Tensor& candidates,
const at::Tensor& retrive_index,
const at::Tensor& retrive_next_token,
const at::Tensor& retrive_next_sibling,
const at::Tensor& target_predict) {
CHECK_INPUT(candidates);
CHECK_DIM(2, candidates);
CHECK_DIM(2, accept_index);
const auto index_dtype = retrive_index.scalar_type();
int64_t batch_size = candidates.size(0);
int64_t num_draft_tokens = candidates.size(1);
int64_t num_spec_step = accept_index.size(1);
CHECK_EQ(candidates.scalar_type(), index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(predicts, {batch_size * num_draft_tokens}, at::kInt);
CHECK_INPUT_SHAPE_DTYPE<false>(accept_index, {batch_size, num_spec_step}, at::kInt);
CHECK_INPUT_SHAPE_DTYPE<false>(accept_token_num, {batch_size}, at::kInt);
CHECK_INPUT_SHAPE_DTYPE<false>(retrive_index, {batch_size, num_draft_tokens}, index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(retrive_next_token, {batch_size, num_draft_tokens}, index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(retrive_next_sibling, {batch_size, num_draft_tokens}, index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(target_predict, {batch_size, num_draft_tokens}, index_dtype);
AT_DISPATCH_INDEX_TYPES(index_dtype, "verify_tree_greedy_indices", [&] {
verify_tree_greedy_kernel_impl<index_t>(
predicts.data_ptr<int32_t>(),
accept_index.data_ptr<int32_t>(),
accept_token_num.data_ptr<int32_t>(),
candidates.data_ptr<index_t>(),
retrive_index.data_ptr<index_t>(),
retrive_next_token.data_ptr<index_t>(),
retrive_next_sibling.data_ptr<index_t>(),
target_predict.data_ptr<index_t>(),
batch_size,
num_spec_step,
num_draft_tokens);
});
}
// Build the draft token tree consumed by target verify: tree attention mask,
// per-token positions, and the retrieval linkage (index / first child /
// next sibling) used by verify_tree_greedy.
//
// parent_list: [bs, topk * (depth - 1) + 1] int32 or int64
// (empty [bs, 0] when depth == 1, e.g. MTP steps=1)
// selected_index: [bs, draft_token_num - 1] int32 or int64
// verified_seq_len: [bs] int32 or int64; committed prefix length per request
// tree_mask: out, bool.
// QLEN_ONLY: [bs * draft_token_num * draft_token_num]; rows
// are fully overwritten here.
// FULL_MASK: [sum_i(seq_len_i * draft_token_num) + bs * draft_token_num^2];
// only each row's qlen block is written -- the caller must
// pre-fill the seq_len prefix columns with true.
// positions: [bs * draft_token_num]; out, same dtype as parent_list
// retrive_index: [bs, draft_token_num]; out
// retrive_next_token: [bs, draft_token_num]; out, pre-filled with -1
// retrive_next_sibling:[bs, draft_token_num]; out, pre-filled with -1
// tree_mask_mode: 0 = FULL_MASK, 1 = QLEN_ONLY (2 = QLEN_ONLY_BITPACKING is rejected)
void build_tree_kernel_efficient_cpu(
const at::Tensor& parent_list,
const at::Tensor& selected_index,
const at::Tensor& verified_seq_len,
at::Tensor tree_mask,
at::Tensor positions,
at::Tensor retrive_index,
at::Tensor retrive_next_token,
at::Tensor retrive_next_sibling,
int64_t topk,
int64_t depth,
int64_t draft_token_num,
int64_t tree_mask_mode) {
CHECK_INPUT(parent_list);
CHECK_DIM(2, parent_list);
// CPU workers always use FULL_MASK (0) or QLEN_ONLY (1); QLEN_ONLY_BITPACKING
// (2) has no CPU producer and any other value is a caller bug.
TORCH_CHECK(
tree_mask_mode == 0 || tree_mask_mode == 1,
"build_tree_kernel_efficient_cpu: only FULL_MASK (0) and QLEN_ONLY (1) are supported, got ",
tree_mask_mode);
const auto index_dtype = parent_list.scalar_type();
int64_t bs = parent_list.size(0);
// depth == 1 (e.g. MTP steps=1) has no non-root parents, so
// organize_draft_results emits an empty (bs, 0) parent_list that the kernel
// never indexes; only the multi-step layout is width topk*(depth-1)+1.
if (depth > 1) {
CHECK_EQ(parent_list.size(1), topk * (depth - 1) + 1);
}
CHECK_INPUT_SHAPE_DTYPE<false>(selected_index, {bs, draft_token_num - 1}, index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(verified_seq_len, {bs}, index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(positions, {bs * draft_token_num}, index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(retrive_index, {bs, draft_token_num}, index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(retrive_next_token, {bs, draft_token_num}, index_dtype);
CHECK_INPUT_SHAPE_DTYPE<false>(retrive_next_sibling, {bs, draft_token_num}, index_dtype);
CHECK_INPUT(tree_mask);
CHECK_EQ(tree_mask.scalar_type(), at::kBool);
if (tree_mask_mode == 1) {
CHECK_EQ(tree_mask.numel(), bs * draft_token_num * draft_token_num);
} else {
int64_t seq_len_sum = verified_seq_len.sum().item<int64_t>();
CHECK_EQ(tree_mask.numel(), (seq_len_sum + bs * draft_token_num) * draft_token_num);
}
AT_DISPATCH_INDEX_TYPES(index_dtype, "build_tree_kernel_efficient_indices", [&] {
build_tree_kernel_efficient_impl<index_t>(
parent_list.data_ptr<index_t>(),
selected_index.data_ptr<index_t>(),
verified_seq_len.data_ptr<index_t>(),
tree_mask.data_ptr<bool>(),
positions.data_ptr<index_t>(),
retrive_index.data_ptr<index_t>(),
retrive_next_token.data_ptr<index_t>(),
retrive_next_sibling.data_ptr<index_t>(),
bs,
topk,
depth,
draft_token_num,
tree_mask_mode);
});
}
// Scatter freshly allocated KV slots into the request-to-token map:
// req_to_token[req_pool_indices[i], start_offset[i]:end_offset[i]] =
// out_cache_loc[prefix[i]:prefix[i+1]].
//
// req_pool_indices: [bs] int32 or int64
// req_to_token: [max_num_reqs, pool_len] int32; out
// start_offset: [bs] int32 or int64 (independent of req_pool_indices;
// eagle_prepare_for_decode passes int64 indices with int32 kv lens)
// end_offset: [bs] same dtype as start_offset
// out_cache_loc: [sum_i(end_offset[i] - start_offset[i])] int64
void assign_req_to_token_pool_cpu(
const at::Tensor& req_pool_indices,
at::Tensor req_to_token,
const at::Tensor& start_offset,
const at::Tensor& end_offset,
const at::Tensor& out_cache_loc,
int64_t pool_len) {
CHECK_INPUT(req_pool_indices);
CHECK_INPUT(req_to_token);
CHECK_INPUT(start_offset);
CHECK_INPUT(end_offset);
CHECK_INPUT(out_cache_loc);
CHECK_DIM(2, req_to_token);
CHECK_EQ(req_to_token.scalar_type(), at::kInt);
CHECK_EQ(out_cache_loc.scalar_type(), at::kLong);
CHECK_EQ(end_offset.scalar_type(), start_offset.scalar_type());
CHECK_EQ(req_to_token.size(1), pool_len);
int64_t batch_size = req_pool_indices.size(0);
CHECK_EQ(start_offset.numel(), batch_size);
CHECK_EQ(end_offset.numel(), batch_size);
AT_DISPATCH_INDEX_TYPES(req_pool_indices.scalar_type(), "assign_req_to_token_pool_rpi", [&] {
using rpi_t = index_t;
const rpi_t* rpi_ptr = req_pool_indices.data_ptr<rpi_t>();
AT_DISPATCH_INDEX_TYPES(start_offset.scalar_type(), "assign_req_to_token_pool_offsets", [&] {
assign_req_to_token_pool_kernel_impl<rpi_t, index_t>(
rpi_ptr,
req_to_token.data_ptr<int32_t>(),
start_offset.data_ptr<index_t>(),
end_offset.data_ptr<index_t>(),
out_cache_loc.data_ptr<int64_t>(),
out_cache_loc.numel(),
batch_size,
pool_len);
});
});
}
// Expand req_to_token for multi-step draft decode: row b*topk+tk holds the
// committed prefix of request b followed by candidate tk's draft slots
// (which assign_draft_cache_locs_contiguous laid out at sl + tk*num_steps).
//
// req_to_token: [max_num_reqs, pool_len] int32
// req_pool_indices: [num_seqs] int32 or int64
// seq_lens: [num_seqs] int32 or int64 (independent of req_pool_indices)
// returns: [num_seqs * topk, pool_len] int32; only the first
// seq_lens[b] + num_steps entries of each row are defined
at::Tensor build_draft_decode_metadata_cpu(
const at::Tensor& req_to_token,
const at::Tensor& req_pool_indices,
const at::Tensor& seq_lens,
int64_t topk,
int64_t num_steps,
int64_t pool_len) {
CHECK_INPUT(req_to_token);
CHECK_INPUT(req_pool_indices);
CHECK_INPUT(seq_lens);
CHECK_DIM(2, req_to_token);
CHECK_EQ(req_to_token.scalar_type(), at::kInt);
CHECK_EQ(req_to_token.size(1), pool_len);
int64_t num_seqs = req_pool_indices.size(0);
int64_t bs = num_seqs * topk;
CHECK_EQ(seq_lens.numel(), num_seqs);
auto req_to_token_draft = at::empty({bs, pool_len}, req_to_token.options());
auto* rtt_ptr = req_to_token.data_ptr<int32_t>();
auto* draft_ptr = req_to_token_draft.data_ptr<int32_t>();
AT_DISPATCH_INDEX_TYPES(req_pool_indices.scalar_type(), "build_draft_decode_metadata_rpi", [&] {
using rpi_t = index_t;
const rpi_t* rpi_ptr = req_pool_indices.data_ptr<rpi_t>();
AT_DISPATCH_INDEX_TYPES(seq_lens.scalar_type(), "build_draft_decode_metadata_lens", [&] {
const index_t* sl_ptr = seq_lens.data_ptr<index_t>();
at::parallel_for(0, num_seqs, 0, [&](int64_t begin, int64_t end) {
for (int64_t b = begin; b < end; ++b) {
int64_t idx = rpi_ptr[b];
int64_t sl = sl_ptr[b];
const int32_t* src_row = rtt_ptr + idx * pool_len;
for (int64_t tk = 0; tk < topk; ++tk) {
int64_t flat = b * topk + tk;
int32_t* dst_row = draft_ptr + flat * pool_len;
// Copy prefix
std::memcpy(dst_row, src_row, sl * sizeof(int32_t));
// Copy draft tokens for this candidate
int64_t draft_start = sl + tk * num_steps;
for (int64_t s = 0; s < num_steps; ++s) {
dst_row[sl + s] = src_row[draft_start + s];
}
}
}
});
});
});
return req_to_token_draft;
}
// Pick the last accepted token of each request as its bonus token.
//
// accept_tokens: [bs, accept_stride] int32; row-major, accept_stride = accept_index.shape[1]
// accept_lens: [bs] int32; number of accepted tokens per request (bonus included)
// bonus_tokens: [bs] int32; out
void fill_bonus_tokens_cpu(
const at::Tensor& accept_tokens, const at::Tensor& accept_lens, at::Tensor bonus_tokens, int64_t accept_stride) {
CHECK_INPUT(accept_tokens);
CHECK_INPUT(accept_lens);
CHECK_INPUT(bonus_tokens);
CHECK_EQ(accept_tokens.scalar_type(), at::kInt);
CHECK_EQ(accept_lens.scalar_type(), at::kInt);
CHECK_EQ(bonus_tokens.scalar_type(), at::kInt);
int64_t bs = accept_lens.size(0);
CHECK_EQ(accept_tokens.numel(), bs * accept_stride);
CHECK_EQ(bonus_tokens.numel(), bs);
auto* accept_ptr = accept_tokens.data_ptr<int32_t>();
auto* al_ptr = accept_lens.data_ptr<int32_t>();
auto* out_ptr = bonus_tokens.data_ptr<int32_t>();
at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
for (int64_t pid = begin; pid < end; ++pid) {
int64_t idx = accept_stride * pid + al_ptr[pid] - 1;
out_ptr[pid] = accept_ptr[idx];
}
});
}
// Compact the accepted tokens' KV slots: gather out_cache_loc at the accepted
// indices, skipping -1 (rejected) entries. Sequential by design: the output
// write position depends on how many prior entries were accepted.
//
// accept_index: [bs * num_spec_step] int32 or int64; flat, -1 = rejected
// out_cache_loc: [bs * num_draft_tokens] int64
// accept_out_cache_loc: [>= num_accept] int64; out, only the first num_accept
// entries are written
void fill_accept_out_cache_loc_cpu(
const at::Tensor& accept_index, const at::Tensor& out_cache_loc, at::Tensor accept_out_cache_loc) {
CHECK_INPUT(accept_index);
CHECK_INPUT(out_cache_loc);
CHECK_INPUT(accept_out_cache_loc);
CHECK_EQ(out_cache_loc.scalar_type(), at::kLong);
CHECK_EQ(accept_out_cache_loc.scalar_type(), at::kLong);
// num_accept <= accept_index.numel(), so this bounds every write below.
CHECK_GE(accept_out_cache_loc.numel(), accept_index.numel());
int64_t num_indices = accept_index.numel();
int64_t num_cache_locs = out_cache_loc.numel();
auto* ocl_ptr = out_cache_loc.data_ptr<int64_t>();
auto* out_ptr = accept_out_cache_loc.data_ptr<int64_t>();
AT_DISPATCH_INDEX_TYPES(accept_index.scalar_type(), "fill_accept_out_cache_loc_indices", [&] {
const index_t* ai_ptr = accept_index.data_ptr<index_t>();
int64_t dst = 0;
for (int64_t i = 0; i < num_indices; ++i) {
int64_t src = static_cast<int64_t>(ai_ptr[i]);
if (src > -1) {
TORCH_CHECK(src < num_cache_locs, "fill_accept_out_cache_loc: accept_index ", src, " out of range");
out_ptr[dst++] = ocl_ptr[src];
}
}
});
}
// Read back the draft KV slots reserved by the allocator: for each request,
// copy the topk*num_steps slots starting at seq_lens[pid] out of req_to_token.
//
// req_pool_indices: [bs] int32 or int64
// req_to_token: [max_num_reqs, pool_len] int32
// seq_lens: [bs] int32 or int64 (independent of req_pool_indices)
// out_cache_loc: [bs * topk * num_steps] int64; out
void assign_draft_cache_locs_contiguous_cpu(
const at::Tensor& req_pool_indices,
const at::Tensor& req_to_token,
const at::Tensor& seq_lens,
at::Tensor out_cache_loc,
int64_t pool_len,
int64_t topk,
int64_t num_steps) {
// Contiguous slot layout: requires page_size == 1 or topk == 1 (see prepare_for_v2_draft guard).
CHECK_INPUT(req_pool_indices);
CHECK_INPUT(req_to_token);
CHECK_INPUT(seq_lens);
CHECK_INPUT(out_cache_loc);
CHECK_DIM(2, req_to_token);
CHECK_EQ(req_to_token.scalar_type(), at::kInt);
CHECK_EQ(out_cache_loc.scalar_type(), at::kLong);
CHECK_EQ(req_to_token.size(1), pool_len);
CHECK_EQ(out_cache_loc.numel(), req_pool_indices.numel() * topk * num_steps);
int64_t bs = req_pool_indices.size(0);
int64_t copy_len = topk * num_steps;
CHECK_EQ(seq_lens.numel(), bs);
auto* rtt_ptr = req_to_token.data_ptr<int32_t>();
auto* out_ptr = out_cache_loc.data_ptr<int64_t>();
AT_DISPATCH_INDEX_TYPES(req_pool_indices.scalar_type(), "assign_draft_cache_locs_contiguous_rpi", [&] {
using rpi_t = index_t;
const rpi_t* rpi_ptr = req_pool_indices.data_ptr<rpi_t>();
AT_DISPATCH_INDEX_TYPES(seq_lens.scalar_type(), "assign_draft_cache_locs_contiguous_lens", [&] {
const index_t* sl_ptr = seq_lens.data_ptr<index_t>();
at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
for (int64_t pid = begin; pid < end; ++pid) {
int64_t kv_start = sl_ptr[pid];
int64_t req_idx = rpi_ptr[pid];
const int32_t* src = rtt_ptr + req_idx * pool_len + kv_start;
int64_t* dst = out_ptr + pid * copy_len;
for (int64_t j = 0; j < copy_len; ++j) {
dst[j] = static_cast<int64_t>(src[j]);
}
}
});
});
});
}
// Gather each request's KV slots in [start_offset, end_offset) out of
// req_to_token into a dense int64 vector (verify/extend cache locations).
//
// req_pool_indices: [bs] int32 or int64
// req_to_token: [max_num_reqs, pool_len] int32
// start_offset: [bs] int32 or int64 (independent of req_pool_indices)
// end_offset: [bs] same dtype as start_offset
// out_cache_loc: [sum_i(end_offset[i] - start_offset[i])] int64; out
void assign_extend_cache_locs_cpu(
const at::Tensor& req_pool_indices,
const at::Tensor& req_to_token,
const at::Tensor& start_offset,
const at::Tensor& end_offset,
at::Tensor out_cache_loc,
int64_t pool_len) {
CHECK_INPUT(req_pool_indices);
CHECK_INPUT(req_to_token);
CHECK_INPUT(start_offset);
CHECK_INPUT(end_offset);
CHECK_INPUT(out_cache_loc);
CHECK_DIM(2, req_to_token);
CHECK_EQ(req_to_token.scalar_type(), at::kInt);
CHECK_EQ(out_cache_loc.scalar_type(), at::kLong);
CHECK_EQ(end_offset.scalar_type(), start_offset.scalar_type());
CHECK_EQ(req_to_token.size(1), pool_len);
int64_t bs = req_pool_indices.size(0);
CHECK_EQ(start_offset.numel(), bs);
CHECK_EQ(end_offset.numel(), bs);
auto* rtt_ptr = req_to_token.data_ptr<int32_t>();
auto* out_ptr = out_cache_loc.data_ptr<int64_t>();
AT_DISPATCH_INDEX_TYPES(req_pool_indices.scalar_type(), "assign_extend_cache_locs_rpi", [&] {
using rpi_t = index_t;
const rpi_t* rpi_ptr = req_pool_indices.data_ptr<rpi_t>();
AT_DISPATCH_INDEX_TYPES(start_offset.scalar_type(), "assign_extend_cache_locs_offsets", [&] {
const index_t* start_ptr = start_offset.data_ptr<index_t>();
const index_t* end_ptr = end_offset.data_ptr<index_t>();
// Compute prefix sum for output offsets (sequential)
std::vector<int64_t> out_offsets(bs + 1, 0);
for (int64_t i = 0; i < bs; ++i) {
out_offsets[i + 1] = out_offsets[i] + (end_ptr[i] - start_ptr[i]);
}
// Callers may size out_cache_loc at max capacity (e.g. bs * num_spec_step
// in move_accept_tokens) and leave the tail untouched, hence <= not ==.
TORCH_CHECK(
out_offsets[bs] <= out_cache_loc.numel(),
"assign_extend_cache_locs: out_cache_loc has ",
out_cache_loc.numel(),
" entries but offsets require ",
out_offsets[bs]);
at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
for (int64_t pid = begin; pid < end; ++pid) {
int64_t kv_start = start_ptr[pid];
int64_t kv_end = end_ptr[pid];
int64_t req_idx = rpi_ptr[pid];
int64_t length = kv_end - kv_start;
const int32_t* src = rtt_ptr + req_idx * pool_len + kv_start;
int64_t* dst = out_ptr + out_offsets[pid];
for (int64_t j = 0; j < length; ++j) {
dst[j] = static_cast<int64_t>(src[j]);
}
}
});
});
});
}
// Recover tree linkage from a QLEN-layout boolean tree mask (NGRAM path):
// depth/position, retrieval index, first child and next sibling per node.
//
// tree_mask: [bs * draft_token_num * draft_token_num] bool
// verified_seq_len: [bs] int32 or int64
// positions: [bs * draft_token_num]; out, same dtype as verified_seq_len
// retrive_index: [bs, draft_token_num]; out
// retrive_next_token: [bs, draft_token_num]; out
// retrive_next_sibling:[bs, draft_token_num]; out
void reconstruct_indices_from_tree_mask_cpu(
const at::Tensor& tree_mask,
const at::Tensor& verified_seq_len,
at::Tensor positions,
at::Tensor retrive_index,
at::Tensor retrive_next_token,
at::Tensor retrive_next_sibling,
int64_t batch_size,
int64_t draft_token_num) {
CHECK_INPUT(tree_mask);
CHECK_INPUT(verified_seq_len);
CHECK_INPUT(positions);
CHECK_INPUT(retrive_index);
CHECK_INPUT(retrive_next_token);
CHECK_INPUT(retrive_next_sibling);
CHECK_EQ(tree_mask.scalar_type(), at::kBool);
CHECK_EQ(tree_mask.numel(), batch_size * draft_token_num * draft_token_num);
CHECK_EQ(verified_seq_len.numel(), batch_size);
CHECK_EQ(positions.numel(), batch_size * draft_token_num);
CHECK_EQ(retrive_index.numel(), batch_size * draft_token_num);
CHECK_EQ(retrive_next_token.numel(), batch_size * draft_token_num);
CHECK_EQ(retrive_next_sibling.numel(), batch_size * draft_token_num);
const auto index_dtype = verified_seq_len.scalar_type();
CHECK_EQ(positions.scalar_type(), index_dtype);
CHECK_EQ(retrive_index.scalar_type(), index_dtype);
CHECK_EQ(retrive_next_token.scalar_type(), index_dtype);
CHECK_EQ(retrive_next_sibling.scalar_type(), index_dtype);
const bool* mask_ptr = tree_mask.data_ptr<bool>();
int64_t base_offset = draft_token_num * draft_token_num;
AT_DISPATCH_INDEX_TYPES(index_dtype, "reconstruct_indices_from_tree_mask_indices", [&] {
const index_t* seq_len_ptr = verified_seq_len.data_ptr<index_t>();
index_t* pos_ptr = positions.data_ptr<index_t>();
index_t* ri_ptr = retrive_index.data_ptr<index_t>();
index_t* rnt_ptr = retrive_next_token.data_ptr<index_t>();
index_t* rns_ptr = retrive_next_sibling.data_ptr<index_t>();
at::parallel_for(0, batch_size * draft_token_num, 0, [&](int64_t begin, int64_t end) {
for (int64_t idx = begin; idx < end; ++idx) {
int64_t bid = idx / draft_token_num;
int64_t tid = idx % draft_token_num;
int64_t token_idx = bid * draft_token_num;
int64_t tree_mask_offset = bid * base_offset;
// Step 1: depth and parent via backward scan
int64_t depth = 0;
int64_t parent_idx = -1;
for (int64_t i = tid - 1, start_idx = tree_mask_offset + tid * draft_token_num; i >= 0; --i) {
if (mask_ptr[start_idx + i]) {
depth++;
if (parent_idx == -1) {
parent_idx = i;
}
}
}
// Step 2: retrive_index (identity)
ri_ptr[token_idx + tid] = token_idx + tid;
// Step 3: position = depth + verified_seq_len
pos_ptr[token_idx + tid] = depth + seq_len_ptr[bid];
// Step 4: first child (next_token)
int64_t next_token_idx = -1;
for (int64_t i = tid + 1; i < draft_token_num; ++i) {
if (mask_ptr[tree_mask_offset + i * draft_token_num + tid]) {
next_token_idx = i;
break;
}
}
rnt_ptr[token_idx + tid] = next_token_idx;
// Step 5: next sibling (shares parent, no intervening ancestors)
int64_t next_sibling_idx = -1;
if (parent_idx != -1) {
for (int64_t i = tid + 1; i < draft_token_num; ++i) {
int64_t si = tree_mask_offset + i * draft_token_num + parent_idx;
if (mask_ptr[si]) {
bool is_sibling = true;
int64_t ei = tree_mask_offset + i * draft_token_num + i;
for (int64_t j = si + 1; j < ei; ++j) {
if (mask_ptr[j]) {
is_sibling = false;
break;
}
}
if (is_sibling) {
next_sibling_idx = i;
break;
}
}
}
}
rns_ptr[token_idx + tid] = next_sibling_idx;
}
});
});
}
// Shift each request's extend segment left by one token and write the new
// draft token at the end (or at select_index when given). Mutates input_ids
// in place; callers rely on this.
//
// input_ids: [num_extend_tokens] int64; in/out
// extend_start_loc: [bs] int32 or int64
// extend_seq_lens: [bs] int32 or int64 (independent of extend_start_loc; the
// spec decode-extend batch pairs int64 lens with int32 locs)
// topk_index: [bs] int64; new draft token per request
// select_index: [bs] int64 or None; global slot for the new token
void rotate_input_ids_cpu(
at::Tensor input_ids,
const at::Tensor& extend_start_loc,
const at::Tensor& extend_seq_lens,
const at::Tensor& topk_index,
const std::optional<at::Tensor>& select_index_opt) {
CHECK_INPUT(input_ids);
CHECK_INPUT(extend_start_loc);
CHECK_INPUT(extend_seq_lens);
CHECK_INPUT(topk_index);
CHECK_EQ(input_ids.scalar_type(), at::kLong);
CHECK_EQ(topk_index.scalar_type(), at::kLong);
int64_t bs = extend_seq_lens.size(0);
CHECK_EQ(extend_start_loc.numel(), bs);
CHECK_EQ(topk_index.numel(), bs);
if (select_index_opt.has_value()) {
CHECK_INPUT(select_index_opt.value());
CHECK_EQ(select_index_opt.value().scalar_type(), at::kLong);
CHECK_EQ(select_index_opt.value().numel(), bs);
}
auto* ids_ptr = input_ids.data_ptr<int64_t>();
auto* topk_ptr = topk_index.data_ptr<int64_t>();
const int64_t* select_ptr = conditional_data_ptr<int64_t>(select_index_opt);
AT_DISPATCH_INDEX_TYPES(extend_start_loc.scalar_type(), "rotate_input_ids_start", [&] {
using start_t = index_t;
const start_t* start_ptr = extend_start_loc.data_ptr<start_t>();
AT_DISPATCH_INDEX_TYPES(extend_seq_lens.scalar_type(), "rotate_input_ids_lens", [&] {
const index_t* lens_ptr = extend_seq_lens.data_ptr<index_t>();
at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
for (int64_t pid = begin; pid < end; ++pid) {
int64_t start = start_ptr[pid];
int64_t seq_len = lens_ptr[pid];
int64_t new_token = topk_ptr[pid];
// Shift left by 1
if (seq_len > 1) {
std::memmove(ids_ptr + start, ids_ptr + start + 1, (seq_len - 1) * sizeof(int64_t));
}
// Write new token
if (seq_len > 0) {
if (select_ptr != nullptr) {
ids_ptr[select_ptr[pid]] = new_token;
} else {
ids_ptr[start + seq_len - 1] = new_token;
}
}
}
});
});
});
}
python/sglang/kernels/aot/csrc/cpu/topk.cpp
+693
View File
@@ -0,0 +1,693 @@
#include "common.h"
#include "vec.h"
namespace {
template <typename scalar_t, int SIZE>
inline void softmax(float* __restrict__ out, const scalar_t* __restrict__ input) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
// step 1: get max
fVec max_fvec = fVec(-std::numeric_limits<float>::infinity());
if constexpr (SIZE < kVecSize) {
// SIZE = 1, 2, 4, 8, 16; only the top half is used
bVec x_bvec = bVec::loadu(input, SIZE);
fVec x_fvec0, x_fvec1;
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
x_fvec0 = fVec::set(max_fvec, x_fvec0, SIZE);
max_fvec = at::vec::maximum(max_fvec, x_fvec0);
x_fvec0.store(out, SIZE);
} else {
for (int d = 0; d < SIZE; d += kVecSize) {
bVec x_bvec = bVec::loadu(input + d);
fVec x_fvec0, x_fvec1;
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
max_fvec = at::vec::maximum(max_fvec, x_fvec0);
max_fvec = at::vec::maximum(max_fvec, x_fvec1);
x_fvec0.store(out + d);
x_fvec1.store(out + d + fVec::size());
}
}
float max_val = vec_reduce_max(max_fvec);
max_fvec = fVec(max_val);
// step 2: sum of (x - max).exp()
fVec sum_fvec = fVec(float(0));
if constexpr (SIZE < fVec::size()) {
// SIZE = 1, 2, 4, 8
fVec x_fvec = (fVec::loadu(out, SIZE) - max_fvec).exp_u20();
x_fvec = fVec::set(sum_fvec, x_fvec, SIZE);
sum_fvec += x_fvec;
x_fvec.store(out, SIZE);
} else {
for (int d = 0; d < SIZE; d += fVec::size()) {
fVec x_fvec = (fVec::loadu(out + d) - max_fvec).exp_u20();
sum_fvec += x_fvec;
x_fvec.store(out + d);
}
}
float sum_val = vec_reduce_sum(sum_fvec);
// step 3: x * (1 / sum)
sum_fvec = fVec(1.f / sum_val);
if constexpr (SIZE < fVec::size()) {
// SIZE = 1, 2, 4, 8
fVec out_fvec = fVec::loadu(out, SIZE) * sum_fvec;
out_fvec.store(out, SIZE);
} else {
for (int d = 0; d < SIZE; d += fVec::size()) {
fVec out_fvec = fVec::loadu(out + d) * sum_fvec;
out_fvec.store(out + d);
}
}
}
template <typename scalar_t, int NUM_EXPERTS>
void grouped_topk_kernel_impl(
float* __restrict__ topk_weights,
int32_t* __restrict__ topk_ids,
const scalar_t* __restrict__ gating_output,
int64_t num_tokens,
int64_t topk,
int64_t num_groups,
int64_t topk_group,
bool renormalize) {
const int64_t num_experts_per_group = NUM_EXPERTS / num_groups;
at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
alignas(64) float scores[NUM_EXPERTS];
using elem_t = std::pair<float, int32_t>;
std::vector<elem_t> queue(num_groups);
std::vector<elem_t> queue2(topk_group * num_experts_per_group);
for (int64_t i = begin; i < end; ++i) {
// do softmax to get scores
softmax<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
// find max score per group
for (int64_t g = 0; g < num_groups; ++g) {
float gmax = -std::numeric_limits<float>::infinity();
for (int64_t e = 0; e < num_experts_per_group; ++e) {
gmax = std::max(gmax, scores[g * num_experts_per_group + e]);
}
queue[g] = {gmax, g};
}
// find group topk
std::partial_sort(
queue.begin(), queue.begin() + topk_group, queue.end(), [](const elem_t& x, const elem_t& y) -> bool {
return x.first > y.first;
});
for (int64_t g = 0; g < topk_group; ++g) {
int32_t group_idx = queue[g].second;
for (int64_t e = 0; e < num_experts_per_group; ++e) {
int32_t expert_idx = group_idx * num_experts_per_group + e;
queue2[g * num_experts_per_group + e] = {scores[expert_idx], expert_idx};
}
}
// find global topk
std::partial_sort(
queue2.begin(), queue2.begin() + topk, queue2.end(), [](const elem_t& x, const elem_t& y) -> bool {
return x.first > y.first;
});
for (int64_t j = 0; j < topk; ++j) {
topk_weights[i * topk + j] = queue2[j].first;
topk_ids[i * topk + j] = queue2[j].second;
}
if (renormalize) {
float sum = 0.f;
for (int64_t j = 0; j < topk; ++j) {
sum += topk_weights[i * topk + j];
}
float scale = 1.f / sum;
for (int64_t j = 0; j < topk; ++j) {
topk_weights[i * topk + j] *= scale;
}
}
}
});
}
template <typename scalar_t, int SIZE, std::enable_if_t<!std::is_same_v<scalar_t, float>, int> = 0>
inline void sigmoid(float* __restrict__ out, const scalar_t* __restrict__ input) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
constexpr int kVecSize = bVec::size();
for (int d = 0; d < SIZE; d += kVecSize) {
bVec x_bvec = bVec::loadu(input + d);
fVec x_fvec0, x_fvec1;
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
x_fvec0 = one / (one + x_fvec0.neg().exp_u20());
x_fvec1 = one / (one + x_fvec1.neg().exp_u20());
x_fvec0.store(out + d);
x_fvec1.store(out + d + fVec::size());
}
}
template <typename scalar_t, int SIZE, std::enable_if_t<std::is_same_v<scalar_t, float>, int> = 0>
inline void sigmoid(float* __restrict__ out, const float* __restrict__ input) {
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
constexpr int kVecSize = fVec::size();
for (int d = 0; d < SIZE; d += kVecSize) {
fVec in_fvec = fVec::loadu(input + d);
in_fvec = one / (one + in_fvec.neg().exp_u20());
in_fvec.store(out + d);
}
}
template <typename scalar_t, int NUM_EXPERTS>
void topk_sigmoid_kernel_impl(
float* __restrict__ topk_weights,
int32_t* __restrict__ topk_ids,
const scalar_t* __restrict__ gating_output,
int64_t num_tokens,
int64_t topk,
bool renormalize) {
using Vec = at::vec::Vectorized<float>;
const int64_t num_experts_per_group = NUM_EXPERTS;
at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
alignas(64) float scores[NUM_EXPERTS];
using elem_t = std::pair<float, int32_t>;
std::vector<elem_t> queue(num_experts_per_group);
for (int64_t i = begin; i < end; ++i) {
at::vec::convert<scalar_t, float>(gating_output + i * NUM_EXPERTS, scores, NUM_EXPERTS);
float gmax = at::vec::reduce_all<float>(
[](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, scores, num_experts_per_group);
// find position of first max,
// note that we may have multiple max values.
int first_max_idx = -1;
for (int64_t e = 0; e < num_experts_per_group; ++e) {
if (scores[e] == gmax) {
first_max_idx = e;
break;
}
}
// scalar sigmoid
topk_weights[i] = 1.0 / (1.0 + exp(0.0 - gmax));
topk_ids[i] = first_max_idx;
if (renormalize) {
float sum = 0.f;
for (int64_t j = 0; j < topk; ++j) {
sum += topk_weights[i * topk + j];
}
float scale = 1.f / sum;
for (int64_t j = 0; j < topk; ++j) {
topk_weights[i * topk + j] *= scale;
}
}
}
});
}
template <typename scalar_t, int NUM_EXPERTS>
void topk_softmax_kernel_impl(
float* __restrict__ topk_weights,
int32_t* __restrict__ topk_ids,
const scalar_t* __restrict__ gating_output,
int64_t num_tokens,
int64_t topk,
bool renormalize) {
const int64_t num_experts_per_group = NUM_EXPERTS;
at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
alignas(64) float scores[NUM_EXPERTS];
using elem_t = std::pair<float, int32_t>;
std::vector<elem_t> queue(num_experts_per_group);
for (int64_t i = begin; i < end; ++i) {
softmax<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
for (int64_t e = 0; e < num_experts_per_group; ++e) {
queue[e] = {scores[e], e};
}
std::partial_sort(queue.begin(), queue.begin() + topk, queue.end(), [](const elem_t& x, const elem_t& y) -> bool {
return x.first > y.first;
});
for (int64_t j = 0; j < topk; ++j) {
topk_weights[i * topk + j] = queue[j].first;
topk_ids[i * topk + j] = queue[j].second;
}
if (renormalize) {
float sum = 0.f;
for (int64_t j = 0; j < topk; ++j) {
sum += topk_weights[i * topk + j];
}
float scale = 1.f / sum;
for (int64_t j = 0; j < topk; ++j) {
topk_weights[i * topk + j] *= scale;
}
}
}
});
}
template <typename param_t, int SIZE>
inline void
apply_bias(float* __restrict__ scores2, const float* __restrict__ scores, const param_t* __restrict__ bias) {
using fVec = at::vec::Vectorized<float>;
auto vec_size = fVec::size() * 2;
int d = 0;
for (; d <= SIZE - vec_size; d += vec_size) {
fVec bias0, bias1, x0, x1;
std::tie(bias0, bias1) = load_float_vec2(bias + d);
std::tie(x0, x1) = load_float_vec2(scores + d);
x0 = x0 + bias0;
x1 = x1 + bias1;
x0.store(scores2 + d);
x1.store(scores2 + d + fVec::size());
}
for (; d < SIZE; d++) {
scores2[d] = scores[d] + (float)bias[d];
}
}
template <typename scalar_t, typename param_t, int NUM_EXPERTS, int TOPK>
void biased_grouped_topk_kernel_impl(
float* __restrict__ topk_weights,
int32_t* __restrict__ topk_ids,
scalar_t* __restrict__ gating_output,
const param_t* __restrict__ bias,
float scaling_factor_value,
int64_t num_tokens,
int64_t num_groups,
int64_t topk_group,
bool renormalize) {
using Vec = at::vec::Vectorized<float>;
bool apply_scaling_factor = scaling_factor_value != 1.0f;
const int64_t num_experts_per_group = NUM_EXPERTS / num_groups;
at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
// scores: sigmoid
alignas(64) float scores[NUM_EXPERTS];
// scores for choice: sigmoid + bias
alignas(64) float scores2[NUM_EXPERTS];
using elem_t = std::pair<float, int32_t>;
std::vector<elem_t> queue(num_groups);
std::vector<elem_t> queue2(topk_group * num_experts_per_group);
for (int64_t i = begin; i < end; ++i) {
// do sigmoid to get scores
sigmoid<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
apply_bias<param_t, NUM_EXPERTS>(scores2, scores, bias);
for (int64_t g = 0; g < num_groups; ++g) {
// find the max
float gmax = at::vec::reduce_all<float>(
[](Vec& x, Vec& y) { return at::vec::maximum(x, y); },
scores2 + g * num_experts_per_group,
num_experts_per_group);
// find position of first max,
// note that we may have multiple max values.
int first_max_idx = -1;
for (int64_t e = 0; e < num_experts_per_group; ++e) {
if (scores2[g * num_experts_per_group + e] == gmax) {
first_max_idx = g * num_experts_per_group + e;
break;
}
}
// find the 2nd max
scores2[first_max_idx] = -std::numeric_limits<float>::infinity();
float gmax2 = at::vec::reduce_all<float>(
[](Vec& x, Vec& y) { return at::vec::maximum(x, y); },
scores2 + g * num_experts_per_group,
num_experts_per_group);
// restore scores for choice
scores2[first_max_idx] = gmax;
queue[g] = {gmax + gmax2, g};
}
// find group topk
std::partial_sort(
queue.begin(), queue.begin() + topk_group, queue.end(), [](const elem_t& x, const elem_t& y) -> bool {
return x.first > y.first;
});
for (int64_t g = 0; g < topk_group; ++g) {
int32_t group_idx = queue[g].second;
for (int64_t e = 0; e < num_experts_per_group; ++e) {
int32_t expert_idx = group_idx * num_experts_per_group + e;
queue2[g * num_experts_per_group + e] = {scores2[expert_idx], expert_idx};
}
}
// find global topk
std::partial_sort(
queue2.begin(), queue2.begin() + TOPK, queue2.end(), [](const elem_t& x, const elem_t& y) -> bool {
return x.first > y.first;
});
for (int j = 0; j < TOPK; ++j) {
int32_t index = queue2[j].second;
topk_ids[i * TOPK + j] = index;
topk_weights[i * TOPK + j] = scores[index];
}
#if defined(CPU_CAPABILITY_AVX512)
if (renormalize || apply_scaling_factor) {
__mmask16 mask = (1ULL << TOPK) - 1;
__m512 x = _mm512_maskz_loadu_ps(mask, topk_weights + i * TOPK);
if (renormalize) {
float sum = _mm512_reduce_add_ps(x);
__m512 vscale = _mm512_set1_ps(scaling_factor_value / sum);
__m512 y = _mm512_mul_ps(x, vscale);
_mm512_mask_storeu_ps(topk_weights + i * TOPK, mask, y);
} else {
__m512 vscale = _mm512_set1_ps(scaling_factor_value);
__m512 y = _mm512_mul_ps(x, vscale);
_mm512_mask_storeu_ps(topk_weights + i * TOPK, mask, y);
}
}
#else
if (renormalize || apply_scaling_factor){
if (renormalize) {
float sum = 0.f;
for (int64_t j = 0; j < TOPK; ++j) {
sum += topk_weights[i * TOPK + j];
}
float scale = scaling_factor_value / sum;
for (int64_t j = 0; j < TOPK; ++j) {
topk_weights[i * TOPK + j] *= scale;
}
}else{
for (int64_t j = 0; j < TOPK; ++j) {
topk_weights[i * TOPK + j] *= scaling_factor_value;
}
}
}
#endif
}
});
}
#define LAUNCH_GROUPED_TOPK_KERNEL(NE) \
grouped_topk_kernel_impl<scalar_t, NE>( \
topk_weights.data_ptr<float>(), \
topk_ids.data_ptr<int32_t>(), \
gating_output.data_ptr<scalar_t>(), \
num_tokens, \
topk, \
num_expert_group, \
topk_group, \
renormalize);
#define LAUNCH_TOPK_SIGMOID_KERNEL(NE) \
topk_sigmoid_kernel_impl<scalar_t, NE>( \
topk_weights.data_ptr<float>(), \
topk_ids.data_ptr<int32_t>(), \
gating_output.data_ptr<scalar_t>(), \
num_tokens, \
topk, \
renormalize);
#define LAUNCH_TOPK_SOFTMAX_KERNEL(NE) \
topk_softmax_kernel_impl<scalar_t, NE>( \
topk_weights.data_ptr<float>(), \
topk_ids.data_ptr<int32_t>(), \
gating_output.data_ptr<scalar_t>(), \
num_tokens, \
topk, \
renormalize);
#define LAUNCH_BIASED_GROUPED_TOPK_KERNEL(NE, NTOPK) \
biased_grouped_topk_kernel_impl<scalar_t, param_t, NE, NTOPK>( \
topk_weights.data_ptr<float>(), \
topk_ids.data_ptr<int32_t>(), \
gating_output.data_ptr<scalar_t>(), \
correction_bias.data_ptr<param_t>(), \
scaling_factor_value, \
num_tokens, \
num_expert_group, \
topk_group, \
renormalize);
} // anonymous namespace
std::tuple<at::Tensor, at::Tensor>
topk_sigmoid_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize) {
CHECK_INPUT(gating_output);
const auto st = hidden_states.scalar_type();
CHECK_EQ(gating_output.scalar_type(), st);
int64_t num_tokens = hidden_states.size(0);
int64_t num_experts = gating_output.size(1);
TORCH_CHECK(gating_output.size(0) == num_tokens, "Number of tokens mismatch");
TORCH_CHECK(topk == 1, "topk_sigmoid only supports topk=1 case");
at::Tensor topk_weights = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kFloat));
at::Tensor topk_ids = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kInt));
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "topk_sigmoid_kernel", [&] {
switch (num_experts) {
case 1:
LAUNCH_TOPK_SIGMOID_KERNEL(1);
break;
case 2:
LAUNCH_TOPK_SIGMOID_KERNEL(2);
break;
case 4:
LAUNCH_TOPK_SIGMOID_KERNEL(4);
break;
case 8:
LAUNCH_TOPK_SIGMOID_KERNEL(8);
break;
case 16:
LAUNCH_TOPK_SIGMOID_KERNEL(16);
break;
case 32:
LAUNCH_TOPK_SIGMOID_KERNEL(32);
break;
case 64:
LAUNCH_TOPK_SIGMOID_KERNEL(64);
break;
case 128:
LAUNCH_TOPK_SIGMOID_KERNEL(128);
break;
case 160:
LAUNCH_TOPK_SIGMOID_KERNEL(160);
break;
case 256:
LAUNCH_TOPK_SIGMOID_KERNEL(256);
break;
default:
TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
}
});
return std::make_tuple(topk_weights, topk_ids);
}
std::tuple<at::Tensor, at::Tensor>
topk_softmax_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize) {
CHECK_INPUT(gating_output);
const auto st = hidden_states.scalar_type();
CHECK_EQ(gating_output.scalar_type(), st);
int64_t num_tokens = hidden_states.size(0);
int64_t num_experts = gating_output.size(1);
TORCH_CHECK(gating_output.size(0) == num_tokens, "Number of tokens mismatch");
at::Tensor topk_weights = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kFloat));
at::Tensor topk_ids = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kInt));
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "topk_softmax_cpu", [&] {
switch (num_experts) {
case 1:
LAUNCH_TOPK_SOFTMAX_KERNEL(1);
break;
case 2:
LAUNCH_TOPK_SOFTMAX_KERNEL(2);
break;
case 4:
LAUNCH_TOPK_SOFTMAX_KERNEL(4);
break;
case 8:
LAUNCH_TOPK_SOFTMAX_KERNEL(8);
break;
case 16:
LAUNCH_TOPK_SOFTMAX_KERNEL(16);
break;
case 32:
LAUNCH_TOPK_SOFTMAX_KERNEL(32);
break;
case 64:
LAUNCH_TOPK_SOFTMAX_KERNEL(64);
break;
case 128:
LAUNCH_TOPK_SOFTMAX_KERNEL(128);
break;
case 160:
LAUNCH_TOPK_SOFTMAX_KERNEL(160);
break;
case 256:
LAUNCH_TOPK_SOFTMAX_KERNEL(256);
break;
case 384:
LAUNCH_TOPK_SOFTMAX_KERNEL(384);
break;
case 512:
LAUNCH_TOPK_SOFTMAX_KERNEL(512);
break;
default:
TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
}
});
return std::make_tuple(topk_weights, topk_ids);
}
// grouped topk for DeepSeek V2
std::tuple<at::Tensor, at::Tensor> grouped_topk_cpu(
at::Tensor& hidden_states,
at::Tensor& gating_output,
int64_t topk,
bool renormalize,
int64_t num_expert_group,
int64_t topk_group,
int64_t num_fused_shared_experts,
std::optional<double> routed_scaling_factor,
std::optional<at::Tensor> num_token_non_padded) {
// TODO: Will support num_fused_shared_experts, routed_scaling_factor and num_token_non_padded.
// For now, we just check them as default value.
TORCH_CHECK(
num_fused_shared_experts == 0,
"num_fused_shared_experts must be 0 default value, got: ",
num_fused_shared_experts);
TORCH_CHECK(
!routed_scaling_factor.has_value() || routed_scaling_factor.value() == 1.0f,
"routed_scaling_factor must be None or 1.0f default value, got: ",
routed_scaling_factor.value());
TORCH_CHECK(
!num_token_non_padded.has_value(),
"num_token_non_padded must be None default value, got: ",
num_token_non_padded.value());
CHECK_INPUT(gating_output);
const auto st = hidden_states.scalar_type();
CHECK_EQ(gating_output.scalar_type(), st);
int64_t num_tokens = hidden_states.size(0);
int64_t num_experts = gating_output.size(1);
TORCH_CHECK(gating_output.size(0) == num_tokens, "Number of tokens mismatch");
at::Tensor topk_weights = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kFloat));
at::Tensor topk_ids = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kInt));
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "grouped_topk_kernel", [&] {
switch (num_experts) {
case 1:
LAUNCH_GROUPED_TOPK_KERNEL(1);
break;
case 2:
LAUNCH_GROUPED_TOPK_KERNEL(2);
break;
case 4:
LAUNCH_GROUPED_TOPK_KERNEL(4);
break;
case 8:
LAUNCH_GROUPED_TOPK_KERNEL(8);
break;
case 16:
LAUNCH_GROUPED_TOPK_KERNEL(16);
break;
case 32:
LAUNCH_GROUPED_TOPK_KERNEL(32);
break;
case 64:
LAUNCH_GROUPED_TOPK_KERNEL(64);
break;
case 128:
LAUNCH_GROUPED_TOPK_KERNEL(128);
break;
case 160:
LAUNCH_GROUPED_TOPK_KERNEL(160);
break;
case 256:
LAUNCH_GROUPED_TOPK_KERNEL(256);
break;
default:
TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
}
});
return std::make_tuple(topk_weights, topk_ids);
}
// biased grouped topk DeepSeek V3/R1
std::tuple<at::Tensor, at::Tensor> biased_grouped_topk_cpu(
at::Tensor& hidden_states,
at::Tensor& gating_output,
at::Tensor& correction_bias,
int64_t topk,
bool renormalize,
int64_t num_expert_group,
int64_t topk_group,
int64_t num_fused_shared_experts,
std::optional<double> routed_scaling_factor,
std::optional<at::Tensor> num_token_non_padded) {
// TODO: Will support num_fused_shared_experts and num_token_non_padded.
// For now, we just check them as default value.
TORCH_CHECK(
num_fused_shared_experts == 0,
"num_fused_shared_experts must be 0 default value, got: ",
num_fused_shared_experts);
TORCH_CHECK(
!num_token_non_padded.has_value(),
"num_token_non_padded must be None default value, got: ",
num_token_non_padded.value());
CHECK_INPUT(gating_output);
CHECK_INPUT(correction_bias);
const auto st = gating_output.scalar_type();
int64_t num_tokens = hidden_states.size(0);
int64_t num_experts = gating_output.size(1);
TORCH_CHECK(gating_output.size(0) == num_tokens, "Number of tokens mismatch");
TORCH_CHECK(correction_bias.numel() == num_experts, "Bias shape mismatch");
at::Tensor topk_weights = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kFloat));
at::Tensor topk_ids = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kInt));
float scaling_factor_value = routed_scaling_factor.has_value() ? routed_scaling_factor.value() : 1.0f;
CPU_DISPATCH_FLOATING_TYPES_EXT(st, correction_bias.scalar_type(), "biased_grouped_topk_kernel", [&] {
TORCH_CHECK(topk == 8, "Unexpected topk: ", topk);
switch (num_experts) {
case 128:
LAUNCH_BIASED_GROUPED_TOPK_KERNEL(128, 8);
break;
case 192:
LAUNCH_BIASED_GROUPED_TOPK_KERNEL(192, 8);
break;
case 256:
LAUNCH_BIASED_GROUPED_TOPK_KERNEL(256, 8);
break;
case 384:
LAUNCH_BIASED_GROUPED_TOPK_KERNEL(384, 8);
break;
default:
TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
}
});
return std::make_tuple(topk_weights, topk_ids);
}
python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp
+892
View File
@@ -0,0 +1,892 @@
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/ATen.h>
#include <torch/all.h>
#include <torch/library.h>
#include "sgl_kernel_ops.h"
#include "shm.h"
// silu_and_mul
at::Tensor silu_and_mul_cpu(at::Tensor& input);
// gelu_and_mul
at::Tensor gelu_tanh_and_mul_cpu(const at::Tensor& input);
at::Tensor gelu_and_mul_cpu(const at::Tensor& input);
// fused_sigmoid_mul
at::Tensor fused_sigmoid_mul_cpu(at::Tensor& input, const at::Tensor& gate, bool inplace);
// l2norm
at::Tensor l2norm_cpu(at::Tensor& input, double eps);
// rmsnorm
at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
at::Tensor gemma4_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps, double scale_shift, bool with_scale);
// layernorm
at::Tensor
layernorm_cpu(const at::Tensor& input, const at::Tensor& weight, const std::optional<at::Tensor>& bias, double eps);
// qwen3_next_rmsnorm_gated
at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Tensor& gate, double eps);
// fused_add_rmsnorm
void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps);
void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps);
// fused_add_layernorm
at::Tensor fused_add_layernorm_cpu(
const at::Tensor& input,
at::Tensor& residual,
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
double eps);
// fused_qk_gemma_rmsnorm
std::tuple<at::Tensor, at::Tensor> fused_qk_gemma_rmsnorm_cpu(
const at::Tensor& q,
const at::Tensor& k,
const at::Tensor& q_weight,
const at::Tensor& k_weight,
double eps,
int64_t head_dim);
std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qk_gemma_rmsnorm_with_gate_cpu(
const at::Tensor& q_gate,
const at::Tensor& k,
const at::Tensor& q_weight,
const at::Tensor& k_weight,
double eps,
int64_t head_dim,
int64_t num_head);
// speculative decoding
void verify_tree_greedy_cpu(
at::Tensor predicts,
at::Tensor accept_index,
at::Tensor accept_token_num,
const at::Tensor& candidates,
const at::Tensor& retrive_index,
const at::Tensor& retrive_next_token,
const at::Tensor& retrive_next_sibling,
const at::Tensor& target_predict);
void build_tree_kernel_efficient_cpu(
const at::Tensor& parent_list,
const at::Tensor& selected_index,
const at::Tensor& verified_seq_len,
at::Tensor tree_mask,
at::Tensor positions,
at::Tensor retrive_index,
at::Tensor retrive_next_token,
at::Tensor retrive_next_sibling,
int64_t topk,
int64_t depth,
int64_t draft_token_num,
int64_t tree_mask_mode);
void assign_req_to_token_pool_cpu(
const at::Tensor& req_pool_indices,
at::Tensor req_to_token,
const at::Tensor& start_offset,
const at::Tensor& end_offset,
const at::Tensor& out_cache_loc,
int64_t pool_len);
at::Tensor build_draft_decode_metadata_cpu(
const at::Tensor& req_to_token,
const at::Tensor& req_pool_indices,
const at::Tensor& seq_lens,
int64_t topk,
int64_t num_steps,
int64_t pool_len);
void fill_bonus_tokens_cpu(
const at::Tensor& accept_tokens, const at::Tensor& accept_lens, at::Tensor bonus_tokens, int64_t accept_stride);
void fill_accept_out_cache_loc_cpu(
const at::Tensor& accept_index, const at::Tensor& out_cache_loc, at::Tensor accept_out_cache_loc);
void assign_draft_cache_locs_contiguous_cpu(
const at::Tensor& req_pool_indices,
const at::Tensor& req_to_token,
const at::Tensor& seq_lens,
at::Tensor out_cache_loc,
int64_t pool_len,
int64_t topk,
int64_t num_steps);
void assign_extend_cache_locs_cpu(
const at::Tensor& req_pool_indices,
const at::Tensor& req_to_token,
const at::Tensor& start_offset,
const at::Tensor& end_offset,
at::Tensor out_cache_loc,
int64_t pool_len);
void reconstruct_indices_from_tree_mask_cpu(
const at::Tensor& tree_mask,
const at::Tensor& verified_seq_len,
at::Tensor positions,
at::Tensor retrive_index,
at::Tensor retrive_next_token,
at::Tensor retrive_next_sibling,
int64_t batch_size,
int64_t draft_token_num);
void rotate_input_ids_cpu(
at::Tensor input_ids,
const at::Tensor& extend_start_loc,
const at::Tensor& extend_seq_lens,
const at::Tensor& topk_index,
const std::optional<at::Tensor>& select_index_opt);
// topk
std::tuple<at::Tensor, at::Tensor>
topk_sigmoid_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize);
std::tuple<at::Tensor, at::Tensor>
topk_softmax_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize);
std::tuple<at::Tensor, at::Tensor> grouped_topk_cpu(
at::Tensor& hidden_states,
at::Tensor& gating_output,
int64_t topk,
bool renormalize,
int64_t num_expert_group,
int64_t topk_group,
int64_t num_fused_shared_experts,
std::optional<double> routed_scaling_factor,
std::optional<at::Tensor> num_token_non_padded);
std::tuple<at::Tensor, at::Tensor> biased_grouped_topk_cpu(
at::Tensor& hidden_states,
at::Tensor& gating_output,
at::Tensor& correction_bias,
int64_t topk,
bool renormalize,
int64_t num_expert_group,
int64_t topk_group,
int64_t num_fused_shared_experts,
std::optional<double> routed_scaling_factor,
std::optional<at::Tensor> num_token_non_padded);
// attention
void decode_attention_cpu(
at::Tensor& query,
at::Tensor& k_cache,
at::Tensor& v_cache,
at::Tensor& output,
const std::optional<at::Tensor>& key,
const std::optional<at::Tensor>& value,
at::Tensor& loc,
at::Tensor& attn_logits,
at::Tensor& req_to_token,
at::Tensor& req_pool_indices,
at::Tensor& seq_lens,
double sm_scale,
double logit_cap,
bool is_cross_attn,
int64_t slidling_window_size,
std::optional<at::Tensor> encoder_lens,
std::optional<at::Tensor> sinks);
void extend_attention_cpu(
at::Tensor& q_extend,
const std::optional<at::Tensor>& k_extend,
const std::optional<at::Tensor>& v_extend,
at::Tensor& o_extend,
at::Tensor& k_buffer,
at::Tensor& v_buffer,
at::Tensor& req_to_token,
at::Tensor& req_pool_indices,
at::Tensor& seq_lens,
at::Tensor& extend_seq_lens,
at::Tensor& extend_start_loc,
int64_t max_len_extend,
double sm_scale,
double logit_cap,
bool is_cross_attn,
int64_t sliding_window_size,
std::optional<at::Tensor> encoder_lens,
std::optional<at::Tensor> sinks,
std::optional<at::Tensor> tree_mask);
// flash attention
at::Tensor flash_attn_varlen_func(
const at::Tensor& q,
const at::Tensor& k,
const at::Tensor& v,
const at::Tensor& cu_seqlens_q,
const at::Tensor& cu_seqlens_k,
int64_t max_seqlen_q,
int64_t max_seqlen_k,
bool causal);
// linear attention
std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
const at::Tensor& query,
const at::Tensor& key,
const at::Tensor& value,
const at::Tensor& g,
const at::Tensor& beta,
const at::Tensor& initial_state,
bool output_final_state,
const at::Tensor& cu_seqlens,
bool head_first,
bool use_qk_l2norm_in_kernel,
const at::Tensor& initial_state_indices,
double eps = 1e-6);
// weight prepack
at::Tensor convert_weight_packed(at::Tensor& weight);
// scale prepack for mxfp4
at::Tensor convert_scale_packed(at::Tensor& scale);
// quant
std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A);
// igemm
at::Tensor int8_scaled_mm_cpu(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales1,
at::Tensor& scales2,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni);
// fp8 gemm
at::Tensor fp8_scaled_mm_cpu(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales2,
std::vector<int64_t> block_size,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni);
// mxfp4 gemm
at::Tensor mxfp4_scaled_mm_cpu(
at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2, const std::optional<at::Tensor>& bias, bool is_vnni);
// quant + igemm
at::Tensor int8_scaled_mm_with_quant(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales2,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool is_vnni);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// int4 gemm
at::Tensor int4_scaled_mm_cpu(
at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias);
// weight prepack for int4 weights
std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32
at::Tensor qzeros, // awq: (*, K / group_size, N / 8) || gptq: (*, K / group_size, N / 8) , int32
at::Tensor scales, // awq: (*, K / group_size, N) || gptq: (*, K / group_size, N) , bfloat16
int64_t quant_method_4bit);
#endif
// gemm
at::Tensor
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni);
// gemm fusion
at::Tensor fused_linear_sigmoid_mul(
at::Tensor& mat1,
at::Tensor& mat2,
const std::optional<at::Tensor>& bias,
bool is_vnni,
const at::Tensor& post_mul_mat);
// bmm
void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale);
// fused moe
at::Tensor fused_experts_cpu(
at::Tensor& hidden_states,
at::Tensor& w1,
at::Tensor& w2,
at::Tensor& topk_weights,
at::Tensor& topk_ids,
bool inplace,
int64_t moe_comp_method,
const std::optional<at::Tensor>& w1_scale,
const std::optional<at::Tensor>& w2_scale,
const std::optional<at::Tensor>& w1_zero,
const std::optional<at::Tensor>& w2_zero,
const std::optional<std::vector<int64_t>> block_size,
const std::optional<at::Tensor>& w1_bias,
const std::optional<at::Tensor>& w2_bias,
const std::optional<double>& alpha,
const std::optional<double>& limit,
bool is_vnni);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
at::Tensor shared_expert_cpu(
at::Tensor& hidden_states,
at::Tensor& w1,
at::Tensor& w2,
const std::optional<at::Tensor>& fused_experts_out,
const std::optional<double> routed_scaling_factor,
bool inplace,
bool use_int8_w8a8,
bool use_fp8_w8a16,
const std::optional<at::Tensor>& w1_scale,
const std::optional<at::Tensor>& w2_scale,
const std::optional<std::vector<int64_t>> block_size,
bool is_vnni);
// weight absorption
std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope(
at::Tensor& hidden_states,
at::Tensor& q_a_proj_weight,
at::Tensor& q_b_proj_weight,
at::Tensor& kv_a_proj_weight,
at::Tensor& w_kc,
at::Tensor& q_a_layernorm_weight,
at::Tensor& kv_a_layernorm_weight,
at::Tensor& positions,
at::Tensor& cos_sin_cache,
double eps,
bool use_int8_w8a8,
bool use_fp8_w8a16,
std::optional<at::Tensor> q_a_proj_scale,
std::optional<at::Tensor> q_b_proj_scale,
std::optional<at::Tensor> kv_a_proj_scale,
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size);
std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope_fused_weight(
at::Tensor& hidden_states,
at::Tensor& qkv_a_proj_weight,
at::Tensor& q_b_proj_weight,
at::Tensor& w_kc,
at::Tensor& q_a_layernorm_weight,
at::Tensor& kv_a_layernorm_weight,
at::Tensor& positions,
at::Tensor& cos_sin_cache,
double eps,
bool use_int8_w8a8,
bool use_fp8_w8a16,
std::optional<at::Tensor> qkv_a_proj_scale,
std::optional<at::Tensor> q_b_proj_scale,
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size,
int64_t q_lora_rank,
int64_t kv_lora_rank,
int64_t qk_rope_head_dim);
// mamba causal conv1d
at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight);
at::Tensor causal_conv1d_fwd_cpu(
const at::Tensor& x,
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
const std::optional<at::Tensor>& conv_states,
const std::optional<at::Tensor>& query_start_loc,
const std::optional<at::Tensor>& cache_indices,
const std::optional<at::Tensor>& has_initial_state,
bool silu_activation,
int64_t pad_slot_id,
bool is_vnni);
at::Tensor causal_conv1d_update_cpu(
const at::Tensor& x,
const at::Tensor& conv_states,
const at::Tensor& weight,
const std::optional<at::Tensor>& bias,
bool silu_activation,
const std::optional<at::Tensor>& cache_seqlens,
const std::optional<at::Tensor>& conv_state_indices,
int64_t pad_slot_id,
bool is_vnni);
#endif
// conv3d fast path for patch embedding
at::Tensor conv3d_embed_weight_pack(const at::Tensor& weight);
at::Tensor conv3d_embed_cpu(const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias, bool is_vnni);
// shared memory init
void initialize(int64_t size, int64_t rank);
// shared mmeory all_reduce
void shm_allreduce(at::Tensor& data, int64_t op);
// shared memory all_gather
at::Tensor shm_allgather(at::Tensor& data, int64_t dim);
// shared memory all_gather_into_tensor
void shm_allgather_into_tensor(at::Tensor& output_tensor, at::Tensor& data);
// shared memory reduce_scatter_tensor
void shm_reduce_scatter_tensor(at::Tensor& output_tensor, at::Tensor& data, int64_t op);
// rope
std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
at::Tensor& positions,
at::Tensor& query,
at::Tensor& key,
int64_t head_size,
at::Tensor& cos_sin_cache,
bool is_neox);
std::tuple<at::Tensor, at::Tensor>
apply_rotary_pos_emb_cpu(at::Tensor& query, at::Tensor& key, at::Tensor& cos, at::Tensor& sin);
// mrope
std::tuple<at::Tensor, at::Tensor> multimodal_rotary_embedding_cpu(
at::Tensor& positions,
at::Tensor& query,
at::Tensor& key,
int64_t head_size,
at::Tensor& cos_sin_cache,
const std::optional<std::vector<int64_t>>& mrope_section,
bool mrope_interleaved,
bool is_neox);
// CPU and memory binding
std::string init_cpu_threads_env(const std::string& cpu_ids);
// fused_sigmoid_gating_delta_rule_update
at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
const at::Tensor& A_log,
const at::Tensor& dt_bias,
const at::Tensor& q,
const at::Tensor& k,
const at::Tensor& v,
const at::Tensor& a,
const at::Tensor& b,
at::Tensor& initial_state_source,
const at::Tensor& initial_state_indices,
const at::Tensor& cu_seqlens,
bool use_qk_l2norm_in_kernel,
double softplus_beta = 1.0,
double softplus_threshold = 20.0);
// fused_gdn_gating
std::tuple<at::Tensor, at::Tensor>
fused_gdn_gating_cpu(const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b, const at::Tensor& dt_bias);
// fused_qkvzba_split_reshape_cat_cpu
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_reshape_cat_cpu(
const at::Tensor& mixed_qkvz,
const at::Tensor& mixed_ba,
int64_t num_heads_qk,
int64_t num_heads_v,
int64_t head_qk,
int64_t head_v);
// fused_qkvzba_split_reshape_cat_cpu_contiguous
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_reshape_cat_contiguous_cpu(
const at::Tensor& mixed_qkvz,
const at::Tensor& mixed_ba,
int64_t num_heads_qk,
int64_t num_heads_v,
int64_t head_qk,
int64_t head_v);
// fused_input_proj_cpu
std::tuple<at::Tensor, at::Tensor>
fused_input_proj_cpu(at::Tensor& hidden_states, at::Tensor& qkvz_weight, at::Tensor& ba_weight, bool is_vnni);
// image preprocessor
std::tuple<at::Tensor, at::Tensor> image_preprocess_cpu(
at::TensorList images,
bool do_convert_rgb,
bool do_resize,
int64_t shortest_edge,
int64_t longest_edge,
const std::string& interpolation,
bool do_rescale,
double rescale_factor,
bool do_normalize,
c10::ArrayRef<double> image_mean,
c10::ArrayRef<double> image_std,
int64_t patch_size,
int64_t temporal_patch_size,
int64_t merge_size,
bool disable_grouping,
at::ScalarType out_dtype);
// kvcache
void store_cache_cpu(
const at::Tensor& k,
const at::Tensor& v,
const at::Tensor& k_cache,
const at::Tensor& v_cache,
const at::Tensor& indices,
std::optional<int64_t> row_dim);
void copy_all_layer_kv_cache_cpu(
const at::Tensor& data_ptrs, const at::Tensor& strides, const at::Tensor& tgt_loc, const at::Tensor& src_loc);
// [NOTE] When registering kernels, we should accurately describe the in-place information.
// Taking fused_add_rmsnorm_cpu as an example, add `Tensor(a!)` modifier to all tensors that
// will be modified in-place to avoid incorrect fusing and execution order on graph mode.
TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
// activation
m.def("silu_and_mul_cpu(Tensor input) -> Tensor");
m.impl("silu_and_mul_cpu", torch::kCPU, &silu_and_mul_cpu);
m.def("gelu_tanh_and_mul_cpu(Tensor input) -> Tensor");
m.impl("gelu_tanh_and_mul_cpu", torch::kCPU, &gelu_tanh_and_mul_cpu);
m.def("gelu_and_mul_cpu(Tensor input) -> Tensor");
m.impl("gelu_and_mul_cpu", torch::kCPU, &gelu_and_mul_cpu);
m.def("fused_sigmoid_mul_cpu(Tensor(a!) input, Tensor gate, bool inplace) -> Tensor(a!)");
m.impl("fused_sigmoid_mul_cpu", torch::kCPU, &fused_sigmoid_mul_cpu);
// norm
m.def("rmsnorm_cpu(Tensor input, Tensor weight, float eps) -> Tensor");
m.impl("rmsnorm_cpu", torch::kCPU, &rmsnorm_cpu);
m.def("gemma_rmsnorm_cpu(Tensor input, Tensor weight, float eps) -> Tensor");
m.impl("gemma_rmsnorm_cpu", torch::kCPU, &gemma_rmsnorm_cpu);
m.def("gemma3_rmsnorm_cpu(Tensor input, Tensor weight, float eps) -> Tensor");
m.impl("gemma3_rmsnorm_cpu", torch::kCPU, &gemma3_rmsnorm_cpu);
m.def("gemma4_rmsnorm_cpu(Tensor input, Tensor weight, float eps, float scale_shift, bool with_scale) -> Tensor");
m.impl("gemma4_rmsnorm_cpu", torch::kCPU, &gemma4_rmsnorm_cpu);
m.def("layernorm_cpu(Tensor input, Tensor weight, Tensor? bias, float eps) -> Tensor");
m.impl("layernorm_cpu", torch::kCPU, &layernorm_cpu);
m.def("l2norm_cpu(Tensor input, float eps) -> Tensor");
m.impl("l2norm_cpu", torch::kCPU, &l2norm_cpu);
m.def("fused_rmsnorm_gated_cpu(Tensor input, Tensor weight, Tensor gate, float eps) -> Tensor");
m.impl("fused_rmsnorm_gated_cpu", torch::kCPU, &fused_rmsnorm_gated_cpu);
m.def("fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()");
m.impl("fused_add_rmsnorm_cpu", torch::kCPU, &fused_add_rmsnorm_cpu);
m.def("gemma_fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()");
m.impl("gemma_fused_add_rmsnorm_cpu", torch::kCPU, &gemma_fused_add_rmsnorm_cpu);
m.def(
"fused_add_layernorm_cpu(Tensor input, Tensor residual, Tensor weight, Tensor? bias, float eps) -> "
"Tensor");
m.impl("fused_add_layernorm_cpu", torch::kCPU, &fused_add_layernorm_cpu);
m.def(
"fused_qk_gemma_rmsnorm_cpu(Tensor q, Tensor k, Tensor q_weight, Tensor k_weight, float eps, int head_dim) -> "
"(Tensor, Tensor)");
m.impl("fused_qk_gemma_rmsnorm_cpu", torch::kCPU, &fused_qk_gemma_rmsnorm_cpu);
m.def(
"fused_qk_gemma_rmsnorm_with_gate_cpu(Tensor q_gate, Tensor k, Tensor q_weight, Tensor k_weight, float eps, int "
"head_dim, int num_head) -> "
"(Tensor, Tensor, Tensor)");
m.impl("fused_qk_gemma_rmsnorm_with_gate_cpu", torch::kCPU, &fused_qk_gemma_rmsnorm_with_gate_cpu);
// speculative decoding
m.def(
"verify_tree_greedy_cpu(Tensor(a!) predicts, Tensor(a!) accept_index, "
"Tensor(a!) accept_token_num, Tensor candidates, Tensor retrive_index, "
"Tensor retrive_next_token, Tensor retrive_next_sibling, Tensor target_predict) -> ()");
m.impl("verify_tree_greedy_cpu", torch::kCPU, &verify_tree_greedy_cpu);
m.def(
"build_tree_kernel_efficient_cpu(Tensor parent_list, Tensor selected_index, "
"Tensor verified_seq_len, Tensor(a!) tree_mask, Tensor(a!) positions, "
"Tensor(a!) retrive_index, Tensor(a!) retrive_next_token, "
"Tensor(a!) retrive_next_sibling, int topk, int depth, "
"int draft_token_num, int tree_mask_mode) -> ()");
m.impl("build_tree_kernel_efficient_cpu", torch::kCPU, &build_tree_kernel_efficient_cpu);
m.def(
"assign_req_to_token_pool_cpu(Tensor req_pool_indices, Tensor(a!) req_to_token, "
"Tensor start_offset, Tensor end_offset, Tensor out_cache_loc, "
"int pool_len) -> ()");
m.impl("assign_req_to_token_pool_cpu", torch::kCPU, &assign_req_to_token_pool_cpu);
m.def(
"build_draft_decode_metadata_cpu(Tensor req_to_token, Tensor req_pool_indices, "
"Tensor seq_lens, int topk, int num_steps, int pool_len) -> Tensor");
m.impl("build_draft_decode_metadata_cpu", torch::kCPU, &build_draft_decode_metadata_cpu);
m.def(
"fill_bonus_tokens_cpu(Tensor accept_tokens, Tensor accept_lens, "
"Tensor(a!) bonus_tokens, int accept_stride) -> ()");
m.impl("fill_bonus_tokens_cpu", torch::kCPU, &fill_bonus_tokens_cpu);
m.def(
"fill_accept_out_cache_loc_cpu(Tensor accept_index, Tensor out_cache_loc, "
"Tensor(a!) accept_out_cache_loc) -> ()");
m.impl("fill_accept_out_cache_loc_cpu", torch::kCPU, &fill_accept_out_cache_loc_cpu);
m.def(
"assign_draft_cache_locs_contiguous_cpu(Tensor req_pool_indices, Tensor req_to_token, "
"Tensor seq_lens, Tensor(a!) out_cache_loc, int pool_len, int topk, int num_steps) -> ()");
m.impl("assign_draft_cache_locs_contiguous_cpu", torch::kCPU, &assign_draft_cache_locs_contiguous_cpu);
m.def(
"assign_extend_cache_locs_cpu(Tensor req_pool_indices, Tensor req_to_token, "
"Tensor start_offset, Tensor end_offset, Tensor(a!) out_cache_loc, int pool_len) -> ()");
m.impl("assign_extend_cache_locs_cpu", torch::kCPU, &assign_extend_cache_locs_cpu);
m.def(
"rotate_input_ids_cpu(Tensor(a!) input_ids, Tensor extend_start_loc, "
"Tensor extend_seq_lens, Tensor topk_index, Tensor? select_index=None) -> ()");
m.impl("rotate_input_ids_cpu", torch::kCPU, &rotate_input_ids_cpu);
m.def(
"reconstruct_indices_from_tree_mask_cpu(Tensor tree_mask, Tensor verified_seq_len, "
"Tensor(a!) positions, Tensor(a!) retrive_index, "
"Tensor(a!) retrive_next_token, Tensor(a!) retrive_next_sibling, "
"int batch_size, int draft_token_num) -> ()");
m.impl("reconstruct_indices_from_tree_mask_cpu", torch::kCPU, &reconstruct_indices_from_tree_mask_cpu);
// topk
m.def("topk_sigmoid_cpu(Tensor hidden_states, Tensor gating_output, int topk, bool renormalize) -> (Tensor, Tensor)");
m.impl("topk_sigmoid_cpu", torch::kCPU, &topk_sigmoid_cpu);
m.def("topk_softmax_cpu(Tensor hidden_states, Tensor gating_output, int topk, bool renormalize) -> (Tensor, Tensor)");
m.impl("topk_softmax_cpu", torch::kCPU, &topk_softmax_cpu);
m.def(
"grouped_topk_cpu(Tensor hidden_states, Tensor gating_output, int topk, bool renormalize, int num_expert_group, "
"int topk_group, int num_fused_shared_experts, float? routed_scaling_factor, Tensor? num_token_non_padded) -> "
"(Tensor, Tensor)");
m.impl("grouped_topk_cpu", torch::kCPU, &grouped_topk_cpu);
// biased group topk
m.def(
"biased_grouped_topk_cpu(Tensor hidden_states, Tensor gating_output, Tensor correction_bias, int topk, bool "
"renormalize, int num_expert_group, int topk_group, int num_fused_shared_experts, float? routed_scaling_factor, "
"Tensor? num_token_non_padded) -> (Tensor, Tensor)");
m.impl("biased_grouped_topk_cpu", torch::kCPU, &biased_grouped_topk_cpu);
// decode
m.def(
"decode_attention_cpu(Tensor query, Tensor k_cache, Tensor v_cahce, Tensor(a!) output, Tensor? key, Tensor? "
"value, "
"Tensor loc, Tensor attn_logits, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, float sm_scale, "
"float logit_cap, bool is_cross_attn, int sliding_window_size, Tensor? encoder_lens, Tensor? sinks) -> ()");
m.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu);
// extend
m.def(
"extend_attention_cpu(Tensor q_extend, Tensor? k_extend, Tensor? v_extend, Tensor(a!) o_extend, Tensor k_buffer, "
"Tensor v_buffer, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor extend_seq_lens, Tensor "
"extend_start_loc, int max_len_extend, float sm_scale, float logit_cap, bool is_cross_attn, int "
"sliding_window_size, Tensor? "
"encoder_lens, Tensor? sinks, Tensor? tree_mask=None) -> ()");
m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
// flash attn
m.def(
"flash_attn_varlen_func(Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, "
"int max_seqlen_q, int max_seqlen_k, bool causal) -> Tensor");
m.impl("flash_attn_varlen_func", torch::kCPU, &flash_attn_varlen_func);
// linear attn
m.def(
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, Tensor g, Tensor beta, "
"Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool head_first, "
"bool use_qk_l2norm_in_kernel, Tensor initial_state_indices, float eps=1e-6) -> (Tensor, Tensor)");
m.impl("chunk_gated_delta_rule_cpu", torch::kCPU, &chunk_gated_delta_rule_cpu);
// weight prepack
m.def("convert_weight_packed(Tensor weight) -> Tensor");
m.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
// scale prepack for mxfp4
m.def("convert_scale_packed(Tensor scale) -> Tensor");
m.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed);
// quant
m.def("per_token_quant_int8_cpu(Tensor A) -> (Tensor, Tensor)");
m.impl("per_token_quant_int8_cpu", torch::kCPU, &per_token_quant_int8_cpu);
// igemm
m.def(
"int8_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales1, Tensor scales2, Tensor? bias, ScalarType "
"out_dtype, bool is_vnni) -> Tensor");
m.impl("int8_scaled_mm_cpu", torch::kCPU, &int8_scaled_mm_cpu);
// fp8 gemm
m.def(
"fp8_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales2, int[] block_size, Tensor? bias, ScalarType "
"out_dtype, bool is_vnni) -> Tensor");
m.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
// mxfp4 gemm
m.def("mxfp4_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales2, Tensor? bias, bool is_vnni) -> Tensor");
m.impl("mxfp4_scaled_mm_cpu", torch::kCPU, &mxfp4_scaled_mm_cpu);
// quant + igemm
m.def(
"int8_scaled_mm_with_quant(Tensor mat1, Tensor mat2, Tensor scales2, Tensor? bias, ScalarType out_dtype, bool "
"is_vnni) -> Tensor");
m.impl("int8_scaled_mm_with_quant", torch::kCPU, &int8_scaled_mm_with_quant);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// int4 gemm
m.def("int4_scaled_mm_cpu(Tensor x, Tensor w, Tensor w_zeros, Tensor w_scales, Tensor? bias) -> Tensor");
m.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
// weight prepack for int4 weights
m.def(
"convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor scales, int quant_method_4bit) -> (Tensor, "
"Tensor, Tensor)");
m.impl("convert_weight_packed_scale_zp", torch::kCPU, &convert_weight_packed_scale_zp);
#endif
// gemm
m.def("weight_packed_linear(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni) -> Tensor");
m.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
// gemm fusion
m.def(
"fused_linear_sigmoid_mul(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni, Tensor post_mul_mat) -> Tensor");
m.impl("fused_linear_sigmoid_mul", torch::kCPU, &fused_linear_sigmoid_mul);
// bmm
m.def("bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, Tensor? scale) -> ()");
m.impl("bmm_cpu", torch::kCPU, &bmm_cpu);
// moe
m.def(
"fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor topk_weights, Tensor topk_ids, bool "
"inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
"Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? "
"limit, bool is_vnni) -> Tensor");
m.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// weight absorption
m.def(
"qkv_proj_with_rope(Tensor hidden_states, Tensor q_a_proj_weight, Tensor q_b_proj_weight, Tensor "
"kv_a_proj_weight, Tensor w_kc, Tensor q_a_layernorm_weight, Tensor kv_a_layernorm_weight, Tensor positions, "
"Tensor cos_sin_cache, float eps, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? q_a_proj_scale, Tensor? "
"q_b_proj_scale, Tensor? kv_a_proj_scale, Tensor? w_scale, "
"bool is_vnni, int[]? block_size) -> (Tensor, Tensor, Tensor)");
m.impl("qkv_proj_with_rope", torch::kCPU, &qkv_proj_with_rope);
m.def(
"qkv_proj_with_rope_fused_weight(Tensor hidden_states, Tensor qkv_a_proj_weight, Tensor q_b_proj_weight, "
"Tensor w_kc, Tensor q_a_layernorm_weight, Tensor kv_a_layernorm_weight, Tensor positions, "
"Tensor cos_sin_cache, float eps, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? qkv_a_proj_scale, Tensor? "
"q_b_proj_scale, Tensor? w_scale,"
"bool is_vnni, int[]? block_size, int q_lora_rank, int kv_lora_rank,"
"int qk_rope_head_dim) -> (Tensor, Tensor, Tensor)");
m.impl("qkv_proj_with_rope_fused_weight", torch::kCPU, &qkv_proj_with_rope_fused_weight);
// shared expert
m.def(
"shared_expert_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor? fused_experts_out, float? "
"routed_scaling_factor, bool inplace, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? w1_scale, Tensor? "
"w2_scale, int[]? block_size, bool is_vnni) -> Tensor");
m.impl("shared_expert_cpu", torch::kCPU, &shared_expert_cpu);
// causal conv1d
m.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
m.impl("causal_conv1d_weight_pack", torch::kCPU, &causal_conv1d_weight_pack);
m.def(
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? conv_states, Tensor? query_start_loc,"
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, int pad_slot_id, bool is_vnni) -> "
"Tensor");
m.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
m.def(
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, bool is_vnni) -> Tensor");
m.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
// conv3d fast path for patch embedding
m.def("conv3d_embed_weight_pack(Tensor weight) -> Tensor");
m.impl("conv3d_embed_weight_pack", torch::kCPU, &conv3d_embed_weight_pack);
m.def("conv3d_embed_cpu(Tensor input, Tensor weight, Tensor bias, bool is_vnni) -> Tensor");
m.impl("conv3d_embed_cpu", torch::kCPU, &conv3d_embed_cpu);
// all reduce
m.def("initialize(int size, int rank) -> ()");
m.def("shm_allreduce(Tensor(a!) data, int reduce_op) -> ()");
m.impl("shm_allreduce", torch::kCPU, &shm_allreduce);
m.def("shm_allgather(Tensor data, int dim) -> Tensor");
m.impl("shm_allgather", torch::kCPU, &shm_allgather);
m.def("shm_allgather_into_tensor(Tensor(a!) output_tensor, Tensor data) -> ()");
m.impl("shm_allgather_into_tensor", torch::kCPU, &shm_allgather_into_tensor);
m.def("shm_reduce_scatter_tensor(Tensor(a!) output_tensor, Tensor data, int reduce_op) -> ()");
m.impl("shm_reduce_scatter_tensor", torch::kCPU, &shm_reduce_scatter_tensor);
// rope
m.def(
"rotary_embedding_cpu(Tensor positions, Tensor query, Tensor key, int head_size, Tensor cos_sin_cache, "
"bool is_neox) -> (Tensor, Tensor)");
m.impl("rotary_embedding_cpu", torch::kCPU, &rotary_embedding_cpu);
m.def("apply_rotary_pos_emb_cpu(Tensor query, Tensor key, Tensor cos, Tensor sin) -> (Tensor, Tensor)");
m.impl("apply_rotary_pos_emb_cpu", torch::kCPU, &apply_rotary_pos_emb_cpu);
// multimodal rope
m.def(
"multimodal_rotary_embedding_cpu(Tensor positions, Tensor query, Tensor key, int head_size, Tensor "
"cos_sin_cache, int[]? mrope_section, bool mrope_interleaved, bool is_neox) -> (Tensor, Tensor)");
m.impl("multimodal_rotary_embedding_cpu", torch::kCPU, &multimodal_rotary_embedding_cpu);
// CPU and memory binding
m.def("init_cpu_threads_env(str cpu_ids) -> str");
// fused_sigmoid_gating_delta_rule_update
m.def(
"fused_sigmoid_gating_delta_rule_update_cpu(Tensor A_log, Tensor dt_bias, Tensor q, Tensor k, Tensor v, Tensor "
"a, Tensor b, Tensor(a!) initial_state_source, Tensor initial_state_indices, Tensor cu_seqlens, bool "
"use_qk_l2norm_in_kernel, float softplus_beta=1.0, float softplus_threshold=20.0) -> Tensor");
m.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU, &fused_sigmoid_gating_delta_rule_update_cpu);
// fused_gdn_gating
m.def("fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) -> (Tensor, Tensor)");
m.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
// fused_qkvzba_split_reshape_cat_cpu
m.def(
"fused_qkvzba_split_reshape_cat_cpu(Tensor mixed_qkvz, Tensor mixed_ba, int num_heads_qk, int num_heads_v, int "
"head_qk, int head_v) -> (Tensor, Tensor, Tensor, Tensor)");
m.impl("fused_qkvzba_split_reshape_cat_cpu", torch::kCPU, &fused_qkvzba_split_reshape_cat_cpu);
// fused_qkvzba_split_reshape_cat_contiguous_cpu
m.def(
"fused_qkvzba_split_reshape_cat_contiguous_cpu(Tensor mixed_qkvz, Tensor mixed_ba, int num_heads_qk, int "
"num_heads_v, int "
"head_qk, int head_v) -> (Tensor, Tensor, Tensor, Tensor)");
m.impl("fused_qkvzba_split_reshape_cat_contiguous_cpu", torch::kCPU, &fused_qkvzba_split_reshape_cat_contiguous_cpu);
// fused_input_proj_cpu
m.def(
"fused_input_proj_cpu(Tensor hidden_states, Tensor qkvz_weight, Tensor ba_weight, bool is_vnni) -> (Tensor, "
"Tensor)");
m.impl("fused_input_proj_cpu", torch::kCPU, &fused_input_proj_cpu);
// image preprocessor
m.def(
"image_preprocess_cpu(Tensor[] images, bool do_convert_rgb, bool do_resize, int shortest_edge, int longest_edge,"
"str interpolation, bool do_rescale, float rescale_factor, bool do_normalize, float[] image_mean, float[] "
"image_std, int patch_size, int temporal_patch_size, int merge_size, bool disable_grouping, ScalarType "
"out_dtype) -> (Tensor, Tensor)");
m.impl("image_preprocess_cpu", torch::kCPU, &image_preprocess_cpu);
// kvcache
m.def(
"store_cache_cpu(Tensor k, Tensor v, Tensor(a!) k_cache, Tensor(a!) v_cache, Tensor indices, int? row_dim) -> "
"()");
m.impl("store_cache_cpu", torch::kCPU, &store_cache_cpu);
// The copy mutates the K/V buffers addressed via `data_ptrs` (a table of
// raw base pointers), which schema-level alias annotations cannot express.
m.def("copy_all_layer_kv_cache_cpu(Tensor data_ptrs, Tensor strides, Tensor tgt_loc, Tensor src_loc) -> ()");
m.impl("copy_all_layer_kv_cache_cpu", torch::kCPU, &copy_all_layer_kv_cache_cpu);
}
TORCH_LIBRARY_IMPL(sgl_kernel, CatchAll, m) {
m.impl("init_cpu_threads_env", init_cpu_threads_env);
m.impl("initialize", &initialize);
}
REGISTER_EXTENSION(common_ops)
python/sglang/kernels/aot/csrc/cpu/vec.h
+624
View File
@@ -0,0 +1,624 @@
#pragma once
#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
#define CPU_CAPABILITY_AVX512
#endif
#include <ATen/cpu/vec/functional.h>
#include <ATen/cpu/vec/vec.h>
namespace {
using namespace at::vec;
template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0>
inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, const Vectorized<float>& b) {
return at::vec::convert_from_float<scalar_t>(a, b);
}
template <typename scalar_t>
inline void store_from_float_ext(scalar_t* out, const Vectorized<float>& a) {
float out_buffer[Vectorized<float>::size()];
a.store(out_buffer);
for (int i = 0; i < Vectorized<float>::size(); ++i) {
out[i] = static_cast<scalar_t>(out_buffer[i]);
}
}
// allow f16, bf16
template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
bVec x_vec = bVec::loadu(data);
fVec x0, x1;
std::tie(x0, x1) = at::vec::convert_to_float(x_vec);
return std::make_tuple(x0, x1);
}
// allow f32
inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const float* __restrict__ data) {
using fVec = at::vec::Vectorized<float>;
fVec x0 = fVec::loadu(data);
fVec x1 = fVec::loadu(data + fVec::size());
return std::make_tuple(x0, x1);
}
template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
inline at::vec::Vectorized<float> load_float_vec(const scalar_t* __restrict__ data) {
at::vec::Vectorized<float> out;
if constexpr (std::is_same_v<scalar_t, at::BFloat16>) {
at::vec::load_fp32_from_bf16(data, out);
} else {
at::vec::load_fp32_from_fp16(data, out);
}
return out;
}
#if defined(CPU_CAPABILITY_AVX512)
// `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics
// use native instruction for bfloat16->float32 conversion
template <>
inline Vectorized<at::BFloat16>
convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) {
return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
}
template <>
inline void store_from_float_ext<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
_mm256_storeu_si256(reinterpret_cast<__m256i*>(out), (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
}
template <>
inline void store_from_float_ext<at::Half>(at::Half* out, const Vectorized<float>& a) {
_mm256_storeu_si256(
reinterpret_cast<__m256i*>(out), _mm512_cvtps_ph(__m512(a), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
}
#define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
#define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a)
// this doesn't handle NaN.
inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
const __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
__m512i combined = _mm512_add_epi16(x, _mm512_set1_epi16(0x0780));
combined = _mm512_slli_epi16(combined, 4);
combined = _mm512_and_si512(combined, _mm512_set1_epi16(0x87f0));
combined = _mm512_add_epi16(combined, _mm512_set1_epi16(0x3c00));
const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);
}
inline __m512bh cvt_e4m3_bf16_intrinsic_without_denorm(__m256i fp8_vec) {
// The following conversion is without denorm behavior, that is to say,
// Max subnorm : S.0000.111 = 0.875 ∗ 2**(−6)
// Min subnorm : S.0000.001 = 2**(−9)
// 0.0019 ~ 0.0137 cannot be converted correctly.
__m512i x = _mm512_cvtepu8_epi16(fp8_vec);
auto mask = _mm512_cmpneq_epi16_mask(
_mm512_and_si512(x, _mm512_set1_epi16(127)),
_mm512_setzero_si512()); // mask = x & 0x7f
auto mask_nan = _mm512_cmpneq_epi16_mask(
_mm512_and_si512(x, _mm512_set1_epi16(127)),
_mm512_set1_epi16(127)); // mask_nan = x & 0x7f
auto mantissa = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4); // mantissa = (x & 7) << 4
auto exponent = _mm512_add_epi16(
_mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3),
_mm512_set1_epi16(120)); // exponent = (((x >> 3) & 15) + 120)
auto nonsign = _mm512_maskz_mov_epi16(mask, _mm512_or_si512(mantissa, _mm512_slli_epi16(exponent, 7)));
nonsign = _mm512_mask_mov_epi16(_mm512_set1_epi16(0x7fff), mask_nan, nonsign); // deal with Nan
return (__m512bh)(_mm512_or_si512(
nonsign,
_mm512_slli_epi16(
_mm512_and_si512(x, _mm512_set1_epi16(128)),
8))); // add sign (x & 128) << 8
}
inline __m512bh cvt_e4m3_bf16_intrinsic_with_denorm(__m256i fp8_vec) {
__m512i x = _mm512_cvtepu8_epi16(fp8_vec);
__m512i lg2mant = _mm512_mask_mov_epi16(
_mm512_mask_mov_epi16(
_mm512_setzero_si512(), _mm512_test_epi16_mask(x, _mm512_set1_epi16(2)), _mm512_set1_epi16(1)),
_mm512_test_epi16_mask(x, _mm512_set1_epi16(4)),
_mm512_set1_epi16(2));
return (__m512bh)(_mm512_or_si512(
_mm512_maskz_mov_epi16(
_mm512_cmpneq_epi16_mask(_mm512_and_si512(x, _mm512_set1_epi16(127)), _mm512_setzero_si512()),
_mm512_mask_blend_epi16(
_mm512_test_epi16_mask(x, _mm512_set1_epi16(120)),
_mm512_or_si512(
_mm512_and_si512(
_mm512_sllv_epi16(
_mm512_and_si512(x, _mm512_set1_epi16(3)), _mm512_sub_epi16(_mm512_set1_epi16(7), lg2mant)),
_mm512_set1_epi16(0x007f)),
_mm512_slli_epi16(_mm512_add_epi16(lg2mant, _mm512_set1_epi16(118)), 7)),
_mm512_or_si512(
_mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4),
_mm512_slli_epi16(
_mm512_add_epi16(
_mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3), _mm512_set1_epi16(120)),
7)))),
_mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(128)), 8)));
}
inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
#ifdef SGLANG_CPU_FP8_CVT_FTZ
return cvt_e4m3_bf16_intrinsic_no_nan(a);
#else
return cvt_e4m3_bf16_intrinsic_with_denorm(a);
#endif
}
// faster version of float8_e4m3fn conversion to bfloat16
//
// we mapped cuda implementation from below link and vectorized with avx512:
// https://github.com/thu-pacman/chitu/blob/1ed2078ec26581ebdca05b7306d4385f86edaa7c/csrc/cuda/marlin/marlin_gemm/dequant.h#L387
//
inline __attribute__((always_inline)) __m512bh CVT_FP8_TO_BF16_EXT(__m256i a) {
const __m512i mask0 = _mm512_set1_epi16(0x80); // sign bit
const __m512i mask1 = _mm512_set1_epi16(0x7F); // exponent and mantissa
const __m512i mask2 = _mm512_set1_epi16(0x4000);
__m512i x = _mm512_cvtepu8_epi16(a);
__m512i vsign = _mm512_and_si512(x, mask0);
vsign = _mm512_slli_epi16(vsign, 8);
__m512i vexp_and_mant = _mm512_and_si512(x, mask1);
vexp_and_mant = _mm512_slli_epi16(vexp_and_mant, 4);
// _MM_TERNLOG_A | _MM_TERNLOG_B | _MM_TERNLOG_C: 0b11111110
return (__m512bh)(_mm512_ternarylogic_epi32(vsign, mask2, vexp_and_mant, 0b11111110));
}
// bias for conversion of fp8 to bf16 1/256 in float32
#define kFP8_BIAS 0x3b800000
// remove warning: ignoring attributes on template argument ‘__m512bh’ [-Wignored-attributes]
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wignored-attributes"
#define MXFP4_VALUES \
-6.0f, -4.0f, -3.0f, -2.0f, -1.5f, -1.0f, -0.5f, -0.0f, 6.0f, 4.0f, 3.0f, 2.0f, 1.5f, 1.0f, 0.5f, 0.0f
// convert 64 mxfp4 to 2x bf16 vectors, expect input 32-way packing
inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i a, __m512i s0, __m512i s1) {
// LUT
const __m512 values = _mm512_set_ps(MXFP4_VALUES);
const __m512i lut = (__m512i)(_mm512_cvtne2ps_pbh(values, values));
const __m512i abs_mask = _mm512_set1_epi16(0x7FFF);
const __m512i zero = _mm512_setzero_si512();
// expand values to 16-bit integers
__m512i x0 = _mm512_cvtepu8_epi16(a);
__m512i x1 = _mm512_srli_epi32(x0, 4);
// LUT to convert mxfp4 values to bf16
x0 = _mm512_permutexvar_epi16(x0, lut);
x1 = _mm512_permutexvar_epi16(x1, lut);
// check for zeros
__mmask32 mask0 = _mm512_cmp_epi16_mask(_mm512_and_si512(x0, abs_mask), zero, _MM_CMPINT_EQ);
__mmask32 mask1 = _mm512_cmp_epi16_mask(_mm512_and_si512(x1, abs_mask), zero, _MM_CMPINT_EQ);
// emulate bf16 mul with scale factor
x0 = _mm512_add_epi16(x0, s0);
x1 = _mm512_add_epi16(x1, s1);
// blend with zero
x0 = _mm512_mask_blend_epi16(mask0, x0, zero);
x1 = _mm512_mask_blend_epi16(mask1, x1, zero);
return std::make_tuple(__m512bh(x0), __m512bh(x1));
}
#define CVT_MXFP4_TO_BF16(a, s0, s1) cvt_mxfp4_e2m1_bf16_intrinsic_lut(a, s0, s1)
#pragma GCC diagnostic pop
#endif
// vector to scalar reduction
#if defined(CPU_CAPABILITY_AVX512)
inline float vec_reduce_sum(const Vectorized<float>& a) {
return _mm512_reduce_add_ps(__m512(a));
}
inline float vec_reduce_max(const Vectorized<float>& a) {
return _mm512_reduce_max_ps(__m512(a));
}
#else
inline float vec_reduce_sum(const Vectorized<float>& a) {
return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return x + y; }, a);
}
inline float vec_reduce_max(const Vectorized<float>& a) {
return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return maximum(x, y); }, a);
}
#endif
// https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
template <typename scalar_t>
inline void
quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restrict__ A, int64_t K, float eps = 1e-7) {
float amax = 0.f; // absolute max
for (int64_t k = 0; k < K; ++k) {
const float val = static_cast<float>(A[k]);
amax = std::max(amax, std::abs(val));
}
amax = std::max(amax, eps);
const float scale = amax / 127;
const float inv_scale = 127 / amax;
for (int64_t k = 0; k < K; ++k) {
const float val = static_cast<float>(A[k]) * inv_scale;
Aq[k] = (uint8_t)(std::round(val)) + 128;
}
As = scale;
}
#if defined(CPU_CAPABILITY_AVX512)
template <>
inline void quantize_row_int8<at::BFloat16>(
uint8_t* __restrict__ Aq, float& As, const at::BFloat16* __restrict__ A, int64_t K, float eps) {
const __m512 signBit = _mm512_set1_ps(-0.0f);
const __m512i off = _mm512_set1_epi32(128);
// K is 32x, no remainder
float amax = 0.f;
__m512 vamax0 = _mm512_set1_ps(0.f);
__m512 vamax1 = _mm512_set1_ps(0.f);
for (int64_t k = 0; k < K; k += 32) {
__m512i va = _mm512_loadu_si512((void*)(A + k));
__m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0));
__m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1));
vamax0 = _mm512_max_ps(vamax0, _mm512_andnot_ps(signBit, va0));
vamax1 = _mm512_max_ps(vamax1, _mm512_andnot_ps(signBit, va1));
}
amax = _mm512_reduce_max_ps(_mm512_max_ps(vamax0, vamax1));
amax = std::max(amax, eps);
const float scale = amax / 127;
const float inv_scale = 127 / amax;
const __m512 vd = _mm512_set1_ps(inv_scale);
for (int64_t k = 0; k < K; k += 32) {
__m512i va = _mm512_loadu_si512((void*)(A + k));
__m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0));
__m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1));
va0 = _mm512_mul_ps(va0, vd);
va1 = _mm512_mul_ps(va1, vd);
va0 = _mm512_roundscale_ps(va0, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
va1 = _mm512_roundscale_ps(va1, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
__m128i i0 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va0), off));
__m128i i1 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va1), off));
_mm256_storeu_si256(reinterpret_cast<__m256i*>(Aq + k), _mm256_set_m128i(i1, i0));
}
As = scale;
}
#endif
// transpose utils
// taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998
#if defined(CPU_CAPABILITY_AVX512)
inline void transpose_16x16_16bit(__m256i* v) {
__m256i v1[16];
v1[0] = _mm256_unpacklo_epi16(v[0], v[1]);
v1[1] = _mm256_unpackhi_epi16(v[0], v[1]);
v1[2] = _mm256_unpacklo_epi16(v[2], v[3]);
v1[3] = _mm256_unpackhi_epi16(v[2], v[3]);
v1[4] = _mm256_unpacklo_epi16(v[4], v[5]);
v1[5] = _mm256_unpackhi_epi16(v[4], v[5]);
v1[6] = _mm256_unpacklo_epi16(v[6], v[7]);
v1[7] = _mm256_unpackhi_epi16(v[6], v[7]);
v1[8] = _mm256_unpacklo_epi16(v[8], v[9]);
v1[9] = _mm256_unpackhi_epi16(v[8], v[9]);
v1[10] = _mm256_unpacklo_epi16(v[10], v[11]);
v1[11] = _mm256_unpackhi_epi16(v[10], v[11]);
v1[12] = _mm256_unpacklo_epi16(v[12], v[13]);
v1[13] = _mm256_unpackhi_epi16(v[12], v[13]);
v1[14] = _mm256_unpacklo_epi16(v[14], v[15]);
v1[15] = _mm256_unpackhi_epi16(v[14], v[15]);
v[0] = _mm256_unpacklo_epi32(v1[0], v1[2]);
v[1] = _mm256_unpackhi_epi32(v1[0], v1[2]);
v[2] = _mm256_unpacklo_epi32(v1[1], v1[3]);
v[3] = _mm256_unpackhi_epi32(v1[1], v1[3]);
v[4] = _mm256_unpacklo_epi32(v1[4], v1[6]);
v[5] = _mm256_unpackhi_epi32(v1[4], v1[6]);
v[6] = _mm256_unpacklo_epi32(v1[5], v1[7]);
v[7] = _mm256_unpackhi_epi32(v1[5], v1[7]);
v[8] = _mm256_unpacklo_epi32(v1[8], v1[10]);
v[9] = _mm256_unpackhi_epi32(v1[8], v1[10]);
v[10] = _mm256_unpacklo_epi32(v1[9], v1[11]);
v[11] = _mm256_unpackhi_epi32(v1[9], v1[11]);
v[12] = _mm256_unpacklo_epi32(v1[12], v1[14]);
v[13] = _mm256_unpackhi_epi32(v1[12], v1[14]);
v[14] = _mm256_unpacklo_epi32(v1[13], v1[15]);
v[15] = _mm256_unpackhi_epi32(v1[13], v1[15]);
v1[0] = _mm256_unpacklo_epi64(v[0], v[4]);
v1[1] = _mm256_unpackhi_epi64(v[0], v[4]);
v1[2] = _mm256_unpacklo_epi64(v[1], v[5]);
v1[3] = _mm256_unpackhi_epi64(v[1], v[5]);
v1[4] = _mm256_unpacklo_epi64(v[2], v[6]);
v1[5] = _mm256_unpackhi_epi64(v[2], v[6]);
v1[6] = _mm256_unpacklo_epi64(v[3], v[7]);
v1[7] = _mm256_unpackhi_epi64(v[3], v[7]);
v1[8] = _mm256_unpacklo_epi64(v[8], v[12]);
v1[9] = _mm256_unpackhi_epi64(v[8], v[12]);
v1[10] = _mm256_unpacklo_epi64(v[9], v[13]);
v1[11] = _mm256_unpackhi_epi64(v[9], v[13]);
v1[12] = _mm256_unpacklo_epi64(v[10], v[14]);
v1[13] = _mm256_unpackhi_epi64(v[10], v[14]);
v1[14] = _mm256_unpacklo_epi64(v[11], v[15]);
v1[15] = _mm256_unpackhi_epi64(v[11], v[15]);
v[0] = _mm256_permute2x128_si256(v1[0], v1[8], 0x20);
v[1] = _mm256_permute2x128_si256(v1[1], v1[9], 0x20);
v[2] = _mm256_permute2x128_si256(v1[2], v1[10], 0x20);
v[3] = _mm256_permute2x128_si256(v1[3], v1[11], 0x20);
v[4] = _mm256_permute2x128_si256(v1[4], v1[12], 0x20);
v[5] = _mm256_permute2x128_si256(v1[5], v1[13], 0x20);
v[6] = _mm256_permute2x128_si256(v1[6], v1[14], 0x20);
v[7] = _mm256_permute2x128_si256(v1[7], v1[15], 0x20);
v[8] = _mm256_permute2x128_si256(v1[0], v1[8], 0x31);
v[9] = _mm256_permute2x128_si256(v1[1], v1[9], 0x31);
v[10] = _mm256_permute2x128_si256(v1[2], v1[10], 0x31);
v[11] = _mm256_permute2x128_si256(v1[3], v1[11], 0x31);
v[12] = _mm256_permute2x128_si256(v1[4], v1[12], 0x31);
v[13] = _mm256_permute2x128_si256(v1[5], v1[13], 0x31);
v[14] = _mm256_permute2x128_si256(v1[6], v1[14], 0x31);
v[15] = _mm256_permute2x128_si256(v1[7], v1[15], 0x31);
}
inline void transpose_16x16_32bit(__m512i* v) {
__m512i v1[16];
v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
v1[2] = _mm512_unpacklo_epi32(v[2], v[3]);
v1[3] = _mm512_unpackhi_epi32(v[2], v[3]);
v1[4] = _mm512_unpacklo_epi32(v[4], v[5]);
v1[5] = _mm512_unpackhi_epi32(v[4], v[5]);
v1[6] = _mm512_unpacklo_epi32(v[6], v[7]);
v1[7] = _mm512_unpackhi_epi32(v[6], v[7]);
v1[8] = _mm512_unpacklo_epi32(v[8], v[9]);
v1[9] = _mm512_unpackhi_epi32(v[8], v[9]);
v1[10] = _mm512_unpacklo_epi32(v[10], v[11]);
v1[11] = _mm512_unpackhi_epi32(v[10], v[11]);
v1[12] = _mm512_unpacklo_epi32(v[12], v[13]);
v1[13] = _mm512_unpackhi_epi32(v[12], v[13]);
v1[14] = _mm512_unpacklo_epi32(v[14], v[15]);
v1[15] = _mm512_unpackhi_epi32(v[14], v[15]);
v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]);
v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]);
v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]);
v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]);
v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]);
v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]);
v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]);
v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]);
v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]);
v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]);
v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]);
v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]);
v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]);
v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]);
v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]);
v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]);
v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88);
v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88);
v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88);
v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88);
v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd);
v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd);
v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd);
v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd);
v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88);
v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88);
v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88);
v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88);
v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd);
v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd);
v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd);
v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd);
v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
}
// remove warning : ignoring attributes on template argument ‘__m512i’ [-Wignored-attributes]
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wignored-attributes"
// transpose from [2, 32] to [32, 2]
inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1) {
// r0: {a0, a1, ..., a31}
// r1: {b0, b1, ..., b31}
//
// d0: {a0, b0, ..., a15, b15}
// d1: {a16, b16, ..., a31, b31}
//
__m512i d0 = _mm512_unpacklo_epi16(r0, r1);
__m512i d1 = _mm512_unpackhi_epi16(r0, r1);
r0 = _mm512_shuffle_i32x4(d0, d1, 0x88);
r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd);
d0 = _mm512_shuffle_i32x4(r0, r1, 0x88);
d1 = _mm512_shuffle_i32x4(r0, r1, 0xdd);
return std::make_tuple(d0, d1);
}
#pragma GCC diagnostic pop
// Note: mapped from aten exp_u20
inline __attribute__((always_inline)) __m512 _mm512_exp_u20_ps(const __m512 values) {
const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f);
const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f);
const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f);
const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f);
const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f);
const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e)
const __m512 vec_half = _mm512_set1_ps(0.5f);
const __m512 vec_one = _mm512_set1_ps(1.f);
const __m512 vec_zero = _mm512_set1_ps(0.f);
const __m512 vec_two = _mm512_set1_ps(2.f);
const __m512 vec_ln2f = _mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218));
const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50));
const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218));
const __m512i vec_127 = _mm512_set1_epi32(0x0000007f);
const int n_mantissa_bits = 23;
// exp(x) =
// = exp(n * ln(2) + r) // divide x by ln(2) and get quot and rem
// = 2^n * exp(r) // simplify the exp(n*ln(2)) expression
auto less_ln_flt_min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/);
auto vec_src = _mm512_min_ps(values, vec_ln_flt_max);
vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min);
// fx = floorf(x * log2ef + 0.5)
auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half);
auto vec_fx_i = _mm512_cvt_roundps_epi32(vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC);
vec_fx = _mm512_cvtepi32_ps(vec_fx_i);
// x = x - fx * ln2
auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src);
// compute polynomial
auto vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4);
vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3);
vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2);
vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1);
vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one);
// compute 2^(n-1)
auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one);
auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number);
auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127);
vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits);
auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i);
vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, vec_two_pow_n, vec_zero);
// y = y * 2^n
vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n);
vec_res = _mm512_mul_ps(vec_res, vec_two);
return vec_res;
}
// Note: mapped from aten fexp_u20
inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) {
const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f);
const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f);
const __m512 vec_c2 = _mm512_set1_ps(-0.22433836478672356);
const __m512 vec_c3 = _mm512_set1_ps(-0.079204240219773236);
const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e)
const __m512 vec_a = _mm512_set1_ps(std::pow(2, 23) / std::log2(2));
const __m512 vec_b = _mm512_set1_ps(std::pow(2, 23) * 127.f);
const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50));
const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218));
__m512i vec_infinity = _mm512_set1_epi32(0x7F800000);
__m512i vec_zero = _mm512_setzero_epi32();
// Fast Exponential Computation on SIMD Architectures
// A. Cristiano I. Malossi, Yves Ineichen, Costas Bekas, and Alessandro
// Curioni exp(x) = 2**(x * log2(e))
// = 2**xi * 2**xf - TIPS we are using the EEEE floating point
// representation with identification to the exponent and the
// mentissa
// 2**xf will be approximated to a polynomial of degree 3 computed with
// Horner method
// mask for the boundary condition
auto min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, _CMP_LT_OS);
auto max_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_max, _CMP_GT_OS);
// transformation with log2(e)
auto vec_src = _mm512_mul_ps(values, vec_exp_log2ef);
auto vec_fractional = _mm512_sub_ps(vec_src, _mm512_floor_ps(vec_src));
// compute polynomial using Horner Scheme, for superscalar processor
auto vec_res = _mm512_fmadd_ps(vec_fractional, vec_c3, vec_c2);
vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c1);
vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c0);
vec_src = _mm512_sub_ps(vec_src, vec_res);
// the tips is here, headache in perspective
auto tmp = _mm512_fmadd_ps(vec_a, vec_src, vec_b);
// headache bis - we loose precision with the cast but it "fits", but ok
// after f32 -> f16 later
__m512i casted_integer = _mm512_cvttps_epi32(tmp);
// boundary condition, lower than the min -> 0
casted_integer = _mm512_mask_mov_epi32(casted_integer, min_mask, vec_zero);
// boundary condition, larger than the max -> +oo
casted_integer = _mm512_mask_mov_epi32(casted_integer, max_mask, vec_infinity);
// final interpretation to float
return _mm512_castsi512_ps(casted_integer);
}
// sigmoid(x) = 1 / (1 + exp(-x)); avoid vdivps via rcp14
inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_ps(__m512 x) {
__m512 minus_x = _mm512_xor_ps(_mm512_set1_ps(-0.f), x);
__m512 denom = _mm512_add_ps(_mm512_exp_u20_ps(minus_x), _mm512_set1_ps(1.f));
return _mm512_rcp14_ps(denom);
}
// SiLU(x) = x * sigmoid(x)
inline __attribute__((always_inline)) __m512 _mm512_rcp14_silu_ps(__m512 x) {
return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(x));
}
// x * sigmoid(x * alpha) for clamped SwiGLU
inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_glu_ps(__m512 x, __m512 alpha) {
__m512 xa = _mm512_mul_ps(x, alpha);
return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(xa));
}
#endif
inline at::vec::Vectorized<float> fast_sigmoid(const at::vec::Vectorized<float>& x) {
#if defined(CPU_CAPABILITY_AVX512)
return at::vec::Vectorized<float>(_mm512_rcp14_sigmoid_ps(x));
#else
const auto one = at::vec::Vectorized<float>(1.f);
return one / (one + x.neg().exp_u20());
#endif
}
inline at::vec::Vectorized<float> fast_silu(const at::vec::Vectorized<float>& x) {
#if defined(CPU_CAPABILITY_AVX512)
return at::vec::Vectorized<float>(_mm512_rcp14_silu_ps(x));
#else
return x * fast_sigmoid(x);
#endif
}
inline at::vec::Vectorized<float>
fast_sigmoid_glu(const at::vec::Vectorized<float>& x, const at::vec::Vectorized<float>& alpha) {
#if defined(CPU_CAPABILITY_AVX512)
return at::vec::Vectorized<float>(_mm512_rcp14_sigmoid_glu_ps(x, alpha));
#else
return x * fast_sigmoid(x * alpha);
#endif
}
} // anonymous namespace
python/sglang/kernels/aot/csrc/cpu/vec_pack.h
+294
View File
@@ -0,0 +1,294 @@
// To use the transpose functions
#include <ATen/native/cpu/utils.h>
#include "vec.h"
namespace {
using namespace at::vec;
template <typename index_t>
inline index_t get_index(index_t* ind, int i) {
return (ind == nullptr) ? (index_t)i : ind[i];
}
#if defined(CPU_CAPABILITY_AVX512)
// key: from [N, 32] to [32/2, N, 2]
template <typename scalar_t, typename index_t>
inline void pack_vnni_Nx32(
scalar_t* __restrict__ dst,
const scalar_t* __restrict__ src,
const index_t* __restrict__ ind,
int N,
int ld_src,
int ld_dst) {
__m512i vinputs[16];
int n = 0;
for (; n < N; ++n) {
index_t index = get_index(ind, n);
vinputs[n] = _mm512_loadu_si512(src + index * ld_src);
}
// padding with zero to avoid uninitialized vectors
for (; n < 16; ++n) {
vinputs[n] = _mm512_set1_epi32(0);
}
// pack key
transpose_16x16_32bit(vinputs);
const __mmask16 vmask = (1 << N) - 1;
for (int k = 0; k < 16; ++k) {
_mm512_mask_storeu_epi32(dst + k * ld_dst * 2, vmask, vinputs[k]);
}
}
template <typename scalar_t, typename index_t>
inline void pack_vnni_N_remainder(
scalar_t* __restrict__ dst,
const scalar_t* __restrict__ src,
const index_t* __restrict__ ind,
int N,
int K,
int ld_src,
int ld_dst) {
__m512i vinputs[16];
int K2 = K >> 1;
const __mmask16 vmask = (1 << K2) - 1;
int n = 0;
for (; n < N; ++n) {
index_t index = get_index(ind, n);
vinputs[n] = _mm512_maskz_loadu_epi32(vmask, src + index * ld_src);
}
// padding with zero to avoid uninitialized vectors
for (; n < 16; ++n) {
vinputs[n] = _mm512_set1_epi32(0);
}
// pack key
transpose_16x16_32bit(vinputs);
const __mmask16 vmask2 = (1 << N) - 1;
for (int k = 0; k < K2; ++k) {
_mm512_mask_storeu_epi32(dst + k * ld_dst * 2, vmask2, vinputs[k]);
}
}
// value: from [K, 32] to [K/2, 32, 2]
template <typename scalar_t, typename index_t>
inline void pack_vnni_Kx32(
scalar_t* __restrict__ dst,
const scalar_t* __restrict__ src,
const index_t* __restrict__ ind,
int K,
int ld_src,
int ld_dst) {
__m512i vinputs[2];
int k = 0;
for (; k < K; ++k) {
index_t index = get_index(ind, k);
vinputs[k] = _mm512_loadu_si512(src + index * ld_src);
}
// padding with zero to avoid uninitialized vectors
for (; k < 2; ++k) {
vinputs[k] = _mm512_set1_epi32(0);
}
// pack value
__m512i d0, d1;
std::tie(d0, d1) = transpose_2x32_16bit(vinputs[0], vinputs[1]);
_mm512_storeu_si512(dst + 0 * ld_dst * 2, d0);
_mm512_storeu_si512(dst + 0 * ld_dst * 2 + 32, d1);
}
template <typename scalar_t, typename index_t>
inline void pack_vnni_K_remainder(
scalar_t* __restrict__ dst,
const scalar_t* __restrict__ src,
const index_t* __restrict__ ind,
int K,
int N,
int ld_src,
int ld_dst) {
__m512i vinputs[2];
const __mmask32 vmask = (1 << N) - 1;
int k = 0;
for (; k < K; ++k) {
index_t index = get_index(ind, k);
vinputs[k] = _mm512_maskz_loadu_epi16(vmask, src + index * ld_src);
}
// padding with zero to avoid uninitialized vectors
for (; k < 2; ++k) {
vinputs[k] = _mm512_set1_epi32(0);
}
// pack value
__m512i d0, d1;
std::tie(d0, d1) = transpose_2x32_16bit(vinputs[0], vinputs[1]);
if (N <= 16) {
// 2N * 16bits: N * 32bits
const __mmask16 vmask2 = (1 << N) - 1;
_mm512_mask_storeu_epi32(dst + 0 * ld_dst * 2, vmask2, d0);
} else {
// 2(N-16) * 16bits: (N-16) * 32bits
const __mmask16 vmask2 = (1 << (N - 16)) - 1;
_mm512_storeu_epi32(dst + 0 * ld_dst * 2, d0);
_mm512_mask_storeu_epi32(dst + 0 * ld_dst * 2 + 32, vmask2, d1);
}
}
#endif
// convert to vnni format
// from [N, K/2, 2] to [K/2, N, 2] for bfloat16 and float16
template <typename scalar_t, typename index_t, bool is_indexed>
void pack_vnni(
scalar_t* __restrict__ dst,
const scalar_t* __restrict__ src,
const index_t* __restrict__ ind,
int N,
int K,
int ld_src,
int ld_dst) {
#if defined(CPU_CAPABILITY_AVX512)
const int NB = div_up(N, 16);
const int KB = K / 32;
const int K_remainder = K - KB * 32;
for (int nb = 0; nb < NB; ++nb) {
int nb_size = std::min(N - nb * 16, 16);
for (int kb = 0; kb < KB; ++kb) {
// handle 16x512bits each block
pack_vnni_Nx32<scalar_t, index_t>(
/* dst */ dst + ((kb * 32) >> 1) * ld_dst * 2 + nb * 16 * 2,
/* src */ src + kb * 32 + (is_indexed ? 0 : nb * 16 * ld_src),
/* ind */ is_indexed ? ind + nb * 16 : nullptr,
/* N */ nb_size,
/* ld_src */ ld_src,
/* ld_dst */ ld_dst);
}
if (K_remainder > 0) {
pack_vnni_N_remainder<scalar_t, index_t>(
/* dst */ dst + ((KB * 32) >> 1) * ld_dst * 2 + nb * 16 * 2,
/* src */ src + KB * 32 + (is_indexed ? 0 : nb * 16 * ld_src),
/* ind */ is_indexed ? ind + nb * 16 : nullptr,
/* N */ nb_size,
/* K */ K_remainder,
/* ld_src */ ld_src,
/* ld_dst */ ld_dst);
}
}
#else
for (int n = 0; n < N; ++n) {
index_t index = get_index(ind, n);
for (int k = 0; k < K / 2; ++k) {
for (int d = 0; d < 2; ++d) {
dst[k * ld_dst * 2 + n * 2 + d] = src[index * ld_src + k * 2 + d];
}
}
}
#endif
}
template <typename scalar_t>
void pack_vnni(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int N, int K, int ld_src, int ld_dst) {
pack_vnni<scalar_t, int32_t, false>(dst, src, nullptr, N, K, ld_src, ld_dst);
}
template <typename scalar_t, typename index_t>
void pack_vnni(
scalar_t* __restrict__ dst,
const scalar_t* __restrict__ src,
const index_t* __restrict__ ind,
int N,
int K,
int ld_src,
int ld_dst) {
assert(ind != nullptr);
pack_vnni<scalar_t, index_t, true>(dst, src, ind, N, K, ld_src, ld_dst);
}
// convert to vnni format
// from [K/2, 2, N] to [K/2, N, 2] for bfloat16 and float16
template <typename scalar_t, typename index_t, bool is_indexed>
void pack_vnni2(
scalar_t* __restrict__ dst,
const scalar_t* __restrict__ src,
const index_t* __restrict__ ind,
int K,
int N,
int ld_src,
int ld_dst) {
#if defined(CPU_CAPABILITY_AVX512)
const int KB = div_up(K, 2);
const int NB = N / 32;
const int N_remainder = N - NB * 32;
for (int kb = 0; kb < KB; ++kb) {
int kb_size = std::min(K - kb * 2, 2);
for (int nb = 0; nb < NB; ++nb) {
// handle 2x512bits each block
pack_vnni_Kx32<scalar_t, index_t>(
/* dst */ dst + ((kb * 2) >> 1) * ld_dst * 2 + nb * 32 * 2,
/* src */ src + (is_indexed ? 0 : kb * 2 * ld_src) + nb * 32,
/* ind */ is_indexed ? ind + kb * 2 : nullptr,
/* K */ kb_size,
/* ld_src */ ld_src,
/* ld_dst */ ld_dst);
}
if (N_remainder > 0) {
pack_vnni_K_remainder(
/* dst */ dst + ((kb * 2) >> 1) * ld_dst * 2 + NB * 32 * 2,
/* src */ src + (is_indexed ? 0 : kb * 2 * ld_src) + NB * 32,
/* ind */ is_indexed ? ind + kb * 2 : nullptr,
/* K */ kb_size,
/* N */ N_remainder,
/* ld_src */ ld_src,
/* ld_dst */ ld_dst);
}
}
#else
int k = 0;
for (; k < (K >> 1) * 2; k += 2) {
index_t index0 = get_index(ind, k + 0);
index_t index1 = get_index(ind, k + 1);
for (int n = 0; n < N; ++n) {
dst[(k >> 1) * ld_dst * 2 + n * 2 + 0] = src[index0 * ld_src + n];
dst[(k >> 1) * ld_dst * 2 + n * 2 + 1] = src[index1 * ld_src + n];
}
}
if (K % 2 != 0) {
index_t index = get_index(ind, K - 1);
for (int n = 0; n < N; ++n) {
dst[(K >> 1) * ld_dst * 2 + n * 2 + 0] = src[index * ld_src + n];
dst[(K >> 1) * ld_dst * 2 + n * 2 + 1] = 0;
}
k += 2;
}
#endif
}
template <typename scalar_t>
void pack_vnni2(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int K, int N, int ld_src, int ld_dst) {
pack_vnni2<scalar_t, int32_t, false>(dst, src, nullptr, K, N, ld_src, ld_dst);
}
template <typename scalar_t, typename index_t>
void pack_vnni2(
scalar_t* __restrict__ dst,
const scalar_t* __restrict__ src,
const index_t* __restrict__ ind,
int K,
int N,
int ld_src,
int ld_dst) {
assert(ind != nullptr);
pack_vnni2<scalar_t, index_t, true>(dst, src, ind, K, N, ld_src, ld_dst);
}
} // anonymous namespace
python/sglang/kernels/aot/csrc/cpu/x86_64/shm.h
+280
View File
@@ -0,0 +1,280 @@
#pragma once
#include <immintrin.h>
// Reduce functions down below use vectorized algorithm, the number of bytes
// processed each iteration depends on vector length. 256bit vector ==> 32
// bytes, 512bit vector ==> 64 bytes If you change implementation of
// reduce_bf16_buffers, etc. , check whether this number needs to be changed
#define VECTOR_LENGTH_IN_BYTES 32
inline __m512 cvt_bf16_to_fp32(const __m256i src) __attribute__((target("avx512bw")));
inline __m512 cvt_bf16_to_fp32(const __m256i src) {
auto y = _mm512_cvtepu16_epi32(src);
return _mm512_castsi512_ps(_mm512_bslli_epi128(y, 2));
}
inline __m256i cvt_fp32_to_bf16(const __m512 src) __attribute__((target("avx512bw")));
inline __m256i cvt_fp32_to_bf16(const __m512 src) {
__m512i value = _mm512_castps_si512(src);
__m512i nan = _mm512_set1_epi32(0xffff);
auto mask_value = _mm512_cmp_ps_mask(src, src, _CMP_ORD_Q);
__m512i ones = _mm512_set1_epi32(0x1);
__m512i vec_bias = _mm512_set1_epi32(0x7fff);
// uint32_t lsb = (input >> 16) & 1;
auto t_value = _mm512_and_si512(_mm512_srli_epi32(value, 16), ones);
// uint32_t rounding_bias = 0x7fff + lsb;
t_value = _mm512_add_epi32(t_value, vec_bias);
// input += rounding_bias;
t_value = _mm512_add_epi32(t_value, value);
// input = input >> 16;
t_value = _mm512_srli_epi32(t_value, 16);
// Check NaN before converting back to bf16
t_value = _mm512_mask_blend_epi32(mask_value, nan, t_value);
return _mm512_cvtusepi32_epi16(t_value);
}
inline __m512 cvt_fp16_to_fp32(const __m256i src) __attribute__((target("avx512bw")));
inline __m512 cvt_fp16_to_fp32(const __m256i src) {
return _mm512_cvtph_ps(src);
}
inline __m256i cvt_fp32_to_fp16(const __m512 src) __attribute__((target("avx512bw")));
inline __m256i cvt_fp32_to_fp16(const __m512 src) {
return _mm512_cvtps_ph(src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
}
#define CVT_ADD_BF16(x) \
do { \
auto in##x##_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \
inout_val = _mm512_add_ps(inout_val, in##x##_val); \
} while (0)
__attribute__((target("avx512bw"))) inline void
reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) {
const int element_size = 2;
const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size;
int main_elements = num_elements - (num_elements % vector_length);
int remain_elements = num_elements % vector_length;
// process aligned part
#pragma omp parallel for
for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size;
i += VECTOR_LENGTH_IN_BYTES) {
auto inout_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i)));
switch (world_size) {
case 16:
CVT_ADD_BF16(15);
case 15:
CVT_ADD_BF16(14);
case 14:
CVT_ADD_BF16(13);
case 13:
CVT_ADD_BF16(12);
case 12:
CVT_ADD_BF16(11);
case 11:
CVT_ADD_BF16(10);
case 10:
CVT_ADD_BF16(9);
case 9:
CVT_ADD_BF16(8);
case 8:
CVT_ADD_BF16(7);
case 7:
CVT_ADD_BF16(6);
case 6:
CVT_ADD_BF16(5);
case 5:
CVT_ADD_BF16(4);
case 4:
CVT_ADD_BF16(3);
case 3:
CVT_ADD_BF16(2);
case 2:
CVT_ADD_BF16(1);
case 1:
break;
default:
for (int j = 1; j < world_size; j++) {
auto in_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i)));
inout_val = _mm512_add_ps(inout_val, in_val);
}
}
_mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_bf16(inout_val));
}
// process remaining part
int i = (start_elements + main_elements) * element_size;
while (remain_elements > 0) {
float val = 0.0f;
for (int j = 0; j < world_size; j++) {
val += *(at::BFloat16*)(buffers[j] + i);
}
*(at::BFloat16*)(to_buffer + i) = val;
remain_elements--;
i += element_size;
}
}
#define CVT_ADD_FP16(x) \
do { \
auto in##x##_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \
inout_val = _mm512_add_ps(inout_val, in##x##_val); \
} while (0)
__attribute__((target("avx512bw"))) inline void
reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) {
const int element_size = 2;
const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size;
int main_elements = num_elements - (num_elements % vector_length);
int remain_elements = num_elements % vector_length;
// process aligned part
#pragma omp parallel for
for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size;
i += VECTOR_LENGTH_IN_BYTES) {
auto inout_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i)));
switch (world_size) {
case 16:
CVT_ADD_FP16(15);
case 15:
CVT_ADD_FP16(14);
case 14:
CVT_ADD_FP16(13);
case 13:
CVT_ADD_FP16(12);
case 12:
CVT_ADD_FP16(11);
case 11:
CVT_ADD_FP16(10);
case 10:
CVT_ADD_FP16(9);
case 9:
CVT_ADD_FP16(8);
case 8:
CVT_ADD_FP16(7);
case 7:
CVT_ADD_FP16(6);
case 6:
CVT_ADD_FP16(5);
case 5:
CVT_ADD_FP16(4);
case 4:
CVT_ADD_FP16(3);
case 3:
CVT_ADD_FP16(2);
case 2:
CVT_ADD_FP16(1);
case 1:
break;
default:
for (int j = 1; j < world_size; j++) {
auto in_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i)));
inout_val = _mm512_add_ps(inout_val, in_val);
}
}
_mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_fp16(inout_val));
}
// process remaining part
int i = (start_elements + main_elements) * element_size;
while (remain_elements > 0) {
float val = 0.0f;
for (int j = 0; j < world_size; j++) {
val += *(at::Half*)(buffers[j] + i);
}
*(at::Half*)(to_buffer + i) = val;
remain_elements--;
i += element_size;
}
}
#define CVT_ADD_F32(x) \
do { \
auto in##x##_val = _mm256_loadu_ps((float*)(buffers[x] + i)); \
inout_val = _mm256_add_ps(inout_val, in##x##_val); \
} while (0)
__attribute__((target("avx512bw"))) inline void
reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers, int world_size) {
const int element_size = 4;
const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size;
int main_elements = num_elements - (num_elements % vector_length);
int remain_elements = num_elements % vector_length;
// process aligned part
#pragma omp parallel for
for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size;
i += VECTOR_LENGTH_IN_BYTES) {
auto inout_val = _mm256_loadu_ps((float*)(buffers[0] + i));
switch (world_size) {
case 16:
CVT_ADD_F32(15);
case 15:
CVT_ADD_F32(14);
case 14:
CVT_ADD_F32(13);
case 13:
CVT_ADD_F32(12);
case 12:
CVT_ADD_F32(11);
case 11:
CVT_ADD_F32(10);
case 10:
CVT_ADD_F32(9);
case 9:
CVT_ADD_F32(8);
case 8:
CVT_ADD_F32(7);
case 7:
CVT_ADD_F32(6);
case 6:
CVT_ADD_F32(5);
case 5:
CVT_ADD_F32(4);
case 4:
CVT_ADD_F32(3);
case 3:
CVT_ADD_F32(2);
case 2:
CVT_ADD_F32(1);
case 1:
break;
default:
for (int j = 1; j < world_size; j++) {
auto in_val = _mm256_loadu_ps((float*)(buffers[j] + i));
inout_val = _mm256_add_ps(inout_val, in_val);
}
}
_mm256_storeu_ps((float*)(to_buffer + i), inout_val);
}
// process remaining part
int i = (start_elements + main_elements) * element_size;
while (remain_elements > 0) {
float val = 0.0f;
for (int j = 0; j < world_size; j++) {
val += *(float*)(buffers[j] + i);
}
*(float*)(to_buffer + i) = val;
remain_elements--;
i += element_size;
}
}
__attribute__((target("avx512bw"))) inline void parallel_memcpy(void* to, void* from, size_t n_bytes) {
auto aligned_bytes = n_bytes - (n_bytes % VECTOR_LENGTH_IN_BYTES);
// process aligned part
#pragma omp parallel for
for (size_t i = 0; i < aligned_bytes; i += VECTOR_LENGTH_IN_BYTES) {
auto val = _mm256_loadu_si256((__m256i*)((char*)from + i));
_mm256_storeu_si256((__m256i*)((char*)to + i), val);
}
// process remaining part
for (size_t i = aligned_bytes; i < n_bytes; i++) {
*((char*)to + i) = *((char*)from + i);
}
}
#undef VECTOR_LENGTH_IN_BYTES
python/sglang/kernels/aot/csrc/cutlass_extensions/common.hpp
+21
View File
@@ -0,0 +1,21 @@
#pragma once
#include "cuda_runtime.h"
#include "cutlass/cutlass.h"
/**
* A wrapper for a kernel that is used to guard against compilation on
* architectures that will never use the kernel. The purpose of this is to
* reduce the size of the compiled binary.
* __CUDA_ARCH__ is not defined in host code, so this lets us smuggle the ifdef
* into code that will be executed on the device where it is defined.
*/
template <typename Kernel>
struct enable_sm90_or_later : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 900
Kernel::operator()(std::forward<Args>(args)...);
#endif
}
};
python/sglang/kernels/aot/csrc/cutlass_extensions/detail/collective/mixed_input_utils.hpp
+482
View File
@@ -0,0 +1,482 @@
/*
* Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once
#include "cute/arch/copy_sm90.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cute/util/type_traits.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/numeric_conversion.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective::detail {
template <class Collective>
struct MixedGroupedGemmInputUtils {
private:
using KernelSchedule = typename Collective::KernelSchedule;
using ConversionMode = typename Collective::ConversionMode;
using SmemLayoutA = typename Collective::SmemLayoutA;
using SmemLayoutB = typename Collective::SmemLayoutB;
using SmemLayoutScale = typename Collective::SmemLayoutScale;
using SwappedElementA = typename Collective::SwappedElementA;
using SwappedElementB = typename Collective::SwappedElementB;
using RealSwappedElementA = typename Collective::RealSwappedElementA;
using RealSwappedElementB = typename Collective::RealSwappedElementB;
using ElementScale = typename Collective::ElementScale;
using ElementZero = typename Collective::ElementZero;
using SmemCopyAtomScale = typename Collective::SmemCopyAtomScale;
static constexpr auto KernelConversionMode = Collective::KernelConversionMode;
static constexpr auto ModeHasScales = Collective::ModeHasScales;
static constexpr auto UseScaleLookupTable = Collective::UseScaleLookupTable;
public:
static constexpr auto elements_per_smem_scale() {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
return 0;
} else if constexpr (ModeHasScales) {
return cute::cosize_v<SmemLayoutScale>;
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in scale smem allocation.");
}
}
static constexpr auto elements_per_smem_zero() {
if constexpr (
KernelConversionMode == ConversionMode::DirectConvert ||
KernelConversionMode == ConversionMode::ConvertAndScale) {
return 0;
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
return cute::cosize_v<SmemLayoutScale>;
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in scale smem allocation.");
}
}
// These methods use some the public members of the class. For that reason, we define them after the public section.
static constexpr uint32_t compute_tma_transaction_bytes_mk() {
return cutlass::bits_to_bytes(
size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(cute::sizeof_bits_v<SwappedElementA>));
}
static constexpr uint32_t compute_tma_transaction_bytes_nk() {
return cutlass::bits_to_bytes(
size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(cute::sizeof_bits_v<SwappedElementB>));
}
static constexpr uint32_t compute_tma_transaction_bytes_extra() {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
return 0;
} else if constexpr (ModeHasScales) {
constexpr uint32_t scale_tx_bytes = cutlass::bits_to_bytes(
size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) *
static_cast<uint32_t>(cute::sizeof_bits_v<ElementScale>));
static_assert(scale_tx_bytes % 128 == 0, "Each scale stage must be 128B aligned."); // required by TMA
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
return scale_tx_bytes;
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
// Scale and zero share smem layout
constexpr uint32_t zero_tx_bytes = cutlass::bits_to_bytes(
size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) *
static_cast<uint32_t>(cute::sizeof_bits_v<ElementZero>));
static_assert(zero_tx_bytes % 128 == 0, "Each zero stage must be 128B aligned."); // required by TMA
return scale_tx_bytes + zero_tx_bytes;
} else {
static_assert(
cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in tma transaction bytes computation.");
}
} else {
static_assert(
cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in tma transaction bytes computation.");
}
}
/// Utilities to copy A and extra inputs from smem to RF
template <class SmemTiledCopyA, class TensorASmemView, class TensorACopyView, class... Ts, class... Us>
CUTLASS_DEVICE static void copy_tensors_MK(
SmemTiledCopyA const& smem_tiled_copy_A,
TensorASmemView const& tCsA,
TensorACopyView& tCrA_copy_view,
cute::tuple<Ts...> const& partitioned_mma_extra_info,
cute::tuple<Us...> const& tiled_copy_and_views,
int k_block,
int read_stage) {
copy(smem_tiled_copy_A, tCsA(_, _, k_block, read_stage), tCrA_copy_view(_, _, k_block));
if (k_block == 0) {
// We are starting a new k-tile so copy the scale
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
// nothing to do
} else if constexpr (ModeHasScales) {
auto smem_tiled_copy_S = cute::get<0>(tiled_copy_and_views);
auto tCrS_copy_view = cute::get<1>(tiled_copy_and_views);
auto tCsS = cute::get<0>(partitioned_mma_extra_info);
copy(smem_tiled_copy_S, tCsS(_, _, k_block, read_stage), tCrS_copy_view(_, _, k_block));
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
// Nothing extra to do
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
auto tCsZ = cute::get<2>(partitioned_mma_extra_info);
auto tCrZ_copy_view = cute::get<2>(tiled_copy_and_views);
copy(smem_tiled_copy_S, tCsZ(_, _, k_block, read_stage), tCrZ_copy_view(_, _, k_block));
} else {
static_assert(
cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
}
}
// The core converter uses a lookup table to converts i4 -> 8 bit value.
template <
class EngineIn,
class LayoutIn,
class EngineOut,
class LayoutOut,
class EngineScale,
class LayoutScale>
CUTLASS_DEVICE static void lookup_table_convert( // Accept mutable temporaries
Tensor<EngineIn, LayoutIn> const& src,
Tensor<EngineOut, LayoutOut>&& dst,
Tensor<EngineScale, LayoutScale> const& scales_neg,
Tensor<EngineScale, LayoutScale> const& scales_pos) {
lookup_table_convert(src, dst, scales_neg, scales_pos);
}
template <class EngineIn, class LayoutIn, class EngineOut, class LayoutOut, class EngineScale, class LayoutScale>
CUTLASS_DEVICE static void lookup_table_convert(
Tensor<EngineIn, LayoutIn> const& src,
Tensor<EngineOut, LayoutOut>& dst,
Tensor<EngineScale, LayoutScale> const& scales_neg,
Tensor<EngineScale, LayoutScale> const& scales_pos) {
constexpr int N = cute::cosize(LayoutIn{});
static_assert(N == 4 || N == 8);
static_assert(cosize(LayoutScale{}) <= N / 4, "at least 4 consecutive weights must share the same scale.");
using SrcArray = cutlass::Array<cutlass::int4b_t, 8>;
using DstArray = cutlass::Array<RealSwappedElementB, 8>;
using RegArray = cutlass::AlignedArray<uint32_t, N / 4, sizeof(DstArray)>;
// View the input as reg
auto&& src_reg = cute::recast<uint32_t>(src)(0);
auto&& r = cute::recast<RegArray>(dst)(0);
// Determines if to get from the signed or unsigned candidates
static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
uint32_t sign; // ((reg & 0x88888888) | 0x64206420) >> 1
asm volatile(
"{\n"
" lop3.b32 %0, %1, %2, %3, %4;\n"
"}\n"
: "=r"(sign)
: "r"(src_reg), "n"(0x88888888), "n"(0x64206420), "n"(immLut));
sign = sign >> 1;
// Ignore sign bit when indexing into LUT
uint32_t lut_idx = src_reg & 0x77777777;
Tensor scales_neg_ = cute::filter(scales_neg);
Tensor scales_pos_ = cute::filter(scales_pos);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N / 4; ++i, lut_idx >>= 16, sign >>= 16) {
auto&& scale_neg_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_neg_(i));
auto&& scale_pos_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_pos_(i));
asm volatile(
"{\n"
" .reg .b32 pos, neg ;\n"
" prmt .b32 neg, %3, %4, %1 ;\n"
" prmt .b32 pos, %5, %6, %1 ;\n"
" prmt .b32 %0, pos, neg, %2 ;\n"
"}\n"
: "=r"(r[i])
: "r"(lut_idx), "r"(sign), "r"(scale_neg_[0]), "r"(scale_neg_[1]), "r"(scale_pos_[0]), "r"(scale_pos_[1]));
}
}
/// Utilities to dequantize A.
template <class Layout>
CUTLASS_DEVICE static void static_check_scale(Layout const& tensor) {
static_assert(
shape<0>(Layout{}) >= 4 && stride<0>(Layout{}) == 0,
"At least 4 adjacent weights in a thread must share the same scale.");
}
template <class Engine, class Layout>
CUTLASS_DEVICE static void static_check_scale(Tensor<Engine, Layout> const& tensor) {
static_check_scale(flatten(Layout{}));
}
template <class EngineIn, class EngineOut, class LayoutIn, class LayoutOut, class... Ts>
CUTLASS_DEVICE static void dequantize_A_kblock(
Tensor<EngineIn, LayoutIn> const& tCrA_load,
Tensor<EngineOut, LayoutOut>& tCrA_mma,
cute::tuple<Ts...>& partitioned_extra_info,
int const k_block) {
static_assert(is_rmem<EngineIn>::value, "Input tensor for A conversion must come from registers");
static_assert(is_rmem<EngineOut>::value, "Output tensor for A conversion must come from registers");
static_assert(cosize_v<LayoutIn> == cosize_v<LayoutOut>);
static_assert(size_v<LayoutIn> == cosize_v<LayoutIn>);
static_assert(size_v<LayoutOut> == cosize_v<LayoutOut>);
using SrcType = typename EngineIn::value_type;
using DstType = typename EngineOut::value_type;
Tensor src = tCrA_load(_, _, k_block);
Tensor dst = tCrA_mma(_, _, k_block);
CUTE_STATIC_ASSERT_V(
size(src(_, 0)) == cosize(src(_, 0).layout()), "The first mode of tensor src must be contiguous in memory");
// try to make the size of the first mode equal to 32bit
int constexpr NumValPerSrcReg = cute::min(decltype(size(src(_, 0)))::value, ceil_div(32, sizeof_bits_v<SrcType>));
Tensor src_vm = cute::group_modes<1, -1>(cute::zipped_divide(src, Int<NumValPerSrcReg>{}));
Tensor dst_vm = cute::group_modes<1, -1>(cute::zipped_divide(dst, Int<NumValPerSrcReg>{}));
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size<1>(dst_vm); ++i) {
LayoutAwareConvert(src_vm(_, i), dst_vm(_, i));
}
} else if constexpr (UseScaleLookupTable) {
constexpr int num_elements = decltype(size(src))::value;
static_assert(
is_same_v<RealSwappedElementA, cutlass::int4b_t>,
"Lookup table only supports int4 being the quant type now.");
static_assert(sizeof_bits_v<ElementScale> == 64, "Lookup table only supports 8 8bit scale values now.");
static_assert(
num_elements % 4 == 0 && num_elements >= 4, "Lookup table requires a vector size of 4x when converting.");
Tensor tCrS_neg = cute::get<1>(partitioned_extra_info);
auto&& tCrS_pos = cute::get<2>(partitioned_extra_info); // modification to its value is needed
Tensor scales_neg = tCrS_neg(_, _, k_block);
Tensor scales_pos = tCrS_pos(_, _, k_block);
CUTE_STATIC_ASSERT_V(cute::size(src) == cute::size(scales_neg));
static_check_scale(scales_neg);
static_check_scale(scales_pos);
Tensor scales_neg_vm = cute::group_modes<1, -1>(cute::zipped_divide(scales_neg, Int<NumValPerSrcReg>{}));
Tensor scales_pos_vm = cute::group_modes<1, -1>(cute::zipped_divide(scales_pos, Int<NumValPerSrcReg>{}));
if (k_block == 0) {
Tensor scales_neg_vm_ = filter(scales_neg_vm);
Tensor scales_pos_vm_ = filter(scales_pos_vm);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(scales_neg_vm_.layout()); ++i) {
auto&& scale_neg_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_neg_vm_(i));
auto&& scale_pos_ = reinterpret_cast<cutlass::Array<uint32_t, 2>&>(scales_pos_vm_(i));
constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
asm volatile(
"{\n"
" lop3 .b32 %0, %2, %4, %5, %6;\n"
" xor .b32 %1, %3, %5; \n"
"}\n"
: "=r"(scale_pos_[0]), "=r"(scale_pos_[1])
: "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0xFFFFFF00), "n"(0x80808080), "n"(immLut));
}
}
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size<1>(dst_vm); ++i) {
lookup_table_convert(src_vm(_, i), dst_vm(_, i), scales_neg_vm(_, i), scales_pos_vm(_, i));
}
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
Tensor scales = cute::get<1>(partitioned_extra_info)(_, _, k_block);
CUTE_STATIC_ASSERT_V(size(src) == size(scales));
Tensor scales_vm = cute::group_modes<1, -1>(cute::zipped_divide(scales, Int<NumValPerSrcReg>{}));
if constexpr (is_same_v<DstType, ElementScale>) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size<1>(dst_vm); ++i) {
LayoutAwareConvert(src_vm(_, i), dst_vm(_, i));
CUTLASS_PRAGMA_UNROLL
for (int j = 0; j < size<0>(dst_vm); ++j) {
dst_vm(j, i) *= scales_vm(j, i);
}
}
} else {
auto stage = make_tensor_like<ElementScale>(src_vm(_, 0));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size<1>(dst_vm); ++i) {
LayoutAwareConvert(src_vm(_, i), stage);
CUTLASS_PRAGMA_UNROLL
for (int j = 0; j < size<0>(dst_vm); ++j) {
stage(j) *= scales_vm(j, i);
}
LayoutAwareConvert(stage, dst_vm(_, i));
}
}
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
static_assert(is_same_v<ElementScale, ElementZero>, "ElementScale and ElementZero must be the same.");
Tensor scales = cute::get<1>(partitioned_extra_info)(_, _, k_block);
Tensor zeros = cute::get<3>(partitioned_extra_info)(_, _, k_block);
CUTE_STATIC_ASSERT_V(size(src) == size(scales));
CUTE_STATIC_ASSERT_V(size(src) == size(zeros));
Tensor scales_vm = cute::group_modes<1, -1>(cute::zipped_divide(scales, Int<NumValPerSrcReg>{}));
Tensor zeros_vm = cute::group_modes<1, -1>(cute::zipped_divide(zeros, Int<NumValPerSrcReg>{}));
if constexpr (is_same_v<DstType, ElementScale>) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size<1>(dst_vm); ++i) {
LayoutAwareConvert(src_vm(_, i), dst_vm(_, i));
CUTLASS_PRAGMA_UNROLL
for (int j = 0; j < size<0>(dst_vm); ++j) {
dst_vm(j, i) = dst_vm(j, i) * scales_vm(j, i) + zeros_vm(j, i);
}
}
} else {
auto stage = make_tensor_like<ElementScale>(src_vm(_, 0));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size<1>(dst_vm); ++i) {
LayoutAwareConvert(src_vm(_, i), stage);
CUTLASS_PRAGMA_UNROLL
for (int j = 0; j < size<0>(dst_vm); ++j) {
stage(j) = stage(j) * scales_vm(j, i) + zeros_vm(j, i);
}
LayoutAwareConvert(stage, dst_vm(_, i));
}
}
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "No A data is loaded.");
}
}
template <class EngineIn, class EngineOut, class LayoutIn, class LayoutOut, class... Ts>
CUTLASS_DEVICE static void convert_A_kblock(
Tensor<EngineIn, LayoutIn> const& tCrA_load, Tensor<EngineOut, LayoutOut>& tCrA_mma, int const k_block) {
static_assert(is_rmem<EngineIn>::value, "Input tensor for A conversion must come from registers");
static_assert(is_rmem<EngineOut>::value, "Output tensor for A conversion must come from registers");
static_assert(cosize_v<LayoutIn> == cosize_v<LayoutOut>);
static_assert(size_v<LayoutIn> == cosize_v<LayoutIn>);
static_assert(size_v<LayoutOut> == cosize_v<LayoutOut>);
using SrcType = typename EngineIn::value_type;
Tensor src = tCrA_load(_, _, k_block);
Tensor dst = tCrA_mma(_, _, k_block);
CUTE_STATIC_ASSERT_V(
size(src(_, 0)) == cosize(src(_, 0).layout()), "The first mode of tensor src must be contiguous in memory");
// try to make the size of the first mode equal to 32bit
int constexpr NumValPerSrcReg = cute::min(decltype(size(src(_, 0)))::value, ceil_div(32, sizeof_bits_v<SrcType>));
Tensor src_vm = cute::group_modes<1, -1>(cute::zipped_divide(src, Int<NumValPerSrcReg>{}));
Tensor dst_vm = cute::group_modes<1, -1>(cute::zipped_divide(dst, Int<NumValPerSrcReg>{}));
// KernelConversionMode == ConversionMode::DirectConvert
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size<1>(dst_vm); ++i) {
LayoutAwareConvert(src_vm(_, i), dst_vm(_, i));
}
}
/// Utilities for any additional inputs inside of the TMA load
template <class Params, class TensorStorage, class... Ts>
CUTLASS_DEVICE static auto partition_extra_tma_inputs(
Params const& mainloop_params,
cute::tuple<Ts...> const& load_inputs,
TensorStorage& shared_tensors,
uint2 const& cluster_local_block_id,
int const m_coord,
int const l_coord) {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
return cute::make_tuple();
} else if constexpr (ModeHasScales) {
Tensor sS =
make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE)
Tensor gS_mkl = get<2>(load_inputs);
auto block_tma_s = mainloop_params.tma_load_scale.get_slice(cluster_local_block_id.y);
Tensor gS = gS_mkl(_, _, m_coord, _, l_coord); // (BLK_M,BLK_K,k)
Tensor tSgS = block_tma_s.partition_S(gS); // (TMA,TMA_M,TMA_K,k)
Tensor tSsS = block_tma_s.partition_D(sS); // (TMA,TMA_M,TMA_K,PIPE)
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
return cute::make_tuple(tSgS, tSsS);
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
Tensor sZ =
make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE)
Tensor gZ_mkl = get<3>(load_inputs);
auto block_tma_z = mainloop_params.tma_load_zero.get_slice(cluster_local_block_id.y);
Tensor gZ = gZ_mkl(_, _, m_coord, _, l_coord); // (BLK_M,BLK_K,k)
Tensor tZgZ = block_tma_z.partition_S(gZ); // (TMA,TMA_M,TMA_K,k)
Tensor tZsZ = block_tma_z.partition_D(sZ); // (TMA,TMA_M,TMA_K,PIPE)
return cute::make_tuple(tSgS, tSsS, tZgZ, tZsZ);
} else {
static_assert(
cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for input partitioning.");
}
} else {
static_assert(
cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for input partitioning.");
}
}
/// Utilities for partitioning extra inputs for loading from smem in the mainloop.
template <class ThreadMma, class TensorStorage>
CUTLASS_DEVICE static auto
partition_extra_mma_info(ThreadMma const& mma_thread_slice, TensorStorage& shared_tensors) {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
// nothing to do
return cute::make_tuple();
} else if constexpr (UseScaleLookupTable) {
Tensor sS =
make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{}); // (BLK_M,BLK_SCALE_K,PIPE)
Tensor tCsS = mma_thread_slice.partition_A(sS);
Tensor tCrS = make_tensor<ElementScale>(mma_thread_slice.partition_fragment_A(sS(_, _, Int<0>{})).layout());
return cute::make_tuple(tCsS, tCrS);
} else if constexpr (ModeHasScales) {
Tensor sS =
make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{}); // (BLK_M,BLK_SCALE_K,PIPE)
Tensor tCsS = mma_thread_slice.partition_A(sS);
Tensor tCrS = make_tensor<ElementScale>(mma_thread_slice.partition_fragment_A(sS(_, _, Int<0>{})).layout());
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
return cute::make_tuple(tCsS, tCrS);
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
Tensor sZ = make_tensor(
make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{}); // (BLK_M,BLK_SCALE_K,PIPE)
Tensor tCsZ = mma_thread_slice.partition_A(sZ);
Tensor tCrZ = make_tensor<ElementZero>(mma_thread_slice.partition_fragment_A(sZ(_, _, Int<0>{})).layout());
return cute::make_tuple(tCsS, tCrS, tCsZ, tCrZ);
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
}
/// Returns the tiled copy and copy views for the extra inputs.
template <class TiledMma, class... Ts>
CUTLASS_DEVICE static auto retile_extra_mma_info(
TiledMma const& tiled_mma, cute::tuple<Ts...>& partitioned_extra_info, int const warp_group_thread_idx) {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
// nothing to do
return cute::make_tuple();
} else if constexpr (ModeHasScales) {
auto smem_tiled_copy_S = make_tiled_copy_A(SmemCopyAtomScale{}, tiled_mma);
auto smem_thr_copy_S = smem_tiled_copy_S.get_thread_slice(warp_group_thread_idx);
Tensor tCrS_copy_view = smem_thr_copy_S.retile_D(cute::get<1>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K)
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view);
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
Tensor tCrZ_copy_view = smem_thr_copy_S.retile_D(cute::get<3>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K)
return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view, tCrZ_copy_view);
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
}
};
} // namespace cutlass::gemm::collective::detail
python/sglang/kernels/aot/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
+450
View File
@@ -0,0 +1,450 @@
// Adapted from
// https://github.com/vllm-project/vllm/blob/16bff144be6739c9f773968ace0b9cd239f67f19/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp
/***************************************************************************************************
* Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
*reserved. SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice,
*this list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
*ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
*LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
*CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
*SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
*INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
*CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
*ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
*POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
//
// This file is a modified excerpt of
// include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
// from https://github.com/NVIDIA/cutlass v3.5.0
// It has been modified to support either row/column or scalar broadcasting
// where the tensor being loaded from is always passed in via a device pointer.
// This lets one compiled kernel handle all cases of per-tensor or
// per-channel/per-token quantization.
//
// This interface also allows the scales to be passed in as tensors that
// consistently reside on the device, which avoids an issue with a previous
// implementation where scalars needed to be on the CPU since they
// were passed in via float values. This created a potential performance hazard
// if scales were initially on the device, and caused torch.compile graphs
// breaks when moving scales to the CPU.
//
#pragma once
// Turn off clang-format for the entire file to keep it close to upstream
// clang-format off
#include "cutlass/cutlass.h"
#include "cutlass/arch/barrier.h"
#include "cute/tensor.hpp"
#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp"
namespace cutlass::epilogue::fusion {
using namespace cute;
using namespace detail;
// Row vector broadcast
template<
int Stages,
class CtaTileShapeMNK,
class Element,
class StrideMNL = Stride<_0,_1,_0>,
int Alignment = 128 / sizeof_bits_v<Element>
>
struct Sm90RowOrScalarBroadcast {
static_assert(Stages == 0, "Row broadcast doesn't support smem usage");
static_assert(is_static_v<decltype(take<0,2>(StrideMNL{}))>); // batch stride can be dynamic or static
static_assert(take<0,2>(StrideMNL{}) == Stride<_0,_1>{});
struct SharedStorage {
array_aligned<Element, size<1>(CtaTileShapeMNK{})> smem;
};
// This struct has been modified to have a bool indicating that ptr_row is a
// scalar that must be broadcast, instead of containing a scalar that is
// valid if ptr_row is null.
struct Arguments {
Element const* ptr_row = nullptr;
bool row_broadcast = true;
StrideMNL dRow = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
template <class ProblemShape>
static bool
can_implement(ProblemShape const& problem_shape, Arguments const& args) {
return true;
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
return 0;
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
CudaHostAdapter* cuda_adapter = nullptr) {
return cutlass::Status::kSuccess;
}
CUTLASS_HOST_DEVICE
Sm90RowOrScalarBroadcast() { }
CUTLASS_HOST_DEVICE
Sm90RowOrScalarBroadcast(Params const& params, SharedStorage const& shared_storage)
: params(params)
, smem(const_cast<Element*>(shared_storage.smem.data())) { }
Params params;
Element *smem = nullptr;
CUTLASS_DEVICE bool
is_producer_load_needed() const {
return false;
}
CUTLASS_DEVICE bool
is_C_load_needed() const {
return false;
}
CUTLASS_DEVICE bool
is_zero() const {
return (!params.row_broadcast && *(params.ptr_row) == Element(0));
}
template <class... Args>
CUTLASS_DEVICE auto
get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
return EmptyProducerLoadCallbacks{};
}
template <class GS_GTensor, class GS_STensor, class GS_CTensor, class Tiled_G2S, class SR_STensor, class SR_RTensor, class CTensor, class ThrResidue, class ThrNum>
struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
CUTLASS_DEVICE
ConsumerStoreCallbacks(
GS_GTensor tGS_gRow_, GS_STensor tGS_sRow_,
GS_CTensor tGS_cRow_, Tiled_G2S tiled_g2s_,
SR_STensor tSR_sRow_, SR_RTensor tSR_rRow_,
CTensor tCcRow_, ThrResidue residue_tCcRow_, ThrNum thr_num_, Params const& params_)
: tGS_gRow(tGS_gRow_)
, tGS_sRow(tGS_sRow_)
, tGS_cRow(tGS_cRow_)
, tiled_G2S(tiled_g2s_)
, tSR_sRow(tSR_sRow_)
, tSR_rRow(tSR_rRow_)
, tCcRow(tCcRow_)
, residue_tCcRow(residue_tCcRow_)
, params(params_) {}
GS_GTensor tGS_gRow; // (CPY,CPY_M,CPY_N)
GS_STensor tGS_sRow; // (CPY,CPY_M,CPY_N)
GS_CTensor tGS_cRow; // (CPY,CPY_M,CPY_N)
Tiled_G2S tiled_G2S;
SR_STensor tSR_sRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
SR_RTensor tSR_rRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
CTensor tCcRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
ThrResidue residue_tCcRow; // (m, n)
ThrNum thr_num;
Params const& params;
CUTLASS_DEVICE void
begin() {
if (!params.row_broadcast) {
fill(tSR_rRow, *(params.ptr_row));
return;
}
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
for (int i = 0; i < size(tGS_gRow_flt); ++i) {
if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
continue; // OOB of SMEM,
}
if (elem_less(tGS_cRow_flt(i), make_coord(get<0>(residue_tCcRow), get<1>(residue_tCcRow)))) {
tGS_sRow_flt(i) = tGS_gRow_flt(i);
}
else {
tGS_sRow_flt(i) = Element(0); // Set to Zero when OOB so LDS could be issue without any preds.
}
}
synchronize();
}
CUTLASS_DEVICE void
begin_loop(int epi_m, int epi_n) {
if (epi_m == 0) {
if (!params.row_broadcast) return; // Do not issue LDS when row is scalar
Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
copy(tSR_sRow_flt, tSR_rRow_flt);
}
}
template <typename ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE Array<Element, FragmentSize>
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
Array<Element, FragmentSize> frg_row;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
frg_row[i] = tSR_rRow(epi_v * FragmentSize + i);
}
return frg_row;
}
};
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class... Args
>
CUTLASS_DEVICE auto
get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
auto [M, N, K, L] = args.problem_shape_mnkl;
auto [m, n, k, l] = args.tile_coord_mnkl;
using ThreadCount = decltype(size(args.tiled_copy));
Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
Tensor sRow = make_tensor(make_smem_ptr(smem),
make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N)
//// G2S: Gmem to Smem
auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},
Layout< Shape<_1, ThreadCount>,
Stride<_0, _1>>{},
Layout<_1>{});
auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
Tensor tGS_gRow = thr_g2s.partition_S(gRow);
Tensor tGS_sRow = thr_g2s.partition_D(sRow);
//// G2S: Coord
auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})));
Tensor tGS_cRow = thr_g2s.partition_S(cRow);
//// S2R: Smem to Reg
Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>(
tGS_gRow,
tGS_sRow,
tGS_cRow, tiled_g2s,
tSR_sRow,
tSR_rRow,
args.tCcD,
args.residue_cD,
ThreadCount{},
params);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Column vector broadcast
template<
int Stages,
class CtaTileShapeMNK,
class Element,
class StrideMNL = Stride<_1,_0,_0>,
int Alignment = 128 / sizeof_bits_v<Element>
>
struct Sm90ColOrScalarBroadcast {
static_assert(Stages == 0, "Column broadcast doesn't support smem usage yet");
static_assert(Alignment * sizeof_bits_v<Element> % 128 == 0, "sub-16B alignment not supported yet");
static_assert(
(cute::is_same_v<StrideMNL, Stride<_1,_0, _0>>) || // col vector broadcast, e.g. per-row alpha/bias
(cute::is_same_v<StrideMNL, Stride<_1,_0,int>>)); // batched col vector broadcast, e.g. batched per-row bias
// Accumulator distributes col elements evenly amongst threads so we can just directly load from gmem
struct SharedStorage { };
// This struct has been modified to have a bool indicating that ptr_col is a
// scalar that must be broadcast, instead of containing a scalar that is
// valid if ptr_col is null.
struct Arguments {
Element const* ptr_col = nullptr;
bool col_broadcast = true;
StrideMNL dCol = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
template <class ProblemShape>
static bool
can_implement(ProblemShape const& problem_shape, Arguments const& args) {
return true;
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
return 0;
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
CudaHostAdapter* cuda_adapter = nullptr) {
return cutlass::Status::kSuccess;
}
CUTLASS_DEVICE bool
is_producer_load_needed() const {
return false;
}
CUTLASS_DEVICE bool
is_C_load_needed() const {
return false;
}
CUTLASS_DEVICE bool
is_zero() const {
return (!params.col_broadcast && *(params.ptr_col) == Element(0));
}
CUTLASS_HOST_DEVICE
Sm90ColOrScalarBroadcast() { }
CUTLASS_HOST_DEVICE
Sm90ColOrScalarBroadcast(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params params;
template <class... Args>
CUTLASS_DEVICE auto
get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
return EmptyProducerLoadCallbacks{};
}
template<class GTensor, class RTensor, class CTensor, class ProblemShape>
struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
CUTLASS_DEVICE
ConsumerStoreCallbacks(
GTensor&& tCgCol,
RTensor&& tCrCol,
CTensor&& tCcCol,
ProblemShape problem_shape,
Params const& params
):
tCgCol(cute::forward<GTensor>(tCgCol)),
tCrCol(cute::forward<RTensor>(tCrCol)),
tCcCol(cute::forward<CTensor>(tCcCol)),
m(get<0>(problem_shape)),
params(params) {}
GTensor tCgCol; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
RTensor tCrCol;
CTensor tCcCol; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
Params const& params;
int m;
CUTLASS_DEVICE void
begin() {
Tensor pred = make_tensor<bool>(shape(tCgCol));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tCcCol(i)) < m;
}
if (!params.col_broadcast) {
fill(tCrCol, *(params.ptr_col));
return;
}
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_if(pred, filter(tCgCol), filter(tCrCol));
}
template <typename ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE Array<Element, FragmentSize>
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
Array<Element, FragmentSize> frg_col;
Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
frg_col[i] = tCrCol_mn(epi_v * FragmentSize + i);
}
return frg_col;
}
};
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class... Args
>
CUTLASS_DEVICE auto
get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
auto [M, N, K, L] = args.problem_shape_mnkl;
Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol);
Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
// Generate an identity tensor matching the shape of the global tensor and
// partition the same way, this will be used to generate the predicate
// tensor for loading
Tensor cCol = make_identity_tensor(mCol.shape());
Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
return ConsumerStoreCallbacks(
cute::move(tCgCol),
cute::move(tCrCol),
cute::move(tCcCol),
args.problem_shape_mnkl,
params
);
}
};
}
python/sglang/kernels/aot/csrc/cutlass_extensions/epilogue/epilogue_per_row_per_col_scale.h
+309
View File
@@ -0,0 +1,309 @@
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// Adapted from
// https://github.com/NVIDIA/TensorRT-LLM/blob/be1788106245496872d18e702978e59b6bfd50e0/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale.h
#pragma once
#include <cutlass/arch/memory.h>
#include <cutlass/numeric_conversion.h>
namespace cutlass {
namespace epilogue {
namespace threadblock {
template <
typename ThreadblockShape_,
int ThreadCount,
typename ScaleTileIterator_,
typename OutputTileIterator_,
typename ElementAccumulator_,
typename ElementCompute_,
typename ElementwiseFunctor_,
bool UseMasking_ = false>
class EpilogueVisitorPerRowPerCol {
public:
using ThreadblockShape = ThreadblockShape_;
static int const kThreadCount = ThreadCount;
using ScaleTileIterator = ScaleTileIterator_;
using OutputTileIterator = OutputTileIterator_;
using ElementwiseFunctor = ElementwiseFunctor_;
static int const kIterations = OutputTileIterator::kIterations;
static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
using ElementOutput = typename OutputTileIterator::Element;
using LayoutOutput = cutlass::layout::RowMajor;
using ElementAccumulator = ElementAccumulator_;
using AlphaScaleElementType = typename ScaleTileIterator::Element;
using ElementCompute = ElementCompute_;
using AccumulatorFragment = Array<ElementAccumulator, kElementsPerAccess>;
using ComputeFragment = Array<ElementCompute_, kElementsPerAccess>;
using OutputVector = Array<ElementOutput, kElementsPerAccess>;
static int const kThreadsPerRow = OutputTileIterator::ThreadMap::Detail::kAccessWidth;
static bool const kHasMultiStepsInRow = (OutputTileIterator::ThreadMap::Iterations::kColumn > 1);
/// Argument structure
struct Arguments {
typename ElementwiseFunctor::Params elementwise;
int64_t batch_stride_alpha;
int64_t batch_stride_C;
int64_t batch_stride_D;
//
// Methods
//
Arguments() : batch_stride_alpha(0), batch_stride_C(0), batch_stride_D(0) {}
Arguments(typename ElementwiseFunctor::Params elementwise_)
: elementwise(elementwise_), batch_stride_alpha(0), batch_stride_C(0), batch_stride_D(0) {}
Arguments(
typename ElementwiseFunctor::Params elementwise_,
int64_t batch_stride_alpha_,
int64_t batch_stride_C_,
int64_t batch_stride_D_)
: elementwise(elementwise_),
batch_stride_alpha(batch_stride_alpha_),
batch_stride_C(batch_stride_C_),
batch_stride_D(batch_stride_D_) {}
};
struct Params {
typename ElementwiseFunctor::Params elementwise;
int64_t batch_stride_alpha;
int64_t batch_stride_C;
int64_t batch_stride_D;
//
// Methods
//
CUTLASS_HOST_DEVICE
Params() {}
CUTLASS_HOST_DEVICE
Params(Arguments const& args)
: elementwise(args.elementwise),
batch_stride_alpha(args.batch_stride_alpha),
batch_stride_C(args.batch_stride_C),
batch_stride_D(args.batch_stride_D) {}
};
/// Shared storage
struct SharedStorage {};
private:
Params const& params_;
SharedStorage& shared_storage_;
MatrixCoord extent_;
MatrixCoord extent_real_;
ElementwiseFunctor elementwise_;
bool const with_bias_;
bool const per_token_quant_;
bool const per_channel_quant_;
AlphaScaleElementType* ptr_alpha_row_;
AlphaScaleElementType* ptr_alpha_col_;
ScaleTileIterator iterator_alpha_col_;
OutputTileIterator iterator_C_;
OutputTileIterator iterator_D_;
AlphaScaleElementType element_alpha_row_ = 1.0f;
AlphaScaleElementType element_alpha_col_ = 1.0f;
typename ScaleTileIterator::Fragment fragment_alpha_col_;
typename OutputTileIterator::Fragment fragment_C_;
typename OutputTileIterator::Fragment fragment_D_;
ElementAccumulator beta_;
int column_offset_;
MatrixCoord thread_offset_;
public:
CUTLASS_DEVICE
EpilogueVisitorPerRowPerCol(
Params const& params,
SharedStorage& shared_storage,
cutlass::MatrixCoord const& problem_size,
int thread_idx,
int warp_idx,
int lane_idx,
typename ScaleTileIterator::Params params_alpha_col,
typename OutputTileIterator::Params params_C,
typename OutputTileIterator::Params params_D,
bool with_bias,
bool per_token_quant,
bool per_channel_quant,
AlphaScaleElementType* ptr_alpha_row,
AlphaScaleElementType* ptr_alpha_col,
typename OutputTileIterator::Element* ptr_C,
typename OutputTileIterator::Element* ptr_D,
cutlass::MatrixCoord const& threadblock_offset = cutlass::MatrixCoord(0, 0),
int column_offset = 0,
cutlass::MatrixCoord const& problem_size_real = cutlass::MatrixCoord(0, 0))
: params_(params),
shared_storage_(shared_storage),
extent_(problem_size),
elementwise_(params.elementwise),
with_bias_(with_bias),
per_token_quant_(per_token_quant),
per_channel_quant_(per_channel_quant),
ptr_alpha_row_(ptr_alpha_row),
ptr_alpha_col_(ptr_alpha_col),
iterator_alpha_col_(params_alpha_col, ptr_alpha_col, problem_size, thread_idx, threadblock_offset),
iterator_C_(params_C, ptr_C, problem_size, thread_idx, threadblock_offset),
iterator_D_(params_D, ptr_D, problem_size, thread_idx, threadblock_offset),
extent_real_(problem_size_real) {
if (!per_channel_quant_ && (ptr_alpha_col_ != nullptr)) {
element_alpha_col_ = *ptr_alpha_col_;
}
if (!per_token_quant_ && (ptr_alpha_row_ != nullptr)) {
element_alpha_row_ = *ptr_alpha_row_;
}
}
/// Helper to indicate split-K behavior
CUTLASS_DEVICE
void set_k_partition(
int split_k_index, ///< Index of this threadblock within split-K partitioned scheme
int split_k_slices) { ///< Total number of split-K slices
}
/// Called to set the batch index
CUTLASS_DEVICE
void set_batch_index(int batch_idx) {
iterator_alpha_col_.add_pointer_offset(batch_idx * params_.batch_stride_alpha);
iterator_C_.add_pointer_offset(batch_idx * params_.batch_stride_C);
iterator_D_.add_pointer_offset(batch_idx * params_.batch_stride_D);
}
/// Called at the start of the epilogue just before iterating over accumulator slices
CUTLASS_DEVICE
void begin_epilogue() {
if (per_channel_quant_) {
iterator_alpha_col_.load(fragment_alpha_col_);
}
if (with_bias_) {
iterator_C_.load(fragment_C_);
}
}
/// Called at the start of one step before starting accumulator exchange
CUTLASS_DEVICE
void begin_step(int step_idx) {
fragment_D_.clear();
}
/// Called at the start of a row
CUTLASS_DEVICE
void begin_row(int row_idx) {
// load alpha_row in begin_step only when per token(row) scaling is used
if (per_token_quant_) {
int thread_offset_row =
iterator_D_.thread_start_row() + OutputTileIterator::ThreadMap::iteration_offset(row_idx).row();
arch::global_load<AlphaScaleElementType, sizeof(AlphaScaleElementType)>(
element_alpha_row_, ptr_alpha_row_ + thread_offset_row, thread_offset_row < extent_.row());
}
}
/// Called after accumulators have been exchanged for each accumulator vector
CUTLASS_DEVICE
void visit(int iter_idx, int row_idx, int column_idx, int frag_idx, AccumulatorFragment const& accum) {
NumericArrayConverter<ElementCompute, ElementAccumulator, kElementsPerAccess> source_converter;
ComputeFragment result = source_converter(accum);
if (per_channel_quant_) {
ComputeFragment alpha_col = reinterpret_cast<ComputeFragment*>(&fragment_alpha_col_)[column_idx];
result = per_token_channel_scale_accumulator_(result, alpha_col, element_alpha_row_);
} else {
result = per_token_scale_accumulator_(result, element_alpha_col_, element_alpha_row_);
}
if (with_bias_) {
NumericArrayConverter<ElementCompute, ElementOutput, kElementsPerAccess> bias_converter;
OutputVector bias = reinterpret_cast<OutputVector*>(&fragment_C_)[column_idx];
result = bias_accumulator_(result, bias_converter(bias));
}
// Convert to the output
NumericArrayConverter<ElementOutput, ElementCompute, kElementsPerAccess> output_converter;
OutputVector& output = reinterpret_cast<OutputVector*>(&fragment_D_)[frag_idx];
output = output_converter(result);
}
/// Called at the end of a row
CUTLASS_DEVICE
void end_row(int row_idx) {}
/// Called after all accumulator elements have been visited
CUTLASS_DEVICE
void end_step(int step_idx) {
iterator_D_.store(fragment_D_);
++iterator_D_;
}
/// Called after all steps have been completed
CUTLASS_DEVICE
void end_epilogue() {}
private:
CUTLASS_DEVICE
ComputeFragment per_token_channel_scale_accumulator_(
ComputeFragment const& accum, ComputeFragment const& scale_col, AlphaScaleElementType const& scale_row) {
ComputeFragment result;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < ComputeFragment::kElements; ++i) {
result[i] = accum[i] * (scale_col[i] * scale_row);
}
return result;
}
CUTLASS_DEVICE
ComputeFragment per_token_scale_accumulator_(
ComputeFragment const& accum, AlphaScaleElementType const& scale_col, AlphaScaleElementType const& scale_row) {
ComputeFragment result;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < ComputeFragment::kElements; ++i) {
result[i] = accum[i] * (scale_col * scale_row);
}
return result;
}
CUTLASS_DEVICE
ComputeFragment bias_accumulator_(ComputeFragment const& accum, ComputeFragment const& bias) {
ComputeFragment result;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < OutputVector::kElements; ++i) {
result[i] = accum[i] + bias[i];
}
return result;
}
};
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass

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