diff --git a/.vscode/c_cpp_properties.json b/.vscode/c_cpp_properties.json new file mode 100644 index 00000000..c2098a2d --- /dev/null +++ b/.vscode/c_cpp_properties.json @@ -0,0 +1,18 @@ +{ + "configurations": [ + { + "name": "linux-gcc-x64", + "includePath": [ + "${workspaceFolder}/**" + ], + "compilerPath": "/usr/bin/gcc", + "cStandard": "${default}", + "cppStandard": "${default}", + "intelliSenseMode": "linux-gcc-x64", + "compilerArgs": [ + "" + ] + } + ], + "version": 4 +} \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 00000000..7163c030 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,24 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "C/C++ Runner: Debug Session", + "type": "cppdbg", + "request": "launch", + "args": [], + "stopAtEntry": false, + "externalConsole": false, + "cwd": "/home/liu/InfiniTensor/Learning-CUDA/src", + "program": "/home/liu/InfiniTensor/Learning-CUDA/src/build/Debug/outDebug", + "MIMode": "gdb", + "miDebuggerPath": "gdb", + "setupCommands": [ + { + "description": "Enable pretty-printing for gdb", + "text": "-enable-pretty-printing", + "ignoreFailures": true + } + ] + } + ] +} \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..3e5eb956 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,59 @@ +{ + "C_Cpp_Runner.cCompilerPath": "gcc", + "C_Cpp_Runner.cppCompilerPath": "g++", + "C_Cpp_Runner.debuggerPath": "gdb", + "C_Cpp_Runner.cStandard": "", + "C_Cpp_Runner.cppStandard": "", + "C_Cpp_Runner.msvcBatchPath": "", + "C_Cpp_Runner.useMsvc": false, + "C_Cpp_Runner.warnings": [ + "-Wall", + "-Wextra", + "-Wpedantic", + "-Wshadow", + "-Wformat=2", + "-Wcast-align", + "-Wconversion", + "-Wsign-conversion", + "-Wnull-dereference" + ], + "C_Cpp_Runner.msvcWarnings": [ + "/W4", + "/permissive-", + "/w14242", + "/w14287", + "/w14296", + "/w14311", + "/w14826", + "/w44062", + "/w44242", + "/w14905", + "/w14906", + "/w14263", + "/w44265", + "/w14928" + ], + "C_Cpp_Runner.enableWarnings": true, + "C_Cpp_Runner.warningsAsError": false, + "C_Cpp_Runner.compilerArgs": [], + "C_Cpp_Runner.linkerArgs": [], + "C_Cpp_Runner.includePaths": [], + "C_Cpp_Runner.includeSearch": [ + "*", + "**/*" + ], + "C_Cpp_Runner.excludeSearch": [ + "**/build", + "**/build/**", + "**/.*", + "**/.*/**", + "**/.vscode", + "**/.vscode/**" + ], + "C_Cpp_Runner.useAddressSanitizer": false, + "C_Cpp_Runner.useUndefinedSanitizer": false, + "C_Cpp_Runner.useLeakSanitizer": false, + "C_Cpp_Runner.showCompilationTime": false, + "C_Cpp_Runner.useLinkTimeOptimization": false, + "C_Cpp_Runner.msvcSecureNoWarnings": false +} \ No newline at end of file diff --git a/Makefile b/Makefile index 883d452a..540a4518 100644 --- a/Makefile +++ b/Makefile @@ -98,6 +98,6 @@ $(TARGET): $(STUDENT_OBJ) $(TEST_OBJ) $(CC) $(CFLAGS) $(PLATFORM_DEFINE) -o $@ $^ $(EXTRA_LIBS) # Generate src object: Compile kernels.cu (triggers template instantiation) -$(STUDENT_OBJ): $(STUDENT_SRC) +$(addsuffix .o,$(basename $(STUDENT_SRC))): $(STUDENT_SRC) @echo "=== Compiling student code ($(STUDENT_SRC)) ===" $(CC) $(CFLAGS) $(PLATFORM_DEFINE) -c $< -o $@ diff --git a/code_prob_summary.md b/code_prob_summary.md new file mode 100644 index 00000000..5e6b66f0 --- /dev/null +++ b/code_prob_summary.md @@ -0,0 +1,55 @@ +## 1. + +```bash +~/InfiniTensor/Learning-CUDA$ make VERBOSE=true +``` + +通过此来通过 nvidia的版本的检测。 + +``` +/data/C500/Learning-CUDA# LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH make VERBOSE=true PLATFORM=metax +``` + +通过此来通过沐曦metax的版本的检测。 + + + +## 2.存在的问题和困惑 + +online softmax FA2的编写部分，出现了e-5级别的误差导致难以使用真正的FA2来去通过编译。 + +只能使用相对原始一些的attention计算逻辑。对于metax和nvidia都是类似的情况。 + +| Case | Type | Avg Time (ms) | Max Diff | Max Tolerance | Result | +| ------ | --------- | ------------- | ------------- | ------------- | -------- | +| 1 | float | 0.452434 | 0.0000000 | 0.0000000 | Pass | +| 1 | half | 0.384831 | 0.0000000 | 0.0000000 | Pass | +| 2 | float | 0.408454 | 0.0000000 | 0.0000000 | Pass | +| 2 | half | 0.445536 | 0.0000000 | 0.0000000 | Pass | +| 3 | float | 0.387116 | 0.0000010 | 0.0007855 | Pass | +| 3 | half | 0.415170 | 0.0000000 | 0.0000000 | Pass | +| 4 | float | 0.450248 | 0.0000237 | 0.0001501 | Pass | +| 4 | half | 0.448288 | 0.0000076 | 0.0050435 | Pass | +| 5 | float | 0.584307 | 0.0000048 | 0.0009049 | Pass | +| 5 | half | 0.557689 | 0.0019531 | 0.0231250 | Pass | +| **6** | **float** | **2.688137** | **0.0000257** | **0.0000146** | **Fail** | +| 6 | half | 1.377309 | 0.0078125 | 0.0451562 | Pass | +| 7 | float | 0.391246 | 0.0000057 | 0.0008215 | Pass | +| 7 | half | 0.459092 | 0.0000000 | 0.0000000 | Pass | +| 8 | float | 0.438399 | 0.0000057 | 0.0008215 | Pass | +| 8 | half | 0.443422 | 0.0000000 | 0.0000000 | Pass | +| 9 | float | 0.429803 | 0.0000172 | 0.0002611 | Pass | +| 9 | half | 0.447457 | 0.0039062 | 0.0431250 | Pass | +| 10 | float | 0.558301 | 0.0000207 | 0.0002537 | Pass | +| 10 | half | 0.524745 | 0.0039062 | 0.0313086 | Pass | +| 11 | float | 0.400683 | 0.0000019 | 0.0009848 | Pass | +| 11 | half | 0.471791 | 0.0000305 | 0.0052132 | Pass | +| 12 | float | 0.427469 | 0.0000029 | 0.0009080 | Pass | +| 12 | half | 0.472397 | 0.0039062 | 0.0410937 | Pass | +| 13 | float | 15.444282 | 0.0000048 | 0.0008913 | Pass | +| 13 | half | 3.981441 | 0.0078125 | 0.0469141 | Pass | +| **14** | **float** | **63.248356** | **0.0000133** | **0.0000123** | **Fail** | +| 14 | half | 18.820628 | 0.0078125 | 0.0470703 | Pass | + + + diff --git a/onlinesoftmax_prob/FA2.cu b/onlinesoftmax_prob/FA2.cu new file mode 100644 index 00000000..e90179cd --- /dev/null +++ b/onlinesoftmax_prob/FA2.cu @@ -0,0 +1,293 @@ +#include + +#include + +#define WARPSIZE 32 + +// 感觉用了才觉得这个是必须的，不然要对外写很多模板的细节要求或者写很多的宏定义 +// 且这些要求一般都是内含计算规律和要求的，以便应用到kernel中的计算，内存申请中 +template +struct Fa2Traits +{ + static constexpr int kWarpSize = WARPSIZE; + static constexpr int kThreads = NumWarps * kWarpSize; + + static_assert(HeadDim > 0 && HeadDim % kWarpSize == 0); + static_assert(Br > 0 && Br % NumWarps == 0); + static_assert(Bc > 0 && Bc % kWarpSize == 0); + + static constexpr int kDPerLane = HeadDim / kWarpSize; + static constexpr int kRowsPerWarp = Br / NumWarps; + static constexpr int kKeysPerLane = Bc / kWarpSize; + static constexpr int kKvTileElements = Bc * HeadDim; + static constexpr std::size_t kSmemBytes = + 2 * kKvTileElements * sizeof(float); +}; + +struct Fa2Params +{ + const float *q; + const float *k; + const float *v; + float *o; + + int batch_size; + int q_heads; + int kv_heads; + int q_lens; + int kv_lens; + + float qk_scale; // 一般是1/sqrt(head_Dim) +}; + +// 用xor有更好的效果 +__device__ __forceinline__ float +warpReduceSum(float v) +{ + unsigned Mask = 0xffffffffu; + v += __shfl_xor_sync(Mask, v, 16); + v += __shfl_xor_sync(Mask, v, 8); + v += __shfl_xor_sync(Mask, v, 4); + v += __shfl_xor_sync(Mask, v, 2); + v += __shfl_xor_sync(Mask, v, 1); + + return v; +} + +__device__ __forceinline__ float warpReduceMax(float v) +{ + unsigned Mask = 0xffffffffu; + + v = fmaxf(v, __shfl_xor_sync(Mask, v, 16)); + v = fmaxf(v, __shfl_xor_sync(Mask, v, 8)); + v = fmaxf(v, __shfl_xor_sync(Mask, v, 4)); + v = fmaxf(v, __shfl_xor_sync(Mask, v, 2)); + v = fmaxf(v, __shfl_xor_sync(Mask, v, 1)); + + return v; +} + +// 模板是为了让编译器在编译期间就确定好代码的基本结构 +// 比如可以让smem可以使用模板来去创建内存大小 +// 比如可以根据模板的情况 来去删去/保留某些代码流 +//***一个Warp处理一个Row的Dim个数据(一般一个Warp负责多个Rows，但是是一个Row处理完在处理下一个) +// 每个Lane负责处理和保存的元素 = HeadDim / WARPSIZE + +template +__global__ void FlashAttentionKernel(Fa2Params params) +{ + using Traits = Fa2Traits; + // params.qk_scale = 1 / sqrt(HeadDim); + int tid = threadIdx.x; + + // Block的并行维度 + int bid = blockIdx.x; + int q_head = blockIdx.y; + int batch = blockIdx.z; + + int query_heads = params.q_heads; + int kv_heads = params.kv_heads; + + int group_size = query_heads / kv_heads; // GQA/MHA + int kv_head = q_head / group_size; // 让多个q_head共享使用第i个kv_head来实现GQA 体现在blockIdx.y + + int warp_id = tid / WARPSIZE; + int warp_lane = tid % WARPSIZE; // 当前线程属于warp内的哪一个线程 + + // 外层需要遍历KV_tile 同时需要知道当前Block需要处理的Q_tile是哪一个内存偏移 + // KV_tile的地址是会变化读取的，而Q_tile在读取后对当前Block是静态的 + + // Block内的warp处理好几个rows，需要为每个warp自己开始处理的row进行内部定位 + // 后续在KV tile循环的过程中，他们的rows会一次一行的处理完，这时候我们开始记录内部的Q_row的offset + int Q_row_base = bid * Br; + int Q_row_offset; + + __shared__ float smem[2 * Traits::kKvTileElements]; + float *k_tile = smem; + float *v_tile = smem + Traits::kKvTileElements; + + float m_reg[Traits::kRowsPerWarp]; // 每行的最大值(Score中的最大值) + float l_reg[Traits::kRowsPerWarp]; // softmax每行的指数和分母 + float acc_reg[Traits::kRowsPerWarp][Traits::kDPerLane]; + float q_tile[Traits::kRowsPerWarp][Traits::kDPerLane]; + + // 1.进行Q tile的加载等初始化过程 + for (int rowIdx = 0; rowIdx < Traits::kRowsPerWarp; rowIdx++) + { + // 计算自己warp负责的rows的起始偏移行(Block tile内部的row坐标计算了) + // 正确地址写法（沿用当前变量名）： + // int row_in_warp = warp_id * Traits::kRowsPerWarp + rowIdx; + // Q_row_offset = Q_row_base + row_in_warp; + // std::size_t Q_base = + // ((static_cast(batch) * query_heads + q_head) * + // params.q_lens + + // Q_row_offset) * + // HeadDim; + int row_in_warp = warp_id * Traits::kRowsPerWarp + rowIdx; + // 主要这个是在计算Q[,,Len,] 第三维度对应的token的 + Q_row_offset = row_in_warp + Q_row_base; // Q[B,Hq,Len,D] Stride[Hq*Len*D,Len*D,D,1] + const std::size_t Q_base = + ((static_cast(batch) * query_heads + q_head) * + params.q_lens + + Q_row_offset) * + HeadDim; + // 基于Q_base的移动就是在遍历某个特定Head的Dim + + for (int colIdx = 0; colIdx < Traits::kDPerLane; colIdx++) + { + int d = warp_lane + WARPSIZE * colIdx; // 对于线程角度下的k这个实际索引 lane0:0--32--64 lane1:1--33--65 + q_tile[rowIdx][colIdx] = params.q[Q_base + d]; + acc_reg[rowIdx][colIdx] = 0.0f; + } + m_reg[rowIdx] = -INFINITY; + l_reg[rowIdx] = 0.0f; + } + + // 目前是要求需要整除， 没太做一些冗余一小部分的padding处理 + int kv_tile_nums = params.kv_lens / Bc; + // 2.开始主要的KV tile的循环 + + for (int kv_tile_idx = 0; kv_tile_idx < kv_tile_nums; kv_tile_idx++) + { + + int kv_tile_base = kv_tile_idx * Bc; + // Block内的thread需要协作从Bc x d的KV tile中联合读取数据 + // 就采用Block内线程联合读取到smem的方式来去处理 + // 正确地址写法（K/V 的布局相同，共用同一个起始偏移）： + // std::size_t KV_base = + // ((static_cast(batch) * kv_heads + kv_head) * + // params.kv_lens + + // kv_tile_base) * + // HeadDim; + const std::size_t KV_base = + ((static_cast(batch) * kv_heads + kv_head) * + params.kv_lens + + kv_tile_base) * + HeadDim; + for (int i = tid; i < Traits::kKvTileElements; i += Traits::kThreads) + { + k_tile[i] = params.k[KV_base + i]; + v_tile[i] = params.v[KV_base + i]; + // 当前KV tile循环下所需要的数据都已加载到smem中 + } + __syncthreads(); + for (int rowIdx = 0; rowIdx < Traits::kRowsPerWarp; rowIdx++) // 这里的是为了计算BrxBc大小的score + { + float score_and_p[Traits::kKeysPerLane]; // 每行的数据都进行一次记录，他主要是作为中转后续存到acc中 + // 同一 warp 内各 lane 分散持有不同 key 的 score/p；它们不是相同副本。 + for (int i = 0; i < Traits::kKeysPerLane; i++) + score_and_p[i] = -INFINITY; // 这里不是0 + // 该 + + for (int k = 0; k < Bc; k++) // 内部的一次for循环求一个dot的score，这里的循环是为了计算一行的dot的score + { + float partial_dot = 0.0f; + for (int colIdx = 0; colIdx < Traits::kDPerLane; colIdx++) + // k---这个代表着在处理一行中的第几个score，这个score的位置需要由对此负责的warp中的一个线程来去认领和处理 + { + + int d = warp_lane + colIdx * WARPSIZE; + partial_dot += q_tile[rowIdx][colIdx] * k_tile[k * HeadDim + d]; // k在K_tile中是row坐标 + } + + float qk_dot = warpReduceSum(partial_dot); + bool is_Mask = true; // true时对应的是q位置合法 + if constexpr (Causal) + { + // + int row_In_warp = warp_id * Traits::kRowsPerWarp + rowIdx; + int q_postion = Q_row_base + row_In_warp; // 目前是这个pos的设置 + int k_postion = k + kv_tile_base; + is_Mask = k_postion <= q_postion; // causal：只能访问当前位置及以前的 key + } + int owner_slot = k / WARPSIZE; // 这里slot用来表示 负责处理这个lane的线程 将其保存到他的寄存器的第一个槽位中 + int owner_lane = k % WARPSIZE; // 这里并行的是多个warp来去协调处理当前与KV_tile的结果 + if (owner_lane == warp_lane) + { + score_and_p[owner_slot] = is_Mask ? params.qk_scale * qk_dot : -INFINITY; // 此时还没有经过softmax的指数处理 + // -INFINITY是指数后正好为0 + } + } + + float local_tile_max = -INFINITY; + + for (int colIdx = 0; colIdx < Traits::kKeysPerLane; colIdx++) + { + local_tile_max = fmaxf(local_tile_max, score_and_p[colIdx]); + } + + // 这里有点关于需不需要专门线程来去存储/处理最大值的一个小问题 + float tile_max = warpReduceMax(local_tile_max); + float m_old = m_reg[rowIdx]; + float m_new = fmaxf(m_old, tile_max); + float alpha = m_old == -INFINITY ? 0.0f : __expf(m_old - m_new); // 这里的条件是在预防什么情况 + m_reg[rowIdx] = m_new; + + float local_tile_sum = 0.0f; + + for (int colIdx = 0; colIdx < Traits::kKeysPerLane; colIdx++) + { + float score = score_and_p[colIdx]; + float p = (score == -INFINITY) ? 0.0f : __expf(score - m_new); // 这部分的公式稍微确认一下 + score_and_p[colIdx] = p; + local_tile_sum += p; + } + float tile_sum = warpReduceSum(local_tile_sum); + l_reg[rowIdx] = alpha * l_reg[rowIdx] + tile_sum; + + for (int colIdx = 0; colIdx < Traits::kDPerLane; colIdx++) + { + acc_reg[rowIdx][colIdx] = alpha * acc_reg[rowIdx][colIdx]; + } + + for (int colIdx = 0; colIdx < Traits::kKeysPerLane; colIdx++) + { + for (int owner_lane = 0; owner_lane < WARPSIZE; owner_lane++) + { + int key = colIdx * WARPSIZE + owner_lane; + float p = __shfl_sync(0xffffffffu, score_and_p[colIdx], owner_lane); + + for (int d_slot = 0; d_slot < Traits::kDPerLane; d_slot++) + { + int d = warp_lane + d_slot * WARPSIZE; + acc_reg[rowIdx][d_slot] += (p * v_tile[key * HeadDim + d]); + } + } + } + } + __syncthreads(); + } + // 此时KV tile迭代完毕 所有的acc_reg中存有的最终的o的结果，需要各自将其写入对应的位置中 + // 一行一行的写入 + for (int row_Idx = 0; row_Idx < Traits::kRowsPerWarp; row_Idx++) + { // O[B,Hq,Len,D] Stride[Hq*Len*D,Len*D,D,1] + + // 正确地址写法（当前 row_Idx 对应的全局 Q/O 行）： + // int row_in_warp = + // warp_id * Traits::kRowsPerWarp + row_Idx; + // Q_row_offset = Q_row_base + row_in_warp; + // std::size_t O_base = + // ((static_cast(batch) * query_heads + q_head) * + // params.q_lens + + // Q_row_offset) * + // HeadDim; + int row_in_warp = warp_id * Traits::kRowsPerWarp + row_Idx; + Q_row_offset = Q_row_base + row_in_warp; + const std::size_t O_base = + ((static_cast(batch) * query_heads + q_head) * + params.q_lens + + Q_row_offset) * + HeadDim; + + float inv_l = l_reg[row_Idx] > 0.0f ? 1.0f / l_reg[row_Idx] : 0.0f; // 分母指数和 后面直接作为系数调整acc即可 + // l_reg 在所有 lane 中是相同副本；每个 lane 用它归一化自己负责的 d。 + + for (int col_Idx = 0; col_Idx < Traits::kDPerLane; col_Idx++) + { + int d = warp_lane + col_Idx * WARPSIZE; + params.o[O_base + d] = acc_reg[row_Idx][col_Idx] * inv_l; + } + } +} + +// Q/O 使用 [B,Hq,Lq,D] 坐标；K/V 使用 [B,Hkv,Lkv,D] 坐标。 diff --git a/onlinesoftmax_prob/FA2_compatible.cu b/onlinesoftmax_prob/FA2_compatible.cu new file mode 100644 index 00000000..03e02fff --- /dev/null +++ b/onlinesoftmax_prob/FA2_compatible.cu @@ -0,0 +1,522 @@ +/* + * A correctness-first adaptation of src/FA2.cu for the llaisys attention + * interface. The original work and this adaptation are covered by the MIT + * license in the repository root. + * + * Supported interface: + * Q/O: [batch_size, q_len, q_heads, head_dim] + * K/V: [batch_size, kv_len, kv_heads, head_dim] + * + * P0 change map (search "[P0]" in this file to locate each adjustment): + * The tiled kernel below keeps the same FA2 online-softmax update order as + * src/FA2.cu and adds: + * - token-major addressing used by llaisys; + * - GQA/MHA head mapping; + * - PyTorch-compatible upper-left causal masking; + * - q_len / kv_len tail handling; + * - arbitrary positive head_dim values; + * - float/half I/O, plus bfloat16 when the backend exposes cuda_bf16.h; + * - FP32 score and output accumulation. + */ + +#if __has_include() +#include +#define FA2_COMPATIBLE_HAS_BFLOAT16 1 +#endif +#include +#include + +#include "../include/FA2_compatible.cuh" +#include "../include/platform_warp.h" + +#include +#if defined(FA2_TRACE_DEVICE) +#include +#endif +#include +#include +#include + +namespace fa2_compatible +{ + namespace + { + + // HeadDim is a compile-time capacity bucket. params.head_dim remains + // the actual runtime width, so non-power-of-two tails are still valid. + template + struct Fa2Traits + { + static constexpr int kWarpSize = platform::WARP_SIZE; + static constexpr int kWarpsPerCta = 4; + static constexpr int kRowsPerWarp = 4; + static constexpr int kRowsPerCta = + kWarpsPerCta * kRowsPerWarp; + static constexpr int kThreads = kWarpSize * kWarpsPerCta; + static constexpr int kBc = 16; + static constexpr int kDPerLane = + (HeadDim + kWarpSize - 1) / kWarpSize; + static constexpr int kKvTileElements = kBc * HeadDim; + }; + + using DefaultTraits = Fa2Traits<64>; + static_assert(DefaultTraits::kThreads % DefaultTraits::kWarpSize == 0, + "Full-mask reductions require complete warps"); + static_assert(DefaultTraits::kThreads <= 1024, + "Attention CTA exceeds the supported thread count"); + static_assert(DefaultTraits::kBc <= DefaultTraits::kWarpSize, + "Each tile key is owned by one physical warp lane"); + + template + __device__ __forceinline__ float toFloat(T value); + + template <> + __device__ __forceinline__ float toFloat(float value) + { + return value; + } + + template <> + __device__ __forceinline__ float toFloat<__half>(__half value) + { + return __half2float(value); + } + +#if defined(FA2_COMPATIBLE_HAS_BFLOAT16) + template <> + __device__ __forceinline__ float toFloat<__nv_bfloat16>( + __nv_bfloat16 value) + { + return __bfloat162float(value); + } +#endif + + template + __device__ __forceinline__ T fromFloat(float value); + + template <> + __device__ __forceinline__ float fromFloat(float value) + { + return value; + } + + template <> + __device__ __forceinline__ __half fromFloat<__half>(float value) + { + return __float2half(value); + } + +#if defined(FA2_COMPATIBLE_HAS_BFLOAT16) + template <> + __device__ __forceinline__ __nv_bfloat16 + fromFloat<__nv_bfloat16>(float value) + { + return __float2bfloat16(value); + } +#endif + + template + struct Fa2Params + { + T *output; + const T *query; + const T *key; + const T *value; + + // [P0] 原基线缺少 batch/causal；这里集中保存完整接口元数据，避免 + // kernel 内再假设 batch=1 或 causal 恒开启。 + int batch_size; + int q_len; + int kv_len; + int q_heads; + int kv_heads; + int head_dim; + bool is_causal; + float scale; + }; + + template + __global__ void fa2CompatibleKernel(Fa2Params params) + { + using Traits = Fa2Traits; + + const int tid = threadIdx.x; + const int warp_id = tid / Traits::kWarpSize; + const int warp_lane = tid % Traits::kWarpSize; + // [P0 CTA mapping] Four warps share one KV tile. Each warp owns + // four query rows, so one CTA processes 16 consecutive rows. + const int batch = blockIdx.z; + const int q_head = blockIdx.y; + const int cta_q_base = blockIdx.x * Traits::kRowsPerCta; + + // [P0 GQA] q_head 按组映射到共享的 kv_head。 + const int group_size = params.q_heads / params.kv_heads; + const int kv_head = q_head / group_size; + + // K/V are the only CTA-wide tiles. Q and the online-softmax + // accumulators stay in per-thread local arrays, matching FA2.cu. + __shared__ float smem[2 * Traits::kKvTileElements]; + float *k_tile = smem; + float *v_tile = smem + Traits::kKvTileElements; + + float m_reg[Traits::kRowsPerWarp]; + float l_reg[Traits::kRowsPerWarp]; + float acc_reg[Traits::kRowsPerWarp][Traits::kDPerLane]; + float q_tile[Traits::kRowsPerWarp][Traits::kDPerLane]; +#pragma unroll + for (int row_slot = 0; row_slot < Traits::kRowsPerWarp; + ++row_slot) + { + m_reg[row_slot] = -INFINITY; + l_reg[row_slot] = 0.0f; +#pragma unroll + for (int d_slot = 0; d_slot < Traits::kDPerLane; ++d_slot) + { + const int d = warp_lane + d_slot * Traits::kWarpSize; + const int row_in_cta = + row_slot * Traits::kWarpsPerCta + warp_id; + const int q_row = cta_q_base + row_in_cta; + const bool valid_q = q_row < params.q_len; + if (valid_q && d < params.head_dim) + { + const std::size_t q_index = + ((static_cast(batch) * params.q_len + + q_row) * + params.q_heads + + q_head) * + params.head_dim + + d; + q_tile[row_slot][d_slot] = + toFloat(params.query[q_index]); + } + else + { + q_tile[row_slot][d_slot] = 0.0f; + } + acc_reg[row_slot][d_slot] = 0.0f; + } + } + + const int kv_tile_nums = + (params.kv_len + Traits::kBc - 1) / Traits::kBc; + + // FA2-style online softmax: each KV tile updates the running row max, + // denominator and output accumulator in one pass. This preserves the + // compatible layout/tail handling without materializing a score matrix. + for (int kv_tile_idx = 0; kv_tile_idx < kv_tile_nums; + ++kv_tile_idx) + { + const int kv_tile_base = kv_tile_idx * Traits::kBc; + const int cta_last_q = + min(cta_q_base + Traits::kRowsPerCta - 1, + params.q_len - 1); + + // This predicate is uniform for the whole CTA. Once a causal + // KV tile starts after its last valid query row, every key in + // the tile is masked for all four warps, so no load or online + // softmax update is needed. + if (params.is_causal && kv_tile_base > cta_last_q) + { + continue; + } + + const int kv_tile_elements = Traits::kBc * params.head_dim; + + // All threads participate because the block-wide barriers below + // must also be reached by warps in the final partial Q tile. + for (int i = tid; i < kv_tile_elements; i += Traits::kThreads) + { + const int key_in_tile = i / params.head_dim; + const int d = i % params.head_dim; + const int key_row = kv_tile_base + key_in_tile; + const int tile_index = key_in_tile * HeadDim + d; + + if (key_row < params.kv_len) + { + // [P0 layout] K/V 地址统一为 + // (((b * kv_len + k) * kv_heads + kv_head) * D + d)。 + const std::size_t kv_index = + ((static_cast(batch) * params.kv_len + + key_row) * + params.kv_heads + + kv_head) * + params.head_dim + + d; + k_tile[tile_index] = toFloat(params.key[kv_index]); + v_tile[tile_index] = toFloat(params.value[kv_index]); + } + else + { + k_tile[tile_index] = 0.0f; + v_tile[tile_index] = 0.0f; + } + } + __syncthreads(); + + // Each warp updates its four independent online-softmax rows + // while the same K/V tile remains resident in shared memory. +#pragma unroll + for (int row_slot = 0; row_slot < Traits::kRowsPerWarp; + ++row_slot) + { + const int row_in_cta = + row_slot * Traits::kWarpsPerCta + warp_id; + const int q_row = cta_q_base + row_in_cta; + const bool valid_q = q_row < params.q_len; +#if defined(FA2_TRACE_DEVICE) + const bool trace_row = + std::is_same::value && + ((params.batch_size == 4 && params.q_len == 64 && + batch == 1 && q_row == 39 && q_head == 61) || + (params.batch_size == 4 && params.q_len == 512 && + batch == 2 && q_row == 199 && q_head == 62)); +#endif + // Traits::kBc <= physical warp size, so lane k owns p[k]. + float score_and_p = -INFINITY; + for (int key_in_tile = 0; + key_in_tile < Traits::kBc; ++key_in_tile) + { + const int key_row = kv_tile_base + key_in_tile; + const bool valid_key = + valid_q && key_row < params.kv_len && + (!params.is_causal || key_row <= q_row); + const float *k_row = k_tile + key_in_tile * HeadDim; + float qk_dot = 0.0f; +#if defined(PLATFORM_NVIDIA) + if constexpr (std::is_same::value) + { + if (valid_key && warp_lane == 0) + { + const std::size_t q_base = + ((static_cast(batch) * + params.q_len + + q_row) * + params.q_heads + + q_head) * + params.head_dim; + for (int d = 0; d < params.head_dim; ++d) + { + qk_dot = fmaf(params.query[q_base + d], + k_row[d], qk_dot); + } + } + qk_dot = platform::warpBroadcast(qk_dot, 0); + } + else +#endif + { +#pragma unroll + for (int d_slot = 0; + d_slot < Traits::kDPerLane; ++d_slot) + { + const int d = + warp_lane + d_slot * Traits::kWarpSize; + if (valid_key && d < params.head_dim) + { + qk_dot = fmaf(q_tile[row_slot][d_slot], + k_row[d], qk_dot); + } + } + qk_dot = platform::warpReduceSum(qk_dot); + } + + if (valid_key && warp_lane == key_in_tile) + { + score_and_p = qk_dot * params.scale; +#if defined(FA2_TRACE_DEVICE) + if (trace_row) + { + printf("FA2_SCORE B=%d Q=%d H=%d TILE=%d K=%d " + "score=%.9e\n", + batch, q_row, q_head, kv_tile_base, + key_row, score_and_p); + } +#endif + } + } + + const float tile_max = + platform::warpReduceMax(score_and_p); + const float m_old = m_reg[row_slot]; + const float m_new = fmaxf(m_old, tile_max); + const float alpha = m_old == -INFINITY + ? 0.0f + : __expf(m_old - m_new); + m_reg[row_slot] = m_new; + + score_and_p = score_and_p == -INFINITY + ? 0.0f + : __expf(score_and_p - m_new); + const float tile_sum = + platform::warpReduceSum(score_and_p); + l_reg[row_slot] = + fmaf(alpha, l_reg[row_slot], tile_sum); + +#if defined(FA2_TRACE_DEVICE) + if (trace_row && warp_lane == 0) + { + printf("FA2_TILE B=%d Q=%d H=%d TILE=%d " + "m_old=%.9e m_new=%.9e alpha=%.9e " + "tile_sum=%.9e l=%.9e\n", + batch, q_row, q_head, kv_tile_base, m_old, + m_new, alpha, tile_sum, l_reg[row_slot]); + } +#endif + +#pragma unroll + for (int d_slot = 0; + d_slot < Traits::kDPerLane; ++d_slot) + { + const int d = + warp_lane + d_slot * Traits::kWarpSize; + float tile_acc = 0.0f; + for (int key_in_tile = 0; + key_in_tile < Traits::kBc; ++key_in_tile) + { + const float p = platform::warpBroadcast( + score_and_p, key_in_tile); + if (d < params.head_dim) + { + tile_acc = fmaf( + p, + v_tile[key_in_tile * HeadDim + d], + tile_acc); + } + } + if (d < params.head_dim) + { + acc_reg[row_slot][d_slot] = fmaf( + alpha, acc_reg[row_slot][d_slot], tile_acc); +#if defined(FA2_TRACE_DEVICE) + const int trace_d = + params.q_len == 64 ? 14 : 28; + if (trace_row && d == trace_d) + { + printf("FA2_ACC B=%d Q=%d H=%d D=%d TILE=%d " + "tile_acc=%.9e acc=%.9e norm=%.9e\n", + batch, q_row, q_head, d, kv_tile_base, + tile_acc, + acc_reg[row_slot][d_slot], + acc_reg[row_slot][d_slot] / + l_reg[row_slot]); + } +#endif + } + } + } + __syncthreads(); + } + +#pragma unroll + for (int row_slot = 0; row_slot < Traits::kRowsPerWarp; + ++row_slot) + { + const int row_in_cta = + row_slot * Traits::kWarpsPerCta + warp_id; + const int q_row = cta_q_base + row_in_cta; + if (q_row < params.q_len) + { + const float inv_l = l_reg[row_slot] > 0.0f + ? 1.0f / l_reg[row_slot] + : 0.0f; +#pragma unroll + for (int d_slot = 0; + d_slot < Traits::kDPerLane; ++d_slot) + { + const int d = + warp_lane + d_slot * Traits::kWarpSize; + if (d >= params.head_dim) + { + continue; + } + const std::size_t output_index = + ((static_cast(batch) * params.q_len + + q_row) * + params.q_heads + + q_head) * + params.head_dim + + d; + params.output[output_index] = + fromFloat(acc_reg[row_slot][d_slot] * inv_l); + } + } + } + } + + template + cudaError_t launchKernel(const Fa2Params ¶ms, + cudaStream_t stream) + { + using Traits = Fa2Traits; + const dim3 block(Traits::kThreads); + const dim3 grid((params.q_len + Traits::kRowsPerCta - 1) / + Traits::kRowsPerCta, + params.q_heads, params.batch_size); + fa2CompatibleKernel<<>>(params); + return cudaGetLastError(); + } + + } // namespace + + template + cudaError_t launch(T *output, const T *query, const T *key, + const T *value, int batch_size, int q_len, int kv_len, + int q_heads, int kv_heads, int head_dim, bool is_causal, + float scale, cudaStream_t stream) + { + if (batch_size == 0 || q_len == 0) + { + return cudaSuccess; + } + if (output == nullptr || query == nullptr || key == nullptr || + value == nullptr || batch_size < 0 || q_len < 0 || kv_len <= 0 || + q_heads <= 0 || kv_heads <= 0 || head_dim <= 0 || + q_heads % kv_heads != 0) + { + return cudaErrorInvalidValue; + } + + Fa2Params params{output, query, key, value, batch_size, q_len, kv_len, + q_heads, kv_heads, head_dim, is_causal, scale}; + + if (head_dim <= 1) + return launchKernel(params, stream); + if (head_dim <= 2) + return launchKernel(params, stream); + if (head_dim <= 4) + return launchKernel(params, stream); + if (head_dim <= 8) + return launchKernel(params, stream); + if (head_dim <= 16) + return launchKernel(params, stream); + if (head_dim <= 32) + return launchKernel(params, stream); + if (head_dim <= 64) + return launchKernel(params, stream); + if (head_dim <= 128) + return launchKernel(params, stream); + if (head_dim <= 256) + return launchKernel(params, stream); + return cudaErrorInvalidValue; + } + + template cudaError_t launch( + float *, const float *, const float *, const float *, int, int, int, + int, int, int, bool, float, cudaStream_t); + + template cudaError_t launch<__half>( + __half *, const __half *, const __half *, const __half *, int, int, + int, int, int, int, bool, float, cudaStream_t); + +#if defined(FA2_COMPATIBLE_HAS_BFLOAT16) + template cudaError_t launch<__nv_bfloat16>( + __nv_bfloat16 *, const __nv_bfloat16 *, const __nv_bfloat16 *, + const __nv_bfloat16 *, int, int, int, int, int, int, bool, float, + cudaStream_t); +#endif + +} // namespace fa2_compatible + +#if defined(FA2_COMPATIBLE_HAS_BFLOAT16) +#undef FA2_COMPATIBLE_HAS_BFLOAT16 +#endif diff --git a/onlinesoftmax_prob/RMSNorm.cu b/onlinesoftmax_prob/RMSNorm.cu new file mode 100644 index 00000000..03961711 --- /dev/null +++ b/onlinesoftmax_prob/RMSNorm.cu @@ -0,0 +1,135 @@ +#include +#include + +#ifdef PLATFORM_METAX +static constexpr int WARP_SIZE = 64; +#elif defined(PLATFORM_ILUVATAR) +static constexpr int WARP_SIZE = 64; +#elif defined(PLATFORM_NVIDIA) +static constexpr int WARP_SIZE = 32; +#endif + +#if defined(PLATFORM_METAX) || defined(PLATFORM_ILUVATAR) +using WarpMask = unsigned long long; +#elif defined(PLATFORM_NVIDIA) +using WarpMask = unsigned int; +#endif + +static constexpr WarpMask FULL_MASK = ~WarpMask{0}; + +namespace +{ + + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = (kBlockSize + WARP_SIZE - 1) / WARP_SIZE; + + static_assert(kBlockSize % WARP_SIZE == 0, + "Full-mask reductions require complete warps"); + static_assert(kWarpsPerBlock <= WARP_SIZE, + "The second reduction stage must fit in one warp"); + +} // namespace +// 这里支持的是矩阵形式的RMSNorm +template +__device__ __forceinline__ float warpReduceSumXor(float value) +{ + static_assert(Width > 0 && (Width & (Width - 1)) == 0, + "Width must be a power of two"); + static_assert(Width <= WARP_SIZE, + "Width cannot exceed the hardware warp size"); + +#pragma unroll + for (int offset = Width / 2; offset > 0; offset >>= 1) + { + value += __shfl_xor_sync(FULL_MASK, value, offset, Width); + } + return value; +} + +__device__ float blockReduceSquareSum(float v) +{ + __shared__ float warp_sums[kWarpsPerBlock]; + const int warp_id = threadIdx.x / WARP_SIZE; + const int lane_id = threadIdx.x % WARP_SIZE; + + v = warpReduceSumXor(v); + if (lane_id == 0) + warp_sums[warp_id] = v; + __syncthreads(); + + if (warp_id == 0) + { + // Every lane in warp 0 participates; unused lanes contribute zero. + v = lane_id < kWarpsPerBlock ? warp_sums[lane_id] : 0.0f; + v = warpReduceSumXor(v); + } + return v; +} + +__device__ __forceinline__ float toFloat(float x) +{ + return x; +} +__device__ __forceinline__ float toFloat(half x) +{ + return __half2float(x); +} + +template +__device__ __forceinline__ + T + fromFloat(float value); + +template <> +__device__ __forceinline__ float fromFloat(float value) +{ + return value; +} + +template <> +__device__ __forceinline__ half fromFloat(float value) +{ + return __float2half_rn(value); +} + +template +__global__ void RMSNorm_kernel(const T *input, const T *weight, T *output, int hidden_dim, float eps = 1e-5) +{ + // 1.这里是处理一个Block(也就是一行的kernel) + int tid = threadIdx.x; + int gridStride = blockDim.x; + float partial_sum = 0.0f; + int rowoffset = blockIdx.x * hidden_dim; + + // 2.将可能超过线程总数的部分 提前压缩到线程数可以处理的部分 + for (int i = tid; i < hidden_dim; i += gridStride) + { + float val = toFloat(input[rowoffset + i]); + partial_sum += val * val; + } + // 只有tid == 0时存有正确的sum + float tmp = blockReduceSquareSum(partial_sum); + __shared__ float sum; + if (tid == 0) + sum = rsqrtf(tmp / hidden_dim + eps); // 全局知道了RMS作为系数 + + __syncthreads(); + for (int i = tid; i < hidden_dim; i += gridStride) + { + float value = toFloat(input[rowoffset + i]); + float scale = toFloat(weight[i]); + output[rowoffset + i] = fromFloat(value * sum * scale); + } +} + +template +void launchRMSNorm(const T *input, const T *weight, T *output, int rows, int hidden_dim, float eps) +{ + dim3 block(kBlockSize); + dim3 grid(rows); + + RMSNorm_kernel<<>>(input, weight, output, hidden_dim, eps); +} + +template void launchRMSNorm(const float *, const float *, float *, int, int, float); +template void launchRMSNorm(const half *, const half *, half *, int, int, float); diff --git a/src/kernels.cu b/src/kernels.cu index 2cc53e7e..78de384d 100644 --- a/src/kernels.cu +++ b/src/kernels.cu @@ -1,15 +1,19 @@ +// #include <__clang_cuda_builtin_vars.h> +// #include <__clang_cuda_math.h> +#include +#include #include #include #include "../tester/utils.h" - +#define WARPSIZE 32 /** * @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: + * i and column col: * - * output[i, j] = input[i, j] * rsqrt(mean(input[i, :]^2) + eps) * weight[j] + * output[i, col] = input[i, col] * rsqrt(mean(input[i, :]^2) + eps) * weight[col] * * The output vector is preallocated with rows * hidden_dim elements. * @@ -21,48 +25,401 @@ * @param[in] hidden_dim Size of the normalized dimension. * @param[in] eps Numerical stability epsilon. */ + +// 核心目的是：参数内部都已FP32的精度进行处理，保证精度的同时来去适应half，float类型 + +template +__device__ float toFloat(T x) +{ + return static_cast(x); +} + +template <> +__device__ float toFloat(half x) +{ + return __half2float(x); +} + 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 +__device__ T fromFloat(float x) +{ + return static_cast(x); +} + +template <> +__device__ half fromFloat(float x) +{ + return __float2half(x); +} + +template +__device__ __forceinline__ float warpReduceSum(float val) +{ +#pragma unroll + for (int offset = WARPSIZE / 2; offset > 0; offset >>= 1) + { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +template +__global__ void rmsNormKernel(const T *input, const T *weight, + T *output, size_t rows, size_t hidden_dim, + float eps) +{ + const size_t row = blockIdx.x; + const int tid = threadIdx.x; + + if (row >= rows) + return; + + const int lane = tid % 32; + const int warp_id = tid / 32; + const int num_warps = blockDim.x / WARPSIZE; + + // 最多支持 1024 threads = 32 warps + __shared__ float warp_sums[32]; + __shared__ float rms_scale; // 一行全局共享 + + const T *row_input = input + row * hidden_dim; + T *row_output = output + row * hidden_dim; + + float square_sum = 0.0f; + + // 每个线程覆盖多个元素，如果一行的数据比较长，这里的stride = blockDim.x + for (size_t col = tid; col < hidden_dim; col += blockDim.x) + { + float x = toFloat(row_input[col]); + square_sum += x * x; + } + + square_sum = warpReduceSum(square_sum); + + // 每个 warp 只写一个值到 shared memory，完成第一层归约 + if (lane == 0) + { + warp_sums[warp_id] = square_sum; + } + + __syncthreads(); + + // warp 0 对各个 warp 的结果继续 reduction，二层归约 + if (warp_id == 0) + { + float block_sum = + lane < num_warps ? warp_sums[lane] : 0.0f; + + block_sum = warpReduceSum(block_sum); + + if (lane == 0) + { + rms_scale = rsqrtf(block_sum / static_cast(hidden_dim) + eps); + } + } + + __syncthreads(); + + // 类似的线程处理自己负责行的多个数据 + for (size_t col = tid; col < hidden_dim; col += blockDim.x) + { + float x = toFloat(row_input[col]); + float w = toFloat(weight[col]); + + row_output[col] = fromFloat(x * rms_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) +{ + size_t input_bytes = rows * hidden_dim * sizeof(T); + + size_t weight_bytes = hidden_dim * sizeof(T); + + // 分配指针 + T *d_input = nullptr; + T *d_weight = nullptr; + T *d_output = nullptr; + + // 给指针分配内存空间 + cudaMalloc(&d_input, input_bytes); + cudaMalloc(&d_weight, weight_bytes); + cudaMalloc(&d_output, input_bytes); + + // 复制数据 + cudaMemcpy(d_input, h_input.data(), input_bytes, cudaMemcpyHostToDevice); + + cudaMemcpy(d_weight, h_weight.data(), weight_bytes, cudaMemcpyHostToDevice); + + // 一个block负责一行 + dim3 block(256); + dim3 grid(rows); + + size_t shared_bytes = block.x * sizeof(float); + + rmsNormKernel<<>>(d_input, d_weight, d_output, rows, hidden_dim, eps); + + cudaDeviceSynchronize(); + + h_output.resize(rows * hidden_dim); + + cudaMemcpy(h_output.data(), d_output, input_bytes, cudaMemcpyDeviceToHost); + + 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] + * + * @tparam T Data type for input/output tensors + * @param[in] h_q Query tensor of shape + * [batch_size, target_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 tensor of shape + * [batch_size, target_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] 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] kv_heads Number of key/value heads + * @param[in] head_dim Dimension of each attention head * @param[in] is_causal Whether to apply causal masking */ + +// 一个Block负责一个输出向量: +// +// O[batch_id, q_pos, q_head, :] +// +// 当前版本主要以测试正确性和浮点结果一致性为目标， +// 不使用Online Softmax，也不保存完整Score。 +// +// 计算过程分成三步: +// +// 1. 第一次QK：找到当前Attention Row的row_max +// 2. 第二次QK：计算Softmax归一化和 +// 3. 第三次QK：计算weight并直接累加P * V +// +// QK存在重复计算，后续再考虑性能优化。 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 +__global__ void flashAttentionKernel( + const T *q, + const T *k, + const T *v, + T *o, + int batch_size, + int target_seq_len, + int src_seq_len, + int query_heads, + int kv_heads, + int head_dim, + bool is_causal) +{ + extern __shared__ float smem[]; + + int task_id = blockIdx.x; + + // 当前Block负责哪个Q head + int q_head = task_id % query_heads; + // 当前Block负责哪个Query Row + int q_pos = (task_id / query_heads) % target_seq_len; + // 当前Block属于哪个Batch + int batch_id = task_id / (query_heads * target_seq_len); + + if (batch_id >= batch_size) + { + return; + } + // query_heads == kv_heads: + // 一个Q head对应一个KV head --- MHA + // query_heads > kv_heads: + // 多个Q head共享一个KV head --- GQA + int kv_head = q_head * kv_heads / query_heads; + + // Causal Attention只能访问当前位置以及之前的Key + int key_count = is_causal ? min(q_pos + 1, src_seq_len) : src_seq_len; + + float attn_scale = rsqrtf(static_cast(head_dim)); + + const std::size_t q_offset = ((static_cast(batch_id) * target_seq_len + q_pos) * query_heads + q_head) * head_dim; + // 计算QK：寻找当前Attention Row中的最大即row_max + if (threadIdx.x == 0) + { + float row_max = -INFINITY; + for (int k_pos = 0; k_pos < key_count; k_pos++) + { + const std::size_t k_offset = ((static_cast(batch_id) * src_seq_len + k_pos) * kv_heads + kv_head) * head_dim; + float logit = 0.0f; + + // 一个Q Row和一个K Row进行dot + for (int dim = 0; dim < head_dim; dim++) + { + float q_val = toFloat(q[q_offset + dim]); + float k_val = toFloat(k[k_offset + dim]); + + logit += q_val * k_val; + } + + logit *= attn_scale; + + row_max = fmaxf(row_max, logit); + } + // 计算QK，并记录Softmax归一化和 + + float softmax_sum = 0.0f; + + for (int k_pos = 0; k_pos < key_count; k_pos++) + { + const std::size_t k_offset = ((static_cast(batch_id) * src_seq_len + k_pos) * kv_heads + kv_head) * head_dim; + float logit = 0.0f; + + for (int dim = 0; dim < head_dim; dim++) + { + float q_val = toFloat(q[q_offset + dim]); + float k_val = toFloat(k[k_offset + dim]); + + logit += q_val * k_val; + } + + logit *= attn_scale; + + softmax_sum += expf(logit - row_max); // safe softmax + } + + // 后续Block内部线程都需要使用 + smem[0] = row_max; + smem[1] = softmax_sum; + } + + // 等待thread 0完成当前Attention Row的row_max和softmax_sum + __syncthreads(); + + // 每个线程在寄存器中独立持有 + float row_max = smem[0]; + float softmax_sum = smem[1]; + + // 第三次计算QK，同时完成 P * V + // 每个线程负责若干个head_dim位置 + for (int dim = threadIdx.x; dim < head_dim; dim += blockDim.x) + { + float out_val = 0.0f; + + // 当前输出维度需要遍历所有有效的KV Row + for (int k_pos = 0; k_pos < key_count; k_pos++) + { + // K/V布局相同: + // [B, Lkv, Hkv, D] + // 因此可以共用同一个KV起始位置 + const std::size_t kv_offset = ((static_cast(batch_id) * src_seq_len + k_pos) * kv_heads + kv_head) * head_dim; + // 重新计算当前Q Row与K Row的logit + float logit = 0.0f; + + for (int k_dim = 0; k_dim < head_dim; k_dim++) + { + float q_val = toFloat(q[q_offset + k_dim]); + float k_val = toFloat(k[kv_offset + k_dim]); + + logit += q_val * k_val; + } + + logit *= attn_scale; + + // weight = softmax(QK / sqrt(head_dim)) + float weight = expf(logit - row_max) / softmax_sum; + float v_val = toFloat(v[kv_offset + dim]); + + out_val += weight * v_val; + } + + // [B, Lq, Hq, D] + o[q_offset + dim] = fromFloat(out_val); + } +} + +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) +{ + const size_t query_elements = static_cast(batch_size) * target_seq_len * query_heads * head_dim; + const size_t kv_elements = static_cast(batch_size) * src_seq_len * kv_heads * head_dim; + const size_t query_bytes = query_elements * sizeof(T); + const size_t kv_bytes = kv_elements * sizeof(T); + + T *d_q = nullptr; + T *d_k = nullptr; + T *d_v = nullptr; + T *d_o = nullptr; + + // 分配GPU内存 + cudaMalloc(&d_q, query_bytes); + cudaMalloc(&d_k, kv_bytes); + cudaMalloc(&d_v, kv_bytes); + cudaMalloc(&d_o, query_bytes); + + // Host -> Device + cudaMemcpy(d_q, h_q.data(), query_bytes, cudaMemcpyHostToDevice); + cudaMemcpy(d_k, h_k.data(), kv_bytes, cudaMemcpyHostToDevice); + cudaMemcpy(d_v, h_v.data(), kv_bytes, cudaMemcpyHostToDevice); + + // 一个Block负责: + // O[batch_id, q_pos, q_head, :] + int total_blocks = batch_size * target_seq_len * query_heads; + dim3 grid(total_blocks); + + // 当前主要的线程并行发生在输出的head_dim上。 + // 这里保留原版本的block_size选择方式， + // 后续再单独进行性能上的调整。 + int parallel_size = src_seq_len > head_dim ? src_seq_len : head_dim; + int block_size = 32; + + // 保证并行维度 足够 + while (block_size < parallel_size && block_size < 256) + { + block_size <<= 1; + } + + dim3 block(block_size); + + // 暂时使用非online的形式处理结果 + // smem[0] = row_max + // smem[1] = softmax_sum + size_t shared_bytes = 2 * sizeof(float); + + 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); + + h_o.resize(query_elements); + + // Device -> Host + // 记录计算结果，便于后续处理 + cudaMemcpy(h_o.data(), d_o, query_bytes, cudaMemcpyDeviceToHost); + + cudaFree(d_q); + cudaFree(d_k); + cudaFree(d_v); + cudaFree(d_o); } // ********************************************************************* // Explicit Template Instantiations (REQUIRED FOR LINKING WITH TESTER.O) // DO NOT MODIFY THIS SECTION // ********************************************************************* -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&, - std::vector&, size_t, size_t, float); -template void flashAttention(const std::vector&, const std::vector&, - const std::vector&, std::vector&, - int, int, int, int, int, int, bool); -template void flashAttention(const std::vector&, const std::vector&, - const std::vector&, std::vector&, - int, int, int, int, int, int, bool); +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 &, + std::vector &, size_t, size_t, float); +template void flashAttention(const std::vector &, const std::vector &, + const std::vector &, std::vector &, + int, int, int, int, int, int, bool); +template void flashAttention(const std::vector &, const std::vector &, + const std::vector &, std::vector &, + int, int, int, int, int, int, bool); \ No newline at end of file diff --git a/src/kernels.maca b/src/kernels.maca index 4c320f21..de427c61 100644 --- a/src/kernels.maca +++ b/src/kernels.maca @@ -1,15 +1,20 @@ +#include +#include #include -#include +#include +#include #include "../tester/utils.h" +#define WARPSIZE 64 + /** * @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: + * i and column col: * - * output[i, j] = input[i, j] * rsqrt(mean(input[i, :]^2) + eps) * weight[j] + * output[i, col] = input[i, col] * rsqrt(mean(input[i, :]^2) + eps) * weight[col] * * The output vector is preallocated with rows * hidden_dim elements. * @@ -21,48 +26,386 @@ * @param[in] hidden_dim Size of the normalized dimension. * @param[in] eps Numerical stability epsilon. */ + +// 核心目的是：参数内部都以FP32的精度进行处理，保证精度的同时来去适应half，float类型 + +template +__device__ float toFloat(T x) +{ + return static_cast(x); +} + +template <> +__device__ float toFloat(half x) +{ + return __half2float(x); +} + +template +__device__ T fromFloat(float x) +{ + return static_cast(x); +} + +template <> +__device__ half fromFloat(float x) +{ + return __float2half(x); +} + +template +__device__ __forceinline__ float warpReduceSum(float val) +{ +#pragma unroll + for (int offset = WARPSIZE / 2; offset > 0; offset >>= 1) + { + val += __shfl_down_sync(0xffffffffffffffff, val, offset); + } + + return val; +} + +template +__global__ void rmsNormKernel(const T *input, const T *weight, + T *output, size_t rows, size_t hidden_dim, + float eps) +{ + const size_t row = blockIdx.x; + const int tid = threadIdx.x; + + if (row >= rows) + { + return; + } + + // MACA一个warp中有64个线程 + const int lane = tid % WARPSIZE; + const int warp_id = tid / WARPSIZE; + const int num_warps = blockDim.x / WARPSIZE; + + // 保存每个warp归约后的局部平方和 + __shared__ float warp_sums[WARPSIZE]; + // 一行所有线程共享RMSNorm缩放系数 + __shared__ float rms_scale; + + const T *row_input = input + row * hidden_dim; + T *row_output = output + row * hidden_dim; + + float square_sum = 0.0f; + + // 每个线程覆盖多个元素，如果一行的数据比较长，这里的stride = blockDim.x + for (size_t col = tid; col < hidden_dim; col += blockDim.x) + { + float x = toFloat(row_input[col]); + square_sum += x * x; + } + + // 每个warp内部完成第一层归约 + square_sum = warpReduceSum(square_sum); + + // 每个warp只写一个值到shared memory，完成第一层归约 + if (lane == 0) + { + warp_sums[warp_id] = square_sum; + } + + __syncthreads(); + + // warp 0对各个warp的结果继续reduction，完成第二层归约 + if (warp_id == 0) + { + float block_sum = + lane < num_warps ? warp_sums[lane] : 0.0f; + + block_sum = warpReduceSum(block_sum); + + if (lane == 0) + { + rms_scale = + rsqrtf(block_sum / static_cast(hidden_dim) + eps); + } + } + + __syncthreads(); + + // 类似的线程处理自己负责行的多个数据 + for (size_t col = tid; col < hidden_dim; col += blockDim.x) + { + float x = toFloat(row_input[col]); + float w = toFloat(weight[col]); + + row_output[col] = + fromFloat(x * rms_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 +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) +{ + size_t input_bytes = rows * hidden_dim * sizeof(T); + + size_t weight_bytes = hidden_dim * sizeof(T); + + // 分配指针 + T *d_input = nullptr; + T *d_weight = nullptr; + T *d_output = nullptr; + + // 给指针分配内存空间 + mcMalloc(&d_input, input_bytes); + mcMalloc(&d_weight, weight_bytes); + mcMalloc(&d_output, input_bytes); + + // 复制数据 + mcMemcpy(d_input, h_input.data(), input_bytes, mcMemcpyHostToDevice); + + mcMemcpy(d_weight, h_weight.data(), weight_bytes, mcMemcpyHostToDevice); + + // 一个block负责一行 + // MACA的WARPSIZE = 64，block = 256时一个Block中共有4个warp + dim3 block(256); + dim3 grid(rows); + + rmsNormKernel + <<>>(d_input, d_weight, d_output, rows, hidden_dim, eps); + + mcDeviceSynchronize(); + + h_output.resize(rows * hidden_dim); + + mcMemcpy(h_output.data(), d_output, input_bytes, mcMemcpyDeviceToHost); + + mcFree(d_input); + mcFree(d_weight); + mcFree(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] + * + * @tparam T Data type for input/output tensors + * @param[in] h_q Query tensor of shape + * [batch_size, target_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 tensor of shape + * [batch_size, target_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] 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] kv_heads Number of key/value heads + * @param[in] head_dim Dimension of each attention head * @param[in] is_causal Whether to apply causal masking */ + +// 一个Block负责一个输出向量: +// 目前基本的attention计算过程分成三步: +// 1. 第一次QK：找到当前Attention Row的row_max +// 2. 第二次QK：计算Softmax归一化和 +// 3. 第三次QK：计算weight并直接累加P * V +// QK存在重复计算，后续再考虑性能优化。 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 +__global__ void flashAttentionKernel(const T *q, const T *k, const T *v, T *o, int batch_size, int target_seq_len, + int src_seq_len, int query_heads, int kv_heads, int head_dim, bool is_causal) +{ + extern __shared__ float smem[]; + + int task_id = blockIdx.x; + + int q_head = task_id % query_heads; + int q_pos = (task_id / query_heads) % target_seq_len; + int batch_id = task_id / (query_heads * target_seq_len); + + if (batch_id >= batch_size) + { + return; + } + + // query_heads == kv_heads: + // 一个Q head对应一个KV head + // query_heads > kv_heads: + // 多个Q head共享一个KV head + int kv_head = q_head * kv_heads / query_heads; + + // Causal Attention只能访问当前位置以及之前的Key + int key_count = is_causal ? min(q_pos + 1, src_seq_len) : src_seq_len; + float attn_scale = rsqrtf(static_cast(head_dim)); + + const std::size_t q_offset = ((static_cast(batch_id) * target_seq_len + q_pos) * query_heads + q_head) * head_dim; + + // 计算QK：寻找Row中的最大logit + if (threadIdx.x == 0) + { + float row_max = -INFINITY; + + for (int k_pos = 0; k_pos < key_count; k_pos++) + { + const std::size_t k_offset = ((static_cast(batch_id) * src_seq_len + k_pos) * kv_heads + kv_head) * head_dim; + float logit = 0.0f; + + // 一个Q Row和一个K Row进行dot + for (int dim = 0; dim < head_dim; dim++) + { + float q_val = toFloat(q[q_offset + dim]); + float k_val = toFloat(k[k_offset + dim]); + + logit += q_val * k_val; + } + + logit *= attn_scale; + row_max = fmaxf(row_max, logit); + } + + // 计算QK：计算Softmax归一化和 + float softmax_sum = 0.0f; + + for (int k_pos = 0; k_pos < key_count; k_pos++) + { + const std::size_t k_offset = ((static_cast(batch_id) * src_seq_len + k_pos) * kv_heads + kv_head) * head_dim; + float logit = 0.0f; + + for (int dim = 0; dim < head_dim; dim++) + { + float q_val = toFloat(q[q_offset + dim]); + float k_val = toFloat(k[k_offset + dim]); + + logit += q_val * k_val; + } + + logit *= attn_scale; + + // 所有logit减去row_max， + // 避免exp过程中产生数值溢出 + softmax_sum += expf(logit - row_max); + } + + // 后续Block内部线程都需要使用 + smem[0] = row_max; + smem[1] = softmax_sum; + } + + __syncthreads(); + + float row_max = smem[0]; + float softmax_sum = smem[1]; + = + // 计算QK，同时完成 P * V + // 每个线程负责若干个head_dim位置 + for (int dim = threadIdx.x; dim < head_dim; dim += blockDim.x) + { + float out_val = 0.0f; + // 当前输出维度需要遍历所有有效的KV Row + for (int k_pos = 0; k_pos < key_count; k_pos++) + { + // K/V布局相同: + // + // [B, Lkv, Hkv, D] + // + // 因此可以共用同一个KV起始位置 + const std::size_t kv_offset = ((static_cast(batch_id) * src_seq_len + k_pos) * kv_heads + kv_head) * head_dim; + + // 重新计算当前Q Row与K Row的logit + float logit = 0.0f; + + for (int k_dim = 0; k_dim < head_dim; k_dim++) + { + float q_val = toFloat(q[q_offset + k_dim]); + float k_val = toFloat(k[kv_offset + k_dim]); + + logit += q_val * k_val; + } + + logit *= attn_scale; + + float weight = expf(logit - row_max) / softmax_sum; + float v_val = toFloat(v[kv_offset + dim]); + + out_val += weight * v_val; + } + + o[q_offset + dim] = fromFloat(out_val); + } +} + +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) +{ + const size_t query_elements = static_cast(batch_size) * target_seq_len * query_heads * head_dim; + const size_t kv_elements = static_cast(batch_size) * src_seq_len * kv_heads * head_dim; + const size_t query_bytes = query_elements * sizeof(T); + const size_t kv_bytes = kv_elements * sizeof(T); + + T *d_q = nullptr; + T *d_k = nullptr; + T *d_v = nullptr; + T *d_o = nullptr; + + // 分配GPU内存 + mcMalloc(&d_q, query_bytes); + mcMalloc(&d_k, kv_bytes); + mcMalloc(&d_v, kv_bytes); + mcMalloc(&d_o, query_bytes); + + // Host -> Device + mcMemcpy(d_q, h_q.data(), query_bytes, mcMemcpyHostToDevice); + mcMemcpy(d_k, h_k.data(), kv_bytes, mcMemcpyHostToDevice); + mcMemcpy(d_v, h_v.data(), kv_bytes, mcMemcpyHostToDevice); + + // 一个Block负责: + // + // O[batch_id, q_pos, q_head, :] + int total_blocks = batch_size * target_seq_len * query_heads; + + dim3 grid(total_blocks); + + // 当前主要的线程并行发生在输出的head_dim上。 + // 后续再单独进行性能上的调整。 + int parallel_size = src_seq_len > head_dim ? src_seq_len : head_dim; + int block_size = WARPSIZE; + + while (block_size < parallel_size && block_size < 256) + { + block_size <<= 1; + } + + dim3 block(block_size); + + // 需要保存:smem[0] = row_max smem[1] = softmax_sum + size_t shared_bytes = 2 * sizeof(float); + + 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); + + h_o.resize(query_elements); + + // Device -> Host + mcMemcpy(h_o.data(), d_o, query_bytes, mcMemcpyDeviceToHost); + + mcFree(d_q); + mcFree(d_k); + mcFree(d_v); + mcFree(d_o); } // ********************************************************************* // Explicit Template Instantiations (REQUIRED FOR LINKING WITH TESTER.O) // DO NOT MODIFY THIS SECTION // ********************************************************************* -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&, - std::vector&, size_t, size_t, float); -template void flashAttention(const std::vector&, const std::vector&, - const std::vector&, std::vector&, - int, int, int, int, int, int, bool); -template void flashAttention(const std::vector&, const std::vector&, - const std::vector&, std::vector&, - int, int, int, int, int, int, bool); +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 &, + std::vector &, size_t, size_t, float); +template void flashAttention(const std::vector &, const std::vector &, + const std::vector &, std::vector &, + int, int, int, int, int, int, bool); +template void flashAttention(const std::vector &, const std::vector &, + const std::vector &, std::vector &, + int, int, int, int, int, int, bool); diff --git a/src/kernels.o b/src/kernels.o new file mode 100644 index 00000000..c58a6733 Binary files /dev/null and b/src/kernels.o differ diff --git a/test_kernels b/test_kernels new file mode 100755 index 00000000..013431cc Binary files /dev/null and b/test_kernels differ diff --git a/tester/tester_metax.o b/tester/tester_metax.o index 0112a230..07c79675 100644 Binary files a/tester/tester_metax.o and b/tester/tester_metax.o differ diff --git a/tester/tester_metax.o:Zone.Identifier b/tester/tester_metax.o:Zone.Identifier new file mode 100644 index 00000000..d6c1ec68 Binary files /dev/null and b/tester/tester_metax.o:Zone.Identifier differ