rustytorch/crates/training/rtx-flash-attention/cuda/flash_attention_forward.cu

/*
 * Flash Attention Forward CUDA Kernel
 * 
 * Implements the Flash Attention algorithm with O(n) memory complexity
 * using SRAM tiling and online softmax computation.
 * 
 * Reference: "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
 * https://arxiv.org/abs/2205.14135
 */

#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>
#include <cooperative_groups.h>

// Use the configuration defines from Rust
#ifndef BLOCK_SIZE_Q
#define BLOCK_SIZE_Q 64
#endif

#ifndef BLOCK_SIZE_KV
#define BLOCK_SIZE_KV 64
#endif

#ifndef HEAD_DIM
#define HEAD_DIM 128
#endif

#ifndef NUM_HEADS
#define NUM_HEADS 32
#endif

#ifndef MAX_SEQ_LEN
#define MAX_SEQ_LEN 32768
#endif

// Constants
#define WARP_SIZE 32
#define MAX_THREADS_PER_BLOCK 1024
#define SHARED_MEM_ALIGNMENT 16

using namespace nvcuda;
namespace cg = cooperative_groups;

// Utility functions for half precision arithmetic
__device__ __forceinline__ float half_to_float(half x) {
    return __half2float(x);
}

__device__ __forceinline__ half float_to_half(float x) {
    return __float2half(x);
}

// Online softmax state for numerical stability
struct OnlineSoftmaxState {
    float m; // running max
    float l; // running sum
    
    __device__ OnlineSoftmaxState() : m(-INFINITY), l(0.0f) {}
    
    __device__ void update(float x) {
        float m_new = fmaxf(m, x);
        float l_new = l * expf(m - m_new) + expf(x - m_new);
        m = m_new;
        l = l_new;
    }
    
    __device__ void merge(const OnlineSoftmaxState& other) {
        float m_new = fmaxf(m, other.m);
        float l_new = l * expf(m - m_new) + other.l * expf(other.m - m_new);
        m = m_new;
        l = l_new;
    }
    
    __device__ float get_log_sum_exp() const {
        return m + logf(l);
    }
};

// Warp-level reduction for online softmax
__device__ OnlineSoftmaxState warp_reduce_softmax(OnlineSoftmaxState state) {
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        float other_m = __shfl_down_sync(0xffffffff, state.m, offset);
        float other_l = __shfl_down_sync(0xffffffff, state.l, offset);
        OnlineSoftmaxState other;
        other.m = other_m;
        other.l = other_l;
        state.merge(other);
    }
    return state;
}

// Block-level reduction for online softmax
__device__ OnlineSoftmaxState block_reduce_softmax(OnlineSoftmaxState state) {
    __shared__ OnlineSoftmaxState shared_states[32]; // Max warps per block
    
    int warp_id = threadIdx.x / WARP_SIZE;
    int lane_id = threadIdx.x % WARP_SIZE;
    
    // Warp-level reduce
    state = warp_reduce_softmax(state);
    
    // Store warp result
    if (lane_id == 0) {
        shared_states[warp_id] = state;
    }
    
    __syncthreads();
    
    // Block-level reduce
    if (warp_id == 0) {
        OnlineSoftmaxState final_state;
        if (lane_id < (blockDim.x + WARP_SIZE - 1) / WARP_SIZE) {
            final_state = shared_states[lane_id];
        }
        final_state = warp_reduce_softmax(final_state);
        
        // Broadcast result
        shared_states[0] = final_state;
    }
    
    __syncthreads();
    
    return shared_states[0];
}

// Tensor core matrix multiplication for FP16
__device__ void tensor_core_gemm_16x16x16(
    const half* a, const half* b, float* c,
    int lda, int ldb, int ldc
) {
    // Use Tensor Core WMMA API for high performance
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> 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);
}

// Main Flash Attention forward kernel
extern "C" __global__ void flash_attention_forward_kernel(
    const half* __restrict__ q,      // Query [batch * heads, seq_len, head_dim]
    const half* __restrict__ k,      // Key [batch * heads, seq_len, head_dim]
    const half* __restrict__ v,      // Value [batch * heads, seq_len, head_dim]
    half* __restrict__ o,            // Output [batch * heads, seq_len, head_dim]
    float* __restrict__ lse,         // Log-sum-exp [batch * heads, seq_len]
    int seq_len,
    int head_dim,
    float softmax_scale,
    unsigned int causal,
    int block_size_q,
    int block_size_kv
) {
    // Shared memory for tiling
    extern __shared__ char shared_mem[];
    
    // Partition shared memory
    half* q_shared = reinterpret_cast<half*>(shared_mem);
    half* k_shared = q_shared + BLOCK_SIZE_Q * HEAD_DIM;
    half* v_shared = k_shared + BLOCK_SIZE_KV * HEAD_DIM;
    float* scores_shared = reinterpret_cast<float*>(v_shared + BLOCK_SIZE_KV * HEAD_DIM);
    
    // Thread and block indices
    int batch_head_idx = blockIdx.x;
    int q_block_idx = blockIdx.y;
    int tid = threadIdx.x;
    int warp_id = tid / WARP_SIZE;
    int lane_id = tid % WARP_SIZE;
    
    // Calculate offsets
    int q_offset = batch_head_idx * seq_len * head_dim;
    int k_offset = batch_head_idx * seq_len * head_dim;
    int v_offset = batch_head_idx * seq_len * head_dim;
    int o_offset = batch_head_idx * seq_len * head_dim;
    int lse_offset = batch_head_idx * seq_len;
    
    // Q block range
    int q_start = q_block_idx * BLOCK_SIZE_Q;
    int q_end = min(q_start + BLOCK_SIZE_Q, seq_len);
    int q_size = q_end - q_start;
    
    // Load Q block into shared memory
    for (int i = tid; i < q_size * head_dim; i += blockDim.x) {
        int q_row = i / head_dim;
        int q_col = i % head_dim;
        if (q_start + q_row < seq_len) {
            q_shared[q_row * head_dim + q_col] = q[q_offset + (q_start + q_row) * head_dim + q_col];
        }
    }
    __syncthreads();
    
    // Initialize output and online softmax state
    float o_local[HEAD_DIM] = {0.0f};
    OnlineSoftmaxState softmax_state;
    
    // Iterate over KV blocks
    for (int kv_block = 0; kv_block * BLOCK_SIZE_KV < seq_len; kv_block++) {
        int kv_start = kv_block * BLOCK_SIZE_KV;
        int kv_end = min(kv_start + BLOCK_SIZE_KV, seq_len);
        int kv_size = kv_end - kv_start;
        
        // Load K and V blocks into shared memory
        for (int i = tid; i < kv_size * head_dim; i += blockDim.x) {
            int kv_row = i / head_dim;
            int kv_col = i % head_dim;
            if (kv_start + kv_row < seq_len) {
                k_shared[kv_row * head_dim + kv_col] = k[k_offset + (kv_start + kv_row) * head_dim + kv_col];
                v_shared[kv_row * head_dim + kv_col] = v[v_offset + (kv_start + kv_row) * head_dim + kv_col];
            }
        }
        __syncthreads();
        
        // Compute attention scores: Q @ K^T
        for (int q_local_idx = 0; q_local_idx < q_size; q_local_idx++) {
            if (tid < kv_size) {
                float score = 0.0f;
                
                // Dot product using vectorized loads
                for (int d = 0; d < head_dim; d += 4) {
                    float4 q_vec = reinterpret_cast<const float4*>(&q_shared[q_local_idx * head_dim + d])[0];
                    float4 k_vec = reinterpret_cast<const float4*>(&k_shared[tid * head_dim + d])[0];
                    
                    score += q_vec.x * k_vec.x + q_vec.y * k_vec.y + q_vec.z * k_vec.z + q_vec.w * k_vec.w;
                }
                
                score *= softmax_scale;
                
                // Apply causal mask
                int q_pos = q_start + q_local_idx;
                int k_pos = kv_start + tid;
                if (causal && k_pos > q_pos) {
                    score = -INFINITY;
                }
                
                scores_shared[q_local_idx * BLOCK_SIZE_KV + tid] = score;
                
                // Update online softmax
                if (q_local_idx == 0) { // Only update for one Q position per thread
                    softmax_state.update(score);
                }
            }
        }
        __syncthreads();
        
        // Apply softmax and accumulate values
        for (int q_local_idx = 0; q_local_idx < q_size; q_local_idx++) {
            // Get softmax state for this Q position
            OnlineSoftmaxState local_state;
            for (int kv_idx = 0; kv_idx < kv_size; kv_idx++) {
                float score = scores_shared[q_local_idx * BLOCK_SIZE_KV + kv_idx];
                local_state.update(score);
            }
            
            // Reduce across block to get global softmax state
            OnlineSoftmaxState global_state = block_reduce_softmax(local_state);
            
            // Compute softmax probabilities and accumulate values
            for (int kv_idx = tid; kv_idx < kv_size; kv_idx += blockDim.x) {
                float score = scores_shared[q_local_idx * BLOCK_SIZE_KV + kv_idx];
                float prob = expf(score - global_state.m) / global_state.l;
                
                // Accumulate to output
                for (int d = 0; d < head_dim; d++) {
                    float v_val = half_to_float(v_shared[kv_idx * head_dim + d]);
                    o_local[d] += prob * v_val;
                }
            }
        }
        __syncthreads();
    }
    
    // Store output
    for (int q_local_idx = 0; q_local_idx < q_size; q_local_idx++) {
        int q_global_idx = q_start + q_local_idx;
        if (q_global_idx < seq_len) {
            for (int d = tid; d < head_dim; d += blockDim.x) {
                o[o_offset + q_global_idx * head_dim + d] = float_to_half(o_local[d]);
            }
            
            // Store log-sum-exp (only one thread per Q position)
            if (tid == 0) {
                lse[lse_offset + q_global_idx] = softmax_state.get_log_sum_exp();
            }
        }
    }
}

// Specialized kernel for small sequences (optimization)
extern "C" __global__ void flash_attention_forward_small_kernel(
    const half* __restrict__ q,
    const half* __restrict__ k,
    const half* __restrict__ v,
    half* __restrict__ o,
    float* __restrict__ lse,
    int seq_len,
    int head_dim,
    float softmax_scale,
    unsigned int causal
) {
    // For small sequences, we can fit everything in shared memory
    extern __shared__ char shared_mem[];
    
    half* q_all = reinterpret_cast<half*>(shared_mem);
    half* k_all = q_all + seq_len * head_dim;
    half* v_all = k_all + seq_len * head_dim;
    float* scores_all = reinterpret_cast<float*>(v_all + seq_len * head_dim);
    
    int batch_head_idx = blockIdx.x;
    int tid = threadIdx.x;
    
    // Load all Q, K, V into shared memory
    int offset = batch_head_idx * seq_len * head_dim;
    for (int i = tid; i < seq_len * head_dim; i += blockDim.x) {
        q_all[i] = q[offset + i];
        k_all[i] = k[offset + i];
        v_all[i] = v[offset + i];
    }
    __syncthreads();
    
    // Compute all attention scores
    for (int i = tid; i < seq_len * seq_len; i += blockDim.x) {
        int q_idx = i / seq_len;
        int k_idx = i % seq_len;
        
        float score = 0.0f;
        for (int d = 0; d < head_dim; d++) {
            score += half_to_float(q_all[q_idx * head_dim + d]) * half_to_float(k_all[k_idx * head_dim + d]);
        }
        score *= softmax_scale;
        
        // Apply causal mask
        if (causal && k_idx > q_idx) {
            score = -INFINITY;
        }
        
        scores_all[i] = score;
    }
    __syncthreads();
    
    // Apply softmax and compute output
    for (int q_idx = 0; q_idx < seq_len; q_idx++) {
        if (tid == 0) {
            // Compute softmax for this query
            float max_score = -INFINITY;
            for (int k_idx = 0; k_idx < seq_len; k_idx++) {
                max_score = fmaxf(max_score, scores_all[q_idx * seq_len + k_idx]);
            }
            
            float sum_exp = 0.0f;
            for (int k_idx = 0; k_idx < seq_len; k_idx++) {
                float exp_score = expf(scores_all[q_idx * seq_len + k_idx] - max_score);
                scores_all[q_idx * seq_len + k_idx] = exp_score;
                sum_exp += exp_score;
            }
            
            // Normalize and store LSE
            lse[batch_head_idx * seq_len + q_idx] = max_score + logf(sum_exp);
            
            for (int k_idx = 0; k_idx < seq_len; k_idx++) {
                scores_all[q_idx * seq_len + k_idx] /= sum_exp;
            }
        }
    }
    __syncthreads();
    
    // Compute output
    for (int i = tid; i < seq_len * head_dim; i += blockDim.x) {
        int q_idx = i / head_dim;
        int d = i % head_dim;
        
        float output_val = 0.0f;
        for (int k_idx = 0; k_idx < seq_len; k_idx++) {
            float prob = scores_all[q_idx * seq_len + k_idx];
            float v_val = half_to_float(v_all[k_idx * head_dim + d]);
            output_val += prob * v_val;
        }
        
        o[batch_head_idx * seq_len * head_dim + i] = float_to_half(output_val);
    }
}