381 lines
11 KiB
Plaintext
381 lines
11 KiB
Plaintext
/*
|
|
* Flash Attention CUDA Utilities
|
|
*
|
|
* RTX 5090 (sm_89/sm_90) optimized utility functions and kernels
|
|
* for maximum performance and memory efficiency.
|
|
*/
|
|
|
|
#include <cuda_runtime.h>
|
|
#include <cuda_fp16.h>
|
|
#include <cuda_bf16.h>
|
|
#include <mma.h>
|
|
#include <cooperative_groups.h>
|
|
#include <cooperative_groups/memcpy_async.h>
|
|
#include <cuda/atomic>
|
|
|
|
// RTX 5090 specific optimizations
|
|
#if __CUDA_ARCH__ >= 890
|
|
#define RTX_5090_OPTIMIZATION_ENABLED
|
|
#define USE_TENSOR_CORE_4TH_GEN
|
|
#define USE_ASYNC_COPY
|
|
#define MAX_SHARED_MEMORY_PER_BLOCK 163840 // 160KB for RTX 5090
|
|
#endif
|
|
|
|
// Constants optimized for RTX 5090
|
|
#define WARP_SIZE 32
|
|
#define MAX_THREADS_PER_BLOCK 1024
|
|
#define SHARED_MEM_ALIGNMENT 16
|
|
#define CACHE_LINE_SIZE 128
|
|
#define L2_CACHE_SIZE (96 * 1024 * 1024) // 96MB L2 cache on RTX 5090
|
|
|
|
using namespace nvcuda;
|
|
namespace cg = cooperative_groups;
|
|
|
|
// Advanced data types for RTX 5090
|
|
#ifdef USE_TENSOR_CORE_4TH_GEN
|
|
using precision_t = __nv_bfloat16;
|
|
using precision2_t = __nv_bfloat162;
|
|
using accumulator_t = float;
|
|
#else
|
|
using precision_t = half;
|
|
using precision2_t = half2;
|
|
using accumulator_t = float;
|
|
#endif
|
|
|
|
// Memory access patterns optimized for RTX 5090
|
|
template<int N>
|
|
struct alignas(16) VectorLoad {
|
|
precision_t data[N];
|
|
|
|
__device__ VectorLoad() {}
|
|
|
|
__device__ void load(const precision_t* ptr) {
|
|
#pragma unroll
|
|
for (int i = 0; i < N; ++i) {
|
|
data[i] = ptr[i];
|
|
}
|
|
}
|
|
|
|
__device__ void store(precision_t* ptr) const {
|
|
#pragma unroll
|
|
for (int i = 0; i < N; ++i) {
|
|
ptr[i] = data[i];
|
|
}
|
|
}
|
|
};
|
|
|
|
using float4_t = VectorLoad<4>;
|
|
using float8_t = VectorLoad<8>;
|
|
using float16_t = VectorLoad<16>;
|
|
|
|
// RTX 5090 optimized async memory copy
|
|
template<typename T>
|
|
__device__ void async_copy_shared_global(T* shared_ptr, const T* global_ptr, size_t count) {
|
|
#ifdef USE_ASYNC_COPY
|
|
auto block = cg::this_thread_block();
|
|
cg::memcpy_async(block, shared_ptr, global_ptr, sizeof(T) * count);
|
|
#else
|
|
// Fallback for older architectures
|
|
for (int i = threadIdx.x; i < count; i += blockDim.x) {
|
|
shared_ptr[i] = global_ptr[i];
|
|
}
|
|
#endif
|
|
}
|
|
|
|
// Advanced warp-level reduction with shuffle optimization
|
|
template<typename T>
|
|
__device__ T warp_reduce_sum_advanced(T val) {
|
|
// RTX 5090 has improved shuffle performance
|
|
#pragma unroll
|
|
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
|
T other = __shfl_down_sync(0xffffffff, val, offset);
|
|
val += other;
|
|
}
|
|
return val;
|
|
}
|
|
|
|
// Block-level reduction optimized for RTX 5090
|
|
template<typename T>
|
|
__device__ T block_reduce_sum_optimized(T val) {
|
|
__shared__ T shared[32]; // Max warps per block
|
|
|
|
int warp_id = threadIdx.x / WARP_SIZE;
|
|
int lane_id = threadIdx.x % WARP_SIZE;
|
|
|
|
// Warp-level reduce with advanced shuffle
|
|
val = warp_reduce_sum_advanced(val);
|
|
|
|
// Store warp result
|
|
if (lane_id == 0) {
|
|
shared[warp_id] = val;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// Final reduction using first warp
|
|
if (warp_id == 0) {
|
|
val = (lane_id < (blockDim.x + WARP_SIZE - 1) / WARP_SIZE) ? shared[lane_id] : T(0);
|
|
val = warp_reduce_sum_advanced(val);
|
|
|
|
// Broadcast result
|
|
shared[0] = val;
|
|
}
|
|
|
|
__syncthreads();
|
|
return shared[0];
|
|
}
|
|
|
|
// Tensor Core WMMA operations for RTX 5090
|
|
#ifdef USE_TENSOR_CORE_4TH_GEN
|
|
__device__ void tensor_core_mma_bf16(
|
|
const precision_t* a, const precision_t* b, accumulator_t* c,
|
|
int m, int n, int k, int lda, int ldb, int ldc
|
|
) {
|
|
// Use 4th generation Tensor Cores with BF16
|
|
wmma::fragment<wmma::matrix_a, 16, 16, 16, precision_t, wmma::row_major> a_frag;
|
|
wmma::fragment<wmma::matrix_b, 16, 16, 16, precision_t, wmma::col_major> b_frag;
|
|
wmma::fragment<wmma::accumulator, 16, 16, 16, accumulator_t> c_frag;
|
|
|
|
wmma::fill_fragment(c_frag, 0.0f);
|
|
wmma::load_matrix_sync(a_frag, a, lda);
|
|
wmma::load_matrix_sync(b_frag, b, ldb);
|
|
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
|
|
wmma::store_matrix_sync(c, c_frag, ldc, wmma::mem_row_major);
|
|
}
|
|
#endif
|
|
|
|
// Memory bandwidth optimization for attention computation
|
|
__device__ void prefetch_data(const void* ptr, size_t bytes) {
|
|
#if __CUDA_ARCH__ >= 890
|
|
// RTX 5090 has advanced prefetch capabilities
|
|
// Use CUDA's prefetch instead of __builtin_prefetch
|
|
if (ptr != nullptr) {
|
|
// Manual prefetch by touching memory
|
|
volatile char dummy = *((const volatile char*)ptr);
|
|
(void)dummy;
|
|
}
|
|
#endif
|
|
}
|
|
|
|
// Optimized softmax computation using RTX 5090 features
|
|
__device__ void fast_softmax_inplace(float* logits, int len) {
|
|
// Find maximum for numerical stability
|
|
float max_val = -INFINITY;
|
|
#pragma unroll 8
|
|
for (int i = 0; i < len; i++) {
|
|
max_val = fmaxf(max_val, logits[i]);
|
|
}
|
|
|
|
// Compute exponentials and sum
|
|
float sum = 0.0f;
|
|
#pragma unroll 8
|
|
for (int i = 0; i < len; i++) {
|
|
float exp_val = expf(logits[i] - max_val);
|
|
logits[i] = exp_val;
|
|
sum += exp_val;
|
|
}
|
|
|
|
// Normalize
|
|
float inv_sum = 1.0f / sum;
|
|
#pragma unroll 8
|
|
for (int i = 0; i < len; i++) {
|
|
logits[i] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
// Cache-aware data layout transformation
|
|
template<int TILE_SIZE>
|
|
__device__ void transpose_tile_shared(
|
|
const precision_t* src, precision_t* dst,
|
|
int src_stride, int dst_stride,
|
|
int rows, int cols
|
|
) {
|
|
__shared__ precision_t tile[TILE_SIZE][TILE_SIZE + 1]; // +1 to avoid bank conflicts
|
|
|
|
int tx = threadIdx.x;
|
|
int ty = threadIdx.y;
|
|
|
|
// Load tile from source
|
|
if (tx < cols && ty < rows) {
|
|
tile[ty][tx] = src[ty * src_stride + tx];
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// Store transposed tile to destination
|
|
if (tx < rows && ty < cols) {
|
|
dst[ty * dst_stride + tx] = tile[tx][ty];
|
|
}
|
|
}
|
|
|
|
// RTX 5090 specific memory access patterns
|
|
__device__ void coalesced_load_fp16(
|
|
const half* src, half* dst, int count
|
|
) {
|
|
// Use 128-bit loads for maximum bandwidth
|
|
const int vec_size = 8; // 8 half values = 128 bits
|
|
const int vec_count = count / vec_size;
|
|
|
|
using vec_t = float4; // Represents 8 half values
|
|
|
|
const vec_t* src_vec = reinterpret_cast<const vec_t*>(src);
|
|
vec_t* dst_vec = reinterpret_cast<vec_t*>(dst);
|
|
|
|
for (int i = threadIdx.x; i < vec_count; i += blockDim.x) {
|
|
dst_vec[i] = src_vec[i];
|
|
}
|
|
|
|
// Handle remaining elements
|
|
int remaining = count - vec_count * vec_size;
|
|
if (threadIdx.x < remaining) {
|
|
int idx = vec_count * vec_size + threadIdx.x;
|
|
dst[idx] = src[idx];
|
|
}
|
|
}
|
|
|
|
// Occupancy optimization calculator
|
|
__device__ int calculate_optimal_block_size(int seq_len, int head_dim) {
|
|
// RTX 5090 has 128 SMs, optimize for high occupancy
|
|
int max_threads = 2048; // Max threads per SM
|
|
|
|
// Calculate optimal block size based on problem size
|
|
int optimal_threads = min(1024, (seq_len * head_dim + 31) / 32 * 32);
|
|
return optimal_threads;
|
|
}
|
|
|
|
// Advanced attention pattern detection for optimization
|
|
__device__ bool is_attention_pattern_sparse(
|
|
const float* attention_weights,
|
|
int seq_len,
|
|
float sparsity_threshold = 0.1f
|
|
) {
|
|
int non_zero_count = 0;
|
|
|
|
#pragma unroll 4
|
|
for (int i = 0; i < seq_len; i++) {
|
|
if (attention_weights[i] > sparsity_threshold) {
|
|
non_zero_count++;
|
|
}
|
|
}
|
|
|
|
float sparsity = 1.0f - (float)non_zero_count / seq_len;
|
|
return sparsity > 0.8f; // 80% sparsity threshold
|
|
}
|
|
|
|
// Memory-efficient gradient accumulation
|
|
__device__ void atomic_add_half(half* address, half val) {
|
|
#if __CUDA_ARCH__ >= 700
|
|
atomicAdd(address, val);
|
|
#else
|
|
// Fallback for older architectures
|
|
unsigned int* base_address = (unsigned int*)((size_t)address & ~3);
|
|
unsigned int old = *base_address;
|
|
unsigned int assumed;
|
|
|
|
do {
|
|
assumed = old;
|
|
half* h_ptr = (half*)&old + ((size_t)address & 3) / sizeof(half);
|
|
*h_ptr = __hadd(*h_ptr, val);
|
|
old = atomicCAS(base_address, assumed, old);
|
|
} while (assumed != old);
|
|
#endif
|
|
}
|
|
|
|
// RTX 5090 L2 cache optimization
|
|
__device__ void optimize_l2_access_pattern(
|
|
const void* data_ptr, size_t data_size
|
|
) {
|
|
#if __CUDA_ARCH__ >= 890
|
|
// Hint to keep frequently accessed data in L2 cache
|
|
if (data_size < L2_CACHE_SIZE / 4 && data_ptr != nullptr) {
|
|
// Manual prefetch by touching memory
|
|
volatile char dummy = *((const volatile char*)data_ptr);
|
|
(void)dummy;
|
|
}
|
|
#endif
|
|
}
|
|
|
|
// Performance counters for profiling (RTX 5090 specific)
|
|
struct PerformanceCounters {
|
|
unsigned long long clock_start;
|
|
unsigned long long clock_end;
|
|
unsigned int active_warps;
|
|
unsigned int memory_transactions;
|
|
|
|
__device__ void start() {
|
|
clock_start = clock64();
|
|
active_warps = __ballot_sync(0xffffffff, true);
|
|
}
|
|
|
|
__device__ void end() {
|
|
clock_end = clock64();
|
|
}
|
|
|
|
__device__ unsigned long long get_cycles() const {
|
|
return clock_end - clock_start;
|
|
}
|
|
};
|
|
|
|
// Kernel launch parameter optimization for RTX 5090
|
|
extern "C" __device__ void calculate_optimal_launch_params(
|
|
int batch_size, int num_heads, int seq_len, int head_dim,
|
|
int* optimal_grid_x, int* optimal_grid_y, int* optimal_block_x, int* optimal_block_y
|
|
) {
|
|
// RTX 5090 has 128 SMs, optimize grid dimensions
|
|
int total_attention_heads = batch_size * num_heads;
|
|
int seq_blocks = (seq_len + 63) / 64; // 64 is optimal block size for seq dimension
|
|
|
|
*optimal_grid_x = min(total_attention_heads, 128);
|
|
*optimal_grid_y = seq_blocks;
|
|
|
|
// Optimize block dimensions for maximum occupancy
|
|
int threads_per_block = calculate_optimal_block_size(seq_len, head_dim);
|
|
*optimal_block_x = min(threads_per_block, 1024);
|
|
*optimal_block_y = 1;
|
|
}
|
|
|
|
// Shared memory banking optimization
|
|
template<int BANK_SIZE = 32>
|
|
__device__ int avoid_bank_conflicts(int index, int offset = 1) {
|
|
return index + (index / BANK_SIZE) * offset;
|
|
}
|
|
|
|
// RTX 5090 specific numerical precision optimization
|
|
__device__ float high_precision_accumulate(float a, float b) {
|
|
#ifdef RTX_5090_OPTIMIZATION_ENABLED
|
|
// Use fused multiply-add for better precision
|
|
return __fmaf_rn(a, 1.0f, b);
|
|
#else
|
|
return a + b;
|
|
#endif
|
|
}
|
|
|
|
// Advanced memory management for Flash Attention
|
|
class FlashAttentionMemoryManager {
|
|
private:
|
|
void* shared_memory_pool;
|
|
size_t pool_size;
|
|
size_t allocated_bytes;
|
|
|
|
public:
|
|
__device__ FlashAttentionMemoryManager(void* pool, size_t size)
|
|
: shared_memory_pool(pool), pool_size(size), allocated_bytes(0) {}
|
|
|
|
__device__ void* allocate(size_t bytes) {
|
|
size_t aligned_bytes = (bytes + SHARED_MEM_ALIGNMENT - 1) & ~(SHARED_MEM_ALIGNMENT - 1);
|
|
if (allocated_bytes + aligned_bytes <= pool_size) {
|
|
void* ptr = (char*)shared_memory_pool + allocated_bytes;
|
|
allocated_bytes += aligned_bytes;
|
|
return ptr;
|
|
}
|
|
return nullptr; // Out of memory
|
|
}
|
|
|
|
__device__ void reset() {
|
|
allocated_bytes = 0;
|
|
}
|
|
|
|
__device__ size_t available_bytes() const {
|
|
return pool_size - allocated_bytes;
|
|
}
|
|
}; |