\u8fd9\u662f\u6839\u636e Tri Dao \u7684 Flash Attention v2 \u7b97\u6cd5<\/a>\u7684 Triton \u5b9e\u73b0\u3002\u81f4\u8c22\uff1aOpenAI \u6838\u5fc3\u56e2\u961f<\/p>\n \u7279\u522b\u9e23\u8c22\uff1a<\/p>\n fused-attention-batch4-head32-d64-fwd-causal=True: Download Jupyter notebook: 06-fused-attention.ipynb<\/a><\/p>\n Download Python source code: 06-fused-attention.py<\/a><\/p>\n\n
Out:<\/p>\n
fused-attention-batch4-head32-d64-fwd-causal=False:
fused-attention-batch4-head32-d64-bwd-causal=True:<\/p>\n
import pytest\nimport torch\n\n\nimport triton\nimport triton.language as tl\n\n\n\n\ndef is_hip():\n return triton.runtime.driver.active.get_current_target().backend == "hip"\n\n\n\n\n@triton.jit\ndef _attn_fwd_inner(acc, l_i, m_i, q, #\n K_block_ptr, V_block_ptr, #\n start_m, qk_scale, #\n BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, #\n STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr, #\n N_CTX: tl.constexpr, fp8_v: tl.constexpr):\n # range of values handled by this stage\n # \u8fd9\u4e2a\u9636\u6bb5\u5904\u7406\u7684\u6570\u503c\u8303\u56f4\n if STAGE == 1:\n lo, hi = 0, start_m * BLOCK_M\n elif STAGE == 2:\n lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M\n lo = tl.multiple_of(lo, BLOCK_M)\n # causal = False\n else:\n lo, hi = 0, N_CTX\n K_block_ptr = tl.advance(K_block_ptr, (0, lo))\n V_block_ptr = tl.advance(V_block_ptr, (lo, 0))\n # loop over k, v and update accumulator\n # \u5faa\u73af\u904d\u5386 k \u548c v\uff0c\u5e76\u66f4\u65b0\u7d2f\u52a0\u5668\u3002\n for start_n in range(lo, hi, BLOCK_N):\n start_n = tl.multiple_of(start_n, BLOCK_N)\n # -- compute qk ----\n # -- \u8ba1\u7b97 qk ----\n k = tl.load(K_block_ptr)\n qk = tl.dot(q, k)\n if STAGE == 2:\n mask = offs_m[:, None] >= (start_n + offs_n[None, :])\n qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)\n m_ij = tl.maximum(m_i, tl.max(qk, 1))\n qk -= m_ij[:, None]\n else:\n m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)\n qk = qk * qk_scale - m_ij[:, None]\n p = tl.math.exp2(qk)\n l_ij = tl.sum(p, 1)\n # -- update m_i and l_i\n # -- \u66f4\u65b0 m_i \u548c l_i\n alpha = tl.math.exp2(m_i - m_ij)\n l_i = l_i * alpha + l_ij\n # -- update output accumulator --\n # -- \u66f4\u65b0\u8f93\u51fa\u7d2f\u52a0\u5668 --\n acc = acc * alpha[:, None]\n # update acc\n # \u66f4\u65b0 acc\n v = tl.load(V_block_ptr)\n if fp8_v:\n p = p.to(tl.float8e5)\n else:\n p = p.to(tl.float16)\n acc = tl.dot(p, v, acc)\n # update m_i and l_i\n # \u66f4\u65b0 m_i \u548c l_i\n m_i = m_ij\n V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))\n K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))\n return acc, l_i, m_i\n\n\n\n\n# We don't run auto-tuning every time to keep the tutorial fast. Keeping\n# the code below and commenting out the equivalent parameters is convenient for\n# re-tuning.\n# \u4e3a\u4e86\u91cd\u65b0\u8c03\u6574\uff0c\u6211\u4eec\u4e0d\u4f1a\u6bcf\u6b21\u90fd\u8fd0\u884c\u81ea\u52a8\u8c03\u4f18\u4ee5\u4fdd\u6301\u6559\u7a0b\u7684\u5feb\u901f\u6027\u3002\n# \u4fdd\u7559\u4e0b\u9762\u7684\u4ee3\u7801\u5e76\u6ce8\u91ca\u6389\u7b49\u6548\u7684\u53c2\u6570\u662f\u65b9\u4fbf\u7684\u505a\u6cd5\u3002\nconfigs = [\n triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \/\n for BM in [64, 128]\/\n for BN in [32, 64]\/\n for s in ([1] if is_hip() else [3, 4, 7])\/\n for w in [4, 8]\/\n]\n\n\n\n\ndef keep(conf):\n BLOCK_M = conf.kwargs["BLOCK_M"]\n BLOCK_N = conf.kwargs["BLOCK_N"]\n if BLOCK_M * BLOCK_N < 128 * 128 and conf.num_warps == 8:\n return False\n return True\n\n\n\n\n@triton.autotune(list(filter(keep, configs)), key=["N_CTX", "HEAD_DIM"])\n@triton.jit\ndef _attn_fwd(Q, K, V, sm_scale, M, Out, #\n stride_qz, stride_qh, stride_qm, stride_qk, #\n stride_kz, stride_kh, stride_kn, stride_kk, #\n stride_vz, stride_vh, stride_vk, stride_vn, #\n stride_oz, stride_oh, stride_om, stride_on, #\n Z, H, N_CTX, #\n HEAD_DIM: tl.constexpr, #\n BLOCK_M: tl.constexpr, #\n BLOCK_N: tl.constexpr, #\n STAGE: tl.constexpr #\n ):\n tl.static_assert(BLOCK_N <= HEAD_DIM)\n start_m = tl.program_id(0)\n off_hz = tl.program_id(1)\n off_z = off_hz \/\/ H\n off_h = off_hz % H\n qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh\n\n\n # block pointers\n # \u5757\u6307\u9488\n Q_block_ptr = tl.make_block_ptr(\n base=Q + qvk_offset,\n shape=(N_CTX, HEAD_DIM),\n strides=(stride_qm, stride_qk),\n offsets=(start_m * BLOCK_M, 0),\n block_shape=(BLOCK_M, HEAD_DIM),\n order=(1, 0),\n )\n v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)\n V_block_ptr = tl.make_block_ptr(\n base=V + qvk_offset,\n shape=(N_CTX, HEAD_DIM),\n strides=(stride_vk, stride_vn),\n offsets=(0, 0),\n block_shape=(BLOCK_N, HEAD_DIM),\n order=v_order,\n )\n K_block_ptr = tl.make_block_ptr(\n base=K + qvk_offset,\n shape=(HEAD_DIM, N_CTX),\n strides=(stride_kk, stride_kn),\n offsets=(0, 0),\n block_shape=(HEAD_DIM, BLOCK_N),\n order=(0, 1),\n )\n O_block_ptr = tl.make_block_ptr(\n base=Out + qvk_offset,\n shape=(N_CTX, HEAD_DIM),\n strides=(stride_om, stride_on),\n offsets=(start_m * BLOCK_M, 0),\n block_shape=(BLOCK_M, HEAD_DIM),\n order=(1, 0),\n )\n # initialize offsets\n # \u521d\u59cb\u5316\u504f\u79fb\u91cf\n offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)\n offs_n = tl.arange(0, BLOCK_N)\n # initialize pointer to m and l\u3001\n # \u521d\u59cb\u5316 m \u548c l \u7684\u6307\u9488\n m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")\n l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0\n acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)\n # load scales\n # \u52a0\u8f7d\u7cfb\u6570\n qk_scale = sm_scale\n qk_scale *= 1.44269504 # 1\/log(2)\n # load q: it will stay in SRAM throughout\n # \u52a0\u8f7d q: \u5b83\u5c06\u59cb\u7ec8\u4fdd\u7559\u5728 SRAM \u4e2d\n q = tl.load(Q_block_ptr)\n # stage 1: off-band\n # \u9636\u6bb5 1: off-band\n # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE\n # \u5bf9\u4e8e causal = True\uff0cSTAGE = 3\uff0c\u800c _attn_fwd_inner \u7684 STAGE \u53c2\u6570\u8bbe\u5b9a\u4e3a 1\u3002\n # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE\n # \u5bf9\u4e8e causal = False, STAGE = 1, \u800c _attn_fwd_inner \u7684 STAGE \u53c2\u6570\u8bbe\u5b9a\u4e3a 3\u3002\n if STAGE & 1:\n acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr, #\n start_m, qk_scale, #\n BLOCK_M, HEAD_DIM, BLOCK_N, #\n 4 - STAGE, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 #\n )\n # stage 2: on-band\n # \u9636\u6bb5 2: on-band\n if STAGE & 2:\n # barrier makes it easier for compielr to schedule the\n # two loops independently\n acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr, #\n start_m, qk_scale, #\n BLOCK_M, HEAD_DIM, BLOCK_N, #\n 2, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 #\n )\n # epilogue\n # \u7ed3\u5c3e\n m_i += tl.math.log2(l_i)\n acc = acc \/ l_i[:, None]\n m_ptrs = M + off_hz * N_CTX + offs_m\n tl.store(m_ptrs, m_i)\n tl.store(O_block_ptr, acc.to(Out.type.element_ty))\n\n\n\n\n@triton.jit\ndef _attn_bwd_preprocess(O, DO, #\n Delta, #\n Z, H, N_CTX, #\n BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr #\n ):\n off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)\n off_hz = tl.program_id(1)\n off_n = tl.arange(0, HEAD_DIM)\n # load\n # \u52a0\u8f7d\n o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])\n do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)\n delta = tl.sum(o * do, axis=1)\n # write-back\n # \u5199\u56de\n tl.store(Delta + off_hz * N_CTX + off_m, delta)\n\n\n\n\n# The main inner-loop logic for computing dK and dV.\n# \u5185\u5faa\u73af\u7684\u4e3b\u8981\u903b\u8f91\u662f\u8ba1\u7b97 dK \u548c dV\n@triton.jit\ndef _attn_bwd_dkdv(dk, dv, #\n Q, k, v, sm_scale, #\n DO, #\n M, D, #\n # shared by Q\/K\/V\/DO.\n # \u7531 Q\/K\/V\/DO \u5171\u4eab\n stride_tok, stride_d, #\n H, N_CTX, BLOCK_M1: tl.constexpr, #\n BLOCK_N1: tl.constexpr, #\n HEAD_DIM: tl.constexpr, #\n # Filled in by the wrapper.\n # \u7531 wrapper \u586b\u5145\n start_n, start_m, num_steps, #\n MASK: tl.constexpr):\n offs_m = start_m + tl.arange(0, BLOCK_M1)\n offs_n = start_n + tl.arange(0, BLOCK_N1)\n offs_k = tl.arange(0, HEAD_DIM)\n qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d\n do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d\n # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.\n # BLOCK_N1 \u5fc5\u987b\u662f BLOCK_M1 \u7684\u500d\u6570\uff0c\u5426\u5219\u4ee3\u7801\u5c06\u65e0\u6cd5\u6b63\u5e38\u5de5\u4f5c\u3002\n tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)\n curr_m = start_m\n step_m = BLOCK_M1\n for blk_idx in range(num_steps):\n qT = tl.load(qT_ptrs)\n # Load m before computing qk to reduce pipeline stall.\n # \u5728\u8ba1\u7b97 qk \u524d\u5148\u52a0\u8f7d m\uff0c\u4ee5\u51cf\u5c11\u6d41\u6c34\u7ebf\u505c\u987f\u3002\n offs_m = curr_m + tl.arange(0, BLOCK_M1)\n m = tl.load(M + offs_m)\n qkT = tl.dot(k, qT)\n pT = tl.math.exp2(qkT - m[None, :])\n # Autoregressive masking.\n # \u81ea\u56de\u5f52\u63a9\u7801\u3002\n if MASK:\n mask = (offs_m[None, :] >= offs_n[:, None])\n pT = tl.where(mask, pT, 0.0)\n do = tl.load(do_ptrs)\n # Compute dV.\n # \u8ba1\u7b97 dV.\n ppT = pT\n ppT = ppT.to(tl.float16)\n dv += tl.dot(ppT, do)\n # D (= delta) is pre-divided by ds_scale.\n # D\uff08= delta\uff09\u5df2\u7ecf\u901a\u8fc7 ds_scale \u8fdb\u884c\u4e86\u9884\u9664\u3002\n Di = tl.load(D + offs_m)\n # Compute dP and dS.\n # \u8ba1\u7b97 dP \u548c dS.\n dpT = tl.dot(v, tl.trans(do)).to(tl.float32)\n dsT = pT * (dpT - Di[None, :])\n dsT = dsT.to(tl.float16)\n dk += tl.dot(dsT, tl.trans(qT))\n # Increment pointers.\n # \u589e\u52a0\u6307\u9488\u3002\n curr_m += step_m\n qT_ptrs += step_m * stride_tok\n do_ptrs += step_m * stride_tok\n return dk, dv\n\n\n\n\n# the main inner-loop logic for computing dQ\n# \u5185\u5faa\u73af\u7684\u4e3b\u8981\u903b\u8f91\u662f\u8ba1\u7b97 dQ\n@triton.jit\ndef _attn_bwd_dq(dq, q, K, V, #\n do, m, D,\n # shared by Q\/K\/V\/DO.\n # \u7531 Q\/K\/V\/DO \u5171\u4eab\n stride_tok, stride_d, #\n H, N_CTX, #\n BLOCK_M2: tl.constexpr, #\n BLOCK_N2: tl.constexpr, #\n HEAD_DIM: tl.constexpr,\n # Filled in by the wrapper.\n # \u7531 wrapper \u586b\u5145\n start_m, start_n, num_steps, #\n MASK: tl.constexpr):\n offs_m = start_m + tl.arange(0, BLOCK_M2)\n offs_n = start_n + tl.arange(0, BLOCK_N2)\n offs_k = tl.arange(0, HEAD_DIM)\n kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d\n vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d\n # D (= delta) is pre-divided by ds_scale.\n # D\uff08\u5373 delta\uff09\u5728\u4f7f\u7528\u524d\u5df2\u7ecf\u88ab ds_scale \u9884\u5148\u9664\u4ee5\u3002\n Di = tl.load(D + offs_m)\n # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.\n # BLOCK_M2 \u5fc5\u987b\u662f BLOCK_N2 \u7684\u500d\u6570\uff0c\u5426\u5219\u4ee3\u7801\u5c06\u65e0\u6cd5\u6b63\u5e38\u5de5\u4f5c\u3002\n tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)\n curr_n = start_n\n step_n = BLOCK_N2\n for blk_idx in range(num_steps):\n kT = tl.load(kT_ptrs)\n vT = tl.load(vT_ptrs)\n qk = tl.dot(q, kT)\n p = tl.math.exp2(qk - m)\n # Autoregressive masking.\n # \u81ea\u56de\u5f52\u63a9\u7801\u3002\n if MASK:\n offs_n = curr_n + tl.arange(0, BLOCK_N2)\n mask = (offs_m[:, None] >= offs_n[None, :])\n p = tl.where(mask, p, 0.0)\n # Compute dP and dS.\n # \u8ba1\u7b97 dP \u548c dS.\n dp = tl.dot(do, vT).to(tl.float32)\n ds = p * (dp - Di[:, None])\n ds = ds.to(tl.float16)\n # Compute dQ.\n # \u8ba1\u7b97 dQ.\n # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.\n # \u6ce8\u610f\uff1a\u6211\u4eec\u9700\u8981\u5728\u6700\u540e\u5bf9 dq \u8fdb\u884c\u53cd\u7f29\u653e\uff0c\u56e0\u4e3a kT \u5df2\u7ecf\u8fdb\u884c\u4e86\u9884\u7f29\u653e\u3002\n dq += tl.dot(ds, tl.trans(kT))\n # \u589e\u52a0\u6307\u9488\u3002\n curr_n += step_n\n kT_ptrs += step_n * stride_tok\n vT_ptrs += step_n * stride_tok\n return dq\n\n\n\n\n@triton.jit\ndef _attn_bwd(Q, K, V, sm_scale, #\n DO, #\n DQ, DK, DV, #\n M, D,\n # shared by Q\/K\/V\/DO.\n # \u7531 Q\/K\/V\/DO \u5171\u4eab\n stride_z, stride_h, stride_tok, stride_d, #\n H, N_CTX, #\n BLOCK_M1: tl.constexpr, #\n BLOCK_N1: tl.constexpr, #\n BLOCK_M2: tl.constexpr, #\n BLOCK_N2: tl.constexpr, #\n BLK_SLICE_FACTOR: tl.constexpr, #\n HEAD_DIM: tl.constexpr):\n LN2: tl.constexpr = 0.6931471824645996 # = ln(2)\n\n\n bhid = tl.program_id(2)\n off_chz = (bhid * N_CTX).to(tl.int64)\n adj = (stride_h * (bhid % H) + stride_z * (bhid \/\/ H)).to(tl.int64)\n pid = tl.program_id(0)\n\n\n # offset pointers for batch\/head\n # \u4e3a batch\/head \u504f\u79fb\u6307\u9488\u3002\n Q += adj\n K += adj\n V += adj\n DO += adj\n DQ += adj\n DK += adj\n DV += adj\n M += off_chz\n D += off_chz\n\n\n # load scales\n # \u52a0\u8f7d\u7cfb\u6570\n offs_k = tl.arange(0, HEAD_DIM)\n\n\n start_n = pid * BLOCK_N1\n start_m = start_n\n\n\n MASK_BLOCK_M1: tl.constexpr = BLOCK_M1 \/\/ BLK_SLICE_FACTOR\n offs_n = start_n + tl.arange(0, BLOCK_N1)\n\n\n dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)\n dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)\n\n\n # load K and V: they stay in SRAM throughout the inner loop.\n # \u52a0\u8f7d K \u548c V\uff1a\u5b83\u4eec\u5728\u6574\u4e2a\u5185\u90e8\u5faa\u73af\u8fc7\u7a0b\u4e2d\u4fdd\u6301\u5728 SRAM \u4e2d\u3002\n k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)\n v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)\n\n\n num_steps = BLOCK_N1 \/\/ MASK_BLOCK_M1\n\n\n dk, dv = _attn_bwd_dkdv(dk, dv, #\n Q, k, v, sm_scale, #\n DO, #\n M, D, #\n stride_tok, stride_d, #\n H, N_CTX, #\n MASK_BLOCK_M1, BLOCK_N1, HEAD_DIM, #\n start_n, start_m, num_steps, #\n MASK=True #\n )\n\n\n start_m += num_steps * MASK_BLOCK_M1\n num_steps = (N_CTX - start_m) \/\/ BLOCK_M1\n\n\n # Compute dK and dV for non-masked blocks.\n # \u8ba1\u7b97\u975e\u63a9\u7801\u5757\u7684 dK \u548c dV\u3002\n dk, dv = _attn_bwd_dkdv( #\n dk, dv, #\n Q, k, v, sm_scale, #\n DO, #\n M, D, #\n stride_tok, stride_d, #\n H, N_CTX, #\n BLOCK_M1, BLOCK_N1, HEAD_DIM, #\n start_n, start_m, num_steps, #\n MASK=False #\n )\n\n\n dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d\n tl.store(dv_ptrs, dv)\n\n\n # Write back dK.\n # \u5199\u56de dK\n dk *= sm_scale\n dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d\n tl.store(dk_ptrs, dk)\n\n\n # THIS BLOCK DOES DQ:\n # \u8be5\u6bb5\u8ba1\u7b97 DQ\n start_m = pid * BLOCK_M2\n end_n = start_m + BLOCK_M2\n\n\n MASK_BLOCK_N2: tl.constexpr = BLOCK_N2 \/\/ BLK_SLICE_FACTOR\n offs_m = start_m + tl.arange(0, BLOCK_M2)\n\n\n q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)\n dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)\n do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)\n\n\n m = tl.load(M + offs_m)\n m = m[:, None]\n\n\n # Compute dQ for masked (diagonal) blocks.\n # \u8ba1\u7b97\u63a9\u7801\uff08\u5bf9\u89d2\u7ebf\uff09\u5757\u7684 dQ\u3002\n # NOTE: This code scans each row of QK^T backward (from right to left,\n # but inside each call to _attn_bwd_dq, from left to right), but that's\n # not due to anything important. I just wanted to reuse the loop\n # structure for dK & dV above as much as possible.\n # \u6ce8\u610f\uff1a\u8fd9\u6bb5\u4ee3\u7801\u9006\u5411\u626b\u63cf\u6bcf\u884c QK^T\uff08\u4ece\u53f3\u5411\u5de6\uff0c\u4f46\u5728\u6bcf\u6b21\u8c03\u7528 _attn_bwd_dq \u5185\u90e8\uff0c\n # \u4ece\u5de6\u5411\u53f3\uff09\uff0c\u4f46\u8fd9\u5e76\u4e0d\u91cd\u8981\u3002\u6211\u53ea\u662f\u5e0c\u671b\u5c3d\u53ef\u80fd\u591a\u5730\u91cd\u7528\u4e0a\u8ff0 dK \u548c dV \u7684\u5faa\u73af\u7ed3\u6784\u3002\n num_steps = BLOCK_M2 \/\/ MASK_BLOCK_N2\n dq = _attn_bwd_dq(dq, q, K, V, #\n do, m, D, #\n stride_tok, stride_d, #\n H, N_CTX, #\n BLOCK_M2, MASK_BLOCK_N2, HEAD_DIM, #\n start_m, end_n - num_steps * MASK_BLOCK_N2, num_steps, #\n MASK=True #\n )\n end_n -= num_steps * MASK_BLOCK_N2\n # stage 2\n # \u9636\u6bb5 2\n num_steps = end_n \/\/ BLOCK_N2\n dq = _attn_bwd_dq(dq, q, K, V, #\n do, m, D, #\n stride_tok, stride_d, #\n H, N_CTX, #\n BLOCK_M2, BLOCK_N2, HEAD_DIM, #\n start_m, end_n - num_steps * BLOCK_N2, num_steps, #\n MASK=False #\n )\n # Write back dQ.\n # \u5199\u56de dQ\n dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d\n dq *= LN2\n tl.store(dq_ptrs, dq)\n\n\n\n\nclass _attention(torch.autograd.Function):\n\n\n @staticmethod\n def forward(ctx, q, k, v, causal, sm_scale):\n # shape constraints\n # \u5f62\u72b6\u7ea6\u675f\n HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]\n # when v is in float8_e5m2 it is transposed.\n # \u5f53 v \u5728 float8_e5m2 \u683c\u5f0f\u4e0b\u65f6\uff0c\u8fdb\u884c\u8f6c\u7f6e\u3002\n HEAD_DIM_V = v.shape[-1]\n assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V\n assert HEAD_DIM_K in {16, 32, 64, 128, 256}\n o = torch.empty_like(q)\n stage = 3 if causal else 1\n extra_kern_args = {}\n # Tuning for AMD target\n # \u4e3a AMD \u8bbe\u5907\u8c03\u6574\n if is_hip():\n waves_per_eu = 3 if HEAD_DIM_K <= 64 else 2\n extra_kern_args = {"waves_per_eu": waves_per_eu, "allow_flush_denorm": True}\n\n\n grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)\n M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)\n _attn_fwd[grid](\n q, k, v, sm_scale, M, o, #\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 q.shape[0], q.shape[1], #\n N_CTX=q.shape[2], #\n HEAD_DIM=HEAD_DIM_K, #\n STAGE=stage, #\n **extra_kern_args)\n\n\n ctx.save_for_backward(q, k, v, o, M)\n ctx.grid = grid\n ctx.sm_scale = sm_scale\n ctx.HEAD_DIM = HEAD_DIM_K\n ctx.causal = causal\n return o\n\n\n @staticmethod\n def backward(ctx, do):\n q, k, v, o, M = ctx.saved_tensors\n assert do.is_contiguous()\n assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()\n dq = torch.empty_like(q)\n dk = torch.empty_like(k)\n dv = torch.empty_like(v)\n BATCH, N_HEAD, N_CTX = q.shape[:3]\n PRE_BLOCK = 128\n NUM_WARPS, NUM_STAGES = 4, 5\n BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 128, 128, 32\n BLK_SLICE_FACTOR = 2\n RCP_LN2 = 1.4426950408889634 # = 1.0 \/ ln(2)\n arg_k = k\n arg_k = arg_k * (ctx.sm_scale * RCP_LN2)\n PRE_BLOCK = 128\n assert N_CTX % PRE_BLOCK == 0\n pre_grid = (N_CTX \/\/ PRE_BLOCK, BATCH * N_HEAD)\n delta = torch.empty_like(M)\n _attn_bwd_preprocess[pre_grid](\n o, do, #\n delta, #\n BATCH, N_HEAD, N_CTX, #\n BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM #\n )\n grid = (N_CTX \/\/ BLOCK_N1, 1, BATCH * N_HEAD)\n _attn_bwd[grid](\n q, arg_k, v, ctx.sm_scale, do, dq, dk, dv, #\n M, delta, #\n q.stride(0), q.stride(1), q.stride(2), q.stride(3), #\n N_HEAD, N_CTX, #\n BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, #\n BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, #\n BLK_SLICE_FACTOR=BLK_SLICE_FACTOR, #\n HEAD_DIM=ctx.HEAD_DIM, #\n num_warps=NUM_WARPS, #\n num_stages=NUM_STAGES #\n )\n\n\n return dq, dk, dv, None, None\n\n\n\n\nattention = _attention.apply\n\n\n\n\n@pytest.mark.parametrize("Z, H, N_CTX, HEAD_DIM", [(1, 2, 1024, 64)])\n@pytest.mark.parametrize("causal", [True])\ndef test_op(Z, H, N_CTX, HEAD_DIM, causal, dtype=torch.float16):\n torch.manual_seed(20)\n q = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())\n k = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())\n v = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device="cuda").normal_(mean=0.0, std=0.5).requires_grad_())\n sm_scale = 0.5\n dout = torch.randn_like(q)\n # reference implementation\n # \u53c2\u8003\u5b9e\u73b0\n M = torch.tril(torch.ones((N_CTX, N_CTX), device="cuda"))\n p = torch.matmul(q, k.transpose(2, 3)) * sm_scale\n if causal:\n p[:, :, M == 0] = float("-inf")\n p = torch.softmax(p.float(), dim=-1).half()\n # p = torch.exp(p)\n ref_out = torch.matmul(p, v)\n ref_out.backward(dout)\n ref_dv, v.grad = v.grad.clone(), None\n ref_dk, k.grad = k.grad.clone(), None\n ref_dq, q.grad = q.grad.clone(), None\n # triton implementation\n # triton \u5b9e\u73b0\n tri_out = attention(q, k, v, causal, sm_scale).half()\n tri_out.backward(dout)\n tri_dv, v.grad = v.grad.clone(), None\n tri_dk, k.grad = k.grad.clone(), None\n tri_dq, q.grad = q.grad.clone(), None\n # compare\n # \u6bd4\u8f83\n assert torch.allclose(ref_out, tri_out, atol=1e-2, rtol=0)\n rtol = 0.0\n # Relative tolerance workaround for known hardware limitation of MI200 GPU.\n # \u9488\u5bf9 MI200 GPU \u5df2\u77e5\u7684\u786c\u4ef6\u9650\u5236\uff0c\u4f7f\u7528\u76f8\u5bf9\u5bb9\u5dee\u7684\u89e3\u51b3\u65b9\u6cd5\u3002\n # For details see https:\/\/pytorch.org\/docs\/stable\/notes\/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices\n # \u8be6\u60c5\u89c1 https:\/\/pytorch.org\/docs\/stable\/notes\/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices\n if torch.version.hip is not None and triton.runtime.driver.active.get_current_target().arch == "gfx90a":\n rtol = 1e-2\n assert torch.allclose(ref_dv, tri_dv, atol=1e-2, rtol=rtol)\n assert torch.allclose(ref_dk, tri_dk, atol=1e-2, rtol=rtol)\n assert torch.allclose(ref_dq, tri_dq, atol=1e-2, rtol=rtol)\n\n\n\n\ntry:\n from flash_attn.flash_attn_interface import \/\n flash_attn_qkvpacked_func as flash_attn_func\n HAS_FLASH = True\nexcept BaseException:\n HAS_FLASH = False\n\n\nTORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2')\nBATCH, N_HEADS, HEAD_DIM = 4, 32, 64\n# vary seq length for fixed head and batch=4\n# \u5bf9\u4e8e\u56fa\u5b9a\u7684 head \u548c batch \u4e3a 4\uff0c\u53d8\u5316\u5e8f\u5217\u957f\u5ea6\u3002\nconfigs = []\nfor mode in ["fwd", "bwd"]:\n for causal in [True, False]:\n if mode == "bwd" and not causal:\n continue\n configs.append(\n triton.testing.Benchmark(\n x_names=["N_CTX"],\n x_vals=[2**i for i in range(10, 15)],\n line_arg="provider",\n line_vals=["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else []) +\n (["flash"] if HAS_FLASH else []),\n line_names=["Triton [FP16]"] + (["Triton [FP8]"] if TORCH_HAS_FP8 else []) +\n (["Flash-2"] if HAS_FLASH else []),\n styles=[("red", "-"), ("blue", "-")],\n ylabel="ms",\n plot_name=f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}-causal={causal}",\n args={\n "H": N_HEADS,\n "BATCH": BATCH,\n "HEAD_DIM": HEAD_DIM,\n "mode": mode,\n "causal": causal,\n },\n ))\n\n\n\n\n@triton.testing.perf_report(configs)\ndef bench_flash_attention(BATCH, H, N_CTX, HEAD_DIM, causal, mode, provider, device="cuda"):\n assert mode in ["fwd", "bwd"]\n warmup = 25\n rep = 100\n dtype = torch.float16\n if "triton" in provider:\n q = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)\n k = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)\n v = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)\n if mode == "fwd" and "fp8" in provider:\n q = q.to(torch.sfloat8_e5m2)\n k = k.to(torch.float8_e5m2)\n v = v.permute(0, 1, 3, 2).contiguous()\n v = v.permute(0, 1, 3, 2)\n v = v.to(torch.float8_e5m2)\n sm_scale = 1.3\n fn = lambda: attention(q, k, v, causal, sm_scale)\n if mode == "bwd":\n o = fn()\n do = torch.randn_like(o)\n fn = lambda: o.backward(do, retain_graph=True)\n ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)\n if provider == "flash":\n qkv = torch.randn((BATCH, N_CTX, 3, H, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)\n fn = lambda: flash_attn_func(qkv, causal=causal)\n if mode == "bwd":\n o = fn()\n do = torch.randn_like(o)\n fn = lambda: o.backward(do, retain_graph=True)\n ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)\n flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * HEAD_DIM\n total_flops = 2 * flops_per_matmul\n if causal:\n total_flops *= 0.5\n if mode == "bwd":\n total_flops *= 2.5 # 2.0(bwd) + 0.5(recompute)\n return total_flops \/ ms * 1e-9\n\n\n\n\nif __name__ == "__main__":\n # only works on post-Ampere GPUs right now\n # \u76ee\u524d\u53ea\u9002\u7528\u4e8e\u5b89\u57f9\u67b6\u6784 GPU\u3002\n bench_flash_attention.run(save_path=".", print_data=True)<\/code><\/pre>\n