feat(sgl-kernel): add InfLLM v2 attention kernels (#29383)
Co-authored-by: Size Wang <paulgeorge13hhhhh@gmail.com> Co-authored-by: lijiayi <lijiayi@modelbest.cn> Co-authored-by: suhmily10 <suhmily@gmail.com> Co-authored-by: Xiaoyue Xu <xiaoyue.xu.me@gmail.com> Co-authored-by: hansjohn <74091612+hansjohn@users.noreply.github.com> Co-authored-by: zhangyan <1762895426@qq.com>
This commit is contained in:
cauphe
co-authored by
Size Wang
lijiayi
suhmily10
Xiaoyue Xu
hansjohn
zhangyan
Size Wang
lijiayi
suhmily10
Xiaoyue Xu
hansjohn
zhangyan
parent
be70bfbdbb
commit
9bd02dc5b9
+1
-1
@@ -1,3 +1,3 @@
|
||||
[codespell]
|
||||
ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, fo, visibles, nothink, renderD, ond, tbe, CopyIn, notin, subtile, subtiles, IST
|
||||
ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, fo, visibles, nothink, renderD, ond, tbe, CopyIn, notin, subtile, subtiles, dout, IST
|
||||
skip = *.json, *.jsonl, *.patch, *.txt, *.lock
|
||||
|
||||
@@ -275,6 +275,8 @@ set(SOURCES
|
||||
"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"
|
||||
@@ -475,6 +477,68 @@ if (SGL_KERNEL_ENABLE_FA3)
|
||||
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"
|
||||
|
||||
@@ -49,6 +49,15 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
|
||||
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
|
||||
*/
|
||||
|
||||
@@ -0,0 +1,485 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2024, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
// Include these 2 headers instead of torch/extension.h since we don't need all of the torch headers.
|
||||
#include <ATen/cuda/CUDAGeneratorImpl.h> // For at::Generator and at::PhiloxCudaState
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <cutlass/numeric_types.h>
|
||||
#include <torch/nn/functional.h>
|
||||
#include <torch/python.h>
|
||||
|
||||
#include "flash.h"
|
||||
#include "hardware_info.h"
|
||||
#include "philox_unpack.cuh" // For at::cuda::philox::unpack
|
||||
#include "static_switch.h"
|
||||
|
||||
#define CHECK_DEVICE(x) TORCH_CHECK(x.is_cuda(), #x " must be on CUDA")
|
||||
#define CHECK_SHAPE(x, ...) \
|
||||
TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
|
||||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
||||
|
||||
void set_params_fprop(
|
||||
Flash_fwd_params& params,
|
||||
// sizes
|
||||
const size_t b,
|
||||
const size_t seqlen_q,
|
||||
const size_t seqlen_k,
|
||||
const size_t seqlen_q_rounded,
|
||||
const size_t seqlen_k_rounded,
|
||||
const size_t h,
|
||||
const size_t h_k,
|
||||
const size_t d,
|
||||
const size_t d_rounded,
|
||||
// device pointers
|
||||
const at::Tensor q,
|
||||
const at::Tensor k,
|
||||
const at::Tensor v,
|
||||
at::Tensor out,
|
||||
void* cu_seqlens_q_d,
|
||||
void* cu_seqlens_k_d,
|
||||
void* seqused_k,
|
||||
void* p_d,
|
||||
void* softmax_lse_d,
|
||||
float p_dropout,
|
||||
float softmax_scale,
|
||||
int window_size_left,
|
||||
int window_size_right,
|
||||
const float softcap,
|
||||
bool seqlenq_ngroups_swapped = false,
|
||||
const bool unpadded_lse = false) {
|
||||
// Reset the parameters
|
||||
params = {};
|
||||
|
||||
params.is_bf16 = q.dtype() == torch::kBFloat16;
|
||||
|
||||
// Set the pointers and strides.
|
||||
params.q_ptr = q.data_ptr();
|
||||
params.k_ptr = k.data_ptr();
|
||||
params.v_ptr = v.data_ptr();
|
||||
// All stride are in elements, not bytes.
|
||||
params.q_row_stride = q.stride(-3);
|
||||
params.k_row_stride = k.stride(-3);
|
||||
params.v_row_stride = v.stride(-3);
|
||||
params.q_head_stride = q.stride(-2);
|
||||
params.k_head_stride = k.stride(-2);
|
||||
params.v_head_stride = v.stride(-2);
|
||||
params.o_ptr = out.data_ptr();
|
||||
params.o_row_stride = params.o_ptr ? out.stride(-3) : 0;
|
||||
params.o_head_stride = params.o_ptr ? out.stride(-2) : 0;
|
||||
|
||||
if (cu_seqlens_q_d == nullptr) {
|
||||
params.q_batch_stride = q.stride(0);
|
||||
params.k_batch_stride = k.stride(0);
|
||||
params.v_batch_stride = v.stride(0);
|
||||
params.o_batch_stride = params.o_ptr ? out.stride(0) : 0;
|
||||
if (seqlenq_ngroups_swapped) {
|
||||
params.q_batch_stride *= seqlen_q;
|
||||
params.o_batch_stride *= seqlen_q;
|
||||
}
|
||||
}
|
||||
|
||||
params.cu_seqlens_q = static_cast<int*>(cu_seqlens_q_d);
|
||||
params.cu_seqlens_k = static_cast<int*>(cu_seqlens_k_d);
|
||||
params.seqused_k = static_cast<int*>(seqused_k);
|
||||
|
||||
// P = softmax(QK^T)
|
||||
params.p_ptr = p_d;
|
||||
|
||||
// Softmax sum
|
||||
params.softmax_lse_ptr = softmax_lse_d;
|
||||
|
||||
// Set the dimensions.
|
||||
params.b = b;
|
||||
params.h = h;
|
||||
params.h_k = h_k;
|
||||
params.h_h_k_ratio = h / h_k;
|
||||
params.seqlen_q = seqlen_q;
|
||||
params.seqlen_k = seqlen_k;
|
||||
params.seqlen_q_rounded = seqlen_q_rounded;
|
||||
params.seqlen_k_rounded = seqlen_k_rounded;
|
||||
params.d = d;
|
||||
params.d_rounded = d_rounded;
|
||||
|
||||
// Set the different scale values.
|
||||
#ifdef FLASHATTENTION_DISABLE_SOFTCAP
|
||||
TORCH_CHECK(softcap <= 0.0, "This flash attention build does not support softcap.");
|
||||
#endif
|
||||
if (softcap > 0.0) {
|
||||
params.softcap = softmax_scale / softcap;
|
||||
params.scale_softmax = softcap;
|
||||
params.scale_softmax_log2 = softcap * M_LOG2E;
|
||||
} else {
|
||||
// Remove potential NaN
|
||||
params.softcap = 0.0;
|
||||
params.scale_softmax = softmax_scale;
|
||||
params.scale_softmax_log2 = softmax_scale * M_LOG2E;
|
||||
}
|
||||
|
||||
// Set this to probability of keeping an element to simplify things.
|
||||
params.p_dropout = 1.f - p_dropout;
|
||||
// Convert p from float to int so we don't have to convert the random uint to float to compare.
|
||||
// [Minor] We want to round down since when we do the comparison we use <= instead of <
|
||||
// params.p_dropout_in_uint = uint32_t(std::floor(params.p_dropout * 4294967295.0));
|
||||
// params.p_dropout_in_uint16_t = uint16_t(std::floor(params.p_dropout * 65535.0));
|
||||
params.p_dropout_in_uint8_t = uint8_t(std::floor(params.p_dropout * 255.0));
|
||||
params.rp_dropout = 1.f / params.p_dropout;
|
||||
params.scale_softmax_rp_dropout = params.rp_dropout * params.scale_softmax;
|
||||
TORCH_CHECK(p_dropout < 1.f);
|
||||
#ifdef FLASHATTENTION_DISABLE_DROPOUT
|
||||
TORCH_CHECK(p_dropout == 0.0f, "This flash attention build does not support dropout.");
|
||||
#endif
|
||||
|
||||
// Causal is the special case where window_size_right == 0 and window_size_left < 0.
|
||||
// Local is the more general case where window_size_right >= 0 or window_size_left >= 0.
|
||||
params.is_causal = window_size_left < 0 && window_size_right == 0;
|
||||
|
||||
if (window_size_left < 0 && window_size_right >= 0) {
|
||||
window_size_left = seqlen_k;
|
||||
}
|
||||
if (window_size_left >= 0 && window_size_right < 0) {
|
||||
window_size_right = seqlen_k;
|
||||
}
|
||||
params.window_size_left = window_size_left;
|
||||
params.window_size_right = window_size_right;
|
||||
|
||||
#ifdef FLASHATTENTION_DISABLE_LOCAL
|
||||
TORCH_CHECK(
|
||||
params.is_causal || (window_size_left < 0 && window_size_right < 0),
|
||||
"This flash attention build does not support local attention.");
|
||||
#endif
|
||||
|
||||
params.is_seqlens_k_cumulative = true;
|
||||
|
||||
#ifdef FLASHATTENTION_DISABLE_UNEVEN_K
|
||||
TORCH_CHECK(d == d_rounded, "This flash attention build does not support headdim not being a multiple of 32.");
|
||||
#endif
|
||||
|
||||
params.unpadded_lse = unpadded_lse;
|
||||
params.seqlenq_ngroups_swapped = seqlenq_ngroups_swapped;
|
||||
}
|
||||
|
||||
void run_mha_fwd_split_stage1(Flash_fwd_params& params, cudaStream_t stream) {
|
||||
FP16_SWITCH(!params.is_bf16, [&] {
|
||||
HEADDIM_SWITCH(params.d, [&] {
|
||||
BOOL_SWITCH(params.is_causal, Is_causal, [&] {
|
||||
run_mha_fwd_splitkv_dispatch<elem_type, kHeadDim, Is_causal>(params, stream);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
void set_params_alibi(
|
||||
Flash_fwd_params& params, c10::optional<at::Tensor>& alibi_slopes_, int batch_size, int num_heads) {
|
||||
#ifdef FLASHATTENTION_DISABLE_ALIBI
|
||||
TORCH_CHECK(!alibi_slopes_.has_value(), "This flash attention build does not support alibi.");
|
||||
params.alibi_slopes_ptr = nullptr;
|
||||
#else
|
||||
if (alibi_slopes_.has_value()) {
|
||||
auto alibi_slopes = alibi_slopes_.value();
|
||||
TORCH_CHECK(alibi_slopes.dtype() == torch::kFloat32, "ALiBi slopes must have dtype fp32");
|
||||
CHECK_DEVICE(alibi_slopes);
|
||||
TORCH_CHECK(alibi_slopes.stride(-1) == 1, "ALiBi slopes tensor must have contiguous last dimension");
|
||||
TORCH_CHECK(
|
||||
alibi_slopes.sizes() == torch::IntArrayRef({num_heads}) ||
|
||||
alibi_slopes.sizes() == torch::IntArrayRef({batch_size, num_heads}));
|
||||
params.alibi_slopes_ptr = alibi_slopes.data_ptr();
|
||||
params.alibi_slopes_batch_stride = alibi_slopes.dim() == 2 ? alibi_slopes.stride(0) : 0;
|
||||
} else {
|
||||
params.alibi_slopes_ptr = nullptr;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> mha_varlen_fwd_stage1(
|
||||
at::Tensor& q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
|
||||
const at::Tensor& k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x
|
||||
// page_block_size x num_heads_k x head_size if there's a block_table.
|
||||
const at::Tensor& v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x
|
||||
// page_block_size x num_heads_k x head_size if there's a block_table.
|
||||
c10::optional<at::Tensor>& out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
|
||||
const at::Tensor& cu_seqlens_q, // b+1
|
||||
const at::Tensor& cu_seqlens_k, // b+1
|
||||
const at::Tensor& cu_seqlens_v, // b+1
|
||||
c10::optional<at::Tensor>&
|
||||
seqused_k, // b. If given, only this many elements of each batch element's keys are used.
|
||||
c10::optional<const at::Tensor>& leftpad_k_, // batch_size
|
||||
c10::optional<at::Tensor>& block_table_, // batch_size x max_num_blocks_per_seq
|
||||
c10::optional<at::Tensor>& alibi_slopes_, // num_heads or b x num_heads
|
||||
int max_seqlen_q,
|
||||
const int max_seqlen_k,
|
||||
const float p_dropout,
|
||||
const float softmax_scale,
|
||||
const bool zero_tensors,
|
||||
bool is_causal,
|
||||
int window_size_left,
|
||||
int window_size_right,
|
||||
const float softcap,
|
||||
const bool return_softmax,
|
||||
c10::optional<at::Generator> gen_) {
|
||||
// Otherwise the kernel will be launched from cuda:0 device
|
||||
at::cuda::CUDAGuard device_guard{q.device()};
|
||||
|
||||
auto [cc_major, cc_minor] = get_compute_capability(get_current_device());
|
||||
// bool is_sm75 = cc_major == 7 && cc_minor == 5;
|
||||
bool is_sm8x = cc_major == 8 && cc_minor >= 0;
|
||||
bool is_sm90 = cc_major == 9 && cc_minor == 0;
|
||||
// TORCH_CHECK(is_sm90 || is_sm8x, "FlashAttention only supports Ampere GPUs or newer.");
|
||||
// We will support Turing in the near future
|
||||
// TORCH_CHECK(is_sm90 || is_sm8x || is_sm75, "FlashAttention only supports Turing GPUs or newer.");
|
||||
|
||||
auto q_dtype = q.dtype();
|
||||
TORCH_CHECK(
|
||||
q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16, "FlashAttention only support fp16 and bf16 data type");
|
||||
if (q_dtype == torch::kBFloat16) {
|
||||
// TORCH_CHECK(is_sm90 || is_sm8x, "bfloat16 is only supported on Ampere GPUs or newer");
|
||||
}
|
||||
TORCH_CHECK(k.dtype() == q_dtype, "query and key must have the same dtype");
|
||||
TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype");
|
||||
TORCH_CHECK(cu_seqlens_q.dtype() == torch::kInt32, "cu_seqlens_q must have dtype int32");
|
||||
TORCH_CHECK(cu_seqlens_k.dtype() == torch::kInt32, "cu_seqlens_k must have dtype int32");
|
||||
TORCH_CHECK(cu_seqlens_v.dtype() == torch::kInt32, "cu_seqlens_v must have dtype int32");
|
||||
|
||||
CHECK_DEVICE(q);
|
||||
CHECK_DEVICE(k);
|
||||
CHECK_DEVICE(v);
|
||||
CHECK_DEVICE(cu_seqlens_q);
|
||||
CHECK_DEVICE(cu_seqlens_k);
|
||||
CHECK_DEVICE(cu_seqlens_v);
|
||||
|
||||
at::Tensor block_table;
|
||||
const bool paged_KV = block_table_.has_value();
|
||||
if (paged_KV) {
|
||||
block_table = block_table_.value();
|
||||
CHECK_DEVICE(block_table);
|
||||
TORCH_CHECK(block_table.dtype() == torch::kInt32, "block_table must have dtype torch.int32");
|
||||
TORCH_CHECK(block_table.stride(-1) == 1, "block_table must have contiguous last dimension");
|
||||
}
|
||||
|
||||
TORCH_CHECK(q.stride(-1) == 1, "Input tensor must have contiguous last dimension");
|
||||
TORCH_CHECK(k.stride(-1) == 1, "Input tensor must have contiguous last dimension");
|
||||
TORCH_CHECK(v.stride(-1) == 1, "Input tensor must have contiguous last dimension");
|
||||
CHECK_CONTIGUOUS(cu_seqlens_q);
|
||||
CHECK_CONTIGUOUS(cu_seqlens_k);
|
||||
CHECK_CONTIGUOUS(cu_seqlens_v);
|
||||
|
||||
const auto sizes = q.sizes();
|
||||
|
||||
const int batch_size = cu_seqlens_q.numel() - 1;
|
||||
int num_heads = sizes[1];
|
||||
const int head_size = sizes[2];
|
||||
const int num_heads_k = paged_KV ? k.size(2) : k.size(1);
|
||||
|
||||
if (softcap > 0.f) {
|
||||
TORCH_CHECK(p_dropout == 0.f, "Softcapping does not support dropout for now");
|
||||
}
|
||||
|
||||
const int max_num_blocks_per_seq = !paged_KV ? 0 : block_table.size(1);
|
||||
const int num_blocks = !paged_KV ? 0 : k.size(0);
|
||||
const int page_block_size = !paged_KV ? 1 : k.size(1);
|
||||
TORCH_CHECK(!paged_KV || page_block_size % 256 == 0, "Paged KV cache block size must be divisible by 256");
|
||||
|
||||
if (max_seqlen_q == 1 && !alibi_slopes_.has_value()) {
|
||||
is_causal = false;
|
||||
} // causal=true is the same as causal=false in this case
|
||||
if (is_causal) {
|
||||
window_size_right = 0;
|
||||
}
|
||||
|
||||
void* cu_seqlens_q_d = cu_seqlens_q.data_ptr();
|
||||
|
||||
// Faster to transpose q from (b, 1, (nheads_kv ngroups), d) to (b, ngroups, nheads_kv, d) in this case
|
||||
// H/t Daniel Haziza
|
||||
const int seqlenq_ngroups_swapped = max_seqlen_q == 1 && num_heads > num_heads_k && window_size_left < 0 &&
|
||||
window_size_right < 0 && p_dropout == 0.f && head_size % 8 == 0 &&
|
||||
!alibi_slopes_.has_value();
|
||||
const int ngroups = num_heads / num_heads_k;
|
||||
if (seqlenq_ngroups_swapped) {
|
||||
q = q.reshape({batch_size, num_heads_k, ngroups, head_size})
|
||||
.transpose(1, 2)
|
||||
.reshape({batch_size * ngroups, num_heads_k, head_size});
|
||||
max_seqlen_q = ngroups;
|
||||
num_heads = num_heads_k;
|
||||
cu_seqlens_q_d = nullptr;
|
||||
}
|
||||
|
||||
const int total_q = q.sizes()[0];
|
||||
|
||||
TORCH_CHECK(batch_size > 0, "batch size must be positive");
|
||||
TORCH_CHECK(head_size <= 256, "FlashAttention forward only supports head dimension at most 256");
|
||||
TORCH_CHECK(head_size % 8 == 0, "query, key, value, and out_ must have a head_size that is a multiple of 8");
|
||||
TORCH_CHECK(num_heads % num_heads_k == 0, "Number of heads in key/value must divide number of heads in query");
|
||||
|
||||
if (window_size_left >= max_seqlen_k) {
|
||||
window_size_left = -1;
|
||||
}
|
||||
if (window_size_right >= max_seqlen_k) {
|
||||
window_size_right = -1;
|
||||
}
|
||||
|
||||
CHECK_SHAPE(q, total_q, num_heads, head_size);
|
||||
if (!paged_KV) {
|
||||
const int total_k = k.size(0);
|
||||
CHECK_SHAPE(k, total_k, num_heads_k, head_size);
|
||||
// CHECK_SHAPE(v, total_k, num_heads_k, head_size);
|
||||
} else {
|
||||
CHECK_SHAPE(k, num_blocks, page_block_size, num_heads_k, head_size);
|
||||
// CHECK_SHAPE(v, num_blocks, page_block_size, num_heads_k, head_size);
|
||||
CHECK_SHAPE(block_table, batch_size, max_num_blocks_per_seq);
|
||||
}
|
||||
|
||||
CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
|
||||
CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
|
||||
CHECK_SHAPE(cu_seqlens_v, batch_size + 1);
|
||||
if (seqused_k.has_value()) {
|
||||
auto seqused_k_ = seqused_k.value();
|
||||
TORCH_CHECK(seqused_k_.dtype() == torch::kInt32, "seqused_k must have dtype int32");
|
||||
TORCH_CHECK(seqused_k_.is_cuda(), "seqused_k must be on CUDA device");
|
||||
TORCH_CHECK(seqused_k_.is_contiguous(), "seqused_k must be contiguous");
|
||||
CHECK_SHAPE(seqused_k_, batch_size);
|
||||
}
|
||||
|
||||
auto opts = q.options();
|
||||
at::Tensor out;
|
||||
out = torch::empty({0}, opts);
|
||||
// if (out_.has_value()) {
|
||||
// out = out_.value();
|
||||
// TORCH_CHECK(out.dtype() == q_dtype, "Output must have the same dtype as inputs");
|
||||
// CHECK_DEVICE(out);
|
||||
// TORCH_CHECK(out.stride(-1) == 1, "Output tensor must have contiguous last dimension");
|
||||
// CHECK_SHAPE(out, sizes[0], sizes[1], head_size);
|
||||
// if (seqlenq_ngroups_swapped) {
|
||||
// out = out.reshape({batch_size, num_heads_k, ngroups, head_size}).transpose(1, 2).reshape({batch_size *
|
||||
// ngroups, num_heads_k, head_size});
|
||||
// }
|
||||
// } else {
|
||||
// out = torch::empty_like(q);
|
||||
// }
|
||||
|
||||
auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
|
||||
const int head_size_rounded = head_size <= 192 ? round_multiple(head_size, 32) : 256;
|
||||
const int seqlen_q_rounded = round_multiple(max_seqlen_q, 128);
|
||||
const int seqlen_k_rounded = round_multiple(max_seqlen_k, 128);
|
||||
|
||||
// auto softmax_lse = torch::empty({num_heads, total_q}, opts.dtype(at::kFloat));
|
||||
at::Tensor p;
|
||||
// Only return softmax if there's dropout to reduce compilation time
|
||||
if (return_softmax) {
|
||||
// Return tensor with shape (num_heads_k, total_q, max_seqlen_k)
|
||||
p = torch::full({num_heads_k, total_q / 16, seqlen_k_rounded}, 0, opts);
|
||||
} else {
|
||||
p = torch::empty({0}, opts);
|
||||
}
|
||||
|
||||
if (zero_tensors) {
|
||||
// out.zero_();
|
||||
// softmax_lse.fill_(-std::numeric_limits<float>::infinity());
|
||||
if (return_softmax) {
|
||||
p.zero_();
|
||||
}
|
||||
}
|
||||
|
||||
Flash_fwd_params params;
|
||||
set_params_fprop(
|
||||
params,
|
||||
batch_size,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
seqlen_q_rounded,
|
||||
seqlen_k_rounded,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
head_size,
|
||||
head_size_rounded,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
out,
|
||||
cu_seqlens_q_d,
|
||||
cu_seqlens_k.data_ptr(),
|
||||
seqused_k.has_value() ? seqused_k.value().data_ptr() : nullptr,
|
||||
return_softmax ? p.data_ptr() : nullptr,
|
||||
nullptr, // softmax_lse.data_ptr(),
|
||||
p_dropout,
|
||||
softmax_scale,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
softcap,
|
||||
seqlenq_ngroups_swapped,
|
||||
/*unpadded_lse*/ true);
|
||||
|
||||
params.cu_seqlens_v = static_cast<int*>(cu_seqlens_v.data_ptr());
|
||||
params.is_seqlens_v_cumulative = true; // Treat cu_seqlens_v as cumulative sequence lengths
|
||||
// {
|
||||
// // Copy cu_seqlens_v to CPU for printing
|
||||
// at::Tensor cu_seqlens_v_cpu = cu_seqlens_v.to(torch::kCPU);
|
||||
// const int* cu_seqlens_v_data = cu_seqlens_v_cpu.data_ptr<int>();
|
||||
// printf("params.cu_seqlens_v: ");
|
||||
// for (int i = 0; i < batch_size + 1; ++i) {
|
||||
// printf("%d ", cu_seqlens_v_data[i]);
|
||||
// }
|
||||
// printf("\n");
|
||||
// }
|
||||
params.total_q = total_q;
|
||||
|
||||
params.m_block_dim = 16;
|
||||
params.n_block_dim = 1;
|
||||
|
||||
if (paged_KV) {
|
||||
params.block_table = block_table.data_ptr<int>();
|
||||
params.block_table_batch_stride = block_table.stride(0);
|
||||
params.k_batch_stride = k.stride(0);
|
||||
// params.v_batch_stride = v.stride(0);
|
||||
}
|
||||
params.page_block_size = page_block_size;
|
||||
// Keep references to these tensors to extend their lifetime
|
||||
|
||||
if (leftpad_k_.has_value()) {
|
||||
auto leftpad_k = leftpad_k_.value();
|
||||
TORCH_CHECK(!paged_KV, "We don't support Paged KV and leftpad_k running at the same time yet");
|
||||
TORCH_CHECK(leftpad_k.dtype() == torch::kInt32, "leftpad_k must have dtype int32");
|
||||
CHECK_DEVICE(leftpad_k);
|
||||
CHECK_CONTIGUOUS(leftpad_k);
|
||||
CHECK_SHAPE(leftpad_k, batch_size);
|
||||
params.leftpad_k = static_cast<int*>(leftpad_k.data_ptr());
|
||||
}
|
||||
|
||||
// number of times random will be generated per thread, to offset philox counter in thc random
|
||||
// state
|
||||
// We use a custom RNG that increases the offset by batch_size * nheads * 32.
|
||||
int64_t counter_offset = params.b * params.h * 32;
|
||||
auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
auto rng_state = torch::empty({2}, options.dtype(torch::kInt64));
|
||||
// Forward kernel will populate memory with the seed and offset.
|
||||
params.rng_state = reinterpret_cast<uint64_t*>(rng_state.data_ptr());
|
||||
|
||||
if (p_dropout > 0.0) {
|
||||
auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(gen_, at::cuda::detail::getDefaultCUDAGenerator());
|
||||
// See Note [Acquire lock when using random generators]
|
||||
std::lock_guard<std::mutex> lock(gen->mutex_);
|
||||
params.philox_args = gen->philox_cuda_state(counter_offset);
|
||||
}
|
||||
|
||||
set_params_alibi(params, alibi_slopes_, batch_size, num_heads);
|
||||
|
||||
if (max_seqlen_k > 0) {
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
params.num_splits = 1;
|
||||
run_mha_fwd_split_stage1(params, stream);
|
||||
} else {
|
||||
// If seqlen_k == 0, then we have an empty tensor. We need to set the output to 0.
|
||||
// out.zero_();
|
||||
// softmax_lse.fill_(std::numeric_limits<float>::infinity());
|
||||
}
|
||||
|
||||
if (seqlenq_ngroups_swapped) {
|
||||
int64_t size_before[] = {batch_size, max_seqlen_q, num_heads_k, head_size};
|
||||
int64_t size_after[] = {batch_size, num_heads_k * max_seqlen_q, head_size};
|
||||
// out = out.reshape(size_before).transpose(1, 2).reshape(size_after);
|
||||
q = q.reshape(size_before).transpose(1, 2).reshape(size_after);
|
||||
// softmax_lse = softmax_lse.reshape({num_heads * max_seqlen_q, batch_size});
|
||||
}
|
||||
|
||||
return {p};
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
namespace flash {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <bool Varlen = true>
|
||||
struct BlockInfo {
|
||||
template <typename Params>
|
||||
__device__ BlockInfo(const Params& params, const int bidb)
|
||||
: sum_s_q(!Varlen || params.cu_seqlens_q == nullptr ? -1 : params.cu_seqlens_q[bidb]),
|
||||
sum_s_k(
|
||||
!Varlen || params.cu_seqlens_k == nullptr || !params.is_seqlens_k_cumulative ? -1
|
||||
: params.cu_seqlens_k[bidb]),
|
||||
sum_s_v(
|
||||
!Varlen || params.cu_seqlens_v == nullptr || !params.is_seqlens_v_cumulative ? -1
|
||||
: params.cu_seqlens_v[bidb]),
|
||||
actual_seqlen_q(
|
||||
!Varlen || params.cu_seqlens_q == nullptr ? params.seqlen_q : params.cu_seqlens_q[bidb + 1] - sum_s_q)
|
||||
// If is_seqlens_k_cumulative, then seqlen_k is cu_seqlens_k[bidb + 1] - cu_seqlens_k[bidb].
|
||||
// Otherwise it's cu_seqlens_k[bidb], i.e., we use cu_seqlens_k to store the sequence lengths of K.
|
||||
,
|
||||
leftpad_k(params.leftpad_k == nullptr ? 0 : params.leftpad_k[bidb]),
|
||||
seqlen_k_cache(
|
||||
(!Varlen || params.cu_seqlens_k == nullptr
|
||||
? params.seqlen_k
|
||||
: (params.is_seqlens_k_cumulative ? params.cu_seqlens_k[bidb + 1] - sum_s_k
|
||||
: params.cu_seqlens_k[bidb])) -
|
||||
leftpad_k),
|
||||
actual_seqlen_k(
|
||||
params.seqused_k ? params.seqused_k[bidb] - leftpad_k
|
||||
: seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew))
|
||||
// If is_seqlens_v_cumulative, then seqlen_v is cu_seqlens_v[bidb + 1] - cu_seqlens_v[bidb].
|
||||
// Otherwise it's cu_seqlens_v[bidb], i.e., we use cu_seqlens_v to store the sequence lengths of V.
|
||||
,
|
||||
leftpad_v(params.leftpad_v == nullptr ? 0 : params.leftpad_v[bidb]),
|
||||
seqlen_v_cache(
|
||||
(!Varlen || params.cu_seqlens_v == nullptr
|
||||
? params.seqlen_v
|
||||
: (params.is_seqlens_v_cumulative ? params.cu_seqlens_v[bidb + 1] - sum_s_v
|
||||
: params.cu_seqlens_v[bidb])) -
|
||||
leftpad_v),
|
||||
actual_seqlen_c(
|
||||
params.seqused_v ? params.seqused_v[bidb] - leftpad_v
|
||||
: seqlen_v_cache + (params.vnew_ptr == nullptr ? 0 : params.seqlen_vnew)) {}
|
||||
|
||||
template <typename index_t>
|
||||
__forceinline__ __device__ index_t
|
||||
q_offset(const index_t batch_stride, const index_t row_stride, const int bidb) const {
|
||||
return sum_s_q == -1 ? bidb * batch_stride : uint32_t(sum_s_q) * row_stride;
|
||||
}
|
||||
|
||||
template <typename index_t>
|
||||
__forceinline__ __device__ index_t
|
||||
k_offset(const index_t batch_stride, const index_t row_stride, const int bidb) const {
|
||||
return sum_s_k == -1 ? bidb * batch_stride + leftpad_k * row_stride : uint32_t(sum_s_k + leftpad_k) * row_stride;
|
||||
}
|
||||
|
||||
template <typename index_t>
|
||||
__forceinline__ __device__ index_t
|
||||
v_offset(const index_t batch_stride, const index_t row_stride, const int bidb) const {
|
||||
return sum_s_v == -1 ? bidb * batch_stride + leftpad_v * row_stride : uint32_t(sum_s_v + leftpad_v) * row_stride;
|
||||
}
|
||||
|
||||
template <typename index_t>
|
||||
inline __device__ index_t blockmask_q_offset(const index_t m_block_dim, const int bidb) const {
|
||||
return sum_s_q == -1 ? bidb * (actual_seqlen_q / m_block_dim) : uint32_t(sum_s_q) / m_block_dim;
|
||||
}
|
||||
|
||||
const int sum_s_q;
|
||||
const int sum_s_k;
|
||||
const int sum_s_v;
|
||||
const int actual_seqlen_q;
|
||||
// We have to have seqlen_k_cache declared before actual_seqlen_k, otherwise actual_seqlen_k is set to 0.
|
||||
const int leftpad_k;
|
||||
const int seqlen_k_cache;
|
||||
const int actual_seqlen_k;
|
||||
// We have to have seqlen_v_cache declared before actual_seqlen_c, otherwise actual_seqlen_c is set to 0.
|
||||
const int leftpad_v;
|
||||
const int seqlen_v_cache;
|
||||
const int actual_seqlen_c;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace flash
|
||||
@@ -0,0 +1,103 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2024, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "philox.cuh"
|
||||
#include "utils.h"
|
||||
|
||||
namespace flash {
|
||||
|
||||
struct Dropout {
|
||||
const unsigned long long seed, offset;
|
||||
const uint8_t p_dropout_in_uint8_t;
|
||||
|
||||
__forceinline__ __device__ Dropout(
|
||||
const unsigned long long seed,
|
||||
const unsigned long long offset,
|
||||
const uint8_t p_dropout_in_uint8_t,
|
||||
const int bid,
|
||||
const int hid,
|
||||
const int tid,
|
||||
const int nheads)
|
||||
: seed(seed), offset(offset + (bid * nheads + hid) * 32 + tid % 32), p_dropout_in_uint8_t(p_dropout_in_uint8_t) {}
|
||||
|
||||
template <bool encode_dropout_in_sign_bit = false, typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void
|
||||
apply_dropout(Tensor<Engine, Layout>& tensor_, int block_row_start, int block_col_start, int block_row_stride) {
|
||||
// convert shape from (4, MMA_M, MMA_N) to (8, MMA_M, MMA_N / 2)
|
||||
Tensor tensor = make_tensor(tensor_.data(), flash::convert_layout_acc_dropout(tensor_.layout()));
|
||||
using T = typename Engine::value_type;
|
||||
auto encode_dropout = [](bool keep, T val) { return keep ? val : (encode_dropout_in_sign_bit ? -val : T(0)); };
|
||||
static_assert(decltype(size<2>(tensor))::value % 2 == 0);
|
||||
const uint16_t p_dropout_8bit_in_uint16_t = uint16_t(p_dropout_in_uint8_t);
|
||||
const uint32_t p_dropout_8bit_in_uint32_t =
|
||||
(uint32_t(p_dropout_8bit_in_uint16_t) << 16) | uint32_t(p_dropout_8bit_in_uint16_t);
|
||||
// if (cute::thread0()) { printf("threshold2 = 0x%x\n", p_dropout_8bit_in_uint32_t); }
|
||||
#pragma unroll
|
||||
for (int m = 0; m < size<1>(tensor); ++m, block_row_start += block_row_stride) {
|
||||
uint2 rowcol = make_uint2(block_row_start, block_col_start);
|
||||
#pragma unroll
|
||||
for (int n = 0; n < size<2>(tensor) / 2; ++n, ++rowcol.y) {
|
||||
// if (cute::thread(32, 0)) { printf("m = %d, n = %d, row = %d, col = %d\n", m, n, int(rowcol.x),
|
||||
// int(rowcol.y));}
|
||||
uint4 random_uint4 = flash::philox(seed, reinterpret_cast<unsigned long long&>(rowcol), offset);
|
||||
// if (cute::thread0()) { printf("philox = %u, %d, %d, %d\n", random_uint4.x, random_uint4.y, random_uint4.z,
|
||||
// random_uint4.w);}
|
||||
uint8_t (&rnd_8)[16] = reinterpret_cast<uint8_t (&)[16]>(random_uint4);
|
||||
// Special implementation for 16-bit types: we duplicate the threshold to the
|
||||
// low and high 16 bits of a 32-bit value, then use the f16x2 comparison instruction
|
||||
// to get a mask. The low 16 bits of the mask will be either 0xffff or 0x0000,
|
||||
// and the high 16 bits will be either 0xffff or 0x0000, depending on whether
|
||||
// the random value is less than the threshold.
|
||||
// We then do a bit-wise AND between the mask and the original value (in 32-bit).
|
||||
// We're exploiting the fact that floating point comparison is equivalent to integer
|
||||
// comparison, since we're comparing unsigned integers whose top 8-bits are zero.
|
||||
if (!encode_dropout_in_sign_bit &&
|
||||
(std::is_same<T, cutlass::half_t>::value || std::is_same<T, cutlass::bfloat16_t>::value)) {
|
||||
uint16_t rnd_16[16];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; i++) {
|
||||
rnd_16[i] = uint16_t(rnd_8[i]);
|
||||
}
|
||||
uint32_t (&rnd_32)[8] = reinterpret_cast<uint32_t (&)[8]>(rnd_16);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 2; j++) {
|
||||
Tensor tensor_uint32 = recast<uint32_t>(tensor(_, m, n * 2 + j));
|
||||
// if (cute::thread0()) { printf("random = 0x%x, 0x%x, 0x%x, 0x%x\n", rnd_32[j * 4 + 0], rnd_32[j * 4 + 1], rnd_32[j * 4
|
||||
// + 2], rnd_32[j * 4 + 3]); } if (cute::thread0()) { printf("tensor_uint32 = 0x%x, 0x%x, 0x%x, 0x%x\n",
|
||||
// tensor_uint32(0), tensor_uint32(1), tensor_uint32(2), tensor_uint32(3)); }
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
uint32_t mask;
|
||||
asm volatile("set.le.u32.f16x2 %0, %1, %2;\n"
|
||||
: "=r"(mask)
|
||||
: "r"(rnd_32[j * 4 + i]), "r"(p_dropout_8bit_in_uint32_t));
|
||||
tensor_uint32(i) &= mask;
|
||||
}
|
||||
// if (cute::thread0()) { printf("tensor_uint32 = 0x%x, 0x%x, 0x%x, 0x%x\n", tensor_uint32(0),
|
||||
// tensor_uint32(1), tensor_uint32(2), tensor_uint32(3)); }
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 2; j++) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i++) {
|
||||
tensor(i, m, n * 2 + j) =
|
||||
encode_dropout(rnd_8[j * 8 + i] <= p_dropout_in_uint8_t, tensor(i, m, n * 2 + j));
|
||||
}
|
||||
Tensor tensor_uint32 = recast<uint32_t>(tensor(_, m, n * 2 + j));
|
||||
// if (cute::thread0()) { printf("tensor_uint32 = 0x%x, 0x%x, 0x%x, 0x%x\n", tensor_uint32(0),
|
||||
// tensor_uint32(1), tensor_uint32(2), tensor_uint32(3)); }
|
||||
}
|
||||
}
|
||||
// // if ((threadIdx.x == 0) && (blockIdx.x == 0) && (blockIdx.y == 0)) {
|
||||
// // printf("n = %d, ph Philox: %u, %u, %u, %u\n", n, rnd_8.x, rnd_8.y, rnd_8.z, rnd_8.w);
|
||||
// // }
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace flash
|
||||
@@ -0,0 +1,152 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/cuda/CUDAGeneratorImpl.h> // For at::Generator and at::PhiloxCudaState
|
||||
#include <cuda.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
constexpr int TOTAL_DIM = 0;
|
||||
constexpr int H_DIM = 1;
|
||||
constexpr int D_DIM = 2;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Qkv_params {
|
||||
using index_t = int64_t;
|
||||
// The QKV matrices.
|
||||
void* __restrict__ q_ptr;
|
||||
void* __restrict__ k_ptr;
|
||||
void* __restrict__ v_ptr;
|
||||
|
||||
// The stride between rows of the Q, K and V matrices.
|
||||
index_t q_batch_stride;
|
||||
index_t k_batch_stride;
|
||||
index_t v_batch_stride;
|
||||
index_t q_row_stride;
|
||||
index_t k_row_stride;
|
||||
index_t v_row_stride;
|
||||
index_t q_head_stride;
|
||||
index_t k_head_stride;
|
||||
index_t v_head_stride;
|
||||
|
||||
// The number of heads.
|
||||
int h, h_k;
|
||||
// In the case of multi-query and grouped-query attention (MQA/GQA), nheads_k could be
|
||||
// different from nheads (query).
|
||||
int h_h_k_ratio; // precompute h / h_k,
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Flash_fwd_params : public Qkv_params {
|
||||
// The O matrix (output).
|
||||
void* __restrict__ o_ptr;
|
||||
void* __restrict__ oaccum_ptr;
|
||||
|
||||
// The stride between rows of O.
|
||||
index_t o_batch_stride;
|
||||
index_t o_row_stride;
|
||||
index_t o_head_stride;
|
||||
|
||||
// The pointer to the P matrix.
|
||||
void* __restrict__ p_ptr;
|
||||
|
||||
// The pointer to the softmax sum.
|
||||
void* __restrict__ softmax_lse_ptr;
|
||||
void* __restrict__ softmax_lseaccum_ptr;
|
||||
|
||||
// The dimensions.
|
||||
int b, seqlen_q, seqlen_k, seqlen_v, seqlen_knew, seqlen_vnew, d, seqlen_q_rounded, seqlen_k_rounded, d_rounded,
|
||||
rotary_dim, total_q;
|
||||
|
||||
// The scaling factors for the kernel.
|
||||
float scale_softmax;
|
||||
float scale_softmax_log2;
|
||||
|
||||
// array of length b+1 holding starting offset of each sequence.
|
||||
int* __restrict__ cu_seqlens_q;
|
||||
int* __restrict__ cu_seqlens_k;
|
||||
int* __restrict__ cu_seqlens_v;
|
||||
int* __restrict__ leftpad_k;
|
||||
int* __restrict__ leftpad_v;
|
||||
|
||||
// If provided, the actual length of each k sequence.
|
||||
int* __restrict__ seqused_k;
|
||||
int* __restrict__ seqused_v;
|
||||
uint64_t* __restrict__ blockmask;
|
||||
int m_block_dim, n_block_dim, num_k_heads;
|
||||
int num_blocks_m, num_blocks_n;
|
||||
|
||||
// The K_new and V_new matrices.
|
||||
void* __restrict__ knew_ptr;
|
||||
void* __restrict__ vnew_ptr;
|
||||
|
||||
// The stride between rows of the Q, K and V matrices.
|
||||
index_t knew_batch_stride;
|
||||
index_t vnew_batch_stride;
|
||||
index_t knew_row_stride;
|
||||
index_t vnew_row_stride;
|
||||
index_t knew_head_stride;
|
||||
index_t vnew_head_stride;
|
||||
|
||||
// The cos and sin matrices for rotary embedding.
|
||||
void* __restrict__ rotary_cos_ptr;
|
||||
void* __restrict__ rotary_sin_ptr;
|
||||
|
||||
// The indices to index into the KV cache.
|
||||
int* __restrict__ cache_batch_idx;
|
||||
|
||||
// Paged KV cache
|
||||
int* __restrict__ block_table;
|
||||
index_t block_table_batch_stride;
|
||||
int page_block_size;
|
||||
|
||||
// The dropout probability (probability of keeping an activation).
|
||||
float p_dropout;
|
||||
// uint32_t p_dropout_in_uint;
|
||||
// uint16_t p_dropout_in_uint16_t;
|
||||
uint8_t p_dropout_in_uint8_t;
|
||||
|
||||
// Scale factor of 1 / (1 - p_dropout).
|
||||
float rp_dropout;
|
||||
float scale_softmax_rp_dropout;
|
||||
|
||||
// Local window size
|
||||
int window_size_left, window_size_right;
|
||||
float softcap;
|
||||
|
||||
// Random state.
|
||||
at::PhiloxCudaState philox_args;
|
||||
|
||||
// Pointer to the RNG seed (idx 0) and offset (idx 1).
|
||||
uint64_t* rng_state;
|
||||
|
||||
bool is_bf16;
|
||||
bool is_causal;
|
||||
|
||||
// If is_seqlens_k_cumulative, then seqlen_k is cu_seqlens_k[bidb + 1] - cu_seqlens_k[bidb].
|
||||
// Otherwise it's cu_seqlens_k[bidb], i.e., we use cu_seqlens_k to store the sequence lengths of K.
|
||||
bool is_seqlens_k_cumulative;
|
||||
|
||||
// If is_seqlens_v_cumulative, then seqlen_v is cu_seqlens_v[bidb + 1] - cu_seqlens_v[bidb].
|
||||
// Otherwise it's cu_seqlens_v[bidb], i.e., we use cu_seqlens_v to store the sequence lengths of V.
|
||||
bool is_seqlens_v_cumulative;
|
||||
|
||||
bool is_rotary_interleaved;
|
||||
|
||||
int num_splits; // For split-KV version
|
||||
|
||||
void* __restrict__ alibi_slopes_ptr;
|
||||
index_t alibi_slopes_batch_stride;
|
||||
|
||||
bool unpadded_lse; // For varlen paths: LSE is in [nheads, total_seqlen_q] format instead of [b, nheads, seqlen_q].
|
||||
bool seqlenq_ngroups_swapped; // q has been transposed from (b, 1, (nheads_kv ngroups), d) to (b, ngroups, nheads_kv,
|
||||
// d).
|
||||
};
|
||||
|
||||
template <typename T, int Headdim, bool Is_causal>
|
||||
void run_mha_fwd_splitkv_dispatch(Flash_fwd_params& params, cudaStream_t stream);
|
||||
@@ -0,0 +1,108 @@
|
||||
#pragma once
|
||||
|
||||
namespace flash {
|
||||
|
||||
class fwdIterator {
|
||||
public:
|
||||
template <typename Params, typename BlockInfo>
|
||||
__device__ fwdIterator(
|
||||
const Params& params,
|
||||
const BlockInfo& binfo,
|
||||
const int kBlockM,
|
||||
const int kBlockN,
|
||||
const int batch_idx,
|
||||
const int head_idx,
|
||||
const int loop_step_idx,
|
||||
int n_block_min,
|
||||
int n_block_max) { // row first
|
||||
if (params.blockmask == nullptr) {
|
||||
blockmask_ptr = nullptr;
|
||||
return;
|
||||
}
|
||||
this->cache_seqlen_k = binfo.actual_seqlen_k - binfo.actual_seqlen_q / params.m_block_dim;
|
||||
this->max_block_idx = cute::ceil_div(binfo.actual_seqlen_k, params.n_block_dim);
|
||||
this->m_block_dim = params.m_block_dim;
|
||||
this->n_block_dim = params.n_block_dim;
|
||||
this->n_block_min = n_block_min;
|
||||
this->n_block_max = n_block_max;
|
||||
this->batch_idx = batch_idx; // Store batch_idx for debugging
|
||||
this->head_idx = head_idx;
|
||||
|
||||
// Calculate the offset for the uint64 blockmask
|
||||
const int num_blocks_m = params.num_blocks_m;
|
||||
const int num_blocks_n = params.num_blocks_n;
|
||||
const int uint64_per_row = (num_blocks_n + 64 - 1) / 64;
|
||||
const int row_offset = params.cu_seqlens_q != nullptr ? binfo.blockmask_q_offset(m_block_dim, batch_idx)
|
||||
: batch_idx * params.num_k_heads * params.num_blocks_m;
|
||||
|
||||
blockmask_ptr = params.blockmask + head_idx * params.num_blocks_m * uint64_per_row + row_offset * uint64_per_row +
|
||||
loop_step_idx * uint64_per_row;
|
||||
|
||||
// printf("blockmask_ptr = %d\n", blockmask_ptr);
|
||||
|
||||
const int q_block_idx = loop_step_idx + cache_seqlen_k;
|
||||
}
|
||||
|
||||
__device__ int max_no_larger(int target) const {
|
||||
if (blockmask_ptr == nullptr) {
|
||||
// printf("blockmask_ptr is nullptr\n");
|
||||
return target;
|
||||
}
|
||||
// printf("blockmask_ptr is NOT!!!! nullptr\n");
|
||||
if (max_block_idx == 0) {
|
||||
return -1;
|
||||
};
|
||||
|
||||
// 目标值不能超过最大块索引
|
||||
target = min(target, max_block_idx - 1);
|
||||
|
||||
// 计算相对于当前q_bit_position的实际位置
|
||||
int target_bit_pos = target;
|
||||
|
||||
// 确定此块在哪个uint64中
|
||||
int uint64_offset = target_bit_pos / 64;
|
||||
|
||||
// 确定此块在uint64中的哪一位
|
||||
int bit_pos = target_bit_pos % 64;
|
||||
|
||||
// 创建一个掩码，保留target及更低位的所有位
|
||||
uint64_t mask = bit_pos != 63 ? (1ULL << (bit_pos + 1)) - 1 : 0xFFFFFFFFFFFFFFFFULL;
|
||||
|
||||
// 检查当前uint64中target及以下的位
|
||||
uint64_t value = blockmask_ptr[uint64_offset] & mask;
|
||||
|
||||
// 如果当前uint64中有设置的位
|
||||
int result = -1;
|
||||
if (value != 0) {
|
||||
// 找到最高位的1（即不大于target的最大设置位）
|
||||
int highest_bit = 63 - __clzll(value); // __clzll计算前导0的数量
|
||||
result = highest_bit + (uint64_offset * 64);
|
||||
} else {
|
||||
// 如果当前uint64中没有找到，检查更低的uint64块
|
||||
for (int i = uint64_offset - 1; i >= 0; i--) {
|
||||
value = blockmask_ptr[i];
|
||||
if (value != 0) {
|
||||
// 找到最高位的1
|
||||
int highest_bit = 63 - __clzll(value);
|
||||
// 计算相对于q_bit_position的偏移
|
||||
result = highest_bit + (i * 64);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 没有找到设置位
|
||||
return result;
|
||||
}
|
||||
|
||||
uint64_t* blockmask_ptr;
|
||||
int row_offset; // 行偏移量
|
||||
int uint64_per_row; // 每行使用的uint64数量
|
||||
int cache_seqlen_k;
|
||||
int max_block_idx;
|
||||
int m_block_dim, n_block_dim;
|
||||
int n_block_min, n_block_max;
|
||||
int batch_idx, head_idx;
|
||||
};
|
||||
|
||||
} // namespace flash
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,335 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
|
||||
|
||||
#include "flash.h"
|
||||
#include "flash_fwd_kernel.h"
|
||||
#include "hardware_info.h"
|
||||
#include "static_switch.h"
|
||||
|
||||
// Determine if the architecture supports FLASH and define a macro to handle parameter modifiers
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
#define ARCH_SUPPORTS_FLASH
|
||||
#define KERNEL_PARAM_MODIFIER __grid_constant__
|
||||
#else
|
||||
#define KERNEL_PARAM_MODIFIER
|
||||
#endif
|
||||
|
||||
// Define a macro for unsupported architecture handling to centralize the error message
|
||||
#define FLASH_UNSUPPORTED_ARCH \
|
||||
printf("FATAL: FlashAttention requires building with sm version sm80-sm90, but was built for < 8.0!");
|
||||
|
||||
// Use a macro to clean up kernel definitions
|
||||
#define DEFINE_FLASH_FORWARD_KERNEL(kernelName, ...) \
|
||||
template <typename Kernel_traits, __VA_ARGS__> \
|
||||
__global__ void kernelName(KERNEL_PARAM_MODIFIER const Flash_fwd_params params)
|
||||
|
||||
DEFINE_FLASH_FORWARD_KERNEL(
|
||||
flash_fwd_kernel,
|
||||
bool Is_dropout,
|
||||
bool Is_causal,
|
||||
bool Is_local,
|
||||
bool Has_alibi,
|
||||
bool Is_even_MN,
|
||||
bool Is_even_K,
|
||||
bool Is_softcap,
|
||||
bool Return_softmax) {
|
||||
#if defined(ARCH_SUPPORTS_FLASH)
|
||||
static_assert(!(Is_causal && Is_local)); // Enforce constraints
|
||||
flash::compute_attn<
|
||||
Kernel_traits,
|
||||
Is_dropout,
|
||||
Is_causal,
|
||||
Is_local,
|
||||
Has_alibi,
|
||||
Is_even_MN,
|
||||
Is_even_K,
|
||||
Is_softcap,
|
||||
Return_softmax>(params);
|
||||
#else
|
||||
FLASH_UNSUPPORTED_ARCH
|
||||
#endif
|
||||
}
|
||||
|
||||
DEFINE_FLASH_FORWARD_KERNEL(
|
||||
flash_fwd_splitkv_kernel,
|
||||
bool Is_causal,
|
||||
bool Is_local,
|
||||
bool Has_alibi,
|
||||
bool Is_even_MN,
|
||||
bool Is_even_K,
|
||||
bool Is_softcap,
|
||||
bool Split,
|
||||
bool Append_KV) {
|
||||
#if defined(ARCH_SUPPORTS_FLASH)
|
||||
flash::compute_attn_splitkv<
|
||||
Kernel_traits,
|
||||
Is_causal,
|
||||
Is_local,
|
||||
Has_alibi,
|
||||
Is_even_MN,
|
||||
Is_even_K,
|
||||
Is_softcap,
|
||||
Split,
|
||||
Append_KV>(params);
|
||||
#else
|
||||
FLASH_UNSUPPORTED_ARCH
|
||||
#endif
|
||||
}
|
||||
|
||||
DEFINE_FLASH_FORWARD_KERNEL(
|
||||
flash_fwd_splitkv_stage1_kernel,
|
||||
bool Is_causal,
|
||||
bool Is_local,
|
||||
bool Has_alibi,
|
||||
bool Is_even_MN,
|
||||
bool Is_even_K,
|
||||
bool Is_softcap,
|
||||
bool Split,
|
||||
bool Append_KV) {
|
||||
#if defined(ARCH_SUPPORTS_FLASH)
|
||||
flash::compute_attn_splitkv_stage1<
|
||||
Kernel_traits,
|
||||
Is_causal,
|
||||
Is_local,
|
||||
Has_alibi,
|
||||
Is_even_MN,
|
||||
Is_even_K,
|
||||
Is_softcap,
|
||||
Split,
|
||||
Append_KV>(params);
|
||||
#else
|
||||
FLASH_UNSUPPORTED_ARCH
|
||||
#endif
|
||||
}
|
||||
|
||||
DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_splitkv_combine_kernel, int kBlockM, int Log_max_splits, bool Is_even_K) {
|
||||
static_assert(Log_max_splits >= 1);
|
||||
flash::combine_attn_seqk_parallel<Kernel_traits, kBlockM, Log_max_splits, Is_even_K>(params);
|
||||
}
|
||||
|
||||
template <typename Kernel_traits, bool Is_dropout, bool Is_causal>
|
||||
void run_flash_fwd(Flash_fwd_params& params, cudaStream_t stream) {
|
||||
constexpr size_t smem_size = Kernel_traits::kSmemSize;
|
||||
// printf("smem_size = %d\n", smem_size);
|
||||
|
||||
// Work-around for gcc 7. It doesn't like nested BOOL_SWITCH.
|
||||
// https://github.com/kokkos/kokkos-kernels/issues/349
|
||||
// https://github.com/HazyResearch/flash-attention/issues/21
|
||||
|
||||
const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
|
||||
dim3 grid(num_m_block, params.b, params.h);
|
||||
const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr &&
|
||||
params.seqlen_k % Kernel_traits::kBlockN == 0 &&
|
||||
params.seqlen_q % Kernel_traits::kBlockM == 0;
|
||||
const bool is_even_K = params.d == Kernel_traits::kHeadDim;
|
||||
// const bool return_softmax = params.p_ptr != nullptr;
|
||||
BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] {
|
||||
EVENK_SWITCH(is_even_K, IsEvenKConst, [&] {
|
||||
// LOCAL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_local, [&] {
|
||||
constexpr static bool Is_local = false;
|
||||
{ // TODO remove debug info
|
||||
// BOOL_SWITCH(return_softmax, ReturnSoftmaxConst, [&] {
|
||||
constexpr static bool ReturnSoftmaxConst = false;
|
||||
{ // TODO remove debug info
|
||||
// ALIBI_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] {
|
||||
constexpr static bool Has_alibi = false;
|
||||
{ // TODO remove debug info
|
||||
// SOFTCAP_SWITCH(params.softcap > 0.0, Is_softcap, [&] {
|
||||
constexpr static bool Is_softcap = false;
|
||||
{
|
||||
// Will only return softmax if dropout, to reduce compilation time.
|
||||
// If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates.
|
||||
// If return_softmax, set IsEvenMNConst to false to reduce number of templates
|
||||
// If head dim > 128, set IsEvenMNConst to false to reduce number of templates
|
||||
// If Is_local, set Is_causal to false
|
||||
auto kernel = &flash_fwd_kernel < Kernel_traits, Is_dropout && !Is_softcap, Is_causal,
|
||||
Is_local && !Is_causal, Has_alibi,
|
||||
IsEvenMNConst && IsEvenKConst && !Is_local && !ReturnSoftmaxConst && Kernel_traits::kHeadDim <= 128,
|
||||
IsEvenKConst, Is_softcap, ReturnSoftmaxConst && Is_dropout && !Is_softcap > ;
|
||||
// auto kernel = &flash_fwd_kernel<Kernel_traits, false, Is_causal, false, false, true, true, false>;
|
||||
// printf("IsEvenMNConst = %d, IsEvenKConst = %d, Is_local = %d, Is_causal = %d, ReturnSoftmaxConst = %d,
|
||||
// Is_dropout = %d\n", int(IsEvenMNConst), int(IsEvenKConst), int(Is_local), int(Is_causal),
|
||||
// int(ReturnSoftmaxConst), int(Is_dropout)); auto kernel = &flash_fwd_kernel<Kernel_traits, false,
|
||||
// Is_causal, false, true, true, false>;
|
||||
if (smem_size >= 48 * 1024) {
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
|
||||
}
|
||||
// int ctas_per_sm;
|
||||
// cudaError status_ = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
// &ctas_per_sm, kernel, Kernel_traits::kNThreads, smem_size);
|
||||
// printf("smem_size = %d, CTAs per SM = %d\n", int(smem_size), ctas_per_sm);
|
||||
kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
template <typename Kernel_traits, bool Is_causal>
|
||||
void run_flash_splitkv_fwd(Flash_fwd_params& params, cudaStream_t stream) {
|
||||
static_assert(!Kernel_traits::Is_Q_in_regs, "SplitKV implementation does not support Is_Q_in_regs");
|
||||
static_assert(!Kernel_traits::Share_Q_K_smem, "SplitKV implementation does not support Share_Q_K_smem");
|
||||
constexpr size_t smem_size = Kernel_traits::kSmemSize;
|
||||
const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
|
||||
dim3 grid(
|
||||
num_m_block,
|
||||
params.num_splits > 1 ? params.num_splits : params.b,
|
||||
params.num_splits > 1 ? params.b * params.h : params.h);
|
||||
const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr &&
|
||||
params.seqlen_k % Kernel_traits::kBlockN == 0 &&
|
||||
params.seqlen_q % Kernel_traits::kBlockM == 0;
|
||||
const bool is_even_K = params.d == Kernel_traits::kHeadDim;
|
||||
BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] {
|
||||
EVENK_SWITCH(is_even_K, IsEvenKConst, [&] {
|
||||
// LOCAL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_local, [&] {
|
||||
constexpr static bool Is_local = false;
|
||||
{ // TODO remove debug info
|
||||
BOOL_SWITCH(params.num_splits > 1, Split, [&] {
|
||||
BOOL_SWITCH(params.knew_ptr != nullptr, Append_KV, [&] {
|
||||
// ALIBI_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] {
|
||||
constexpr static bool Has_alibi = false;
|
||||
{ // TODO remove debug info
|
||||
// SOFTCAP_SWITCH(params.softcap > 0.0, Is_softcap, [&] {
|
||||
constexpr static bool Is_softcap = false;
|
||||
{ // TODO remove debug info
|
||||
// If Append_KV, then we must have seqlen_offsets, which means cu_seqlens_k != nullptr.
|
||||
// If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates.
|
||||
// If Is_local, set Is_causal to false
|
||||
auto kernel = &flash_fwd_splitkv_kernel < Kernel_traits, Is_causal, Is_local && !Is_causal, Has_alibi,
|
||||
IsEvenMNConst && !Append_KV && IsEvenKConst && !Is_local && Kernel_traits::kHeadDim <= 128,
|
||||
IsEvenKConst, Is_softcap, Split, Append_KV > ;
|
||||
// auto kernel = &flash_fwd_splitkv_kernel<Kernel_traits, Is_causal, false, true, Split, Append_KV>;
|
||||
// auto kernel = &flash_fwd_splitkv_kernel<Kernel_traits, Is_causal, false, IsEvenKConst>;
|
||||
if (smem_size >= 48 * 1024) {
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
|
||||
}
|
||||
kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
});
|
||||
});
|
||||
if (params.num_splits > 1) {
|
||||
// We want kBlockM to be as small as possible for more parallelism.
|
||||
// With 128 threads we can load 512 elements at a time, so if headdim is divisible by 128, kBlockM = 4.
|
||||
// If headdim is divisible by 64, then we set kBlockM = 8, etc.
|
||||
constexpr static int kBlockM =
|
||||
Kernel_traits::kHeadDim % 128 == 0 ? 4 : (Kernel_traits::kHeadDim % 64 == 0 ? 8 : 16);
|
||||
dim3 grid_combine((params.b * params.h * params.seqlen_q + kBlockM - 1) / kBlockM);
|
||||
EVENK_SWITCH(is_even_K, IsEvenKConst, [&] {
|
||||
if (params.num_splits <= 2) {
|
||||
flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 1, IsEvenKConst>
|
||||
<<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
|
||||
} else if (params.num_splits <= 4) {
|
||||
flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 2, IsEvenKConst>
|
||||
<<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
|
||||
} else if (params.num_splits <= 8) {
|
||||
flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 3, IsEvenKConst>
|
||||
<<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
|
||||
} else if (params.num_splits <= 16) {
|
||||
flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 4, IsEvenKConst>
|
||||
<<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
|
||||
} else if (params.num_splits <= 32) {
|
||||
flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 5, IsEvenKConst>
|
||||
<<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
|
||||
} else if (params.num_splits <= 64) {
|
||||
flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 6, IsEvenKConst>
|
||||
<<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
|
||||
} else if (params.num_splits <= 128) {
|
||||
flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 7, IsEvenKConst>
|
||||
<<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
|
||||
}
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Kernel_traits, bool Is_causal>
|
||||
void run_flash_splitkv_fwd_stage1(Flash_fwd_params& params, cudaStream_t stream) {
|
||||
static_assert(!Kernel_traits::Is_Q_in_regs, "SplitKV implementation does not support Is_Q_in_regs");
|
||||
static_assert(!Kernel_traits::Share_Q_K_smem, "SplitKV implementation does not support Share_Q_K_smem");
|
||||
constexpr size_t smem_size = Kernel_traits::kSmemSize;
|
||||
const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
|
||||
dim3 grid(
|
||||
num_m_block,
|
||||
params.num_splits > 1 ? params.num_splits : params.b,
|
||||
params.num_splits > 1 ? params.b * params.h : params.h);
|
||||
const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr &&
|
||||
params.seqlen_k % Kernel_traits::kBlockN == 0 &&
|
||||
params.seqlen_q % Kernel_traits::kBlockM == 0;
|
||||
const bool is_even_K = params.d == Kernel_traits::kHeadDim;
|
||||
BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] {
|
||||
EVENK_SWITCH(is_even_K, IsEvenKConst, [&] {
|
||||
// LOCAL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_local, [&] {
|
||||
constexpr static bool Is_local = false;
|
||||
{ // TODO remove debug info
|
||||
// BOOL_SWITCH(params.num_splits > 1, Split, [&] {
|
||||
constexpr static bool Split = false;
|
||||
{ // TODO remove debug info
|
||||
// BOOL_SWITCH(params.knew_ptr != nullptr, Append_KV, [&] {
|
||||
constexpr static bool Append_KV = false;
|
||||
{ // TODO remove debug info
|
||||
// ALIBI_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] {
|
||||
constexpr static bool Has_alibi = false;
|
||||
{ // TODO remove debug info
|
||||
// SOFTCAP_SWITCH(params.softcap > 0.0, Is_softcap, [&] {
|
||||
constexpr static bool Is_softcap = false;
|
||||
{ // TODO remove debug info
|
||||
// If Append_KV, then we must have seqlen_offsets, which means cu_seqlens_k != nullptr.
|
||||
// If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates.
|
||||
// If Is_local, set Is_causal to false
|
||||
auto kernel = &flash_fwd_splitkv_stage1_kernel < Kernel_traits, Is_causal, Is_local && !Is_causal,
|
||||
Has_alibi,
|
||||
IsEvenMNConst && !Append_KV && IsEvenKConst && !Is_local && Kernel_traits::kHeadDim <= 128,
|
||||
IsEvenKConst, Is_softcap, Split, Append_KV > ;
|
||||
if (smem_size >= 48 * 1024) {
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
|
||||
}
|
||||
kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
template <typename T, int Headdim, bool Is_causal>
|
||||
void run_mha_fwd_splitkv_dispatch(Flash_fwd_params& params, cudaStream_t stream) {
|
||||
if (params.blockmask == nullptr) {
|
||||
constexpr static int kBlockM = 64; // Fixed for all head dimensions
|
||||
// TD [2023-08-28]: nvcc segfaults for headdim 96 with block size 64 x 256,
|
||||
// and for headdim 192 with block size 64 x 128.
|
||||
// Also for headdim 160 with block size 64 x 128 after the rotary addition.
|
||||
constexpr static int kBlockN = Headdim <= 64 ? 256 : (Headdim <= 128 ? 128 : 64);
|
||||
if (params.m_block_dim == 1) {
|
||||
run_flash_splitkv_fwd<Flash_fwd_kernel_traits<Headdim, kBlockM, kBlockN, 4, false, false, T>, Is_causal>(
|
||||
params, stream);
|
||||
} else {
|
||||
run_flash_splitkv_fwd_stage1<Flash_fwd_kernel_traits<Headdim, 16, 64, 1, false, false, T>, Is_causal>(
|
||||
params, stream);
|
||||
}
|
||||
} else if (params.cu_seqlens_q != nullptr) {
|
||||
constexpr static int kBlockM = 16;
|
||||
constexpr static int kBlockN = 64;
|
||||
run_flash_splitkv_fwd<Flash_fwd_kernel_traits<Headdim, kBlockM, kBlockN, 1, false, false, T>, Is_causal>(
|
||||
params, stream);
|
||||
} else {
|
||||
constexpr static int kBlockM = 64;
|
||||
constexpr static int kBlockN = 64;
|
||||
run_flash_splitkv_fwd<Flash_fwd_kernel_traits<Headdim, kBlockM, kBlockN, 4, false, false, T>, Is_causal>(
|
||||
params, stream);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
// Copyright (c) 2024, Tri Dao.
|
||||
// Splitting the different head dimensions to different files to speed up compilation.
|
||||
// This file is auto-generated. See "generate_kernels.py"
|
||||
|
||||
#include "flash_fwd_launch_template.h"
|
||||
|
||||
template void
|
||||
run_mha_fwd_splitkv_dispatch<cutlass::bfloat16_t, 128, true>(Flash_fwd_params& params, cudaStream_t stream);
|
||||
@@ -0,0 +1,8 @@
|
||||
// Copyright (c) 2024, Tri Dao.
|
||||
// Splitting the different head dimensions to different files to speed up compilation.
|
||||
// This file is auto-generated. See "generate_kernels.py"
|
||||
|
||||
#include "flash_fwd_launch_template.h"
|
||||
|
||||
template void
|
||||
run_mha_fwd_splitkv_dispatch<cutlass::bfloat16_t, 128, false>(Flash_fwd_params& params, cudaStream_t stream);
|
||||
@@ -0,0 +1,8 @@
|
||||
// Copyright (c) 2024, Tri Dao.
|
||||
// Splitting the different head dimensions to different files to speed up compilation.
|
||||
// This file is auto-generated. See "generate_kernels.py"
|
||||
|
||||
#include "flash_fwd_launch_template.h"
|
||||
|
||||
template void
|
||||
run_mha_fwd_splitkv_dispatch<cutlass::bfloat16_t, 64, true>(Flash_fwd_params& params, cudaStream_t stream);
|
||||
@@ -0,0 +1,8 @@
|
||||
// Copyright (c) 2024, Tri Dao.
|
||||
// Splitting the different head dimensions to different files to speed up compilation.
|
||||
// This file is auto-generated. See "generate_kernels.py"
|
||||
|
||||
#include "flash_fwd_launch_template.h"
|
||||
|
||||
template void
|
||||
run_mha_fwd_splitkv_dispatch<cutlass::bfloat16_t, 64, false>(Flash_fwd_params& params, cudaStream_t stream);
|
||||
@@ -0,0 +1,39 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2024, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <tuple>
|
||||
|
||||
#if !defined(__CUDACC_RTC__)
|
||||
#include "cuda_runtime.h"
|
||||
#endif
|
||||
|
||||
#define CHECK_CUDA(call) \
|
||||
do { \
|
||||
cudaError_t status_ = call; \
|
||||
if (status_ != cudaSuccess) { \
|
||||
fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, cudaGetErrorString(status_)); \
|
||||
exit(1); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
inline int get_current_device() {
|
||||
int device;
|
||||
CHECK_CUDA(cudaGetDevice(&device));
|
||||
return device;
|
||||
}
|
||||
|
||||
inline std::tuple<int, int> get_compute_capability(int device) {
|
||||
int capability_major, capability_minor;
|
||||
CHECK_CUDA(cudaDeviceGetAttribute(&capability_major, cudaDevAttrComputeCapabilityMajor, device));
|
||||
CHECK_CUDA(cudaDeviceGetAttribute(&capability_minor, cudaDevAttrComputeCapabilityMinor, device));
|
||||
return {capability_major, capability_minor};
|
||||
}
|
||||
|
||||
inline int get_num_sm(int device) {
|
||||
int multiprocessor_count;
|
||||
CHECK_CUDA(cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device));
|
||||
return multiprocessor_count;
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2024, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cutlass/numeric_types.h>
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/layout.h"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
template <int kHeadDim_, int kBlockM_, int kBlockN_, int kNWarps_, typename elem_type = cutlass::half_t>
|
||||
struct Flash_kernel_traits {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
using Element = elem_type;
|
||||
static constexpr bool Has_cp_async = true;
|
||||
#else
|
||||
using Element = cutlass::half_t;
|
||||
static constexpr bool Has_cp_async = false;
|
||||
#endif
|
||||
|
||||
using ElementAccum = float;
|
||||
using index_t = int64_t;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
using MMA_Atom_Arch = std::conditional_t<
|
||||
std::is_same_v<elem_type, cutlass::half_t>,
|
||||
MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
|
||||
MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>>;
|
||||
#else
|
||||
using MMA_Atom_Arch = MMA_Atom<SM75_16x8x8_F32F16F16F32_TN>;
|
||||
#endif
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750
|
||||
using SmemCopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, elem_type>;
|
||||
using SmemCopyAtomTransposed = Copy_Atom<SM75_U16x8_LDSM_T, elem_type>;
|
||||
#else
|
||||
using SmemCopyAtom = Copy_Atom<DefaultCopy, elem_type>;
|
||||
using SmemCopyAtomTransposed = Copy_Atom<DefaultCopy, elem_type>;
|
||||
#endif
|
||||
};
|
||||
|
||||
// If Share_Q_K_smem is true, that forces Is_Q_in_regs to be true
|
||||
template <
|
||||
int kHeadDim_,
|
||||
int kBlockM_,
|
||||
int kBlockN_,
|
||||
int kNWarps_,
|
||||
bool Is_Q_in_regs_ = false,
|
||||
bool Share_Q_K_smem_ = false,
|
||||
typename elem_type = cutlass::half_t,
|
||||
typename Base = Flash_kernel_traits<kHeadDim_, kBlockM_, kBlockN_, kNWarps_, elem_type>>
|
||||
struct Flash_fwd_kernel_traits : public Base {
|
||||
using Element = typename Base::Element;
|
||||
using ElementAccum = typename Base::ElementAccum;
|
||||
using index_t = typename Base::index_t;
|
||||
static constexpr bool Has_cp_async = Base::Has_cp_async;
|
||||
using SmemCopyAtom = typename Base::SmemCopyAtom;
|
||||
using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed;
|
||||
|
||||
static constexpr bool Share_Q_K_smem = Share_Q_K_smem_;
|
||||
static constexpr bool Is_Q_in_regs = Is_Q_in_regs_ || Share_Q_K_smem;
|
||||
|
||||
// The number of threads.
|
||||
static constexpr int kNWarps = kNWarps_;
|
||||
static constexpr int kNThreads = kNWarps * 32;
|
||||
|
||||
static constexpr int kBlockM = kBlockM_;
|
||||
static constexpr int kBlockN = kBlockN_;
|
||||
static constexpr int kHeadDim = kHeadDim_;
|
||||
static_assert(kHeadDim % 32 == 0);
|
||||
static constexpr int kBlockKSmem = kHeadDim % 64 == 0 ? 64 : 32;
|
||||
static constexpr int kBlockKGmem = kHeadDim % 128 == 0 ? 128 : (kHeadDim % 64 == 0 ? 64 : 32);
|
||||
static constexpr int kSwizzle = kBlockKSmem == 32 ? 2 : 3;
|
||||
|
||||
using TiledMma = TiledMMA<
|
||||
typename Base::MMA_Atom_Arch,
|
||||
Layout<Shape<Int<kNWarps>, _1, _1>>, // 4x1x1 or 8x1x1 thread group
|
||||
Tile<Int<16 * kNWarps>, _16, _16>>;
|
||||
|
||||
using SmemLayoutAtomQ = decltype(composition(
|
||||
Swizzle<kSwizzle, 3, 3>{},
|
||||
// This has to be kBlockKSmem, using kHeadDim gives wrong results for d=128
|
||||
Layout<Shape<_8, Int<kBlockKSmem>>, Stride<Int<kBlockKSmem>, _1>>{}));
|
||||
using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, Shape<Int<kBlockM>, Int<kHeadDim>>{}));
|
||||
|
||||
using SmemLayoutKV = decltype(tile_to_shape(SmemLayoutAtomQ{}, Shape<Int<kBlockN>, Int<kHeadDim>>{}));
|
||||
|
||||
// https://github.com/ColfaxResearch/cutlass-kernels/blob/a222587e6d59b93ba704853d3946fb686d8b8892/src/fmha/fmha_forward.cu#L434
|
||||
using SmemLayoutVtransposed =
|
||||
decltype(composition(SmemLayoutKV{}, make_layout(Shape<Int<kHeadDim>, Int<kBlockN>>{}, GenRowMajor{})));
|
||||
using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{}));
|
||||
|
||||
using SmemLayoutAtomO = decltype(composition(
|
||||
Swizzle<kSwizzle, 3, 3>{}, Layout<Shape<Int<8>, Int<kBlockKSmem>>, Stride<Int<kBlockKSmem>, _1>>{}));
|
||||
using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, Shape<Int<kBlockM>, Int<kHeadDim>>{}));
|
||||
using SmemCopyAtomO = Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>;
|
||||
using SmemCopyAtomOaccum = Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAccum>;
|
||||
|
||||
static constexpr int kSmemQSize = size(SmemLayoutQ{}) * sizeof(Element);
|
||||
static constexpr int kSmemKVSize = size(SmemLayoutKV{}) * 2 * sizeof(Element);
|
||||
static constexpr int kSmemSize = Share_Q_K_smem ? std::max(kSmemQSize, kSmemKVSize) : kSmemQSize + kSmemKVSize;
|
||||
|
||||
static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element);
|
||||
static_assert(kHeadDim % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad");
|
||||
// Using kBlockKSmem here is 6-10% faster than kBlockKGmem for d=128 because of bank conflicts.
|
||||
// For example, for d=128, smem is split into 2 "pages", each page takes care of columns
|
||||
// 0-63 and 64-127. If we have 16 threads per row for gmem read, when we write to smem,
|
||||
// thread 0 - 7 will write to the first page and thread 8 - 15 will write to the second page,
|
||||
// to the same banks.
|
||||
static constexpr int kGmemThreadsPerRow = kBlockKSmem / kGmemElemsPerLoad;
|
||||
static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow");
|
||||
using GmemLayoutAtom =
|
||||
Layout<Shape<Int<kNThreads / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>, Stride<Int<kGmemThreadsPerRow>, _1>>;
|
||||
|
||||
// We use CACHEGLOBAL instead of CACHEALWAYS for both Q and K/V, since we won't be reading
|
||||
// from the same address by the same threadblock. This is slightly faster.
|
||||
using Gmem_copy_struct = std::conditional_t<
|
||||
Has_cp_async,
|
||||
SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>,
|
||||
AutoVectorizingCopyWithAssumedAlignment<128>>;
|
||||
using GmemTiledCopyQKV = decltype(make_tiled_copy(
|
||||
Copy_Atom<Gmem_copy_struct, Element>{},
|
||||
GmemLayoutAtom{},
|
||||
Layout<Shape<_1, _8>>{})); // Val layout, 8 vals per read
|
||||
using GmemTiledCopyO = decltype(make_tiled_copy(
|
||||
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
|
||||
GmemLayoutAtom{},
|
||||
Layout<Shape<_1, _8>>{})); // Val layout, 8 vals per store
|
||||
|
||||
using GmemLayoutAtomOaccum = std::conditional_t<
|
||||
kBlockKSmem == 32,
|
||||
Layout<
|
||||
Shape<_16, _8>, // Thread layout, 8 threads per row
|
||||
Stride<_8, _1>>,
|
||||
Layout<
|
||||
Shape<_8, _16>, // Thread layout, 16 threads per row
|
||||
Stride<_16, _1>>>;
|
||||
using GmemTiledCopyOaccum = decltype(make_tiled_copy(
|
||||
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAccum>{},
|
||||
GmemLayoutAtomOaccum{},
|
||||
Layout<Shape<_1, _4>>{})); // Val layout, 4 vals per store
|
||||
using GmemLayoutAtomRotcossin = GmemLayoutAtom;
|
||||
using GmemTiledCopyRotcossin = decltype(make_tiled_copy(
|
||||
Copy_Atom<UniversalCopy<uint64_t>, Element>{},
|
||||
GmemLayoutAtomRotcossin{},
|
||||
Layout<Shape<_1, _4>>{})); // Val layout, 4 vals per load
|
||||
using GmemTiledCopyRotcossinCont = decltype(make_tiled_copy(
|
||||
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
|
||||
GmemLayoutAtomRotcossin{},
|
||||
Layout<Shape<_1, _8>>{})); // Val layout, 8 vals per load
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,365 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2024, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
namespace flash {
|
||||
|
||||
using namespace cute;
|
||||
|
||||
template <typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void
|
||||
apply_mask(Tensor<Engine, Layout>& tensor, const int max_seqlen_k, const int col_idx_offset_ = 0) {
|
||||
// tensor has shape (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
static_assert(Layout::rank == 2, "Only support 2D Tensor");
|
||||
const int lane_id = threadIdx.x % 32;
|
||||
const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2;
|
||||
#pragma unroll
|
||||
for (int nj = 0; nj < size<1, 1>(tensor); ++nj) {
|
||||
const int col_idx_base = col_idx_offset + nj * 8;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < size<1, 0>(tensor); ++j) {
|
||||
const int col_idx = col_idx_base + j;
|
||||
if (col_idx >= max_seqlen_k) {
|
||||
// Without the "make_coord" we get wrong results
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); ++mi) {
|
||||
tensor(mi, make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <bool HasWSLeft = true, typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void apply_mask_local(
|
||||
Tensor<Engine, Layout>& tensor,
|
||||
const int col_idx_offset_,
|
||||
const int max_seqlen_k,
|
||||
const int row_idx_offset,
|
||||
const int max_seqlen_q,
|
||||
const int warp_row_stride,
|
||||
const int window_size_left,
|
||||
const int window_size_right,
|
||||
const int m_block_dim = 1) {
|
||||
// tensor has shape (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
static_assert(Layout::rank == 2, "Only support 2D Tensor");
|
||||
const int lane_id = threadIdx.x % 32;
|
||||
const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2;
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0, 1>(tensor); ++mi) {
|
||||
const int row_idx_base = row_idx_offset + mi * warp_row_stride;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size<0, 0>(tensor); ++i) {
|
||||
const int row_idx = row_idx_base + i * 8;
|
||||
|
||||
// Apply m_block_dim scaling to get logical indices
|
||||
const int orig_row_idx = row_idx / m_block_dim;
|
||||
const int orig_max_seqlen_q = max_seqlen_q / m_block_dim;
|
||||
|
||||
const int col_idx_limit_left = std::max(0, orig_row_idx + max_seqlen_k - orig_max_seqlen_q - window_size_left);
|
||||
const int col_idx_limit_right =
|
||||
std::min(max_seqlen_k, orig_row_idx + 1 + max_seqlen_k - orig_max_seqlen_q + window_size_right);
|
||||
#pragma unroll
|
||||
for (int nj = 0; nj < size<1, 1>(tensor); ++nj) {
|
||||
const int col_idx_base = col_idx_offset + nj * 8;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < size<1, 0>(tensor); ++j) {
|
||||
const int col_idx = col_idx_base + j;
|
||||
if (col_idx >= col_idx_limit_right || (HasWSLeft && col_idx < col_idx_limit_left)) {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
// if (cute::thread0()) {
|
||||
// printf("mi = %d, i = %d, row_idx = %d, max_seqlen_k = %d\n", mi, i, row_idx, max_seqlen_k);
|
||||
// print(tensor(make_coord(i, mi), _));
|
||||
// // print(tensor(_, j + nj * size<1, 0>(tensor)));
|
||||
// }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void apply_mask_causal(
|
||||
Tensor<Engine, Layout>& tensor,
|
||||
const int col_idx_offset_,
|
||||
const int max_seqlen_k,
|
||||
const int row_idx_offset,
|
||||
const int max_seqlen_q,
|
||||
const int warp_row_stride,
|
||||
const int m_block_dim = 1) {
|
||||
// Causal masking is equivalent to local masking with window_size_left = infinity and window_size_right = 0
|
||||
apply_mask_local</*HasWSLeft=*/false>(
|
||||
tensor, col_idx_offset_, max_seqlen_k, row_idx_offset, max_seqlen_q, warp_row_stride, -1, 0, m_block_dim);
|
||||
}
|
||||
|
||||
template <typename Engine0, typename Layout0, typename Engine1, typename Layout1>
|
||||
__forceinline__ __device__ void apply_mask_causal_w_idx(
|
||||
Tensor<Engine0, Layout0>& tensor,
|
||||
Tensor<Engine1, Layout1> const& idx_rowcol,
|
||||
const int col_idx_offset_,
|
||||
const int max_seqlen_k,
|
||||
const int row_idx_offset) {
|
||||
// tensor has shape (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
static_assert(Layout0::rank == 2, "Only support 2D Tensor");
|
||||
static_assert(Layout1::rank == 2, "Only support 2D Tensor");
|
||||
CUTE_STATIC_ASSERT_V(size<0>(tensor) == size<0>(idx_rowcol));
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tensor) == size<1>(idx_rowcol));
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); ++mi) {
|
||||
const int col_idx_limit = std::min(max_seqlen_k, 1 + row_idx_offset + get<0>(idx_rowcol(mi, 0)));
|
||||
#pragma unroll
|
||||
for (int ni = 0; ni < size<1, 1>(tensor); ++ni) {
|
||||
if (col_idx_offset_ + get<1>(idx_rowcol(0, ni)) >= col_idx_limit) {
|
||||
tensor(mi, ni) = -INFINITY;
|
||||
}
|
||||
}
|
||||
// if (cute::thread0()) {
|
||||
// printf("ni = %d, j = %d, col_idx = %d, max_seqlen_k = %d\n", ni, j, col_idx, max_seqlen_k);
|
||||
// print(tensor(_, make_coord(j, ni)));
|
||||
// // print(tensor(_, j + ni * size<1, 0>(tensor)));
|
||||
// }
|
||||
}
|
||||
}
|
||||
|
||||
template <bool Is_causal, bool Is_local, bool Has_alibi>
|
||||
struct Mask {
|
||||
const int max_seqlen_k, max_seqlen_q;
|
||||
const int window_size_left, window_size_right;
|
||||
const float alibi_slope;
|
||||
const int m_block_dim;
|
||||
|
||||
__forceinline__ __device__ Mask(
|
||||
const int max_seqlen_k,
|
||||
const int max_seqlen_q,
|
||||
const int window_size_left,
|
||||
const int window_size_right,
|
||||
const float alibi_slope = 0.f,
|
||||
const int m_block_dim = 1)
|
||||
: max_seqlen_k(max_seqlen_k),
|
||||
max_seqlen_q(max_seqlen_q),
|
||||
window_size_left(window_size_left),
|
||||
window_size_right(window_size_right),
|
||||
alibi_slope(!Has_alibi ? 0.0 : alibi_slope),
|
||||
m_block_dim(m_block_dim) {};
|
||||
|
||||
// Causal_mask: whether this particular iteration needs causal masking
|
||||
template <bool Causal_mask = false, bool Is_even_MN = true, typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void apply_mask(
|
||||
Tensor<Engine, Layout>& tensor_, const int col_idx_offset_, const int row_idx_offset, const int warp_row_stride) {
|
||||
static_assert(!(Causal_mask && Is_local), "Cannot be both causal and local");
|
||||
static_assert(Layout::rank == 3, "Only support 3D Tensor");
|
||||
static_assert(decltype(size<0>(tensor_))::value == 4, "First dimension must be 4");
|
||||
static constexpr bool Need_masking = Has_alibi || Causal_mask || Is_local || !Is_even_MN;
|
||||
// if (cute::thread0()) { printf("Has_alibi = %d, Causal_mask=%d, Is_local=%d, Is_even_MN = %d, Need_masking =
|
||||
// %d\n", Has_alibi, Causal_mask, Is_local, Is_even_MN, Need_masking); }
|
||||
if constexpr (Need_masking) {
|
||||
// Reshape tensor_ from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
Tensor tensor = make_tensor(tensor_.data(), flash::convert_layout_acc_rowcol(tensor_.layout()));
|
||||
// Do we need both row and column indices, or just column incides?
|
||||
static constexpr bool Col_idx_only = !(Has_alibi && !Is_causal) && !Is_local && !Causal_mask;
|
||||
const int lane_id = threadIdx.x % 32;
|
||||
const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2;
|
||||
if constexpr (Col_idx_only) {
|
||||
#pragma unroll
|
||||
for (int nj = 0; nj < size<1, 1>(tensor); ++nj) {
|
||||
const int col_idx_base = col_idx_offset + nj * 8;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < size<1, 0>(tensor); ++j) {
|
||||
const int col_idx = col_idx_base + j;
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); ++mi) {
|
||||
// No causal, no local
|
||||
if constexpr (Has_alibi) {
|
||||
tensor(mi, make_coord(j, nj)) += alibi_slope * col_idx;
|
||||
}
|
||||
if constexpr (!Is_even_MN) {
|
||||
if (col_idx >= max_seqlen_k) {
|
||||
tensor(mi, make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0, 1>(tensor); ++mi) {
|
||||
const int row_idx_base = row_idx_offset + mi * warp_row_stride;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size<0, 0>(tensor); ++i) {
|
||||
const int row_idx = row_idx_base + i * 8;
|
||||
|
||||
const int orig_row_idx = row_idx / this->m_block_dim;
|
||||
const int orig_max_seqlen_q = max_seqlen_q / this->m_block_dim;
|
||||
|
||||
const int col_idx_limit_left =
|
||||
std::max(0, orig_row_idx + max_seqlen_k - orig_max_seqlen_q - window_size_left);
|
||||
const int col_idx_limit_right =
|
||||
std::min(max_seqlen_k, orig_row_idx + 1 + max_seqlen_k - orig_max_seqlen_q + window_size_right);
|
||||
#pragma unroll
|
||||
for (int nj = 0; nj < size<1, 1>(tensor); ++nj) {
|
||||
const int col_idx_base = col_idx_offset + nj * 8;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < size<1, 0>(tensor); ++j) {
|
||||
const int col_idx = col_idx_base + j;
|
||||
if constexpr (Has_alibi) {
|
||||
if constexpr (Is_causal) {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) += alibi_slope * col_idx;
|
||||
} else {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) -=
|
||||
alibi_slope * abs(orig_row_idx + max_seqlen_k - orig_max_seqlen_q - col_idx);
|
||||
}
|
||||
}
|
||||
if constexpr (Causal_mask) {
|
||||
if (col_idx >= col_idx_limit_right) {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
if constexpr (Is_local) {
|
||||
if (col_idx >= col_idx_limit_right || col_idx < col_idx_limit_left) {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
if constexpr (!Causal_mask && !Is_local && !Is_even_MN) {
|
||||
// Causal and Local already handles MN masking
|
||||
if (col_idx >= max_seqlen_k) {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Causal_mask: whether this particular iteration needs causal masking
|
||||
template <bool Causal_mask = false, bool Is_even_MN = true, typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void apply_mask_stage1(
|
||||
Tensor<Engine, Layout>& tensor_,
|
||||
const int col_idx_offset_,
|
||||
const int row_idx_offset,
|
||||
const int warp_row_stride,
|
||||
const int stride = 16) {
|
||||
static_assert(!(Causal_mask && Is_local), "Cannot be both causal and local");
|
||||
static_assert(Layout::rank == 3, "Only support 3D Tensor");
|
||||
static_assert(decltype(size<0>(tensor_))::value == 4, "First dimension must be 4");
|
||||
static constexpr bool Need_masking = Has_alibi || Causal_mask || Is_local || !Is_even_MN;
|
||||
// if (cute::thread0()) { printf("Has_alibi = %d, Causal_mask=%d, Is_local=%d, Is_even_MN = %d, Need_masking =
|
||||
// %d\n", Has_alibi, Causal_mask, Is_local, Is_even_MN, Need_masking); }
|
||||
if constexpr (Need_masking) {
|
||||
// Reshape tensor_ from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
Tensor tensor = make_tensor(tensor_.data(), flash::convert_layout_acc_rowcol(tensor_.layout()));
|
||||
// Do we need both row and column indices, or just column incides?
|
||||
static constexpr bool Col_idx_only = !(Has_alibi && !Is_causal) && !Is_local && !Causal_mask;
|
||||
const int lane_id = threadIdx.x % 32;
|
||||
const int col_idx_offset = col_idx_offset_ + (lane_id % 4) * 2;
|
||||
if constexpr (Col_idx_only) {
|
||||
#pragma unroll
|
||||
for (int nj = 0; nj < size<1, 1>(tensor); ++nj) {
|
||||
const int col_idx_base = col_idx_offset + nj * 8;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < size<1, 0>(tensor); ++j) {
|
||||
const int col_idx = col_idx_base + j;
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); ++mi) {
|
||||
// No causal, no local
|
||||
if constexpr (Has_alibi) {
|
||||
tensor(mi, make_coord(j, nj)) += alibi_slope * col_idx;
|
||||
}
|
||||
if constexpr (!Is_even_MN) {
|
||||
if (col_idx >= max_seqlen_k) {
|
||||
tensor(mi, make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0, 1>(tensor); ++mi) {
|
||||
const int row_idx_base = row_idx_offset + mi * warp_row_stride;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size<0, 0>(tensor); ++i) {
|
||||
const int row_idx = row_idx_base + i * 8;
|
||||
|
||||
const int orig_row_idx = row_idx / this->m_block_dim;
|
||||
const int orig_max_seqlen_q = max_seqlen_q / this->m_block_dim;
|
||||
|
||||
// 计算压缩后的max_seqlen_q
|
||||
const int compress_stride = stride;
|
||||
const int compressed_max_seqlen_q = (orig_max_seqlen_q - compress_stride + 1) / compress_stride;
|
||||
const int compressed_row_idx = (orig_row_idx - stride + 1) / stride;
|
||||
int _max_seqlen_k = max_seqlen_k; // compressed lse的时候max_seqlen_k是压缩后的长度
|
||||
const int offset_row_idx =
|
||||
std::max(0, (orig_row_idx + 1) / stride - 1 + _max_seqlen_k - compressed_max_seqlen_q);
|
||||
|
||||
const int col_idx_limit_left = std::max(0, (orig_row_idx - stride + 1) / stride - window_size_left);
|
||||
const int col_idx_limit_right = std::min(_max_seqlen_k, (offset_row_idx + window_size_right));
|
||||
// const int col_idx_limit_right = std::min(max_seqlen_k, (orig_row_idx - stride + 1) / stride +
|
||||
// window_size_right);
|
||||
|
||||
// if (cute::thread0()) {
|
||||
// if (stride == 64) {
|
||||
// printf("orig_row_idx = %d, orig_max_seqlen_q = %d, compressed_max_seqlen_q = %d, compressed_row_idx = %d,
|
||||
// _max_seqlen_k = %d, offset_row_idx = %d, col_idx_limit_left = %d, col_idx_limit_right = %d\n",
|
||||
// orig_row_idx, orig_max_seqlen_q, compressed_max_seqlen_q, compressed_row_idx, _max_seqlen_k,
|
||||
// offset_row_idx, col_idx_limit_left, col_idx_limit_right);
|
||||
// }
|
||||
// }
|
||||
// flash::cp_async_wait<0>(); __syncthreads();
|
||||
|
||||
#pragma unroll
|
||||
for (int nj = 0; nj < size<1, 1>(tensor); ++nj) {
|
||||
const int col_idx_base = col_idx_offset + nj * 8;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < size<1, 0>(tensor); ++j) {
|
||||
const int col_idx = col_idx_base + j;
|
||||
if constexpr (Causal_mask) {
|
||||
if (col_idx >= col_idx_limit_right) {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
if constexpr (Is_local) {
|
||||
if (col_idx >= col_idx_limit_right || col_idx < col_idx_limit_left) {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
if constexpr (!Causal_mask && !Is_local && !Is_even_MN) {
|
||||
// Causal and Local already handles MN masking
|
||||
if (col_idx >= max_seqlen_k) {
|
||||
tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void all_mask(Tensor<Engine, Layout>& tensor_) {
|
||||
static_assert(Layout::rank == 3, "Only support 3D Tensor");
|
||||
static_assert(decltype(size<0>(tensor_))::value == 4, "First dimension must be 4");
|
||||
Tensor tensor = make_tensor(tensor_.data(), flash::convert_layout_acc_rowcol(tensor_.layout()));
|
||||
#pragma unroll
|
||||
for (int nj = 0; nj < size<1, 1>(tensor); ++nj) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < size<1, 0>(tensor); ++j) {
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); ++mi) {
|
||||
tensor(mi, make_coord(j, nj)) = -INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
} // namespace flash
|
||||
@@ -0,0 +1,49 @@
|
||||
// Pytorch also has an implementation of Philox RNG:
|
||||
// https://github.com/pytorch/pytorch/blob/8ca3c881db3e3510fcb7725389f6a0633c9b992c/torch/csrc/jit/tensorexpr/cuda_random.h
|
||||
#pragma once
|
||||
// Philox CUDA.
|
||||
|
||||
namespace flash {
|
||||
|
||||
struct ull2 {
|
||||
unsigned long long x;
|
||||
unsigned long long y;
|
||||
};
|
||||
|
||||
__forceinline__ __device__ uint2 mulhilo32(const unsigned int a, const unsigned int b) {
|
||||
uint2* res;
|
||||
unsigned long long tmp;
|
||||
asm("mul.wide.u32 %0, %1, %2;\n\t" : "=l"(tmp) : "r"(a), "r"(b));
|
||||
res = (uint2*)(&tmp);
|
||||
return *res;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ uint4 philox_single_round(const uint4 ctr, const uint2 key) {
|
||||
constexpr unsigned long kPhiloxSA = 0xD2511F53;
|
||||
constexpr unsigned long kPhiloxSB = 0xCD9E8D57;
|
||||
uint2 res0 = mulhilo32(kPhiloxSA, ctr.x);
|
||||
uint2 res1 = mulhilo32(kPhiloxSB, ctr.z);
|
||||
uint4 ret = {res1.y ^ ctr.y ^ key.x, res1.x, res0.y ^ ctr.w ^ key.y, res0.x};
|
||||
return ret;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ uint4
|
||||
philox(unsigned long long seed, unsigned long long subsequence, unsigned long long offset) {
|
||||
constexpr unsigned long kPhilox10A = 0x9E3779B9;
|
||||
constexpr unsigned long kPhilox10B = 0xBB67AE85;
|
||||
uint2 key = reinterpret_cast<uint2&>(seed);
|
||||
uint4 counter;
|
||||
ull2* tmp = reinterpret_cast<ull2*>(&counter);
|
||||
tmp->x = offset;
|
||||
tmp->y = subsequence;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 6; i++) {
|
||||
counter = philox_single_round(counter, key);
|
||||
key.x += (kPhilox10A);
|
||||
key.y += (kPhilox10B);
|
||||
}
|
||||
uint4 output = philox_single_round(counter, key);
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace flash
|
||||
@@ -0,0 +1,4 @@
|
||||
// This is purely so that it works with torch 2.1. For torch 2.2+ we can include ATen/cuda/PhiloxUtils.cuh
|
||||
|
||||
#pragma once
|
||||
#include <ATen/cuda/detail/UnpackRaw.cuh>
|
||||
@@ -0,0 +1,175 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2024, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
#include "utils.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace flash {
|
||||
|
||||
using namespace cute;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
bool Is_even_K = true,
|
||||
bool Clear_OOB_K = true,
|
||||
typename Engine0,
|
||||
typename Layout0,
|
||||
typename Engine1,
|
||||
typename Layout1,
|
||||
typename Engine2,
|
||||
typename Layout2,
|
||||
typename Engine3,
|
||||
typename Layout3>
|
||||
__forceinline__ __device__ void copy_rotary_interleaved(
|
||||
Tensor<Engine0, Layout0> const& S,
|
||||
Tensor<Engine1, Layout1>& D,
|
||||
Tensor<Engine2, Layout2> const& Cos,
|
||||
Tensor<Engine2, Layout2> const& Sin,
|
||||
Tensor<Engine3, Layout3> const& identity_MN,
|
||||
const int max_MN,
|
||||
const int min_MN,
|
||||
const int dim,
|
||||
const int rotary_dim) {
|
||||
CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});
|
||||
CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA
|
||||
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K
|
||||
CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); // MMA_K
|
||||
static_assert(decltype(size<0>(S))::value == decltype(size<0>(Cos))::value * 2);
|
||||
static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32
|
||||
Tensor rCos = make_fragment_like(Cos);
|
||||
Tensor rSin = make_fragment_like(Sin);
|
||||
Tensor rS = make_fragment_like(S);
|
||||
#pragma unroll
|
||||
for (int m = 0; m < size<1>(S); ++m) {
|
||||
if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < size<2>(S); ++k) {
|
||||
if (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim) {
|
||||
cute::copy(S(_, m, k), rS(_, m, k));
|
||||
if (get<1>(identity_MN(0, 0, k)) < rotary_dim) {
|
||||
cute::copy(Cos(_, m, k), rCos(_, m, k));
|
||||
cute::copy(Sin(_, m, k), rSin(_, m, k));
|
||||
Tensor S_fp32 = convert_type<float>(rS(_, m, k));
|
||||
Tensor cos_fp32 = convert_type<float>(rCos(_, m, k));
|
||||
Tensor sin_fp32 = convert_type<float>(rSin(_, m, k));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size<0>(rS) / 2; ++i) {
|
||||
float real = S_fp32(2 * i) * cos_fp32(i) - S_fp32(2 * i + 1) * sin_fp32(i);
|
||||
float imag = S_fp32(2 * i) * sin_fp32(i) + S_fp32(2 * i + 1) * cos_fp32(i);
|
||||
S_fp32(2 * i) = real;
|
||||
S_fp32(2 * i + 1) = imag;
|
||||
}
|
||||
// Idk but I need to copy for the convert_type to work
|
||||
Tensor S_fp32_copy = make_fragment_like(S_fp32);
|
||||
cute::copy(S_fp32, S_fp32_copy);
|
||||
using T = typename Engine0::value_type;
|
||||
Tensor S_og_type = convert_type<T>(S_fp32_copy);
|
||||
cute::copy(S_og_type, rS(_, m, k));
|
||||
}
|
||||
cute::copy(rS(_, m, k), D(_, m, k));
|
||||
} else if (Clear_OOB_K) {
|
||||
cute::clear(D(_, m, k));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
bool Is_even_K = true,
|
||||
bool Clear_OOB_K = true,
|
||||
typename Engine0,
|
||||
typename Layout0,
|
||||
typename Engine1,
|
||||
typename Layout1,
|
||||
typename Engine2,
|
||||
typename Layout2,
|
||||
typename Engine3,
|
||||
typename Layout3>
|
||||
__forceinline__ __device__ void copy_rotary_contiguous(
|
||||
Tensor<Engine0, Layout0> const& S,
|
||||
Tensor<Engine1, Layout1>& D,
|
||||
Tensor<Engine2, Layout2> const& Cos,
|
||||
Tensor<Engine2, Layout2> const& Sin,
|
||||
Tensor<Engine3, Layout3> const& identity_MN,
|
||||
const int max_MN,
|
||||
const int min_MN,
|
||||
const int dim,
|
||||
const int rotary_dim) {
|
||||
CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});
|
||||
CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA
|
||||
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K
|
||||
CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(Cos)); // MMA
|
||||
CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin));
|
||||
static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32
|
||||
Tensor rCos = make_fragment_like(Cos);
|
||||
Tensor rSin = make_fragment_like(Sin);
|
||||
Tensor rS = make_fragment_like(S);
|
||||
Tensor rS_other = make_fragment_like(rS(_, 0, 0));
|
||||
#pragma unroll
|
||||
for (int m = 0; m < size<1>(S); ++m) {
|
||||
if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < size<2>(S); ++k) {
|
||||
if (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim) {
|
||||
cute::copy(S(_, m, k), rS(_, m, k));
|
||||
if (get<1>(identity_MN(0, 0, k)) < rotary_dim) {
|
||||
const bool is_left = get<1>(identity_MN(0, 0, k)) < rotary_dim / 2;
|
||||
Tensor gS_other =
|
||||
make_tensor(S(_, m, k).data() + (is_left ? rotary_dim / 2 : -rotary_dim / 2), S(_, m, k).layout());
|
||||
cute::copy(gS_other, rS_other);
|
||||
// if (cute::thread0()) { print_tensor(rS(_, m, k)); print_tensor(rS_other); }
|
||||
Tensor gCos = make_tensor(Cos(_, m, k).data() + (is_left ? 0 : -rotary_dim / 2), Cos(_, m, k).layout());
|
||||
Tensor gSin = make_tensor(Sin(_, m, k).data() + (is_left ? 0 : -rotary_dim / 2), Sin(_, m, k).layout());
|
||||
cute::copy(gCos, rCos(_, m, k));
|
||||
cute::copy(gSin, rSin(_, m, k));
|
||||
// if (cute::thread0()) { print_tensor(rCos(_, m, k)); print_tensor(rSin(_, m, k)); }
|
||||
Tensor S_fp32 = convert_type<float>(rS(_, m, k));
|
||||
Tensor S_other_fp32 = convert_type<float>(rS_other);
|
||||
Tensor cos_fp32 = convert_type<float>(rCos(_, m, k));
|
||||
Tensor sin_fp32 = convert_type<float>(rSin(_, m, k));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size<0>(rS); ++i) {
|
||||
S_fp32(i) = S_fp32(i) * cos_fp32(i) + S_other_fp32(i) * (is_left ? -sin_fp32(i) : sin_fp32(i));
|
||||
}
|
||||
// Idk but I need to copy for the convert_type to work
|
||||
Tensor S_fp32_copy = make_fragment_like(S_fp32);
|
||||
cute::copy(S_fp32, S_fp32_copy);
|
||||
using T = typename Engine0::value_type;
|
||||
Tensor S_og_type = convert_type<T>(S_fp32_copy);
|
||||
cute::copy(S_og_type, rS(_, m, k));
|
||||
// if (cute::thread0()) { print_tensor(rS(_, m, k)); }
|
||||
}
|
||||
cute::copy(rS(_, m, k), D(_, m, k));
|
||||
} else if (Clear_OOB_K) {
|
||||
cute::clear(D(_, m, k));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace flash
|
||||
@@ -0,0 +1,279 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2024, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cutlass/numeric_types.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
#include "philox.cuh"
|
||||
#include "utils.h"
|
||||
|
||||
namespace flash {
|
||||
|
||||
using namespace cute;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
bool zero_init = true,
|
||||
typename Engine0,
|
||||
typename Layout0,
|
||||
typename Engine1,
|
||||
typename Layout1,
|
||||
typename Operator>
|
||||
__device__ __forceinline__ void
|
||||
thread_reduce_(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1>& summary, Operator& op) {
|
||||
static_assert(Layout0::rank == 2, "Only support 2D Tensor");
|
||||
static_assert(Layout1::rank == 1, "Only support 1D Tensor");
|
||||
CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); mi++) {
|
||||
summary(mi) = zero_init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0));
|
||||
#pragma unroll
|
||||
for (int ni = 1; ni < size<1>(tensor); ni++) {
|
||||
summary(mi) = op(summary(mi), tensor(mi, ni));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
|
||||
__device__ __forceinline__ void
|
||||
quad_allreduce_(Tensor<Engine0, Layout0>& dst, Tensor<Engine1, Layout1>& src, Operator& op) {
|
||||
CUTE_STATIC_ASSERT_V(size(dst) == size(src));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size(dst); i++) {
|
||||
dst(i) = Allreduce<4>::run(src(i), op);
|
||||
}
|
||||
}
|
||||
|
||||
template <
|
||||
bool zero_init = true,
|
||||
typename Engine0,
|
||||
typename Layout0,
|
||||
typename Engine1,
|
||||
typename Layout1,
|
||||
typename Operator>
|
||||
__device__ __forceinline__ void
|
||||
reduce_(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1>& summary, Operator& op) {
|
||||
thread_reduce_<zero_init>(tensor, summary, op);
|
||||
quad_allreduce_(summary, summary, op);
|
||||
}
|
||||
|
||||
template <bool zero_init = true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
|
||||
__device__ __forceinline__ void reduce_max(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1>& max) {
|
||||
MaxOp<float> max_op;
|
||||
reduce_<zero_init>(tensor, max, max_op);
|
||||
}
|
||||
|
||||
template <bool zero_init = true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
|
||||
__device__ __forceinline__ void reduce_sum(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1>& sum) {
|
||||
SumOp<float> sum_op;
|
||||
thread_reduce_<zero_init>(tensor, sum, sum_op);
|
||||
}
|
||||
|
||||
// Apply the exp to all the elements.
|
||||
template <bool Scale_max = true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
|
||||
__forceinline__ __device__ void
|
||||
scale_apply_exp2(Tensor<Engine0, Layout0>& tensor, Tensor<Engine1, Layout1> const& max, const float scale) {
|
||||
static_assert(Layout0::rank == 2, "Only support 2D Tensor");
|
||||
static_assert(Layout1::rank == 1, "Only support 1D Tensor");
|
||||
CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); ++mi) {
|
||||
// If max is -inf, then all elements must have been -inf (possibly due to masking).
|
||||
// We don't want (-inf - (-inf)) since that would give NaN.
|
||||
// If we don't have float around M_LOG2E the multiplication is done in fp64.
|
||||
const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * (Scale_max ? scale : float(M_LOG2E));
|
||||
#pragma unroll
|
||||
for (int ni = 0; ni < size<1>(tensor); ++ni) {
|
||||
// Instead of computing exp(x - max), we compute exp2(x * log_2(e) -
|
||||
// max * log_2(e)) This allows the compiler to use the ffma
|
||||
// instruction instead of fadd and fmul separately.
|
||||
// The following macro will disable the use of fma.
|
||||
// See: https://github.com/pytorch/pytorch/issues/121558 for more details
|
||||
// This macro is set in PyTorch and not FlashAttention
|
||||
#ifdef UNFUSE_FMA
|
||||
tensor(mi, ni) = exp2f(__fmul_rn(tensor(mi, ni), scale) - max_scaled);
|
||||
#else
|
||||
tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply the exp to all the elements.
|
||||
template <bool Scale_max = true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
|
||||
__forceinline__ __device__ void get_softmax(
|
||||
Tensor<Engine0, Layout0>& tensor,
|
||||
Tensor<Engine1, Layout1> const& max,
|
||||
Tensor<Engine1, Layout1> const& sum,
|
||||
const float scale) {
|
||||
static_assert(Layout0::rank == 2, "Only support 2D Tensor");
|
||||
static_assert(Layout1::rank == 1, "Only support 1D Tensor");
|
||||
CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); ++mi) {
|
||||
// If max is -inf, then all elements must have been -inf (possibly due to masking).
|
||||
// We don't want (-inf - (-inf)) since that would give NaN.
|
||||
// If we don't have float around M_LOG2E the multiplication is done in fp64.
|
||||
const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * (Scale_max ? scale : float(M_LOG2E));
|
||||
const float sum_scaled = 1. / sum(mi);
|
||||
#pragma unroll
|
||||
for (int ni = 0; ni < size<1>(tensor); ++ni) {
|
||||
// Instead of computing exp(x - max), we compute exp2(x * log_2(e) -
|
||||
// max * log_2(e)) This allows the compiler to use the ffma
|
||||
// instruction instead of fadd and fmul separately.
|
||||
// The following macro will disable the use of fma.
|
||||
// See: https://github.com/pytorch/pytorch/issues/121558 for more details
|
||||
// This macro is set in PyTorch and not FlashAttention
|
||||
#ifdef UNFUSE_FMA
|
||||
tensor(mi, ni) = exp2f(__fmul_rn(tensor(mi, ni), scale) - max_scaled) * sum_scaled;
|
||||
#else
|
||||
tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled) * sum_scaled;
|
||||
#endif
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply the exp to all the elements.
|
||||
template <bool zero_init = true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
|
||||
__forceinline__ __device__ void max_scale_exp2_sum(
|
||||
Tensor<Engine0, Layout0>& tensor, Tensor<Engine1, Layout1>& max, Tensor<Engine1, Layout1>& sum, const float scale) {
|
||||
static_assert(Layout0::rank == 2, "Only support 2D Tensor");
|
||||
static_assert(Layout1::rank == 1, "Only support 1D Tensor");
|
||||
CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(tensor); ++mi) {
|
||||
MaxOp<float> max_op;
|
||||
max(mi) = zero_init ? tensor(mi, 0) : max_op(max(mi), tensor(mi, 0));
|
||||
#pragma unroll
|
||||
for (int ni = 1; ni < size<1>(tensor); ni++) {
|
||||
max(mi) = max_op(max(mi), tensor(mi, ni));
|
||||
}
|
||||
max(mi) = Allreduce<4>::run(max(mi), max_op);
|
||||
// If max is -inf, then all elements must have been -inf (possibly due to masking).
|
||||
// We don't want (-inf - (-inf)) since that would give NaN.
|
||||
const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * scale;
|
||||
sum(mi) = 0;
|
||||
#pragma unroll
|
||||
for (int ni = 0; ni < size<1>(tensor); ++ni) {
|
||||
// Instead of computing exp(x - max), we compute exp2(x * log_2(e) -
|
||||
// max * log_2(e)) This allows the compiler to use the ffma
|
||||
// instruction instead of fadd and fmul separately.
|
||||
tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
|
||||
sum(mi) += tensor(mi, ni);
|
||||
}
|
||||
SumOp<float> sum_op;
|
||||
sum(mi) = Allreduce<4>::run(sum(mi), sum_op);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <int kNRows>
|
||||
struct Softmax {
|
||||
using TensorT = decltype(make_tensor<float>(Shape<Int<kNRows>>{}));
|
||||
TensorT row_max, row_sum;
|
||||
|
||||
__forceinline__ __device__ Softmax() {};
|
||||
|
||||
template <bool Is_first, bool Check_inf = false, typename Tensor0, typename Tensor1>
|
||||
__forceinline__ __device__ void softmax_rescale_o(Tensor0& acc_s, Tensor1& acc_o, float softmax_scale_log2) {
|
||||
// Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout()));
|
||||
static_assert(decltype(size<0>(scores))::value == kNRows);
|
||||
if (Is_first) {
|
||||
flash::template reduce_max</*zero_init=*/true>(scores, row_max);
|
||||
flash::scale_apply_exp2(scores, row_max, softmax_scale_log2);
|
||||
flash::reduce_sum</*zero_init=*/true>(scores, row_sum);
|
||||
} else {
|
||||
Tensor scores_max_prev = make_fragment_like(row_max);
|
||||
cute::copy(row_max, scores_max_prev);
|
||||
flash::template reduce_max</*zero_init=*/false>(scores, row_max);
|
||||
// Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))
|
||||
Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout()));
|
||||
static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size(row_max); ++mi) {
|
||||
float scores_max_cur = !Check_inf ? row_max(mi) : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi));
|
||||
float scores_scale = exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2);
|
||||
row_sum(mi) *= scores_scale;
|
||||
#pragma unroll
|
||||
for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
|
||||
acc_o_rowcol(mi, ni) *= scores_scale;
|
||||
}
|
||||
}
|
||||
flash::scale_apply_exp2(scores, row_max, softmax_scale_log2);
|
||||
// We don't do the reduce across threads here since we don't need to use the row_sum.
|
||||
// We do that reduce at the end when we need to normalize the softmax.
|
||||
flash::reduce_sum</*zero_init=*/false>(scores, row_sum);
|
||||
}
|
||||
};
|
||||
|
||||
template <bool Is_first, bool Check_inf = false, typename Tensor0>
|
||||
__forceinline__ __device__ void softmax_rescale_simple(Tensor0& acc_s, float softmax_scale_log2) {
|
||||
// Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout()));
|
||||
static_assert(decltype(size<0>(scores))::value == kNRows);
|
||||
if (Is_first) {
|
||||
flash::template reduce_max</*zero_init=*/true>(scores, row_max);
|
||||
flash::scale_apply_exp2(scores, row_max, softmax_scale_log2);
|
||||
flash::reduce_sum</*zero_init=*/true>(scores, row_sum);
|
||||
} else {
|
||||
Tensor scores_max_prev = make_fragment_like(row_max);
|
||||
cute::copy(row_max, scores_max_prev);
|
||||
flash::template reduce_max</*zero_init=*/false>(scores, row_max);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size(row_max); ++mi) {
|
||||
float scores_max_cur = !Check_inf ? row_max(mi) : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi));
|
||||
float scores_scale = exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2);
|
||||
row_sum(mi) *= scores_scale;
|
||||
}
|
||||
flash::scale_apply_exp2(scores, row_max, softmax_scale_log2);
|
||||
// We don't do the reduce across threads here since we don't need to use the row_sum.
|
||||
// We do that reduce at the end when we need to normalize the softmax.
|
||||
flash::reduce_sum</*zero_init=*/false>(scores, row_sum);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Tensor0>
|
||||
__forceinline__ __device__ void softmax_rescale_gt(Tensor0& acc_s, float softmax_scale_log2) {
|
||||
// Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout()));
|
||||
static_assert(decltype(size<0>(scores))::value == kNRows);
|
||||
flash::get_softmax(scores, row_max, row_sum, softmax_scale_log2);
|
||||
};
|
||||
|
||||
__forceinline__ __device__ void get_row_sum() {
|
||||
SumOp<float> sum_op;
|
||||
quad_allreduce_(row_sum, row_sum, sum_op);
|
||||
}
|
||||
|
||||
template <bool Is_dropout = false, bool Split = false, typename Tensor0>
|
||||
__forceinline__ __device__ TensorT
|
||||
normalize_softmax_lse(Tensor0& acc_o, float softmax_scale, float rp_dropout = 1.0) {
|
||||
SumOp<float> sum_op;
|
||||
quad_allreduce_(row_sum, row_sum, sum_op);
|
||||
TensorT lse = make_fragment_like(row_sum);
|
||||
Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout()));
|
||||
static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) {
|
||||
float sum = row_sum(mi);
|
||||
float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum;
|
||||
lse(mi) = (sum == 0.f || sum != sum) ? (Split ? -INFINITY : INFINITY) : row_max(mi) * softmax_scale + __logf(sum);
|
||||
float scale = !Is_dropout ? inv_sum : inv_sum * rp_dropout;
|
||||
#pragma unroll
|
||||
for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
|
||||
acc_o_rowcol(mi, ni) *= scale;
|
||||
}
|
||||
}
|
||||
return lse;
|
||||
};
|
||||
};
|
||||
|
||||
} // namespace flash
|
||||
@@ -0,0 +1,94 @@
|
||||
// Inspired by
|
||||
// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
|
||||
// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
|
||||
|
||||
#pragma once
|
||||
|
||||
/// @param COND - a boolean expression to switch by
|
||||
/// @param CONST_NAME - a name given for the constexpr bool variable.
|
||||
/// @param ... - code to execute for true and false
|
||||
///
|
||||
/// Usage:
|
||||
/// ```
|
||||
/// BOOL_SWITCH(flag, BoolConst, [&] {
|
||||
/// some_function<BoolConst>(...);
|
||||
/// });
|
||||
/// ```
|
||||
|
||||
#define BOOL_SWITCH(COND, CONST_NAME, ...) \
|
||||
[&] { \
|
||||
if (COND) { \
|
||||
constexpr static bool CONST_NAME = true; \
|
||||
return __VA_ARGS__(); \
|
||||
} else { \
|
||||
constexpr static bool CONST_NAME = false; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
}()
|
||||
|
||||
#ifdef FLASHATTENTION_DISABLE_DROPOUT
|
||||
#define DROPOUT_SWITCH(COND, CONST_NAME, ...) \
|
||||
[&] { \
|
||||
constexpr static bool CONST_NAME = false; \
|
||||
return __VA_ARGS__(); \
|
||||
}()
|
||||
#else
|
||||
#define DROPOUT_SWITCH BOOL_SWITCH
|
||||
#endif
|
||||
|
||||
#ifdef FLASHATTENTION_DISABLE_ALIBI
|
||||
#define ALIBI_SWITCH(COND, CONST_NAME, ...) \
|
||||
[&] { \
|
||||
constexpr static bool CONST_NAME = false; \
|
||||
return __VA_ARGS__(); \
|
||||
}()
|
||||
#else
|
||||
#define ALIBI_SWITCH BOOL_SWITCH
|
||||
#endif
|
||||
|
||||
#ifdef FLASHATTENTION_DISABLE_UNEVEN_K
|
||||
#define EVENK_SWITCH(COND, CONST_NAME, ...) \
|
||||
[&] { \
|
||||
constexpr static bool CONST_NAME = true; \
|
||||
return __VA_ARGS__(); \
|
||||
}()
|
||||
#else
|
||||
#define EVENK_SWITCH BOOL_SWITCH
|
||||
#endif
|
||||
|
||||
#ifdef FLASHATTENTION_DISABLE_SOFTCAP
|
||||
#define SOFTCAP_SWITCH(COND, CONST_NAME, ...) \
|
||||
[&] { \
|
||||
constexpr static bool CONST_NAME = false; \
|
||||
return __VA_ARGS__(); \
|
||||
}()
|
||||
#else
|
||||
#define SOFTCAP_SWITCH BOOL_SWITCH
|
||||
#endif
|
||||
|
||||
#ifdef FLASHATTENTION_DISABLE_LOCAL
|
||||
#define LOCAL_SWITCH(COND, CONST_NAME, ...) \
|
||||
[&] { \
|
||||
constexpr static bool CONST_NAME = false; \
|
||||
return __VA_ARGS__(); \
|
||||
}()
|
||||
#else
|
||||
#define LOCAL_SWITCH BOOL_SWITCH
|
||||
#endif
|
||||
|
||||
#define FP16_SWITCH(COND, ...) \
|
||||
[&] { \
|
||||
using elem_type = cutlass::bfloat16_t; \
|
||||
return __VA_ARGS__(); \
|
||||
}()
|
||||
|
||||
#define HEADDIM_SWITCH(HEADDIM, ...) \
|
||||
[&] { \
|
||||
if (HEADDIM == 64) { \
|
||||
constexpr static int kHeadDim = 64; \
|
||||
return __VA_ARGS__(); \
|
||||
} else { \
|
||||
constexpr static int kHeadDim = 128; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
}()
|
||||
@@ -0,0 +1,483 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <assert.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
#include <cuda_bf16.h>
|
||||
#endif
|
||||
|
||||
#include <cutlass/array.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cutlass/numeric_conversion.h>
|
||||
#include <cutlass/numeric_types.h>
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
using namespace cute;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace flash {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ uint32_t relu2(const uint32_t x);
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ uint32_t relu2<cutlass::half_t>(const uint32_t x) {
|
||||
uint32_t res;
|
||||
const uint32_t zero = 0u;
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
asm volatile("max.f16x2 %0, %1, %2;\n" : "=r"(res) : "r"(x), "r"(zero));
|
||||
#else
|
||||
asm volatile(
|
||||
"{\n"
|
||||
"\t .reg .f16x2 sela;\n"
|
||||
"\t set.gtu.u32.f16x2 sela, %1, %2;\n"
|
||||
"\t and.b32 %0, sela, %1;\n"
|
||||
"}\n"
|
||||
: "=r"(res)
|
||||
: "r"(x), "r"(zero));
|
||||
#endif
|
||||
return res;
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
template <>
|
||||
__forceinline__ __device__ uint32_t relu2<cutlass::bfloat16_t>(const uint32_t x) {
|
||||
uint32_t res;
|
||||
const uint32_t zero = 0u;
|
||||
asm volatile("max.bf16x2 %0, %1, %2;\n" : "=r"(res) : "r"(x), "r"(zero));
|
||||
return res;
|
||||
}
|
||||
#endif
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ uint32_t convert_relu2(const float2 x);
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ uint32_t convert_relu2<cutlass::half_t>(const float2 x) {
|
||||
uint32_t res;
|
||||
const uint32_t a = reinterpret_cast<const uint32_t&>(x.x);
|
||||
const uint32_t b = reinterpret_cast<const uint32_t&>(x.y);
|
||||
asm volatile("cvt.rn.relu.f16x2.f32 %0, %1, %2;\n" : "=r"(res) : "r"(b), "r"(a));
|
||||
return res;
|
||||
}
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ uint32_t convert_relu2<cutlass::bfloat16_t>(const float2 x) {
|
||||
uint32_t res;
|
||||
const uint32_t a = reinterpret_cast<const uint32_t&>(x.x);
|
||||
const uint32_t b = reinterpret_cast<const uint32_t&>(x.y);
|
||||
asm volatile("cvt.rn.relu.bf16x2.f32 %0, %1, %2;\n" : "=r"(res) : "r"(b), "r"(a));
|
||||
return res;
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
struct MaxOp {
|
||||
__device__ __forceinline__ T operator()(T const& x, T const& y) {
|
||||
return x > y ? x : y;
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct MaxOp<float> {
|
||||
// This is slightly faster
|
||||
__device__ __forceinline__ float operator()(float const& x, float const& y) {
|
||||
return max(x, y);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
struct SumOp {
|
||||
__device__ __forceinline__ T operator()(T const& x, T const& y) {
|
||||
return x + y;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <int THREADS>
|
||||
struct Allreduce {
|
||||
static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
|
||||
template <typename T, typename Operator>
|
||||
static __device__ __forceinline__ T run(T x, Operator& op) {
|
||||
constexpr int OFFSET = THREADS / 2;
|
||||
x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
|
||||
return Allreduce<OFFSET>::run(x, op);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
struct Allreduce<2> {
|
||||
template <typename T, typename Operator>
|
||||
static __device__ __forceinline__ T run(T x, Operator& op) {
|
||||
x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
bool A_in_regs = false,
|
||||
bool B_in_regs = false,
|
||||
typename Tensor0,
|
||||
typename Tensor1,
|
||||
typename Tensor2,
|
||||
typename Tensor3,
|
||||
typename Tensor4,
|
||||
typename TiledMma,
|
||||
typename TiledCopyA,
|
||||
typename TiledCopyB,
|
||||
typename ThrCopyA,
|
||||
typename ThrCopyB>
|
||||
__forceinline__ __device__ void gemm(
|
||||
Tensor0& acc,
|
||||
Tensor1& tCrA,
|
||||
Tensor2& tCrB,
|
||||
Tensor3 const& tCsA,
|
||||
Tensor4 const& tCsB,
|
||||
TiledMma tiled_mma,
|
||||
TiledCopyA smem_tiled_copy_A,
|
||||
TiledCopyB smem_tiled_copy_B,
|
||||
ThrCopyA smem_thr_copy_A,
|
||||
ThrCopyB smem_thr_copy_B) {
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K
|
||||
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // M
|
||||
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N
|
||||
if (!A_in_regs) {
|
||||
cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
|
||||
}
|
||||
if (!B_in_regs) {
|
||||
cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size<2>(tCrA); ++i) {
|
||||
if (i < size<2>(tCrA) - 1) {
|
||||
if (!A_in_regs) {
|
||||
cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
|
||||
}
|
||||
if (!B_in_regs) {
|
||||
cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
|
||||
}
|
||||
}
|
||||
cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename Tensor0,
|
||||
typename Tensor1,
|
||||
typename Tensor2,
|
||||
typename Tensor3,
|
||||
typename TiledMma,
|
||||
typename TiledCopy,
|
||||
typename ThrCopy>
|
||||
__forceinline__ __device__ void gemm_rs(
|
||||
Tensor0& acc,
|
||||
Tensor1& tCrA,
|
||||
Tensor2& tCrB,
|
||||
Tensor3 const& tCsB,
|
||||
TiledMma tiled_mma,
|
||||
TiledCopy smem_tiled_copy_B,
|
||||
ThrCopy smem_thr_copy_B) {
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K
|
||||
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N
|
||||
cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size<2>(tCrA); ++i) {
|
||||
if (i < size<2>(tCrA) - 1) {
|
||||
cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
|
||||
}
|
||||
cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))
|
||||
template <typename Layout>
|
||||
__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
|
||||
static_assert(decltype(size<0>(acc_layout))::value == 4);
|
||||
static_assert(decltype(rank(acc_layout))::value == 3);
|
||||
auto l = logical_divide(acc_layout, Shape<_2>{}); // ((2, 2), MMA_M, MMA_N)
|
||||
return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2)
|
||||
// if using m16n8k16, or to (4, MMA_M, MMA_N) if using m16n8k8.
|
||||
template <typename MMA_traits, typename Layout>
|
||||
__forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
|
||||
using X = Underscore;
|
||||
static_assert(decltype(size<0>(acc_layout))::value == 4);
|
||||
static_assert(decltype(rank(acc_layout))::value == 3);
|
||||
constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});
|
||||
static_assert(mma_shape_K == 8 || mma_shape_K == 16);
|
||||
if constexpr (mma_shape_K == 8) {
|
||||
return acc_layout;
|
||||
} else {
|
||||
auto l = logical_divide(acc_layout, Shape<X, X, _2>{}); // (4, MMA_M, (2, MMA_N / 2)))
|
||||
return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2)
|
||||
template <typename Layout>
|
||||
__forceinline__ __device__ auto convert_layout_acc_dropout(Layout acc_layout) {
|
||||
using X = Underscore;
|
||||
static_assert(decltype(size<0>(acc_layout))::value == 4);
|
||||
static_assert(decltype(rank(acc_layout))::value == 3);
|
||||
auto l = logical_divide(acc_layout, Shape<X, X, _2>{}); // (4, MMA_M, (2, MMA_N / 2)))
|
||||
return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename To_type, typename Engine, typename Layout>
|
||||
__forceinline__ __device__ auto convert_type(Tensor<Engine, Layout> const& tensor) {
|
||||
using From_type = typename Engine::value_type;
|
||||
constexpr int numel = decltype(size(tensor))::value;
|
||||
cutlass::NumericArrayConverter<To_type, From_type, numel> convert_op;
|
||||
// HACK: this requires tensor to be "contiguous"
|
||||
auto frag = convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel>*>(tensor.data()));
|
||||
return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void relu_(Tensor<Engine, Layout>& tensor) {
|
||||
constexpr int numel = decltype(size(tensor))::value;
|
||||
static_assert(numel % 2 == 0);
|
||||
using value_t = typename Engine::value_type;
|
||||
// HACK: this requires tensor to be "contiguous"
|
||||
Tensor tensor_uint32 = recast<uint32_t>(tensor);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size(tensor_uint32); ++i) {
|
||||
tensor_uint32(i) = relu2<value_t>(tensor_uint32(i));
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// On SM80 and above, we can fuse fp32 -> fp16/bf16 conversion and relu into 1 instruction
|
||||
template <typename To_type, typename Engine, typename Layout>
|
||||
__forceinline__ __device__ auto convert_type_relu(Tensor<Engine, Layout> const& tensor) {
|
||||
using From_type = typename Engine::value_type;
|
||||
static_assert(std::is_same_v<To_type, cutlass::half_t> || std::is_same_v<To_type, cutlass::bfloat16_t>);
|
||||
static_assert(std::is_same_v<float, From_type>);
|
||||
constexpr int numel = decltype(size(tensor))::value;
|
||||
static_assert(numel % 2 == 0);
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
// HACK: this requires tensor to be "contiguous"
|
||||
Tensor tensor_float2 = recast<float2>(tensor);
|
||||
Tensor out_uint32 = make_tensor<uint32_t>(tensor_float2.layout());
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size(out_uint32); ++i) {
|
||||
out_uint32(i) = convert_relu2<To_type>(tensor_float2(i));
|
||||
}
|
||||
Tensor out = make_tensor(make_rmem_ptr<To_type>(out_uint32.data()), tensor.layout());
|
||||
#else
|
||||
Tensor out = flash::convert_type<To_type>(tensor);
|
||||
flash::relu_(out);
|
||||
#endif
|
||||
return out;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Blocks until all but N previous cp.async.commit_group operations have committed.
|
||||
// This differs from cute::cp_async_wait in that when N = 0 we don't call cp.async.wait_all
|
||||
// (which is equivalent to commit_group then wait_group 0).
|
||||
// Instead we just call cp.async.wait_group 0, which is slightly faster.
|
||||
// https://github.com/NVIDIA/cutlass/blob/master/include/cute/arch/copy_sm80.hpp#L113
|
||||
template <int N>
|
||||
CUTE_HOST_DEVICE void cp_async_wait() {
|
||||
#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)
|
||||
asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
|
||||
#endif
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
bool Is_even_MN = true,
|
||||
bool Is_even_K = true,
|
||||
bool Clear_OOB_MN = false,
|
||||
bool Clear_OOB_K = true,
|
||||
typename TiledCopy,
|
||||
typename Engine0,
|
||||
typename Layout0,
|
||||
typename Engine1,
|
||||
typename Layout1,
|
||||
typename Engine2,
|
||||
typename Layout2,
|
||||
typename Engine3,
|
||||
typename Layout3>
|
||||
__forceinline__ __device__ void copy(
|
||||
TiledCopy tiled_copy,
|
||||
Tensor<Engine0, Layout0> const& S,
|
||||
Tensor<Engine1, Layout1>& D,
|
||||
Tensor<Engine2, Layout2> const& identity_MN,
|
||||
Tensor<Engine3, Layout3> const& predicate_K,
|
||||
const int max_MN = 0) {
|
||||
CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});
|
||||
CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA
|
||||
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K
|
||||
// There's no case where !Clear_OOB_K && Clear_OOB_MN
|
||||
static_assert(!(Clear_OOB_MN && !Clear_OOB_K));
|
||||
#pragma unroll
|
||||
for (int m = 0; m < size<1>(S); ++m) {
|
||||
if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < size<2>(S); ++k) {
|
||||
if (Is_even_K || predicate_K(k)) {
|
||||
cute::copy(tiled_copy, S(_, m, k), D(_, m, k));
|
||||
} else if (Clear_OOB_K) {
|
||||
cute::clear(D(_, m, k));
|
||||
}
|
||||
}
|
||||
} else if (Clear_OOB_MN) {
|
||||
cute::clear(D(_, m, _));
|
||||
}
|
||||
}
|
||||
// TD [2023-04-13]: Strange that the code below can cause race condition.
|
||||
// I think it's because the copies are under an if statement.
|
||||
// if (Is_even_K) {
|
||||
// #pragma unroll
|
||||
// for (int m = 0; m < size<1>(S); ++m) {
|
||||
// if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) {
|
||||
// copy(tiled_copy, S(_, m, _), D(_, m, _));
|
||||
// } else if (Clear_OOB_MN) {
|
||||
// clear(D(_, m, _));
|
||||
// }
|
||||
// }
|
||||
// } else { // It's slightly faster in this case if iterate over K first
|
||||
// #pragma unroll
|
||||
// for (int k = 0; k < size<2>(S); ++k) {
|
||||
// if (predicate_K(k)) {
|
||||
// #pragma unroll
|
||||
// for (int m = 0; m < size<1>(S); ++m) {
|
||||
// if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) {
|
||||
// copy(tiled_copy, S(_, m, k), D(_, m, k));
|
||||
// } else if (Clear_OOB_MN) {
|
||||
// clear(D(_, m, k));
|
||||
// }
|
||||
// }
|
||||
// } else if (Clear_OOB_K) { // There's no case where !Clear_OOB_K && Clear_OOB_MN
|
||||
// if (Clear_OOB_MN || Is_even_MN) {
|
||||
// clear(D(_, _, k));
|
||||
// } else {
|
||||
// #pragma unroll
|
||||
// for (int m = 0; m < size<1>(S); ++m) {
|
||||
// if (!(Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN)) {
|
||||
// clear(D(_, m, k));
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
bool Is_even_K = true,
|
||||
typename Engine0,
|
||||
typename Layout0,
|
||||
typename Engine1,
|
||||
typename Layout1,
|
||||
typename Engine2,
|
||||
typename Layout2,
|
||||
typename Engine3,
|
||||
typename Layout3>
|
||||
__forceinline__ __device__ void copy_w_min_idx(
|
||||
Tensor<Engine0, Layout0> const& S,
|
||||
Tensor<Engine1, Layout1>& D,
|
||||
Tensor<Engine2, Layout2> const& identity_MN,
|
||||
Tensor<Engine3, Layout3> const& predicate_K,
|
||||
const int max_MN = 0,
|
||||
const int min_MN = 0) {
|
||||
CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});
|
||||
CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});
|
||||
CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA
|
||||
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K
|
||||
// if (threadIdx.x == 0 && blockIdx.z == 0) { printf("blockIdx.y = %d, max_MN = %d, min_MN = %d\n", blockIdx.y, max_MN,
|
||||
// min_MN); }
|
||||
#pragma unroll
|
||||
for (int m = 0; m < size<1>(S); ++m) {
|
||||
// if (threadIdx.x == 0 && blockIdx.z == 0) { printf("blockIdx.y = %d, m = %d\n", blockIdx.y, get<0>(identity_MN(0,
|
||||
// m, 0))); }
|
||||
if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) {
|
||||
// if (threadIdx.x == 0 && blockIdx.z == 0) { printf("Inner loop, blockIdx.y = %d, m = %d\n", blockIdx.y,
|
||||
// get<0>(identity_MN(0, m, 0))); }
|
||||
#pragma unroll
|
||||
for (int k = 0; k < size<2>(S); ++k) {
|
||||
if (Is_even_K || predicate_K(k)) {
|
||||
cute::copy(S(_, m, k), D(_, m, k));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename Engine, typename Layout>
|
||||
__forceinline__ __device__ void apply_softcap(Tensor<Engine, Layout>& tensor, const float softcap) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size(tensor); ++i) {
|
||||
tensor(i) = cutlass::fast_tanh(tensor(i) * softcap);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Engine0, typename Layout0, typename Engine1, typename Layout1>
|
||||
__forceinline__ __device__ void
|
||||
calculate_dtanh(Tensor<Engine0, Layout0>& src_tensor, Tensor<Engine1, Layout1>& dst_tensor, const float softcap) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < size(src_tensor); ++i) {
|
||||
dst_tensor(i) = (1.f - (src_tensor(i) * src_tensor(i))) * softcap;
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace flash
|
||||
@@ -0,0 +1,57 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
|
||||
// Pybind entry for the InfLLM-V2 FlashAttention backend (vendored from
|
||||
// 3rdparty/infllmv2_cuda_impl). This builds as a standalone extension module
|
||||
// `infllm_ops` so its `flash::` symbols stay isolated from sgl-kernel's own
|
||||
// flash attention (`flash_ops` / `common_ops`).
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <c10/util/Optional.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
// Forward declarations of the FlashAttention entry points implemented in
|
||||
// flash_attn/flash_api.cpp. Signatures must match exactly.
|
||||
std::vector<at::Tensor> mha_varlen_fwd_stage1(
|
||||
at::Tensor& q,
|
||||
const at::Tensor& k,
|
||||
const at::Tensor& v,
|
||||
c10::optional<at::Tensor>& out_,
|
||||
const at::Tensor& cu_seqlens_q,
|
||||
const at::Tensor& cu_seqlens_k,
|
||||
const at::Tensor& cu_seqlens_v,
|
||||
c10::optional<at::Tensor>& seqused_k,
|
||||
c10::optional<const at::Tensor>& leftpad_k_,
|
||||
c10::optional<at::Tensor>& block_table_,
|
||||
c10::optional<at::Tensor>& alibi_slopes_,
|
||||
int max_seqlen_q,
|
||||
const int max_seqlen_k,
|
||||
const float p_dropout,
|
||||
const float softmax_scale,
|
||||
const bool zero_tensors,
|
||||
bool is_causal,
|
||||
int window_size_left,
|
||||
int window_size_right,
|
||||
const float softcap,
|
||||
const bool return_softmax,
|
||||
c10::optional<at::Generator> gen_);
|
||||
|
||||
PYBIND11_MODULE(infllm_ops, m) {
|
||||
m.doc() = "InfLLM V2 FlashAttention backend (vendored into sgl-kernel)";
|
||||
m.def("varlen_fwd_stage1", &mha_varlen_fwd_stage1, "Forward pass (variable length) NSA stage 1");
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
// InfLLM-V2 variable-length max pooling, AOT build.
|
||||
//
|
||||
// Migrated from `3rdparty/infllmv2_cuda_impl/csrc/max_pooling_1d.cuh`. The
|
||||
// device kernels are kept faithful to the original; only the host-side
|
||||
// launchers are rewritten from the raw `cudaStream_t` + `data_ptr` pybind
|
||||
// interface to the sgl-kernel `at::Tensor` + torch.ops convention.
|
||||
//
|
||||
// Notes vs. the original implementation:
|
||||
// * `TypeTraits<T>::inf()` is replaced by `static_cast<T>(INFINITY)`, and the
|
||||
// pooling max is accumulated in fp32 so we don't rely on half/bf16
|
||||
// comparison operators.
|
||||
// * Outputs are pre-allocated on the Python side and passed in (the original
|
||||
// wrappers also allocated a zero-filled output before launching).
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "utils.h"
|
||||
|
||||
namespace {
|
||||
|
||||
// input: [num_heads, total_q, max_seqlen_k]
|
||||
// output: [num_heads, total_q, out_len]
|
||||
template <typename T>
|
||||
__global__ void max_pooling_1d_varlen_kernel(
|
||||
const T* input,
|
||||
T* output,
|
||||
const int* cu_seqlens_q,
|
||||
const int* cu_seqlens_k,
|
||||
const int* cache_lens,
|
||||
int batch_size,
|
||||
int num_heads,
|
||||
int max_seqlen_k,
|
||||
int out_len,
|
||||
int kernel_size,
|
||||
int stride,
|
||||
int padding,
|
||||
int block_size,
|
||||
int local_blocks,
|
||||
int init_blocks) {
|
||||
const int bidh = blockIdx.y; // head index
|
||||
const int bidq_global = blockIdx.x; // global query index across all batches
|
||||
|
||||
int batch_idx = 0;
|
||||
int q_start = 0, q_end = 0, k_start = 0, k_end = 0;
|
||||
for (int b = 0; b < batch_size; b++) {
|
||||
q_start = cu_seqlens_q[b];
|
||||
q_end = cu_seqlens_q[b + 1];
|
||||
k_start = cu_seqlens_k[b];
|
||||
k_end = cu_seqlens_k[b + 1];
|
||||
if (bidq_global >= q_start && bidq_global < q_end) {
|
||||
batch_idx = b;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const int bidq_local = bidq_global - q_start;
|
||||
const int seqlen_q = q_end - q_start;
|
||||
const int seqlen_k = k_end - k_start;
|
||||
if (bidq_local >= seqlen_q) return;
|
||||
|
||||
const size_t total_q_all = static_cast<size_t>(cu_seqlens_q[batch_size]);
|
||||
const size_t in_offset =
|
||||
static_cast<size_t>(bidh) * total_q_all * max_seqlen_k + static_cast<size_t>(bidq_global) * max_seqlen_k;
|
||||
const T* in = input + in_offset;
|
||||
const size_t out_offset =
|
||||
static_cast<size_t>(bidh) * total_q_all * out_len + static_cast<size_t>(bidq_global) * out_len;
|
||||
T* out = output + out_offset;
|
||||
|
||||
const int cache_len = cache_lens[batch_idx];
|
||||
const int off_bq = (bidq_local + cache_len) / block_size;
|
||||
const T pos_inf = static_cast<T>(static_cast<float>(INFINITY));
|
||||
|
||||
for (int k = threadIdx.x; k < out_len; k += blockDim.x) {
|
||||
const int off_bk = k;
|
||||
const bool should_mask_inf = (off_bk < init_blocks) || ((off_bq >= off_bk) && (off_bq <= off_bk + local_blocks));
|
||||
|
||||
if (should_mask_inf) {
|
||||
out[k] = pos_inf;
|
||||
} else {
|
||||
int start = k * stride - padding;
|
||||
int end = start + kernel_size;
|
||||
start = max(start, 0);
|
||||
end = min(end, seqlen_k);
|
||||
|
||||
float max_val = -INFINITY;
|
||||
for (int i = start; i < end; i++) {
|
||||
const float v = static_cast<float>(in[i]);
|
||||
if (v > max_val) max_val = v;
|
||||
}
|
||||
out[k] = static_cast<T>(max_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
void infllm_v2_max_pooling_1d_varlen(
|
||||
at::Tensor input,
|
||||
at::Tensor output,
|
||||
at::Tensor cu_seqlens_q,
|
||||
at::Tensor cu_seqlens_k,
|
||||
at::Tensor cache_lens,
|
||||
int64_t max_seqlen_q,
|
||||
int64_t max_seqlen_k,
|
||||
int64_t kernel_size,
|
||||
int64_t stride,
|
||||
int64_t padding,
|
||||
int64_t block_size,
|
||||
int64_t local_blocks,
|
||||
int64_t init_blocks,
|
||||
int64_t total_q) {
|
||||
TORCH_CHECK(input.dim() == 3, "input must be 3D [num_heads, total_q, max_k]");
|
||||
TORCH_CHECK(output.dim() == 3, "output must be 3D [num_heads, total_q, out_len]");
|
||||
TORCH_CHECK(cu_seqlens_q.scalar_type() == at::kInt, "cu_seqlens_q must be int32");
|
||||
TORCH_CHECK(cu_seqlens_k.scalar_type() == at::kInt, "cu_seqlens_k must be int32");
|
||||
TORCH_CHECK(cache_lens.scalar_type() == at::kInt, "cache_lens must be int32");
|
||||
|
||||
const int batch_size = static_cast<int>(cu_seqlens_q.size(0)) - 1;
|
||||
const int num_heads = static_cast<int>(input.size(0));
|
||||
const int out_len = static_cast<int>(output.size(2));
|
||||
const int grid_q = static_cast<int>(total_q > 0 ? total_q : input.size(1));
|
||||
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
const dim3 grid(grid_q, num_heads);
|
||||
const dim3 block(256);
|
||||
|
||||
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] {
|
||||
max_pooling_1d_varlen_kernel<c_type><<<grid, block, 0, stream>>>(
|
||||
static_cast<const c_type*>(input.data_ptr()),
|
||||
static_cast<c_type*>(output.data_ptr()),
|
||||
cu_seqlens_q.data_ptr<int>(),
|
||||
cu_seqlens_k.data_ptr<int>(),
|
||||
cache_lens.data_ptr<int>(),
|
||||
batch_size,
|
||||
num_heads,
|
||||
static_cast<int>(max_seqlen_k),
|
||||
out_len,
|
||||
static_cast<int>(kernel_size),
|
||||
static_cast<int>(stride),
|
||||
static_cast<int>(padding),
|
||||
static_cast<int>(block_size),
|
||||
static_cast<int>(local_blocks),
|
||||
static_cast<int>(init_blocks));
|
||||
return true;
|
||||
});
|
||||
}
|
||||
@@ -112,6 +112,25 @@ int64_t cutlass_mla_get_workspace_size(
|
||||
int64_t sm_count = 0,
|
||||
int64_t num_kv_splits = 1 /* Set to 1 to avoid cuda_graph issue by default. */);
|
||||
|
||||
/*
|
||||
* From csrc/infllm_v2
|
||||
*/
|
||||
void infllm_v2_max_pooling_1d_varlen(
|
||||
at::Tensor input,
|
||||
at::Tensor output,
|
||||
at::Tensor cu_seqlens_q,
|
||||
at::Tensor cu_seqlens_k,
|
||||
at::Tensor cache_lens,
|
||||
int64_t max_seqlen_q,
|
||||
int64_t max_seqlen_k,
|
||||
int64_t kernel_size,
|
||||
int64_t stride,
|
||||
int64_t padding,
|
||||
int64_t block_size,
|
||||
int64_t local_blocks,
|
||||
int64_t init_blocks,
|
||||
int64_t total_q);
|
||||
|
||||
/*
|
||||
* From csrc/elementwise
|
||||
*/
|
||||
|
||||
@@ -71,6 +71,10 @@ else:
|
||||
shuffle_rows,
|
||||
)
|
||||
from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda
|
||||
from sgl_kernel.infllm_v2 import (
|
||||
infllmv2_attn_stage1,
|
||||
max_pooling_1d_varlen,
|
||||
)
|
||||
from sgl_kernel.kvcacheio import (
|
||||
transfer_kv_all_layer,
|
||||
transfer_kv_all_layer_mla,
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
from sgl_kernel.infllm_v2.attention import infllmv2_attn_stage1
|
||||
from sgl_kernel.infllm_v2.max_pooling import max_pooling_1d_varlen
|
||||
|
||||
__all__ = [
|
||||
"infllmv2_attn_stage1",
|
||||
"max_pooling_1d_varlen",
|
||||
]
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Robust loader for the standalone ``infllm_ops`` pybind extension.
|
||||
|
||||
The InfLLM-V2 FlashAttention backend is built as its own module ``infllm_ops``
|
||||
(installed into the ``sgl_kernel`` package directory). Under editable installs
|
||||
the compiled ``.so`` may live in ``site-packages/sgl_kernel`` while the imported
|
||||
``sgl_kernel`` package resolves to the source tree, so a plain ``from sgl_kernel
|
||||
import infllm_ops`` is not always sufficient. This loader searches the known
|
||||
candidate locations and loads the extension by file path.
|
||||
"""
|
||||
|
||||
import glob
|
||||
import importlib.util
|
||||
import site
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
_infllm_ops = None
|
||||
|
||||
|
||||
def _candidate_dirs() -> List[Path]:
|
||||
dirs: List[Path] = []
|
||||
|
||||
# 1) The directory of the sgl_kernel package as currently imported.
|
||||
try:
|
||||
import sgl_kernel
|
||||
|
||||
dirs.append(Path(sgl_kernel.__file__).parent)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 2) This module's parent package directory (source tree).
|
||||
dirs.append(Path(__file__).resolve().parent.parent)
|
||||
|
||||
# 3) Every ``sgl_kernel`` directory found on the install paths.
|
||||
search_roots: List[str] = []
|
||||
try:
|
||||
search_roots.extend(site.getsitepackages())
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
search_roots.append(site.getusersitepackages())
|
||||
except Exception:
|
||||
pass
|
||||
search_roots.extend(p for p in sys.path if p)
|
||||
for root in search_roots:
|
||||
dirs.append(Path(root) / "sgl_kernel")
|
||||
|
||||
# De-duplicate while preserving order.
|
||||
seen = set()
|
||||
unique: List[Path] = []
|
||||
for d in dirs:
|
||||
key = str(d)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique.append(d)
|
||||
return unique
|
||||
|
||||
|
||||
def _find_so() -> Optional[Path]:
|
||||
for d in _candidate_dirs():
|
||||
if not d.is_dir():
|
||||
continue
|
||||
matches = sorted(glob.glob(str(d / "infllm_ops*.so")))
|
||||
if matches:
|
||||
return Path(matches[0])
|
||||
return None
|
||||
|
||||
|
||||
def load_infllm_ops():
|
||||
"""Import and return the ``infllm_ops`` extension module (cached)."""
|
||||
global _infllm_ops
|
||||
if _infllm_ops is not None:
|
||||
return _infllm_ops
|
||||
|
||||
# Fast path: a normal import may already work.
|
||||
try:
|
||||
from sgl_kernel import infllm_ops as _mod # type: ignore
|
||||
|
||||
_infllm_ops = _mod
|
||||
return _infllm_ops
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
so_path = _find_so()
|
||||
if so_path is None:
|
||||
raise ImportError(
|
||||
"[sgl_kernel] Could not locate the 'infllm_ops' extension (infllm_ops*.so). "
|
||||
"Ensure sgl-kernel was built with the InfLLM-V2 FlashAttention backend."
|
||||
)
|
||||
|
||||
spec = importlib.util.spec_from_file_location("infllm_ops", str(so_path))
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError(f"[sgl_kernel] Could not create module spec for {so_path}")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
_infllm_ops = module
|
||||
return _infllm_ops
|
||||
@@ -0,0 +1,82 @@
|
||||
"""InfLLM-V2 sparse FlashAttention public API.
|
||||
|
||||
Ported (drop-in) from ``3rdparty/infllmv2_cuda_impl/infllm_v2/infllmv2_sparse_attention.py``.
|
||||
The CUDA backend now lives in the standalone ``infllm_ops`` extension.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from sgl_kernel.infllm_v2._loader import load_infllm_ops
|
||||
|
||||
|
||||
def maybe_contiguous(x):
|
||||
return x.contiguous() if x is not None and x.stride(-1) != 1 else x
|
||||
|
||||
|
||||
def infllmv2_attn_stage1(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
cu_seqlens_v,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
dropout_p=0.0,
|
||||
softmax_scale=None,
|
||||
causal=False,
|
||||
window_size=(-1, -1),
|
||||
softcap=0.0,
|
||||
alibi_slopes=None,
|
||||
deterministic=False,
|
||||
return_attn_probs=True,
|
||||
block_table=None,
|
||||
):
|
||||
"""Neighborhood Sparse Attention (NSA) Stage 1 with varlen support.
|
||||
|
||||
Drop-in replacement for ``infllm_v2.infllmv2_attn_stage1``. Returns the
|
||||
attention-score matrix with the NSA sparsity pattern, shape
|
||||
``(num_heads_k, total_q, max_seqlen_k)``.
|
||||
"""
|
||||
infllm_ops = load_infllm_ops()
|
||||
if softmax_scale is None:
|
||||
softmax_scale = q.shape[-1] ** (-0.5)
|
||||
|
||||
q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
|
||||
|
||||
total_q, nheads, head_dim = q.shape
|
||||
nheads_k = k.shape[1]
|
||||
nheads_per_group = nheads // nheads_k
|
||||
|
||||
q = q.reshape(total_q, nheads_k, nheads_per_group, head_dim)
|
||||
q = (
|
||||
q.transpose(1, 2)
|
||||
.reshape(total_q * nheads_per_group, nheads_k, head_dim)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
result = infllm_ops.varlen_fwd_stage1(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
None,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
cu_seqlens_v,
|
||||
None,
|
||||
None,
|
||||
block_table,
|
||||
alibi_slopes,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
dropout_p,
|
||||
softmax_scale,
|
||||
True,
|
||||
causal,
|
||||
window_size[0],
|
||||
window_size[1],
|
||||
softcap,
|
||||
True,
|
||||
None,
|
||||
)
|
||||
|
||||
return result[0]
|
||||
@@ -0,0 +1,61 @@
|
||||
import torch
|
||||
|
||||
|
||||
def max_pooling_1d_varlen(
|
||||
input: torch.Tensor, # num_heads x total_q x max_k
|
||||
cu_seqlens_q: torch.Tensor, # batch_size + 1
|
||||
cu_seqlens_k: torch.Tensor, # batch_size + 1
|
||||
cache_lens: torch.Tensor, # batch_size
|
||||
max_seqlen_q: int,
|
||||
max_context_len: int,
|
||||
local_blocks: int,
|
||||
init_blocks: int,
|
||||
block_size: int = 64,
|
||||
stride: int = 16,
|
||||
total_q: int = -1,
|
||||
) -> torch.Tensor:
|
||||
"""Variable-length 1D max pooling over packed sequences.
|
||||
|
||||
Drop-in replacement for ``infllm_v2.max_pooling_1d_varlen``.
|
||||
"""
|
||||
assert input.dtype in (torch.float16, torch.bfloat16)
|
||||
assert cu_seqlens_q.dtype == torch.int32
|
||||
assert cu_seqlens_k.dtype == torch.int32
|
||||
assert cache_lens.dtype == torch.int32
|
||||
assert input.dim() == 3, f"Expected 3D input, got {input.dim()}D"
|
||||
|
||||
input = input.contiguous()
|
||||
cu_seqlens_q = cu_seqlens_q.contiguous()
|
||||
cu_seqlens_k = cu_seqlens_k.contiguous()
|
||||
cache_lens = cache_lens.contiguous()
|
||||
|
||||
max_seqlen_k = max_context_len // stride
|
||||
out_len = (max_context_len + block_size - 1) // block_size
|
||||
|
||||
stride = block_size // stride
|
||||
kernel_size = stride + 1
|
||||
padding = 1
|
||||
|
||||
num_heads = input.shape[0]
|
||||
total_q = input.shape[1]
|
||||
|
||||
output = torch.zeros(
|
||||
num_heads, total_q, out_len, device=input.device, dtype=input.dtype
|
||||
)
|
||||
torch.ops.sgl_kernel.infllm_v2_max_pooling_1d_varlen.default(
|
||||
input,
|
||||
output,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
cache_lens,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
block_size,
|
||||
local_blocks,
|
||||
init_blocks,
|
||||
total_q,
|
||||
)
|
||||
return output
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Equivalence tests for the migrated InfLLM-V2 FlashAttention API.
|
||||
|
||||
These compare the ``sgl_kernel.infllm_v2`` implementations against the original
|
||||
``infllm_v2`` package (3rdparty/infllmv2_cuda_impl). Both call the same CUDA
|
||||
kernels, so outputs are expected to match closely. The whole module is skipped
|
||||
if the reference ``infllm_v2`` package is not importable.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
sgl = pytest.importorskip("sgl_kernel.infllm_v2")
|
||||
ref = pytest.importorskip("infllm_v2")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="CUDA is required for InfLLM-V2 kernels"
|
||||
)
|
||||
|
||||
|
||||
def _assert_close(a, b, name):
|
||||
a = a.float()
|
||||
b = b.float()
|
||||
assert a.shape == b.shape, f"{name}: shape mismatch {a.shape} vs {b.shape}"
|
||||
max_diff = (a - b).abs().max().item()
|
||||
assert torch.allclose(a, b, atol=1e-2, rtol=1e-2), f"{name}: max diff {max_diff}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_dim", [64, 128])
|
||||
@pytest.mark.parametrize("causal", [False, True])
|
||||
@pytest.mark.parametrize("seqlen_q,seqlen_k", [(256, 16), (64, 17)])
|
||||
def test_stage1_matches_reference(head_dim, causal, seqlen_q, seqlen_k):
|
||||
torch.manual_seed(0)
|
||||
n_heads, n_kv_heads = 32, 2
|
||||
dtype = torch.bfloat16
|
||||
|
||||
q = torch.randn(n_heads, seqlen_q, head_dim, dtype=dtype, device="cuda")
|
||||
k = torch.randn(n_kv_heads, seqlen_k, head_dim, dtype=dtype, device="cuda")
|
||||
|
||||
cu_seqlens_q = torch.tensor([0, seqlen_q], dtype=torch.int32, device="cuda")
|
||||
cu_seqlens_k = torch.tensor([0, seqlen_k], dtype=torch.int32, device="cuda")
|
||||
|
||||
q = q.transpose(0, 1).contiguous()
|
||||
k = k.transpose(0, 1).contiguous()
|
||||
|
||||
common = dict(
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_k=cu_seqlens_k,
|
||||
cu_seqlens_v=cu_seqlens_k,
|
||||
max_seqlen_q=seqlen_q,
|
||||
max_seqlen_k=seqlen_k,
|
||||
causal=causal,
|
||||
)
|
||||
out_ref = ref.infllmv2_attn_stage1(q, k, k, **common)
|
||||
out_sgl = sgl.infllmv2_attn_stage1(q, k, k, **common)
|
||||
_assert_close(out_sgl, out_ref, "stage1")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,109 @@
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel import max_pooling_1d_varlen
|
||||
|
||||
|
||||
def _ref_varlen(
|
||||
score: torch.Tensor, # [num_heads, total_q, max_k]
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
cache_lens: torch.Tensor,
|
||||
max_context_len: int,
|
||||
local_blocks: int,
|
||||
init_blocks: int,
|
||||
block_size: int,
|
||||
kernel_stride: int,
|
||||
) -> torch.Tensor:
|
||||
"""Pure-torch reference mirroring the CUDA kernel exactly (fp32 math)."""
|
||||
num_heads, total_q, _ = score.shape
|
||||
out_len = (max_context_len + block_size - 1) // block_size
|
||||
stride = block_size // kernel_stride
|
||||
kernel_size = stride + 1
|
||||
padding = 1
|
||||
|
||||
cu_q = cu_seqlens_q.tolist()
|
||||
cu_k = cu_seqlens_k.tolist()
|
||||
cache = cache_lens.tolist()
|
||||
batch_size = len(cache)
|
||||
|
||||
out = torch.zeros(num_heads, total_q, out_len, dtype=torch.float32)
|
||||
s = score.float().cpu()
|
||||
for q in range(total_q):
|
||||
b = 0
|
||||
for bb in range(batch_size):
|
||||
if cu_q[bb] <= q < cu_q[bb + 1]:
|
||||
b = bb
|
||||
break
|
||||
bidq_local = q - cu_q[b]
|
||||
seqlen_k = cu_k[b + 1] - cu_k[b]
|
||||
off_bq = (bidq_local + cache[b]) // block_size
|
||||
for h in range(num_heads):
|
||||
for k in range(out_len):
|
||||
if (k < init_blocks) or (off_bq >= k and off_bq <= k + local_blocks):
|
||||
out[h, q, k] = float("inf")
|
||||
else:
|
||||
start = max(k * stride - padding, 0)
|
||||
end = min(start + kernel_size, seqlen_k)
|
||||
if end > start:
|
||||
out[h, q, k] = s[h, q, start:end].max()
|
||||
else:
|
||||
out[h, q, k] = float("-inf")
|
||||
return out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("num_heads", [1, 4])
|
||||
@pytest.mark.parametrize("seq_lens", [[37], [16, 48], [8, 8, 24]])
|
||||
def test_max_pooling_varlen_matches_reference(dtype, num_heads, seq_lens):
|
||||
torch.manual_seed(0)
|
||||
block_size = 64
|
||||
kernel_stride = 16
|
||||
local_blocks = 1
|
||||
init_blocks = 1
|
||||
max_context_len = 512
|
||||
|
||||
total_q = sum(seq_lens)
|
||||
max_k = max_context_len // kernel_stride
|
||||
cu = [0]
|
||||
for n in seq_lens:
|
||||
cu.append(cu[-1] + n)
|
||||
cu_seqlens_q = torch.tensor(cu, dtype=torch.int32, device="cuda")
|
||||
cu_seqlens_k = torch.tensor(cu, dtype=torch.int32, device="cuda")
|
||||
cache_lens = torch.zeros(len(seq_lens), dtype=torch.int32, device="cuda")
|
||||
|
||||
score = torch.randn(num_heads, total_q, max_k, dtype=dtype, device="cuda")
|
||||
|
||||
out = max_pooling_1d_varlen(
|
||||
score,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
cache_lens,
|
||||
max_seqlen_q=max(seq_lens),
|
||||
max_context_len=max_context_len,
|
||||
local_blocks=local_blocks,
|
||||
init_blocks=init_blocks,
|
||||
block_size=block_size,
|
||||
stride=kernel_stride,
|
||||
total_q=total_q,
|
||||
)
|
||||
ref = _ref_varlen(
|
||||
score,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
cache_lens,
|
||||
max_context_len,
|
||||
local_blocks,
|
||||
init_blocks,
|
||||
block_size,
|
||||
kernel_stride,
|
||||
).to(out.device)
|
||||
|
||||
assert torch.equal(torch.isinf(out) & (out > 0), torch.isinf(ref) & (ref > 0))
|
||||
finite = torch.isfinite(ref)
|
||||
torch.testing.assert_close(out[finite].float(), ref[finite], rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
Reference in New Issue
Block a user