04-算子融合与优化技术

// 未融合版本
__global__ void add_kernel(float* c, const float* a, const float* b, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) c[idx] = a[idx] + b[idx];
}

__global__ void relu_kernel(float* y, const float* x, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) y[idx] = fmaxf(x[idx], 0.0f);
}

// 融合版本：Add + ReLU
__global__ void fused_add_relu(
    float* __restrict__ output,
    const float* __restrict__ a,
    const float* __restrict__ b,
    int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float val = a[idx] + b[idx];
        output[idx] = fmaxf(val, 0.0f);
    }
}

// 更复杂的融合：Mul + Add + Sigmoid
__global__ void fused_mul_add_sigmoid(
    float* __restrict__ output,
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    int n
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float val = x[idx] * w[idx] + b[idx];
        output[idx] = 1.0f / (1.0f + expf(-val));
    }
}

// 使用 float4 向量化的融合 Kernel
__global__ void fused_add_relu_vec4(
    float4* __restrict__ output,
    const float4* __restrict__ a,
    const float4* __restrict__ b,
    int n  // n 是 float4 的数量
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        float4 va = a[idx];
        float4 vb = b[idx];

        float4 result;
        result.x = fmaxf(va.x + vb.x, 0.0f);
        result.y = fmaxf(va.y + vb.y, 0.0f);
        result.z = fmaxf(va.z + vb.z, 0.0f);
        result.w = fmaxf(va.w + vb.w, 0.0f);

        output[idx] = result;
    }
}

// LayerNorm + Dropout + Add (残差连接) 融合
__global__ void fused_layernorm_dropout_add(
    float* __restrict__ output,
    const float* __restrict__ input,
    const float* __restrict__ residual,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ dropout_mask,
    int batch_size,
    int hidden_size,
    float dropout_prob,
    float eps,
    unsigned long long seed
) {
    int batch_idx = blockIdx.x;
    if (batch_idx >= batch_size) return;

    const float* in_row = input + batch_idx * hidden_size;
    const float* res_row = residual + batch_idx * hidden_size;
    float* out_row = output + batch_idx * hidden_size;
    float* mask_row = dropout_mask + batch_idx * hidden_size;

    // Step 1: Compute mean
    float local_sum = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        local_sum += in_row[i];
    }
    float mean = block_reduce_sum<256>(local_sum) / hidden_size;

    __shared__ float s_mean;
    if (threadIdx.x == 0) s_mean = mean;
    __syncthreads();
    mean = s_mean;

    // Step 2: Compute variance
    float local_var = 0.0f;
    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        float diff = in_row[i] - mean;
        local_var += diff * diff;
    }
    float var = block_reduce_sum<256>(local_var) / hidden_size;

    __shared__ float s_rstd;
    if (threadIdx.x == 0) s_rstd = rsqrtf(var + eps);
    __syncthreads();
    float rstd = s_rstd;

    // Step 3: Normalize + Dropout + Add residual
    // 使用 curand 生成 dropout mask
    curandStatePhilox4_32_10_t state;
    curand_init(seed, batch_idx * hidden_size + threadIdx.x, 0, &state);

    float scale = 1.0f / (1.0f - dropout_prob);

    for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
        // LayerNorm
        float normalized = (in_row[i] - mean) * rstd;
        float ln_out = normalized * gamma[i] + beta[i];

        // Dropout
        float rand = curand_uniform(&state);
        float mask = (rand > dropout_prob) ? 1.0f : 0.0f;
        mask_row[i] = mask;
        float dropout_out = ln_out * mask * scale;

        // Add residual
        out_row[i] = dropout_out + res_row[i];
    }
}

// GEMM 结果的后处理融合
// D = activation(alpha * A @ B + beta * C + bias)

template<typename Activation>
__global__ void gemm_epilogue_fused(
    const float* __restrict__ gemm_output,  // GEMM 原始输出
    const float* __restrict__ bias,
    float* __restrict__ output,
    int M, int N,
    float alpha,
    float beta
) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < N) {
        int idx = row * N + col;
        float val = alpha * gemm_output[idx] + bias[col];
        output[idx] = Activation::apply(val);
    }
}

// 激活函数仿函数
struct ReLU {
    __device__ __forceinline__ static float apply(float x) {
        return fmaxf(x, 0.0f);
    }
};

struct GELU {
    __device__ __forceinline__ static float apply(float x) {
        // GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
        const float c = 0.7978845608f;  // sqrt(2/π)
        const float k = 0.044715f;
        float x3 = x * x * x;
        return 0.5f * x * (1.0f + tanhf(c * (x + k * x3)));
    }
};

struct SiLU {
    __device__ __forceinline__ static float apply(float x) {
        return x / (1.0f + expf(-x));
    }
};

// 在 CUTLASS 中定义 Epilogue
#include <cutlass/epilogue/thread/linear_combination_relu.h>
#include <cutlass/epilogue/thread/linear_combination_gelu.h>

using EpilogueWithReLU = cutlass::epilogue::thread::LinearCombinationRelu<
    float,                                  // Element type
    128 / cutlass::sizeof_bits<float>::value,  // Elements per access
    float,                                  // Accumulator type
    float                                   // Compute type
>;

using EpilogueWithGELU = cutlass::epilogue::thread::LinearCombinationGELU<
    float,
    128 / cutlass::sizeof_bits<float>::value,
    float,
    float
>;

/*
 * 融合 GEMM + Bias + GELU 的完整实现
 * 使用 Tensor Core
 */

#define BM 128
#define BN 128
#define BK 32
#define TM 8
#define TN 8

__global__ void fused_gemm_bias_gelu(
    const half* __restrict__ A,
    const half* __restrict__ B,
    const half* __restrict__ bias,
    half* __restrict__ C,
    int M, int N, int K
) {
    __shared__ half As[BK][BM];
    __shared__ half Bs[BK][BN];

    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;
    int tid = ty * blockDim.x + tx;

    // WMMA fragments
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag[2];
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_frag[2];
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag[2][2];

    // 初始化累加器
    #pragma unroll
    for (int i = 0; i < 2; i++) {
        #pragma unroll
        for (int j = 0; j < 2; j++) {
            wmma::fill_fragment(c_frag[i][j], 0.0f);
        }
    }

    int warp_id = tid / 32;
    int warp_row = warp_id / 2;
    int warp_col = warp_id % 2;

    // K 维度迭代
    for (int bk = 0; bk < K; bk += BK) {
        // 加载 A, B 到共享内存
        load_shared_a(A, As, by, bk, tid, M, K);
        load_shared_b(B, Bs, bx, bk, tid, K, N);
        __syncthreads();

        // WMMA 计算
        #pragma unroll
        for (int k = 0; k < BK; k += 16) {
            #pragma unroll
            for (int m = 0; m < 2; m++) {
                wmma::load_matrix_sync(a_frag[m],
                    &As[k][warp_row * 64 + m * 16], BM);
            }
            #pragma unroll
            for (int n = 0; n < 2; n++) {
                wmma::load_matrix_sync(b_frag[n],
                    &Bs[k][warp_col * 64 + n * 16], BN);
            }

            #pragma unroll
            for (int m = 0; m < 2; m++) {
                #pragma unroll
                for (int n = 0; n < 2; n++) {
                    wmma::mma_sync(c_frag[m][n], a_frag[m], b_frag[n], c_frag[m][n]);
                }
            }
        }
        __syncthreads();
    }

    // Epilogue: Bias + GELU
    int c_row_base = by * BM + warp_row * 64;
    int c_col_base = bx * BN + warp_col * 64;

    #pragma unroll
    for (int m = 0; m < 2; m++) {
        #pragma unroll
        for (int n = 0; n < 2; n++) {
            // 应用 bias 和 GELU 到每个 fragment 元素
            #pragma unroll
            for (int i = 0; i < c_frag[m][n].num_elements; i++) {
                // 计算该元素的列索引
                int elem_col = c_col_base + n * 16 + (i % 16);
                if (elem_col < N) {
                    float val = c_frag[m][n].x[i];
                    // Add bias
                    val += __half2float(bias[elem_col]);
                    // GELU
                    val = gelu_forward(val);
                    c_frag[m][n].x[i] = val;
                }
            }

            // 转换为 half 并存储
            wmma::fragment<wmma::accumulator, 16, 16, 16, half> c_half_frag;
            #pragma unroll
            for (int i = 0; i < c_frag[m][n].num_elements; i++) {
                c_half_frag.x[i] = __float2half(c_frag[m][n].x[i]);
            }

            int c_row = c_row_base + m * 16;
            int c_col = c_col_base + n * 16;
            if (c_row < M && c_col < N) {
                wmma::store_matrix_sync(&C[c_row * N + c_col], c_half_frag, N,
                    wmma::mem_row_major);
            }
        }
    }
}

__device__ __forceinline__ float gelu_forward(float x) {
    const float c = 0.7978845608f;
    const float k = 0.044715f;
    return 0.5f * x * (1.0f + tanhf(c * (x + k * x * x * x)));
}

/*
 * Flash Attention 简化实现
 * 仅用于说明原理，实际实现更复杂
 */

#define BR 64   // Q block size
#define BC 64   // K/V block size
#define d  64   // head dimension

__global__ void flash_attention_kernel(
    const float* __restrict__ Q,  // [N, d]
    const float* __restrict__ K,  // [N, d]
    const float* __restrict__ V,  // [N, d]
    float* __restrict__ O,        // [N, d]
    float* __restrict__ L,        // [N] softmax 归一化因子
    float* __restrict__ M,        // [N] 行最大值
    int N,
    float scale
) {
    // 共享内存
    __shared__ float Qi[BR][d];   // Q 块
    __shared__ float Kj[BC][d];   // K 块
    __shared__ float Vj[BC][d];   // V 块
    __shared__ float Sij[BR][BC]; // 注意力分数块

    int batch_head = blockIdx.x;
    int q_block = blockIdx.y;

    int q_start = q_block * BR;

    // 加载 Q 块到共享内存
    load_block(Q + batch_head * N * d, Qi, q_start, N, d);
    __syncthreads();

    // 初始化输出
    float Oi[d] = {0.0f};
    float mi = -INFINITY;
    float li = 0.0f;

    int thread_row = threadIdx.x / BC;
    int thread_col = threadIdx.x % BC;

    // 遍历 K, V 块
    for (int j = 0; j < (N + BC - 1) / BC; j++) {
        int kv_start = j * BC;

        // 加载 K, V 块
        load_block(K + batch_head * N * d, Kj, kv_start, N, d);
        load_block(V + batch_head * N * d, Vj, kv_start, N, d);
        __syncthreads();

        // 计算 S_ij = Q_i @ K_j^T * scale
        if (thread_row < BR && thread_col < BC) {
            float sum = 0.0f;
            #pragma unroll
            for (int k = 0; k < d; k++) {
                sum += Qi[thread_row][k] * Kj[thread_col][k];
            }
            Sij[thread_row][thread_col] = sum * scale;
        }
        __syncthreads();

        // 对于该线程负责的 Q 行，计算 online softmax
        if (threadIdx.x < BR) {
            int i = threadIdx.x;

            // 找该行最大值
            float row_max = -INFINITY;
            #pragma unroll
            for (int jj = 0; jj < BC && kv_start + jj < N; jj++) {
                row_max = fmaxf(row_max, Sij[i][jj]);
            }

            // Online softmax 更新
            float mi_new = fmaxf(mi, row_max);
            float li_new = li * expf(mi - mi_new);

            // 更新 O
            float scale_old = expf(mi - mi_new);
            #pragma unroll
            for (int k = 0; k < d; k++) {
                Oi[k] *= scale_old;
            }

            // 累加当前块的贡献
            #pragma unroll
            for (int jj = 0; jj < BC && kv_start + jj < N; jj++) {
                float pij = expf(Sij[i][jj] - mi_new);
                li_new += pij;

                #pragma unroll
                for (int k = 0; k < d; k++) {
                    Oi[k] += pij * Vj[jj][k];
                }
            }

            mi = mi_new;
            li = li_new;
        }
        __syncthreads();
    }

    // 归一化并写回
    if (threadIdx.x < BR && q_start + threadIdx.x < N) {
        int row = q_start + threadIdx.x;
        float inv_l = 1.0f / li;

        #pragma unroll
        for (int k = 0; k < d; k++) {
            O[(batch_head * N + row) * d + k] = Oi[k] * inv_l;
        }

        L[batch_head * N + row] = li;
        M[batch_head * N + row] = mi;
    }
}

/*
 * Flash Attention 2 优化点：
 * 1. 更好的并行策略：在 sequence 维度并行
 * 2. 减少非 matmul 操作
 * 3. 更好的 warp 专业化
 */

template<int BR, int BC, int d>
__global__ void flash_attention_2_kernel(
    const half* __restrict__ Q,
    const half* __restrict__ K,
    const half* __restrict__ V,
    half* __restrict__ O,
    int N,
    float scale
) {
    // Thread block 处理一个 Q 块
    // 使用 Tensor Core

    __shared__ half Qi[BR][d];
    __shared__ half Kj[BC][d];
    __shared__ half Vj[BC][d];

    // WMMA fragments
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> q_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> k_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> v_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> s_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> o_frag;

    wmma::fill_fragment(o_frag, 0.0f);

    int warp_id = threadIdx.x / 32;
    int q_block = blockIdx.y;
    int batch_head = blockIdx.x;

    float mi = -INFINITY;
    float li = 0.0f;

    // 加载 Q 块 (只加载一次)
    load_q_block(Q + batch_head * N * d, Qi, q_block * BR, N, d);
    __syncthreads();

    // 遍历 K/V 块
    for (int kv_block = 0; kv_block < (N + BC - 1) / BC; kv_block++) {
        // 加载 K, V
        load_kv_block(K, V, Kj, Vj, batch_head, kv_block * BC, N, d);
        __syncthreads();

        // 使用 Tensor Core 计算 S = Q @ K^T
        wmma::fill_fragment(s_frag, 0.0f);

        for (int k = 0; k < d; k += 16) {
            wmma::load_matrix_sync(q_frag, &Qi[warp_id * 16][k], d);
            wmma::load_matrix_sync(k_frag, &Kj[0][k], d);
            wmma::mma_sync(s_frag, q_frag, k_frag, s_frag);
        }

        // Scale
        #pragma unroll
        for (int i = 0; i < s_frag.num_elements; i++) {
            s_frag.x[i] *= scale;
        }

        // Causal mask (如果需要)
        apply_causal_mask(s_frag, q_block * BR + warp_id * 16, kv_block * BC);

        // Online softmax
        float new_max = reduce_max(s_frag);
        float mi_new = fmaxf(mi, new_max);

        // Rescale previous output
        float scale_factor = expf(mi - mi_new);
        #pragma unroll
        for (int i = 0; i < o_frag.num_elements; i++) {
            o_frag.x[i] *= scale_factor;
        }

        // Compute softmax and accumulate
        float row_sum = 0.0f;
        wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> p_frag;

        #pragma unroll
        for (int i = 0; i < s_frag.num_elements; i++) {
            float p = expf(s_frag.x[i] - mi_new);
            row_sum += p;
            p_frag.x[i] = __float2half(p);
        }

        // O += P @ V
        wmma::load_matrix_sync(v_frag, &Vj[0][0], d);
        wmma::mma_sync(o_frag, p_frag, v_frag, o_frag);

        // Update stats
        li = li * scale_factor + row_sum;
        mi = mi_new;

        __syncthreads();
    }

    // Normalize and store
    float inv_li = 1.0f / li;
    #pragma unroll
    for (int i = 0; i < o_frag.num_elements; i++) {
        o_frag.x[i] *= inv_li;
    }

    wmma::store_matrix_sync(
        &O[(batch_head * N + q_block * BR + warp_id * 16) * d],
        o_frag, d, wmma::mem_row_major
    );
}