ViperEkura
d6f757dc13
- delete the warp-tiled kernel and both per-shape (N,K) selector tables; block size is 256 threads everywhere except M=8 with N*K <= 12 MiB, which keeps a 128-thread CTA
- HBM-streaming measurements (weight copies rotated through L2, the real decode regime) show the variants within ~3% on L20 because the kernel is bandwidth-bound; the retired tables were tuned against an L2-resident loop and sometimes picked the slowest variant ((2048,8192) M=8: coop128 6% slower than coop256)
- a shape no longer switches kernels (and accumulation order) with M, removing one shape-dependent nondeterminism source
- remove the stale split-K launcher comment
- move bf16_gemv.cu and bf16_swiglu.cu from csrc/kernels/gemv/ to csrc/kernels/ beside rotary_emb.cu; the family keeps no shared headers
- rename test_bf16_gemv_matches_half_cta_edge_bands to test_bf16_gemv_matches_m8_edge_bands and update docs/developer/cuda_kernels.md
Benchmark: L20 (sm_89), PyTorch 2.11.0+cu128, interleaved CUDA-event timing with rotated weight copies exceeding the 96MB L2; variant spread <=3% across 14 shapes x M in {1,2,4,8}, and the retained rule wins 5-9% at M=8 small weights ((512,3584), (1536,1536), (6912,1536))
114 lines
3.4 KiB
CMake
114 lines
3.4 KiB
CMake
cmake_minimum_required(VERSION 3.18)
|
|
project(astrai_kernels LANGUAGES CUDA CXX)
|
|
|
|
set(CMAKE_CXX_STANDARD 17)
|
|
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
|
set(CMAKE_CUDA_STANDARD 17)
|
|
|
|
find_package(CUDAToolkit REQUIRED)
|
|
|
|
if(NOT DEFINED TORCH_HOME)
|
|
set(TORCH_HOME "$ENV{TORCH_HOME}")
|
|
endif()
|
|
if(NOT TORCH_HOME)
|
|
message(FATAL_ERROR "TORCH_HOME must point at the torch install dir (site-packages/torch)")
|
|
endif()
|
|
|
|
if(NOT DEFINED PYTHON_INCLUDE_DIR)
|
|
set(PYTHON_INCLUDE_DIR "/usr/include/python${PYTHON_VERSION_MAJOR}.${PYTHON_VERSION_MINOR}")
|
|
endif()
|
|
|
|
if(NOT DEFINED ASTRAI_CUDA_ARCH)
|
|
if(DEFINED ENV{ASTRAI_CUDA_ARCH})
|
|
set(ASTRAI_CUDA_ARCH "$ENV{ASTRAI_CUDA_ARCH}")
|
|
else()
|
|
set(ASTRAI_CUDA_ARCH 80)
|
|
endif()
|
|
endif()
|
|
|
|
set(TORCH_LIB_DIR "${TORCH_HOME}/lib")
|
|
set(CUDA_LIB_DIR "/usr/local/cuda/lib64")
|
|
|
|
set(CXX_FLAGS -O3 -funroll-loops)
|
|
set(NVCC_FLAGS -O3
|
|
--expt-relaxed-constexpr
|
|
--use_fast_math
|
|
"--ptxas-options=-O3,-v"
|
|
--extra-device-vectorization
|
|
--threads=16)
|
|
|
|
set(TORCH_LIBS
|
|
"${TORCH_LIB_DIR}/libtorch_python.so"
|
|
"${TORCH_LIB_DIR}/libtorch_cuda.so"
|
|
"${TORCH_LIB_DIR}/libc10_cuda.so"
|
|
"${TORCH_LIB_DIR}/libtorch_cpu.so"
|
|
"${TORCH_LIB_DIR}/libtorch.so"
|
|
"${TORCH_LIB_DIR}/libc10.so"
|
|
CUDA::cudart)
|
|
|
|
set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}")
|
|
|
|
# Kernel registry — parallel lists of module names (.so / pybind names,
|
|
# globally unique across families) and their per-family source paths under
|
|
# kernels/. `loader.py` auto-discovers the .so files in astrai/extension/lib/,
|
|
# so this CMake registry is the single place to register a new kernel.
|
|
#
|
|
# FP8 MMA instructions require sm_89+. Keep the target out of the build on
|
|
# older architectures instead of instantiating templates that cannot compile.
|
|
# The remaining kernels are still useful on sm_80+ (including sm_86).
|
|
set(KERNEL_NAMES
|
|
attn_decode
|
|
attn_prefill
|
|
attn_paged_decode
|
|
attn_paged_prefill
|
|
bf16_gemv
|
|
bf16_swiglu
|
|
rotary_emb
|
|
)
|
|
set(KERNEL_SRCS
|
|
attention/decode.cu
|
|
attention/prefill.cu
|
|
attention/paged_decode.cu
|
|
attention/paged_prefill.cu
|
|
bf16_gemv.cu
|
|
bf16_swiglu.cu
|
|
rotary_emb.cu
|
|
)
|
|
|
|
if(ASTRAI_CUDA_ARCH GREATER_EQUAL 89)
|
|
list(APPEND KERNEL_NAMES fp8_ops)
|
|
list(APPEND KERNEL_SRCS fp8/ops.cu)
|
|
else()
|
|
message(WARNING
|
|
"FP8 operator disabled: ASTRAI_CUDA_ARCH=${ASTRAI_CUDA_ARCH} "
|
|
"requires compute capability 89 or newer")
|
|
endif()
|
|
|
|
list(LENGTH KERNEL_NAMES _kernel_count)
|
|
math(EXPR _kernel_last "${_kernel_count} - 1")
|
|
foreach(i RANGE ${_kernel_last})
|
|
list(GET KERNEL_NAMES ${i} name)
|
|
list(GET KERNEL_SRCS ${i} src)
|
|
add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${src}")
|
|
|
|
target_compile_definitions(${name} PRIVATE TORCH_EXTENSION_NAME=${name})
|
|
|
|
target_include_directories(${name} PRIVATE
|
|
"${CMAKE_CURRENT_SOURCE_DIR}/kernels"
|
|
"${TORCH_HOME}/include"
|
|
"${TORCH_HOME}/include/torch/csrc/api/include"
|
|
"${PYTHON_INCLUDE_DIR}")
|
|
|
|
target_link_libraries(${name} PRIVATE ${TORCH_LIBS})
|
|
target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}")
|
|
|
|
target_compile_options(${name} PRIVATE
|
|
$<$<COMPILE_LANGUAGE:CXX>:${CXX_FLAGS}>
|
|
$<$<COMPILE_LANGUAGE:CUDA>:${NVCC_FLAGS}>)
|
|
|
|
set_target_properties(${name} PROPERTIES
|
|
PREFIX ""
|
|
SUFFIX ".${PY_SOABI}.so"
|
|
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/../astrai/extension/lib")
|
|
endforeach()
|