#include "kittens.cuh"
#include "utils.cpp"

#if ATTN_D == 64
#define DQ_L4 "2"
#define DQ_L10 "5"
#define DQ_L2 "1"
#define DQ_L6 "3"
#define DQ_L6_LOOP "3"
#else
#define DQ_L4 "4"
#define DQ_L10 "10"
#define DQ_L2 "2"
#define DQ_L6 "6"
#define DQ_L6_LOOP "6"
#endif

#ifndef ATTN_B
constexpr int ATTN_B = 16; // batch size
#endif

#ifndef ATTN_H
constexpr int ATTN_H = 64; // number of query heads
#endif

#ifndef ATTN_H_KV
constexpr int ATTN_H_KV = 8; // number of key/value heads (for GQA)
#endif

constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV; // queries per KV head group
constexpr int HEADS_PER_WG = (ATTN_D == 128 && GROUP_SIZE % 2 == 0) ? 2 : 1;

#ifndef ATTN_N
constexpr int ATTN_N = 1024; // sequence length
#endif

#ifndef ATTN_D
constexpr int ATTN_D = 128; // dimension
#endif
constexpr int STEP_QO = 64; // block size for QO
constexpr int BLOCK_SIZE_KV = 256; // block size for KV
constexpr int SLICE_QO = 32;
constexpr int DOT_SLICE_QO = 16;
constexpr int WARP_SIZE_KV = 64; // warp size for KV
constexpr bool causal = true;
// WINDOW>0: sliding-window backward (query i sees keys in [i-WINDOW+1, i])
#ifndef WINDOW
#define WINDOW 0
#endif

#define NUM_WARPS 4
#define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS)

using G = kittens::group<NUM_WARPS>;

using namespace kittens;

using _gl_QdO  = gl<bf16, ATTN_B, ATTN_N, ATTN_H, ATTN_D>;
using _gl_KV   = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
using _gl_dQ   = gl<bf16, ATTN_B, ATTN_H, ATTN_N, ATTN_D>;
using _gl_dKV  = gl<bf16, ATTN_B * (GROUP_SIZE / HEADS_PER_WG), ATTN_N, ATTN_H_KV, ATTN_D>;
using _gl_Lvec = gl<float, ATTN_B, ATTN_H, 1, ATTN_N>;

template<int D> struct attn_bwd_combined_globals {
  _gl_QdO Q;
  _gl_KV K, V;
  _gl_QdO dOg;
  _gl_dQ dQg;
  _gl_dKV dKg, dVg;
  _gl_Lvec L_vec, delta_vec;
  dim3 grid() { return dim3(ATTN_H / HEADS_PER_WG, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
  dim3 block() { return dim3(NUM_THREADS); }
  size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; }
};

template<int D> __launch_bounds__(NUM_THREADS, 1)
__global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr, bf16 *dO_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr, float *L_vec_ptr, float *delta_vec_ptr) {

  const int q_head_idx_fixed = blockIdx.x * HEADS_PER_WG;  // First query head handled by this workgroup.
  const int kv_head_idx = q_head_idx_fixed / GROUP_SIZE;
  const int q_head_in_group = q_head_idx_fixed % GROUP_SIZE;
  const int seq_idx = blockIdx.y;
  const int batch_idx = blockIdx.z;
  const int first_q_head = q_head_idx_fixed;

  const int warpid = kittens::warpid();
  const int j = seq_idx * NUM_WARPS + warpid;

  // optimization on loops bounds
  const int total_steps_per_head = ATTN_N / STEP_QO;
  const int j_min = seq_idx * NUM_WARPS;
  const int k_start_min = j_min * WARP_SIZE_KV;
  // first Q step that can overlap this K_span:
  const int first_step = max(0, k_start_min / STEP_QO);
#if WINDOW
  // cap the Q loop, padded by 2 masked steps: the epilogue's deferred dq path miscomputes in-window tail queries
  const int num_steps_per_head = min(total_steps_per_head - first_step, (BLOCK_SIZE_KV + WINDOW) / STEP_QO + 2);
#else
  const int num_steps_per_head = total_steps_per_head - first_step;
#endif
  const int num_steps = num_steps_per_head * HEADS_PER_WG;
  const int k_pos = j * WARP_SIZE_KV;

  constexpr float L_SCALE_FACTOR = 1.44269504089f;
  constexpr float P_SCALE_FACTOR = (D == 128) ? 0.08838834764f*1.44269504089f : 0.125f*1.44269504089f;
  constexpr float dP_SCALE_FACTOR = (D == 128) ? 0.08838834764f : 0.125f;

  // Shared tiles
  extern __shared__ alignment_dummy __shm[];
  shared_allocator al((int*)&__shm[0]);

  st_bf<BLOCK_SIZE_KV, D, st_16x16_s> (&K_j_smem) = al.allocate<st_bf<BLOCK_SIZE_KV, D, st_16x16_s>>();
  st_bf<SLICE_QO, D, st_16x32_s> (&Q_i_smem)[2][2] = al.allocate<st_bf<SLICE_QO, D, st_16x32_s>, 2, 2>();
  st_bf<SLICE_QO, D, st_16x32_s> (&dO_i_smem)[2][2] = al.allocate<st_bf<SLICE_QO, D, st_16x32_s>, 2, 2>();
  st_bf<BLOCK_SIZE_KV, DOT_SLICE_QO, st_16x16_swizzled_s> (&attn_i_smem) = al.allocate<st_bf<BLOCK_SIZE_KV, DOT_SLICE_QO, st_16x16_swizzled_s>>();
  sv_fl<STEP_QO> (&L_smem)[2] = al.allocate<sv_fl<STEP_QO>, 2>();
  sv_fl<STEP_QO> (&delta_smem)[2] = al.allocate<sv_fl<STEP_QO>, 2>();

  // Register tiles
#if ATTN_D == 64
  using Q_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<368, 375>>, 4>;
  using dO_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<78, 85>>, 4>;
  using dO_col_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<94, 101>>, 4>;
  using K_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<256, 263>, ducks::art::range<272, 279>, ducks::art::range<288, 295>, ducks::art::range<62, 69>>, 4>;
  using V_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<304, 311>, ducks::art::range<316, 323>, ducks::art::range<328, 335>, ducks::art::range<340, 347>>, 4>;
#else
  using Q_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<368, 383>>, 4>; // 16 registers - a[112:127]
  using dO_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<78, 93>>, 4>; // 16 registers - v[72:87]
  using dO_col_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<94, 109>>, 4>; // 16 registers - v[88:103]
  using K_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<256, 303>, ducks::art::range<62, 77>>, 4>; // 64 registers - a[0:47] & v[56:71]
  using V_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<304, 367>>, 4>; // 64 registers - a[48:111]
#endif
  using P_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<46, 61>>, 4>; // 16 registers - v[40:55]
  using dP_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<62, 77>>, 4>; // 16 registers - v[56:71]
  using P_bf16_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<118, 125>>, 2>; // 8 registers - v[116:123]
  using dP_bf16_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<62, 69>>, 2>; // 8 registers - v[56:63]
  using P_bf16_col_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<118, 125>>, 4>; // 8 registers
  using dP_bf16_col_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<62, 69>>, 4>; // 8 registers
  using dS_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<30, 61>>, 4>; // 32 registers - v[24:55]
#if ATTN_D == 64
  using dQ_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<110, 113>>, 4>;
#else
  using dQ_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<110, 117>>, 4>; // 8 registers - v[108:115]
#endif
  ducks::art::clobber<Q_ranges>();
  ducks::art::clobber<dO_ranges>();
  ducks::art::clobber<dO_col_ranges>();
  ducks::art::clobber<K_ranges>();
  ducks::art::clobber<V_ranges>();
  ducks::art::clobber<P_ranges>();
  ducks::art::clobber<dP_ranges>();
  ducks::art::clobber<P_bf16_ranges>();
  ducks::art::clobber<dP_bf16_ranges>();
  ducks::art::clobber<dS_ranges>();
  ducks::art::clobber<dQ_ranges>();


#if ATTN_D == 64
  using dV_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<128, 191>>, 16>;
  using dK_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<384, 447>>, 16>;
  using Kcol_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<192, 223>>, 4>;
#else
  using dV_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<128, 255>>, 16>; // 128 registers v[128:255]
  using dK_ranges = ducks::art::split_many_t<ducks::art::type_list<ducks::art::range<384, 511>>, 16>; // 128 registers a[128:255]
  using Kcol_ranges = K_ranges;
#endif
  ducks::art::clobber<dV_ranges>();
  ducks::art::clobber<dK_ranges>();
#if ATTN_D == 64
  ducks::art::clobber<Kcol_ranges>();
#endif

  art<bf16, DOT_SLICE_QO, D, row_l, rt_16x32_s, Q_ranges> Q_i; // 16 registers
  art<bf16, DOT_SLICE_QO, D, row_l, rt_16x32_s, dO_ranges> dO_i; // 16 registers
  art<bf16, DOT_SLICE_QO, D, col_l, rt_16x32_s, Q_ranges> Q_i_col; // 16 registers
  art<bf16, DOT_SLICE_QO, D, col_l, rt_16x32_s, dO_col_ranges> dO_i_col; // 16 registers
  art<bf16, WARP_SIZE_KV, D, row_l, rt_16x32_s, K_ranges> K_j; // 64 registers
  art<bf16, WARP_SIZE_KV, D, row_l, rt_16x32_s, V_ranges> V_j; // 64 registers
  constexpr int L_i = 126;
  constexpr int delta_i = 127;
  constexpr int neg_inf_v = 29;
  // Move -inf to VGPR neg_inf_v
  kittens::macros::clobber_gpr<neg_inf_v>();
  kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000);

  art<float, DOT_SLICE_QO, WARP_SIZE_KV, col_l, rt_16x16_s, P_ranges> P_ij; // 16 registers
  art<float, DOT_SLICE_QO, WARP_SIZE_KV, col_l, rt_16x16_s, dP_ranges> dP_ij; // 16 registers
  art<bf16, DOT_SLICE_QO, WARP_SIZE_KV, col_l, rt_16x16_s, P_bf16_ranges> P_ij_bf16; // 8 registers
  art<bf16, DOT_SLICE_QO, WARP_SIZE_KV, col_l, rt_16x16_s, dP_bf16_ranges> dP_ij_bf16; // 8 registers
  art<bf16, WARP_SIZE_KV, DOT_SLICE_QO, row_l, rt_16x16_s, ducks::art::transpose_2d<dP_bf16_ranges, 1, 4>> dP_ij_bf16_accum_row; // 8 registers

  art<bf16, DOT_SLICE_QO, WARP_SIZE_KV, col_l, rt_16x32_s, P_bf16_col_ranges> P_ij_bf16_col; // 8 registers
  art<bf16, DOT_SLICE_QO, WARP_SIZE_KV, col_l, rt_16x32_s, dP_bf16_col_ranges> dP_ij_bf16_col; // 8 registers

  art<bf16, 256, ATTN_D/4, col_l, rt_32x16_4_s, Kcol_ranges> K_j_col;
  art<bf16, 256, 16, col_l, rt_32x16_4_s, dS_ranges> dP_ij_bf16_col_T; // 32 registers // for dq

  art<float, D, WARP_SIZE_KV, col_l, rt_32x32_s, dK_ranges> dK_j_T; // 128 registers
  art<float, D, WARP_SIZE_KV, col_l, rt_32x32_s, dV_ranges> dV_j_T; // 128 registers
  art<float, ATTN_D/4, 16, col_l, rt_16x16_s, dQ_ranges> dQ_i_T;
  art<float, 16, ATTN_D/4, row_l, rt_16x16_s, ducks::art::transpose_2d<dQ_ranges, ATTN_D/64, 1>> dQ_i;

  // This is used for both dK_j_T and dV_j_T
  art<float, WARP_SIZE_KV, D, row_l, rt_32x32_s, ducks::art::transpose_2d<dV_ranges, ATTN_D/32, 2>> dV_j;

  // Construct gl objects with compile-time dims AFTER clobbers so compiler knows which VGPRs are taken
  _gl_dQ  dQg{dQ_ptr, nullptr, nullptr, nullptr, nullptr};
  _gl_dKV dKg{dK_ptr, nullptr, nullptr, nullptr, nullptr};
  _gl_dKV dVg{dV_ptr, nullptr, nullptr, nullptr, nullptr};
  _gl_QdO dOg{dO_ptr, nullptr, nullptr, nullptr, nullptr};
  _gl_QdO Q{Q_ptr, nullptr, nullptr, nullptr, nullptr};
  _gl_KV  K{K_ptr, nullptr, nullptr, nullptr, nullptr};
  _gl_KV  V{V_ptr, nullptr, nullptr, nullptr, nullptr};
  _gl_Lvec L_vec_gl{L_vec_ptr, nullptr, nullptr, nullptr, nullptr};
  _gl_Lvec delta_vec_gl{delta_vec_ptr, nullptr, nullptr, nullptr, nullptr};
  attn_bwd_combined_globals<D> g{Q, K, V, dOg, dQg, dKg, dVg, L_vec_gl, delta_vec_gl};

  // Swizzled offsets for Q and dO
  constexpr int bytes_per_thread = st_16x32_s::template bytes_per_thread<bf16>();
  constexpr int bytes_per_warp = bytes_per_thread * kittens::WARP_THREADS;
  constexpr int memcpy_per_tile = BLOCK_SIZE_KV * DOT_SLICE_QO * sizeof(bf16) / (bytes_per_thread * NUM_THREADS);
  static_assert(BLOCK_SIZE_KV * DOT_SLICE_QO * sizeof(bf16) >= bytes_per_warp, "shared tile must be at least 1024 bytes");
  uint32_t swizzled_offsets_Q_dO[memcpy_per_tile];
  G::prefill_swizzled_offsets<1, false>(Q_i_smem[0][0], g.Q, swizzled_offsets_Q_dO);

  int tic = 0, toc = 1;

  // Load K_j from HBM to shared memory
  G::load<1, false>(K_j_smem, g.K, {batch_idx, seq_idx, kv_head_idx, 0});

  // Load V_j from HBM to registers
  load<1>(V_j, g.V, {batch_idx, 0, kv_head_idx, 0}, {0, j, 0, 0});

  // Load Q, dO, L, delta for this specific query head
  load(L_smem[tic], g.L_vec, {batch_idx, first_q_head, 0, first_step});
  load(delta_smem[tic], g.delta_vec, {batch_idx, first_q_head, 0, first_step});
  G::load<1, false>(Q_i_smem[tic][0],  g.Q,   {batch_idx, first_step * 2 + 0, first_q_head, 0}, swizzled_offsets_Q_dO);
  G::load<1, false>(dO_i_smem[tic][0], g.dOg, {batch_idx, first_step * 2 + 0, first_q_head, 0}, swizzled_offsets_Q_dO);
  G::load<1, false>(Q_i_smem[tic][1],  g.Q,   {batch_idx, first_step * 2 + 1, first_q_head, 0}, swizzled_offsets_Q_dO);
  G::load<1, false>(dO_i_smem[tic][1], g.dOg, {batch_idx, first_step * 2 + 1, first_q_head, 0}, swizzled_offsets_Q_dO);
  __builtin_amdgcn_s_waitcnt(0);
  __builtin_amdgcn_s_barrier();
  __builtin_amdgcn_sched_barrier(0);

  // Addresses
  const uint32_t K_j_addr = get_address(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
  // Compute K_j_col_addr
  // uint32_t K_j_col_addr = get_address(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
  const uint32_t K_j_col_addr = [&] {
    const int laneid = kittens::laneid();
    const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}).data[0]);
    const int row_offset = (laneid % 16) / 4 + (laneid / 16) * 4;
    const int col_offset = ((laneid % 4) * 4);
    const int lane_byte_offset = (row_offset * 16 + col_offset) * sizeof(bf16);
    const uint32_t addr = src_ptr + lane_byte_offset;
    return addr;
  }();

  auto attn_i_smem_subtile = subtile_inplace<WARP_SIZE_KV, DOT_SLICE_QO>(attn_i_smem, {warpid, 0});
  const uint32_t dP_ij_bf16_accum_row_addr = get_address(attn_i_smem_subtile, dP_ij_bf16_accum_row);

  uint32_t Q_i_addr;
  uint32_t dO_i_addr;
  uint32_t dO_i_col_addr;
  uint32_t Q_i_col_addr;

  // Compute dP_ij_bf16_col_T_addr
  // const uint32_t dP_ij_bf16_col_T_addr = [&] {
  //   const int laneid = kittens::laneid();
  //   const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&attn_i_smem.data[0]);
  //   const int row_offset = (laneid % 16) / 4 + (laneid / 16) * 4;
  //   const int col_offset = ((laneid % 4) * 4);
  //   const int lane_byte_offset = (row_offset * 16 + col_offset) * sizeof(bf16);
  //   const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 7) << 3);
  //   const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
  //   return addr;
  // }();
  uint32_t dP_ij_bf16_col_T_addr = get_address(dP_ij_bf16_col_T, attn_i_smem);

  if (num_steps > 1) {
    // Prologue
    {
      const int q_head_idx = first_q_head;
      const int q_seq_idx = first_step;
      const int q_pos = q_seq_idx * STEP_QO;

      const int next_q_head_idx = first_q_head;
      const int next_q_seq_idx = first_step + 1;

      // dot slice 0
      {
        load(L_smem[toc], g.L_vec, {batch_idx, next_q_head_idx, 0, next_q_seq_idx});
        G::load<1, false>(Q_i_smem[toc][0], g.Q, {batch_idx, next_q_seq_idx * 2, next_q_head_idx, 0});

        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 0));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 0));
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();

        // 10. S_ij = Q_i K_j^T * scale
        // 11. P_ij = exp2(S_ij - L_i)
        // 13. dP_ij = dO_i @ V_j^T
        // 14. dS_ij = P_ij o (dP_ij - delta_i)
        // mma_ABt(P_ij, Q_i, K_j);
        mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
        load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
        mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
        load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
        // Load dO_i from shared memory to registers
        // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 0, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
        // Load dO_i_col from shared memory to registers
        // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        // Compute dO_i_col_addr
        // dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_col_addr = [&] {
          const int laneid = kittens::laneid();
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 1, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
        // Dot slice 0
        kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
          // If the query position is less than the key position, set P_ij to -inf
          if (q_pos < k_pos) {
            mov<neg_inf_v>(P_ij);
          // If the query position is equal to the key position, we need to apply a causal mask
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 0] and set [0, 1:4] to -inf
            make_causal<0, 0, neg_inf_v>(P_ij, P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
          // window lower boundary, mirror of the causal edge
          } else if (q_pos - k_pos == WINDOW) {
            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
          } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_ABt(dP_ij, dO_i, V_j);
        mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
        sub_row<0, 2, L_i>(P_ij, P_ij);
        mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 0>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
        // Load Q_i_col from shared memory to registers
        // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        // Compute Q_i_col_addr
        // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_col_addr = [&] {
          const int laneid = kittens::laneid();  
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const int addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 1>(P_ij, P_ij);
        mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
        load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
        mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
        sub_row<0, 3, L_i>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 0>(P_ij_bf16, P_ij);
        mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
        exp2<0, 2>(P_ij, P_ij);
        mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 1>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
        if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 3>(P_ij, P_ij);
        mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
        if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 2>(P_ij_bf16, P_ij);
        copy<0, 3>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
        swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
        mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col);
        // Load K_j_col from shared memory to registers
        // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
        load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col);
        sub_row<0, 0, delta_i>(dP_ij, dP_ij);
        sub_row<0, 1, delta_i>(dP_ij, dP_ij);
        mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col);
        load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col);
        mul<0, 0>(dP_ij, dP_ij, P_ij);
        mul<0, 1>(dP_ij, dP_ij, P_ij);
        copy<0, 0>(dP_ij_bf16, dP_ij);
        copy<0, 1>(dP_ij_bf16, dP_ij);
        sub_row<0, 2, delta_i>(dP_ij, dP_ij);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col);
        load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        // 12. dV_j += P_ij^T @ dO_i
        // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)

        // Store dP_ij_bf16_accum_row to shared memory
        // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
        store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col);
        sub_row<0, 3, delta_i>(dP_ij, dP_ij);
        mul<0, 2>(dP_ij, dP_ij, P_ij);
        mul<0, 3>(dP_ij, dP_ij, P_ij);
        copy<0, 2>(dP_ij_bf16, dP_ij);
        copy<0, 3>(dP_ij_bf16, dP_ij);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col);

        // dot slice 1
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 1));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 1));
        
        store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col);
        swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
        asm volatile("s_waitcnt lgkmcnt(12)");
        // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col);
        if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();
        mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col);
        // Load dP_ij_bf16_col_T from shared memory to registers
        // load(dP_ij_bf16_col_T, attn_i_smem);
        load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col);
        load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col);
        load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col);
        asm volatile("s_waitcnt vmcnt(0) lgkmcnt(" DQ_L6 ")");
        __builtin_amdgcn_s_barrier();
        // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
        // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        G::load<1, false>(dO_i_smem[toc][0], g.dOg, {batch_idx, next_q_seq_idx * 2, next_q_head_idx, 0}, swizzled_offsets_Q_dO);
        mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load(delta_smem[toc], g.delta_vec, {batch_idx, next_q_head_idx, 0, next_q_seq_idx});
        mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {1, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
        __builtin_amdgcn_s_barrier();
        mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
        if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
      }

      // dot slice 1
      {
        // 10. S_ij = Q_i K_j^T * scale
        // 11. P_ij = exp2(S_ij - L_i)
        // 13. dP_ij = dO_i @ V_j^T
        // 14. dS_ij = P_ij o (dP_ij - delta_i)
        // mma_ABt(P_ij, Q_i, K_j);
        mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
        load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
        load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
        // Load dO_i from shared memory to registers
        // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {1, 0}));
        load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 0, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
        // Load dO_i_col from shared memory to registers
        // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        // Compute dO_i_col_addr
        // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_col_addr = [&] {
          const int laneid = kittens::laneid();
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {1, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 1, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
        // Dot slice 1
        kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
          // If the query position is less than the key position, set P_ij to -inf
          if (q_pos < k_pos) {
            mov<neg_inf_v>(P_ij);
          // If the query position is equal to the key position, we need to apply a causal mask
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 1] and set [0, 2:4] to -inf
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
          } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
          } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_ABt(dP_ij, dO_i, V_j);
        mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
        sub_row<0, 2, L_i>(P_ij, P_ij);
        mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 0>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
        // Load Q_i_col from shared memory to registers
        // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        // Compute Q_i_col_addr
        // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_col_addr = [&] {
          const int laneid = kittens::laneid();  
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {1, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const int addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 1>(P_ij, P_ij);
        mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
        load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
        mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
        sub_row<0, 3, L_i>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 0>(P_ij_bf16, P_ij);
        mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
        exp2<0, 2>(P_ij, P_ij);
        mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 1>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
        if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 3>(P_ij, P_ij);
        mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
        if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 2>(P_ij_bf16, P_ij);
        copy<0, 3>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
        swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
        mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        // Load K_j_col from shared memory to registers
        // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
        load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 0, delta_i>(dP_ij, dP_ij);
        sub_row<0, 1, delta_i>(dP_ij, dP_ij);
        mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        mul<0, 0>(dP_ij, dP_ij, P_ij);
        mul<0, 1>(dP_ij, dP_ij, P_ij);
        copy<0, 0>(dP_ij_bf16, dP_ij);
        copy<0, 1>(dP_ij_bf16, dP_ij);
        sub_row<0, 2, delta_i>(dP_ij, dP_ij);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        // 12. dV_j += P_ij^T @ dO_i
        // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
        // Store dP_ij_bf16_accum_row to shared memory
        // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
        store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 3, delta_i>(dP_ij, dP_ij);
        mul<0, 2>(dP_ij, dP_ij, P_ij);
        mul<0, 3>(dP_ij, dP_ij, P_ij);
        copy<0, 2>(dP_ij_bf16, dP_ij);
        copy<0, 3>(dP_ij_bf16, dP_ij);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

        // dot slice 2
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 2));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 2));
        
        store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
        asm volatile("s_waitcnt lgkmcnt(12)");
        // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();
        mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        // Load dP_ij_bf16_col_T from shared memory to registers
        // load(dP_ij_bf16_col_T, attn_i_smem);
        load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4, 0}, warpid);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4, 0}, warpid);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt vmcnt(4) lgkmcnt(" DQ_L6_LOOP ")");
        __builtin_amdgcn_s_barrier();
        // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
        // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        G::load<1, false>(Q_i_smem[toc][1], g.Q, {batch_idx, next_q_seq_idx * 2 + 1, next_q_head_idx, 0}, swizzled_offsets_Q_dO);
        mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {0, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
        __builtin_amdgcn_s_barrier();
        mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
        if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
      }

      // dot slice 2
      {
        // 10. S_ij = Q_i K_j^T * scale
        // 11. P_ij = exp2(S_ij - L_i)
        // 13. dP_ij = dO_i @ V_j^T
        // 14. dS_ij = P_ij o (dP_ij - delta_i)
        // mma_ABt(P_ij, Q_i, K_j);
        mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
        load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
        load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
        // Load dO_i from shared memory to registers
        // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {0, 0}));
        load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 0, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
        // Load dO_i_col from shared memory to registers
        // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        // Compute dO_i_col_addr
        // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_col_addr = [&] {
          const int laneid = kittens::laneid();
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {0, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 1, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
        // Dot slice 2
        kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
          // If the query position is less than the key position, set P_ij to -inf
          if (q_pos < k_pos) {
            mov<neg_inf_v>(P_ij);
          // If the query position is equal to the key position, we need to apply a causal mask
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
          } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
          } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_ABt(dP_ij, dO_i, V_j);
        mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
        sub_row<0, 2, L_i>(P_ij, P_ij);
        mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 0>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
        // Load Q_i_col from shared memory to registers
        // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        // Compute Q_i_col_addr
        // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_col_addr = [&] {
          const int laneid = kittens::laneid();  
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {0, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const int addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 1>(P_ij, P_ij);
        mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
        load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
        mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
        sub_row<0, 3, L_i>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 0>(P_ij_bf16, P_ij);
        mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
        exp2<0, 2>(P_ij, P_ij);
        mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 1>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
        if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 3>(P_ij, P_ij);
        mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
        if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 2>(P_ij_bf16, P_ij);
        copy<0, 3>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
        swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
        mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        // Load K_j_col from shared memory to registers
        // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
        load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 0, delta_i>(dP_ij, dP_ij);
        sub_row<0, 1, delta_i>(dP_ij, dP_ij);
        mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        mul<0, 0>(dP_ij, dP_ij, P_ij);
        mul<0, 1>(dP_ij, dP_ij, P_ij);
        copy<0, 0>(dP_ij_bf16, dP_ij);
        copy<0, 1>(dP_ij_bf16, dP_ij);
        sub_row<0, 2, delta_i>(dP_ij, dP_ij);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        // 12. dV_j += P_ij^T @ dO_i
        // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
        // Store dP_ij_bf16_accum_row to shared memory
        // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
        store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 3, delta_i>(dP_ij, dP_ij);
        mul<0, 2>(dP_ij, dP_ij, P_ij);
        mul<0, 3>(dP_ij, dP_ij, P_ij);
        copy<0, 2>(dP_ij_bf16, dP_ij);
        copy<0, 3>(dP_ij_bf16, dP_ij);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

        // dot slice 3
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 3));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 3));
        
        store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
        asm volatile("s_waitcnt lgkmcnt(12)");
        // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();
        mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        // Load dP_ij_bf16_col_T from shared memory to registers
        // load(dP_ij_bf16_col_T, attn_i_smem);
        load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 1, 0}, warpid);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 1, 0}, warpid);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt vmcnt(4) lgkmcnt(" DQ_L6_LOOP ")");
        __builtin_amdgcn_s_barrier();
        // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
        // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        G::load<1, false>(dO_i_smem[toc][1], g.dOg, {batch_idx, next_q_seq_idx * 2 + 1, next_q_head_idx, 0}, swizzled_offsets_Q_dO);
        mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {1, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
        __builtin_amdgcn_s_barrier();
        mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
        if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
      }

      // dot slice 3
      {
        // 10. S_ij = Q_i K_j^T * scale
        // 11. P_ij = exp2(S_ij - L_i)
        // 13. dP_ij = dO_i @ V_j^T
        // 14. dS_ij = P_ij o (dP_ij - delta_i)
        // mma_ABt(P_ij, Q_i, K_j);
        mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
        load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
        load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
        // Load dO_i from shared memory to registers
        // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {1, 0}));
        load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 0, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
        // Load dO_i_col from shared memory to registers
        // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        // Compute dO_i_col_addr
        // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_col_addr = [&] {
          const int laneid = kittens::laneid();
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {1, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 1, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
        // Dot slice 3
        kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
          // If the query position is less than the key position, set P_ij to -inf
          if (q_pos < k_pos) {
            mov<neg_inf_v>(P_ij);
          // If the query position is equal to the key position, we need to apply a causal mask
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
#if WINDOW
          } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
          } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_ABt(dP_ij, dO_i, V_j);
        mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
        sub_row<0, 2, L_i>(P_ij, P_ij);
        mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 0>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
        // Load Q_i_col from shared memory to registers
        // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        // Compute Q_i_col_addr
        // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_col_addr = [&] {
          const int laneid = kittens::laneid();  
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {1, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const int addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 1>(P_ij, P_ij);
        mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
        load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
        mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
        sub_row<0, 3, L_i>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 0>(P_ij_bf16, P_ij);
        mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
        exp2<0, 2>(P_ij, P_ij);
        mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 1>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
        if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 3>(P_ij, P_ij);
        mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
        if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 2>(P_ij_bf16, P_ij);
        copy<0, 3>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
        swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
        mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        // Load K_j_col from shared memory to registers
        // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
        load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 0, delta_i>(dP_ij, dP_ij);
        sub_row<0, 1, delta_i>(dP_ij, dP_ij);
        mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        mul<0, 0>(dP_ij, dP_ij, P_ij);
        mul<0, 1>(dP_ij, dP_ij, P_ij);
        copy<0, 0>(dP_ij_bf16, dP_ij);
        copy<0, 1>(dP_ij_bf16, dP_ij);
        sub_row<0, 2, delta_i>(dP_ij, dP_ij);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        // 12. dV_j += P_ij^T @ dO_i
        // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
        // Store dP_ij_bf16_accum_row to shared memory
        // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
        store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 3, delta_i>(dP_ij, dP_ij);
        mul<0, 2>(dP_ij, dP_ij, P_ij);
        mul<0, 3>(dP_ij, dP_ij, P_ij);
        copy<0, 2>(dP_ij_bf16, dP_ij);
        copy<0, 3>(dP_ij_bf16, dP_ij);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

        // dot slice 0 - next iteration
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[toc], 0));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[toc], 0));
        
        store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
        asm volatile("s_waitcnt lgkmcnt(12)");
        // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();
        mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        // Load dP_ij_bf16_col_T from shared memory to registers
        // load(dP_ij_bf16_col_T, attn_i_smem);
        load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 2, 0}, warpid);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 2, 0}, warpid);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt vmcnt(4) lgkmcnt(" DQ_L6_LOOP ")");
        __builtin_amdgcn_s_barrier();
        // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
        // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
        __builtin_amdgcn_s_barrier();
        mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
        if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
      }
      tic ^= 1; toc ^= 1;
    }

    // 9. for 1 <= i <= T_r (1024 / 32 = 32)  
    for (int i = 1; i < num_steps - 1; ++i, tic ^= 1, toc ^= 1) {
      const int last_head_offset = (i - 1) >= num_steps_per_head;
      const int last_q_head_idx = last_head_offset + first_q_head;
      const int last_q_seq_idx = i - 1 - last_head_offset * num_steps_per_head + first_step;

      const int head_offset = i >= num_steps_per_head;
      const int q_head_idx = head_offset + first_q_head;
      const int q_seq_idx = i - head_offset * num_steps_per_head + first_step;
      const int q_pos = q_seq_idx * STEP_QO;

      const int next_head_offset = (i + 1) >= num_steps_per_head;
      const int next_q_head_idx = next_head_offset + first_q_head;
      const int next_q_seq_idx = i + 1 - next_head_offset * num_steps_per_head + first_step;

      // dot slice 0
      {
        // 10. S_ij = Q_i K_j^T * scale
        // 11. P_ij = exp2(S_ij - L_i)
        // 13. dP_ij = dO_i @ V_j^T
        // 14. dS_ij = P_ij o (dP_ij - delta_i)
        // mma_ABt(P_ij, Q_i, K_j);
        G::load<1, false>(Q_i_smem[toc][0], g.Q, {batch_idx, next_q_seq_idx * 2, next_q_head_idx, 0}, swizzled_offsets_Q_dO);
        mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
        load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
        load(L_smem[toc], g.L_vec, {batch_idx, next_q_head_idx, 0, next_q_seq_idx});
        if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
        load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
        // Load dO_i from shared memory to registers
        // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 0, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
        // Load dO_i_col from shared memory to registers
        // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        // Compute dO_i_col_addr
        // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_col_addr = [&] {
          const int laneid = kittens::laneid();
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 1, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
        // Dot slice 0
        kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
          // If the query position is less than the key position, set P_ij to -inf
          if (q_pos < k_pos) {
            mov<neg_inf_v>(P_ij);
          // If the query position is equal to the key position, we need to apply a causal mask
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 0] and set [0, 1:4] to -inf
            make_causal<0, 0, neg_inf_v>(P_ij, P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
          // window lower boundary, mirror of the causal edge
          } else if (q_pos - k_pos == WINDOW) {
            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
          } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_ABt(dP_ij, dO_i, V_j);
        mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
        sub_row<0, 2, L_i>(P_ij, P_ij);
        mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 0>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
        // Load Q_i_col from shared memory to registers
        // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        // Compute Q_i_col_addr
        // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_col_addr = [&] {
          const int laneid = kittens::laneid();  
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const int addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 1>(P_ij, P_ij);
        mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
        load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
        mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
        sub_row<0, 3, L_i>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 0>(P_ij_bf16, P_ij);
        mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
        exp2<0, 2>(P_ij, P_ij);
        mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 1>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
        if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 3>(P_ij, P_ij);
        mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
        if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 2>(P_ij_bf16, P_ij);
        copy<0, 3>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
        swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
        mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        // Load K_j_col from shared memory to registers
        // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
        load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 0, delta_i>(dP_ij, dP_ij);
        sub_row<0, 1, delta_i>(dP_ij, dP_ij);
        mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        mul<0, 0>(dP_ij, dP_ij, P_ij);
        mul<0, 1>(dP_ij, dP_ij, P_ij);
        copy<0, 0>(dP_ij_bf16, dP_ij);
        copy<0, 1>(dP_ij_bf16, dP_ij);
        sub_row<0, 2, delta_i>(dP_ij, dP_ij);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        // 12. dV_j += P_ij^T @ dO_i
        // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
        // Store dP_ij_bf16_accum_row to shared memory
        // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
        store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 3, delta_i>(dP_ij, dP_ij);
        mul<0, 2>(dP_ij, dP_ij, P_ij);
        mul<0, 3>(dP_ij, dP_ij, P_ij);
        copy<0, 2>(dP_ij_bf16, dP_ij);
        copy<0, 3>(dP_ij_bf16, dP_ij);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

        // dot slice 1
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 1));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 1));
        
        store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
        asm volatile("s_waitcnt lgkmcnt(12)");
        // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();
        mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        // Load dP_ij_bf16_col_T from shared memory to registers
        // load(dP_ij_bf16_col_T, attn_i_smem);
        load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, last_q_head_idx, last_q_seq_idx * 4 + 3, 0}, warpid);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, last_q_head_idx, last_q_seq_idx * 4 + 3, 0}, warpid);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt vmcnt(4) lgkmcnt(" DQ_L6_LOOP ")");
        __builtin_amdgcn_s_barrier();
        // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
        // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        G::load<1, false>(dO_i_smem[toc][0], g.dOg, {batch_idx, next_q_seq_idx * 2, next_q_head_idx, 0}, swizzled_offsets_Q_dO);
        mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load(delta_smem[toc], g.delta_vec, {batch_idx, next_q_head_idx, 0, next_q_seq_idx});
        mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {1, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
        __builtin_amdgcn_s_barrier();
        mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
        if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);

        asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
      }

      // dot slice 1
      {
        // 10. S_ij = Q_i K_j^T * scale
        // 11. P_ij = exp2(S_ij - L_i)
        // 13. dP_ij = dO_i @ V_j^T
        // 14. dS_ij = P_ij o (dP_ij - delta_i)
        // mma_ABt(P_ij, Q_i, K_j);
        mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
        load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
        load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
        // Load dO_i from shared memory to registers
        // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {1, 0}));
        load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 0, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
        // Load dO_i_col from shared memory to registers
        // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        // Compute dO_i_col_addr
        // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_col_addr = [&] {
          const int laneid = kittens::laneid();
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {1, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 1, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
        // Dot slice 1
        kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
          // If the query position is less than the key position, set P_ij to -inf
          if (q_pos < k_pos) {
            mov<neg_inf_v>(P_ij);
          // If the query position is equal to the key position, we need to apply a causal mask
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 1] and set [0, 2:4] to -inf
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
          } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
          } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_ABt(dP_ij, dO_i, V_j);
        mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
        sub_row<0, 2, L_i>(P_ij, P_ij);
        mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 0>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
        // Load Q_i_col from shared memory to registers
        // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        // Compute Q_i_col_addr
        // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_col_addr = [&] {
          const int laneid = kittens::laneid();  
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {1, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const int addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 1>(P_ij, P_ij);
        mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
        load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
        mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
        sub_row<0, 3, L_i>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 0>(P_ij_bf16, P_ij);
        mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
        exp2<0, 2>(P_ij, P_ij);
        mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 1>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
        if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 3>(P_ij, P_ij);
        mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
        if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 2>(P_ij_bf16, P_ij);
        copy<0, 3>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
        swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
        mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        // Load K_j_col from shared memory to registers
        // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
        load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 0, delta_i>(dP_ij, dP_ij);
        sub_row<0, 1, delta_i>(dP_ij, dP_ij);
        mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        mul<0, 0>(dP_ij, dP_ij, P_ij);
        mul<0, 1>(dP_ij, dP_ij, P_ij);
        copy<0, 0>(dP_ij_bf16, dP_ij);
        copy<0, 1>(dP_ij_bf16, dP_ij);
        sub_row<0, 2, delta_i>(dP_ij, dP_ij);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        // 12. dV_j += P_ij^T @ dO_i
        // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
        // Store dP_ij_bf16_accum_row to shared memory
        // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
        store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 3, delta_i>(dP_ij, dP_ij);
        mul<0, 2>(dP_ij, dP_ij, P_ij);
        mul<0, 3>(dP_ij, dP_ij, P_ij);
        copy<0, 2>(dP_ij_bf16, dP_ij);
        copy<0, 3>(dP_ij_bf16, dP_ij);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

        // dot slice 2
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 2));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 2));
        
        store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
        asm volatile("s_waitcnt lgkmcnt(12)");
        // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();
        mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        // Load dP_ij_bf16_col_T from shared memory to registers
        // load(dP_ij_bf16_col_T, attn_i_smem);
        load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 0, 0}, warpid);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 0, 0}, warpid);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt vmcnt(4) lgkmcnt(" DQ_L6_LOOP ")");
        __builtin_amdgcn_s_barrier();
        // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
        // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        G::load<1, false>(Q_i_smem[toc][1], g.Q, {batch_idx, next_q_seq_idx * 2 + 1, next_q_head_idx, 0}, swizzled_offsets_Q_dO);
        mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {0, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
        __builtin_amdgcn_s_barrier();
        mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
        if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
      }

      // dot slice 2
      {
        // 10. S_ij = Q_i K_j^T * scale
        // 11. P_ij = exp2(S_ij - L_i)
        // 13. dP_ij = dO_i @ V_j^T
        // 14. dS_ij = P_ij o (dP_ij - delta_i)
        // mma_ABt(P_ij, Q_i, K_j);
        mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
        load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
        load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
        // Load dO_i from shared memory to registers
        // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {0, 0}));
        load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 0, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
        // Load dO_i_col from shared memory to registers
        // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        // Compute dO_i_col_addr
        // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_col_addr = [&] {
          const int laneid = kittens::laneid();
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {0, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 1, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
        // Dot slice 2
        kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
          // If the query position is less than the key position, set P_ij to -inf
          if (q_pos < k_pos) {
            mov<neg_inf_v>(P_ij);
          // If the query position is equal to the key position, we need to apply a causal mask
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
          } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
          } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_ABt(dP_ij, dO_i, V_j);
        mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
        sub_row<0, 2, L_i>(P_ij, P_ij);
        mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 0>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
        // Load Q_i_col from shared memory to registers
        // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        // Compute Q_i_col_addr
        // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_col_addr = [&] {
          const int laneid = kittens::laneid();  
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {0, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const int addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 1>(P_ij, P_ij);
        mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
        load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
        mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
        sub_row<0, 3, L_i>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 0>(P_ij_bf16, P_ij);
        mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
        exp2<0, 2>(P_ij, P_ij);
        mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 1>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
        if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 3>(P_ij, P_ij);
        mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
        if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 2>(P_ij_bf16, P_ij);
        copy<0, 3>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
        swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
        mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        // Load K_j_col from shared memory to registers
        // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
        load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 0, delta_i>(dP_ij, dP_ij);
        sub_row<0, 1, delta_i>(dP_ij, dP_ij);
        mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        mul<0, 0>(dP_ij, dP_ij, P_ij);
        mul<0, 1>(dP_ij, dP_ij, P_ij);
        copy<0, 0>(dP_ij_bf16, dP_ij);
        copy<0, 1>(dP_ij_bf16, dP_ij);
        sub_row<0, 2, delta_i>(dP_ij, dP_ij);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        // 12. dV_j += P_ij^T @ dO_i
        // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
        // Store dP_ij_bf16_accum_row to shared memory
        // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
        store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 3, delta_i>(dP_ij, dP_ij);
        mul<0, 2>(dP_ij, dP_ij, P_ij);
        mul<0, 3>(dP_ij, dP_ij, P_ij);
        copy<0, 2>(dP_ij_bf16, dP_ij);
        copy<0, 3>(dP_ij_bf16, dP_ij);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

        // dot slice 3
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 3));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 3));
        
        store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
        asm volatile("s_waitcnt lgkmcnt(12)");
        // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();
        mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        // Load dP_ij_bf16_col_T from shared memory to registers
        // load(dP_ij_bf16_col_T, attn_i_smem);
        load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 1, 0}, warpid);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 1, 0}, warpid);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt vmcnt(4) lgkmcnt(" DQ_L6_LOOP ")");
        __builtin_amdgcn_s_barrier();
        // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
        // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        G::load<1, false>(dO_i_smem[toc][1], g.dOg, {batch_idx, next_q_seq_idx * 2 + 1, next_q_head_idx, 0}, swizzled_offsets_Q_dO);
        mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {1, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
        __builtin_amdgcn_s_barrier();
        mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
        if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
      }

      // dot slice 3
      {
        // 10. S_ij = Q_i K_j^T * scale
        // 11. P_ij = exp2(S_ij - L_i)
        // 13. dP_ij = dO_i @ V_j^T
        // 14. dS_ij = P_ij o (dP_ij - delta_i)
        // mma_ABt(P_ij, Q_i, K_j);
        mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
        load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
        load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
        // Load dO_i from shared memory to registers
        // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {1, 0}));
        load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 0, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(6)");
        if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
        mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(6)");
        mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
        // Load dO_i_col from shared memory to registers
        // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        // Compute dO_i_col_addr
        // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
        dO_i_col_addr = [&] {
          const int laneid = kittens::laneid();
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {1, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
        sub_row<0, 1, L_i>(P_ij, P_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
        if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
        // Dot slice 3
        kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
          // If the query position is less than the key position, set P_ij to -inf
          if (q_pos < k_pos) {
            mov<neg_inf_v>(P_ij);
          // If the query position is equal to the key position, we need to apply a causal mask
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
#if WINDOW
          } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
          } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_ABt(dP_ij, dO_i, V_j);
        mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
        sub_row<0, 2, L_i>(P_ij, P_ij);
        mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 0>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
        // Load Q_i_col from shared memory to registers
        // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        // Compute Q_i_col_addr
        // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
        Q_i_col_addr = [&] {
          const int laneid = kittens::laneid();  
          const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {1, 0}).data[0]);
          const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
          const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
          const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
          const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
          const int addr = src_ptr + swizzled_lane_byte_offset;
          return addr;
        }();
        load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 1>(P_ij, P_ij);
        mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
        load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
        mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
        if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
        sub_row<0, 3, L_i>(P_ij, P_ij);
        if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 0>(P_ij_bf16, P_ij);
        mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
        exp2<0, 2>(P_ij, P_ij);
        mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 1>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
        if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
        exp2<0, 3>(P_ij, P_ij);
        mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
        if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
        mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
        copy<0, 2>(P_ij_bf16, P_ij);
        copy<0, 3>(P_ij_bf16, P_ij);
        if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
        swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
        if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
        asm volatile("s_waitcnt lgkmcnt(8)");
        // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
        mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        // Load K_j_col from shared memory to registers
        // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
        load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 0, delta_i>(dP_ij, dP_ij);
        sub_row<0, 1, delta_i>(dP_ij, dP_ij);
        mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        mul<0, 0>(dP_ij, dP_ij, P_ij);
        mul<0, 1>(dP_ij, dP_ij, P_ij);
        copy<0, 0>(dP_ij_bf16, dP_ij);
        copy<0, 1>(dP_ij_bf16, dP_ij);
        sub_row<0, 2, delta_i>(dP_ij, dP_ij);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        // 12. dV_j += P_ij^T @ dO_i
        // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
        // Store dP_ij_bf16_accum_row to shared memory
        // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
        store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        sub_row<0, 3, delta_i>(dP_ij, dP_ij);
        mul<0, 2>(dP_ij, dP_ij, P_ij);
        mul<0, 3>(dP_ij, dP_ij, P_ij);
        copy<0, 2>(dP_ij_bf16, dP_ij);
        copy<0, 3>(dP_ij_bf16, dP_ij);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

        // dot slice 0 - next iteration
        load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[toc], 0));
        load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[toc], 0));
        
        store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
        swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
        asm volatile("s_waitcnt lgkmcnt(12)");
        // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
        mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt lgkmcnt(8)");
        if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
        __builtin_amdgcn_s_barrier();
        mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        // Load dP_ij_bf16_col_T from shared memory to registers
        // load(dP_ij_bf16_col_T, attn_i_smem);
        load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        mul<L_i, L_i>(L_SCALE_FACTOR);
        atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 2, 0}, warpid);
        if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 2, 0}, warpid);
        if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
        if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
        asm volatile("s_waitcnt vmcnt(4) lgkmcnt(" DQ_L6_LOOP ")");
        __builtin_amdgcn_s_barrier();
        // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
        // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
        mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // Load Q_i from shared memory to registers
        // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}));
        Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}));
        load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
        __builtin_amdgcn_s_barrier();
        mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
        load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
        if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
        if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
        asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
      }
    }
  }

  const int last_q_head_idx = first_q_head + HEADS_PER_WG - 1;
  const int last_q_seq_idx = first_step + num_steps_per_head - 2;

  const int q_head_idx = first_q_head + HEADS_PER_WG - 1;
  const int q_seq_idx = first_step + num_steps_per_head - 1;
  const int q_pos = q_seq_idx * STEP_QO;
  // Epilogue
  {
    // dot slice 0
    {

      // 10. S_ij = Q_i K_j^T * scale
      // 11. P_ij = exp2(S_ij - L_i)
      // 13. dP_ij = dO_i @ V_j^T
      // 14. dS_ij = P_ij o (dP_ij - delta_i)
      // mma_ABt(P_ij, Q_i, K_j);
      mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
      load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
      load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
      mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(6)");
      mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
      // Load dO_i from shared memory to registers
      // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
      sub_row<0, 0, L_i>(P_ij, P_ij);
      asm volatile("s_waitcnt lgkmcnt(6)");
      if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
      mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(6)");
      mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
      // Load dO_i_col from shared memory to registers
      // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      // Compute dO_i_col_addr
      // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      dO_i_col_addr = [&] {
        const int laneid = kittens::laneid();
        const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}).data[0]);
        const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
        const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
        const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
        const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
        const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
        return addr;
      }();
      load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
      sub_row<0, 1, L_i>(P_ij, P_ij);
      asm volatile("s_waitcnt lgkmcnt(8)");
      if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
      // Dot slice 0
      kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
        // If the query position is less than the key position, set P_ij to -inf
        if (q_pos < k_pos) {
          mov<neg_inf_v>(P_ij);
        // If the query position is equal to the key position, we need to apply a causal mask
        } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 0] and set [0, 1:4] to -inf
            make_causal<0, 0, neg_inf_v>(P_ij, P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
        } else if (q_pos - k_pos == WINDOW) {
            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
        } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(8)");
      // mma_ABt(dP_ij, dO_i, V_j);
      mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
      sub_row<0, 2, L_i>(P_ij, P_ij);
      mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 0>(P_ij, P_ij);
      if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
      // Load Q_i_col from shared memory to registers
      // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      // Compute Q_i_col_addr
      // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      Q_i_col_addr = [&] {
        const int laneid = kittens::laneid();  
        const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}).data[0]);
        const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
        const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
        const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
        const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
        const int addr = src_ptr + swizzled_lane_byte_offset;
        return addr;
      }();
      load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 1>(P_ij, P_ij);
      mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
      load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
      mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
      if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
      sub_row<0, 3, L_i>(P_ij, P_ij);
      if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 0>(P_ij_bf16, P_ij);
      mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
      exp2<0, 2>(P_ij, P_ij);
      mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 1>(P_ij_bf16, P_ij);
      if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
      if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 3>(P_ij, P_ij);
      mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
      if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 2>(P_ij_bf16, P_ij);
      copy<0, 3>(P_ij_bf16, P_ij);
      if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
      swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
      if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
      asm volatile("s_waitcnt lgkmcnt(8)");
      // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
      mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      // Load K_j_col from shared memory to registers
      // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
      load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      sub_row<0, 0, delta_i>(dP_ij, dP_ij);
      sub_row<0, 1, delta_i>(dP_ij, dP_ij);
      mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      mul<0, 0>(dP_ij, dP_ij, P_ij);
      mul<0, 1>(dP_ij, dP_ij, P_ij);
      copy<0, 0>(dP_ij_bf16, dP_ij);
      copy<0, 1>(dP_ij_bf16, dP_ij);
      sub_row<0, 2, delta_i>(dP_ij, dP_ij);
      if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      // 12. dV_j += P_ij^T @ dO_i
      // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
      // Store dP_ij_bf16_accum_row to shared memory
      // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
      store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      sub_row<0, 3, delta_i>(dP_ij, dP_ij);
      mul<0, 2>(dP_ij, dP_ij, P_ij);
      mul<0, 3>(dP_ij, dP_ij, P_ij);
      copy<0, 2>(dP_ij_bf16, dP_ij);
      copy<0, 3>(dP_ij_bf16, dP_ij);
      if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

      // dot slice 1
      load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 1));
      load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 1));
      
      store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
      asm volatile("s_waitcnt lgkmcnt(12)");
      // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
      mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      asm volatile("s_waitcnt lgkmcnt(8)");
      if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
      __builtin_amdgcn_s_barrier();
      mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      // Load dP_ij_bf16_col_T from shared memory to registers
      // load(dP_ij_bf16_col_T, attn_i_smem);
      load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      mul<L_i, L_i>(L_SCALE_FACTOR);
      if (num_steps > 1) {
        atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, last_q_head_idx, last_q_seq_idx * 4 + 3, 0}, warpid);
      }
      if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      if (num_steps > 1) {
        if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, last_q_head_idx, last_q_seq_idx * 4 + 3, 0}, warpid);
      }
      if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L6 ")");
      __builtin_amdgcn_s_barrier();
      // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
      // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      // Load Q_i from shared memory to registers
      // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {1, 0}));
      load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
      __builtin_amdgcn_s_barrier();
      mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      // Load K_j from shared memory to registers
      // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
      load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
      if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
    }

    // dot slice 1
    {
      // 10. S_ij = Q_i K_j^T * scale
      // 11. P_ij = exp2(S_ij - L_i)
      // 13. dP_ij = dO_i @ V_j^T
      // 14. dS_ij = P_ij o (dP_ij - delta_i)
      // mma_ABt(P_ij, Q_i, K_j);
      mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
      load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
      load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
      mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(6)");
      mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
      // Load dO_i from shared memory to registers
      // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {1, 0}));
      load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
      sub_row<0, 0, L_i>(P_ij, P_ij);
      asm volatile("s_waitcnt lgkmcnt(6)");
      if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
      mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(6)");
      mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
      // Load dO_i_col from shared memory to registers
      // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      // Compute dO_i_col_addr
      // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      dO_i_col_addr = [&] {
        const int laneid = kittens::laneid();
        const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {1, 0}).data[0]);
        const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
        const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
        const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
        const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
        const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
        return addr;
      }();
      load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
      sub_row<0, 1, L_i>(P_ij, P_ij);
      asm volatile("s_waitcnt lgkmcnt(8)");
      if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
      // Dot slice 1
      kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
        // If the query position is less than the key position, set P_ij to -inf
        if (q_pos < k_pos) {
          mov<neg_inf_v>(P_ij);
        // If the query position is equal to the key position, we need to apply a causal mask
        } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 1] and set [0, 2:4] to -inf
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
        } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
        } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(8)");
      // mma_ABt(dP_ij, dO_i, V_j);
      mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
      sub_row<0, 2, L_i>(P_ij, P_ij);
      mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 0>(P_ij, P_ij);
      if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
      // Load Q_i_col from shared memory to registers
      // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      // Compute Q_i_col_addr
      // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      Q_i_col_addr = [&] {
        const int laneid = kittens::laneid();  
        const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {1, 0}).data[0]);
        const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
        const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
        const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
        const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
        const int addr = src_ptr + swizzled_lane_byte_offset;
        return addr;
      }();
      load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 1>(P_ij, P_ij);
      mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
      load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
      mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
      if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
      sub_row<0, 3, L_i>(P_ij, P_ij);
      if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 0>(P_ij_bf16, P_ij);
      mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
      exp2<0, 2>(P_ij, P_ij);
      mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 1>(P_ij_bf16, P_ij);
      if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
      if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 3>(P_ij, P_ij);
      mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
      if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 2>(P_ij_bf16, P_ij);
      copy<0, 3>(P_ij_bf16, P_ij);
      if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
      swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
      if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
      asm volatile("s_waitcnt lgkmcnt(8)");
      // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
      mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      // Load K_j_col from shared memory to registers
      // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
      load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      sub_row<0, 0, delta_i>(dP_ij, dP_ij);
      sub_row<0, 1, delta_i>(dP_ij, dP_ij);
      mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      mul<0, 0>(dP_ij, dP_ij, P_ij);
      mul<0, 1>(dP_ij, dP_ij, P_ij);
      copy<0, 0>(dP_ij_bf16, dP_ij);
      copy<0, 1>(dP_ij_bf16, dP_ij);
      sub_row<0, 2, delta_i>(dP_ij, dP_ij);
      if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      // 12. dV_j += P_ij^T @ dO_i
      // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
      // Store dP_ij_bf16_accum_row to shared memory
      // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
      store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      sub_row<0, 3, delta_i>(dP_ij, dP_ij);
      mul<0, 2>(dP_ij, dP_ij, P_ij);
      mul<0, 3>(dP_ij, dP_ij, P_ij);
      copy<0, 2>(dP_ij_bf16, dP_ij);
      copy<0, 3>(dP_ij_bf16, dP_ij);
      if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

      // dot slice 2
      load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 2));
      load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 2));
      
      store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
      asm volatile("s_waitcnt lgkmcnt(12)");
      // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
      mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      asm volatile("s_waitcnt lgkmcnt(8)");
      if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
      __builtin_amdgcn_s_barrier();
      mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      // Load dP_ij_bf16_col_T from shared memory to registers
      // load(dP_ij_bf16_col_T, attn_i_smem);
      load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      mul<L_i, L_i>(L_SCALE_FACTOR);
      atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4, 0}, warpid);
      if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4, 0}, warpid);
      if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L6 ")");
      __builtin_amdgcn_s_barrier();
      // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
      // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      // Load Q_i from shared memory to registers
      // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {0, 0}));
      load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
      __builtin_amdgcn_s_barrier();
      mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      // Load K_j from shared memory to registers
      // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
      load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
      if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
    }

    // dot slice 2
    {
      // 10. S_ij = Q_i K_j^T * scale
      // 11. P_ij = exp2(S_ij - L_i)
      // 13. dP_ij = dO_i @ V_j^T
      // 14. dS_ij = P_ij o (dP_ij - delta_i)
      // mma_ABt(P_ij, Q_i, K_j);
      mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
      load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
      load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
      mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(6)");
      mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
      // Load dO_i from shared memory to registers
      // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {0, 0}));
      load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
      sub_row<0, 0, L_i>(P_ij, P_ij);
      asm volatile("s_waitcnt lgkmcnt(6)");
      if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
      mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(6)");
      mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
      // Load dO_i_col from shared memory to registers
      // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      // Compute dO_i_col_addr
      // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      dO_i_col_addr = [&] {
        const int laneid = kittens::laneid();
        const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {0, 0}).data[0]);
        const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
        const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
        const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
        const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
        const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
        return addr;
      }();
      load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
      sub_row<0, 1, L_i>(P_ij, P_ij);
      asm volatile("s_waitcnt lgkmcnt(8)");
      if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
      // Dot slice 2
      kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
        // If the query position is less than the key position, set P_ij to -inf
        if (q_pos < k_pos) {
          mov<neg_inf_v>(P_ij);
        // If the query position is equal to the key position, we need to apply a causal mask
        } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
#if WINDOW
        } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
        } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(8)");
      // mma_ABt(dP_ij, dO_i, V_j);
      mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
      sub_row<0, 2, L_i>(P_ij, P_ij);
      mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 0>(P_ij, P_ij);
      if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
      // Load Q_i_col from shared memory to registers
      // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      // Compute Q_i_col_addr
      // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      Q_i_col_addr = [&] {
        const int laneid = kittens::laneid();  
        const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {0, 0}).data[0]);
        const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
        const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
        const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
        const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
        const int addr = src_ptr + swizzled_lane_byte_offset;
        return addr;
      }();
      load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 1>(P_ij, P_ij);
      mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
      load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
      mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
      if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
      sub_row<0, 3, L_i>(P_ij, P_ij);
      if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 0>(P_ij_bf16, P_ij);
      mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
      exp2<0, 2>(P_ij, P_ij);
      mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 1>(P_ij_bf16, P_ij);
      if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
      if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 3>(P_ij, P_ij);
      mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
      if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 2>(P_ij_bf16, P_ij);
      copy<0, 3>(P_ij_bf16, P_ij);
      if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
      swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
      if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
      asm volatile("s_waitcnt lgkmcnt(8)");
      // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
      mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      // Load K_j_col from shared memory to registers
      // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
      load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      sub_row<0, 0, delta_i>(dP_ij, dP_ij);
      sub_row<0, 1, delta_i>(dP_ij, dP_ij);
      mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      mul<0, 0>(dP_ij, dP_ij, P_ij);
      mul<0, 1>(dP_ij, dP_ij, P_ij);
      copy<0, 0>(dP_ij_bf16, dP_ij);
      copy<0, 1>(dP_ij_bf16, dP_ij);
      sub_row<0, 2, delta_i>(dP_ij, dP_ij);
      if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      // 12. dV_j += P_ij^T @ dO_i
      // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
      // Store dP_ij_bf16_accum_row to shared memory
      // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
      store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      sub_row<0, 3, delta_i>(dP_ij, dP_ij);
      mul<0, 2>(dP_ij, dP_ij, P_ij);
      mul<0, 3>(dP_ij, dP_ij, P_ij);
      copy<0, 2>(dP_ij_bf16, dP_ij);
      copy<0, 3>(dP_ij_bf16, dP_ij);
      if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

      // dot slice 3
      load<L_i>(subvec_inplace<DOT_SLICE_QO>(L_smem[tic], 3));
      load<delta_i>(subvec_inplace<DOT_SLICE_QO>(delta_smem[tic], 3));
      
      store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
      asm volatile("s_waitcnt lgkmcnt(12)");
      // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
      mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      asm volatile("s_waitcnt lgkmcnt(8)");
      if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
      __builtin_amdgcn_s_barrier();
      mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      // Load dP_ij_bf16_col_T from shared memory to registers
      // load(dP_ij_bf16_col_T, attn_i_smem);
      load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      mul<L_i, L_i>(L_SCALE_FACTOR);
      atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 1, 0}, warpid);
      if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 1, 0}, warpid);
      if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L6 ")");
      __builtin_amdgcn_s_barrier();
      // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
      // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);

      mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      // Load Q_i from shared memory to registers
      // load(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      Q_i_addr = get_address(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {1, 0}));
      load<0, 0>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      load<0, 1>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
      __builtin_amdgcn_s_barrier();
      mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
      mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      // Load K_j from shared memory to registers
      // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
      load<0, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<0, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<0, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<0, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mul<0, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mul<0, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      load<1, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<1, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mul<0, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mul<0, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
      if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) load<1, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<1, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
    }

    // dot slice 3
    {
      // 10. S_ij = Q_i K_j^T * scale
      // 11. P_ij = exp2(S_ij - L_i)
      // 13. dP_ij = dO_i @ V_j^T
      // 14. dS_ij = P_ij o (dP_ij - delta_i)
      // mma_ABt(P_ij, Q_i, K_j);
      mma_ABt<0, 0, 0>(P_ij, Q_i, K_j);
      load<2, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<2, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      mma_ABt<0, 0, 1>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mma_ABt<0, 0, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<2, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<2, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_ABt<0, 0, 3>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mul<1, 0, 0>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mul<1, 0, 1>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      mma_ABt<0, 1, 0>(P_ij, Q_i, K_j);
      load<3, 0>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      load<3, 1>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      mma_ABt<0, 1, 1>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) mul<1, 0, 2>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mul<1, 0, 3>(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
      if constexpr (D == 128) mma_ABt<0, 1, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<3, 2>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) load<3, 3>(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}), K_j_addr);
      if constexpr (D == 128) mma_ABt<0, 1, 3>(P_ij, Q_i, K_j, P_ij);
      mul<0, 0>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(6)");
      mma_ABt<0, 2, 0>(P_ij, Q_i, K_j);
      // Load dO_i from shared memory to registers
      // load(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      dO_i_addr = get_address(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {1, 0}));
      load<0, 0>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      load<0, 1>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      mma_ABt<0, 2, 1>(P_ij, Q_i, K_j, P_ij);
      sub_row<0, 0, L_i>(P_ij, P_ij);
      asm volatile("s_waitcnt lgkmcnt(6)");
      if constexpr (D == 128) mma_ABt<0, 2, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<0, 2>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      if constexpr (D == 128) load<0, 3>(dO_i, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_addr);
      if constexpr (D == 128) mma_ABt<0, 2, 3>(P_ij, Q_i, K_j, P_ij);
      mul<0, 1>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(6)");
      mma_ABt<0, 3, 0>(P_ij, Q_i, K_j);
      // Load dO_i_col from shared memory to registers
      // load(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      // Compute dO_i_col_addr
      // uint32_t dO_i_col_addr = get_address(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}));
      dO_i_col_addr = [&] {
        const int laneid = kittens::laneid();
        const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][1], {1, 0}).data[0]);
        const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
        const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
        const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
        const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
        const uint32_t addr = src_ptr + swizzled_lane_byte_offset;
        return addr;
      }();
      load<0, 0>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      load<0, 1>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      mma_ABt<0, 3, 1>(P_ij, Q_i, K_j, P_ij);
      sub_row<0, 1, L_i>(P_ij, P_ij);
      asm volatile("s_waitcnt lgkmcnt(8)");
      if constexpr (D == 128) mma_ABt<0, 3, 2>(P_ij, Q_i, K_j, P_ij);
      if constexpr (D == 128) load<0, 2>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      if constexpr (D == 128) load<0, 3>(dO_i_col, subtile_inplace<DOT_SLICE_QO, D>(dO_i_smem[tic][0], {0, 0}), dO_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 3, 3>(P_ij, Q_i, K_j, P_ij);
      // Dot slice 3
      kittens::macros::v_mov_b32_up2p<neg_inf_v>(0xff800000); if constexpr (causal) {
        // If the query position is less than the key position, set P_ij to -inf
        if (q_pos < k_pos) {
          mov<neg_inf_v>(P_ij);
        // If the query position is equal to the key position, we need to apply a causal mask
        } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
#if WINDOW
        } else if (q_pos - k_pos == WINDOW) {
            mov<0, 0, neg_inf_v>(P_ij);
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
        } else if (q_pos - k_pos > WINDOW) {
            mov<neg_inf_v>(P_ij);
#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
      asm volatile("s_waitcnt lgkmcnt(8)");
      // mma_ABt(dP_ij, dO_i, V_j);
      mma_ABt<0, 0, 0>(dP_ij, dO_i, V_j);
      sub_row<0, 2, L_i>(P_ij, P_ij);
      mma_ABt<0, 0, 1>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 0>(P_ij, P_ij);
      if constexpr (D == 128) mma_ABt<0, 0, 2>(dP_ij, dO_i, V_j, dP_ij);
      // Load Q_i_col from shared memory to registers
      // load(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      // Compute Q_i_col_addr
      // uint32_t Q_i_col_addr = get_address(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}));
      Q_i_col_addr = [&] {
        const int laneid = kittens::laneid();  
        const uint32_t src_ptr = reinterpret_cast<uintptr_t>(&subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][1], {1, 0}).data[0]);
        const int row_offset = (laneid % 16) / 4 + (laneid / 32) * 8;
        const int col_offset = ((laneid % 4) * 4) + 16*((laneid % 32)/16);
        const int lane_byte_offset = (row_offset * 32 + col_offset) * sizeof(bf16);
        const int swizzled_lane_byte_offset = lane_byte_offset ^ ((lane_byte_offset >> 9) << 5);
        const int addr = src_ptr + swizzled_lane_byte_offset;
        return addr;
      }();
      load<0, 0>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 0, 3>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 1>(P_ij, P_ij);
      mma_ABt<0, 1, 0>(dP_ij, dO_i, V_j);
      load<0, 1>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      mma_ABt<0, 1, 1>(dP_ij, dO_i, V_j, dP_ij);
      mul<0, 3>(P_ij, P_ij, P_SCALE_FACTOR);
      if constexpr (D == 128) mma_ABt<0, 1, 2>(dP_ij, dO_i, V_j, dP_ij);
      sub_row<0, 3, L_i>(P_ij, P_ij);
      if constexpr (D == 128) mma_ABt<0, 1, 3>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 0>(P_ij_bf16, P_ij);
      mma_ABt<0, 2, 0>(dP_ij, dO_i, V_j);
      exp2<0, 2>(P_ij, P_ij);
      mma_ABt<0, 2, 1>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 1>(P_ij_bf16, P_ij);
      if constexpr (D == 128) mma_ABt<0, 2, 2>(dP_ij, dO_i, V_j, dP_ij);
      if constexpr (D == 128) load<0, 2>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      if constexpr (D == 128) mma_ABt<0, 2, 3>(dP_ij, dO_i, V_j, dP_ij);
      exp2<0, 3>(P_ij, P_ij);
      mma_ABt<0, 3, 0>(dP_ij, dO_i, V_j);
      if constexpr (D == 128) load<0, 3>(Q_i_col, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_col_addr);
      mma_ABt<0, 3, 1>(dP_ij, dO_i, V_j, dP_ij);
      copy<0, 2>(P_ij_bf16, P_ij);
      copy<0, 3>(P_ij_bf16, P_ij);
      if constexpr (D == 128) mma_ABt<0, 3, 2>(dP_ij, dO_i, V_j, dP_ij);
      swap_layout_inplace(P_ij_bf16_col, P_ij_bf16);
      if constexpr (D == 128) mma_ABt<0, 3, 3>(dP_ij, dO_i, V_j, dP_ij);
      asm volatile("s_waitcnt lgkmcnt(8)");
      // mma_AtB(dV_j_T, dO_i_col, P_ij_bf16_col);
      mma_AtB<0, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      // Load K_j_col from shared memory to registers
      // load(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}));
      load<0, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<0, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      sub_row<0, 0, delta_i>(dP_ij, dP_ij);
      sub_row<0, 1, delta_i>(dP_ij, dP_ij);
      mma_AtB<1, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      load<1, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<1, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<1, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      mul<0, 0>(dP_ij, dP_ij, P_ij);
      mul<0, 1>(dP_ij, dP_ij, P_ij);
      copy<0, 0>(dP_ij_bf16, dP_ij);
      copy<0, 1>(dP_ij_bf16, dP_ij);
      sub_row<0, 2, delta_i>(dP_ij, dP_ij);
      if constexpr (D == 128) mma_AtB<2, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      load<2, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      // 12. dV_j += P_ij^T @ dO_i
      // 16. dK_j += dS_ij^T @ Q_i   (128x64)=(128x16)x(16x64)
      // Store dP_ij_bf16_accum_row to shared memory
      // store(attn_i_smem_subtile, dP_ij_bf16_accum_row);
      store<0, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      store<1, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      if constexpr (D == 128) mma_AtB<2, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      sub_row<0, 3, delta_i>(dP_ij, dP_ij);
      mul<0, 2>(dP_ij, dP_ij, P_ij);
      mul<0, 3>(dP_ij, dP_ij, P_ij);
      copy<0, 2>(dP_ij_bf16, dP_ij);
      copy<0, 3>(dP_ij_bf16, dP_ij);
      if constexpr (D == 128) mma_AtB<3, 0, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);

      // dot slice 0 - next iteration
      
      store<2, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      store<3, 0>(attn_i_smem_subtile, dP_ij_bf16_accum_row, dP_ij_bf16_accum_row_addr);
      if constexpr (D == 128) mma_AtB<3, 1, 0>(dV_j_T, dO_i_col, P_ij_bf16_col, dV_j_T);
      swap_layout_inplace(dP_ij_bf16_col, dP_ij_bf16);
      asm volatile("s_waitcnt lgkmcnt(12)");
      // mma_AtB(dK_j_T, Q_i_col, dP_ij_bf16_col);
      mma_AtB<0, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      if constexpr (D == 128) load<2, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<3, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<3, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<4, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      asm volatile("s_waitcnt lgkmcnt(8)");
      if constexpr (D == 64) asm volatile("s_waitcnt lgkmcnt(0)");
      __builtin_amdgcn_s_barrier();
      mma_AtB<1, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      // Load dP_ij_bf16_col_T from shared memory to registers
      // load(dP_ij_bf16_col_T, attn_i_smem);
      load<0, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<1, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<2, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<3, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      mma_AtB<1, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      atomic_pk_add_bf16_with_warpid<2, 0, 0>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 2, 0}, warpid);
      if constexpr (D == 128) mma_AtB<2, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      load<4, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<5, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      if constexpr (D == 128) load<4, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      load<5, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) mma_AtB<2, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      if constexpr (D == 128) atomic_pk_add_bf16_with_warpid<2, 0, 1>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 2, 0}, warpid);
      if constexpr (D == 128) mma_AtB<3, 0, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      load<6, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      load<7, 0>(dP_ij_bf16_col_T, attn_i_smem, dP_ij_bf16_col_T_addr);
      if constexpr (D == 128) load<5, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) mma_AtB<3, 1, 0>(dK_j_T, Q_i_col, dP_ij_bf16_col, dK_j_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L6 ")");
      __builtin_amdgcn_s_barrier();
      // 15. dQ_i += dS_ij @ K_j (32x16)=(32x256)x(256x16)
      // mma_AtB(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      mma_AtB<0, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      load<6, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<6, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mma_AtB<0, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      load<7, 0>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      if constexpr (D == 128) load<7, 1>(K_j_col, subtile_inplace<256, ATTN_D/4>(K_j_smem, {0, warpid}), K_j_col_addr);
      mma_AtB<0, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      mma_AtB<0, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      // ds_read_b128 a[112:115]
      // ds_read_b128 a[116:119]
      mma_AtB<0, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L4 ")");
      __builtin_amdgcn_s_barrier();
      mma_AtB<0, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      // ds_read_b128 a[120:123]
      // ds_read_b128 a[124:127]
      mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
      // ds_read_b128 a[0:3]
      // ds_read_b128 a[4:7]
      if constexpr (D == 128) mma_AtB<1, 0, 1>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 2>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      // ds_read_b128 a[8:11]
      // ds_read_b128 a[12:15]
      if constexpr (D == 128) mma_AtB<1, 0, 3>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      if constexpr (D == 128) mma_AtB<1, 0, 4>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      // ds_read_b128 a[16:19]
      // ds_read_b128 a[20:23]
      if constexpr (D == 128) mma_AtB<1, 0, 5>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L10 ")");
      if constexpr (D == 128) mma_AtB<1, 0, 6>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      // ds_read_b128 a[24:27]
      // ds_read_b128 a[28:31]
      if constexpr (D == 128) mma_AtB<1, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
      asm volatile("s_waitcnt lgkmcnt(" DQ_L2 ")");
    }
  }

  store<1>(g.dVg, dV_j, {batch_idx * (GROUP_SIZE / HEADS_PER_WG) + q_head_in_group / HEADS_PER_WG, 0, kv_head_idx, 0}, {0, j, 0, 0});
  __builtin_amdgcn_s_waitcnt(0);
  __builtin_amdgcn_s_barrier();

  // We first copy dV_j_T from accumulator GPRs to vector GPRs and then perform the store
  accvgpr_read(dV_j_T, dK_j_T);
  mul(dV_j_T, dV_j_T, dP_SCALE_FACTOR);
  store<1>(g.dKg, dV_j, {batch_idx * (GROUP_SIZE / HEADS_PER_WG) + q_head_in_group / HEADS_PER_WG, 0, kv_head_idx, 0}, {0, j, 0, 0});

  // Write out final dQ_i slice
  mul(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
  atomic_pk_add_bf16_with_warpid<2>(g.dQg, dQ_i, {batch_idx, q_head_idx, q_seq_idx * 4 + 3, 0}, warpid);
}

template __global__ void attend_bwd_combined_ker<ATTN_D>(bf16*, bf16*, bf16*, bf16*, bf16*, bf16*, bf16*, float*, float*);