/***************************************************************************************************
 * Copyright (c) 2025  - 2025 NVIDIA CORPORATION & AFFILIATES. All rights
 *reserved. SPDX-License-Identifier: BSD-3-Clause
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/

#pragma once

#include "cutlass/cutlass.h"

#include "cute/arch/simd_sm100.hpp"
#include "cute/tensor.hpp"

#include "cutlass/arch/arch.h"
#include "cutlass/arch/memory_sm80.h"
#include "cutlass/gemm/collective/collective_builder.hpp"

#include "natten/cuda/fna_blackwell/collective/fna_common.hpp"
#include "natten/cuda/fna_blackwell/collective/fna_fusion.hpp"
#include "natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp"

namespace cutlass::fna::kernel {

using namespace cutlass::fna::collective;

using namespace cute;

template <
    class Element,
    class ElementAcc,
    class TileShape,
    class Mask,
    class QTileShape,
    class KVTileShape,
    class NADim>
struct Sm100FnaBwdKernelTmaWarpSpecialized {
  using MultiDimTileShape = cute::tuple<QTileShape, KVTileShape>;

  static_assert(
      size(QTileShape{}) == get<0>(TileShape{}),
      "QTileShape doesn't match the size of Q tile in the FMHA kernel.");
  static_assert(
      size(KVTileShape{}) == get<1>(TileShape{}),
      "KVTileShape doesn't match the size of KV tile in the FMHA kernel.");

  using TileShapeQ = decltype(get<0>(TileShape{}));
  static_assert(std::is_same_v<TileShapeQ, _128>, "tile shape K must be 128");
  using TileShapeK = decltype(get<1>(TileShape{}));
  static_assert(std::is_same_v<TileShapeK, _128>, "tile shape K must be 128");
  using TileShapeDQK = decltype(get<2>(TileShape{}));
  using TileShapeDVO = decltype(get<2>(TileShape{}));

  using TmemAllocator = cute::TMEM::Allocator1Sm;
  struct TmemAllocation {
    static constexpr uint32_t kDK = 0; // TileShapeK x TileShapeDQK x acc
    static constexpr uint32_t kDV =
        kDK + TileShapeDQK{}; // TileShapeK x TileShapeDVO x acc
    static constexpr uint32_t kDQ =
        kDV + TileShapeDVO{}; // TileShapeQ x TileShapeDQK x acc
    static constexpr uint32_t kDP = kDQ; // TileShapeK x TileShapeQ   x inp
    static constexpr uint32_t kS = kDQ + max(TileShapeQ{}, TileShapeDQK{});
    static constexpr uint32_t kP = kS;
    static constexpr uint32_t kTotal = kS + TileShapeQ{};
  };

  static_assert(
      static_cast<int>(TmemAllocation::kTotal) <=
          TmemAllocator::Sm100TmemCapacityColumns,
      "using too much tmem");

  enum class WarpRole {
    Empty = 0x0,
    Load = 0x1,
    Mma = 0x2,
    Compute = 0x3,
    Reduce = 0x4
  };

  static constexpr unsigned long long kWarpAssignment = 0x12'3333'3333'4444ull;
  static constexpr int kNumComputeWarps = 8;
  static constexpr int kNumReduceWarps = 4;
  CUTLASS_DEVICE WarpRole warp_idx_to_role(int warp_idx) {
    return static_cast<WarpRole>((kWarpAssignment >> (4 * warp_idx)) & 0xF);
  }

  struct RegisterAllocation {
    static constexpr int kWarpgroup0 = 160 - 8;
    static constexpr int kWarpgroup1 = 128;
    static constexpr int kWarpgroup2 = 96;
    static constexpr int kReduce = kWarpgroup0;
    static constexpr int kCompute = kWarpgroup1;
    static constexpr int kMma = kWarpgroup2;
    static constexpr int kEmpty = kWarpgroup2;
    static constexpr int kLoad = kWarpgroup2;

    static_assert(kWarpgroup0 + 2 * kWarpgroup1 + kWarpgroup2 <= 512);
  };

  using ArchTag = cutlass::arch::Sm100;

  using ClusterShape = Shape<_1, _1, _1>;
  using Schedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100;

  static constexpr int MinBlocksPerMultiprocessor = 1;
  static constexpr int kNumWarps = kNumComputeWarps + kNumReduceWarps + 4;
  static constexpr int MaxThreadsPerBlock = NumThreadsPerWarp * kNumWarps;

  static constexpr int Alignment = 128 / sizeof_bits_v<Element>;
  static constexpr int kStages = 2;

  using TensorStrideContiguousK = Stride<int, _1, Stride<int, int>>;
  using TensorStrideContiguousMN = Stride<_1, int, Stride<int, int>>;

  // compute S
  using CollectiveMmaKQ = typename cutlass::gemm::collective::CollectiveBuilder<
      cutlass::arch::Sm100,
      cutlass::arch::OpClassTensorOp,
      Element,
      TensorStrideContiguousK,
      Alignment,
      Element,
      TensorStrideContiguousK,
      Alignment,
      ElementAcc,
      Shape<TileShapeK, TileShapeQ, TileShapeDQK>,
      ClusterShape,
      cutlass::gemm::collective::StageCount<kStages>,
      Schedule>::CollectiveOp;
  using TileShapeKQ = typename CollectiveMmaKQ::TileShape;
  using TiledMmaKQ = typename CollectiveMmaKQ::TiledMma;

  // compute dP
  using CollectiveMmaVDO =
      typename cutlass::gemm::collective::CollectiveBuilder<
          cutlass::arch::Sm100,
          cutlass::arch::OpClassTensorOp,
          Element,
          TensorStrideContiguousK,
          Alignment,
          Element,
          TensorStrideContiguousK,
          Alignment,
          ElementAcc,
          Shape<TileShapeK, TileShapeQ, TileShapeDVO>,
          ClusterShape,
          cutlass::gemm::collective::StageCount<kStages>,
          Schedule>::CollectiveOp;
  using TileShapeVDO = typename CollectiveMmaVDO::TileShape;
  using TiledMmaVDO = typename CollectiveMmaVDO::TiledMma;

  // compute dV
  using CollectiveMmaPDO =
      typename cutlass::gemm::collective::CollectiveBuilder<
          cutlass::arch::Sm100,
          cutlass::arch::OpClassTensorOp,
          // needs to match ordering of S calculation
          Element,
          TensorStrideContiguousK,
          Alignment,
          Element,
          TensorStrideContiguousMN,
          Alignment,
          ElementAcc,
          Shape<TileShapeK, TileShapeDVO, TileShapeQ>,
          ClusterShape,
          cutlass::gemm::collective::StageCount<kStages>,
          Schedule>::CollectiveOp;
  using TileShapePDO = typename CollectiveMmaPDO::TileShape;
  using TiledMmaPDO =
      decltype(to_tiled_mma_sm100_ts(typename CollectiveMmaPDO::TiledMma{}));

  // compute dK
  using CollectiveMmaDSQ =
      typename cutlass::gemm::collective::CollectiveBuilder<
          cutlass::arch::Sm100,
          cutlass::arch::OpClassTensorOp,
          // somewhat arbitrary since we dump to smem, need to agree with the
          // next one
          Element,
          TensorStrideContiguousK,
          Alignment,
          Element,
          TensorStrideContiguousMN,
          Alignment,
          ElementAcc,
          Shape<TileShapeK, TileShapeDQK, TileShapeQ>,
          ClusterShape,
          cutlass::gemm::collective::StageCount<kStages>,
          Schedule>::CollectiveOp;
  using TileShapeDSQ = typename CollectiveMmaDSQ::TileShape;
  using TiledMmaDSQ = typename CollectiveMmaDSQ::TiledMma;

  // compute dQ
  using CollectiveMmaDSK =
      typename cutlass::gemm::collective::CollectiveBuilder<
          cutlass::arch::Sm100,
          cutlass::arch::OpClassTensorOp,
          // somewhat arbitrary since we dump to smem, need to agree with the
          // previous one
          Element,
          TensorStrideContiguousMN,
          Alignment,
          Element,
          TensorStrideContiguousMN,
          Alignment,
          ElementAcc,
          Shape<TileShapeQ, TileShapeDQK, TileShapeK>,
          ClusterShape,
          cutlass::gemm::collective::StageCount<kStages>,
          Schedule>::CollectiveOp;
  using TileShapeDSK = typename CollectiveMmaDSK::TileShape;
  using TiledMmaDSK = typename CollectiveMmaDSK::TiledMma;

  // pipelines are named Pipeline<Producer><Consumer><Resource>
  static constexpr int kStagesComputeSmem = 1;
  using PipelineLoadMmaQ = PipelineTmaUmmaAsync<2, ClusterShape>;
  using PipelineLoadMmaDO = PipelineTmaUmmaAsync<1, ClusterShape>;
  using PipelineLoadComputeLSE = PipelineAsync<1>;
  using PipelineLoadComputeSumOdO = PipelineAsync<1>;
  using PipelineMmaComputeS = PipelineUmmaAsync<1>;
  using PipelineMmaComputeDP = PipelineUmmaAsync<1>;
  using PipelineMmaReduceDQ = PipelineUmmaAsync<1>;
  using PipelineComputeMmaP = PipelineUmmaConsumerAsync<1>;
  using PipelineComputeMmaDS = PipelineUmmaConsumerAsync<kStagesComputeSmem>;
  using PipelineMmaComputeDKDV = PipelineUmmaAsync<2>;
  static constexpr int kStagesReduceTmaStore = 2;
  using PipelineReduceTmaStore = PipelineTmaStore<kStagesReduceTmaStore>;

  struct PipelineStorage {
    alignas(16) typename PipelineLoadMmaQ::SharedStorage load_mma_q;
    alignas(16) typename PipelineLoadMmaDO::SharedStorage load_mma_do;
    alignas(16) typename PipelineLoadComputeLSE::SharedStorage load_compute_lse;
    alignas(16)
        typename PipelineLoadComputeSumOdO::SharedStorage load_compute_sum_odo;
    alignas(16) typename PipelineMmaComputeS::SharedStorage mma_compute_s;
    alignas(16) typename PipelineMmaComputeDP::SharedStorage mma_compute_dp;
    alignas(16) typename PipelineMmaReduceDQ::SharedStorage mma_reduce_dq;
    alignas(16) typename PipelineComputeMmaP::SharedStorage compute_mma_p;
    alignas(16) typename PipelineComputeMmaDS::SharedStorage compute_mma_ds;
    alignas(16) typename PipelineMmaComputeDKDV::SharedStorage mma_compute_dkdv;
  };

  template <class Layout, class Stages = _1>
  static CUTE_DEVICE constexpr auto restage(
      Layout const& layout,
      Stages stages = {}) {
    return composition(layout, make_tuple(_, _, _, make_layout(stages)));
  }

  using SmemLayoutK =
      decltype(restage(typename CollectiveMmaKQ::SmemLayoutA{}));
  using SmemLayoutV =
      decltype(restage(typename CollectiveMmaVDO::SmemLayoutA{}));
  using SmemLayoutQ =
      decltype(restage(typename CollectiveMmaKQ::SmemLayoutB{}, _2{}));
  using SmemLayoutDO =
      decltype(restage(typename CollectiveMmaVDO::SmemLayoutB{}, _1{}));
  using SmemLayoutDS = decltype(restage(
      typename CollectiveMmaDSK::SmemLayoutA{},
      Int<kStagesComputeSmem>{}));
  using SmemLayoutLSE = Layout<Shape<TileShapeQ, _1>>;
  using SmemLayoutSumOdO = Layout<Shape<TileShapeQ, _1>>;

  using SmemLayoutQT =
      decltype(restage(typename CollectiveMmaDSQ::SmemLayoutB{}, _2{}));
  using SmemLayoutKT =
      decltype(restage(typename CollectiveMmaDSK::SmemLayoutB{}));
  using SmemLayoutDST = decltype(restage(
      typename CollectiveMmaDSQ::SmemLayoutA{},
      Int<kStagesComputeSmem>{}));
  using SmemLayoutDOT =
      decltype(restage(typename CollectiveMmaPDO::SmemLayoutB{}, _1{}));

  using TileShapeDQ = _32;
  using SmemAtomDQ =
      decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
               cute::UMMA::Major::K,
               ElementAcc,
               TileShapeQ,
               TileShapeDQ>());
  using SmemShapeDQ =
      Shape<TileShapeQ, TileShapeDQ, Int<kStagesReduceTmaStore>>;
  using SmemLayoutDQ =
      decltype(tile_to_shape(SmemAtomDQ{}, SmemShapeDQ{}, Step<_2, _1, _3>{}));

  struct TensorStorage {
    union {
      alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutK>> smem_k;
      alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutKT>> smem_k_t;
    };
    alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutV>> smem_v;
    union {
      alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutQ>> smem_q;
      alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutQT>> smem_q_t;
    };
    union {
      alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutDO>> smem_do;
      alignas(
          2048) cute::array<Element, cute::cosize_v<SmemLayoutDOT>> smem_do_t;
    };
    union {
      alignas(2048) cute::array<Element, cute::cosize_v<SmemLayoutDS>> smem_ds;
      alignas(
          2048) cute::array<Element, cute::cosize_v<SmemLayoutDST>> smem_ds_t;
    };
    alignas(1024) cute::array<ElementAcc, cute::cosize_v<SmemLayoutDQ>> smem_dq;
    alignas(16) cute::array<ElementAcc, cute::cosize_v<SmemLayoutLSE>> smem_lse;
    alignas(16)
        cute::array<ElementAcc, cute::cosize_v<SmemLayoutSumOdO>> smem_sum_odo;
  };

  static constexpr int kTransactionsBytesLoadQ = cutlass::bits_to_bytes(
      cosize(take<0, 3>(SmemLayoutQ{})) * cute::sizeof_bits_v<Element>);
  static constexpr int kTransactionsBytesLoadDO = cutlass::bits_to_bytes(
      cosize(take<0, 3>(SmemLayoutDO{})) * cute::sizeof_bits_v<Element>);

  static constexpr int kTransactionsBytesLoadK = cutlass::bits_to_bytes(
      cosize(take<0, 3>(SmemLayoutK{})) * cute::sizeof_bits_v<Element>);
  static constexpr int kTransactionsBytesLoadV = cutlass::bits_to_bytes(
      cosize(take<0, 3>(SmemLayoutV{})) * cute::sizeof_bits_v<Element>);

  struct SharedStorage {
    TensorStorage tensors;
    PipelineStorage pipelines;
    uint32_t tmem_base_ptr;
  };

  // this is tight enough that it won't work with sizeof due to padding for
  // alignment
  static constexpr int SharedStorageSize =
      offsetof(SharedStorage, tmem_base_ptr) + sizeof(uint32_t);
  static_assert(
      SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes,
      "using too much smem");

  using ProblemShape =
      Shape<int, int, int, Shape<int, int>>; // Q K D (H B), eventuall D =
                                             // (D_QK, D_VO)
  using TensorStride = TensorStrideContiguousK; // S D (H B)
  using RowTensorStride = Stride<_1, Stride<int, int>>; // S (H B)

  struct MainloopArguments {
    const Element* ptr_q;
    TensorStride stride_q;
    const Element* ptr_k;
    TensorStride stride_k;
    const Element* ptr_v;
    TensorStride stride_v;
    const Element* ptr_do;
    TensorStride stride_do;

    const ElementAcc* ptr_lse;
    RowTensorStride stride_lse;

    const ElementAcc* ptr_sum_odo;
    RowTensorStride stride_sum_odo;

    ElementAcc* ptr_dq_acc;
    TensorStride stride_dq_acc;

    ElementAcc softmax_scale = 1.0f / sqrtf(TileShapeDQK{});
  };

  using TMA_K = typename CollectiveMmaKQ::Params::TMA_A;
  using TMA_V = typename CollectiveMmaVDO::Params::TMA_A;
  using TMA_Q = typename CollectiveMmaKQ::Params::TMA_B;
  using TMA_DO = typename CollectiveMmaVDO::Params::TMA_B;

  using TMA_DQ = decltype(make_tma_copy(
      SM90_TMA_REDUCE_ADD{},
      make_tensor(
          (const ElementAcc*)nullptr,
          make_shape(1, 1, make_shape(1, 1)),
          TensorStride{}),
      SmemLayoutDQ{}(_, _, _0{})));

  struct MainloopParams {
    TMA_K tma_load_k;
    TMA_V tma_load_v;
    TMA_Q tma_load_q;
    TMA_DO tma_load_do;
    TMA_DQ tma_red_dq;
  };

  struct EpilogueArguments {
    Element* ptr_dk;
    TensorStride stride_dk;
    Element* ptr_dv;
    TensorStride stride_dv;
  };

  struct FnaArguments {
    NADim q_shape;
    NADim kv_shape;
    NADim qkv_shape;
    NADim window_size;
    NADim stride;
    NADim dilation;
    int num_heads_actual; // heads / size(dilation)
  };

  struct FnaParams {
    NADim qkv_shape;
    NADim q_shape;
    NADim kv_shape;
    cute::tuple<NADim, NADim, NADim, NADim, NADim>
        na_params; // win, win_left, win_right, stride, stride_offset
    bool is_fully_block_sparse;
    bool has_q_padding;
    NADim dilation;
    bool requires_qkv_fixup;
    bool is_dilated;
    int num_heads_actual;
  };

  struct Arguments {
    ProblemShape problem_shape;
    FnaArguments fna;
    MainloopArguments mainloop;
    EpilogueArguments epilogue;
    KernelHardwareInfo hw_info;
  };

  struct Params {
    ProblemShape problem_shape;
    FnaParams fna;
    MainloopArguments mainloop;
    MainloopParams mainloop_params;
    EpilogueArguments epilogue;
    KernelHardwareInfo hw_info;
  };

  static bool can_implement(Arguments const& args) {
    auto [Q, K, D, HB] = args.problem_shape;
    auto [H, B] = HB;
    if (Q <= 0 || K <= 0 || D <= 0 || H <= 0 || B <= 0) {
      return false;
    }
    if (D % Alignment != 0) {
      return false;
    }
    // NOTE (ali): Turns out Blackwell FMHA is sensitive to this, but not Hopper
    // FMHA. In Hopper, the only alignment along seqlen_Q we needed to ensure
    // was for LSE, but this kernel fails TMA descriptor init due to
    // misalignment.
    if (Q % Alignment != 0) {
      return false;
    }
    // dK alone will end up with nans intermittently if this constraint is not
    // met.
    if (Q % TileShapeQ{} != 0) {
      return false;
    }
    return evenly_divides(args.fna.q_shape, QTileShape{}) &&
        evenly_divides(args.fna.kv_shape, KVTileShape{}) &&
        evenly_divides(H, size(args.fna.dilation)) && // dilation groups are
                                                      // folded into heads
        tuple_leq(args.fna.window_size, args.fna.qkv_shape) &&
        tuple_leq(args.fna.stride, args.fna.window_size);
  }

  static Status initialize_workspace(Arguments const&, void*, cudaStream_t) {
    return Status::kSuccess;
  }

  static Params to_underlying_arguments(Arguments const& args, void*) {
    auto [Q, K, D, HB] = args.problem_shape;

    auto params_kq = CollectiveMmaKQ::to_underlying_arguments(
        make_shape(K, Q, D, HB),
        typename CollectiveMmaKQ::Arguments{
            args.mainloop.ptr_k,
            args.mainloop.stride_k,
            args.mainloop.ptr_q,
            args.mainloop.stride_q,
        },
        /*workspace=*/nullptr);

    auto params_vdo = CollectiveMmaVDO::to_underlying_arguments(
        make_shape(K, Q, D, HB),
        typename CollectiveMmaVDO::Arguments{
            args.mainloop.ptr_v,
            args.mainloop.stride_v,
            args.mainloop.ptr_do,
            args.mainloop.stride_do,
        },
        /*workspace=*/nullptr);

    TMA_DQ tma_red_dq = make_tma_copy(
        SM90_TMA_REDUCE_ADD{},
        make_tensor(
            args.mainloop.ptr_dq_acc,
            make_shape(Q, D, HB),
            args.mainloop.stride_dq_acc),
        SmemLayoutDQ{}(_, _, _0{}));

    // FNA params
    auto window_left = get_window_left(args.fna.window_size);
    auto window_right = get_window_right(args.fna.window_size);
    auto stride_offset = get_bwd_stride_offset(args.fna.stride);
    auto na_params = make_tuple(
        args.fna.window_size,
        window_left,
        window_right,
        args.fna.stride,
        stride_offset);
    bool requires_qkv_fixup =
        not evenly_divides(args.fna.qkv_shape, args.fna.dilation);

    auto is_fully_block_sparse = fully_block_sparse<typename Mask::Causal>(
        args.fna.qkv_shape,
        args.fna.window_size,
        args.fna.stride,
        QTileShape{},
        KVTileShape{});

    return Params{
        args.problem_shape,
        FnaParams{
            args.fna.qkv_shape,
            args.fna.q_shape,
            args.fna.kv_shape,
            na_params,
            is_fully_block_sparse,
            /* has_q_padding */
            not evenly_divides(args.fna.qkv_shape, QTileShape{}),
            args.fna.dilation,
            requires_qkv_fixup,
            is_dilated(args.fna.dilation),
            args.fna.num_heads_actual},
        args.mainloop,
        MainloopParams{
            params_kq.tma_load_a,
            params_vdo.tma_load_a,
            params_kq.tma_load_b,
            params_vdo.tma_load_b,
            tma_red_dq},
        args.epilogue,
        args.hw_info};
  }

  template <class T>
  static CUTLASS_DEVICE auto quantize(T const& input) {
    constexpr int AlignmentS = 4;
    auto output = make_tensor<Element>(shape(input));
    auto input_vec = recast<Array<ElementAcc, AlignmentS>>(input);
    auto output_vec = recast<Array<Element, AlignmentS>>(output);

    cutlass::NumericArrayConverter<Element, ElementAcc, AlignmentS> epilogue_op;

    CUTLASS_PRAGMA_UNROLL
    for (int i = 0; i < size(input_vec); i++) {
      output_vec(i) = epilogue_op(input_vec(i));
    }

    return output;
  }

  template <class BlkCoord, class IterToTileMap>
  CUTLASS_DEVICE void load(
      BlkCoord const& blk_coord,
      ProblemShape const& problem_shape,
      int iter_index,
      int iter_count,
      IterToTileMap const& iter_to_tile_map,
      MainloopArguments const& mainloop_args,
      MainloopParams const& mainloop_params,
      TensorStorage& shared_tensors,
      PipelineLoadMmaQ& pipeline_load_mma_q,
      typename PipelineLoadMmaQ::PipelineState&
          pipeline_load_mma_q_producer_state,
      PipelineLoadMmaDO& pipeline_load_mma_do,
      typename PipelineLoadMmaDO::PipelineState&
          pipeline_load_mma_do_producer_state,
      PipelineLoadComputeLSE& pipeline_load_compute_lse,
      typename PipelineLoadComputeLSE::PipelineState&
          pipeline_load_compute_lse_producer_state,
      PipelineLoadComputeSumOdO& pipeline_load_compute_sum_odo,
      typename PipelineLoadComputeSumOdO::PipelineState&
          pipeline_load_compute_sum_odo_producer_state) {
    auto [Q, K, D, HB] = problem_shape;

    using X = Underscore;

    uint16_t mcast_mask = 0;

    auto mK = mainloop_params.tma_load_k.get_tma_tensor(make_shape(K, D, HB));
    auto mQ = mainloop_params.tma_load_q.get_tma_tensor(make_shape(Q, D, HB));
    auto mV = mainloop_params.tma_load_v.get_tma_tensor(make_shape(K, D, HB));
    auto mDO = mainloop_params.tma_load_do.get_tma_tensor(make_shape(Q, D, HB));

    auto gK =
        local_tile(mK, TileShapeKQ{}, make_coord(_, _, _), Step<_1, X, _1>{});
    auto gQ =
        local_tile(mQ, TileShapeKQ{}, make_coord(_, _, _), Step<X, _1, _1>{});
    auto gV =
        local_tile(mV, TileShapeVDO{}, make_coord(_, _, _), Step<_1, X, _1>{});
    auto gDO =
        local_tile(mDO, TileShapeVDO{}, make_coord(_, _, _), Step<X, _1, _1>{});

    ThrMMA cta_mma_kq = TiledMmaKQ{}.get_slice(_0{});
    ThrMMA cta_mma_vdo = TiledMmaVDO{}.get_slice(_0{});

    auto tSTgK = cta_mma_kq.partition_A(gK);
    auto tSTgQ = cta_mma_kq.partition_B(gQ);
    auto tDPTgV = cta_mma_vdo.partition_A(gV);
    auto tDPTgDO = cta_mma_vdo.partition_B(gDO);

    auto sQ = make_tensor(
        make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{});
    auto sK = make_tensor(
        make_smem_ptr(shared_tensors.smem_k.begin()), SmemLayoutK{});
    auto sV = make_tensor(
        make_smem_ptr(shared_tensors.smem_v.begin()), SmemLayoutV{});
    auto sDO = make_tensor(
        make_smem_ptr(shared_tensors.smem_do.begin()), SmemLayoutDO{});

    auto [tKgK_mkl, tKsK] = tma_partition(
        mainloop_params.tma_load_k,
        _0{},
        make_layout(_1{}),
        group_modes<0, 3>(sK),
        group_modes<0, 3>(tSTgK));
    auto [tQgQ_mkl, tQsQ] = tma_partition(
        mainloop_params.tma_load_q,
        _0{},
        make_layout(_1{}),
        group_modes<0, 3>(sQ),
        group_modes<0, 3>(tSTgQ));
    auto [tVgV_mkl, tVsV] = tma_partition(
        mainloop_params.tma_load_v,
        _0{},
        make_layout(_1{}),
        group_modes<0, 3>(sV),
        group_modes<0, 3>(tDPTgV));
    auto [tDOgDO_mkl, tDOsDO] = tma_partition(
        mainloop_params.tma_load_do,
        _0{},
        make_layout(_1{}),
        group_modes<0, 3>(sDO),
        group_modes<0, 3>(tDPTgDO));

    // set up lse and sum_odo

    auto [blk_coord_q, blk_coord_k, blk_coord_batch] = blk_coord;

    pipeline_load_mma_q.producer_acquire(pipeline_load_mma_q_producer_state);
    auto tma_barrier = pipeline_load_mma_q.producer_get_barrier(
        pipeline_load_mma_q_producer_state);

    pipeline_load_mma_q.producer_expect_transaction(
        pipeline_load_mma_q_producer_state, kTransactionsBytesLoadK);

    // load K
    if (cute::elect_one_sync()) {
      cute::copy(
          mainloop_params.tma_load_k.with(*tma_barrier, mcast_mask),
          tKgK_mkl(_, blk_coord_k, _0{}, blk_coord_batch),
          tKsK(_, _0{}));
    }

    // load Q
    if (cute::elect_one_sync()) {
      cute::copy(
          mainloop_params.tma_load_q.with(*tma_barrier, mcast_mask),
          tQgQ_mkl(_, iter_to_tile_map(iter_index), _0{}, blk_coord_batch),
          tQsQ(_, pipeline_load_mma_q_producer_state.index()));
    }

    ++pipeline_load_mma_q_producer_state;

    pipeline_load_compute_lse.producer_acquire(
        pipeline_load_compute_lse_producer_state);

    // load LSE
    // 32 threads loading 128 values of 32b each
    // so 4*32b=128b

    int thread_idx = threadIdx.x % NumThreadsPerWarp;
    int smem_idx =
        TileShapeQ{} * pipeline_load_compute_lse_producer_state.index() +
        thread_idx * 4;
    int gmem_idx = TileShapeQ{} * iter_to_tile_map(iter_index) + thread_idx * 4;
    auto mLSE = make_tensor(
        mainloop_args.ptr_lse, make_shape(Q, HB), mainloop_args.stride_lse);
    cutlass::arch::cp_async_zfill<16>(
        shared_tensors.smem_lse.begin() + smem_idx,
        &mLSE(gmem_idx, blk_coord_batch),
        gmem_idx < Q);

    pipeline_load_compute_lse.producer_commit(
        pipeline_load_compute_lse_producer_state,
        cutlass::arch::cpasync_barrier_arrive);
    ++pipeline_load_compute_lse_producer_state;

    pipeline_load_mma_do.producer_acquire(pipeline_load_mma_do_producer_state);
    tma_barrier = pipeline_load_mma_do.producer_get_barrier(
        pipeline_load_mma_do_producer_state);

    pipeline_load_mma_do.producer_expect_transaction(
        pipeline_load_mma_do_producer_state, kTransactionsBytesLoadV);

    // load V
    if (cute::elect_one_sync()) {
      cute::copy(
          mainloop_params.tma_load_v.with(*tma_barrier, mcast_mask),
          tVgV_mkl(_, blk_coord_k, _0{}, blk_coord_batch),
          tVsV(_, _0{}));
    }

    // load dO
    if (cute::elect_one_sync()) {
      cute::copy(
          mainloop_params.tma_load_do.with(*tma_barrier, mcast_mask),
          tDOgDO_mkl(_, iter_to_tile_map(iter_index), _0{}, blk_coord_batch),
          tDOsDO(_, pipeline_load_mma_do_producer_state.index()));
    }

    ++pipeline_load_mma_do_producer_state;

    pipeline_load_compute_sum_odo.producer_acquire(
        pipeline_load_compute_sum_odo_producer_state);

    // load sum_OdO
    smem_idx =
        TileShapeQ{} * pipeline_load_compute_sum_odo_producer_state.index() +
        thread_idx * 4;
    gmem_idx = TileShapeQ{} * iter_to_tile_map(iter_index) + thread_idx * 4;
    auto mSumOdO = make_tensor(
        mainloop_args.ptr_sum_odo,
        make_shape(Q, HB),
        mainloop_args.stride_sum_odo);
    cutlass::arch::cp_async<16>(
        shared_tensors.smem_sum_odo.begin() + smem_idx,
        &mSumOdO(gmem_idx, blk_coord_batch),
        gmem_idx < Q);

    pipeline_load_compute_sum_odo.producer_commit(
        pipeline_load_compute_sum_odo_producer_state,
        cutlass::arch::cpasync_barrier_arrive);
    ++pipeline_load_compute_sum_odo_producer_state;

    iter_count -= 1;
    iter_index += 1;

    while (iter_count > 0) {
      pipeline_load_mma_q.producer_acquire(pipeline_load_mma_q_producer_state);
      tma_barrier = pipeline_load_mma_q.producer_get_barrier(
          pipeline_load_mma_q_producer_state);

      // load Q
      if (cute::elect_one_sync()) {
        cute::copy(
            mainloop_params.tma_load_q.with(*tma_barrier, mcast_mask),
            tQgQ_mkl(_, iter_to_tile_map(iter_index), _0{}, blk_coord_batch),
            tQsQ(_, pipeline_load_mma_q_producer_state.index()));
      }

      ++pipeline_load_mma_q_producer_state;

      pipeline_load_compute_lse.producer_acquire(
          pipeline_load_compute_lse_producer_state);

      // load LSE
      smem_idx =
          TileShapeQ{} * pipeline_load_compute_lse_producer_state.index() +
          thread_idx * 4;
      gmem_idx = TileShapeQ{} * iter_to_tile_map(iter_index) + thread_idx * 4;
      cutlass::arch::cp_async<16>(
          shared_tensors.smem_lse.begin() + smem_idx,
          &mLSE(gmem_idx, blk_coord_batch),
          gmem_idx < Q);

      pipeline_load_compute_lse.producer_commit(
          pipeline_load_compute_lse_producer_state,
          cutlass::arch::cpasync_barrier_arrive);
      ++pipeline_load_compute_lse_producer_state;

      pipeline_load_mma_do.producer_acquire(
          pipeline_load_mma_do_producer_state);
      tma_barrier = pipeline_load_mma_do.producer_get_barrier(
          pipeline_load_mma_do_producer_state);

      // load dO
      if (cute::elect_one_sync()) {
        cute::copy(
            mainloop_params.tma_load_do.with(*tma_barrier, mcast_mask),
            tDOgDO_mkl(_, iter_to_tile_map(iter_index), _0{}, blk_coord_batch),
            tDOsDO(_, pipeline_load_mma_do_producer_state.index()));
      }

      ++pipeline_load_mma_do_producer_state;

      pipeline_load_compute_sum_odo.producer_acquire(
          pipeline_load_compute_sum_odo_producer_state);

      // load sum_OdO
      smem_idx =
          TileShapeQ{} * pipeline_load_compute_sum_odo_producer_state.index() +
          thread_idx * 4;
      gmem_idx = TileShapeQ{} * iter_to_tile_map(iter_index) + thread_idx * 4;
      cutlass::arch::cp_async_zfill<16>(
          shared_tensors.smem_sum_odo.begin() + smem_idx,
          &mSumOdO(gmem_idx, blk_coord_batch),
          gmem_idx < Q);

      pipeline_load_compute_sum_odo.producer_commit(
          pipeline_load_compute_sum_odo_producer_state,
          cutlass::arch::cpasync_barrier_arrive);
      ++pipeline_load_compute_sum_odo_producer_state;

      iter_count -= 1;
      iter_index += 1;
    }
  }

  template <class BlkCoord>
  CUTLASS_DEVICE void mma(
      BlkCoord const& blk_coord,
      ProblemShape const& problem_shape,
      int iter_index,
      int iter_count,
      MainloopArguments const& mainloop_args,
      TensorStorage& shared_tensors,
      PipelineLoadMmaQ& pipeline_load_mma_q,
      typename PipelineLoadMmaQ::PipelineState&
          pipeline_load_mma_q_consumer_state,
      PipelineLoadMmaDO& pipeline_load_mma_do,
      typename PipelineLoadMmaDO::PipelineState&
          pipeline_load_mma_do_consumer_state,
      PipelineMmaComputeS& pipeline_mma_compute_s,
      typename PipelineMmaComputeS::PipelineState&
          pipeline_mma_compute_s_producer_state,
      PipelineMmaComputeDP& pipeline_mma_compute_dp,
      typename PipelineMmaComputeDP::PipelineState&
          pipeline_mma_compute_dp_producer_state,
      PipelineMmaReduceDQ& pipeline_mma_reduce_dq,
      typename PipelineMmaReduceDQ::PipelineState&
          pipeline_mma_reduce_dq_producer_state,
      PipelineComputeMmaP& pipeline_compute_mma_p,
      typename PipelineComputeMmaP::PipelineState&
          pipeline_compute_mma_p_consumer_state,
      PipelineComputeMmaDS& pipeline_compute_mma_ds,
      typename PipelineComputeMmaDS::PipelineState&
          pipeline_compute_mma_ds_consumer_state,
      PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv,
      typename PipelineMmaComputeDKDV::PipelineState&
          pipeline_mma_compute_dkdv_producer_state) {
    auto [Q, K, D, HB] = problem_shape;

    auto sQ = make_tensor(
        make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{});
    auto sK = make_tensor(
        make_smem_ptr(shared_tensors.smem_k.begin()), SmemLayoutK{});
    auto sV = make_tensor(
        make_smem_ptr(shared_tensors.smem_v.begin()), SmemLayoutV{});
    auto sDO = make_tensor(
        make_smem_ptr(shared_tensors.smem_do.begin()), SmemLayoutDO{});

    auto sQT = make_tensor(
        make_smem_ptr(shared_tensors.smem_q_t.begin()), SmemLayoutQT{});
    auto sKT = make_tensor(
        make_smem_ptr(shared_tensors.smem_k_t.begin()), SmemLayoutKT{});
    auto sDS = make_tensor(
        make_smem_ptr(shared_tensors.smem_ds.begin()), SmemLayoutDS{});
    auto sDST = make_tensor(
        make_smem_ptr(shared_tensors.smem_ds_t.begin()), SmemLayoutDST{});
    auto sP = make_tensor(
        make_smem_ptr((Element*)nullptr),
        typename CollectiveMmaPDO::SmemLayoutA{});
    auto sDOT = make_tensor(
        make_smem_ptr(shared_tensors.smem_do_t.begin()), SmemLayoutDOT{});

    Tensor tSTrK = TiledMmaKQ::make_fragment_A(sK);
    Tensor tSTrQ = TiledMmaKQ::make_fragment_B(sQ);

    Tensor tDPTrV = TiledMmaVDO::make_fragment_A(sV);
    Tensor tDPTrDO = TiledMmaVDO::make_fragment_B(sDO);

    Tensor tDQrDS = TiledMmaDSK::make_fragment_A(sDS);
    Tensor tDQrKT = TiledMmaDSK::make_fragment_B(sKT);

    Tensor tDKrDST = TiledMmaDSQ::make_fragment_A(sDST);
    Tensor tDKrQT = TiledMmaDSQ::make_fragment_B(sQT);

    Tensor tDVrP = TiledMmaPDO::make_fragment_A(sP)(_, _, _, _0{});
    tDVrP.data() = TmemAllocation::kP;
    Tensor tDVrDOT = TiledMmaPDO::make_fragment_B(sDOT);

    TiledMmaKQ tiled_mma_kq;
    TiledMmaVDO tiled_mma_vdo;
    TiledMmaDSK tiled_mma_dsk;
    TiledMmaDSQ tiled_mma_dsq;
    TiledMmaPDO tiled_mma_pdo;

    tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::Zero;
    tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::Zero;

    Tensor tSTtST =
        partition_fragment_C(tiled_mma_kq, select<0, 1>(TileShapeKQ{}));
    tSTtST.data() = TmemAllocation::kS;

    Tensor tDPTtDPT =
        partition_fragment_C(tiled_mma_vdo, select<0, 1>(TileShapeVDO{}));
    tDPTtDPT.data() = TmemAllocation::kDP;

    Tensor tDQtDQ =
        partition_fragment_C(tiled_mma_dsk, select<0, 1>(TileShapeDSK{}));
    tDQtDQ.data() = TmemAllocation::kDQ;

    Tensor tDKtDK =
        partition_fragment_C(tiled_mma_dsq, select<0, 1>(TileShapeDSQ{}));
    tDKtDK.data() = TmemAllocation::kDK;

    Tensor tDVtDV =
        partition_fragment_C(tiled_mma_pdo, select<0, 1>(TileShapePDO{}));
    tDVtDV.data() = TmemAllocation::kDV;

    auto pipeline_load_mma_q_release_state = pipeline_load_mma_q_consumer_state;

    pipeline_load_mma_q.consumer_wait(pipeline_load_mma_q_consumer_state);
    pipeline_mma_compute_s.producer_acquire(
        pipeline_mma_compute_s_producer_state);

    // S = Q*K
    tiled_mma_kq.accumulate_ = UMMA::ScaleOut::Zero;
    CUTLASS_PRAGMA_UNROLL
    for (int k_block = 0; k_block < size<2>(tSTrQ); ++k_block) {
      cute::gemm(
          tiled_mma_kq,
          tSTrK(_, _, k_block, _0{}),
          tSTrQ(_, _, k_block, pipeline_load_mma_q_consumer_state.index()),
          tSTtST);
      tiled_mma_kq.accumulate_ = UMMA::ScaleOut::One;
    }

    ++pipeline_load_mma_q_consumer_state;

    pipeline_mma_compute_s.producer_commit(
        pipeline_mma_compute_s_producer_state);
    ++pipeline_mma_compute_s_producer_state;

    pipeline_load_mma_do.consumer_wait(pipeline_load_mma_do_consumer_state);

    pipeline_mma_compute_dp.producer_acquire(
        pipeline_mma_compute_dp_producer_state);
    pipeline_mma_reduce_dq.producer_acquire(
        pipeline_mma_reduce_dq_producer_state);

    // dP = dO*V
    tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::Zero;
    CUTLASS_PRAGMA_UNROLL
    for (int k_block = 0; k_block < size<2>(tDPTrV); ++k_block) {
      cute::gemm(
          tiled_mma_vdo,
          tDPTrV(_, _, k_block, _0{}),
          tDPTrDO(_, _, k_block, pipeline_load_mma_do_consumer_state.index()),
          tDPTtDPT);
      tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::One;
    }

    pipeline_mma_compute_dp.producer_commit(
        pipeline_mma_compute_dp_producer_state);
    ++pipeline_mma_compute_dp_producer_state;

    pipeline_compute_mma_p.consumer_wait(pipeline_compute_mma_p_consumer_state);

    // dV = P*dO
    CUTLASS_PRAGMA_UNROLL
    for (int k_block = 0; k_block < size<2>(tDVrP); ++k_block) {
      cute::gemm(
          tiled_mma_pdo,
          tDVrP(_, _, k_block),
          tDVrDOT(_, _, k_block, pipeline_load_mma_do_consumer_state.index()),
          tDVtDV);
      tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::One;
    }

    pipeline_compute_mma_p.consumer_release(
        pipeline_compute_mma_p_consumer_state);
    ++pipeline_compute_mma_p_consumer_state;

    pipeline_load_mma_do.consumer_release(pipeline_load_mma_do_consumer_state);
    ++pipeline_load_mma_do_consumer_state;

    iter_count -= 1;

    // in tmem, S & P overlap
    // and dP and dQ overlap
    // so we need to acquire dQ and dP at the same time
    while (iter_count > 0) {
      pipeline_load_mma_q.consumer_wait(pipeline_load_mma_q_consumer_state);
      pipeline_mma_compute_s.producer_acquire(
          pipeline_mma_compute_s_producer_state);

      // S = Q*K
      tiled_mma_kq.accumulate_ = UMMA::ScaleOut::Zero;
      CUTLASS_PRAGMA_UNROLL
      for (int k_block = 0; k_block < size<2>(tSTrQ); ++k_block) {
        cute::gemm(
            tiled_mma_kq,
            tSTrK(_, _, k_block, _0{}),
            tSTrQ(_, _, k_block, pipeline_load_mma_q_consumer_state.index()),
            tSTtST);
        tiled_mma_kq.accumulate_ = UMMA::ScaleOut::One;
      }

      ++pipeline_load_mma_q_consumer_state;

      pipeline_mma_compute_s.producer_commit(
          pipeline_mma_compute_s_producer_state);
      ++pipeline_mma_compute_s_producer_state;

      pipeline_compute_mma_ds.consumer_wait(
          pipeline_compute_mma_ds_consumer_state);

      // we need to acquire dP here, because tmem dQ == tmem dP
      pipeline_mma_compute_dp.producer_acquire(
          pipeline_mma_compute_dp_producer_state);

      // dQ = dS*K
      tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::Zero;
      CUTLASS_PRAGMA_UNROLL
      for (int k_block = 0; k_block < size<2>(tDQrDS); ++k_block) {
        cute::gemm(
            tiled_mma_dsk,
            tDQrDS(
                _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()),
            tDQrKT(_, _, k_block, _0{}),
            tDQtDQ);
        tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::One;
      }

      pipeline_mma_reduce_dq.producer_commit(
          pipeline_mma_reduce_dq_producer_state);
      ++pipeline_mma_reduce_dq_producer_state;

      // dK = dS*Q
      CUTLASS_PRAGMA_UNROLL
      for (int k_block = 0; k_block < size<2>(tDKrDST); ++k_block) {
        cute::gemm(
            tiled_mma_dsq,
            tDKrDST(
                _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()),
            tDKrQT(_, _, k_block, pipeline_load_mma_q_release_state.index()),
            tDKtDK);
        tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::One;
      }

      pipeline_load_mma_q.consumer_release(pipeline_load_mma_q_release_state);
      ++pipeline_load_mma_q_release_state;

      pipeline_compute_mma_ds.consumer_release(
          pipeline_compute_mma_ds_consumer_state);
      ++pipeline_compute_mma_ds_consumer_state;

      // we grab dq here, because in tmem dq == dp
      pipeline_mma_reduce_dq.producer_acquire(
          pipeline_mma_reduce_dq_producer_state);

      pipeline_load_mma_do.consumer_wait(pipeline_load_mma_do_consumer_state);

      // dP = dO*V
      tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::Zero;
      CUTLASS_PRAGMA_UNROLL
      for (int k_block = 0; k_block < size<2>(tDPTrV); ++k_block) {
        cute::gemm(
            tiled_mma_vdo,
            tDPTrV(_, _, k_block, _0{}),
            tDPTrDO(_, _, k_block, pipeline_load_mma_do_consumer_state.index()),
            tDPTtDPT);
        tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::One;
      }

      pipeline_mma_compute_dp.producer_commit(
          pipeline_mma_compute_dp_producer_state);
      ++pipeline_mma_compute_dp_producer_state;

      pipeline_compute_mma_p.consumer_wait(
          pipeline_compute_mma_p_consumer_state);

      // dV = P*dO
      CUTLASS_PRAGMA_UNROLL
      for (int k_block = 0; k_block < size<2>(tDVrP); ++k_block) {
        cute::gemm(
            tiled_mma_pdo,
            tDVrP(_, _, k_block),
            tDVrDOT(_, _, k_block, pipeline_load_mma_do_consumer_state.index()),
            tDVtDV);
        tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::One;
      }

      pipeline_compute_mma_p.consumer_release(
          pipeline_compute_mma_p_consumer_state);
      ++pipeline_compute_mma_p_consumer_state;

      pipeline_load_mma_do.consumer_release(
          pipeline_load_mma_do_consumer_state);
      ++pipeline_load_mma_do_consumer_state;

      iter_count -= 1;
    }

    // signal to the epilogue that dV is ready
    pipeline_mma_compute_dkdv.producer_acquire(
        pipeline_mma_compute_dkdv_producer_state);
    pipeline_mma_compute_dkdv.producer_commit(
        pipeline_mma_compute_dkdv_producer_state);
    ++pipeline_mma_compute_dkdv_producer_state;

    pipeline_mma_compute_dkdv.producer_acquire(
        pipeline_mma_compute_dkdv_producer_state);

    pipeline_compute_mma_ds.consumer_wait(
        pipeline_compute_mma_ds_consumer_state);

    // dK = dS*Q
    CUTLASS_PRAGMA_UNROLL
    for (int k_block = 0; k_block < size<2>(tDKrDST); ++k_block) {
      cute::gemm(
          tiled_mma_dsq,
          tDKrDST(
              _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()),
          tDKrQT(_, _, k_block, pipeline_load_mma_q_release_state.index()),
          tDKtDK);
      tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::One;
    }

    // signal to epilgue that dK is ready
    pipeline_mma_compute_dkdv.producer_commit(
        pipeline_mma_compute_dkdv_producer_state);
    ++pipeline_mma_compute_dkdv_producer_state;

    // we've already acquired mma_reduce_dq in the loop

    // dQ = dS*K
    tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::Zero;
    CUTLASS_PRAGMA_UNROLL
    for (int k_block = 0; k_block < size<2>(tDQrDS); ++k_block) {
      cute::gemm(
          tiled_mma_dsk,
          tDQrDS(_, _, k_block, pipeline_compute_mma_ds_consumer_state.index()),
          tDQrKT(_, _, k_block, _0{}),
          tDQtDQ);
      tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::One;
    }

    pipeline_mma_reduce_dq.producer_commit(
        pipeline_mma_reduce_dq_producer_state);
    ++pipeline_mma_reduce_dq_producer_state;

    pipeline_load_mma_q.consumer_release(pipeline_load_mma_q_release_state);
    ++pipeline_load_mma_q_release_state;

    pipeline_compute_mma_ds.consumer_release(
        pipeline_compute_mma_ds_consumer_state);
    ++pipeline_compute_mma_ds_consumer_state;
  }

  template <class TensorG, class TensorR, class TensorC, class TensorShape>
  CUTLASS_DEVICE void store(
      TensorG gmem,
      TensorR const& regs,
      TensorC const& coord,
      TensorShape const& tensor_shape) {
    Tensor preds = cute::lazy::transform(
        coord, [&](auto const& c) { return elem_less(c, tensor_shape); });

    auto copy_op = make_cotiled_copy(
        Copy_Atom<UniversalCopy<uint128_t>, Element>{},
        make_layout(
            make_shape(_1{}, Int<sizeof(uint128_t) / sizeof(Element)>{})),
        regs.layout());
    auto thr_copy = copy_op.get_slice(_0{});

    Tensor tCg = thr_copy.partition_D(gmem);
    Tensor tCr = thr_copy.partition_S(quantize(regs));
    Tensor tPc = thr_copy.partition_D(preds);

    copy_if(copy_op, tPc, tCr, tCg);
  }

  template <class BlkCoord>
  CUTLASS_DEVICE void epilogue(
      BlkCoord const& blk_coord,
      ProblemShape const& problem_shape,
      MainloopArguments const& mainloop_args,
      EpilogueArguments const& epilogue_args,
      PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv,
      typename PipelineMmaComputeDKDV::PipelineState&
          pipeline_mma_compute_dkdv_consumer_state) {
    auto [Q, K, D, HB] = problem_shape;
    auto [blk_coord_q, blk_coord_k, blk_coord_batch] = blk_coord;

    auto load_op = SM100_TMEM_LOAD_32dp32b16x{};

    auto tDKtDK =
        partition_fragment_C(TiledMmaDSQ{}, select<0, 1>(TileShapeDSQ{}))(
            make_coord(_, _), _0{}, _0{});
    tDKtDK.data() = TmemAllocation::kDK;

    auto mDK = make_tensor(
        make_gmem_ptr(epilogue_args.ptr_dk),
        make_shape(K, TileShapeDQK{}, HB),
        epilogue_args.stride_dk);
    auto gDK =
        local_tile(mDK, TileShapeDSQ{}, make_coord(_, _, _), Step<_1, _1, X>{})(
            _, _, blk_coord_k, _0{}, blk_coord_batch);

    Tensor cDK = domain_offset(
        make_coord(get<1>(blk_coord) * TileShapeK{}, _0{}),
        make_identity_tensor(take<0, 2>(TileShapeDSQ{})));

    constexpr int kNumWarpgroups = kNumComputeWarps / 4;
    int dp_idx = threadIdx.x % 128;
    int wg_idx = (threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp)) / 128;

    auto split_wg = [&](auto const& t) {
      if constexpr (decltype(rank(t))::value == 3) {
        auto p = t.compose(make_layout(make_shape(
            size<0>(t),
            size<1>(t),
            make_shape(
                Int<kNumWarpgroups>{}, size<2>(t) / Int<kNumWarpgroups>{}))));
        return p(_, _, make_coord(wg_idx, _));
      } else {
        auto p = t.compose(make_layout(make_shape(
            size<0>(t),
            size<1>(t),
            size<2>(t),
            make_shape(
                Int<kNumWarpgroups>{}, size<3>(t) / Int<kNumWarpgroups>{}))));
        return p(_, _, _, make_coord(wg_idx, _));
      }
    };

    auto tiled_t2r_dk = make_tmem_copy(load_op, tDKtDK);
    auto thread_t2r_dk = tiled_t2r_dk.get_slice(dp_idx);

    Tensor tTR_cDK = split_wg(thread_t2r_dk.partition_D(cDK));
    Tensor tTR_gDK = split_wg(thread_t2r_dk.partition_D(gDK));
    Tensor tTR_rDK = make_tensor<ElementAcc>(shape(tTR_cDK));
    Tensor tTR_tDK = split_wg(thread_t2r_dk.partition_S(tDKtDK));

    auto tDVtDV =
        partition_fragment_C(TiledMmaDSQ{}, select<0, 1>(TileShapeDSQ{}))(
            make_coord(_, _), _0{}, _0{});
    tDVtDV.data() = TmemAllocation::kDV;

    auto mDV = make_tensor(
        make_gmem_ptr(epilogue_args.ptr_dv),
        make_shape(K, TileShapeDVO{}, HB),
        epilogue_args.stride_dv);
    auto gDV =
        local_tile(mDV, TileShapePDO{}, make_coord(_, _, _), Step<_1, _1, X>{})(
            _, _, blk_coord_k, _0{}, blk_coord_batch);

    Tensor cDV = domain_offset(
        make_coord(get<1>(blk_coord) * TileShapeK{}, _0{}),
        make_identity_tensor(take<0, 2>(TileShapePDO{})));

    auto tiled_t2r_dv = make_tmem_copy(load_op, tDVtDV);
    auto thread_t2r_dv = tiled_t2r_dv.get_slice(dp_idx);

    Tensor tTR_cDV = split_wg(thread_t2r_dv.partition_D(cDV));
    Tensor tTR_gDV = split_wg(thread_t2r_dv.partition_D(gDV));
    Tensor tTR_rDV = make_tensor<ElementAcc>(shape(tTR_cDV));
    Tensor tTR_tDV = split_wg(thread_t2r_dv.partition_S(tDVtDV));

    pipeline_mma_compute_dkdv.consumer_wait(
        pipeline_mma_compute_dkdv_consumer_state);

    // load tDVtDV
    cute::copy(tiled_t2r_dv, tTR_tDV, tTR_rDV);

    // store tDVgDV
    store(tTR_gDV, tTR_rDV, tTR_cDV, select<1, 2>(problem_shape));

    cutlass::arch::fence_view_async_tmem_load();
    pipeline_mma_compute_dkdv.consumer_release(
        pipeline_mma_compute_dkdv_consumer_state);
    ++pipeline_mma_compute_dkdv_consumer_state;

    pipeline_mma_compute_dkdv.consumer_wait(
        pipeline_mma_compute_dkdv_consumer_state);

    // load tDKtDK
    cute::copy(tiled_t2r_dk, tTR_tDK, tTR_rDK);

    CUTLASS_PRAGMA_UNROLL
    for (int i = 0; i < size(tTR_rDK); i++) {
      tTR_rDK(i) = mainloop_args.softmax_scale * tTR_rDK(i);
    }

    // store tDKgDK
    store(tTR_gDK, tTR_rDK, tTR_cDK, select<1, 2>(problem_shape));

    cutlass::arch::fence_view_async_tmem_load();
    pipeline_mma_compute_dkdv.consumer_release(
        pipeline_mma_compute_dkdv_consumer_state);
    ++pipeline_mma_compute_dkdv_consumer_state;
  }

  template <class BlkCoord, class QKVShape, class NAParams>
  CUTLASS_DEVICE void compute(
      BlkCoord const& blk_coord,
      ProblemShape const& problem_shape,
      int iter_index,
      int iter_count,
      QKVShape const& kv_shape,
      QKVShape const& qkv_shape,
      NAParams const& na_params,
      QKVShape const& q_offset,
      QKVShape const& q_tiles,
      bool is_fully_block_sparse,
      bool has_q_padding,
      MainloopArguments const& mainloop_args,
      EpilogueArguments const& epilogue_args,
      TensorStorage& shared_tensors,
      PipelineLoadComputeLSE& pipeline_load_compute_lse,
      typename PipelineLoadComputeLSE::PipelineState&
          pipeline_load_compute_lse_consumer_state,
      PipelineLoadComputeSumOdO& pipeline_load_compute_sum_odo,
      typename PipelineLoadComputeSumOdO::PipelineState&
          pipeline_load_compute_sum_odo_consumer_state,
      PipelineMmaComputeS& pipeline_mma_compute_s,
      typename PipelineMmaComputeS::PipelineState&
          pipeline_mma_compute_s_consumer_state,
      PipelineMmaComputeDP& pipeline_mma_compute_dp,
      typename PipelineMmaComputeDP::PipelineState&
          pipeline_mma_compute_dp_consumer_state,
      PipelineComputeMmaP& pipeline_compute_mma_p,
      typename PipelineComputeMmaP::PipelineState&
          pipeline_compute_mma_p_producer_state,
      PipelineComputeMmaDS& pipeline_compute_mma_ds,
      typename PipelineComputeMmaDS::PipelineState&
          pipeline_compute_mma_ds_producer_state,
      PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv,
      typename PipelineMmaComputeDKDV::PipelineState&
          pipeline_mma_compute_dkdv_consumer_state) {
    auto [Q, K, D, HB] = problem_shape;

    // in tmem, S & P overlap
    // and dP and dQ overlap

    // there are two compute wg's that cooperatively compute softmax
    // they are striped by this tmem atom, i.e. wg0 has 16 elems, then wg1 etc

    auto load_op = SM100_TMEM_LOAD_32dp32b16x{};
    auto store_op = SM100_TMEM_STORE_32dp32b8x{};

    Tensor tSTtST =
        partition_fragment_C(TiledMmaKQ{}, select<0, 1>(TileShapeKQ{}))(
            make_coord(_, _), _0{}, _0{});
    tSTtST.data() = TmemAllocation::kS;

    Tensor tDPTtDPT =
        partition_fragment_C(TiledMmaVDO{}, select<0, 1>(TileShapeVDO{}))(
            make_coord(_, _), _0{}, _0{});
    tDPTtDPT.data() = TmemAllocation::kDP;

    Tensor cST = make_identity_tensor(take<0, 2>(TileShapeKQ{}));
    Tensor cDPT = make_identity_tensor(take<0, 2>(TileShapeVDO{}));

    constexpr int kNumWarpgroups = kNumComputeWarps / 4;
    int dp_idx = threadIdx.x % 128;
    int wg_idx = (threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp)) / 128;
    auto tiled_t2r = make_tmem_copy(load_op, tSTtST);
    auto thread_t2r = tiled_t2r.get_slice(dp_idx);

    auto split_wg = [&](auto const& t) {
      if constexpr (decltype(rank(t))::value == 3) {
        auto p = t.compose(make_layout(make_shape(
            size<0>(t),
            size<1>(t),
            make_shape(
                Int<kNumWarpgroups>{}, size<2>(t) / Int<kNumWarpgroups>{}))));
        return p(_, _, make_coord(wg_idx, _));
      } else {
        auto p = t.compose(make_layout(make_shape(
            size<0>(t),
            size<1>(t),
            size<2>(t),
            make_shape(
                Int<kNumWarpgroups>{}, size<3>(t) / Int<kNumWarpgroups>{}))));
        return p(_, _, _, make_coord(wg_idx, _));
      }
    };

    Tensor tTR_cST = split_wg(thread_t2r.partition_D(cST));
    Tensor tTR_rST = make_tensor<ElementAcc>(shape(tTR_cST));
    Tensor tTR_tST = split_wg(thread_t2r.partition_S(tSTtST));

    Tensor tTR_cDPT_p = thread_t2r.partition_D(cDPT);
    Tensor tTR_cDPT = split_wg(tTR_cDPT_p);
    Tensor tTR_rDPT = make_tensor<ElementAcc>(shape(tTR_cDPT));
    Tensor tTR_tDPT = split_wg(thread_t2r.partition_S(tDPTtDPT));

    Tensor sLSE = make_tensor(
        make_smem_ptr(shared_tensors.smem_lse.begin()), SmemLayoutLSE{});
    Tensor sSumOdO = make_tensor(
        make_smem_ptr(shared_tensors.smem_sum_odo.begin()), SmemLayoutSumOdO{});

    auto sP = make_tensor(
        make_smem_ptr((Element*)nullptr),
        typename CollectiveMmaPDO::SmemLayoutA{});

    auto tDVrP = TiledMmaPDO::make_fragment_A(sP)(_, _, _, _0{});
    auto tDVcST = TiledMmaPDO{}.get_slice(_0{}).partition_A(cST);
    tDVrP.data() = TmemAllocation::kP;

    auto tiled_r2t = make_tmem_copy(store_op, tDVrP);
    auto thread_r2t = tiled_r2t.get_slice(dp_idx);

    auto tRT_tP = split_wg(thread_r2t.partition_D(tDVrP));
    auto tRT_cST = split_wg(thread_r2t.partition_S(tDVcST));

    CUTLASS_PRAGMA_NO_UNROLL
    while (iter_count > 0) {
      // wait for S and P
      pipeline_mma_compute_s.consumer_wait(
          pipeline_mma_compute_s_consumer_state);
      pipeline_compute_mma_p.producer_acquire(
          pipeline_compute_mma_p_producer_state);
      // wait for LSE
      pipeline_load_compute_lse.consumer_wait(
          pipeline_load_compute_lse_consumer_state);

      // compute P = softmax(S, LSE)
      cute::copy(tiled_t2r, tTR_tST, tTR_rST);

      if (not is_fully_block_sparse) {
        Mask{}.apply_mask(
            tTR_rST,
            /* index_q = */
            [&](int i) {
              auto c_transpose = tTR_cST(i);
              // NOTE: we only care about q_idx here!
              return get<1>(c_transpose) + iter_index * TileShapeQ{};
            },
            /* index_qk = */
            [&](int i) {
              auto c_transpose = tTR_cST(i);
              return make_coord(
                  get<1>(c_transpose) + iter_index * TileShapeQ{},
                  get<0>(c_transpose) + get<1>(blk_coord) * TileShapeK{});
            },
            MultiDimTileShape{},
            kv_shape,
            qkv_shape,
            na_params,
            q_offset,
            q_tiles);
      } else if (has_q_padding) {
        Mask{}.apply_padded_mask(
            tTR_rST,
            /* index_q = */
            [&](int i) {
              auto c_transpose = tTR_cST(i);
              // NOTE: we only care about q_idx here!
              return get<1>(c_transpose) + iter_index * TileShapeQ{};
            },
            MultiDimTileShape{},
            qkv_shape,
            na_params,
            q_offset,
            q_tiles);
      }

      ElementAcc log2_e = static_cast<ElementAcc>(M_LOG2E);
      float2 softmax_scale_log2_e;
      softmax_scale_log2_e.x = mainloop_args.softmax_scale * log2_e;
      softmax_scale_log2_e.y = mainloop_args.softmax_scale * log2_e;

      CUTLASS_PRAGMA_UNROLL
      for (int i = 0; i < size(tTR_rST); i += 2) {
        float2 acc;
        float2 lse;
        float2 out;
        acc.x = tTR_rST(i);
        acc.y = tTR_rST(i + 1);
        lse.x = sLSE(
            get<1>(tTR_cST(i)),
            pipeline_load_compute_lse_consumer_state.index());
        lse.y = sLSE(
            get<1>(tTR_cST(i + 1)),
            pipeline_load_compute_lse_consumer_state.index());
        cute::fma(out, softmax_scale_log2_e, acc, lse);
        tTR_rST(i) = ::exp2f(out.x);
        tTR_rST(i + 1) = ::exp2f(out.y);
      }

      auto tRT_rST = quantize(tTR_rST);
      auto tRT_rST_reshaped = make_tensor(tRT_rST.data(), shape(tRT_cST));

      cutlass::arch::fence_view_async_tmem_load();
      cutlass::arch::NamedBarrier(
          kNumComputeWarps * NumThreadsPerWarp,
          cutlass::arch::ReservedNamedBarriers::TransformBarrier)
          .arrive_and_wait();

      cute::copy(tiled_r2t, tRT_rST_reshaped, tRT_tP);

      // notify for P
      cutlass::arch::fence_view_async_tmem_store();
      pipeline_compute_mma_p.producer_commit(
          pipeline_compute_mma_p_producer_state);
      ++pipeline_compute_mma_p_producer_state;
      // release S
      pipeline_mma_compute_s.consumer_release(
          pipeline_mma_compute_s_consumer_state);
      ++pipeline_mma_compute_s_consumer_state;
      // release LSE
      pipeline_load_compute_lse.consumer_release(
          pipeline_load_compute_lse_consumer_state);
      ++pipeline_load_compute_lse_consumer_state;

      // wait for OdO
      pipeline_load_compute_sum_odo.consumer_wait(
          pipeline_load_compute_sum_odo_consumer_state);
      // wait for dP
      pipeline_mma_compute_dp.consumer_wait(
          pipeline_mma_compute_dp_consumer_state);

      // wait for dS
      // in principle, we could defer waiting for dS, and move in the freeing of
      // dP however, that would force us to keep dS in registers longer
      pipeline_compute_mma_ds.producer_acquire(
          pipeline_compute_mma_ds_producer_state);

      // compute dS = dsoftmax(P, dP, sum_OdO)
      cute::copy(tiled_t2r, tTR_tDPT, tTR_rDPT);

      CUTLASS_PRAGMA_UNROLL
      for (int i = 0; i < size(tTR_rDPT); i += 2) {
        float2 st;
        st.x = tTR_rST(i);
        st.y = tTR_rST(i + 1);
        float2 dpt;
        dpt.x = tTR_rDPT(i);
        dpt.y = tTR_rDPT(i + 1);
        float2 odo;
        odo.x = sSumOdO(
            get<1>(tTR_cDPT(i)),
            pipeline_load_compute_sum_odo_consumer_state.index());
        odo.y = sSumOdO(
            get<1>(tTR_cDPT(i + 1)),
            pipeline_load_compute_sum_odo_consumer_state.index());
        float2 dif;
        // sum odo is negated during preprocess
        cute::add(dif, dpt, odo);
        float2 out;
        cute::mul(out, dif, st);
        tTR_rDPT(i) = out.x;
        tTR_rDPT(i + 1) = out.y;
      }

      auto tTR_rDST = quantize(tTR_rDPT);

      // release dP
      cutlass::arch::fence_view_async_tmem_load();
      pipeline_mma_compute_dp.consumer_release(
          pipeline_mma_compute_dp_consumer_state);
      ++pipeline_mma_compute_dp_consumer_state;

      Tensor sDS = make_tensor(
          make_smem_ptr((Element*)shared_tensors.smem_ds.begin()),
          SmemLayoutDS{})(
          _, _, _, pipeline_compute_mma_ds_producer_state.index());

      auto thread_layout = make_ordered_layout(
          make_shape(_128{}, _128{}), make_stride(_1{}, _0{}));

      auto sDS_pi = as_position_independent_swizzle_tensor(sDS);
      auto sDS_pi_slice_p = sDS_pi.compose(thread_layout)(dp_idx, _).compose(
          make_layout(shape(tTR_cDPT_p)));
      auto sDS_pi_slice = split_wg(sDS_pi_slice_p);

      copy_aligned(tTR_rDST, sDS_pi_slice);

      // notify for dS
      cutlass::arch::fence_view_async_shared();
      pipeline_compute_mma_ds.producer_commit(
          pipeline_compute_mma_ds_producer_state);
      ++pipeline_compute_mma_ds_producer_state;
      // release OdO
      pipeline_load_compute_sum_odo.consumer_release(
          pipeline_load_compute_sum_odo_consumer_state);
      ++pipeline_load_compute_sum_odo_consumer_state;

      iter_count -= 1;
      iter_index += 1;
    }

    epilogue(
        blk_coord,
        problem_shape,
        mainloop_args,
        epilogue_args,
        pipeline_mma_compute_dkdv,
        pipeline_mma_compute_dkdv_consumer_state);
  }

  template <class BlkCoord, class IterToTileMap>
  CUTLASS_DEVICE void reduce(
      BlkCoord const& blk_coord,
      ProblemShape const& problem_shape,
      int iter_index,
      int iter_count,
      IterToTileMap const& iter_to_tile_map,
      MainloopArguments const& mainloop_args,
      MainloopParams const& mainloop_params,
      TensorStorage& shared_tensors,
      PipelineMmaReduceDQ& pipeline_mma_reduce_dq,
      typename PipelineMmaReduceDQ::PipelineState&
          pipeline_mma_reduce_dq_consumer_state,
      PipelineReduceTmaStore& pipeline_reduce_tma_store,
      typename PipelineReduceTmaStore::PipelineState&
          pipeline_reduce_tma_store_producer_state) {
    using X = Underscore;

    auto [Q, K, D, HB] = problem_shape;

    auto [blk_coord_q, blk_coord_k, blk_coord_batch] = blk_coord;

    // must match TileShapeDQ
    auto load_op = SM100_TMEM_LOAD_32dp32b32x{};

    auto tDQtDQ =
        partition_fragment_C(TiledMmaDSK{}, select<0, 1>(TileShapeDSK{}))(
            make_coord(_, _), _0{}, _0{});
    tDQtDQ.data() = TmemAllocation::kDQ;

    Tensor mDQ =
        mainloop_params.tma_red_dq.get_tma_tensor(make_shape(Q, D, HB));
    auto gDQ =
        local_tile(mDQ, TileShapeKQ{}, make_coord(_, _, _), Step<X, _1, _1>{})(
            _, _, _, _0{}, blk_coord_batch);

    Tensor cDQ = make_identity_tensor(take<0, 2>(TileShapeDSK{}));

    Tensor sDQ = make_tensor(
        make_smem_ptr(shared_tensors.smem_dq.begin()), SmemLayoutDQ{});

    int thread_idx = threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp);
    auto tiled_t2r = make_tmem_copy(load_op, tDQtDQ);
    auto thread_t2r = tiled_t2r.get_slice(thread_idx);

    Tensor tTR_cDQ = thread_t2r.partition_D(cDQ);
    Tensor tTR_gDQ = thread_t2r.partition_D(gDQ);
    Tensor tTR_sDQ = thread_t2r.partition_D(sDQ);
    Tensor tTR_tDQ = thread_t2r.partition_S(tDQtDQ);

    auto block_tma = mainloop_params.tma_red_dq.get_slice(_0{});

    Tensor tDQsDQ = block_tma.partition_S(sDQ);
    Tensor tDQcDQ = block_tma.partition_S(cDQ);
    Tensor tDQgDQ = block_tma.partition_D(gDQ);

    int lane_predicate =
        (threadIdx.x % (kNumReduceWarps * NumThreadsPerWarp)) == 0;

    while (iter_count > 0) {
      pipeline_mma_reduce_dq.consumer_wait(
          pipeline_mma_reduce_dq_consumer_state);

      Tensor tTR_rDQ = make_tensor<ElementAcc>(shape(tTR_cDQ));

      // load dQ from tmem to rmem
      cute::copy(tiled_t2r, tTR_tDQ, tTR_rDQ);

      cutlass::arch::fence_view_async_tmem_load();
      pipeline_mma_reduce_dq.consumer_release(
          pipeline_mma_reduce_dq_consumer_state);
      ++pipeline_mma_reduce_dq_consumer_state;

      // we don't have enough smem to dump it all to smem, so we do it in stages
      CUTLASS_PRAGMA_UNROLL
      for (int i = 0; i < size<2>(tTR_cDQ); i++) {
        if (lane_predicate) {
          pipeline_reduce_tma_store.producer_acquire(
              pipeline_reduce_tma_store_producer_state);
        }
        // wait in all threads for the acquire to complete
        cutlass::arch::NamedBarrier(
            kNumReduceWarps * NumThreadsPerWarp,
            cutlass::arch::ReservedNamedBarriers::TransposeBarrier)
            .arrive_and_wait();

        cute::copy(
            tTR_rDQ(_, _, i),
            tTR_sDQ(
                _, _, _0{}, pipeline_reduce_tma_store_producer_state.index()));

        // wait for the stores to all be visible to the TMA
        cutlass::arch::fence_view_async_shared();
        cutlass::arch::NamedBarrier(
            kNumReduceWarps * NumThreadsPerWarp,
            cutlass::arch::ReservedNamedBarriers::TransposeBarrier)
            .arrive_and_wait();
        if (lane_predicate) {
          // launch tma store
          copy(
              mainloop_params.tma_red_dq,
              tDQsDQ(
                  _, _, _0{}, pipeline_reduce_tma_store_producer_state.index()),
              tDQgDQ(_, _, i, iter_to_tile_map(iter_index)));
          pipeline_reduce_tma_store.producer_commit(
              pipeline_reduce_tma_store_producer_state);
        }

        ++pipeline_reduce_tma_store_producer_state;
      }

      iter_count -= 1;
      iter_index += 1;
    }
  }

  CUTLASS_DEVICE void operator()(Params const& params, char* smem) {
    int warp_idx = cutlass::canonical_warp_idx_sync();
    auto role = warp_idx_to_role(warp_idx);
    uint32_t lane_predicate = cute::elect_one_sync();

    if (role == WarpRole::Load && lane_predicate) {
      prefetch_tma_descriptor(
          params.mainloop_params.tma_load_q.get_tma_descriptor());
      prefetch_tma_descriptor(
          params.mainloop_params.tma_load_k.get_tma_descriptor());
      prefetch_tma_descriptor(
          params.mainloop_params.tma_load_v.get_tma_descriptor());
      prefetch_tma_descriptor(
          params.mainloop_params.tma_load_do.get_tma_descriptor());
    }

    SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem);

    int initializing_warp = 0;
    typename PipelineLoadMmaQ::Params pipeline_load_mma_q_params;
    if (role == WarpRole::Load) {
      pipeline_load_mma_q_params.role =
          PipelineLoadMmaQ::ThreadCategory::Producer;
    }
    if (role == WarpRole::Mma) {
      pipeline_load_mma_q_params.role =
          PipelineLoadMmaQ::ThreadCategory::Consumer;
    }
    pipeline_load_mma_q_params.is_leader =
        lane_predicate && (role == WarpRole::Load);
    // Also loads K in the first iteration
    pipeline_load_mma_q_params.transaction_bytes = kTransactionsBytesLoadQ;
    pipeline_load_mma_q_params.initializing_warp = initializing_warp++;
    PipelineLoadMmaQ pipeline_load_mma_q(
        shared_storage.pipelines.load_mma_q,
        pipeline_load_mma_q_params,
        ClusterShape{},
        /*barrier init*/ cute::true_type{},
        /*mask calc*/ cute::false_type{});

    typename PipelineLoadMmaDO::Params pipeline_load_mma_do_params;
    if (role == WarpRole::Load) {
      pipeline_load_mma_do_params.role =
          PipelineLoadMmaDO::ThreadCategory::Producer;
    }
    if (role == WarpRole::Mma) {
      pipeline_load_mma_do_params.role =
          PipelineLoadMmaDO::ThreadCategory::Consumer;
    }
    pipeline_load_mma_do_params.is_leader =
        lane_predicate && (role == WarpRole::Load);
    // Also loads V in the first iteration
    pipeline_load_mma_do_params.transaction_bytes = kTransactionsBytesLoadDO;
    pipeline_load_mma_do_params.initializing_warp = initializing_warp++;
    PipelineLoadMmaDO pipeline_load_mma_do(
        shared_storage.pipelines.load_mma_do,
        pipeline_load_mma_do_params,
        ClusterShape{},
        /*barrier init*/ cute::true_type{},
        /*mask calc*/ cute::false_type{});

    typename PipelineLoadComputeLSE::Params pipeline_load_compute_lse_params;
    if (role == WarpRole::Load) {
      pipeline_load_compute_lse_params.role =
          PipelineLoadComputeLSE::ThreadCategory::Producer;
    }
    if (role == WarpRole::Compute) {
      pipeline_load_compute_lse_params.role =
          PipelineLoadComputeLSE::ThreadCategory::Consumer;
    }
    pipeline_load_compute_lse_params.producer_arv_count = NumThreadsPerWarp;
    pipeline_load_compute_lse_params.consumer_arv_count =
        kNumComputeWarps * NumThreadsPerWarp;
    pipeline_load_compute_lse_params.initializing_warp = initializing_warp++;
    PipelineLoadComputeLSE pipeline_load_compute_lse(
        shared_storage.pipelines.load_compute_lse,
        pipeline_load_compute_lse_params,
        /*barrier init*/ cute::true_type{});

    typename PipelineLoadComputeSumOdO::Params
        pipeline_load_compute_sum_odo_params;
    if (role == WarpRole::Load) {
      pipeline_load_compute_sum_odo_params.role =
          PipelineLoadComputeSumOdO::ThreadCategory::Producer;
    }
    if (role == WarpRole::Compute) {
      pipeline_load_compute_sum_odo_params.role =
          PipelineLoadComputeSumOdO::ThreadCategory::Consumer;
    }
    pipeline_load_compute_sum_odo_params.producer_arv_count = NumThreadsPerWarp;
    pipeline_load_compute_sum_odo_params.consumer_arv_count =
        kNumComputeWarps * NumThreadsPerWarp;
    pipeline_load_compute_sum_odo_params.initializing_warp =
        initializing_warp++;
    PipelineLoadComputeSumOdO pipeline_load_compute_sum_odo(
        shared_storage.pipelines.load_compute_sum_odo,
        pipeline_load_compute_sum_odo_params,
        /*barrier init*/ cute::true_type{});

    typename PipelineMmaComputeS::Params pipeline_mma_compute_s_params;
    if (role == WarpRole::Mma) {
      pipeline_mma_compute_s_params.role =
          PipelineMmaComputeS::ThreadCategory::Producer;
    }
    if (role == WarpRole::Compute) {
      pipeline_mma_compute_s_params.role =
          PipelineMmaComputeS::ThreadCategory::Consumer;
    }
    pipeline_mma_compute_s_params.consumer_arv_count =
        kNumComputeWarps * cutlass::NumThreadsPerWarp;
    pipeline_mma_compute_s_params.initializing_warp = initializing_warp++;
    PipelineMmaComputeS pipeline_mma_compute_s(
        shared_storage.pipelines.mma_compute_s,
        pipeline_mma_compute_s_params,
        ClusterShape{},
        /*barrier init*/ cute::true_type{},
        /*mask calc*/ cute::false_type{});

    typename PipelineMmaComputeDP::Params pipeline_mma_compute_dp_params;
    if (role == WarpRole::Mma) {
      pipeline_mma_compute_dp_params.role =
          PipelineMmaComputeDP::ThreadCategory::Producer;
    }
    if (role == WarpRole::Compute) {
      pipeline_mma_compute_dp_params.role =
          PipelineMmaComputeDP::ThreadCategory::Consumer;
    }
    pipeline_mma_compute_dp_params.consumer_arv_count =
        kNumComputeWarps * cutlass::NumThreadsPerWarp;
    pipeline_mma_compute_dp_params.initializing_warp = initializing_warp++;
    PipelineMmaComputeDP pipeline_mma_compute_dp(
        shared_storage.pipelines.mma_compute_dp,
        pipeline_mma_compute_dp_params,
        ClusterShape{},
        /*barrier init*/ cute::true_type{},
        /*mask calc*/ cute::false_type{});

    typename PipelineMmaReduceDQ::Params pipeline_mma_reduce_dq_params;
    if (role == WarpRole::Mma) {
      pipeline_mma_reduce_dq_params.role =
          PipelineMmaReduceDQ::ThreadCategory::Producer;
    }
    if (role == WarpRole::Reduce) {
      pipeline_mma_reduce_dq_params.role =
          PipelineMmaReduceDQ::ThreadCategory::Consumer;
    }
    pipeline_mma_reduce_dq_params.consumer_arv_count =
        kNumReduceWarps * cutlass::NumThreadsPerWarp;
    pipeline_mma_reduce_dq_params.initializing_warp = initializing_warp++;
    PipelineMmaReduceDQ pipeline_mma_reduce_dq(
        shared_storage.pipelines.mma_reduce_dq,
        pipeline_mma_reduce_dq_params,
        ClusterShape{},
        /*barrier init*/ cute::true_type{},
        /*mask calc*/ cute::false_type{});

    typename PipelineComputeMmaP::Params pipeline_compute_mma_p_params;
    if (role == WarpRole::Mma) {
      pipeline_compute_mma_p_params.role =
          PipelineComputeMmaP::ThreadCategory::Consumer;
    }
    if (role == WarpRole::Compute) {
      pipeline_compute_mma_p_params.role =
          PipelineComputeMmaP::ThreadCategory::Producer;
    }
    pipeline_compute_mma_p_params.producer_arv_count =
        kNumComputeWarps * cutlass::NumThreadsPerWarp;
    pipeline_compute_mma_p_params.consumer_arv_count = 1;
    pipeline_compute_mma_p_params.initializing_warp = initializing_warp++;
    PipelineComputeMmaP pipeline_compute_mma_p(
        shared_storage.pipelines.compute_mma_p,
        pipeline_compute_mma_p_params,
        ClusterShape{},
        /*barrier init*/ cute::true_type{},
        /*mask calc*/ cute::false_type{});

    typename PipelineComputeMmaDS::Params pipeline_compute_mma_ds_params;
    if (role == WarpRole::Mma) {
      pipeline_compute_mma_ds_params.role =
          PipelineComputeMmaDS::ThreadCategory::Consumer;
    }
    if (role == WarpRole::Compute) {
      pipeline_compute_mma_ds_params.role =
          PipelineComputeMmaDS::ThreadCategory::Producer;
    }
    pipeline_compute_mma_ds_params.producer_arv_count =
        kNumComputeWarps * cutlass::NumThreadsPerWarp;
    pipeline_compute_mma_ds_params.consumer_arv_count = 1;
    pipeline_compute_mma_ds_params.initializing_warp = initializing_warp++;
    PipelineComputeMmaDS pipeline_compute_mma_ds(
        shared_storage.pipelines.compute_mma_ds,
        pipeline_compute_mma_ds_params,
        ClusterShape{},
        /*barrier init*/ cute::true_type{},
        /*mask calc*/ cute::false_type{});

    typename PipelineMmaComputeDKDV::Params pipeline_mma_compute_dkdv_params;
    if (role == WarpRole::Mma) {
      pipeline_mma_compute_dkdv_params.role =
          PipelineMmaComputeDKDV::ThreadCategory::Producer;
    }
    if (role == WarpRole::Compute) {
      pipeline_mma_compute_dkdv_params.role =
          PipelineMmaComputeDKDV::ThreadCategory::Consumer;
    }
    pipeline_mma_compute_dkdv_params.consumer_arv_count =
        kNumComputeWarps * cutlass::NumThreadsPerWarp;
    pipeline_mma_compute_dkdv_params.initializing_warp = initializing_warp++;
    PipelineMmaComputeDKDV pipeline_mma_compute_dkdv(
        shared_storage.pipelines.mma_compute_dkdv,
        pipeline_mma_compute_dkdv_params,
        ClusterShape{},
        /*barrier init*/ cute::true_type{},
        /*mask calc*/ cute::false_type{});
    PipelineReduceTmaStore pipeline_reduce_tma_store;

    TmemAllocator tmem_allocator;

    pipeline_init_arrive_relaxed(size(ClusterShape{}));

    pipeline_load_mma_q.init_masks(ClusterShape{});
    pipeline_load_mma_do.init_masks(ClusterShape{});
    pipeline_mma_compute_s.init_masks(ClusterShape{});
    pipeline_mma_compute_dp.init_masks(ClusterShape{});
    pipeline_mma_reduce_dq.init_masks(ClusterShape{});
    pipeline_compute_mma_p.init_masks(ClusterShape{});
    pipeline_compute_mma_ds.init_masks(ClusterShape{});
    pipeline_mma_compute_dkdv.init_masks(ClusterShape{});

    typename decltype(pipeline_load_mma_q)::PipelineState
        pipeline_load_mma_q_consumer_state;
    typename decltype(pipeline_load_mma_do)::PipelineState
        pipeline_load_mma_do_consumer_state;
    typename decltype(pipeline_load_compute_lse)::PipelineState
        pipeline_load_compute_lse_consumer_state;
    typename decltype(pipeline_load_compute_sum_odo)::PipelineState
        pipeline_load_compute_sum_odo_consumer_state;
    typename decltype(pipeline_mma_compute_s)::PipelineState
        pipeline_mma_compute_s_consumer_state;
    typename decltype(pipeline_mma_compute_dp)::PipelineState
        pipeline_mma_compute_dp_consumer_state;
    typename decltype(pipeline_mma_reduce_dq)::PipelineState
        pipeline_mma_reduce_dq_consumer_state;
    typename decltype(pipeline_compute_mma_p)::PipelineState
        pipeline_compute_mma_p_consumer_state;
    typename decltype(pipeline_compute_mma_ds)::PipelineState
        pipeline_compute_mma_ds_consumer_state;
    typename decltype(pipeline_mma_compute_dkdv)::PipelineState
        pipeline_mma_compute_dkdv_consumer_state;

    auto pipeline_load_mma_q_producer_state =
        make_producer_start_state<decltype(pipeline_load_mma_q)>();
    auto pipeline_load_mma_do_producer_state =
        make_producer_start_state<decltype(pipeline_load_mma_do)>();
    auto pipeline_load_compute_lse_producer_state =
        make_producer_start_state<decltype(pipeline_load_compute_lse)>();
    auto pipeline_load_compute_sum_odo_producer_state =
        make_producer_start_state<decltype(pipeline_load_compute_sum_odo)>();
    auto pipeline_mma_compute_s_producer_state =
        make_producer_start_state<decltype(pipeline_mma_compute_s)>();
    auto pipeline_mma_compute_dp_producer_state =
        make_producer_start_state<decltype(pipeline_mma_compute_dp)>();
    auto pipeline_mma_reduce_dq_producer_state =
        make_producer_start_state<decltype(pipeline_mma_reduce_dq)>();
    auto pipeline_compute_mma_p_producer_state =
        make_producer_start_state<decltype(pipeline_compute_mma_p)>();
    auto pipeline_compute_mma_ds_producer_state =
        make_producer_start_state<decltype(pipeline_compute_mma_ds)>();
    auto pipeline_mma_compute_dkdv_producer_state =
        make_producer_start_state<decltype(pipeline_mma_compute_dkdv)>();
    auto pipeline_reduce_tma_store_producer_state =
        make_producer_start_state<decltype(pipeline_reduce_tma_store)>();

    pipeline_init_wait(size(ClusterShape{}));

    auto blk_coord =
        make_coord(_0{}, blockIdx.x, make_coord(blockIdx.y, blockIdx.z));
    auto problem_shape = params.problem_shape;

    // FNA
    auto qkv_shape = params.fna.qkv_shape;
    bool is_fully_block_sparse = params.fna.is_fully_block_sparse;
    bool has_q_padding = params.fna.has_q_padding;
    if (params.fna.requires_qkv_fixup) {
      qkv_shape = Mask{}.correct_qkv_shape(
          problem_shape,
          params.fna.qkv_shape,
          blk_coord,
          params.fna.dilation,
          params.fna.num_heads_actual);
      is_fully_block_sparse = fully_block_sparse<typename Mask::Causal>(
          qkv_shape,
          get<0>(params.fna.na_params),
          get<3>(params.fna.na_params),
          QTileShape{},
          KVTileShape{});
      has_q_padding = not evenly_divides(qkv_shape, QTileShape{});
    } else if (params.fna.is_dilated) {
      qkv_shape = ceil_div(params.fna.qkv_shape, params.fna.dilation);
      is_fully_block_sparse = fully_block_sparse<typename Mask::Causal>(
          qkv_shape,
          get<0>(params.fna.na_params),
          get<3>(params.fna.na_params),
          QTileShape{},
          KVTileShape{});
      has_q_padding = not evenly_divides(qkv_shape, QTileShape{});
    }

    auto [q_start, num_q_tiles] = Mask{}.get_trip_count(
        blk_coord,
        MultiDimTileShape{},
        params.fna.kv_shape,
        qkv_shape,
        params.fna.na_params);

    int iter_start = 0;
    int iter_count = size(num_q_tiles);

    auto q_start_tile = ceil_div(q_start, QTileShape{});

    auto q_tiled = ceil_div(params.fna.q_shape, QTileShape{});
    auto ctr = make_identity_tensor(num_q_tiles);
    auto ctr_offset = domain_offset(q_start_tile, ctr);

    auto q_tiled_layout = make_layout(q_tiled);

    auto iter_to_tile_map = [&ctr_offset, &q_tiled_layout](int iter) {
      return crd2idx(ctr_offset(iter), q_tiled_layout);
    };

    // Dilated cases where dilation doesn't evenly divide input shape can
    // sometimes lead to threadblocks with no work to do (similar to varlen). We
    // should exit those immediately instead of letting them proceed because
    // some of the roles (i.e. load) unroll their loops, and this can cause
    // stalls.
    if (iter_count == 0) {
      return;
    }
    // End of FNA

    if (role == WarpRole::Load) {
      warpgroup_reg_set<RegisterAllocation::kLoad>();

      load(
          blk_coord,
          problem_shape,
          iter_start,
          iter_count,
          iter_to_tile_map,
          params.mainloop,
          params.mainloop_params,
          shared_storage.tensors,
          pipeline_load_mma_q,
          pipeline_load_mma_q_producer_state,
          pipeline_load_mma_do,
          pipeline_load_mma_do_producer_state,
          pipeline_load_compute_lse,
          pipeline_load_compute_lse_producer_state,
          pipeline_load_compute_sum_odo,
          pipeline_load_compute_sum_odo_producer_state);

    } else if (role == WarpRole::Mma) {
      warpgroup_reg_set<RegisterAllocation::kMma>();

      tmem_allocator.allocate(
          TmemAllocator::Sm100TmemCapacityColumns,
          &shared_storage.tmem_base_ptr);
      __syncwarp();

      mma(blk_coord,
          problem_shape,
          iter_start,
          iter_count,
          params.mainloop,
          shared_storage.tensors,
          pipeline_load_mma_q,
          pipeline_load_mma_q_consumer_state,
          pipeline_load_mma_do,
          pipeline_load_mma_do_consumer_state,
          pipeline_mma_compute_s,
          pipeline_mma_compute_s_producer_state,
          pipeline_mma_compute_dp,
          pipeline_mma_compute_dp_producer_state,
          pipeline_mma_reduce_dq,
          pipeline_mma_reduce_dq_producer_state,
          pipeline_compute_mma_p,
          pipeline_compute_mma_p_consumer_state,
          pipeline_compute_mma_ds,
          pipeline_compute_mma_ds_consumer_state,
          pipeline_mma_compute_dkdv,
          pipeline_mma_compute_dkdv_producer_state);

    } else if (role == WarpRole::Compute) {
      warpgroup_reg_set<RegisterAllocation::kCompute>();

      compute(
          blk_coord,
          problem_shape,
          iter_start,
          iter_count,
          // FNA params
          params.fna.kv_shape,
          qkv_shape,
          params.fna.na_params,
          q_start,
          num_q_tiles,
          is_fully_block_sparse,
          has_q_padding,
          //
          params.mainloop,
          params.epilogue,
          shared_storage.tensors,
          pipeline_load_compute_lse,
          pipeline_load_compute_lse_consumer_state,
          pipeline_load_compute_sum_odo,
          pipeline_load_compute_sum_odo_consumer_state,
          pipeline_mma_compute_s,
          pipeline_mma_compute_s_consumer_state,
          pipeline_mma_compute_dp,
          pipeline_mma_compute_dp_consumer_state,
          pipeline_compute_mma_p,
          pipeline_compute_mma_p_producer_state,
          pipeline_compute_mma_ds,
          pipeline_compute_mma_ds_producer_state,
          pipeline_mma_compute_dkdv,
          pipeline_mma_compute_dkdv_consumer_state);

      cutlass::arch::NamedBarrier(
          kNumComputeWarps * NumThreadsPerWarp,
          cutlass::arch::ReservedNamedBarriers::EpilogueBarrier)
          .arrive_and_wait();

      if (warp_idx % kNumComputeWarps == 0) {
        uint32_t free_stage_ptr = shared_storage.tmem_base_ptr;
        tmem_allocator.free(
            free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
      }

    } else if (role == WarpRole::Reduce) {
      warpgroup_reg_set<RegisterAllocation::kReduce>();

      reduce(
          blk_coord,
          problem_shape,
          iter_start,
          iter_count,
          iter_to_tile_map,
          params.mainloop,
          params.mainloop_params,
          shared_storage.tensors,
          pipeline_mma_reduce_dq,
          pipeline_mma_reduce_dq_consumer_state,
          pipeline_reduce_tma_store,
          pipeline_reduce_tma_store_producer_state);

      pipeline_reduce_tma_store.producer_tail(
          pipeline_reduce_tma_store_producer_state);
    } else {
      warpgroup_reg_set<RegisterAllocation::kEmpty>();

      /* no-op */
    }
  }

  static dim3 get_block_shape() {
    dim3 block(MaxThreadsPerBlock, 1, 1);
    return block;
  }

  static dim3 get_grid_shape(Params const& params) {
    auto [Q, K, D, HB] = params.problem_shape;
    auto [H, B] = HB;
    dim3 grid(ceil_div(K, TileShapeK{}), H, B);
    return grid;
  }
};

} // namespace cutlass::fna::kernel
