diff --git a/result.md b/result.md new file mode 100644 index 00000000..99a9dd08 --- /dev/null +++ b/result.md @@ -0,0 +1,73 @@ +root@autodl-container-mamykqv3ku-9260f267:~/autodl-tmp/Learning-CUDA# SKIP_ATTENTION=1 make +=== Running tests (output from src/kernels.o) === +=== Verbose mode: Enabled (using '--verbose') === +./test_kernels +Testing on device: NVIDIA GeForce RTX 3080 Ti + +=== rmsNorm Tests === +Test # 1: float | Verification: Passed +Test # 1: half | Verification: Passed +Test # 2: float | Verification: Passed +Test # 2: half | Verification: Passed +Test # 3: float | Verification: Passed +Test # 3: half | Verification: Passed +Test # 4: float | Verification: Passed +Test # 4: half | Verification: Passed +Test # 5: float | Verification: Passed +Test # 5: half | Verification: Passed +Test # 6: float | Verification: Passed +Test # 6: half | Verification: Passed +Test # 7: float | Verification: Passed +Test # 7: half | Verification: Passed +Test # 8: float | Verification: Passed +Test # 8: half | Verification: Passed +Test # 9: float | Verification: Passed +Test # 9: half | Verification: Passed +Test #10: float | Verification: Passed +Test #10: half | Verification: Passed +Test #11: float | Verification: Passed +Test #11: half | Verification: Passed +Test #12: float | Verification: Passed +Test #12: half | Verification: Passed +Test #13: float | Verification: Passed +Test #13: half | Verification: Passed + +root@autodl-container-mamykqv3ku-9260f267:~/autodl-tmp/Learning-CUDA# SKIP_RMS_NORM=1 make +=== Compiling student code (src/kernels.cu ) === +nvcc -std=c++17 -O0 -DPLATFORM_NVIDIA -c src/kernels.cu -o src/kernels.o +=== Linking executable (student code + test logic) === +nvcc -std=c++17 -O0 -DPLATFORM_NVIDIA -o test_kernels src/kernels.o tester/tester_nv.o +=== Running tests (output from src/kernels.o) === +=== Verbose mode: Enabled (using '--verbose') === +./test_kernels +Testing on device: NVIDIA GeForce RTX 3080 Ti + +=== Attention Tests === +Test # 1: float | Verification: Passed +Test # 1: half | Verification: Passed +Test # 2: float | Verification: Passed +Test # 2: half | Verification: Passed +Test # 3: float | Verification: Passed +Test # 3: half | Verification: Passed +Test # 4: float | Verification: Passed +Test # 4: half | Verification: Passed +Test # 5: float | Verification: Passed +Test # 5: half | Verification: Passed +Test # 6: float | Verification: Passed +Test # 6: half | Verification: Passed +Test # 7: float | Verification: Passed +Test # 7: half | Verification: Passed +Test # 8: float | Verification: Passed +Test # 8: half | Verification: Passed +Test # 9: float | Verification: Passed +Test # 9: half | Verification: Passed +Test #10: float | Verification: Passed +Test #10: half | Verification: Passed +Test #11: float | Verification: Passed +Test #11: half | Verification: Passed +Test #12: float | Verification: Passed +Test #12: half | Verification: Passed +Test #13: float | Verification: Passed +Test #13: half | Verification: Passed +Test #14: float | Verification: Passed +Test #14: half | Verification: Passed diff --git a/src/kernels.cu b/src/kernels.cu index 2cc53e7e..9a1197af 100644 --- a/src/kernels.cu +++ b/src/kernels.cu @@ -1,61 +1,300 @@ +// final submission #include #include +#include +#include +#include #include "../tester/utils.h" -/** - * @brief Computes RMSNorm over the last dimension of a 2D tensor. - * - * The input is a row-major matrix with shape [rows, hidden_dim]. For each row - * i and column j: - * - * output[i, j] = input[i, j] * rsqrt(mean(input[i, :]^2) + eps) * weight[j] - * - * The output vector is preallocated with rows * hidden_dim elements. - * - * @tparam T Data type of input, weight, and output tensors. - * @param[in] h_input Flattened input matrix of shape [rows, hidden_dim]. - * @param[in] h_weight Per-column scale vector of shape [hidden_dim]. - * @param[out] h_output Flattened output matrix of shape [rows, hidden_dim]. - * @param[in] rows Number of rows/tokens. - * @param[in] hidden_dim Size of the normalized dimension. - * @param[in] eps Numerical stability epsilon. - */ +#ifndef CHECK_CUDA +#define CHECK_CUDA(call) do { \ + cudaError_t err = call; \ + if(err != cudaSuccess) { \ + printf("CUDA Error at %s:%d - %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) +#endif + +// ===================================================================== +// 修复点：Makefile中传入的是 PLATFORM_ILUVATAR，并且 64线程Warp 需要 64位 的 Mask +// ===================================================================== +#if defined(PLATFORM_ILUVATAR) || defined(__ILUVATAR__) +#define WARP_SIZE 64 +#define FULL_MASK 0xffffffffffffffffULL // 必须是 64 位全 1 掩码 +#define BLOCK_SIZE 64 +#define Bc 128 +#else +#define WARP_SIZE 32 +#define FULL_MASK 0xffffffff +#define BLOCK_SIZE 128 +#define Bc 64 +#endif + +#define MAX_HEAD_DIM 256 +#define MAX_BC 128 + +// Layout Macros +#define QO_OFFSET(b,s,h,d) ((((b)*target_seq_len+(s))*query_heads+(h))*head_dim+(d)) +#define KV_OFFSET(b,s,h,d) ((((b)*src_seq_len+(s))*kv_heads+(h))*head_dim+(d)) + +// ===================================================================== +// Warp / Block Reduce Helpers +// ===================================================================== +__inline__ __device__ float warpReduceSum(float val) { + for(int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(FULL_MASK, val, offset); + } + return val; +} + +__inline__ __device__ float blockReduceSum(float val) { + __shared__ float shared[64]; + int lane = threadIdx.x % WARP_SIZE; + int wid = threadIdx.x / WARP_SIZE; + + val = warpReduceSum(val); + + if (lane == 0) { + shared[wid] = val; + } + __syncthreads(); + + val = (threadIdx.x < (blockDim.x + WARP_SIZE - 1)/WARP_SIZE) ? shared[lane] : 0.0f; + + if (wid == 0) { + val = warpReduceSum(val); + } + return val; +} + +// ===================================================================== +// RMSNorm Kernel +// ===================================================================== +template +__global__ void rmsNormKernel( + const T* __restrict__ input, + const T* __restrict__ weight, + T* __restrict__ output, + int rows, + int hidden_dim, + float eps +) +{ + int row = blockIdx.x; + if(row >= rows) return; + + float sum = 0.0f; + + for(int col = threadIdx.x; col < hidden_dim; col += blockDim.x) { + float value = static_cast(input[row * hidden_dim + col]); + sum += value * value; + } + + float total_sum = blockReduceSum(sum); + + __shared__ float s_scale; + if (threadIdx.x == 0) { + s_scale = rsqrtf(total_sum / hidden_dim + eps); + } + __syncthreads(); + + float scale = s_scale; + + for(int col = threadIdx.x; col < hidden_dim; col += blockDim.x) { + float value = static_cast(input[row * hidden_dim + col]); + float w = static_cast(weight[col]); + output[row * hidden_dim + col] = static_cast(value * scale * w); + } +} + template void rmsNorm(const std::vector& h_input, const std::vector& h_weight, std::vector& h_output, size_t rows, size_t hidden_dim, float eps) { - // TODO: Implement the rmsNorm function + T* d_input; + T* d_weight; + T* d_output; + + size_t input_bytes = rows * hidden_dim * sizeof(T); + size_t weight_bytes = hidden_dim * sizeof(T); + + CHECK_CUDA(cudaMalloc(&d_input, input_bytes)); + CHECK_CUDA(cudaMalloc(&d_output, input_bytes)); + CHECK_CUDA(cudaMalloc(&d_weight, weight_bytes)); + + CHECK_CUDA(cudaMemcpy(d_input, h_input.data(), input_bytes, cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_weight, h_weight.data(), weight_bytes, cudaMemcpyHostToDevice)); + + dim3 grid(rows); + dim3 block(256); + + rmsNormKernel<<>>(d_input, d_weight, d_output, rows, hidden_dim, eps); + + CHECK_CUDA(cudaGetLastError()); + CHECK_CUDA(cudaDeviceSynchronize()); + + CHECK_CUDA(cudaMemcpy(h_output.data(), d_output, input_bytes, cudaMemcpyDeviceToHost)); + + CHECK_CUDA(cudaFree(d_input)); + CHECK_CUDA(cudaFree(d_weight)); + CHECK_CUDA(cudaFree(d_output)); } -/** - * @brief Computes flash attention for given query, key, and value tensors. - * - * @tparam T Data type (float) for input/output tensors - * @param[in] h_q Query tensor of shape [batch_size, tgt_seq_len, query_heads, head_dim] - * @param[in] h_k Key tensor of shape [batch_size, src_seq_len, kv_heads, head_dim] - * @param[in] h_v Value tensor of shape [batch_size, src_seq_len, kv_heads, head_dim] - * @param[out] h_o Output attention tensor of shape [batch_size, tgt_seq_len, query_heads, head_dim] - * @param[in] batch_size Batch dimension size - * @param[in] target_seq_len Target sequence length - * @param[in] src_seq_len Source sequence length - * @param[in] query_heads Number of query attention heads - * @param[in] kv_heads Number of key/value heads (supports grouped query attention) - * @param[in] head_dim Dimension size of each attention head - * @param[in] is_causal Whether to apply causal masking - */ -template -void flashAttention(const std::vector& h_q, const std::vector& h_k, - const std::vector& h_v, std::vector& h_o, - int batch_size, int target_seq_len, int src_seq_len, - int query_heads, int kv_heads, int head_dim, bool is_causal) { - // TODO: Implement the flash attention function +// ===================================================================== +// Reference Attention Kernel +// ===================================================================== +template +__global__ void referenceAttentionKernel( + const T* __restrict__ q, + const T* __restrict__ k, + const T* __restrict__ v, + T* __restrict__ o, + int batch_size, + int target_seq_len, + int src_seq_len, + int query_heads, + int kv_heads, + int head_dim, + bool is_causal +) +{ + int total_queries = batch_size * target_seq_len * query_heads; + int gtid = blockIdx.x * blockDim.x + threadIdx.x; + + if(gtid >= total_queries) return; + + int q_head_idx = gtid % query_heads; + int rem = gtid / query_heads; + int t_idx = rem % target_seq_len; + int b_idx = rem / target_seq_len; + + int kv_head_idx = q_head_idx / (query_heads / kv_heads); + float scale = __frcp_rn(sqrtf((float)head_dim)); + + float q_buf[MAX_HEAD_DIM]; + float o_buf[MAX_HEAD_DIM]; + + int qo_base = QO_OFFSET(b_idx, t_idx, q_head_idx, 0); + + for(int d = 0; d < head_dim; d++) { + q_buf[d] = static_cast(q[qo_base + d]); + o_buf[d] = 0.0f; + } + + float global_max = -1e9f; + + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) { + if(is_causal && s_idx > t_idx) continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + float score = 0.0f; + + for(int d = 0; d < head_dim; d++) { + score = fmaf(q_buf[d], static_cast(k[kv_base + d]), score); + } + + global_max = fmaxf(global_max, score * scale); + } + + float sum = 0.0f; + + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) { + if(is_causal && s_idx > t_idx) continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + float score = 0.0f; + + for(int d = 0; d < head_dim; d++) { + score = fmaf(q_buf[d], static_cast(k[kv_base + d]), score); + } + + sum += expf(score * scale - global_max); + } + + float inv_sum = (sum > 0.0f) ? __frcp_rn(sum) : 0.0f; + + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) { + if(is_causal && s_idx > t_idx) continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + float score = 0.0f; + + for(int d = 0; d < head_dim; d++) { + score = fmaf(q_buf[d], static_cast(k[kv_base + d]), score); + } + + float prob = expf(score * scale - global_max); + float factor = prob * inv_sum; + + for(int d = 0; d < head_dim; d++) { + o_buf[d] = fmaf(factor, static_cast(v[kv_base + d]), o_buf[d]); + } + } + + for(int d = 0; d < head_dim; d++) { + o[qo_base + d] = static_cast(o_buf[d]); + } +} + +// ===================================================================== +// Host Wrapper Function +// ===================================================================== +template +void flashAttention( + const std::vector& h_q, + const std::vector& h_k, + const std::vector& h_v, + std::vector& h_o, + int batch_size, + int target_seq_len, + int src_seq_len, + int query_heads, + int kv_heads, + int head_dim, + bool is_causal +) +{ + size_t q_size = batch_size * target_seq_len * query_heads * head_dim * sizeof(T); + size_t k_size = batch_size * src_seq_len * kv_heads * head_dim * sizeof(T); + size_t v_size = batch_size * src_seq_len * kv_heads * head_dim * sizeof(T); + h_o.resize(batch_size * target_seq_len * query_heads * head_dim); + + T *d_q, *d_k, *d_v, *d_o; + CHECK_CUDA(cudaMalloc(&d_q, q_size)); + CHECK_CUDA(cudaMalloc(&d_k, k_size)); + CHECK_CUDA(cudaMalloc(&d_v, v_size)); + CHECK_CUDA(cudaMalloc(&d_o, q_size)); + + CHECK_CUDA(cudaMemcpy(d_q, h_q.data(), q_size, cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_k, h_k.data(), k_size, cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_v, h_v.data(), v_size, cudaMemcpyHostToDevice)); + + int total_queries = batch_size * target_seq_len * query_heads; + dim3 block(256); + dim3 grid((total_queries + block.x - 1) / block.x); + + referenceAttentionKernel<<>>( + d_q, d_k, d_v, d_o, + batch_size, target_seq_len, src_seq_len, + query_heads, kv_heads, head_dim, is_causal + ); + + CHECK_CUDA(cudaGetLastError()); + CHECK_CUDA(cudaDeviceSynchronize()); + + CHECK_CUDA(cudaMemcpy(h_o.data(), d_o, q_size, cudaMemcpyDeviceToHost)); + + CHECK_CUDA(cudaFree(d_q)); + CHECK_CUDA(cudaFree(d_k)); + CHECK_CUDA(cudaFree(d_v)); + CHECK_CUDA(cudaFree(d_o)); } -// ********************************************************************* -// Explicit Template Instantiations (REQUIRED FOR LINKING WITH TESTER.O) -// DO NOT MODIFY THIS SECTION -// ********************************************************************* +// ===================================================================== +// Explicit Template Instantiations +// ===================================================================== template void rmsNorm(const std::vector&, const std::vector&, std::vector&, size_t, size_t, float); template void rmsNorm(const std::vector&, const std::vector&, diff --git a/src/kernels.maca b/src/kernels.maca index 4c320f21..c732248b 100644 --- a/src/kernels.maca +++ b/src/kernels.maca @@ -1,55 +1,321 @@ +// final submission #include +#include #include #include "../tester/utils.h" -/** - * @brief Computes RMSNorm over the last dimension of a 2D tensor. - * - * The input is a row-major matrix with shape [rows, hidden_dim]. For each row - * i and column j: - * - * output[i, j] = input[i, j] * rsqrt(mean(input[i, :]^2) + eps) * weight[j] - * - * The output vector is preallocated with rows * hidden_dim elements. - * - * @tparam T Data type of input, weight, and output tensors. - * @param[in] h_input Flattened input matrix of shape [rows, hidden_dim]. - * @param[in] h_weight Per-column scale vector of shape [hidden_dim]. - * @param[out] h_output Flattened output matrix of shape [rows, hidden_dim]. - * @param[in] rows Number of rows/tokens. - * @param[in] hidden_dim Size of the normalized dimension. - * @param[in] eps Numerical stability epsilon. - */ +#ifndef CHECK_MC +#define CHECK_MC(call) do { \ + mcError_t err = call; \ + if(err != mcSuccess) { \ + printf("MC Error at %s:%d - %s\n", __FILE__, __LINE__, mcGetErrorString(err)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) +#endif + +// Warp 级别的求和规约 +__inline__ __device__ float warpReduceSum(float val) { + for (int offset = 16; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +// Block 级别的求和规约 +__inline__ __device__ float blockReduceSum(float val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warpReduceSum(val); + + if (lane == 0) { + shared[wid] = val; + } + __syncthreads(); + + val = (threadIdx.x < (blockDim.x / 32)) ? shared[lane] : 0.0f; + + if (wid == 0) { + val = warpReduceSum(val); + } + + return val; +} + +// ********************************************************************* +// RMSNorm Kernel (优化版) +// ********************************************************************* template -void rmsNorm(const std::vector& h_input, const std::vector& h_weight, - std::vector& h_output, size_t rows, size_t hidden_dim, - float eps) { - // TODO: Implement the rmsNorm function +__global__ void rmsNormKernel( + const T* __restrict__ input, + const T* __restrict__ weight, + T* __restrict__ output, + int rows, + int hidden_dim, + float eps +) +{ + int row = blockIdx.x; + if(row >= rows) return; + + float sum = 0.0f; + for(int col = threadIdx.x; col < hidden_dim; col += blockDim.x) + { + float value = static_cast(input[row * hidden_dim + col]); + sum += value * value; + } + + float total_sum = blockReduceSum(sum); + + __shared__ float s_scale; + if (threadIdx.x == 0) + { + s_scale = rsqrtf(total_sum / hidden_dim + eps); + } + __syncthreads(); + + float scale = s_scale; + + for(int col = threadIdx.x; col < hidden_dim; col += blockDim.x) + { + float value = static_cast(input[row * hidden_dim + col]); + float w = static_cast(weight[col]); + output[row * hidden_dim + col] = static_cast(value * scale * w); + } } -/** - * @brief Computes flash attention for given query, key, and value tensors. - * - * @tparam T Data type (float) for input/output tensors - * @param[in] h_q Query tensor of shape [batch_size, tgt_seq_len, query_heads, head_dim] - * @param[in] h_k Key tensor of shape [batch_size, src_seq_len, kv_heads, head_dim] - * @param[in] h_v Value tensor of shape [batch_size, src_seq_len, kv_heads, head_dim] - * @param[out] h_o Output attention tensor of shape [batch_size, tgt_seq_len, query_heads, head_dim] - * @param[in] batch_size Batch dimension size - * @param[in] target_seq_len Target sequence length - * @param[in] src_seq_len Source sequence length - * @param[in] query_heads Number of query attention heads - * @param[in] kv_heads Number of key/value heads (supports grouped query attention) - * @param[in] head_dim Dimension size of each attention head - * @param[in] is_causal Whether to apply causal masking - */ template -void flashAttention(const std::vector& h_q, const std::vector& h_k, - const std::vector& h_v, std::vector& h_o, - int batch_size, int target_seq_len, int src_seq_len, - int query_heads, int kv_heads, int head_dim, bool is_causal) { - // TODO: Implement the flash attention function +void rmsNorm( + const std::vector& h_input, + const std::vector& h_weight, + std::vector& h_output, + size_t rows, + size_t hidden_dim, + float eps +) +{ + T* d_input; + T* d_weight; + T* d_output; + + size_t input_bytes = rows * hidden_dim * sizeof(T); + size_t weight_bytes = hidden_dim * sizeof(T); + h_output.resize(rows * hidden_dim); + + CHECK_MC(mcMalloc(&d_input, input_bytes)); + CHECK_MC(mcMalloc(&d_output, input_bytes)); + CHECK_MC(mcMalloc(&d_weight, weight_bytes)); + + CHECK_MC(mcMemcpy(d_input, h_input.data(), input_bytes, mcMemcpyHostToDevice)); + CHECK_MC(mcMemcpy(d_weight, h_weight.data(), weight_bytes, mcMemcpyHostToDevice)); + + dim3 grid(rows); + dim3 block(256); + + rmsNormKernel<<>>( + d_input, d_weight, d_output, rows, hidden_dim, eps + ); + + CHECK_MC(mcGetLastError()); + CHECK_MC(mcDeviceSynchronize()); + + CHECK_MC(mcMemcpy(h_output.data(), d_output, input_bytes, mcMemcpyDeviceToHost)); + + CHECK_MC(mcFree(d_input)); + CHECK_MC(mcFree(d_weight)); + CHECK_MC(mcFree(d_output)); +} + +// ********************************************************************* +// flashAttention +// ******************************************************************** + +#define LAYOUT_BHSD 0 + +#if LAYOUT_BHSD +#define QO_OFFSET(b,s,h,d) ((((b)*query_heads+(h))*target_seq_len+(s))*head_dim+(d)) +#define KV_OFFSET(b,s,h,d) ((((b)*kv_heads+(h))*src_seq_len+(s))*head_dim+(d)) +#else +#define QO_OFFSET(b,s,h,d) ((((b)*target_seq_len+(s))*query_heads+(h))*head_dim+(d)) +#define KV_OFFSET(b,s,h,d) ((((b)*src_seq_len+(s))*kv_heads+(h))*head_dim+(d)) +#endif + +template +__global__ void flashAttentionKernel( + const T* __restrict__ q, + const T* __restrict__ k, + const T* __restrict__ v, + T* __restrict__ o, + + int batch_size, + int target_seq_len, + int src_seq_len, + + int query_heads, + int kv_heads, + + int head_dim, + + bool is_causal +) +{ + int total = batch_size * target_seq_len * query_heads; + int gtid = blockIdx.x * blockDim.x + threadIdx.x; + if(gtid >= total) + return; + + int q_head_idx = gtid % query_heads; + int rem = gtid / query_heads; + int t_idx = rem % target_seq_len; + int b_idx = rem / target_seq_len; + + // GQA 映射 + int kv_head_idx = q_head_idx / (query_heads / kv_heads); + + float scale = __frcp_rn(sqrtf((float)head_dim)); + + float q_buf[256]; + float o_buf[256]; + + int qo_base = QO_OFFSET(b_idx, t_idx, q_head_idx, 0); + + for(int d = 0; d < head_dim; d++) + q_buf[d] = (float)q[qo_base + d]; + + for(int d = 0; d < head_dim; d++) + o_buf[d] = 0.0f; + + // ===================================================== + // 用 -INFINITY 和 s_idx > t_idx + // ===================================================== + float global_max = -INFINITY; + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) + { + if(is_causal && s_idx > t_idx) + continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + + float score = 0.0f; + for(int d = 0; d < head_dim; d++) + score = fmaf(q_buf[d], (float)k[kv_base + d], score); + + global_max = fmaxf(global_max, score * scale); + } + + // ===================================================== + // s_idx > t_idx + // ===================================================== + float sum = 0.0f; + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) + { + if(is_causal && s_idx > t_idx) + continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + + float score = 0.0f; + for(int d = 0; d < head_dim; d++) + score = fmaf(q_buf[d], (float)k[kv_base + d], score); + + sum += expf(score * scale - global_max); + } + + float inv_sum = (sum != 0.0f) ? __frcp_rn(sum) : 0.0f; + + // ===================================================== + // 用 s_idx > t_idx + // ===================================================== + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) + { + if(is_causal && s_idx > t_idx) + continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + + float score = 0.0f; + for(int d = 0; d < head_dim; d++) + score = fmaf(q_buf[d], (float)k[kv_base + d], score); + + float prob = expf(score * scale - global_max); + float factor = inv_sum * prob; + + for(int d = 0; d < head_dim; d++) + o_buf[d] = fmaf(factor, (float)v[kv_base + d], o_buf[d]); + } + + // Write O + for(int d = 0; d < head_dim; d++) + o[qo_base + d] = (T)o_buf[d]; +} + +template +void flashAttention( + const std::vector& h_q, + const std::vector& h_k, + const std::vector& h_v, + std::vector& h_o, + + int batch_size, + int target_seq_len, + int src_seq_len, + + int query_heads, + int kv_heads, + + int head_dim, + + bool is_causal +) +{ + T *d_q; + T *d_k; + T *d_v; + T *d_o; + + size_t q_size = h_q.size() * sizeof(T); + size_t k_size = h_k.size() * sizeof(T); + size_t v_size = h_v.size() * sizeof(T); + h_o.resize(h_q.size()); + + CHECK_MC(mcMalloc(&d_q, q_size)); + CHECK_MC(mcMalloc(&d_k, k_size)); + CHECK_MC(mcMalloc(&d_v, v_size)); + CHECK_MC(mcMalloc(&d_o, q_size)); + + CHECK_MC(mcMemcpy(d_q, h_q.data(), q_size, mcMemcpyHostToDevice)); + CHECK_MC(mcMemcpy(d_k, h_k.data(), k_size, mcMemcpyHostToDevice)); + CHECK_MC(mcMemcpy(d_v, h_v.data(), v_size, mcMemcpyHostToDevice)); + + int total = batch_size * target_seq_len * query_heads; + + dim3 block(256); + dim3 grid((total + block.x - 1) / block.x); + + flashAttentionKernel<<>>( + d_q, d_k, d_v, d_o, + batch_size, target_seq_len, src_seq_len, + query_heads, kv_heads, + head_dim, + is_causal + ); + + mcError_t err = mcGetLastError(); + if(err != mcSuccess) + { + printf("FlashAttention kernel error: %s\n", mcGetErrorString(err)); + } + + CHECK_MC(mcDeviceSynchronize()); + CHECK_MC(mcMemcpy(h_o.data(), d_o, q_size, mcMemcpyDeviceToHost)); + + CHECK_MC(mcFree(d_q)); + CHECK_MC(mcFree(d_k)); + CHECK_MC(mcFree(d_v)); + CHECK_MC(mcFree(d_o)); } // ********************************************************************* diff --git a/src/kernels.mu b/src/kernels.mu index 1ce371eb..083c4f94 100644 --- a/src/kernels.mu +++ b/src/kernels.mu @@ -1,55 +1,298 @@ +// final submission #include +#include #include +#include // MUSA Runtime API #include "../tester/utils.h" -/** - * @brief Computes RMSNorm over the last dimension of a 2D tensor. - * - * The input is a row-major matrix with shape [rows, hidden_dim]. For each row - * i and column j: - * - * output[i, j] = input[i, j] * rsqrt(mean(input[i, :]^2) + eps) * weight[j] - * - * The output vector is preallocated with rows * hidden_dim elements. - * - * @tparam T Data type of input, weight, and output tensors. - * @param[in] h_input Flattened input matrix of shape [rows, hidden_dim]. - * @param[in] h_weight Per-column scale vector of shape [hidden_dim]. - * @param[out] h_output Flattened output matrix of shape [rows, hidden_dim]. - * @param[in] rows Number of rows/tokens. - * @param[in] hidden_dim Size of the normalized dimension. - * @param[in] eps Numerical stability epsilon. - */ + +extern "C" { + // __attribute__((weak)) 的魔法： + // 如果系统库里有这个函数（v2环境），链接器会忽略我们写的这段代码。 + // 如果系统库里没有这个函数（v1环境），链接器就会使用我们这里的实现兜底。 + __attribute__((weak)) musaError_t musaGetDeviceProperties_v2(musaDeviceProp *prop, int device) { + return musaGetDeviceProperties(prop, device); + } +} + +// ===================================================================== +// MUSA 错误检查宏 +// ===================================================================== +#ifndef CHECK_MUSA +#define CHECK_MUSA(call) do { \ + musaError_t err = call; \ + if(err != musaSuccess) { \ + printf("MUSA Error at %s:%d - %s\n", __FILE__, __LINE__, musaGetErrorString(err)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) +#endif + +// ===================================================================== +// 规约函数：针对 MUSA S4000 WarpSize = 128 做了平台安全适配 +// ===================================================================== +__inline__ __device__ float blockReduceSum(float val) { + // 最大支持 blockDim.x = 1024 + static __shared__ float shared[1024]; + int tid = threadIdx.x; + + // 将所有线程的数据存入共享内存 + shared[tid] = val; + __syncthreads(); + + // 标准树状归约 (Tree Reduction)，不依赖特定的 Warp Size 掩码，S4000 上绝对安全 + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + shared[tid] += shared[tid + stride]; + } + __syncthreads(); + } + return shared[0]; +} + +// ===================================================================== +// RMSNorm +// ===================================================================== template -void rmsNorm(const std::vector& h_input, const std::vector& h_weight, - std::vector& h_output, size_t rows, size_t hidden_dim, - float eps) { - // TODO: Implement the rmsNorm function +__global__ void rmsNormKernel( + const T* __restrict__ input, + const T* __restrict__ weight, + T* __restrict__ output, + int rows, + int hidden_dim, + float eps +) +{ + int row = blockIdx.x; + if(row >= rows) return; + + float sum = 0.0f; + + for(int col = threadIdx.x; col < hidden_dim; col += blockDim.x) + { + float value = static_cast(input[row * hidden_dim + col]); + sum += value * value; + } + + // 调用硬件安全的 Block 归约 + float total_sum = blockReduceSum(sum); + + __shared__ float s_scale; + if (threadIdx.x == 0) + { + s_scale = rsqrtf(total_sum / hidden_dim + eps); + } + __syncthreads(); + + float scale = s_scale; + + for(int col = threadIdx.x; col < hidden_dim; col += blockDim.x) + { + float value = static_cast(input[row * hidden_dim + col]); + float w = static_cast(weight[col]); + output[row * hidden_dim + col] = static_cast(value * scale * w); + } } -/** - * @brief Computes flash attention for given query, key, and value tensors. - * - * @tparam T Data type (float) for input/output tensors - * @param[in] h_q Query tensor of shape [batch_size, tgt_seq_len, query_heads, head_dim] - * @param[in] h_k Key tensor of shape [batch_size, src_seq_len, kv_heads, head_dim] - * @param[in] h_v Value tensor of shape [batch_size, src_seq_len, kv_heads, head_dim] - * @param[out] h_o Output attention tensor of shape [batch_size, tgt_seq_len, query_heads, head_dim] - * @param[in] batch_size Batch dimension size - * @param[in] target_seq_len Target sequence length - * @param[in] src_seq_len Source sequence length - * @param[in] query_heads Number of query attention heads - * @param[in] kv_heads Number of key/value heads (supports grouped query attention) - * @param[in] head_dim Dimension size of each attention head - * @param[in] is_causal Whether to apply causal masking - */ template -void flashAttention(const std::vector& h_q, const std::vector& h_k, - const std::vector& h_v, std::vector& h_o, - int batch_size, int target_seq_len, int src_seq_len, - int query_heads, int kv_heads, int head_dim, bool is_causal) { - // TODO: Implement the flash attention function +void rmsNorm( + const std::vector& h_input, + const std::vector& h_weight, + std::vector& h_output, + size_t rows, + size_t hidden_dim, + float eps +) +{ + T* d_input; + T* d_weight; + T* d_output; + + size_t input_bytes = rows * hidden_dim * sizeof(T); + size_t weight_bytes = hidden_dim * sizeof(T); + + // MUSA 显存分配 + CHECK_MUSA(musaMalloc(&d_input, input_bytes)); + CHECK_MUSA(musaMalloc(&d_output, input_bytes)); + CHECK_MUSA(musaMalloc(&d_weight, weight_bytes)); + + // Host -> Device + CHECK_MUSA(musaMemcpy(d_input, h_input.data(), input_bytes, musaMemcpyHostToDevice)); + CHECK_MUSA(musaMemcpy(d_weight, h_weight.data(), weight_bytes, musaMemcpyHostToDevice)); + + dim3 grid(rows); + dim3 block(256); // 256 threads = S4000 下刚好 2 个 128-Warp，调度优良 + + rmsNormKernel<<>>( + d_input, d_weight, d_output, rows, hidden_dim, eps + ); + + CHECK_MUSA(musaGetLastError()); + CHECK_MUSA(musaDeviceSynchronize()); + + // Device -> Host + CHECK_MUSA(musaMemcpy(h_output.data(), d_output, input_bytes, musaMemcpyDeviceToHost)); + + CHECK_MUSA(musaFree(d_input)); + CHECK_MUSA(musaFree(d_weight)); + CHECK_MUSA(musaFree(d_output)); +} + +// ===================================================================== +// FlashAttention +// ===================================================================== + +#define QO_OFFSET(b,s,h,d) ((((b)*target_seq_len+(s))*query_heads+(h))*head_dim+(d)) +#define KV_OFFSET(b,s,h,d) ((((b)*src_seq_len+(s))*kv_heads+(h))*head_dim+(d)) + +template +__global__ void flashAttentionKernel( + const T* __restrict__ q, + const T* __restrict__ k, + const T* __restrict__ v, + T* __restrict__ o, + int batch_size, + int target_seq_len, + int src_seq_len, + int query_heads, + int kv_heads, + int head_dim, + bool is_causal +) +{ + int total = batch_size * target_seq_len * query_heads; + int gtid = blockIdx.x * blockDim.x + threadIdx.x; + if(gtid >= total) + return; + + int q_head_idx = gtid % query_heads; + int rem = gtid / query_heads; + int t_idx = rem % target_seq_len; + int b_idx = rem / target_seq_len; + + int kv_head_idx = q_head_idx / (query_heads / kv_heads); + + float scale = rsqrtf((float)head_dim); + + float q_buf[256]; + float o_buf[256]; + + int qo_base = QO_OFFSET(b_idx, t_idx, q_head_idx, 0); + + for(int d = 0; d < head_dim; d++) + q_buf[d] = (float)q[qo_base + d]; + + for(int d = 0; d < head_dim; d++) + o_buf[d] = 0.0f; + + // Pass 1: 寻找最大值 (使用 fmaxf 保证 32bit 单精度) + float global_max = -INFINITY; + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) + { + if(is_causal && s_idx > t_idx) continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + + float score = 0.0f; + for(int d = 0; d < head_dim; d++) + score = fmaf(q_buf[d], (float)k[kv_base + d], score); + + global_max = fmaxf(global_max, score * scale); + } + + // Pass 2: 计算 Softmax 分母 (使用 expf) + float sum = 0.0f; + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) + { + if(is_causal && s_idx > t_idx) continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + + float score = 0.0f; + for(int d = 0; d < head_dim; d++) + score = fmaf(q_buf[d], (float)k[kv_base + d], score); + + sum += expf(score * scale - global_max); + } + + float inv_sum = (sum != 0.0f) ? (1.0f / sum) : 0.0f; + + // Pass 3: 计算 O 矩阵 + for(int s_idx = 0; s_idx < src_seq_len; s_idx++) + { + if(is_causal && s_idx > t_idx) continue; + + int kv_base = KV_OFFSET(b_idx, s_idx, kv_head_idx, 0); + + float score = 0.0f; + for(int d = 0; d < head_dim; d++) + score = fmaf(q_buf[d], (float)k[kv_base + d], score); + + float prob = expf(score * scale - global_max); + float factor = inv_sum * prob; + + for(int d = 0; d < head_dim; d++) + o_buf[d] = fmaf(factor, (float)v[kv_base + d], o_buf[d]); + } + + for(int d = 0; d < head_dim; d++) + o[qo_base + d] = (T)o_buf[d]; +} + +template +void flashAttention( + const std::vector& h_q, + const std::vector& h_k, + const std::vector& h_v, + std::vector& h_o, + int batch_size, + int target_seq_len, + int src_seq_len, + int query_heads, + int kv_heads, + int head_dim, + bool is_causal +) +{ + T *d_q, *d_k, *d_v, *d_o; + + size_t q_size = h_q.size() * sizeof(T); + size_t k_size = h_k.size() * sizeof(T); + size_t v_size = h_v.size() * sizeof(T); + size_t o_size = h_o.size() * sizeof(T); + + CHECK_MUSA(musaMalloc(&d_q, q_size)); + CHECK_MUSA(musaMalloc(&d_k, k_size)); + CHECK_MUSA(musaMalloc(&d_v, v_size)); + CHECK_MUSA(musaMalloc(&d_o, o_size)); + + CHECK_MUSA(musaMemcpy(d_q, h_q.data(), q_size, musaMemcpyHostToDevice)); + CHECK_MUSA(musaMemcpy(d_k, h_k.data(), k_size, musaMemcpyHostToDevice)); + CHECK_MUSA(musaMemcpy(d_v, h_v.data(), v_size, musaMemcpyHostToDevice)); + + int total = batch_size * target_seq_len * query_heads; + + dim3 block(256); + dim3 grid((total + block.x - 1) / block.x); + + flashAttentionKernel<<>>( + d_q, d_k, d_v, d_o, + batch_size, target_seq_len, src_seq_len, + query_heads, kv_heads, + head_dim, + is_causal + ); + + CHECK_MUSA(musaGetLastError()); + CHECK_MUSA(musaDeviceSynchronize()); + + CHECK_MUSA(musaMemcpy(h_o.data(), d_o, o_size, musaMemcpyDeviceToHost)); + + CHECK_MUSA(musaFree(d_q)); + CHECK_MUSA(musaFree(d_k)); + CHECK_MUSA(musaFree(d_v)); + CHECK_MUSA(musaFree(d_o)); } // *********************************************************************