Skip to content

CUDA: faster FlashAttention for batch sizes > 1 #6646

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Apr 18, 2024
Merged

CUDA: faster FlashAttention for batch sizes > 1 #6646

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a83d993
CUDA: refactor host code, dyn. par. blocks
JohannesGaessler Apr 9, 2024
359d0f5
fix flash_attn_vec_f16 race condition
JohannesGaessler Apr 13, 2024
049533d
flush softmax exp below threshold to 0
JohannesGaessler Apr 15, 2024
aef96ff
store temp KQ in registers
JohannesGaessler Apr 16, 2024
a9d6591
Calculate KQ as FP32 if KQV has GGML_PREC_F32
JohannesGaessler Apr 16, 2024
4e4d58a
Add __hgt2_mask implementation for CUDA 11
JohannesGaessler Apr 17, 2024
44ca576
fix KQ FP32 precision fpr parallel_blocks > 1
JohannesGaessler Apr 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,7 @@ static ggml_cuda_device_info ggml_cuda_init() {
info.devices[id].cc = 100*prop.major + 10*prop.minor;
#endif // defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__)
info.devices[id].smpb = prop.sharedMemPerBlock;
info.devices[id].nsm = prop.multiProcessorCount;
}

for (int id = 0; id < info.device_count; ++id) {
Expand Down
13 changes: 13 additions & 0 deletions ggml-cuda/common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,13 @@ static __device__ __forceinline__ half2 warp_reduce_max(half2 x) {
#endif // !(defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__)) && __CUDA_ARCH__ >= CC_PASCAL && CUDART_VERSION >= CUDART_HMAX
}

#if CUDART_VERSION < 12000
static __device__ __forceinline__ uint __hgt2_mask(const half2 a, const half2 b) {
const uint mask_low = 0x0000FFFF * ( __low2half(a) > __low2half(b));
const uint mask_high = 0xFFFF0000 * (__high2half(a) > __high2half(b));
return mask_low | mask_high;
}
#endif // CUDART_VERSION < 12000

#if defined(GGML_USE_HIPBLAS)
#define __CUDA_ARCH__ 1300
Expand Down Expand Up @@ -390,6 +397,11 @@ static __device__ __forceinline__ int __dp4a(const int a, const int b, int c) {
}
#endif // defined(GGML_USE_HIPBLAS)

#define FP16_AVAILABLE defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__) ? \
defined(RDNA1) || defined(RDNA2) || defined(RDNA3) : __CUDA_ARCH__ >= CC_PASCAL
#define FP16_MMA_AVAILABLE defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__) ? \
defined(RDNA3) : __CUDA_ARCH__ >= CC_VOLTA

// TODO: move to ggml-common.h
static const __device__ int8_t kvalues_iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113};

Expand All @@ -403,6 +415,7 @@ struct ggml_cuda_device_info {

struct cuda_device_info {
int cc; // compute capability
int nsm; // number of streaming multiprocessors
size_t smpb; // max. shared memory per block
bool vmm; // virtual memory support
size_t vmm_granularity; // granularity of virtual memory
Expand Down
Loading
Loading