Update triton_flash_atn.py

triton_flash_atn.py

@@ -25,27 +25,17 @@ import triton.language as tl
 # TORCH_HAS_FP8E4B8 = hasattr(torch, 'float8_e4m3fnuz')
 
 # AMD E5M2B16
-TORCH_HAS_FP8E5B16 = hasattr(torch, "float8_e5m2fnuz")
+TORCH_HAS_FP8E5B16 = hasattr(torch, 'float8_e5m2fnuz')
 
 
 @triton.jit
-def _attn_fwd_inner(
- acc,
- l_i,
- m_i,
- q,
- K_block_ptr,
- V_block_ptr,
- start_m,
- BLOCK_M: tl.constexpr,
- BLOCK_DMODEL: tl.constexpr,
- BLOCK_N: tl.constexpr,
- STAGE: tl.constexpr,
- offs_m: tl.constexpr,
- offs_n: tl.constexpr,
- N_CTX,
- pre_load_v: tl.constexpr,
-):
+def _attn_fwd_inner(acc, l_i, m_i, q,
+ K_block_ptr, V_block_ptr,
+ start_m,
+ BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
+ STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr,
+ N_CTX,
+ pre_load_v: tl.constexpr):
  # range of values handled by this stage
  if STAGE == 1:
  lo, hi = 0, start_m * BLOCK_M
@@ -93,119 +83,37 @@ def _attn_fwd_inner(
 # re-tuning.
 @triton.autotune(
  configs=[
- triton.Config(
- {
- "BLOCK_M": 64,
- "BLOCK_N": 16,
- "waves_per_eu": 2,
- "slice_k_tile": 0,
- "pre_load_v": False,
- },
- num_stages=1,
- num_warps=2,
- ),
- triton.Config(
- {
- "BLOCK_M": 64,
- "BLOCK_N": 16,
- "waves_per_eu": 2,
- "slice_k_tile": 32,
- "pre_load_v": False,
- },
- num_stages=1,
- num_warps=2,
- ),
- triton.Config(
- {
- "BLOCK_M": 32,
- "BLOCK_N": 32,
- "waves_per_eu": 2,
- "slice_k_tile": 0,
- "pre_load_v": False,
- },
- num_stages=1,
- num_warps=1,
- ),
- triton.Config(
- {
- "BLOCK_M": 32,
- "BLOCK_N": 32,
- "waves_per_eu": 2,
- "slice_k_tile": 32,
- "pre_load_v": False,
- },
- num_stages=1,
- num_warps=1,
- ),
- triton.Config(
- {
- "BLOCK_M": 64,
- "BLOCK_N": 32,
- "waves_per_eu": 2,
- "slice_k_tile": 0,
- "pre_load_v": False,
- },
- num_stages=1,
- num_warps=2,
- ),
- triton.Config(
- {
- "BLOCK_M": 32,
- "BLOCK_N": 16,
- "waves_per_eu": 3,
- "slice_k_tile": 0,
- "pre_load_v": True,
- },
- num_stages=1,
- num_warps=1,
- ),
- triton.Config(
- {
- "BLOCK_M": 32,
- "BLOCK_N": 16,
- "waves_per_eu": 3,
- "slice_k_tile": 0,
- "pre_load_v": False,
- },
- num_stages=1,
- num_warps=1,
- ),
+ triton.Config({'BLOCK_M': 64, 'BLOCK_N': 16, 'waves_per_eu': 2,
+ 'slice_k_tile': 0, 'pre_load_v': False}, num_stages=1, num_warps=2),
+ triton.Config({'BLOCK_M': 64, 'BLOCK_N': 16, 'waves_per_eu': 2,
+ 'slice_k_tile': 32, 'pre_load_v': False}, num_stages=1, num_warps=2),
+ triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 2,
+ 'slice_k_tile': 0, 'pre_load_v': False}, num_stages=1, num_warps=1),
+ triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 2,
+ 'slice_k_tile': 32, 'pre_load_v': False}, num_stages=1, num_warps=1),
+ triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32, 'waves_per_eu': 2,
+ 'slice_k_tile': 0, 'pre_load_v': False}, num_stages=1, num_warps=2),
+ triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 3,
+ 'slice_k_tile': 0, 'pre_load_v': True}, num_stages=1, num_warps=1),
+ triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 3,
+ 'slice_k_tile': 0, 'pre_load_v': False}, num_stages=1, num_warps=1),
  ],
- key=["Z", "H", "N_CTX", "STAGE", "BLOCK_DMODEL"],
+ key=['Z', 'H', 'N_CTX', 'STAGE', 'BLOCK_DMODEL'],
 )
 @triton.jit
-def _attn_fwd(
- Q,
- K,
- V,
- sm_scale,
- M,
- Out,
- stride_qz,
- stride_qh,
- stride_qm,
- stride_qk,
- stride_kz,
- stride_kh,
- stride_kn,
- stride_kk,
- stride_vz,
- stride_vh,
- stride_vk,
- stride_vn,
- stride_oz,
- stride_oh,
- stride_om,
- stride_on,
- Z,
- H,
- N_CTX,
- BLOCK_DMODEL: tl.constexpr,
- STAGE: tl.constexpr,
- BLOCK_M: tl.constexpr,
- BLOCK_N: tl.constexpr,
- pre_load_v: tl.constexpr,
-):
+def _attn_fwd(Q, K, V, sm_scale, M, Out,
+ stride_qz, stride_qh, stride_qm, stride_qk,
+ stride_kz, stride_kh, stride_kn, stride_kk,
+ stride_vz, stride_vh, stride_vk, stride_vn,
+ stride_oz, stride_oh, stride_om, stride_on,
+ Z, H,
+ N_CTX,
+ BLOCK_DMODEL: tl.constexpr,
+ STAGE: tl.constexpr,
+ BLOCK_M: tl.constexpr,
+ BLOCK_N: tl.constexpr,
+ pre_load_v: tl.constexpr,
+ ):
  start_m = tl.program_id(0)
  off_hz = tl.program_id(1)
  qvk_offset = off_hz * stride_qh
@@ -261,45 +169,23 @@ def _attn_fwd(
  # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE
  # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE
  if STAGE & 1:
- acc, l_i, m_i = _attn_fwd_inner(
- acc,
- l_i,
- m_i,
- q,
- K_block_ptr,
- V_block_ptr,
- start_m,
- BLOCK_M,
- BLOCK_DMODEL,
- BLOCK_N,
- 4 - STAGE,
- offs_m,
- offs_n,
- N_CTX,
- pre_load_v,
- )
+ acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr,
+ start_m,
+ BLOCK_M, BLOCK_DMODEL, BLOCK_N,
+ 4 - STAGE, offs_m, offs_n, N_CTX,
+ pre_load_v,
+ )
  # stage 2: on-band
  if STAGE & 2:
  # barrier makes it easier for compielr to schedule the
  # two loops independently
  tl.debug_barrier()
- acc, l_i, m_i = _attn_fwd_inner(
- acc,
- l_i,
- m_i,
- q,
- K_block_ptr,
- V_block_ptr,
- start_m,
- BLOCK_M,
- BLOCK_DMODEL,
- BLOCK_N,
- 2,
- offs_m,
- offs_n,
- N_CTX,
- pre_load_v,
- )
+ acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr,
+ start_m,
+ BLOCK_M, BLOCK_DMODEL, BLOCK_N,
+ 2, offs_m, offs_n, N_CTX,
+ pre_load_v,
+ )
  # epilogue
  # write back m
  acc = acc / l_i[:, None]
@@ -309,46 +195,36 @@ def _attn_fwd(
 
 
 @triton.jit
-def _attn_bwd_preprocess(
- O, DO, Delta, Z, H, N_CTX, BLOCK_M: tl.constexpr, D_HEAD: tl.constexpr
-):
+def _attn_bwd_preprocess(O, DO,
+ Delta,
+ Z, H, N_CTX,
+ BLOCK_M: tl.constexpr, D_HEAD: tl.constexpr
+ ):
  off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
  off_hz = tl.program_id(1)
  off_n = tl.arange(0, D_HEAD)
- o = tl.load(O + off_hz * D_HEAD * N_CTX + off_m[:, None] * D_HEAD + off_n[None, :])
- do = tl.load(
- DO + off_hz * D_HEAD * N_CTX + off_m[:, None] * D_HEAD + off_n[None, :]
- ).to(tl.float32)
+ o = tl.load(O + off_hz * D_HEAD * N_CTX +
+ off_m[:, None] * D_HEAD + off_n[None, :])
+ do = tl.load(DO + off_hz * D_HEAD * N_CTX +
+ off_m[:, None] * D_HEAD + off_n[None, :]).to(tl.float32)
  delta = tl.sum(o * do, axis=1)
  tl.store(Delta + off_hz * N_CTX + off_m, delta)
 
 
 # The main inner-loop logic for computing dK and dV.
 @triton.jit
-def _attn_bwd_dkdv(
- dk,
- dv,
- Q,
- k,
- v,
- sm_scale,
- DO,
- M,
- D,
- # shared by Q/K/V/DO.
- stride_tok,
- stride_d,
- H,
- N_CTX,
- BLOCK_M1: tl.constexpr,
- BLOCK_N1: tl.constexpr,
- BLOCK_DMODEL: tl.constexpr,
- # Filled in by the wrapper.
- start_n,
- start_m,
- num_steps,
- MASK: tl.constexpr,
-):
+def _attn_bwd_dkdv(dk, dv,
+ Q, k, v, sm_scale,
+ DO,
+ M, D,
+ # shared by Q/K/V/DO.
+ stride_tok, stride_d,
+ H, N_CTX, BLOCK_M1: tl.constexpr,
+ BLOCK_N1: tl.constexpr,
+ BLOCK_DMODEL: tl.constexpr,
+ # Filled in by the wrapper.
+ start_n, start_m, num_steps,
+ MASK: tl.constexpr):
  offs_m = start_m + tl.arange(0, BLOCK_M1)
  offs_n = start_n + tl.arange(0, BLOCK_N1)
  offs_k = tl.arange(0, BLOCK_DMODEL)
@@ -358,7 +234,7 @@ def _attn_bwd_dkdv(
  strides=(stride_d, stride_tok),
  offsets=(0, start_m),
  block_shape=(BLOCK_DMODEL, BLOCK_M1),
- order=(0, 1),
+ order=(0, 1)
  )
  DO_block_ptr = tl.make_block_ptr(
  base=DO,
@@ -366,7 +242,7 @@ def _attn_bwd_dkdv(
  strides=(stride_tok, stride_d),
  offsets=(start_m, 0),
  block_shape=(BLOCK_M1, BLOCK_DMODEL),
- order=(1, 0),
+ order=(1, 0)
  )
  # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
  tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
@@ -381,7 +257,7 @@ def _attn_bwd_dkdv(
  pT = tl.math.exp2(qkT - m[None, :])
  # Autoregressive masking.
  if MASK:
- mask = offs_m[None, :] >= offs_n[:, None]
+ mask = (offs_m[None, :] >= offs_n[:, None])
  pT = tl.where(mask, pT, 0.0)
  do = tl.load(DO_block_ptr)
  # Compute dV.
@@ -404,28 +280,17 @@ def _attn_bwd_dkdv(
 
 # the main inner-loop logic for computing dQ
 @triton.jit
-def _attn_bwd_dq(
- dq,
- q,
- K,
- V,
- do,
- m,
- D,
- # shared by Q/K/V/DO.
- stride_tok,
- stride_d,
- H,
- N_CTX,
- BLOCK_M2: tl.constexpr,
- BLOCK_N2: tl.constexpr,
- BLOCK_DMODEL: tl.constexpr,
- # Filled in by the wrapper.
- start_m,
- start_n,
- num_steps,
- MASK: tl.constexpr,
-):
+def _attn_bwd_dq(dq, q, K, V,
+ do, m, D,
+ # shared by Q/K/V/DO.
+ stride_tok, stride_d,
+ H, N_CTX,
+ BLOCK_M2: tl.constexpr,
+ BLOCK_N2: tl.constexpr,
+ BLOCK_DMODEL: tl.constexpr,
+ # Filled in by the wrapper.
+ start_m, start_n, num_steps,
+ MASK: tl.constexpr):
  offs_m = start_m + tl.arange(0, BLOCK_M2)
  offs_n = start_n + tl.arange(0, BLOCK_N2)
  offs_k = tl.arange(0, BLOCK_DMODEL)
@@ -435,7 +300,7 @@ def _attn_bwd_dq(
  strides=(stride_d, stride_tok),
  offsets=(0, start_n),
  block_shape=(BLOCK_DMODEL, BLOCK_N2),
- order=(0, 1),
+ order=(0, 1)
  )
  VT_block_ptr = tl.make_block_ptr(
  base=V,
@@ -443,7 +308,7 @@ def _attn_bwd_dq(
  strides=(stride_d, stride_tok),
  offsets=(0, start_n),
  block_shape=(BLOCK_DMODEL, BLOCK_N2),
- order=(0, 1),
+ order=(0, 1)
  )
  # D (= delta) is pre-divided by ds_scale.
  Di = tl.load(D + offs_m)
@@ -458,7 +323,7 @@ def _attn_bwd_dq(
  # Autoregressive masking.
  if MASK:
  offs_n = curr_n + tl.arange(0, BLOCK_N2)
- mask = offs_m[:, None] >= offs_n[None, :]
+ mask = (offs_m[:, None] >= offs_n[None, :])
  p = tl.where(mask, p, 0.0)
  # Compute dP and dS.
  vT = tl.load(VT_block_ptr)
@@ -477,135 +342,42 @@ def _attn_bwd_dq(
 
 @triton.autotune(
  configs=[
- triton.Config(
- {
- "BLOCK_M1": 32,
- "BLOCK_N1": 64,
- "BLOCK_M2": 64,
- "BLOCK_N2": 32,
- "BLK_SLICE_FACTOR": 1,
- },
- num_stages=1,
- num_warps=4,
- ),
- triton.Config(
- {
- "BLOCK_M1": 32,
- "BLOCK_N1": 64,
- "BLOCK_M2": 64,
- "BLOCK_N2": 32,
- "BLK_SLICE_FACTOR": 2,
- },
- num_stages=1,
- num_warps=4,
- ),
- triton.Config(
- {
- "BLOCK_M1": 64,
- "BLOCK_N1": 128,
- "BLOCK_M2": 128,
- "BLOCK_N2": 64,
- "BLK_SLICE_FACTOR": 1,
- },
- num_stages=1,
- num_warps=4,
- ),
- triton.Config(
- {
- "BLOCK_M1": 64,
- "BLOCK_N1": 128,
- "BLOCK_M2": 128,
- "BLOCK_N2": 64,
- "BLK_SLICE_FACTOR": 2,
- },
- num_stages=1,
- num_warps=4,
- ),
- triton.Config(
- {
- "BLOCK_M1": 64,
- "BLOCK_N1": 64,
- "BLOCK_M2": 64,
- "BLOCK_N2": 64,
- "BLK_SLICE_FACTOR": 1,
- },
- num_stages=1,
- num_warps=4,
- ),
- triton.Config(
- {
- "BLOCK_M1": 64,
- "BLOCK_N1": 64,
- "BLOCK_M2": 64,
- "BLOCK_N2": 64,
- "BLK_SLICE_FACTOR": 2,
- },
- num_stages=1,
- num_warps=4,
- ),
- triton.Config(
- {
- "BLOCK_M1": 32,
- "BLOCK_N1": 128,
- "BLOCK_M2": 128,
- "BLOCK_N2": 32,
- "BLK_SLICE_FACTOR": 1,
- },
- num_stages=1,
- num_warps=4,
- ),
- triton.Config(
- {
- "BLOCK_M1": 32,
- "BLOCK_N1": 128,
- "BLOCK_M2": 128,
- "BLOCK_N2": 32,
- "BLK_SLICE_FACTOR": 2,
- },
- num_stages=1,
- num_warps=4,
- ),
- triton.Config(
- {
- "BLOCK_M1": 32,
- "BLOCK_N1": 128,
- "BLOCK_M2": 128,
- "BLOCK_N2": 32,
- "BLK_SLICE_FACTOR": 2,
- },
- num_stages=1,
- num_warps=8,
- ),
+ triton.Config({'BLOCK_M1': 32, 'BLOCK_N1': 64, 'BLOCK_M2': 64, 'BLOCK_N2': 32, 'BLK_SLICE_FACTOR': 1},
+ num_stages=1, num_warps=4),
+ triton.Config({'BLOCK_M1': 32, 'BLOCK_N1': 64, 'BLOCK_M2': 64, 'BLOCK_N2': 32, 'BLK_SLICE_FACTOR': 2},
+ num_stages=1, num_warps=4),
+ triton.Config({'BLOCK_M1': 64, 'BLOCK_N1': 128, 'BLOCK_M2': 128, 'BLOCK_N2': 64, 'BLK_SLICE_FACTOR': 1},
+ num_stages=1, num_warps=4),
+ triton.Config({'BLOCK_M1': 64, 'BLOCK_N1': 128, 'BLOCK_M2': 128, 'BLOCK_N2': 64, 'BLK_SLICE_FACTOR': 2},
+ num_stages=1, num_warps=4),
+ triton.Config({'BLOCK_M1': 64, 'BLOCK_N1': 64, 'BLOCK_M2': 64, 'BLOCK_N2': 64, 'BLK_SLICE_FACTOR': 1},
+ num_stages=1, num_warps=4),
+ triton.Config({'BLOCK_M1': 64, 'BLOCK_N1': 64, 'BLOCK_M2': 64, 'BLOCK_N2': 64, 'BLK_SLICE_FACTOR': 2},
+ num_stages=1, num_warps=4),
+ triton.Config({'BLOCK_M1': 32, 'BLOCK_N1': 128, 'BLOCK_M2': 128, 'BLOCK_N2': 32, 'BLK_SLICE_FACTOR': 1},
+ num_stages=1, num_warps=4),
+ triton.Config({'BLOCK_M1': 32, 'BLOCK_N1': 128, 'BLOCK_M2': 128, 'BLOCK_N2': 32, 'BLK_SLICE_FACTOR': 2},
+ num_stages=1, num_warps=4),
+ triton.Config({'BLOCK_M1': 32, 'BLOCK_N1': 128, 'BLOCK_M2': 128, 'BLOCK_N2': 32, 'BLK_SLICE_FACTOR': 2},
+ num_stages=1, num_warps=8),
  ],
- key=["H", "N_CTX", "BLOCK_DMODEL"],
+ key=['H', 'N_CTX', 'BLOCK_DMODEL'],
 )
 @triton.jit
-def _attn_bwd(
- Q,
- K,
- V,
- sm_scale,
- DO,
- DQ,
- DK,
- DV,
- M,
- D,
- # shared by Q/K/V/DO.
- stride_z,
- stride_h,
- stride_tok,
- stride_d,
- # H = 16, N_CTX = 1024
- H,
- N_CTX,
- BLOCK_DMODEL: tl.constexpr,
- BLOCK_M1: tl.constexpr,
- BLOCK_N1: tl.constexpr,
- BLOCK_M2: tl.constexpr,
- BLOCK_N2: tl.constexpr,
- BLK_SLICE_FACTOR: tl.constexpr,
-):
+def _attn_bwd(Q, K, V, sm_scale,
+ DO,
+ DQ, DK, DV,
+ M, D,
+ # shared by Q/K/V/DO.
+ stride_z, stride_h, stride_tok, stride_d,
+ # H = 16, N_CTX = 1024
+ H, N_CTX,
+ BLOCK_DMODEL: tl.constexpr,
+ BLOCK_M1: tl.constexpr,
+ BLOCK_N1: tl.constexpr,
+ BLOCK_M2: tl.constexpr,
+ BLOCK_N2: tl.constexpr,
+ BLK_SLICE_FACTOR: tl.constexpr):
  LN2: tl.constexpr = 0.6931471824645996 # = ln(2)
 
  bhid = tl.program_id(2)
@@ -661,54 +433,31 @@ def _attn_bwd(
 
  num_steps = BLOCK_N1 // MASK_BLOCK_M1
 
- dk, dv = _attn_bwd_dkdv(
- dk,
- dv,
- Q,
- k,
- v,
- sm_scale,
- DO,
- M,
- D,
- stride_tok,
- stride_d,
- H,
- N_CTX,
- MASK_BLOCK_M1,
- BLOCK_N1,
- BLOCK_DMODEL,
- start_n,
- start_m,
- num_steps,
- MASK=True,
- )
+ dk, dv = _attn_bwd_dkdv(dk, dv,
+ Q, k, v, sm_scale,
+ DO,
+ M, D,
+ stride_tok, stride_d,
+ H, N_CTX,
+ MASK_BLOCK_M1, BLOCK_N1, BLOCK_DMODEL,
+ start_n, start_m, num_steps,
+ MASK=True
+ )
 
  start_m += num_steps * MASK_BLOCK_M1
  num_steps = (N_CTX - start_m) // BLOCK_M1
 
  # Compute dK and dV for non-masked blocks.
  dk, dv = _attn_bwd_dkdv(
- dk,
- dv,
- Q,
- k,
- v,
- sm_scale,
+ dk, dv,
+ Q, k, v, sm_scale,
  DO,
- M,
- D,
- stride_tok,
- stride_d,
- H,
- N_CTX,
- BLOCK_M1,
- BLOCK_N1,
- BLOCK_DMODEL,
- start_n,
- start_m,
- num_steps,
- MASK=False,
+ M, D,
+ stride_tok, stride_d,
+ H, N_CTX,
+ BLOCK_M1, BLOCK_N1, BLOCK_DMODEL,
+ start_n, start_m, num_steps,
+ MASK=False
  )
 
  DV_block_ptrs = tl.make_block_ptr(
@@ -717,7 +466,7 @@ def _attn_bwd(
  strides=(stride_tok, stride_d),
  offsets=(start_n, 0),
  block_shape=(BLOCK_N1, BLOCK_DMODEL),
- order=(1, 0),
+ order=(1, 0)
  )
  tl.store(DV_block_ptrs, dv.to(tl.float16))
 
@@ -729,7 +478,7 @@ def _attn_bwd(
  strides=(stride_tok, stride_d),
  offsets=(start_n, 0),
  block_shape=(BLOCK_N1, BLOCK_DMODEL),
- order=(1, 0),
+ order=(1, 0)
  )
  tl.store(DK_block_ptrs, dk.to(tl.float16))
 
@@ -746,7 +495,7 @@ def _attn_bwd(
  strides=(stride_tok, stride_d),
  offsets=(start_m, 0),
  block_shape=(BLOCK_M2, BLOCK_DMODEL),
- order=(1, 0),
+ order=(1, 0)
  )
 
  DO_block_ptr = tl.make_block_ptr(
@@ -755,7 +504,7 @@ def _attn_bwd(
  strides=(stride_tok, stride_d),
  offsets=(start_m, 0),
  block_shape=(BLOCK_M2, BLOCK_DMODEL),
- order=(1, 0),
+ order=(1, 0)
  )
  q = tl.load(Q_block_ptr)
  do = tl.load(DO_block_ptr)
@@ -770,49 +519,25 @@ def _attn_bwd(
  # not due to anything important. I just wanted to reuse the loop
  # structure for dK & dV above as much as possible.
  num_steps = BLOCK_M2 // MASK_BLOCK_N2
- dq = _attn_bwd_dq(
- dq,
- q,
- K,
- V,
- do,
- m,
- D,
- stride_tok,
- stride_d,
- H,
- N_CTX,
- BLOCK_M2,
- MASK_BLOCK_N2,
- BLOCK_DMODEL,
- start_m,
- end_n - num_steps * MASK_BLOCK_N2,
- num_steps,
- MASK=True,
- )
+ dq = _attn_bwd_dq(dq, q, K, V,
+ do, m, D,
+ stride_tok, stride_d,
+ H, N_CTX,
+ BLOCK_M2, MASK_BLOCK_N2, BLOCK_DMODEL,
+ start_m, end_n - num_steps * MASK_BLOCK_N2, num_steps,
+ MASK=True
+ )
  end_n -= num_steps * MASK_BLOCK_N2
  # stage 2
  num_steps = end_n // BLOCK_N2
- dq = _attn_bwd_dq(
- dq,
- q,
- K,
- V,
- do,
- m,
- D,
- stride_tok,
- stride_d,
- H,
- N_CTX,
- BLOCK_M2,
- BLOCK_N2,
- BLOCK_DMODEL,
- start_m,
- end_n - num_steps * BLOCK_N2,
- num_steps,
- MASK=False,
- )
+ dq = _attn_bwd_dq(dq, q, K, V,
+ do, m, D,
+ stride_tok, stride_d,
+ H, N_CTX,
+ BLOCK_M2, BLOCK_N2, BLOCK_DMODEL,
+ start_m, end_n - num_steps * BLOCK_N2, num_steps,
+ MASK=False
+ )
  # Write back dQ.
  DQ_block_ptr = tl.make_block_ptr(
  base=DQ,
@@ -820,7 +545,7 @@ def _attn_bwd(
  strides=(stride_tok, stride_d),
  offsets=(start_m, 0),
  block_shape=(BLOCK_M2, BLOCK_DMODEL),
- order=(1, 0),
+ order=(1, 0)
  )
  dq *= LN2
  tl.store(DQ_block_ptr, dq.to(tl.float16))
@@ -849,41 +574,20 @@ class _attention(torch.autograd.Function):
  num_stages = 7 if Lk >= 64 else 3
  stage = 3 if causal else 1
 
- def grid(META):
- return (
- triton.cdiv(q.shape[2], META["BLOCK_M"]),
- q.shape[0] * q.shape[1],
- 1,
- )
-
- M = torch.empty(
- (q.shape[0] * q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32
+ def grid(META): return (
+ triton.cdiv(q.shape[2], META['BLOCK_M']),
+ q.shape[0] * q.shape[1],
+ 1
  )
+ M = torch.empty((q.shape[0] * q.shape[1], q.shape[2]),
+ device=q.device, dtype=torch.float32)
  _attn_fwd[grid](
- q,
- k,
- v,
- sm_scale,
- M,
- o,
- q.stride(0),
- q.stride(1),
- q.stride(2),
- q.stride(3),
- k.stride(0),
- k.stride(1),
- k.stride(2),
- k.stride(3),
- v.stride(0),
- v.stride(1),
- v.stride(2),
- v.stride(3),
- o.stride(0),
- o.stride(1),
- o.stride(2),
- o.stride(3),
- q.shape[0],
- q.shape[1],
+ q, k, v, sm_scale, M, o,
+ q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+ k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+ v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+ o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+ q.shape[0], q.shape[1],
  N_CTX=q.shape[2],
  BLOCK_DMODEL=Lk,
  STAGE=stage,
@@ -925,39 +629,26 @@ class _attention(torch.autograd.Function):
  pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
  delta = torch.empty_like(M)
  _attn_bwd_preprocess[pre_grid](
- o,
- do,
+ o, do,
  delta,
- BATCH,
- N_HEAD,
- N_CTX,
- BLOCK_M=PRE_BLOCK,
- D_HEAD=ctx.BLOCK_DMODEL,
+ BATCH, N_HEAD, N_CTX,
+ BLOCK_M=PRE_BLOCK, D_HEAD=ctx.BLOCK_DMODEL
  )
 
- def grid(META):
- return (triton.cdiv(N_CTX, META["BLOCK_N1"]), 1, BATCH * N_HEAD)
-
+ def grid(META): return (
+ triton.cdiv(N_CTX, META['BLOCK_N1']),
+ 1,
+ BATCH * N_HEAD
+ )
  _attn_bwd[grid](
- q,
- arg_k,
- v,
- ctx.sm_scale,
- do,
- dq,
- dk,
- dv,
- M,
- delta,
- q.stride(0),
- q.stride(1),
- q.stride(2),
- q.stride(3),
- N_HEAD,
- N_CTX,
- BLOCK_DMODEL=ctx.BLOCK_DMODEL,
+ q, arg_k, v, ctx.sm_scale, do, dq, dk, dv,
+ M, delta,
+ q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+ N_HEAD, N_CTX,
+ BLOCK_DMODEL=ctx.BLOCK_DMODEL
  )
 
  return dq, dk, dv, None, None
 
 
+attention = _attention.apply