From 11bfafd2cbd1802a9630a684d60634778a9a6dea Mon Sep 17 00:00:00 2001
From: zhang200210-neu <wj2my4ttgg@privaterelay.appleid.com>
Date: Fri, 7 Aug 2026 15:18:32 +0800
Subject: [PATCH 1/7] Update kernels.cu

---
 src/kernels.cu | 346 ++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 342 insertions(+), 4 deletions(-)

diff --git a/src/kernels.cu b/src/kernels.cu
index 2cc53e7e..8f43dd36 100644
--- a/src/kernels.cu
+++ b/src/kernels.cu
@@ -1,5 +1,7 @@
 #include <vector>
 #include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cmath>
 
 #include "../tester/utils.h"
 
@@ -21,11 +23,91 @@
  * @param[in] hidden_dim Size of the normalized dimension.
  * @param[in] eps Numerical stability epsilon.
  */
+
+
+// CUDA kernel for RMSNorm computation
+template <typename T>
+__global__ void rmsNormKernel(const T* __restrict__ input,
+                               const T* __restrict__ weight,
+                               T* __restrict__ output,
+                               size_t rows,
+                               size_t hidden_dim,
+                               float eps) {
+    // Each block processes one row
+    size_t row = blockIdx.x;
+    if (row >= rows) return;
+    
+    // Shared memory for reduction
+    __shared__ float shared_sum[256];
+    size_t tid = threadIdx.x;
+    
+    float thread_sum = 0.0f;
+    
+    // Step 1: Compute sum of squares for this row
+    for (size_t i = tid; i < hidden_dim; i += blockDim.x) {
+        size_t idx = row * hidden_dim + i;
+        float val = static_cast<float>(input[idx]);
+        thread_sum += val * val;
+    }
+    
+    // Store partial sum to shared memory
+    shared_sum[tid] = thread_sum;
+    __syncthreads();
+    
+    // Parallel reduction in shared memory
+    for (size_t stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            shared_sum[tid] += shared_sum[tid + stride];
+        }
+        __syncthreads();
+    }
+    
+    // Compute RMS normalization factor
+    float mean_square = shared_sum[0] / hidden_dim;
+    float rms = rsqrtf(mean_square + eps);
+    
+    // Step 2: Apply normalization and scaling
+    for (size_t i = tid; i < hidden_dim; i += blockDim.x) {
+        size_t idx = row * hidden_dim + i;
+        float val = static_cast<float>(input[idx]);
+        float w = static_cast<float>(weight[i]);
+        output[idx] = static_cast<T>(val * rms * w);
+    }
+}
+
 template <typename T>
 void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
               std::vector<T>& h_output, size_t rows, size_t hidden_dim,
               float eps) {
-  // TODO: Implement the rmsNorm function
+    // Allocate device memory
+    T *d_input, *d_weight, *d_output;
+    size_t input_size = rows * hidden_dim * sizeof(T);
+    size_t weight_size = hidden_dim * sizeof(T);
+    
+    cudaMalloc(&d_input, input_size);
+    cudaMalloc(&d_weight, weight_size);
+    cudaMalloc(&d_output, input_size);
+    
+    // Copy data to device
+    cudaMemcpy(d_input, h_input.data(), input_size, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_weight, h_weight.data(), weight_size, cudaMemcpyHostToDevice);
+    
+    // Configure kernel launch parameters
+    int threads_per_block = 256;
+    int num_blocks = rows;
+    
+    // Launch kernel
+    rmsNormKernel<<<num_blocks, threads_per_block>>>(
+        d_input, d_weight, d_output, rows, hidden_dim, eps);
+    
+    // Copy result back to host
+    cudaMemcpy(h_output.data(), d_output, input_size, cudaMemcpyDeviceToHost);
+    
+    // Synchronize and free memory
+    cudaDeviceSynchronize();
+    cudaFree(d_input);
+    cudaFree(d_weight);
+    cudaFree(d_output);
 }
 
 /**
@@ -45,11 +127,267 @@ void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
  * @param[in] is_causal Whether to apply causal masking
  */
 template <typename T>
+void cleanup(T* d_q, T* d_k, T* d_v, T* d_o) {
+    if (d_q) cudaFree(d_q);
+    if (d_k) cudaFree(d_k);
+    if (d_v) cudaFree(d_v);
+    if (d_o) cudaFree(d_o);
+}
+
+// 优化但保持正确性的kernel函数
+
+// float版本的优化kernel
+__global__ void flash_attention_kernel_float(
+    const float* Q, const float* K, const float* V, float* O,
+    int batch_size, int target_seq_len, int src_seq_len,
+    int query_heads, int kv_heads, int head_dim,
+    bool is_causal, float scale) {
+    
+    // 每个线程处理一个输出位置
+    int b = blockIdx.x;
+    int t = blockIdx.y;
+    int h = blockIdx.z;
+    int d = threadIdx.x;
+    
+    if (b >= batch_size || t >= target_seq_len || h >= query_heads || d >= head_dim) return;
+    
+    int kvh = h / (query_heads / kv_heads);
+    int valid_len = is_causal ? min(t + 1, src_seq_len) : src_seq_len;
+    
+    size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
+    size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
+    
+    // 使用共享内存存储查询向量
+    extern __shared__ float flash_attn_shared_float[];
+    float* q_shared = flash_attn_shared_float;
+    
+    // 协作加载查询向量到共享内存
+    for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
+        q_shared[i] = Q[q_base + i];
+    }
+    __syncthreads();
+    
+    // 计算最大分数
+    float max_score = -1e10f;
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        
+        // 使用共享内存中的查询向量
+        #pragma unroll(4)
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * K[k_base + i];
+        }
+        
+        float score = dot * scale;
+        if (score > max_score) max_score = score;
+    }
+    
+    // 计算输出
+    float sum_exp = 0.0f;
+    float output = 0.0f;
+    
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        size_t v_base = kv_base + s * kv_heads * head_dim + d;
+        
+        #pragma unroll(4)
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * K[k_base + i];
+        }
+        
+        float score = dot * scale;
+        float exp_val = expf(score - max_score);
+        
+        sum_exp += exp_val;
+        output += exp_val * V[v_base];
+    }
+    
+    // 归一化
+    if (sum_exp > 1e-12f) {
+        output = output / sum_exp;
+    } else if (valid_len > 0) {
+        output = 0.0f;
+        for (int s = 0; s < valid_len; ++s) {
+            size_t v_base = kv_base + s * kv_heads * head_dim + d;
+            output += V[v_base];
+        }
+        output = output / valid_len;
+    }
+    
+    O[q_base + d] = output;
+}
+
+// half版本的优化kernel
+__global__ void flash_attention_kernel_half(
+    const __half* Q, const __half* K, const __half* V, __half* O,
+    int batch_size, int target_seq_len, int src_seq_len,
+    int query_heads, int kv_heads, int head_dim,
+    bool is_causal, __half scale) {
+    
+    int b = blockIdx.x;
+    int t = blockIdx.y;
+    int h = blockIdx.z;
+    int d = threadIdx.x;
+    
+    if (b >= batch_size || t >= target_seq_len || h >= query_heads || d >= head_dim) return;
+    
+    int kvh = h / (query_heads / kv_heads);
+    int valid_len = is_causal ? min(t + 1, src_seq_len) : src_seq_len;
+    
+    size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
+    size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
+    
+    // 使用共享内存存储查询向量
+    extern __shared__ float flash_attn_shared_half[];
+    float* q_shared = flash_attn_shared_half;
+    
+    float scale_f = __half2float(scale);
+    
+    // 协作加载查询向量到共享内存
+    for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
+        q_shared[i] = __half2float(Q[q_base + i]);
+    }
+    __syncthreads();
+    
+    // 计算最大分数
+    float max_score = -1e4f;
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        
+        #pragma unroll(4)
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * __half2float(K[k_base + i]);
+        }
+        
+        float score = dot * scale_f;
+        if (score > max_score) max_score = score;
+    }
+    
+    // 计算softmax和输出
+    float sum_exp = 0.0f;
+    float output_f = 0.0f;
+    
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        size_t v_base = kv_base + s * kv_heads * head_dim + d;
+        
+        #pragma unroll(4)
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * __half2float(K[k_base + i]);
+        }
+        
+        float score = dot * scale_f;
+        float shifted = score - max_score;
+        
+        // 限制范围确保稳定性
+        if (shifted > 10.0f) shifted = 10.0f;
+        if (shifted < -20.0f) shifted = -20.0f;
+        
+        float exp_val = expf(shifted);
+        sum_exp += exp_val;
+        output_f += exp_val * __half2float(V[v_base]);
+    }
+    
+    // 归一化
+    if (sum_exp > 1e-7f) {
+        output_f = output_f / sum_exp;
+    } else if (valid_len > 0) {
+        output_f = 0.0f;
+        for (int s = 0; s < valid_len; ++s) {
+            size_t v_base = kv_base + s * kv_heads * head_dim + d;
+            output_f += __half2float(V[v_base]);
+        }
+        output_f = output_f / valid_len;
+    }
+    
+    O[q_base + d] = __float2half(output_f);
+}
+template <typename T>
 void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
                     const std::vector<T>& h_v, std::vector<T>& 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
+                    int batch_size, int target_seq_len, int src_seq_len,
+                    int query_heads, int kv_heads, int head_dim, bool is_causal) {
+    
+    // 基本检查
+    if (batch_size <= 0 || target_seq_len <= 0 || src_seq_len <= 0 || 
+        query_heads <= 0 || kv_heads <= 0 || head_dim <= 0) {
+        return;
+    }
+    
+    if (query_heads % kv_heads != 0) return;
+    
+    // 计算大小
+    size_t q_size = batch_size * target_seq_len * query_heads * head_dim;
+    size_t kv_size = batch_size * src_seq_len * kv_heads * head_dim;
+    size_t o_size = batch_size * target_seq_len * query_heads * head_dim;
+    
+    // 分配设备内存
+    T *d_q = nullptr, *d_k = nullptr, *d_v = nullptr, *d_o = nullptr;
+    
+    cudaMalloc(&d_q, q_size * sizeof(T));
+    cudaMalloc(&d_k, kv_size * sizeof(T));
+    cudaMalloc(&d_v, kv_size * sizeof(T));
+    cudaMalloc(&d_o, o_size * sizeof(T));
+    
+    // 拷贝数据
+    cudaMemcpy(d_q, h_q.data(), q_size * sizeof(T), cudaMemcpyHostToDevice);
+    cudaMemcpy(d_k, h_k.data(), kv_size * sizeof(T), cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v, h_v.data(), kv_size * sizeof(T), cudaMemcpyHostToDevice);
+    
+    // 缩放因子
+    T scale;
+    if constexpr (std::is_same_v<T, __half>) {
+        float head_dim_f = static_cast<float>(head_dim);
+        float scale_f = 1.0f / sqrtf(head_dim_f);
+        if (scale_f > 5.0f) scale_f = 5.0f;
+        scale = __float2half(scale_f);
+    } else {
+        scale = T(1.0 / sqrt(static_cast<double>(head_dim)));
+    }
+    
+    // 启动对应的kernel
+    dim3 grid(batch_size, target_seq_len, query_heads);
+    
+    // 优化block大小：确保是32的倍数（warp大小）
+    int block_size = 256;
+    if (head_dim < 256) {
+        block_size = ((head_dim + 31) / 32) * 32; // 向上取整到32的倍数
+    }
+    
+    // 计算共享内存大小
+    size_t shared_mem_size = head_dim * sizeof(float);
+    
+    if constexpr (std::is_same_v<T, float>) {
+        flash_attention_kernel_float<<<grid, block_size, shared_mem_size>>>(
+            d_q, d_k, d_v, d_o,
+            batch_size, target_seq_len, src_seq_len,
+            query_heads, kv_heads, head_dim,
+            is_causal, scale);
+    } else {
+        flash_attention_kernel_half<<<grid, block_size, shared_mem_size>>>(
+            d_q, d_k, d_v, d_o,
+            batch_size, target_seq_len, src_seq_len,
+            query_heads, kv_heads, head_dim,
+            is_causal, scale);
+    }
+    
+    // 同步和错误检查
+    cudaDeviceSynchronize();
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        std::cerr << "CUDA error: " << cudaGetErrorString(err) << std::endl;
+    }
+    
+    // 拷贝结果
+    h_o.resize(o_size);
+    cudaMemcpy(h_o.data(), d_o, o_size * sizeof(T), cudaMemcpyDeviceToHost);
+    
+    // 清理
+    cleanup(d_q, d_k, d_v, d_o);
 }
 
 // *********************************************************************

From 4e5f9b0576cbc29eb515d89913358bf1ece8411d Mon Sep 17 00:00:00 2001
From: zhang200210-neu <wj2my4ttgg@privaterelay.appleid.com>
Date: Mon, 10 Aug 2026 19:37:52 +0800
Subject: [PATCH 2/7] Update kernels.mu

---
 src/kernels.mu | 360 +++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 316 insertions(+), 44 deletions(-)

diff --git a/src/kernels.mu b/src/kernels.mu
index 1ce371eb..e799dd81 100644
--- a/src/kernels.mu
+++ b/src/kernels.mu
@@ -1,61 +1,333 @@
 #include <vector>
 #include <musa_fp16.h>
+#include <musa_runtime.h>
+#include <cmath>
+#include <iostream>
+#include <algorithm>
+#include <type_traits>
 
 #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.
- */
+#define MUSA_CHECK(call) \
+    do { \
+        musaError_t err = call; \
+        if (err != musaSuccess) { \
+            std::cerr << "MUSA error at " << __FILE__ << ":" << __LINE__ \
+                      << " - " << musaGetErrorString(err) << std::endl; \
+            exit(EXIT_FAILURE); \
+        } \
+    } while(0)
+
+// ==================== RMSNorm ====================
+template <typename T>
+__global__ void rmsNormKernel(const T* __restrict__ input,
+                               const T* __restrict__ weight,
+                               T* __restrict__ output,
+                               size_t rows, size_t hidden_dim, float eps) {
+    size_t row = blockIdx.x;
+    if (row >= rows) return;
+    
+    __shared__ float shared_sum[256];
+    size_t tid = threadIdx.x;
+    float thread_sum = 0.0f;
+    
+    for (size_t i = tid; i < hidden_dim; i += blockDim.x) {
+        size_t idx = row * hidden_dim + i;
+        float val = static_cast<float>(input[idx]);
+        thread_sum += val * val;
+    }
+    
+    shared_sum[tid] = thread_sum;
+    __syncthreads();
+    
+    for (size_t stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            shared_sum[tid] += shared_sum[tid + stride];
+        }
+        __syncthreads();
+    }
+    
+    float mean_square = shared_sum[0] / hidden_dim;
+    float rms = rsqrtf(mean_square + eps);
+    
+    for (size_t i = tid; i < hidden_dim; i += blockDim.x) {
+        size_t idx = row * hidden_dim + i;
+        float val = static_cast<float>(input[idx]);
+        float w = static_cast<float>(weight[i]);
+        output[idx] = static_cast<T>(val * rms * w);
+    }
+}
+
 template <typename T>
 void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
-              std::vector<T>& h_output, size_t rows, size_t hidden_dim,
-              float eps) {
-  // TODO: Implement the rmsNorm function
+              std::vector<T>& h_output, size_t rows, size_t hidden_dim, float eps) {
+    T *d_input, *d_weight, *d_output;
+    size_t input_size = rows * hidden_dim * sizeof(T);
+    size_t weight_size = hidden_dim * sizeof(T);
+    
+    MUSA_CHECK(musaMalloc(&d_input, input_size));
+    MUSA_CHECK(musaMalloc(&d_weight, weight_size));
+    MUSA_CHECK(musaMalloc(&d_output, input_size));
+    
+    MUSA_CHECK(musaMemcpy(d_input, h_input.data(), input_size, musaMemcpyHostToDevice));
+    MUSA_CHECK(musaMemcpy(d_weight, h_weight.data(), weight_size, musaMemcpyHostToDevice));
+    
+    int threads_per_block = 256;
+    int num_blocks = rows;
+    rmsNormKernel<<<num_blocks, threads_per_block>>>(d_input, d_weight, d_output, rows, hidden_dim, eps);
+    
+    MUSA_CHECK(musaMemcpy(h_output.data(), d_output, input_size, musaMemcpyDeviceToHost));
+    MUSA_CHECK(musaDeviceSynchronize());
+    MUSA_CHECK(musaFree(d_input));
+    MUSA_CHECK(musaFree(d_weight));
+    MUSA_CHECK(musaFree(d_output));
+}
+
+// ==================== Flash Attention ====================
+template <typename T>
+void cleanup(T* d_q, T* d_k, T* d_v, T* d_o) {
+    if (d_q) musaFree(d_q);
+    if (d_k) musaFree(d_k);
+    if (d_v) musaFree(d_v);
+    if (d_o) musaFree(d_o);
+}
+
+// 升级后的 float kernel（双精度累加）
+__global__ void flash_attention_kernel_float(
+    const float* Q, const float* K, const float* V, float* O,
+    int batch_size, int target_seq_len, int src_seq_len,
+    int query_heads, int kv_heads, int head_dim,
+    bool is_causal, float scale) {
+    
+    int b = blockIdx.x;
+    int t = blockIdx.y;
+    int h = blockIdx.z;
+    int d = threadIdx.x;
+    
+    if (b >= batch_size || t >= target_seq_len || h >= query_heads || d >= head_dim) return;
+    
+    int kvh = h / (query_heads / kv_heads);
+    int valid_len = is_causal ? min(t + 1, src_seq_len) : src_seq_len;
+    
+    size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
+    size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
+    
+    extern __shared__ float flash_attn_shared_float[];
+    float* q_shared = flash_attn_shared_float;
+    
+    for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
+        q_shared[i] = Q[q_base + i];
+    }
+    __syncthreads();
+    
+    float max_score = -1e30f;
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        
+        #pragma unroll 4
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * K[k_base + i];
+        }
+        float score = dot * scale;
+        if (score > max_score) max_score = score;
+    }
+    
+    double sum_exp = 0.0;
+    double output = 0.0;
+    
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        size_t v_base = kv_base + s * kv_heads * head_dim + d;
+        
+        #pragma unroll 4
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * K[k_base + i];
+        }
+        
+        float score = dot * scale;
+        float shifted = score - max_score;
+        if (shifted > 10.0f) shifted = 10.0f;
+        if (shifted < -20.0f) shifted = -20.0f;
+        
+        double exp_val = exp((double)shifted);
+        sum_exp += exp_val;
+        output += exp_val * (double)V[v_base];
+    }
+    
+    if (sum_exp > 1e-12) {
+        output = output / sum_exp;
+    } else if (valid_len > 0) {
+        output = 0.0;
+        for (int s = 0; s < valid_len; ++s) {
+            size_t v_base = kv_base + s * kv_heads * head_dim + d;
+            output += (double)V[v_base];
+        }
+        output = output / valid_len;
+    }
+    
+    O[q_base + d] = (float)output;
+}
+
+// half kernel（保持不变）
+__global__ void flash_attention_kernel_half(
+    const __half* Q, const __half* K, const __half* V, __half* O,
+    int batch_size, int target_seq_len, int src_seq_len,
+    int query_heads, int kv_heads, int head_dim,
+    bool is_causal, __half scale) {
+    
+    int b = blockIdx.x;
+    int t = blockIdx.y;
+    int h = blockIdx.z;
+    int d = threadIdx.x;
+    
+    if (b >= batch_size || t >= target_seq_len || h >= query_heads || d >= head_dim) return;
+    
+    int kvh = h / (query_heads / kv_heads);
+    int valid_len = is_causal ? min(t + 1, src_seq_len) : src_seq_len;
+    
+    size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
+    size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
+    
+    extern __shared__ float flash_attn_shared_half[];
+    float* q_shared = flash_attn_shared_half;
+    
+    float scale_f = __half2float(scale);
+    
+    for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
+        q_shared[i] = __half2float(Q[q_base + i]);
+    }
+    __syncthreads();
+    
+    float max_score = -1e4f;
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        
+        #pragma unroll 4
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * __half2float(K[k_base + i]);
+        }
+        float score = dot * scale_f;
+        if (score > max_score) max_score = score;
+    }
+    
+    float sum_exp = 0.0f;
+    float output_f = 0.0f;
+    
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        size_t v_base = kv_base + s * kv_heads * head_dim + d;
+        
+        #pragma unroll 4
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * __half2float(K[k_base + i]);
+        }
+        
+        float score = dot * scale_f;
+        float shifted = score - max_score;
+        if (shifted > 10.0f) shifted = 10.0f;
+        if (shifted < -20.0f) shifted = -20.0f;
+        
+        float exp_val = expf(shifted);
+        sum_exp += exp_val;
+        output_f += exp_val * __half2float(V[v_base]);
+    }
+    
+    if (sum_exp > 1e-7f) {
+        output_f = output_f / sum_exp;
+    } else if (valid_len > 0) {
+        output_f = 0.0f;
+        for (int s = 0; s < valid_len; ++s) {
+            size_t v_base = kv_base + s * kv_heads * head_dim + d;
+            output_f += __half2float(V[v_base]);
+        }
+        output_f = output_f / valid_len;
+    }
+    
+    O[q_base + d] = __float2half(output_f);
+}
+
+// 重载的启动函数
+inline void launch_flash_attention(
+    const float* d_q, const float* d_k, const float* d_v, float* d_o,
+    int batch_size, int target_seq_len, int src_seq_len,
+    int query_heads, int kv_heads, int head_dim, bool is_causal) {
+    
+    float scale = 1.0f / sqrtf(static_cast<float>(head_dim));
+    dim3 grid(batch_size, target_seq_len, query_heads);
+    int block_size = 256;
+    if (head_dim < 256) block_size = ((head_dim + 31) / 32) * 32;
+    size_t shared_mem_size = head_dim * sizeof(float);
+    
+    flash_attention_kernel_float<<<grid, block_size, shared_mem_size>>>(
+        d_q, d_k, d_v, d_o,
+        batch_size, target_seq_len, src_seq_len,
+        query_heads, kv_heads, head_dim,
+        is_causal, scale);
+}
+
+inline void launch_flash_attention(
+    const __half* d_q, const __half* d_k, const __half* d_v, __half* d_o,
+    int batch_size, int target_seq_len, int src_seq_len,
+    int query_heads, int kv_heads, int head_dim, bool is_causal) {
+    
+    float scale_f = 1.0f / sqrtf(static_cast<float>(head_dim));
+    if (scale_f > 5.0f) scale_f = 5.0f;
+    __half scale = __float2half(scale_f);
+    
+    dim3 grid(batch_size, target_seq_len, query_heads);
+    int block_size = 256;
+    if (head_dim < 256) block_size = ((head_dim + 31) / 32) * 32;
+    size_t shared_mem_size = head_dim * sizeof(float);
+    
+    flash_attention_kernel_half<<<grid, block_size, shared_mem_size>>>(
+        d_q, d_k, d_v, d_o,
+        batch_size, target_seq_len, src_seq_len,
+        query_heads, kv_heads, head_dim,
+        is_causal, scale);
 }
 
-/**
- * @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 <typename T>
 void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
                     const std::vector<T>& h_v, std::vector<T>& 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
+                    int batch_size, int target_seq_len, int src_seq_len,
+                    int query_heads, int kv_heads, int head_dim, bool is_causal) {
+    
+    if (batch_size <= 0 || target_seq_len <= 0 || src_seq_len <= 0 || 
+        query_heads <= 0 || kv_heads <= 0 || head_dim <= 0) return;
+    if (query_heads % kv_heads != 0) return;
+    
+    size_t q_size = batch_size * target_seq_len * query_heads * head_dim;
+    size_t kv_size = batch_size * src_seq_len * kv_heads * head_dim;
+    size_t o_size = batch_size * target_seq_len * query_heads * head_dim;
+    
+    T *d_q = nullptr, *d_k = nullptr, *d_v = nullptr, *d_o = nullptr;
+    MUSA_CHECK(musaMalloc(&d_q, q_size * sizeof(T)));
+    MUSA_CHECK(musaMalloc(&d_k, kv_size * sizeof(T)));
+    MUSA_CHECK(musaMalloc(&d_v, kv_size * sizeof(T)));
+    MUSA_CHECK(musaMalloc(&d_o, o_size * sizeof(T)));
+    
+    MUSA_CHECK(musaMemcpy(d_q, h_q.data(), q_size * sizeof(T), musaMemcpyHostToDevice));
+    MUSA_CHECK(musaMemcpy(d_k, h_k.data(), kv_size * sizeof(T), musaMemcpyHostToDevice));
+    MUSA_CHECK(musaMemcpy(d_v, h_v.data(), kv_size * sizeof(T), musaMemcpyHostToDevice));
+    
+    launch_flash_attention(d_q, d_k, d_v, d_o,
+                           batch_size, target_seq_len, src_seq_len,
+                           query_heads, kv_heads, head_dim, is_causal);
+    
+    MUSA_CHECK(musaDeviceSynchronize());
+    MUSA_CHECK(musaGetLastError());
+    
+    h_o.resize(o_size);
+    MUSA_CHECK(musaMemcpy(h_o.data(), d_o, o_size * sizeof(T), musaMemcpyDeviceToHost));
+    
+    cleanup(d_q, d_k, d_v, d_o);
 }
 
-// *********************************************************************
-// Explicit Template Instantiations (REQUIRED FOR LINKING WITH TESTER.O)
-// DO NOT MODIFY THIS SECTION
-// *********************************************************************
+// ==================== 模板显式实例化 ====================
 template void rmsNorm<float>(const std::vector<float>&, const std::vector<float>&,
   std::vector<float>&, size_t, size_t, float);
 template void rmsNorm<half>(const std::vector<half>&, const std::vector<half>&,

From c2e890f2258bcf0008d92c5a2d99d49910523410 Mon Sep 17 00:00:00 2001
From: zhang200210-neu <wj2my4ttgg@privaterelay.appleid.com>
Date: Mon, 10 Aug 2026 23:31:07 +0800
Subject: [PATCH 3/7] Update kernels.maca

---
 src/kernels.maca | 348 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 343 insertions(+), 5 deletions(-)

diff --git a/src/kernels.maca b/src/kernels.maca
index 4c320f21..d48ddc6d 100644
--- a/src/kernels.maca
+++ b/src/kernels.maca
@@ -1,5 +1,7 @@
 #include <vector>
-#include <common/maca_fp16.h>
+#include <mc_fp16.h>
+#include <mc_runtime.h>
+#include <cmath>
 
 #include "../tester/utils.h"
 
@@ -21,11 +23,91 @@
  * @param[in] hidden_dim Size of the normalized dimension.
  * @param[in] eps Numerical stability epsilon.
  */
+
+
+// CUDA kernel for RMSNorm computation
+template <typename T>
+__global__ void rmsNormKernel(const T* __restrict__ input,
+                               const T* __restrict__ weight,
+                               T* __restrict__ output,
+                               size_t rows,
+                               size_t hidden_dim,
+                               float eps) {
+    // Each block processes one row
+    size_t row = blockIdx.x;
+    if (row >= rows) return;
+    
+    // Shared memory for reduction
+    __shared__ float shared_sum[256];
+    size_t tid = threadIdx.x;
+    
+    float thread_sum = 0.0f;
+    
+    // Step 1: Compute sum of squares for this row
+    for (size_t i = tid; i < hidden_dim; i += blockDim.x) {
+        size_t idx = row * hidden_dim + i;
+        float val = static_cast<float>(input[idx]);
+        thread_sum += val * val;
+    }
+    
+    // Store partial sum to shared memory
+    shared_sum[tid] = thread_sum;
+    __syncthreads();
+    
+    // Parallel reduction in shared memory
+    for (size_t stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            shared_sum[tid] += shared_sum[tid + stride];
+        }
+        __syncthreads();
+    }
+    
+    // Compute RMS normalization factor
+    float mean_square = shared_sum[0] / hidden_dim;
+    float rms = rsqrtf(mean_square + eps);
+    
+    // Step 2: Apply normalization and scaling
+    for (size_t i = tid; i < hidden_dim; i += blockDim.x) {
+        size_t idx = row * hidden_dim + i;
+        float val = static_cast<float>(input[idx]);
+        float w = static_cast<float>(weight[i]);
+        output[idx] = static_cast<T>(val * rms * w);
+    }
+}
+
 template <typename T>
 void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
               std::vector<T>& h_output, size_t rows, size_t hidden_dim,
               float eps) {
-  // TODO: Implement the rmsNorm function
+    // Allocate device memory
+    T *d_input, *d_weight, *d_output;
+    size_t input_size = rows * hidden_dim * sizeof(T);
+    size_t weight_size = hidden_dim * sizeof(T);
+    
+    mcMalloc(&d_input, input_size);
+    mcMalloc(&d_weight, weight_size);
+    mcMalloc(&d_output, input_size);
+    
+    // Copy data to device
+    mcMemcpy(d_input, h_input.data(), input_size, mcMemcpyHostToDevice);
+    mcMemcpy(d_weight, h_weight.data(), weight_size, mcMemcpyHostToDevice);
+    
+    // Configure kernel launch parameters
+    int threads_per_block = 256;
+    int num_blocks = rows;
+    
+    // Launch kernel
+    rmsNormKernel<<<num_blocks, threads_per_block>>>(
+        d_input, d_weight, d_output, rows, hidden_dim, eps);
+    
+    // Copy result back to host
+    mcMemcpy(h_output.data(), d_output, input_size, mcMemcpyDeviceToHost);
+    
+    // Synchronize and free memory
+    mcDeviceSynchronize();
+    mcFree(d_input);
+    mcFree(d_weight);
+    mcFree(d_output);
 }
 
 /**
@@ -45,11 +127,267 @@ void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
  * @param[in] is_causal Whether to apply causal masking
  */
 template <typename T>
+void cleanup(T* d_q, T* d_k, T* d_v, T* d_o) {
+    if (d_q) mcFree(d_q);
+    if (d_k) mcFree(d_k);
+    if (d_v) mcFree(d_v);
+    if (d_o) mcFree(d_o);
+}
+
+// 优化但保持正确性的kernel函数
+
+// float版本的优化kernel
+__global__ void flash_attention_kernel_float(
+    const float* Q, const float* K, const float* V, float* O,
+    int batch_size, int target_seq_len, int src_seq_len,
+    int query_heads, int kv_heads, int head_dim,
+    bool is_causal, float scale) {
+    
+    // 每个线程处理一个输出位置
+    int b = blockIdx.x;
+    int t = blockIdx.y;
+    int h = blockIdx.z;
+    int d = threadIdx.x;
+    
+    if (b >= batch_size || t >= target_seq_len || h >= query_heads || d >= head_dim) return;
+    
+    int kvh = h / (query_heads / kv_heads);
+    int valid_len = is_causal ? min(t + 1, src_seq_len) : src_seq_len;
+    
+    size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
+    size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
+    
+    // 使用共享内存存储查询向量
+    extern __shared__ float flash_attn_shared_float[];
+    float* q_shared = flash_attn_shared_float;
+    
+    // 协作加载查询向量到共享内存
+    for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
+        q_shared[i] = Q[q_base + i];
+    }
+    __syncthreads();
+    
+    // 计算最大分数
+    float max_score = -1e10f;
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        
+        // 使用共享内存中的查询向量
+        #pragma unroll(4)
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * K[k_base + i];
+        }
+        
+        float score = dot * scale;
+        if (score > max_score) max_score = score;
+    }
+    
+    // 计算输出
+    float sum_exp = 0.0f;
+    float output = 0.0f;
+    
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        size_t v_base = kv_base + s * kv_heads * head_dim + d;
+        
+        #pragma unroll(4)
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * K[k_base + i];
+        }
+        
+        float score = dot * scale;
+        float exp_val = expf(score - max_score);
+        
+        sum_exp += exp_val;
+        output += exp_val * V[v_base];
+    }
+    
+    // 归一化
+    if (sum_exp > 1e-12f) {
+        output = output / sum_exp;
+    } else if (valid_len > 0) {
+        output = 0.0f;
+        for (int s = 0; s < valid_len; ++s) {
+            size_t v_base = kv_base + s * kv_heads * head_dim + d;
+            output += V[v_base];
+        }
+        output = output / valid_len;
+    }
+    
+    O[q_base + d] = output;
+}
+
+// half版本的优化kernel
+__global__ void flash_attention_kernel_half(
+    const __half* Q, const __half* K, const __half* V, __half* O,
+    int batch_size, int target_seq_len, int src_seq_len,
+    int query_heads, int kv_heads, int head_dim,
+    bool is_causal, __half scale) {
+    
+    int b = blockIdx.x;
+    int t = blockIdx.y;
+    int h = blockIdx.z;
+    int d = threadIdx.x;
+    
+    if (b >= batch_size || t >= target_seq_len || h >= query_heads || d >= head_dim) return;
+    
+    int kvh = h / (query_heads / kv_heads);
+    int valid_len = is_causal ? min(t + 1, src_seq_len) : src_seq_len;
+    
+    size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
+    size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
+    
+    // 使用共享内存存储查询向量
+    extern __shared__ float flash_attn_shared_half[];
+    float* q_shared = flash_attn_shared_half;
+    
+    float scale_f = __half2float(scale);
+    
+    // 协作加载查询向量到共享内存
+    for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
+        q_shared[i] = __half2float(Q[q_base + i]);
+    }
+    __syncthreads();
+    
+    // 计算最大分数
+    float max_score = -1e4f;
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        
+        #pragma unroll(4)
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * __half2float(K[k_base + i]);
+        }
+        
+        float score = dot * scale_f;
+        if (score > max_score) max_score = score;
+    }
+    
+    // 计算softmax和输出
+    float sum_exp = 0.0f;
+    float output_f = 0.0f;
+    
+    for (int s = 0; s < valid_len; ++s) {
+        float dot = 0.0f;
+        size_t k_base = kv_base + s * kv_heads * head_dim;
+        size_t v_base = kv_base + s * kv_heads * head_dim + d;
+        
+        #pragma unroll(4)
+        for (int i = 0; i < head_dim; ++i) {
+            dot += q_shared[i] * __half2float(K[k_base + i]);
+        }
+        
+        float score = dot * scale_f;
+        float shifted = score - max_score;
+        
+        // 限制范围确保稳定性
+        if (shifted > 10.0f) shifted = 10.0f;
+        if (shifted < -20.0f) shifted = -20.0f;
+        
+        float exp_val = expf(shifted);
+        sum_exp += exp_val;
+        output_f += exp_val * __half2float(V[v_base]);
+    }
+    
+    // 归一化
+    if (sum_exp > 1e-7f) {
+        output_f = output_f / sum_exp;
+    } else if (valid_len > 0) {
+        output_f = 0.0f;
+        for (int s = 0; s < valid_len; ++s) {
+            size_t v_base = kv_base + s * kv_heads * head_dim + d;
+            output_f += __half2float(V[v_base]);
+        }
+        output_f = output_f / valid_len;
+    }
+    
+    O[q_base + d] = __float2half(output_f);
+}
+template <typename T>
 void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
                     const std::vector<T>& h_v, std::vector<T>& 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
+                    int batch_size, int target_seq_len, int src_seq_len,
+                    int query_heads, int kv_heads, int head_dim, bool is_causal) {
+    
+    // 基本检查
+    if (batch_size <= 0 || target_seq_len <= 0 || src_seq_len <= 0 || 
+        query_heads <= 0 || kv_heads <= 0 || head_dim <= 0) {
+        return;
+    }
+    
+    if (query_heads % kv_heads != 0) return;
+    
+    // 计算大小
+    size_t q_size = batch_size * target_seq_len * query_heads * head_dim;
+    size_t kv_size = batch_size * src_seq_len * kv_heads * head_dim;
+    size_t o_size = batch_size * target_seq_len * query_heads * head_dim;
+    
+    // 分配设备内存
+    T *d_q = nullptr, *d_k = nullptr, *d_v = nullptr, *d_o = nullptr;
+    
+    mcMalloc(&d_q, q_size * sizeof(T));
+    mcMalloc(&d_k, kv_size * sizeof(T));
+    mcMalloc(&d_v, kv_size * sizeof(T));
+    mcMalloc(&d_o, o_size * sizeof(T));
+    
+    // 拷贝数据
+    mcMemcpy(d_q, h_q.data(), q_size * sizeof(T), mcMemcpyHostToDevice);
+    mcMemcpy(d_k, h_k.data(), kv_size * sizeof(T), mcMemcpyHostToDevice);
+    mcMemcpy(d_v, h_v.data(), kv_size * sizeof(T), mcMemcpyHostToDevice);
+    
+    // 缩放因子
+    T scale;
+    if constexpr (std::is_same_v<T, __half>) {
+        float head_dim_f = static_cast<float>(head_dim);
+        float scale_f = 1.0f / sqrtf(head_dim_f);
+        if (scale_f > 5.0f) scale_f = 5.0f;
+        scale = __float2half(scale_f);
+    } else {
+        scale = T(1.0 / sqrt(static_cast<double>(head_dim)));
+    }
+    
+    // 启动对应的kernel
+    dim3 grid(batch_size, target_seq_len, query_heads);
+    
+    // 优化block大小：确保是32的倍数（warp大小）
+    int block_size = 256;
+    if (head_dim < 256) {
+        block_size = ((head_dim + 31) / 32) * 32; // 向上取整到32的倍数
+    }
+    
+    // 计算共享内存大小
+    size_t shared_mem_size = head_dim * sizeof(float);
+    
+    if constexpr (std::is_same_v<T, float>) {
+        flash_attention_kernel_float<<<grid, block_size, shared_mem_size>>>(
+            d_q, d_k, d_v, d_o,
+            batch_size, target_seq_len, src_seq_len,
+            query_heads, kv_heads, head_dim,
+            is_causal, scale);
+    } else {
+        flash_attention_kernel_half<<<grid, block_size, shared_mem_size>>>(
+            d_q, d_k, d_v, d_o,
+            batch_size, target_seq_len, src_seq_len,
+            query_heads, kv_heads, head_dim,
+            is_causal, scale);
+    }
+    
+    // 同步和错误检查
+    mcDeviceSynchronize();
+    mcError_t err = mcGetLastError();
+    if (err != mcSuccess) {
+        std::cerr << "MC error: " << mcGetErrorString(err) << std::endl;
+    }
+    
+    // 拷贝结果
+    h_o.resize(o_size);
+    mcMemcpy(h_o.data(), d_o, o_size * sizeof(T), mcMemcpyDeviceToHost);
+    
+    // 清理
+    cleanup(d_q, d_k, d_v, d_o);
 }
 
 // *********************************************************************

From 2f93bf9145bdab7285cfc362c890dd889526b356 Mon Sep 17 00:00:00 2001
From: zhang200210-neu <wj2my4ttgg@privaterelay.appleid.com>
Date: Mon, 10 Aug 2026 23:59:58 +0800
Subject: [PATCH 4/7] Update kernels.maca

---
 src/kernels.maca | 96 +++++++++---------------------------------------
 1 file changed, 17 insertions(+), 79 deletions(-)

diff --git a/src/kernels.maca b/src/kernels.maca
index d48ddc6d..93016502 100644
--- a/src/kernels.maca
+++ b/src/kernels.maca
@@ -1,31 +1,19 @@
 #include <vector>
-#include <mc_fp16.h>
-#include <mc_runtime.h>
+#include <maca.h>
+#include <maca_fp16.h>
 #include <cmath>
 
+// 若 half 类型未自动定义，手动提供别名
+#ifndef half
+using half = __half;
+#endif
+
 #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.
+ * ...
  */
-
-
-// CUDA kernel for RMSNorm computation
 template <typename T>
 __global__ void rmsNormKernel(const T* __restrict__ input,
                                const T* __restrict__ weight,
@@ -33,28 +21,23 @@ __global__ void rmsNormKernel(const T* __restrict__ input,
                                size_t rows,
                                size_t hidden_dim,
                                float eps) {
-    // Each block processes one row
     size_t row = blockIdx.x;
     if (row >= rows) return;
     
-    // Shared memory for reduction
     __shared__ float shared_sum[256];
     size_t tid = threadIdx.x;
     
     float thread_sum = 0.0f;
     
-    // Step 1: Compute sum of squares for this row
     for (size_t i = tid; i < hidden_dim; i += blockDim.x) {
         size_t idx = row * hidden_dim + i;
         float val = static_cast<float>(input[idx]);
         thread_sum += val * val;
     }
     
-    // Store partial sum to shared memory
     shared_sum[tid] = thread_sum;
     __syncthreads();
     
-    // Parallel reduction in shared memory
     for (size_t stride = blockDim.x / 2; stride > 0; stride >>= 1) {
         if (tid < stride) {
             shared_sum[tid] += shared_sum[tid + stride];
@@ -62,11 +45,9 @@ __global__ void rmsNormKernel(const T* __restrict__ input,
         __syncthreads();
     }
     
-    // Compute RMS normalization factor
     float mean_square = shared_sum[0] / hidden_dim;
     float rms = rsqrtf(mean_square + eps);
     
-    // Step 2: Apply normalization and scaling
     for (size_t i = tid; i < hidden_dim; i += blockDim.x) {
         size_t idx = row * hidden_dim + i;
         float val = static_cast<float>(input[idx]);
@@ -79,7 +60,6 @@ template <typename T>
 void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
               std::vector<T>& h_output, size_t rows, size_t hidden_dim,
               float eps) {
-    // Allocate device memory
     T *d_input, *d_weight, *d_output;
     size_t input_size = rows * hidden_dim * sizeof(T);
     size_t weight_size = hidden_dim * sizeof(T);
@@ -88,22 +68,17 @@ void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
     mcMalloc(&d_weight, weight_size);
     mcMalloc(&d_output, input_size);
     
-    // Copy data to device
     mcMemcpy(d_input, h_input.data(), input_size, mcMemcpyHostToDevice);
     mcMemcpy(d_weight, h_weight.data(), weight_size, mcMemcpyHostToDevice);
     
-    // Configure kernel launch parameters
     int threads_per_block = 256;
     int num_blocks = rows;
     
-    // Launch kernel
     rmsNormKernel<<<num_blocks, threads_per_block>>>(
         d_input, d_weight, d_output, rows, hidden_dim, eps);
     
-    // Copy result back to host
     mcMemcpy(h_output.data(), d_output, input_size, mcMemcpyDeviceToHost);
     
-    // Synchronize and free memory
     mcDeviceSynchronize();
     mcFree(d_input);
     mcFree(d_weight);
@@ -112,19 +87,7 @@ void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
 
 /**
  * @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 <typename T>
 void cleanup(T* d_q, T* d_k, T* d_v, T* d_o) {
@@ -134,16 +97,13 @@ void cleanup(T* d_q, T* d_k, T* d_v, T* d_o) {
     if (d_o) mcFree(d_o);
 }
 
-// 优化但保持正确性的kernel函数
-
-// float版本的优化kernel
+// float flash attention kernel
 __global__ void flash_attention_kernel_float(
     const float* Q, const float* K, const float* V, float* O,
     int batch_size, int target_seq_len, int src_seq_len,
     int query_heads, int kv_heads, int head_dim,
     bool is_causal, float scale) {
     
-    // 每个线程处理一个输出位置
     int b = blockIdx.x;
     int t = blockIdx.y;
     int h = blockIdx.z;
@@ -157,24 +117,20 @@ __global__ void flash_attention_kernel_float(
     size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
     size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
     
-    // 使用共享内存存储查询向量
     extern __shared__ float flash_attn_shared_float[];
     float* q_shared = flash_attn_shared_float;
     
-    // 协作加载查询向量到共享内存
     for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
         q_shared[i] = Q[q_base + i];
     }
     __syncthreads();
     
-    // 计算最大分数
     float max_score = -1e10f;
     for (int s = 0; s < valid_len; ++s) {
         float dot = 0.0f;
         size_t k_base = kv_base + s * kv_heads * head_dim;
         
-        // 使用共享内存中的查询向量
-        #pragma unroll(4)
+        #pragma unroll 4
         for (int i = 0; i < head_dim; ++i) {
             dot += q_shared[i] * K[k_base + i];
         }
@@ -183,7 +139,6 @@ __global__ void flash_attention_kernel_float(
         if (score > max_score) max_score = score;
     }
     
-    // 计算输出
     float sum_exp = 0.0f;
     float output = 0.0f;
     
@@ -192,7 +147,7 @@ __global__ void flash_attention_kernel_float(
         size_t k_base = kv_base + s * kv_heads * head_dim;
         size_t v_base = kv_base + s * kv_heads * head_dim + d;
         
-        #pragma unroll(4)
+        #pragma unroll 4
         for (int i = 0; i < head_dim; ++i) {
             dot += q_shared[i] * K[k_base + i];
         }
@@ -204,7 +159,6 @@ __global__ void flash_attention_kernel_float(
         output += exp_val * V[v_base];
     }
     
-    // 归一化
     if (sum_exp > 1e-12f) {
         output = output / sum_exp;
     } else if (valid_len > 0) {
@@ -219,7 +173,7 @@ __global__ void flash_attention_kernel_float(
     O[q_base + d] = output;
 }
 
-// half版本的优化kernel
+// half flash attention kernel
 __global__ void flash_attention_kernel_half(
     const __half* Q, const __half* K, const __half* V, __half* O,
     int batch_size, int target_seq_len, int src_seq_len,
@@ -239,25 +193,22 @@ __global__ void flash_attention_kernel_half(
     size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
     size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
     
-    // 使用共享内存存储查询向量
     extern __shared__ float flash_attn_shared_half[];
     float* q_shared = flash_attn_shared_half;
     
     float scale_f = __half2float(scale);
     
-    // 协作加载查询向量到共享内存
     for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
         q_shared[i] = __half2float(Q[q_base + i]);
     }
     __syncthreads();
     
-    // 计算最大分数
     float max_score = -1e4f;
     for (int s = 0; s < valid_len; ++s) {
         float dot = 0.0f;
         size_t k_base = kv_base + s * kv_heads * head_dim;
         
-        #pragma unroll(4)
+        #pragma unroll 4
         for (int i = 0; i < head_dim; ++i) {
             dot += q_shared[i] * __half2float(K[k_base + i]);
         }
@@ -266,7 +217,6 @@ __global__ void flash_attention_kernel_half(
         if (score > max_score) max_score = score;
     }
     
-    // 计算softmax和输出
     float sum_exp = 0.0f;
     float output_f = 0.0f;
     
@@ -275,7 +225,7 @@ __global__ void flash_attention_kernel_half(
         size_t k_base = kv_base + s * kv_heads * head_dim;
         size_t v_base = kv_base + s * kv_heads * head_dim + d;
         
-        #pragma unroll(4)
+        #pragma unroll 4
         for (int i = 0; i < head_dim; ++i) {
             dot += q_shared[i] * __half2float(K[k_base + i]);
         }
@@ -283,7 +233,6 @@ __global__ void flash_attention_kernel_half(
         float score = dot * scale_f;
         float shifted = score - max_score;
         
-        // 限制范围确保稳定性
         if (shifted > 10.0f) shifted = 10.0f;
         if (shifted < -20.0f) shifted = -20.0f;
         
@@ -292,7 +241,6 @@ __global__ void flash_attention_kernel_half(
         output_f += exp_val * __half2float(V[v_base]);
     }
     
-    // 归一化
     if (sum_exp > 1e-7f) {
         output_f = output_f / sum_exp;
     } else if (valid_len > 0) {
@@ -306,13 +254,13 @@ __global__ void flash_attention_kernel_half(
     
     O[q_base + d] = __float2half(output_f);
 }
+
 template <typename T>
 void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
                     const std::vector<T>& h_v, std::vector<T>& 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) {
     
-    // 基本检查
     if (batch_size <= 0 || target_seq_len <= 0 || src_seq_len <= 0 || 
         query_heads <= 0 || kv_heads <= 0 || head_dim <= 0) {
         return;
@@ -320,12 +268,10 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
     
     if (query_heads % kv_heads != 0) return;
     
-    // 计算大小
     size_t q_size = batch_size * target_seq_len * query_heads * head_dim;
     size_t kv_size = batch_size * src_seq_len * kv_heads * head_dim;
     size_t o_size = batch_size * target_seq_len * query_heads * head_dim;
     
-    // 分配设备内存
     T *d_q = nullptr, *d_k = nullptr, *d_v = nullptr, *d_o = nullptr;
     
     mcMalloc(&d_q, q_size * sizeof(T));
@@ -333,12 +279,10 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
     mcMalloc(&d_v, kv_size * sizeof(T));
     mcMalloc(&d_o, o_size * sizeof(T));
     
-    // 拷贝数据
     mcMemcpy(d_q, h_q.data(), q_size * sizeof(T), mcMemcpyHostToDevice);
     mcMemcpy(d_k, h_k.data(), kv_size * sizeof(T), mcMemcpyHostToDevice);
     mcMemcpy(d_v, h_v.data(), kv_size * sizeof(T), mcMemcpyHostToDevice);
     
-    // 缩放因子
     T scale;
     if constexpr (std::is_same_v<T, __half>) {
         float head_dim_f = static_cast<float>(head_dim);
@@ -349,16 +293,13 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
         scale = T(1.0 / sqrt(static_cast<double>(head_dim)));
     }
     
-    // 启动对应的kernel
     dim3 grid(batch_size, target_seq_len, query_heads);
     
-    // 优化block大小：确保是32的倍数（warp大小）
     int block_size = 256;
     if (head_dim < 256) {
-        block_size = ((head_dim + 31) / 32) * 32; // 向上取整到32的倍数
+        block_size = ((head_dim + 31) / 32) * 32;
     }
     
-    // 计算共享内存大小
     size_t shared_mem_size = head_dim * sizeof(float);
     
     if constexpr (std::is_same_v<T, float>) {
@@ -375,18 +316,15 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
             is_causal, scale);
     }
     
-    // 同步和错误检查
     mcDeviceSynchronize();
     mcError_t err = mcGetLastError();
     if (err != mcSuccess) {
         std::cerr << "MC error: " << mcGetErrorString(err) << std::endl;
     }
     
-    // 拷贝结果
     h_o.resize(o_size);
     mcMemcpy(h_o.data(), d_o, o_size * sizeof(T), mcMemcpyDeviceToHost);
     
-    // 清理
     cleanup(d_q, d_k, d_v, d_o);
 }
 

From ae00b8fb82ed28d3f53e555b29b381d2a9979002 Mon Sep 17 00:00:00 2001
From: zhang200210-neu <wj2my4ttgg@privaterelay.appleid.com>
Date: Tue, 11 Aug 2026 14:19:55 +0800
Subject: [PATCH 5/7] Update kernels.cu

---
 src/kernels.cu | 142 ++++++++++++++++++++++++++++---------------------
 1 file changed, 82 insertions(+), 60 deletions(-)

diff --git a/src/kernels.cu b/src/kernels.cu
index 8f43dd36..b6996841 100644
--- a/src/kernels.cu
+++ b/src/kernels.cu
@@ -1,9 +1,14 @@
 #include <vector>
+#include <cmath>
+#include <iostream>
+#include <algorithm>
+
+// NVIDIA 和天数智芯共用 CUDA 兼容头文件
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
-#include <cmath>
 
-#include "../tester/utils.h"
+// 常量定义
+const float INFINITY_F = 1e30f;
 
 /**
  * @brief Computes RMSNorm over the last dimension of a 2D tensor.
@@ -23,9 +28,6 @@
  * @param[in] hidden_dim Size of the normalized dimension.
  * @param[in] eps Numerical stability epsilon.
  */
-
-
-// CUDA kernel for RMSNorm computation
 template <typename T>
 __global__ void rmsNormKernel(const T* __restrict__ input,
                                const T* __restrict__ weight,
@@ -80,13 +82,30 @@ void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
               std::vector<T>& h_output, size_t rows, size_t hidden_dim,
               float eps) {
     // Allocate device memory
-    T *d_input, *d_weight, *d_output;
+    T *d_input = nullptr, *d_weight = nullptr, *d_output = nullptr;
     size_t input_size = rows * hidden_dim * sizeof(T);
     size_t weight_size = hidden_dim * sizeof(T);
     
-    cudaMalloc(&d_input, input_size);
-    cudaMalloc(&d_weight, weight_size);
-    cudaMalloc(&d_output, input_size);
+    cudaError_t err;
+    
+    err = cudaMalloc(&d_input, input_size);
+    if (err != cudaSuccess) {
+        std::cerr << "Failed to allocate d_input: " << cudaGetErrorString(err) << std::endl;
+        return;
+    }
+    err = cudaMalloc(&d_weight, weight_size);
+    if (err != cudaSuccess) {
+        std::cerr << "Failed to allocate d_weight: " << cudaGetErrorString(err) << std::endl;
+        cudaFree(d_input);
+        return;
+    }
+    err = cudaMalloc(&d_output, input_size);
+    if (err != cudaSuccess) {
+        std::cerr << "Failed to allocate d_output: " << cudaGetErrorString(err) << std::endl;
+        cudaFree(d_input);
+        cudaFree(d_weight);
+        return;
+    }
     
     // Copy data to device
     cudaMemcpy(d_input, h_input.data(), input_size, cudaMemcpyHostToDevice);
@@ -97,34 +116,21 @@ void rmsNorm(const std::vector<T>& h_input, const std::vector<T>& h_weight,
     int num_blocks = rows;
     
     // Launch kernel
-    rmsNormKernel<<<num_blocks, threads_per_block>>>(
+    rmsNormKernel<T><<<num_blocks, threads_per_block>>>(
         d_input, d_weight, d_output, rows, hidden_dim, eps);
     
-    // Copy result back to host
+    // Synchronize and copy result back
+    cudaDeviceSynchronize();
     cudaMemcpy(h_output.data(), d_output, input_size, cudaMemcpyDeviceToHost);
     
-    // Synchronize and free memory
-    cudaDeviceSynchronize();
+    // Free memory
     cudaFree(d_input);
     cudaFree(d_weight);
     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
+ * @brief Helper to clean up device memory.
  */
 template <typename T>
 void cleanup(T* d_q, T* d_k, T* d_v, T* d_o) {
@@ -134,16 +140,14 @@ void cleanup(T* d_q, T* d_k, T* d_v, T* d_o) {
     if (d_o) cudaFree(d_o);
 }
 
-// 优化但保持正确性的kernel函数
-
-// float版本的优化kernel
+// float 版本的 flash attention kernel
 __global__ void flash_attention_kernel_float(
     const float* Q, const float* K, const float* V, float* O,
     int batch_size, int target_seq_len, int src_seq_len,
     int query_heads, int kv_heads, int head_dim,
     bool is_causal, float scale) {
     
-    // 每个线程处理一个输出位置
+    // Each thread handles one output position
     int b = blockIdx.x;
     int t = blockIdx.y;
     int h = blockIdx.z;
@@ -157,23 +161,23 @@ __global__ void flash_attention_kernel_float(
     size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
     size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
     
-    // 使用共享内存存储查询向量
+    // Use shared memory to store query vector
     extern __shared__ float flash_attn_shared_float[];
     float* q_shared = flash_attn_shared_float;
     
-    // 协作加载查询向量到共享内存
+    // Cooperative loading of query vector into shared memory
     for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
         q_shared[i] = Q[q_base + i];
     }
     __syncthreads();
     
-    // 计算最大分数
+    // Compute max score
     float max_score = -1e10f;
     for (int s = 0; s < valid_len; ++s) {
         float dot = 0.0f;
         size_t k_base = kv_base + s * kv_heads * head_dim;
         
-        // 使用共享内存中的查询向量
+        // Use query vector from shared memory
         #pragma unroll(4)
         for (int i = 0; i < head_dim; ++i) {
             dot += q_shared[i] * K[k_base + i];
@@ -183,7 +187,7 @@ __global__ void flash_attention_kernel_float(
         if (score > max_score) max_score = score;
     }
     
-    // 计算输出
+    // Compute output
     float sum_exp = 0.0f;
     float output = 0.0f;
     
@@ -204,7 +208,7 @@ __global__ void flash_attention_kernel_float(
         output += exp_val * V[v_base];
     }
     
-    // 归一化
+    // Normalize
     if (sum_exp > 1e-12f) {
         output = output / sum_exp;
     } else if (valid_len > 0) {
@@ -219,7 +223,7 @@ __global__ void flash_attention_kernel_float(
     O[q_base + d] = output;
 }
 
-// half版本的优化kernel
+// half 版本的 flash attention kernel
 __global__ void flash_attention_kernel_half(
     const __half* Q, const __half* K, const __half* V, __half* O,
     int batch_size, int target_seq_len, int src_seq_len,
@@ -239,19 +243,19 @@ __global__ void flash_attention_kernel_half(
     size_t q_base = ((b * target_seq_len + t) * query_heads + h) * head_dim;
     size_t kv_base = b * src_seq_len * kv_heads * head_dim + kvh * head_dim;
     
-    // 使用共享内存存储查询向量
+    // Use shared memory to store query vector
     extern __shared__ float flash_attn_shared_half[];
     float* q_shared = flash_attn_shared_half;
     
     float scale_f = __half2float(scale);
     
-    // 协作加载查询向量到共享内存
+    // Cooperative loading of query vector into shared memory
     for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {
         q_shared[i] = __half2float(Q[q_base + i]);
     }
     __syncthreads();
     
-    // 计算最大分数
+    // Compute max score
     float max_score = -1e4f;
     for (int s = 0; s < valid_len; ++s) {
         float dot = 0.0f;
@@ -266,7 +270,7 @@ __global__ void flash_attention_kernel_half(
         if (score > max_score) max_score = score;
     }
     
-    // 计算softmax和输出
+    // Compute softmax and output
     float sum_exp = 0.0f;
     float output_f = 0.0f;
     
@@ -283,7 +287,7 @@ __global__ void flash_attention_kernel_half(
         float score = dot * scale_f;
         float shifted = score - max_score;
         
-        // 限制范围确保稳定性
+        // Clamp for stability
         if (shifted > 10.0f) shifted = 10.0f;
         if (shifted < -20.0f) shifted = -20.0f;
         
@@ -292,7 +296,7 @@ __global__ void flash_attention_kernel_half(
         output_f += exp_val * __half2float(V[v_base]);
     }
     
-    // 归一化
+    // Normalize
     if (sum_exp > 1e-7f) {
         output_f = output_f / sum_exp;
     } else if (valid_len > 0) {
@@ -306,13 +310,30 @@ __global__ void flash_attention_kernel_half(
     
     O[q_base + d] = __float2half(output_f);
 }
+
+/**
+ * @brief Computes flash attention for given query, key, and value tensors.
+ * 
+ * @tparam T Data type (float or __half) 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 <typename T>
 void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
                     const std::vector<T>& h_v, std::vector<T>& 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) {
     
-    // 基本检查
+    // Basic validation
     if (batch_size <= 0 || target_seq_len <= 0 || src_seq_len <= 0 || 
         query_heads <= 0 || kv_heads <= 0 || head_dim <= 0) {
         return;
@@ -320,12 +341,12 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
     
     if (query_heads % kv_heads != 0) return;
     
-    // 计算大小
+    // Calculate sizes
     size_t q_size = batch_size * target_seq_len * query_heads * head_dim;
     size_t kv_size = batch_size * src_seq_len * kv_heads * head_dim;
     size_t o_size = batch_size * target_seq_len * query_heads * head_dim;
     
-    // 分配设备内存
+    // Allocate device memory
     T *d_q = nullptr, *d_k = nullptr, *d_v = nullptr, *d_o = nullptr;
     
     cudaMalloc(&d_q, q_size * sizeof(T));
@@ -333,12 +354,12 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
     cudaMalloc(&d_v, kv_size * sizeof(T));
     cudaMalloc(&d_o, o_size * sizeof(T));
     
-    // 拷贝数据
+    // Copy data to device
     cudaMemcpy(d_q, h_q.data(), q_size * sizeof(T), cudaMemcpyHostToDevice);
     cudaMemcpy(d_k, h_k.data(), kv_size * sizeof(T), cudaMemcpyHostToDevice);
     cudaMemcpy(d_v, h_v.data(), kv_size * sizeof(T), cudaMemcpyHostToDevice);
     
-    // 缩放因子
+    // Scale factor
     T scale;
     if constexpr (std::is_same_v<T, __half>) {
         float head_dim_f = static_cast<float>(head_dim);
@@ -349,18 +370,19 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
         scale = T(1.0 / sqrt(static_cast<double>(head_dim)));
     }
     
-    // 启动对应的kernel
+    // Launch configuration
     dim3 grid(batch_size, target_seq_len, query_heads);
     
-    // 优化block大小：确保是32的倍数（warp大小）
+    // Optimize block size: ensure multiple of 32 (warp size)
     int block_size = 256;
     if (head_dim < 256) {
-        block_size = ((head_dim + 31) / 32) * 32; // 向上取整到32的倍数
+        block_size = ((head_dim + 31) / 32) * 32; // Round up to multiple of 32
     }
     
-    // 计算共享内存大小
+    // Shared memory size
     size_t shared_mem_size = head_dim * sizeof(float);
     
+    // Launch corresponding kernel
     if constexpr (std::is_same_v<T, float>) {
         flash_attention_kernel_float<<<grid, block_size, shared_mem_size>>>(
             d_q, d_k, d_v, d_o,
@@ -375,18 +397,18 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
             is_causal, scale);
     }
     
-    // 同步和错误检查
+    // Synchronize and error check
     cudaDeviceSynchronize();
     cudaError_t err = cudaGetLastError();
     if (err != cudaSuccess) {
         std::cerr << "CUDA error: " << cudaGetErrorString(err) << std::endl;
     }
     
-    // 拷贝结果
+    // Copy result back
     h_o.resize(o_size);
     cudaMemcpy(h_o.data(), d_o, o_size * sizeof(T), cudaMemcpyDeviceToHost);
     
-    // 清理
+    // Cleanup
     cleanup(d_q, d_k, d_v, d_o);
 }
 
@@ -396,11 +418,11 @@ void flashAttention(const std::vector<T>& h_q, const std::vector<T>& h_k,
 // *********************************************************************
 template void rmsNorm<float>(const std::vector<float>&, const std::vector<float>&,
   std::vector<float>&, size_t, size_t, float);
-template void rmsNorm<half>(const std::vector<half>&, const std::vector<half>&,
-  std::vector<half>&, size_t, size_t, float);
+template void rmsNorm<__half>(const std::vector<__half>&, const std::vector<__half>&,
+  std::vector<__half>&, size_t, size_t, float);
 template void flashAttention<float>(const std::vector<float>&, const std::vector<float>&,
   const std::vector<float>&, std::vector<float>&,
   int, int, int, int, int, int, bool);
-template void flashAttention<half>(const std::vector<half>&, const std::vector<half>&,
-  const std::vector<half>&, std::vector<half>&,
+template void flashAttention<__half>(const std::vector<__half>&, const std::vector<__half>&,
+  const std::vector<__half>&, std::vector<__half>&,
   int, int, int, int, int, int, bool);

From deb8141fdbcd2021c4ee2114d561355f4930973a Mon Sep 17 00:00:00 2001
From: zhang200210-neu <wj2my4ttgg@privaterelay.appleid.com>
Date: Tue, 11 Aug 2026 14:20:20 +0800
Subject: [PATCH 6/7] Update kernels.cu

---
 src/kernels.cu | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/kernels.cu b/src/kernels.cu
index b6996841..958653f4 100644
--- a/src/kernels.cu
+++ b/src/kernels.cu
@@ -3,11 +3,9 @@
 #include <iostream>
 #include <algorithm>
 
-// NVIDIA 和天数智芯共用 CUDA 兼容头文件
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
-// 常量定义
 const float INFINITY_F = 1e30f;
 
 /**

From 4062db4a0cd8dcbc40fdf9452e2e8170dcec9d41 Mon Sep 17 00:00:00 2001
From: zhang200210-neu <wj2my4ttgg@privaterelay.appleid.com>
Date: Tue, 11 Aug 2026 14:21:56 +0800
Subject: [PATCH 7/7] Update kernels.maca

---
 src/kernels.maca | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/kernels.maca b/src/kernels.maca
index 93016502..cf669d10 100644
--- a/src/kernels.maca
+++ b/src/kernels.maca
@@ -3,7 +3,7 @@
 #include <maca_fp16.h>
 #include <cmath>
 
-// 若 half 类型未自动定义，手动提供别名
+
 #ifndef half
 using half = __half;
 #endif