![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131722-69888ce28183f.jpg\")
<\/p>\n\u672c\u6587\u5b9e\u73b0 FlashAttention-2 \u7684\u524d\u5411\u4f20\u64ad，\u5177\u4f53\u5305\u62ec：\u4e3a Q\u3001K\u3001V \u8bbe\u8ba1\u5206\u5757\u7b56\u7565；\u6d41\u5f0f\u5904\u7406 K \u548c V \u5757\u800c\u975e\u7269\u5316\u5b8c\u6574\u6ce8\u610f\u529b\u77e9\u9635；\u5b9e\u73b0\u5728\u7ebf softmax \u7b97\u6cd5\u4fdd\u8bc1\u6570\u503c\u7a33\u5b9a\u6027；\u652f\u6301\u56e0\u679c\u548c\u975e\u56e0\u679c\u4e24\u79cd\u6ce8\u610f\u529b\u6a21\u5f0f；\u7528 Triton autotuner \u81ea\u52a8\u8c03\u4f18\u5185\u6838\u914d\u7f6e；\u6700\u540e\u7528 PyTorch \u9a8c\u8bc1\u6b63\u786e\u6027\u3002 <\/p>\n
FlashAttention vs. standard attention vs torch2.2 (spda flashattn) TFLOP\/s benchmarks<\/p>\n
\u6807\u51c6\u6ce8\u610f\u529b\u7684\u74f6\u9888\u4e0d\u5728\u6d6e\u70b9\u8fd0\u7b97\u91cf\u800c\u5728\u5185\u5b58\u5e26\u5bbd\u3002\u666e\u901a\u6ce8\u610f\u529b\u8ba1\u7b97 S = QK\u1d40 \u4e4b\u540e，\u8981\u628a\u5b8c\u6574\u7684 N \u00d7 N \u77e9\u9635\u5199\u5165 HBM\u518d\u8bfb\u56de\u6765\u7b97 softmax \u5e76\u5b58\u50a8\u7136\u540e\u518d\u8bfb\u4e00\u6b21\u4e58\u4ee5 V，\u6bcf\u4e2a\u5143\u7d20\u88ab\u8bbf\u95ee 2-4 \u6b21\u6bcf\u6b21\u90fd\u8d70 HBM\u3002<\/p>\n
\u5e8f\u5217\u957f\u5ea6 16K \u65f6，\u8fd9\u4e2a\u77e9\u9635\u5305\u542b 16,384\u00b2 \u2248 2.56 \u4ebf\u4e2a\u5143\u7d20\u3002<\/p>\n
\u53cd\u590d\u5728 HBM \u548c\u8ba1\u7b97\u5355\u5143\u4e4b\u95f4\u642c\u8fd0\u8fd9\u51e0\u4ebf\u4e2a\u503c，\u800cHBM \u662f GPU \u4e0a\u5bb9\u91cf\u6700\u5927\u7684\u5185\u5b58\u4e5f\u662f\u6700\u6162\u7684\u3002A100 \u4e0a\u4ece HBM \u8bfb\u6570\u636e\u6bd4\u4ece\u7247\u4e0a SRAM \u8bfb\u5927\u7ea6\u6162 15 \u500d\u3002\u5927\u5f20\u91cf\u548c\u6a21\u578b\u6743\u91cd\u90fd\u653e\u5728\u8fd9\u91cc，\u6240\u4ee5\u5199\u5185\u6838\u7684\u9996\u8981\u76ee\u6807\u5c31\u662f\u51cf\u5c11 HBM \u6d41\u91cf\u628a\u9ad8\u9891\u8bbf\u95ee\u7684\u6570\u636e\u7559\u5728\u5bc4\u5b58\u5668\u6216\u5171\u4eab\u5185\u5b58\u91cc\u3002<\/p>\n
FlashAttention \u7684\u6838\u5fc3\u601d\u60f3\u662f\u8ba9\u6ce8\u610f\u529b\u53d8\u5f97 IO \u611f\u77e5\u3002\u6240\u8c13 IO \u611f\u77e5\u5c31\u662f\u771f\u6b63\u7406\u89e3\u5e76\u5229\u7528\u4e00\u4e2a\u8fd9\u4e2a\u5b9a\u4e49：\u7247\u4e0a SRAM \u6bd4 HBM \u5feb\u51e0\u4e2a\u6570\u91cf\u7ea7\u3002NVIDIA A100 \u6709 40-80GB HBM（\u4e5f\u5c31\u662f\u90a3\u4e2a\u8ba9\u4f60\u9891\u7e41\u906d\u9047 CUDA OOM \u7684\u5168\u5c40\u5185\u5b58）\u5e26\u5bbd 1.5-2.0 TB\/s；\u6bcf\u4e2a SM \u6709 192KB SRAM，\u5171 108 \u4e2a SM，\u5e26\u5bbd\u4f30\u8ba1 19TB\/s \u5de6\u53f3\u3002<\/p>\n
GPU \u786c\u4ef6\u6709\u4e2a\u9ec4\u91d1\u6cd5\u5219：<\/p>\n
\u628a\u6570\u636e\u642c\u5230\u5185\u5b58\u5c42\u6b21\u7684\u4e0a\u5c42\u7136\u540e\u7559\u5728\u90a3\u91cc\u3002\u9664\u975e\u4e07\u4e0d\u5f97\u5df2\u522b\u56de HBM\u3002<\/p>\n
\u6807\u51c6\u6ce8\u610f\u529b\u5b8c\u5168\u65e0\u89c6\u8fd9\u6761\u89c4\u5219，\u628a HBM \u8bfb\u5199\u5f53\u6210\u96f6\u6210\u672c\u64cd\u4f5c\u3002FlashAttention \u8ba1\u7b97\u7684\u7ed3\u679c\u548c\u6807\u51c6\u7f29\u653e\u70b9\u79ef\u6ce8\u610f\u529b\u5b8c\u5168\u4e00\u6837：<\/p>\n
S = QK\u1d40 \u2208 \u211d\u1d3a\u02e3\u1d3a，P = softmax(S) \u2208 \u211d\u1d3a\u02e3\u1d3a，O = PV \u2208 \u211d\u1d3a\u02e3\u1d48<\/p>\n
\u533a\u522b\u5728\u4e8e\u8ba1\u7b97\u7684\u8c03\u5ea6\u65b9\u5f0f\u3002FlashAttention \u4e0d\u5728 HBM \u91cc\u5b58\u50a8\u90a3\u4e2a\u5de8\u5927\u7684 N \u00d7 N \u6ce8\u610f\u529b\u77e9\u9635\u7136\u540e\u518d\u8bfb\u56de\u6765\u7b97 softmax\u800c\u662f\u91cd\u65b0\u7ec4\u7ec7\u8ba1\u7b97：\u5206\u5757\u5904\u7406\u5e8f\u5217\u4ece\u5168\u5c40\u5185\u5b58\u6d41\u5f0f\u8bfb\u53d6 K \u548c V \u5757，\u7528\u5728\u7ebf softmax \u589e\u91cf\u8ba1\u7b97\u6bcf\u4e2a\u5757\u7684\u90e8\u5206\u7ed3\u679c，\u9010\u6b65\u6784\u5efa\u8f93\u51fa\u77e9\u9635 O\u53cd\u5411\u4f20\u64ad\u65f6\u8fd8\u53ef\u4ee5\u9009\u62e9\u91cd\u7b97\u800c\u975e\u5b58\u50a8\u3002<\/p>\n
\u5177\u4f53\u64cd\u4f5c\u662f\u8fd9\u6837\u7684：\u62ff\u4e00\u5757\u67e5\u8be2 Q_block，\u7136\u540e\u5206\u5757\u8fed\u4ee3 K \u548c V \u5e8f\u5217，\u8fb9\u8fed\u4ee3\u8fb9\u505a\u5728\u7ebf softmax \u540c\u65f6\u8ffd\u8e2a\u5fc5\u8981\u7684\u7edf\u8ba1\u91cf，\u7d2f\u79ef\u8f93\u51fa\u5757\u5e76\u5728\u7247\u4e0a\u5f52\u4e00\u5316，\u53ea\u628a\u6700\u7ec8\u7ed3\u679c\u5199\u56de HBM\u3002<\/p>\n
\u8fd9\u6837\u6ce8\u610f\u529b\u7684\u5185\u5b58\u590d\u6742\u5ea6\u5c31\u4ece O(N\u00b2) \u964d\u5230\u4e86 O(N)\u3002<\/p>\n
\u5206\u5757\u77e9\u9635\u4e58\u6cd5\u4e0d\u96be，\u800c\u5206\u5757 softmax \u624d\u662f\u9ebb\u70e6\u4e8b\u3002\u6ce8\u610f\u529b\u4e2d token i \u5bf9\u5176\u4ed6 token \u7684\u5173\u6ce8\u7a0b\u5ea6，\u662f\u5bf9\u8be5\u884c\u6240\u6709\u6ce8\u610f\u529b\u5206\u6570\u505a softmax \u5f97\u5230\u7684： ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131722-69888ce29a903.jpg\")
<\/p>\n
\u666e\u901a\u6ce8\u610f\u529b\u91cc\u8fd9\u5f88\u7b80\u5355，\u56e0\u4e3a\u4e00\u4e2a token \u7684\u5168\u90e8\u6ce8\u610f\u529b\u5206\u6570\u5df2\u7ecf\u7269\u5316\u5728\u5185\u5b58\u4e2d，\u4e00\u6b65\u5c31\u80fd\u7b97\u5b8c\u6700\u5927\u503c\u3001\u5f52\u4e00\u5316\u3001softmax\u3002<\/p>\n
\u800cFlashAttention \u91cc\u60c5\u51b5\u4e0d\u4e00\u6837，\u952e\u548c\u503c\u662f\u5206\u5757\u6d41\u5f0f\u8fdb\u6765\u7684\u5185\u6838\u8fed\u4ee3 K \u548c V \u65f6\u53ea\u80fd\u770b\u5230\u90e8\u5206\u5206\u6570\u5757，\u6c38\u8fdc\u770b\u4e0d\u5230\u5b8c\u6574\u7684\u5206\u6570\u96c6，\u5c31\u6ca1\u6cd5\u4e00\u6b65\u7b97\u5b8c softmax\u3002<\/p>\n
\u89e3\u51b3\u65b9\u6848\u662f\u5728\u7ebf softmax \u516c\u5f0f\u3002\u4e0d\u4e00\u6b65\u7b97\u5b8c，\u800c\u662f\u7ef4\u62a4\u4e09\u4e2a\u9010\u67e5\u8be2\u7684\u72b6\u6001：\u8fd0\u884c\u6700\u5927\u503c m\u1d62（\u4fdd\u8bc1\u6570\u503c\u7a33\u5b9a），\u8fd0\u884c\u5f52\u4e00\u5316\u9879 l\u1d62，\u8fd0\u884c\u8f93\u51fa\u7d2f\u52a0\u5668 O\u1d62\u3002\u6bcf\u6765\u4e00\u4e2a\u65b0\u7684\u6ce8\u610f\u529b\u5206\u6570\u5757，\u5c31\u66f4\u65b0\u8fd9\u4e9b\u503c，\u6700\u540e\u6062\u590d\u7684\u7ed3\u679c\u548c\u5bf9\u6574\u4e2a\u5e8f\u5217\u505a\u5b8c\u6574 softmax \u4e00\u6a21\u4e00\u6837\u3002 ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131722-69888ce2ab7e8.jpg\")
<\/p>\n
\u4ece\u9ad8\u5c42\u770b，\u5b9e\u73b0\u7ed3\u6784\u5982\u4e0b：<\/p>\n
for each (batch, head):
\n for each Q_block:
\n initialize m_i, l_i, O_block
\n for each K\/V block:
\n compute partial scores
\n update online softmax state
\n accumulate output
\n write O_block to memory<\/p>\n
\u6240\u6709\u903b\u8f91\u878d\u5408\u5728\u5185\u6838\u91cc，\u4e2d\u95f4\u72b6\u6001\u5168\u90e8\u9a7b\u7559\u5728\u7247\u4e0a\u5feb\u901f\u5185\u5b58\u3002\u4e0b\u9762\u9010\u6b65\u8bb2\u89e3\u8fd9\u4e2a\u7ed3\u6784\u5982\u4f55\u6620\u5c04\u5230 Triton \u7a0b\u5e8f\u548c GPU \u6267\u884c\u3002<\/p>\n
Python \u5305\u88c5\u5668\u8d1f\u8d23\u51c6\u5907\u8f93\u5165\u5e76\u542f\u52a8 Triton \u5185\u6838，\u505a\u4e09\u4ef6\u4e8b：\u9a8c\u8bc1\u548c\u63d0\u53d6\u8f93\u5165\u5f20\u91cf\u7684\u5f62\u72b6\u4e0e\u6b65\u5e45，\u6784\u5efa\u5185\u6838\u6267\u884c\u7f51\u683c，\u542f\u52a8\u524d\u5411\u6ce8\u610f\u529b\u5185\u6838\u3002\u5305\u88c5\u5668\u672c\u8eab\u4e0d\u542b\u6ce8\u610f\u529b\u903b\u8f91，\u53ea\u5b9a\u4e49\u5de5\u4f5c\u5982\u4f55\u5728 GPU \u4e0a\u8c03\u5ea6\u3002<\/p>\n
# Host wrapper that prepares our inputs and parameters and runs the triton kernel
\nclass TritonFlashAttention(torch.autograd.Function):
\n @staticmethod
\n def flash_attention(Q, K, V, causal):
\n assert Q.is_cuda
\n assert K.is_cuda
\n assert V.is_cuda <\/p>\n
B, H, Lq, D = Q.shape
\n B, H, Lk, D = K.shape
\n B, H, Lk, D = V.shape <\/p>\n
# create the output buffer
\n O = torch.empty_like(Q) <\/p>\n
# we set block_sizes manually for now. We will autotune this later
\n [#BLOCK](#BLOCK)_SIZE_Q = 128
\n [#BLOCK](#BLOCK)_SIZE_KV = 32 <\/p>\n
stage = 3 if causal else 1 <\/p>\n
grid = lambda x: (triton.cdiv(Lq, x["BLOCK_SIZE_Q"]),
\n B * H, 1)
\n M = torch.empty((B, H, Lq), device=Q.device, dtype=torch.float32) <\/p>\n
scaling_factor = 1 \/ math.sqrt(D)
\n fwd_flash_attn_kernel[grid](Q, K, V, O, M, scaling_factor,
\n Q.stride(0), Q.stride(1), Q.stride(2), Q.stride(3),
\n K.stride(0), K.stride(1), K.stride(2), K.stride(3),
\n V.stride(0), V.stride(1), V.stride(2), V.stride(3),
\n O.stride(0), O.stride(1), O.stride(2), O.stride(3),
\n B, NUM_HEADS=H, SEQ_LEN=Lq, HEAD_DIM=D, STAGE=stage,)
\n [#ctx](#ctx).save_for_backward <\/p>\n
return O <\/p>\n
host \u5305\u88c5\u5668\u91cc\u5b9a\u4e49\u4e86\u4e00\u4e2a 2D \u6267\u884c\u7f51\u683c，\u51b3\u5b9a GPU \u5982\u4f55\u5206\u914d\u5de5\u4f5c，\u4e5f\u5c31\u662f\u5e76\u884c\u542f\u52a8\u591a\u5c11\u4e2a Triton \u7a0b\u5e8f\u5b9e\u4f8b\u3002<\/p>\n
grid=lambdax: (triton.cdiv(Lq, x["BLOCK_SIZE_Q"]), B*H, 1) <\/p>\n
\u7b2c\u4e00\u7ef4 program_id(0) \u6807\u8bc6\u7a0b\u5e8f\u5b9e\u4f8b\u5904\u7406\u7684\u67e5\u8be2\u5e8f\u5217\u5757，\u7b2c\u4e8c\u7ef4 program_id(1) \u6807\u8bc6\u5bf9\u5e94\u7684 (batch, head) \u5bf9\u3002<\/p>\n
\u7ef4\u5ea6 0 \u628a\u67e5\u8be2\u5e8f\u5217\u5206\u6210 BLOCK_SIZE_Q \u5927\u5c0f\u7684\u5757，Lq \u662f\u67e5\u8be2\u5e8f\u5217\u957f\u5ea6，\u6bcf\u4e2a\u7a0b\u5e8f\u5b9e\u4f8b\u8d1f\u8d23\u8ba1\u7b97\u8f93\u51fa\u77e9\u9635\u7684\u4e00\u4e2a\u6c34\u5e73"\u6761\u5e26"\u3002\u7ef4\u5ea6 1 \u8de8\u6240\u6709 batch \u548c head \u5e76\u884c，\u6bcf\u4e2a\u7a0b\u5e8f\u5b9e\u4f8b\u5bf9\u5e94\u4e00\u4e2a (batch, head) \u5bf9\u3002\u7ed9\u6bcf\u4e2a\u6ce8\u610f\u529b\u5934\u5206\u914d\u72ec\u7acb\u7a0b\u5e8f\u53ef\u4ee5\u6700\u5927\u5316\u5360\u7528\u7387\u3002\u5185\u6838\u5185\u90e8\u7528 tl.program_id \u914d\u5408\u624b\u52a8\u6b65\u5e45\u7b97\u672f（qb_stride\u3001qh_stride）\u628a\u6bcf\u4e2a worker \u6307\u5411\u5b83\u7684\u5185\u5b58\u5207\u7247\u3002<\/p>\n
\u6bcf\u4e2a\u7a0b\u5e8f\u5b9e\u4f8b\u8d1f\u8d23\u8ba1\u7b97：<\/p>\n
Q[batch, head, q_block : q_block+BLOCK_SIZE_Q]<\/p>\n
\u8fd9\u79cd\u7f51\u683c\u8bbe\u8ba1\u63d0\u4f9b\u4e86\u5e8f\u5217\u7ef4\u5ea6\u5e76\u884c\u3001batch \u548c head \u5e76\u884c，\u800c\u4e14\u7a0b\u5e8f\u95f4\u4e0d\u9700\u8981\u540c\u6b65\u3002\u6bcf\u4e2a\u7a0b\u5e8f\u5728\u7d27\u51d1\u72ec\u7acb\u7684\u5de5\u4f5c\u96c6\u4e0a\u8fd0\u884c，tl.program_id \u7ed3\u5408\u663e\u5f0f\u6b65\u5e45\u7b97\u672f\u628a\u6bcf\u4e2a\u5b9e\u4f8b\u6620\u5c04\u5230\u5bf9\u5e94\u5185\u5b58\u5207\u7247\u3002<\/p>\n
\u524d\u5411\u4f20\u64ad\u5206\u6210\u4e24\u4e2a\u5185\u6838\u3002fwd_flash_attn_kernel \u534f\u8c03\u6267\u884c，\u52a0\u8f7d\u67e5\u8be2\u5757\u3001\u5904\u7406\u56e0\u679c\u903b\u8f91\u3001\u5199\u8f93\u51fa\u3002_attn_fwd_inner \u5b9e\u73b0\u6838\u5fc3 FlashAttention-2 \u8ba1\u7b97，\u6d41\u5f0f\u5904\u7406 K\/V \u5757\u5e76\u6267\u884c\u5728\u7ebf softmax \u66f4\u65b0\u3002\u6bcf\u4e2a Triton \u7a0b\u5e8f\u5b9e\u4f8b\u8ba1\u7b97\u4e00\u4e2a\u67e5\u8be2\u5757 \u00d7 \u4e00\u4e2a\u6ce8\u610f\u529b\u5934 \u00d7 \u4e00\u4e2a batch \u5143\u7d20\u3002<\/p>\n
\u8fd9\u79cd\u5206\u89e3\u628a\u63a7\u5236\u903b\u8f91\u548c\u6d41\u5f0f\u8ba1\u7b97\u5206\u5f00\u5185\u6838\u66f4\u5bb9\u6613\u7406\u89e3\u548c\u4f18\u5316\u3002<\/p>\n
\u8fd9\u4e2a\u5185\u6838\u672c\u8eab\u4e0d\u76f4\u63a5\u5b9e\u73b0\u6ce8\u610f\u529b\u7b97\u6cd5，\u8d1f\u8d23\u7684\u662f\u628a GPU \u7a0b\u5e8f\u5b9e\u4f8b\u6620\u5c04\u5230\u8f93\u5165\u5f20\u91cf\u7684\u5bf9\u5e94\u5757，\u534f\u8c03\u6d41\u5f0f\u6ce8\u610f\u529b\u8ba1\u7b97，\u5904\u7406\u56e0\u679c\u903b\u8f91，\u628a\u6700\u7ec8\u8f93\u51fa\u5199\u56de\u5185\u5b58\u3002<\/p>\n
@triton.jit
\ndef fwd_flash_attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, m_ptr, scale,
\n qb_stride, qh_stride, qn_stride, qd_stride,
\n kb_stride, kh_stride, kn_stride, kd_stride,
\n vb_stride, vh_stride, vn_stride, vd_stride,
\n ob_stride, oh_stride, on_stride, od_stride,
\n BATCH_SIZE, NUM_HEADS:tl.constexpr, SEQ_LEN:tl.constexpr, HEAD_DIM:tl.constexpr,
\n BLOCK_SIZE_Q:tl.constexpr, BLOCK_SIZE_KV:tl.constexpr, STAGE:tl.constexpr): <\/p>\n
# get the id of this program instance
\n block_index_q = tl.program_id(0) # Which chunk of sequence this program is responsible for
\n index_batch_head = tl.program_id(1) # what batch-head to process. zooms out <\/p>\n
# get exact batch
\n index_batch = index_batch_head \/\/ NUM_HEADS <\/p>\n
# get exact head
\n index_head = index_batch_head % NUM_HEADS <\/p>\n
# create offsets to get the index of sequences we are going to process
\n qkv_offset = index_batch * qb_stride + index_head * qh_stride # i.e move from the first to the correct batch then move to the correct head within that batch
\n qkv_offset_K = index_batch * kb_stride + index_head * kh_stride
\n qkv_offset_V = index_batch * vb_stride + index_head * vh_stride
\n qkv_offset_O = index_batch * ob_stride + index_head * oh_stride <\/p>\n
off_q = block_index_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q) # same as off_q (in this head what q block do we need to read )
\n off_kv = tl.arange(0, BLOCK_SIZE_KV)
\n off_head = tl.arange(0, HEAD_DIM) <\/p>\n
# create blocks of pointers to get the address of where the index lives
\n Q_block_ptr = q_ptr + qkv_offset + off_q[:, None] * qn_stride + off_head[None, :] * qd_stride
\n O_block_ptr = o_ptr + qkv_offset_O + off_q[:, None] * on_stride + off_head[None, :] * od_stride <\/p>\n
m_i = tl.zeros((BLOCK_SIZE_Q,), dtype= tl.float32) – float("inf") <\/p>\n
l_i = tl.zeros((BLOCK_SIZE_Q,), dtype=tl.float32) + 1.0
\n O_block = tl.zeros((BLOCK_SIZE_Q, HEAD_DIM), dtype=tl.float32)
\n Q_block = tl.load(Q_block_ptr) # add a mask <\/p>\n
# stage 1: Blocks before the diagonal
\n # stage 2: diagonal block itself
\n # stage 3: for non-causal no masking is needed. For causal mask all the blocks here. <\/p>\n
# runs if causal is True i.e we mask out the future tokens from contributing
\n # this if statement executes for non-causal attention (no masking) or for the blocks to the left of the diagonal in the causal attention
\n # Stage = 3 if causal else 1
\n if STAGE == 1 or STAGE == 3:
\n O_block, l_i, m_i = _attn_fwd_inner(
\n O_block,
\n l_i,
\n m_i,
\n Q_block,
\n block_index_q,
\n scale,
\n BLOCK_SIZE_Q,
\n BLOCK_SIZE_KV,
\n 4 – STAGE,
\n off_kv,
\n off_q,
\n off_head,
\n kn_stride,
\n kd_stride,
\n vd_stride,
\n vn_stride,
\n k_ptr,
\n v_ptr,
\n qkv_offset_K,
\n qkv_offset_V,
\n SEQ_LEN,
\n HEAD_DIM
\n ) <\/p>\n
# this executes for blocks to the right of the diagonal in the causal attention
\n if STAGE == 3:
\n O_block, l_i, m_i = _attn_fwd_inner(
\n O_block,
\n l_i,
\n m_i,
\n Q_block,
\n block_index_q,
\n scale,
\n BLOCK_SIZE_Q,
\n BLOCK_SIZE_KV,
\n 2,
\n off_kv,
\n off_q,
\n off_head,
\n kn_stride,
\n kd_stride,
\n vd_stride,
\n vn_stride,
\n k_ptr,
\n v_ptr,
\n qkv_offset_K,
\n qkv_offset_V,
\n SEQ_LEN,
\n HEAD_DIM
\n ) <\/p>\n
m_i += tl.math.log(l_i)
\n O_block = O_block \/ l_i[:, None]
\n m_ptrs = m_ptr + index_batch_head * SEQ_LEN + off_q
\n tl.store(m_ptrs, m_i)
\n tl.store(O_block_ptr, O_block.to(tl.float16))<\/p>\n
\u56de\u987e Python \u5305\u88c5\u5668\u91cc\u7684\u7f51\u683c：<\/p>\n
grid = (
\n ceil_div(Lq, BLOCK_SIZE_Q),
\n B * H
\n )<\/p>\n
\u8fd9\u4e2a 2D \u7f51\u683c\u6620\u5c04\u63d0\u4f9b\u5e8f\u5217\u7ef4\u5ea6\u5e76\u884c\u548c batch\/head \u5e76\u884c\u3002<\/p>\n
\u5185\u6838\u5185\u90e8：<\/p>\n
block_index_q =tl.program_id(0)
\n index_batch_head =tl.program_id(1)<\/p>\n
\u89e3\u7801\u7b2c\u4e8c\u7ef4：<\/p>\n
index_batch=index_batch_head\/\/NUM_HEADS
\n index_head =index_batch_head%NUM_HEADS<\/p>\n
\u8fd9\u51e0\u4e2a\u53d8\u91cf\u552f\u4e00\u6807\u8bc6\u5f53\u524d\u7a0b\u5e8f\u5b9e\u4f8b\u8d1f\u8d23\u54ea\u4e2a batch \u5143\u7d20\u3001\u54ea\u4e2a\u6ce8\u610f\u529b\u5934\u3001\u54ea\u4e2a\u67e5\u8be2\u5757\u3002<\/p>\n
PyTorch \u6216 numpy \u91cc\u7528\u591a\u7ef4\u8bed\u6cd5\u7d22\u5f15\u5f20\u91cf，\u6bd4\u5982 Q[batch, head, seq_pos, dim]\u3002\u800cTriton \u5185\u6838\u91cc\u6ca1\u6709\u591a\u7ef4\u5f20\u91cf，\u53ea\u6709\u6307\u5411\u8f93\u5165\u7b2c\u4e00\u4e2a\u5143\u7d20\u7684\u88f8\u6307\u9488 q_ptr\u5fc5\u987b\u7528\u6307\u9488\u7b97\u672f\u624b\u52a8\u91cd\u6784\u7d22\u5f15\u3002<\/p>\n
\u67e5\u8be2\u5f20\u91cf Q \u5f62\u72b6\u662f [BATCH, HEADS, SEQ_LEN, HEAD_DIM]，\u786c\u4ef6\u5c42\u9762\u662f\u6241\u5e73\u4e00\u7ef4\u6570\u7ec4\u5b58\u50a8\u3002\u6cbf\u6bcf\u4e2a\u7ef4\u5ea6\u79fb\u52a8\u7528\u6b65\u5e45：qb_stride \u8df3\u4e00\u4e2a batch，qh_stride \u8df3\u4e00\u4e2a head，qn_stride \u8df3\u4e00\u4e2a token，qd_stride \u8df3\u4e00\u4e2a\u7279\u5f81\u3002<\/p>\n
\u6bcf\u4e2a\u7a0b\u5e8f\u5b9e\u4f8b\u5148\u9009\u5b9a\u81ea\u5df1\u8d1f\u8d23\u7684 batch \u548c head \u5207\u7247：<\/p>\n
qkv_offset=index_batch*qb_stride+index_head*qh_stride<\/p>\n
\u8fd9\u4e2a\u504f\u79fb\u4e4b\u540e，\u6307\u9488\u6307\u5411 Q[batch, head, 0, :]\u3002K\u3001V\u3001O \u540c\u7406，\u7528\u5404\u81ea\u7684\u6b65\u5e45\u3002\u7136\u540e\u6784\u5efa\u5f53\u524d\u5757\u7684\u7d22\u5f15\u8303\u56f4：<\/p>\n
off_q =block_index_q*BLOCK_SIZE_Q+tl.arange(0, BLOCK_SIZE_Q)
\n off_head=tl.arange(0, HEAD_DIM)<\/p>\n
\u7528\u8fd9\u4e9b\u504f\u79fb\u52a0\u5e7f\u64ad，\u6784\u5efa\u6307\u5411\u67e5\u8be2\u5757\u7684\u6307\u9488：<\/p>\n
Q_block_ptr=q_ptr+qkv_offset \\\\
\n +off_q[:, None] *qn_stride \\\\
\n +off_head[None, :] *qd_stride<\/p>\n
\u8f93\u51fa O_block_ptr \u4e5f\u7c7b\u4f3c：<\/p>\n
O_block_ptr=o_ptr+qkv_offset_O \\\\
\n +off_q[:, None] *on_stride \\\\
\n +off_head[None, :] *od_stride<\/p>\n
\u5b8c\u5168\u7528\u6307\u9488\u7b97\u672f\u91cd\u73b0\u4e86 4D \u7d22\u5f15 Q[batch, head, q_positions, head_dim]\u3002<\/p>\n
\u8fd9\u79cd\u663e\u5f0f\u6307\u9488\u6784\u5efa\u5f88\u5173\u952e，\u786e\u4fdd\u53ea\u52a0\u8f7d\u6bcf\u4e2a\u7a0b\u5e8f\u5b9e\u4f8b\u9700\u8981\u7684 Q \u5757\u5e76\u9001\u5230 SRAM，\u907f\u514d\u78b0\u4e0d\u76f8\u5173\u7684\u5185\u5b58，\u5b9e\u73b0\u5408\u5e76\u8bbf\u95ee，\u6700\u5927\u5316\u7f13\u5b58\u590d\u7528\u3002<\/p>\n
\u52a0\u8f7d\u67e5\u8be2\u5757\u540e，\u5185\u6838\u521d\u59cb\u5316\u5728\u7ebf softmax \u6240\u9700\u7684\u6bcf\u5757\u72b6\u6001\u5e76\u5206\u6d3e\u6d41\u5f0f\u8ba1\u7b97\u3002\u6d41\u5f0f\u903b\u8f91\u548c\u56e0\u679c\u9636\u6bb5\u7684\u7ec6\u8282\u5728 _attn_fwd_inner \u91cc，\u540e\u9762\u5206\u6790\u3002\u5148\u7406\u89e3\u8fd9\u4e2a\u6bcf\u5757\u72b6\u6001\u4e3a\u4ec0\u4e48\u5b58\u5728\u3001\u4ee3\u8868\u4ec0\u4e48\u3002<\/p>\n
\u4e3a\u4e86\u5728\u8fed\u4ee3 K \u548c V \u5757\u65f6\u6b63\u786e\u589e\u91cf\u8ba1\u7b97 softmax，\u9700\u8981\u8ffd\u8e2a\u4e09\u4e2a\u91cf：\u8fd0\u884c\u6700\u5927\u503c m_i\u3001\u8fd0\u884c softmax \u5206\u6bcd l_i\u3001\u672a\u5f52\u4e00\u5316\u52a0\u6743\u548c O_block\u3002<\/p>\n
\u8fd9\u4e09\u4e2a\u53d8\u91cf\u6784\u6210\u5728\u7ebf softmax \u7b97\u6cd5\u7684\u72b6\u6001\u3002FlashAttention \u5206\u5757\u5904\u7406\u952e\u503c，\u5185\u6838\u6c38\u8fdc\u65e0\u6cd5\u4e00\u6b21\u8bbf\u95ee\u6240\u6709\u6ce8\u610f\u529b\u5206\u6570\u3002\u8981\u5f97\u5230\u548c\u5b8c\u6574 softmax \u4e00\u6837\u7684\u7ed3\u679c，\u5fc5\u987b\u7ef4\u62a4\u6570\u503c\u7a33\u5b9a\u7528\u7684\u8fd0\u884c\u6700\u5927\u503c m_i\u3001\u8fd0\u884c\u5f52\u4e00\u5316\u56e0\u5b50 l_i\u3001\u7d2f\u79ef\u52a0\u6743\u8f93\u51fa O_block\u3002\u8fd9\u4e9b\u72b6\u6001\u5171\u540c\u4f5c\u7528，\u7cbe\u786e\u91cd\u5efa softmax(QK\u1d40) @ V，\u4e0d\u9700\u8981\u7269\u5316\u6ce8\u610f\u529b\u77e9\u9635\u3002<\/p>\n
Softmax \u6d89\u53ca\u6307\u6570\u8fd0\u7b97，FP16\/BF16 \u4e0b\u5bb9\u6613\u6570\u503c\u4e0d\u7a33\u5b9a\u3002\u4e3a\u4e86\u628a\u6307\u6570\u4fdd\u6301\u5728\u5408\u7406\u8303\u56f4，\u6bcf\u4e2a\u67e5\u8be2\u884c\u8ffd\u8e2a\u4e00\u4e2a\u8fd0\u884c\u6700\u5927\u503c m_i\u3002\u5904\u7406\u65b0\u7684 K \u548c V \u5757\u65f6，\u8fd9\u4e2a\u8fd0\u884c\u6700\u5927\u503c\u53ef\u80fd\u589e\u5927\u3002\u4e00\u65e6\u589e\u5927，\u4e4b\u524d\u7528\u65e7\u6700\u5927\u503c\u8ba1\u7b97\u7684\u7d2f\u79ef\u8d21\u732e\u5c31\u4e0d\u5728\u540c\u4e00\u5c3a\u5ea6\u4e0a\u4e86\u3002<\/p>\n
\u7ea0\u6b63\u529e\u6cd5\u662f\u7528\u4e00\u4e2a\u56e0\u5b50\u91cd\u65b0\u7f29\u653e\u7d2f\u79ef\u7684\u5206\u6bcd： ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131722-69888ce2edc80.jpg\")
<\/p>\n
the numerator ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce30d154.jpg\")
<\/p>\n
the scaling factor ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce31d7b2.jpg\")
<\/p>\n
the normalizing denominator<\/p>\n
\u8fd9\u79cd\u91cd\u65b0\u7f29\u653e\u786e\u4fdd\u5206\u6bcd\u91cc\u6240\u6709\u9879\u90fd\u76f8\u5bf9\u540c\u4e00\u4e2a\u6700\u5927\u503c\u3002\u6d41\u5f0f\u5904\u7406\u952e\u503c\u5757\u65f6\u53cd\u590d\u5e94\u7528\u8fd9\u4e2a\u66f4\u65b0\u5c31\u80fd\u6062\u590d\u7cbe\u786e\u7684 softmax \u5f52\u4e00\u5316\u56e0\u5b50，\u4e0d\u9700\u8981\u7269\u5316\u5b8c\u6574\u7684\u6ce8\u610f\u529b\u5206\u6570\u96c6\u3002<\/p>\n
\u5185\u6838\u91cc\u662f\u8fd9\u6837\u5199：<\/p>\n
alpha=exp(m_old-m_new)
\n l_i=l_i*alpha+l_ij<\/p>\n
\u6ce8\u610f\u529b\u8f93\u51fa\u5b9a\u4e49\u4e3a： ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce32cca5.jpg\")
<\/p>\n
Final attention output<\/p>\n
\u6807\u51c6\u5b9e\u73b0\u91cc\u53ef\u4ee5\u76f4\u63a5\u7b97，\u56e0\u4e3a\u5b8c\u6574\u7684 softmax \u5f52\u4e00\u5316\u7cfb\u6570\u4e8b\u5148\u5c31\u77e5\u9053\u3002FlashAttention \u91cc\u952e\u503c\u5206\u5757\u6d41\u5f0f\u8fdb\u6765，\u6700\u7ec8\u5f52\u4e00\u5316\u56e0\u5b50\u8981\u7b49\u6240\u6709 K \u548c V \u5757\u5904\u7406\u5b8c\u624d\u80fd\u786e\u5b9a\u3002<\/p>\n
\u6240\u4ee5\u53ea\u80fd\u7d2f\u79ef\u4e00\u4e2a\u672a\u5f52\u4e00\u5316\u7684\u52a0\u6743\u548c，\u6700\u540e\u518d\u5f52\u4e00\u5316\u3002<\/p>\n
\u6bcf\u6b21\u8fed\u4ee3，\u8ba1\u7b97\u76f8\u5bf9\u4e8e\u5f53\u524d\u8fd0\u884c\u6700\u5927\u503c\u7684\u5757\u7ea7 softmax \u6982\u7387： ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce33c8aa.jpg\")
<\/p>\n
\u7ef4\u62a4\u4e00\u4e2a\u672a\u5f52\u4e00\u5316\u8f93\u51fa\u7d2f\u52a0\u5668： ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce35320c.jpg\")
<\/p>\n
unnormalized softmax output<\/p>\n
\u5904\u7406\u65b0 K\/V \u5757\u65f6\u8fd0\u884c\u6700\u5927\u503c\u53ef\u80fd\u53d8，\u4e4b\u524d\u7d2f\u79ef\u7684\u8f93\u51fa\u5fc5\u987b\u91cd\u65b0\u7f29\u653e\u4ee5\u5339\u914d\u65b0\u6700\u5927\u503c\u3002 ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce364321.jpg\")
<\/p>\n
\u9010\u5757\u66f4\u65b0\u8f93\u51fa\u7d2f\u52a0\u5668：<\/p>\n
O_block=O_block*alpha[:, None]
\n O_block=P_block@V_block+O_block<\/p>\n
\u6240\u6709 K\/V \u5757\u5904\u7406\u5b8c\u540e，\u628a\u7d2f\u79ef\u7684\u672a\u5f52\u4e00\u5316\u8f93\u51fa\u9664\u4ee5\u7d2f\u79ef\u7684 softmax \u5206\u6bcd li \u5f97\u5230\u6700\u7ec8\u6ce8\u610f\u529b\u8f93\u51fa： ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce373e01.jpg\")
<\/p>\n
final normalization<\/p>\n
\u7ed3\u679c\u548c\u6807\u51c6 softmax \u6ce8\u610f\u529b\u5b8c\u5168\u4e00\u6837，\u4f46\u6c38\u8fdc\u4e0d\u4f1a\u5728\u5185\u5b58\u91cc\u7269\u5316\u5b8c\u6574\u6ce8\u610f\u529b\u77e9\u9635\u6216 softmax \u6982\u7387\u3002<\/p>\n
\u6bcf\u4e2a\u7a0b\u5e8f\u5b9e\u4f8b\u4e3a\u6bcf\u4e2a\u67e5\u8be2\u5757\u521d\u59cb\u5316\u8fd9\u4e09\u4e2a\u72b6\u6001\u4e00\u6b21：<\/p>\n
m_i=tl.zeros((BLOCK_SIZE_Q,), dtype=tl.float32) -inf
\n l_i=tl.zeros((BLOCK_SIZE_Q,), dtype=tl.float32) +1
\n O_block=tl.zeros((BLOCK_SIZE_Q, HEAD_DIM), dtype=tl.float32)<\/p>\n
_attn_fwd_inner \u5b9e\u73b0 FlashAttention-2 \u7b97\u6cd5\u6838\u5fc3，\u7531 fwd_flash_attn_kernel \u8c03\u7528，\u4e00\u6b21\u5904\u7406\u4e00\u4e2a\u67e5\u8be2\u5757\u3002<\/p>\n
@triton.jit
\ndef _attn_fwd_inner(O_block, l_i,m_i, Q_block, block_index_q,
\n scale: tl.constexpr,
\n BLOCK_SIZE_Q: tl.constexpr,
\n BLOCK_SIZE_KV: tl.constexpr,
\n STAGE: tl.constexpr,
\n off_kv: tl.constexpr,
\n off_q: tl.constexpr,
\n off_head: tl.constexpr,
\n kn_stride: tl.constexpr,
\n kd_stride: tl.constexpr,
\n vd_stride: tl.constexpr,
\n vn_stride: tl.constexpr,
\n k_ptr,
\n v_ptr,
\n qkv_offset_K: tl.constexpr,
\n qkv_offset_V: tl.constexpr,
\n SEQ_LEN:tl.constexpr,
\n HEAD_DIM: tl.constexpr):<\/p>\n
\u5176\u4e2d Q_block \u5f62\u72b6 [BLOCK_SIZE_Q, HEAD_DIM]，O_block \u662f\u7d2f\u79ef\u8f93\u51fa，m_i \u662f\u6bcf\u67e5\u8be2\u884c\u7684\u8fd0\u884c\u6700\u5927\u503c，l_i \u662f\u8fd0\u884c softmax \u5f52\u4e00\u5316\u3002<\/p>\n
FA \u5185\u6838\u652f\u6301\u56e0\u679c（\u53ea\u770b\u8fc7\u53bb\u548c\u5f53\u524d token）\u548c\u975e\u56e0\u679c\u6ce8\u610f\u529b（\u53cc\u5411，\u53ef\u4ee5\u770b\u672a\u6765）\u3002\u7528\u4e00\u4e2a\u9636\u6bb5\u673a\u5236\u5b9e\u73b0：<\/p>\n
if STAGE == 1:
\n lo, hi = 0, block_index_q * BLOCK_SIZE_Q
\n elif STAGE == 2:
\n lo, hi = block_index_q * BLOCK_SIZE_Q, (block_index_q + 1) * BLOCK_SIZE_Q
\n else:
\n lo, hi = 0, SEQ_LEN<\/p>\n
\u8fd9\u4e2a\u903b\u8f91\u51b3\u5b9a\u5f53\u524d\u5185\u6838\u5904\u7406\u54ea\u4e9b K\/V \u5757\u3002Stage 1 \u662f\u5bf9\u89d2\u7ebf\u5de6\u4fa7\u7684\u5757，K \u548c V \u8303\u56f4\u4ec5\u9650\u4e8e\u6b64\u3002Stage 2 \u662f\u5bf9\u89d2\u7ebf\u5757\u672c\u8eab\u3002Stage 3 \u662f\u975e\u56e0\u679c\u903b\u8f91，K \u548c V \u5173\u6ce8\u6240\u6709 Q\u3002\u8fd9\u6837\u907f\u514d\u8ba1\u7b97\u56e0\u679c\u6ce8\u610f\u529b\u4e2d\u80af\u5b9a\u4f1a\u88ab mask \u6389\u7684\u5206\u6570，\u51cf\u5c11\u4e0d\u5fc5\u8981\u7684 masking \u5de5\u4f5c\u3002<\/p>\n
\u67e5\u8be2\u867d\u7136\u5206\u533a\u5230\u5404\u7a0b\u5e8f\u5b9e\u4f8b，\u4f46\u6bcf\u4e2a\u67e5\u8be2\u5757\u5fc5\u987b\u5173\u6ce8\u6240\u6709\u952e\u503c\u2014\u2014\u8fd9\u662f\u5168\u6ce8\u610f\u529b\u7684\u5b9a\u4e49\u51b3\u5b9a\u7684\u3002\u5b8c\u6574 K \u548c V \u77e9\u9635\u4ece\u4e0d\u4e00\u6b21\u6027\u52a0\u8f7d\u5230 SRAM，\u800c\u662f\u4ee5 BLOCK_SIZE_KV \u5927\u5c0f\u7684\u5757\u6d41\u5f0f\u5904\u7406：<\/p>\n
forstart_kvinrange(lo, hi, BLOCK_SIZE_KV):<\/p>\n
\u52a0\u8f7d BLOCK_SIZE_KV \u4e2a\u952e\u503c，\u8ba1\u7b97\u90e8\u5206\u6ce8\u610f\u529b\u5206\u6570，\u66f4\u65b0\u5728\u7ebf softmax \u72b6\u6001，\u4e22\u5f03\u8be5\u5757，\u5904\u7406\u4e0b\u4e00\u4e2a\u3002\u5185\u5b58\u590d\u6742\u5ea6\u7ef4\u6301 O(N)\u3002<\/p>\n
\u6bcf\u4e2a\u7a0b\u5e8f\u5b9e\u4f8b\u53ea\u52a0\u8f7d\u4e00\u4e2a\u67e5\u8be2\u5757，\u5bf9\u5e94\u5e8f\u5217\u4e2d\u4e00\u5c0f\u90e8\u5206 token\u3002\u4f46\u8fd9\u4e9b token \u8981\u6b63\u786e\u8ba1\u7b97\u6ce8\u610f\u529b\u8f93\u51fa，\u5fc5\u987b\u5173\u6ce8\u5e8f\u5217\u91cc\u6240\u6709\u952e\u503c\u3002\u8fd9\u662f\u81ea\u6ce8\u610f\u529b\u5b9a\u4e49\u51b3\u5b9a\u7684：\u6bcf\u4e2a\u67e5\u8be2\u90fd\u8981\u548c\u6bcf\u4e2a\u952e\u6bd4\u8f83\u3002FlashAttention \u6ca1\u6539\u8fd9\u4e2a\u7b97\u6cd5\u8981\u6c42，\u53ea\u6539\u8ba1\u7b97\u8c03\u5ea6\u65b9\u5f0f\u3002\u952e\u503c\u9010\u5757\u6d41\u5f0f\u8fdb\u6765，\u7d2f\u79ef\u5230\u8f93\u51fa，\u7acb\u523b\u4e22\u5f03，\u5185\u5b58\u5360\u7528\u5c0f，\u7ed3\u679c\u7cbe\u786e\u3002\u4e00\u4e9b\u65b0\u7684\u6ce8\u610f\u529b\u53d8\u4f53（\u5c40\u90e8\u6ce8\u610f\u529b\u3001\u7a00\u758f\u6ce8\u610f\u529b\u3001\u6ed1\u52a8\u7a97\u53e3\u6ce8\u610f\u529b）\u4e0d\u4f1a\u5173\u6ce8\u6240\u6709 token\u3002<\/p>\n
\u548c Q_block \u4e00\u6837，\u8ba1\u7b97\u5f53\u524d\u5757\u7684 token \u7d22\u5f15：<\/p>\n
kv_positions=start_kv+off_kv<\/p>\n
\u7136\u540e\u6784\u5efa\u6307\u9488：<\/p>\n
K_block_ptr = (
\n k_ptr + qkv_offset_K
\n + off_head[:, None] * kd_stride
\n + kv_positions[None, :] * kn_stride
\n) <\/p>\n
V_block_ptr = (
\n v_ptr + qkv_offset_V
\n + kv_positions[:, None] * vn_stride
\n + off_head[None, :] * vd_stride
\n )<\/p>\n
\u5f97\u5230\u5f62\u72b6 [HEAD_DIM, BLOCK_SIZE_KV] \u7684 K \u548c V \u6307\u9488\u3002\u8fb9\u754c mask \u903b\u8f91\u9632\u6b62\u6700\u540e\u4e00\u4e2a\u5757\u8d8a\u754c\u8bbf\u95ee：<\/p>\n
mask_k = kv_positions[None, :] < SEQ_LEN
\n mask_v = kv_positions[:, None] < SEQ_LEN<\/p>\n
\u4ece HBM \u52a0\u8f7d K \u548c V \u5230\u7247\u4e0a SRAM：<\/p>\n
K_block = tl.load(K_block_ptr, mask=mask_k, other=0.0)
\n V_block = tl.load(V_block_ptr, mask=mask_v, other=0.0)<\/p>\n
\u8ba1\u7b97\u5206\u5757\u70b9\u79ef：<\/p>\n
QK_block=tl.dot(Q_block, K_block)<\/p>\n
\u5e94\u7528\u7f29\u653e\u548c mask（\u5982\u679c\u662f\u56e0\u679c\u7684），\u66f4\u65b0\u8fd0\u884c\u6700\u5927\u503c：<\/p>\n
mask = off_q[:, None] >= (start_kv + off_kv[None, :])
\n QK_block = QK_block * scale + tl.where(mask, 0, -1e6)
\n m_ij = tl.maximum(m_i, tl.max(QK_block, 1))
\n QK_block -= m_ij[:, None]
\n m_ij = tl.maximum(m_i, tl.max(QK_block, 1) * scale)
\n QK_block = QK_block * scale – m_ij[:, None]<\/p>\n
\u66f4\u65b0\u5728\u7ebf softmax \u72b6\u6001：<\/p>\n
P_block = exp(QK_block)
\n l_ij = sum(P_block, axis=1)
\n alpha = exp(m_i – m_ij)
\n l_i = l_i * alpha + l_ij<\/p>\n
\u66f4\u65b0\u8f93\u51fa\u7d2f\u52a0\u5668：<\/p>\n
O_block = O_block * alpha[:, None]
\n O_block = dot(P_block, V_block, O_block)<\/p>\n
\u7528\u5f53\u524d\u8fed\u4ee3\u627e\u5230\u7684\u65b0\u6700\u5927\u503c\u66f4\u65b0\u8fd0\u884c\u6700\u5927\u503c：<\/p>\n
m_i=m_ij<\/p>\n
\u66f4\u65b0\u540e\u7684\u72b6\u6001\u8fd4\u56de\u7ed9\u5916\u5c42\u5185\u6838 fwd_flash_attn_kernel\u3002<\/p>\n
\u6240\u6709 K\/V \u5757\u5904\u7406\u5b8c\u540e，\u524d\u5411\u5185\u6838\u5b8c\u6210\u8f93\u51fa：<\/p>\n
O_block=O_block\/l_i[:, None]<\/p>\n
\u7528\u7d2f\u79ef\u7684\u5206\u6bcd\u56e0\u5b50\u5f52\u4e00\u5316\u6ce8\u610f\u529b\u8f93\u51fa\u3002\u5f53\u524d\u67e5\u8be2\u5757\u7684\u6ce8\u610f\u529b\u8f93\u51fa\u5c31\u7b97\u5b8c\u4e86\u3002<\/p>\n
\u524d\u5411\u4f20\u64ad\u5b9e\u73b0\u5b8c\u6bd5\u5e76\u9a8c\u8bc1\u540e，\u53ef\u4ee5\u770b\u770b\u6027\u80fd\u548c\u6807\u51c6\u6ce8\u610f\u529b\u5b9e\u73b0\u6bd4\u8f83\u4e00\u4e0b\u3002 ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce3804ac.jpg\")
<\/p>\n
FlashAttention vs. standard attention vs torch2.2 (spda flashattn) TFLOP\/s benchmarks ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2026\/02\/20260208131723-69888ce393c03.jpg\")
<\/p>\n
\u6240\u6709\u5e8f\u5217\u957f\u5ea6\u4e0a\u6807\u51c6\u6ce8\u610f\u529b\u5728 3-4 TFLOPs\/sec \u5de6\u53f3\u5c31\u5230\u9876\u4e86\u3002\u7406\u8bba\u8ba1\u7b97\u91cf\u867d\u7136\u6309 O(N\u00b2) \u589e\u957f，\u4f46\u6807\u51c6\u6ce8\u610f\u529b\u88ab HBM \u6d41\u91cf\u4e3b\u5bfc\u3002GPU \u5927\u90e8\u5206\u65f6\u95f4\u5728\u642c\u8fd0 N \u00d7 N \u6ce8\u610f\u529b\u77e9\u9635，\u4e0d\u662f\u5728\u505a\u6709\u7528\u8ba1\u7b97\u3002\u5e8f\u5217\u53d8\u957f\u5e76\u4e0d\u80fd\u63d0\u9ad8\u8ba1\u7b97\u5355\u5143\u5229\u7528\u7387，\u53ea\u662f\u5185\u5b58\u538b\u529b\u53d8\u5927\u3002<\/p>\n
Triton FlashAttention \u5185\u6838\u5219\u968f\u5e8f\u5217\u957f\u5ea6\u589e\u52a0\u6fc0\u8fdb\u6269\u5c55\u3002512 token \u65f6\u6027\u80fd\u4e00\u822c，\u8d85\u8fc7 2K token \u540e\u541e\u5410\u91cf\u5feb\u901f\u4e0a\u5347\u300216K token \u65f6\u7ef4\u6301\u5728\u7ea6 190 TFLOPs\/sec\u3002\u8fd9\u6b63\u662f FlashAttention \u8bbe\u8ba1\u8981\u8fbe\u5230\u7684\u6548\u679c：\u963b\u6b62\u6ce8\u610f\u529b\u77e9\u9635\u7269\u5316，\u4e2d\u95f4\u6570\u636e\u9a7b\u7559 SRAM，\u5185\u5b58\u52a0\u8f7d\u5f97\u4ee5\u644a\u9500\u3002\u5e8f\u5217\u8d8a\u957f，\u5185\u6838\u8d8a\u8d8b\u5411\u8ba1\u7b97\u53d7\u9650，GPU \u63a5\u8fd1\u6709\u6548\u5cf0\u503c\u541e\u5410\u91cf\u2014\u2014\u548c\u6807\u51c6\u6ce8\u610f\u529b\u6070\u597d\u76f8\u53cd，\u6807\u51c6\u6ce8\u610f\u529b\u5e8f\u5217\u8d8a\u957f\u8d8a\u5185\u5b58\u53d7\u9650\u3002<\/p>\n
\u7b2c\u4e8c\u5f20\u56fe\u5728 Nvidia A100 \u4e0a\u901a\u8fc7 sdpa API \u6bd4\u8f83\u4e86 Triton FlashAttention \u548c PyTorch \u5b98\u65b9 FlashAttention \u5b9e\u73b0\u3002\u5e8f\u5217\u8f83\u77ed\u65f6 PyTorch \u5b9e\u73b0\u6709\u7ade\u4e89\u529b，\u5e8f\u5217\u957f\u5ea6 \u22654k \u540e，\u81ea\u5b9a\u4e49 Triton \u5185\u6838\u8ffd\u5e73\u5e76\u7565\u5fae\u8d85\u8fc7 PyTorch \u6027\u80fd\u300216k token \u65f6，\u4e24\u8005\u90fd\u6536\u655b\u5230\u7ea6 180-190 TFLOPs\/sec\u3002<\/p>\n
\u6240\u6709\u7ed3\u679c\u5728\u540c\u4e00 GPU（Nvidia A100 SXM）\u76f8\u540c\u6761\u4ef6\u4e0b\u83b7\u5f97\u3002\u541e\u5410\u91cf\u4ee5 TFLOPs\/sec \u62a5\u544a，\u7531\u7f29\u653e\u70b9\u79ef\u6ce8\u610f\u529b\u7684\u7406\u8bba FLOP \u6570\u9664\u4ee5\u5b9e\u6d4b\u5185\u6838\u8fd0\u884c\u65f6\u95f4\u5f97\u51fa\u3002\u5e8f\u5217\u957f\u5ea6\u53d8\u5316，batch \u5927\u5c0f\u3001\u5934\u6570\u3001\u5934\u7ef4\u5ea6\u56fa\u5b9a\u3002<\/p>\n
\u8fd9\u4e9b\u57fa\u51c6\u9a8c\u8bc1\u4e86\u4e09\u4ef6\u4e8b：\u6807\u51c6\u6ce8\u610f\u529b\u4ece\u6839\u672c\u4e0a\u5185\u5b58\u53d7\u9650；FlashAttention \u628a\u74f6\u9888\u4ece\u5185\u5b58\u8f6c\u5230\u8ba1\u7b97；Triton \u63d0\u4f9b\u4e86\u8db3\u591f\u7684\u6570\u636e\u79fb\u52a8\u548c GPU \u5185\u5b58\u5e95\u5c42\u63a7\u5236，\u80fd\u8fbe\u5230\u63a5\u8fd1\u6700\u4f18\u6027\u80fd\u3002<\/p>\n
\u5173\u952e\u662f\u6027\u80fd\u589e\u76ca\u968f\u5e8f\u5217\u957f\u5ea6\u589e\u957f\u3002\u8fd9\u6b63\u662f FlashAttention \u5728\u5b9e\u8df5\u4e2d\u6700\u91cd\u8981\u7684\u5730\u65b9\u3002<\/p>\n
\u73b0\u4ee3 GPU \u4e0a\u6027\u80fd\u7531\u5185\u5b58\u884c\u4e3a\u4e3b\u5bfc，\u4e0d\u662f FLOPs；\u5185\u6838\u878d\u5408\u548c SRAM \u9a7b\u7559\u6bd4\u6570\u5b66\u6280\u5de7\u66f4\u91cd\u8981；\u5728\u7ebf softmax \u662f IO \u611f\u77e5\u6ce8\u610f\u529b\u7684\u5173\u952e；Triton \u66b4\u9732\u4e86\u8db3\u591f\u7684\u786c\u4ef6\u7ec6\u8282\u6765\u5199\u53ef\u8bfb\u53c8\u5feb\u7684\u5185\u6838；\u4ed4\u7ec6\u5206\u5757\u52a0\u81ea\u52a8\u8c03\u4f18，\u81ea\u5b9a\u4e49\u5185\u6838\u80fd\u548c\u5382\u5546\u5b9e\u73b0\u6253\u5e73\u3002<\/p>\n
FlashAttention \u4e0d\u662f\u56e0\u4e3a\u6539\u4e86\u7b97\u6cd5\u624d\u66f4\u5feb，\u662f\u56e0\u4e3a\u5b83\u5c0a\u91cd GPU \u5b9e\u9645\u7684\u5de5\u4f5c\u65b9\u5f0f\u3002<\/p>\n
\u672c\u6587\u53ea\u5b9e\u73b0\u4e86\u524d\u5411\u4f20\u64ad\u3002\u6269\u5c55\u5230\u5b8c\u6574\u7684\u8bad\u7ec3\u7ea7 FlashAttention（\u53cd\u5411\u4f20\u64ad\u3001dropout\u3001\u5404\u79cd mask \u53d8\u4f53）\u7559\u5f85\u540e\u7eed\u5de5\u4f5c\u3002<\/p>\n
\u672c\u6587\u6e90\u4ee3\u7801：<\/p>\n
https:\/\/avoid.overfit.cn\/post\/0ae6fbc34b7f4c1788f6399a7a1fc431<\/p>\n
by Katherine Oluwadarasimi Olowookere<\/p>\n","protected":false},"excerpt":{"rendered":"
\u672c\u6587\u5b9e\u73b0 FlashAttention-2 \u7684\u524d\u5411\u4f20\u64ad，\u5177\u4f53\u5305\u62ec：\u4e3a Q\u3001K\u3001V \u8bbe\u8ba1\u5206\u5757\u7b56\u7565；\u6d41\u5f0f\u5904\u7406 K \u548c V \u5757\u800c\u975e\u7269\u5316\u5b8c\u6574\u6ce8\u610f\u529b\u77e9\u9635；\u5b9e\u73b0\u5728\u7ebf softmax \u7b97\u6cd5\u4fdd\u8bc1\u6570\u503c\u7a33\u5b9a\u6027；\u652f\u6301\u56e0\u679c\u548c\u975e\u56e0\u679c\u4e24\u79cd\u6ce8\u610f\u529b\u6a21\u5f0f；\u7528 Triton autotuner \u81ea\u52a8\u8c03\u4f18\u5185\u6838\u914d\u7f6e；\u6700\u540e\u7528 PyTorch \u9a8c\u8bc1\u6b63\u786e\u6027\u3002
\nFlashAttention vs. standard att<\/p>\n","protected":false},"author":2,"featured_media":73957,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[1],"tags":[7820,2112,50,86],"topic":[],"class_list":["post-73970","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-server","tag-flashattention","tag-triton","tag-50","tag-86"],"yoast_head":"\n