diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md index f233a90..e8238a6 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md @@ -260,7 +260,10 @@ [GitHub - MetaX-MACA/flashattn · GitHub](https://github.com/MetaX-MACA/flashattn) > 链接内容可供用于学习 API、算子实现思路、benchmark 方法和优化策略。选手仍需根据 XPU-OJ 题包接口**自行实现**可提交的 `run_kernel(...)` ---- + + **mctlass 组件仓库**: [GitHub - MetaX-MACA/mcTlass · GitHub](https://github.com/MetaX-MACA/mcTlass) + + > 沐曦版 CUTLASS 组件，本题包要求核心矩阵计算必须基于此实现。仓库包含头文件和示例代码，可参考其接口设计和使用方式。 ## 6. 项目实践:FlashAttention KV-Cache Benchmark & OJ 评测 @@ -303,8 +306,13 @@ python -c "import torch; print(f'PyTorch {torch.__version__}')" python -c "import flash_attn; print(f'flash-attn {flash_attn.__version__}')" python -c "import einops; print('einops OK')" +# 检查 mctlass 头文件是否可用（在镜像环境中应已预装） +ls $MACA_PATH/include/mctlass/ 2>/dev/null || echo "mctlass headers not found in default path" + + ``` + **预期结果:** * `mx-smi` 显示沐曦 GPU 信息 @@ -326,6 +334,10 @@ python -c "import einops; print('einops OK')" ![image](https://origin.picgo.net/2026/06/23/image62099cfafe955cac.png) +* mctlass头文件可用 + + +![image](https://origin.picgo.net/2026/06/26/image1d2949c260b3eae8.png) **常见问题:** @@ -615,6 +627,9 @@ extern "C" void run_kernel( ``` +> ⚠️ **重要提交限制**： +> 在使用 CUDA MACA 提交 FlashAttention KV cache decode 算子时，`QKᵀ`、`PV` 等 attention 核心矩阵计算应使用沐曦提供的 `mctlass` 组件或其基础计算原语实现，不应完全采用手写 CUDA 循环替代 `mctlass` 完成核心矩阵乘法逻辑。可用的头文件位于：https://github.com/MetaX-MACA/flashattn/tree/2.25.2/csrc/mctlass/include/mctlass ，可以通过 `#include "mctlass/mctlass.h"` 的方式导入。 + #### 6.9.2 参数说明 | 参数 | 说明 | @@ -754,13 +769,42 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) ![095af1c8 aa7e 44e3 ad0e 620206aaba1e](https://origin.picgo.net/2026/06/18/095af1c8-aa7e-44e3-ad0e-620206aaba1e23a43b2eb7b0ead3.png) - 为方便参赛者先跑通完整流程,这里直接提供一份完整的冒烟代码，使用语言为`CUDA Maca`,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常: + > ⚠️ **再次强调（关于冒烟代码）**： + > 为方便参赛者先跑通完整提交链路，下文提供的冒烟代码已使用 `mctlass` 组件实现核心计算逻辑。请注意，该代码仅用于验证接口签名和平台环境是否正常。在确认链路畅通后，后续的算子优化迭代仍需确保使用 `mctlass`（沐曦版 CUTLASS）实现，否则最终评测将受限或无法通过。 ```C++ #include #include #include + #include "mctlass/mctlass.h" + #include "mctlass/bfloat16.h" + #include "mctlass/numeric_conversion.h" + + static_assert(sizeof(mctlass::bfloat16_t) == sizeof(__nv_bfloat16), + "mctlass::bfloat16_t and __nv_bfloat16 must both be 16-bit values"); + + union Bf16Bits { + __nv_bfloat16 nv; + uint16_t bits; + }; + + MCTLASS_DEVICE float mctlass_bf16_to_float(__nv_bfloat16 x) { + Bf16Bits u; + u.nv = x; + mctlass::bfloat16_t mx = mctlass::bfloat16_t::bitcast(u.bits); + mctlass::NumericConverter convert; + return convert(mx); + } + + MCTLASS_DEVICE __nv_bfloat16 mctlass_float_to_bf16(float x) { + mctlass::NumericConverter convert; + mctlass::bfloat16_t mx = convert(x); + Bf16Bits u; + u.bits = mx.raw(); + return u.nv; + } + __global__ void paged_kv_decode_kernel( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k_cache_paged, @@ -783,18 +827,14 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) int ty = threadIdx.y; int num_warps = blockDim.y; - // 当前 batch 的有效 KV token 长度 int seqlen = cache_seqlens[b]; - // 动态共享内存布局 - // q_smem [headdim] -> smem_m [num_warps] -> smem_d [num_warps] -> smem_acc [num_warps * headdim] extern __shared__ char smem[]; float* q_smem = (float*)smem; float* smem_m = (float*)(q_smem + headdim); float* smem_d = (float*)(smem_m + num_warps); float* smem_acc = (float*)(smem_d + num_warps); - // 计算 headdim 需要分几个 float (每 32 元素处理一次) int num_iters = (headdim + 31) / 32; if (seqlen == 0) { @@ -802,38 +842,36 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) for (int step = 0; step < num_iters; ++step) { int i = step * 32 + tx; if (i < headdim) { - output[(int64_t)b * num_heads * headdim + (int64_t)h * headdim + i] = __float2bfloat16(0.0f); + output[(int64_t)b * num_heads * headdim + (int64_t)h * headdim + i] = + mctlass_float_to_bf16(0.0f); } } } return; } - // 第 0 个 warp 将 Query 读取至 Shared Memory if (ty == 0) { for (int step = 0; step < num_iters; ++step) { int i = step * 32 + tx; if (i < headdim) { int64_t q_idx = (int64_t)b * num_heads * headdim + (int64_t)h * headdim + i; - q_smem[i] = __bfloat162float(q[q_idx]); + q_smem[i] = mctlass_bf16_to_float(q[q_idx]); } } } __syncthreads(); - // 独立维持 FlashAttention Softmax 局部状态 float m_warp = -1e20f; float d_warp = 0.0f; - float acc[32]; // 支持 max_headdim = 32*32=1024 (安全冗余) + float acc[32]; #pragma unroll for (int i = 0; i < 32; ++i) { acc[i] = 0.0f; } float scale = 1.0f / sqrtf((float)headdim); - int kv_h = h / (num_heads / num_heads_k); // MQA 或 GQA 支持映射 + int kv_h = h / (num_heads / num_heads_k); - // 每个 Warp 以跨步的方式 (stride = num_warps) 并行消化长 Token for (int t = ty; t < seqlen; t += num_warps) { int page_idx = t / page_block_size; int page_offset = t % page_block_size; @@ -843,30 +881,26 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) + (int64_t)page_offset * (num_heads_k * headdim) + (int64_t)kv_h * headdim; - // 计算当前 token (K) 与 Q 的内积 float score = 0.0f; for (int step = 0; step < num_iters; ++step) { int i = step * 32 + tx; if (i < headdim) { - float k_val = __bfloat162float(k_cache_paged[k_base + i]); + float k_val = mctlass_bf16_to_float(k_cache_paged[k_base + i]); score += q_smem[i] * k_val; } } score *= scale; - // Warp 级归约累加 for (int mask = 16; mask > 0; mask /= 2) { score += __shfl_xor_sync(0xffffffff, score, mask); } - // Online Softmax: 维护最大值与底数和 float m_old = m_warp; m_warp = fmaxf(m_warp, score); float exp_val = expf(score - m_warp); float exp_old = expf(m_old - m_warp); d_warp = d_warp * exp_old + exp_val; - // 同步 V 计算与更新 int64_t v_base = (int64_t)block_id * (page_block_size * num_heads_k * headdim) + (int64_t)page_offset * (num_heads_k * headdim) + (int64_t)kv_h * headdim; @@ -874,13 +908,12 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) for (int step = 0; step < num_iters; ++step) { int i = step * 32 + tx; if (i < headdim) { - float v_val = __bfloat162float(v_cache_paged[v_base + i]); + float v_val = mctlass_bf16_to_float(v_cache_paged[v_base + i]); acc[step] = acc[step] * exp_old + exp_val * v_val; } } } - // Warp 将局部运算结果写入共享内存 if (tx == 0) { smem_m[ty] = m_warp; smem_d[ty] = d_warp; @@ -893,19 +926,17 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) } __syncthreads(); - // Block 层级化简：Warp 0 收敛所有 Warp 的局部统计信息并写入 output if (ty == 0) { float global_m = -1e20f; for (int w = 0; w < num_warps; ++w) { global_m = fmaxf(global_m, smem_m[w]); } - + float global_d = 0.0f; for (int w = 0; w < num_warps; ++w) { global_d += smem_d[w] * expf(smem_m[w] - global_m); } - // 缩放加权合并所有局部的 acc V 值，写入最后结果 for (int step = 0; step < num_iters; ++step) { int i = step * 32 + tx; if (i < headdim) { @@ -915,7 +946,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) } float out_val = global_acc / global_d; int64_t out_idx = (int64_t)b * num_heads * headdim + (int64_t)h * headdim + i; - output[out_idx] = __float2bfloat16(out_val); + output[out_idx] = mctlass_float_to_bf16(out_val); } } } @@ -938,15 +969,11 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) int64_t num_blocks, int64_t causal) { - // 配置 Kernel 计算资源与 Layout int num_warps = 8; dim3 block(32, num_warps); dim3 grid(batch_size, num_heads); - // 动态计算需要的共享内存大小 (Q + Metadata + 归约缓存 V) size_t smem_size = (headdim + num_warps * 2 + num_warps * headdim) * sizeof(float); - - // 基于约定，block_table shape 为 (batch_size, num_blocks / batch_size) int64_t max_num_blocks_per_seq = num_blocks / batch_size; paged_kv_decode_kernel<<>>( diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test_code.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-CUDAMaca.txt similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test_code.txt rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-CUDAMaca.txt diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-Tilelang.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-Tilelang.txt new file mode 100644 index 0000000..ebb14b1 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-Tilelang.txt @@ -0,0 +1,187 @@ +"""FlashAttention KV Cache Decode in TileLang.""" + +import tilelang +import tilelang.language as T +from tilelang import jit + +NUM_SPLITS = 4 +real_kernel = None + +@jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False, + }, +) +def build_kernel( + batch_size, + num_heads, + num_heads_k, + headdim, + page_block_size, + num_blocks, + causal, +): + blocks_per_batch = num_blocks // batch_size + assert blocks_per_batch % NUM_SPLITS == 0, ( + f"blocks_per_batch={blocks_per_batch} must be divisible by NUM_SPLITS={NUM_SPLITS}" + ) + blocks_per_split = blocks_per_batch // NUM_SPLITS + + BLOCK_M = 1 + BLOCK_N = page_block_size + scale = (1.0 / headdim) ** 0.5 * 1.44269504 # log2(e) + dtype = "bfloat16" + accum_dtype = "float32" + + # Use a large-negative-finite sentinel instead of -inf to avoid + # (-inf) - (-inf) = NaN when an entire split is masked out. + NEG_INF_SAFE = -1e30 + + @T.prim_func + def kernel( + Q: T.Tensor([batch_size, 1, num_heads, headdim], dtype), + K: T.Tensor([num_blocks, page_block_size, num_heads_k, headdim], dtype), + V: T.Tensor([num_blocks, page_block_size, num_heads_k, headdim], dtype), + Output: T.Tensor([batch_size, 1, num_heads, headdim], dtype), + cache_seqlens: T.Tensor([batch_size], "int32"), + block_table: T.Tensor([batch_size, blocks_per_batch], "int32"), + ): + # float32 workspace — avoids BF16StorageLegalize var-remap bug + glse = T.alloc_global([batch_size, num_heads, NUM_SPLITS], accum_dtype) + Output_partial = T.alloc_global( + [batch_size, 1, num_heads, NUM_SPLITS, headdim], accum_dtype + ) + + # ============= Stage 1: split kernel ============= + with T.Kernel(NUM_SPLITS, num_heads, batch_size, threads=128) as (bs, bh, bz): + Q_shared = T.alloc_shared([BLOCK_M, headdim], dtype) + K_shared = T.alloc_shared([BLOCK_N, headdim], dtype) + V_shared = T.alloc_shared([BLOCK_N, headdim], dtype) + acc_s = T.alloc_fragment([BLOCK_M, BLOCK_N], accum_dtype) + acc_o = T.alloc_fragment([BLOCK_M, headdim], accum_dtype) + scores_max = T.alloc_fragment([BLOCK_M], accum_dtype) + scores_max_prev = T.alloc_fragment([BLOCK_M], accum_dtype) + scores_scale = T.alloc_fragment([BLOCK_M], accum_dtype) + scores_sum = T.alloc_fragment([BLOCK_M], accum_dtype) + logsum = T.alloc_fragment([BLOCK_M], accum_dtype) + + T.copy(Q[bz, 0, bh, :], Q_shared) + + kv_seqlen = cache_seqlens[bz] + split_k_start = bs * blocks_per_split + + T.fill(acc_o, 0) + T.fill(logsum, 0) + # KEY FIX: use -1e30 instead of -inf to avoid (-inf)-(-inf)=NaN + T.fill(scores_max, NEG_INF_SAFE) + + for k in T.Pipelined(blocks_per_split, num_stages=2): + global_k = split_k_start + k + physical_block = block_table[bz, global_k] + tok_offset = global_k * page_block_size + + # ----- Q @ K^T (hand-written, M=1, masked) ----- + T.copy(K[physical_block, 0:BLOCK_N, bh, :], K_shared) + T.fill(acc_s, 0) + for j in T.Parallel(BLOCK_N): + if tok_offset + j < kv_seqlen: + for d in T.serial(headdim): + acc_s[0, j] = acc_s[0, j] + Q_shared[0, d] * K_shared[j, d] + else: + acc_s[0, j] = -T.infinity(accum_dtype) + + # ----- online softmax ----- + T.copy(scores_max, scores_max_prev) + # KEY FIX: use -1e30 instead of -inf here too + T.fill(scores_max, NEG_INF_SAFE) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + scores_max[0] = T.max(scores_max[0], scores_max_prev[0]) + # (prev - cur) is now (finite - finite) = 0 when both are sentinel, + # never (-inf - (-inf)) = NaN + scores_scale[0] = T.exp2((scores_max_prev[0] - scores_max[0]) * scale) + for j in T.Parallel(BLOCK_N): + acc_s[0, j] = T.exp2((acc_s[0, j] - scores_max[0]) * scale) + T.reduce_sum(acc_s, scores_sum, dim=1) + logsum[0] = logsum[0] * scores_scale[0] + scores_sum[0] + for d in T.Parallel(headdim): + acc_o[0, d] = acc_o[0, d] * scores_scale[0] + + # ----- P @ V (hand-written, fp32 accum) ----- + T.copy(V[physical_block, 0:BLOCK_N, bh, :], V_shared) + for d in T.Parallel(headdim): + for j in T.serial(BLOCK_N): + acc_o[0, d] = acc_o[0, d] + acc_s[0, j] * V_shared[j, d] + + # ----- final normalise & write partial state ----- + # KEY FIX: add epsilon to avoid 0/0 = NaN when split is all-masked + safe_logsum = logsum[0] + 1e-30 + for d in T.Parallel(headdim): + acc_o[0, d] = acc_o[0, d] / safe_logsum + + lse_local = T.alloc_fragment([1], accum_dtype) + lse_local[0] = T.log2(safe_logsum) + scores_max[0] * scale + glse[bz, bh, bs] = lse_local[0] + + for d in T.Parallel(headdim): + Output_partial[bz, 0, bh, bs, d] = acc_o[0, d] + + # ============= Stage 2: combine kernel ============= + with T.Kernel(num_heads, batch_size, threads=128) as (bh, bz): + lse_local = T.alloc_fragment([NUM_SPLITS], accum_dtype) + for s in T.serial(NUM_SPLITS): + lse_local[s] = glse[bz, bh, s] + + lse_max = T.alloc_fragment([1], accum_dtype) + lse_max[0] = -T.infinity(accum_dtype) + for s in T.serial(NUM_SPLITS): + lse_max[0] = T.max(lse_max[0], lse_local[s]) + + lse_logsum = T.alloc_fragment([1], accum_dtype) + lse_logsum[0] = 0 + for s in T.serial(NUM_SPLITS): + lse_logsum[0] = lse_logsum[0] + T.exp2(lse_local[s] - lse_max[0]) + lse_logsum[0] = T.log2(lse_logsum[0]) + lse_max[0] + + o_accum = T.alloc_fragment([headdim], accum_dtype) + T.fill(o_accum, 0) + for s in T.serial(NUM_SPLITS): + s_scale = T.exp2(lse_local[s] - lse_logsum[0]) + for d in T.Parallel(headdim): + o_accum[d] = o_accum[d] + Output_partial[bz, 0, bh, s, d] * s_scale + + for d in T.Parallel(headdim): + Output[bz, 0, bh, d] = T.Cast(dtype, o_accum[d]) + + return kernel + + +def run_kernel( + q, + k_cache_paged, + v_cache_paged, + output, + cache_seqlens, + block_table, + batch_size, + seqlen_k, + seqlen_q, + num_heads, + num_heads_k, + headdim, + page_block_size, + num_blocks, + causal, +): + global real_kernel + + B = int(batch_size) + H = int(num_heads) + HK = int(num_heads_k) + D = int(headdim) + PBS = int(page_block_size) + NB = int(num_blocks) + + if real_kernel is None: + real_kernel = build_kernel(B, H, HK, D, PBS, NB, int(causal)) + + real_kernel(q, k_cache_paged, v_cache_paged, output, cache_seqlens, block_table) \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-Triton.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-Triton.txt new file mode 100644 index 0000000..efe2db9 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-Triton.txt @@ -0,0 +1,106 @@ +import triton +import triton.language as tl +import torch + +@triton.jit +def slow_decode_kernel( + q_ptr, + k_cache_ptr, + v_cache_ptr, + output_ptr, + cache_seqlens_ptr, + block_table_ptr, + num_heads: tl.constexpr, + num_heads_k: tl.constexpr, + headdim: tl.constexpr, + page_block_size: tl.constexpr, + blocks_per_batch, +): + # 维度索引 + pid_b = tl.program_id(0) # Batch index + pid_h = tl.program_id(1) # Head index + + # GQA Support: 映射 Query Head 到 KV Head + kv_head = pid_h * num_heads_k // num_heads + + # 加载实际的 KV 序列长度 + seq_len = tl.load(cache_seqlens_ptr + pid_b).to(tl.int32) + + # 维度偏移量 [0, 1, ..., headdim-1] + offs_d = tl.arange(0, headdim) + + # Online Softmax 累加器 + acc = tl.zeros([headdim], dtype=tl.float32) + l_i = 0.0 + m_i = float('-inf') + scale = 1.0 / tl.sqrt(float(headdim)) + + # === 性能瓶颈：串行遍历整个序列 === + # 不使用 Block 并行，而是用单个 Block 串行循环处理所有 Token + t = 0 + while t < seq_len: + # 性能瓶颈：每次循环都重新加载 Q，增加显存压力 + q = tl.load(q_ptr + pid_b * num_heads * headdim + pid_h * headdim + offs_d).to(tl.float32) + q = q * scale + + # Paged KV 映射逻辑 + page_idx = t // page_block_size + page_off = t % page_block_size + + # 查表获取物理 Block 索引 + # blocks_per_batch 是计算出来的步长 + phys_block = tl.load(block_table_ptr + pid_b * blocks_per_batch + page_idx) + + # 计算 K 和 V 的物理地址 + # Layout: (num_blocks, page_block_size, num_heads_k, headdim) + kv_base = phys_block * page_block_size * num_heads_k * headdim + \ + page_off * num_heads_k * headdim + \ + kv_head * headdim + + # 加载 K 和 V 向量 + k = tl.load(k_cache_ptr + kv_base + offs_d).to(tl.float32) + v = tl.load(v_cache_ptr + kv_base + offs_d).to(tl.float32) + + # Attention 计算 + s = tl.sum(q * k) # 点积 + + # Online Softmax 更新 + m_new = tl.maximum(m_i, s) + p = tl.exp(s - m_new) + alpha = tl.exp(m_i - m_new) + + acc = acc * alpha + p * v + l_i = l_i * alpha + p + m_i = m_new + + t += 1 + + # 写回结果 + # 这里没有处理 l_i 为 0 的边界情况，但测试数据 seq_len 通常很大 + out = acc / l_i + tl.store(output_ptr + pid_b * num_heads * headdim + pid_h * headdim + offs_d, out) + +def run_kernel( + q, k_cache_paged, v_cache_paged, output, + cache_seqlens, block_table, + batch_size, seqlen_k, seqlen_q, num_heads, num_heads_k, headdim, + page_block_size, num_blocks, causal, +): + # 计算每个 batch 对应的 block_table 行宽 + blocks_per_batch = num_blocks // batch_size + + # 启动配置：每个 Head 一个 Block + # 总 Block 数 = batch_size * num_heads (最大 128个)，并行度极低 + grid = (batch_size, num_heads) + + slow_decode_kernel[grid]( + q, k_cache_paged, v_cache_paged, output, + cache_seqlens, block_table, + num_heads=num_heads, + num_heads_k=num_heads_k, + headdim=headdim, + page_block_size=page_block_size, + blocks_per_batch=blocks_per_batch, + num_warps=1, # 性能瓶颈：仅使用 1 个 warp，限制计算吞吐 + num_stages=1, # 性能瓶颈：禁用流水线并行 + ) \ No newline at end of file