#pragma once
#include <type_traits>
#include <cstdio>
#include <tuple>
#include "../../cuda_utils/frg_copy.cuh"
#include "../../cuda_utils/pipeline.cuh"
#include "../../cuda_utils/gemm.cuh"
#include "../../cuda_utils/copy.cuh"
#include "../../cuda_utils/perf_cfg.cuh"
#include "../../cuda_utils/allocator.cuh"
#include "../../cuda_utils/launch.cuh"
#include "../../cuda_utils/mask_cfg.cuh"
#include "../../cuda_utils/softmax.cuh"
#include "../../kernels/mma/mma_cfg.cuh"
#include "../../kernels/reduce/reduce_utils.cuh"

namespace vidrial {

using namespace cute;

#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))

/**
 * @brief FlashKernelCfg is a configuration for a flash attention kernel.
 * @tparam MmaCfg1 The configuration for the first mma
 * @tparam MmaCfg2 The configuration for the second mma
 * @tparam MaskCfg The configuration for the mask
 * @tparam SoftmaxCfg The configuration for the softmax
 */
template <typename LT_, typename MmaCfg1, typename MmaCfg2, typename MaskCfg, typename SoftmaxCfg, typename GLSlab>
struct FlashKernelCfg {
    using T = typename MmaCfg1::T;
    using LT = LT_;
    using Mma1_C_frg__2__Mma2_A_frg = decltype(left_inverse(get<0>(MmaCfg1{}.C.mma_FrgThr)).compose(get<0>(MmaCfg2{}.A.mma_FrgThr)));
    using PerfCfg1 = typename MmaCfg1::PerfCfg;
    using PerfCfg2 = typename MmaCfg2::PerfCfg;

    static constexpr PerfCfg1 perf1{};
    static constexpr PerfCfg2 perf2{};
    static constexpr bool q_in_reg = PerfCfg1::q_in_reg;
    static constexpr int threads = size(MmaCfg1::thread_num);
    MmaCfg1 mma1;
    MmaCfg2 mma2;
    MaskCfg mask;
    Mma1_C_frg__2__Mma2_A_frg mma1_C_frg__2__mma2_A_frg;
    SoftmaxCfg softmax;
    GLSlab gLSlab;

    static int smem_size() {
        using A1T = typename MmaCfg1::T;
        using B1T = typename MmaCfg1::T;
        using B2T = typename MmaCfg2::T;
        using C2T = typename MmaCfg2::T;
        using A1_STile = typename MmaCfg1::ACfg::STile;
        using B1_STile = typename MmaCfg1::BCfg::STile;
        using B2_STile = typename MmaCfg2::BCfg::STile;
        using C2_STile = typename MmaCfg2::CCfg::STile;

        int a_prefetch_size = Allocator<16>::total<A1T>(size(A1_STile{}) * MmaCfg1::K_tile_num);

        int abb_smem_size = Allocator<16>::total<A1T, B1T, B2T>(
            size(A1_STile{}) * (q_in_reg ? 0 : perf1.smempipe),
            size(B1_STile{}) * perf1.smempipe,
            size(B2_STile{}) * perf2.smempipe
        );
        int c_smem_size = Allocator<16>::total<C2T>(size(C2_STile{}));
        int l_smem_size = Allocator<16>::total<LT>(size<0>(C2_STile{}));
        if constexpr (q_in_reg) {
            return static_max(a_prefetch_size, static_max(abb_smem_size, c_smem_size + l_smem_size));
        } else {
            return static_max(abb_smem_size, c_smem_size + l_smem_size);
        }
    }
};

/**
 * @brief Guess the missing slab from the given slabs using canonical stride. All slabs conform to the following shape convention:
 *        A: (M, K, P)
 *        B: (N, K, P)
 *        C: (M, N, P)
 *        - P is the batch mode
 * @tparam ASlab The slab of the first mma, or void indicating A is missing
 * @tparam BSlab The slab of the second mma, or void indicating B is missing
 * @tparam CSlab The slab of the third mma, or void indicating C is missing
 * @return The missing slab
 */
template <typename ASlab, typename BSlab, typename CSlab, typename Stride_t=LayoutLeft>
constexpr auto make_missing_slab(Stride_t) {
    if constexpr (is_void_v<ASlab>) { // A is missing
        constexpr int M = size<0>(CSlab{});
        constexpr int K = size<1>(BSlab{});
        constexpr int P = size<2>(CSlab{});
        return static_tree_cast<int64_t>(make_layout(static_tree_cast<int64_t>(Shape<Int<M>, Int<K>, Int<P>>{}), Stride_t{}));
    } else if constexpr (is_void_v<BSlab>) { // B is missing
        constexpr int N = size<1>(CSlab{});
        constexpr int K = size<1>(ASlab{});
        constexpr int P = size<2>(CSlab{});
        return static_tree_cast<int64_t>(make_layout(static_tree_cast<int64_t>(Shape<Int<N>, Int<K>, Int<P>>{}), Stride_t{}));
    } else { // C is missing
        constexpr int M = size<0>(ASlab{});
        constexpr int N = size<0>(BSlab{});
        constexpr int P = size<2>(ASlab{});
        return static_tree_cast<int64_t>(make_layout(static_tree_cast<int64_t>(Shape<Int<M>, Int<N>, Int<P>>{}), Stride_t{}));
    }
}

/**
 * @brief Computes the mask boundary for N dimension
 * @return The start of the causal region in N dimension and the size of the causal region, in unit of N_tile
 */
template <typename FlashKernelCfg, typename MmaMNKCoords, int r=1, int w=1>
CUTE_HOST_DEVICE auto causal_region(FlashKernelCfg &cfg, MmaMNKCoords &tile_coords) {
    constexpr auto offset = cfg.mma1.N / w - cfg.mma1.M / r;
    static_assert(offset >= 0, "Offset must be non-negative for causal mask");
    constexpr int N_tile = cfg.mma1.N_tile;
    constexpr int M_tile = cfg.mma1.M_tile;
    auto M_idx = tile_coords.M_coord() * M_tile;
    auto start = (M_idx / r + offset) / (N_tile / w);
    auto end = CEIL_DIV((M_idx + M_tile) / r + offset, (N_tile / w));
    return std::make_pair(start, end);
}

/**
 * @brief Compute the atom placement for the second mma
 */
template <typename MmaCfg1, typename _MmaAtom2>
constexpr auto get_atom_placement_MNK() {
    constexpr auto cfg = MmaCfg1{};
    using Atom1Shape_t = typename decltype(cfg.mma_Atom)::Shape_MNK;
    using Atom2Shape_t = typename _MmaAtom2::Shape_MNK;
    using Placement1_t = typename MmaCfg1::MNKAtomPlacement;
    constexpr int total_atoms = cfg.thread_num / size(typename _MmaAtom2::ThrID{});
    constexpr int M_placement_1 = size<0>(Placement1_t{});
    constexpr int N_placement_1 = size<1>(Placement1_t{});
    constexpr int K_placement_1 = size<2>(Placement1_t{});
    constexpr int M_placement_2 = M_placement_1 * size<0>(Atom1Shape_t{}) / size<0>(Atom2Shape_t{});
    constexpr int K_placement_2 = N_placement_1 * size<1>(Atom1Shape_t{}) / size<1>(Atom2Shape_t{});
    constexpr int N_placement_2 = total_atoms / (M_placement_2 * K_placement_2);
    static_assert(N_placement_2 >= 1, "N_placement_2 must be at least 1");
    static_assert(K_placement_2 >= 1, "K_placement_2 must be at least 1");
    return make_tile(Int<M_placement_2>{}, Int<N_placement_2>{}, Int<K_placement_2>{});
}

/**
 * @brief Compute the permutation required to make the second mma's B tile's K dimension match the first mma's C atom layout. Note that this function doesn't check inter-atom frgthr layout compatibility.
 */
template <typename MmaCfg1, typename MmaAtom2, typename TileMNK2_t, typename Placement2_t>
constexpr auto get_permutation_MNK() {
    constexpr auto cfg = MmaCfg1{};
    using Atom1Shape_t = typename decltype(cfg.mma_Atom)::Shape_MNK;
    using Atom2Shape_t = typename MmaAtom2::Shape_MNK;
    using Placement1_t = typename MmaCfg1::MNKAtomPlacement;
    constexpr auto AtomShape1_C = select<0, 1>(Atom1Shape_t{});
    constexpr auto AtomShape2_A = select<0, 2>(Atom2Shape_t{});
    // --- first mma's C ---
    constexpr auto OpShape1_tiled_C = make_tile(AtomShape1_C, make_tile(get<0>(Placement1_t{}), get<1>(Placement1_t{}))); // [[M_atom, N_atom], [M_num_atom, N_num_atom]]
    constexpr auto OpShape2_tiled_A = make_tile(AtomShape2_A, make_tile(get<0>(Placement2_t{}), get<2>(Placement2_t{}))); // [[M_atom, K_atom], [M_num_atom, K_num_atom]]
    constexpr auto OpShape1_C = zip(OpShape1_tiled_C); // [[M_atom, M_num_atom], [N_atom, N_num_atom]]
    constexpr auto OpShape1_C_shape = transform(OpShape1_C, [](auto x) {
        return size(x);
    });
    constexpr auto Rest1_C = get<1>(zipped_divide(make_layout(cfg.C.tileShape), OpShape1_C_shape).shape()); // [M_rest, N_rest]
    constexpr auto Tile1_C = make_tile(
        make_tile(get<0, 0>(OpShape1_C), get<0, 1>(OpShape1_C), get<0>(Rest1_C)),
        make_tile(get<1, 0>(OpShape1_C), get<1, 1>(OpShape1_C), get<1>(Rest1_C))
    );

    // --- second mma's A ---
    constexpr auto OpShape2_A = zip(OpShape2_tiled_A); // [[M_atom, M_num_atom], [K_atom, K_num_atom]]
    constexpr auto OpShape2_A_shape = transform(OpShape2_A, [](auto x) {
        return size(x);
    });
    constexpr auto TileShape2_A = select<0, 2>(TileMNK2_t{});
    constexpr auto Rest2_A = get<1>(zipped_divide(make_layout(TileShape2_A), OpShape2_A_shape).shape()); // [M_rest, K_rest]

    constexpr auto Tile2_A = make_tile(
        make_tile(get<0, 0>(OpShape2_A), get<0, 1>(OpShape2_A), get<0>(Rest2_A)),
        make_tile(get<1, 0>(OpShape2_A), get<1, 1>(OpShape2_A), get<1>(Rest2_A))
    );

    if constexpr (is_same_v<decltype(Rest1_C), decltype(Rest2_A)>) {
        return make_tile(_, _, _);
    } else {
        CUTE_STATIC_ASSERT_V(size<0>(Tile1_C) == size<0>(Tile2_A), "M size mismatch");
        CUTE_STATIC_ASSERT_V(size<1>(Tile1_C) == size<1>(Tile2_A), "N and K size mismatch");
        constexpr auto m1 = select<0, 2, 1>(make_layout(get<0>(Tile1_C))); // [M_atom, M_rest, M_num_atom]
        constexpr auto n1 = select<0, 2, 1>(make_layout(get<1>(Tile1_C))); // [N_atom, N_rest, N_num_atom]
        constexpr auto m2 = select<0, 2, 1>(make_layout(get<0>(Tile2_A))); // [M_atom, M_rest, M_num_atom]
        constexpr auto k2 = select<0, 2, 1>(make_layout(get<1>(Tile2_A))); // [K_atom, K_rest, K_num_atom]
        using m_perm_t = decltype(left_inverse(m2).compose(m1));
        using k_perm_t = decltype(left_inverse(k2).compose(n1));
        return make_tile(m_perm_t{}, _, k_perm_t{});
    }

}

template <typename T, typename LT, typename MmaAtom1, typename MNKTileShape1, typename MNKAtomPlacement1, typename MmaAtom2, typename NTile2, typename GQSlab, typename GKSlab, typename GVSlab, typename GOSlab, typename GLSlab, typename Perf1Cfg = DefaultPerfCfg, typename Perf2Cfg = DefaultPerfCfg>
auto make_FlashKernelCfg() {
    constexpr int M_TILE_1 = size<0>(MNKTileShape1{});
    constexpr int N_TILE_1 = size<1>(MNKTileShape1{});
    using MNKTileShape2 = Shape<Int<M_TILE_1>, NTile2, Int<N_TILE_1>>;

    using FakeSSlab = decltype(make_missing_slab<GQSlab, GKSlab, void>(LayoutLeft{}));
    constexpr auto mma1 = MmaKernelCfg<T, MmaAtom1, MNKTileShape1, MNKAtomPlacement1, GQSlab, GKSlab, FakeSSlab, Perf1Cfg>();

    // ---- Compute the atom placement for the second mma ----
    constexpr auto atom_placement_2 = get_atom_placement_MNK<decltype(mma1), MmaAtom2>();

    // ---- Compute the permutation for the second mma ----
    using permutation_2_t = decltype(get_permutation_MNK<decltype(mma1), MmaAtom2, MNKTileShape2, decltype(atom_placement_2)>());
    auto permutation_2 = get_permutation_MNK<decltype(mma1), MmaAtom2, MNKTileShape2, decltype(atom_placement_2)>();

    // ---- Construct the second mma cfg ----
    using FakeASlab = decltype(make_missing_slab<void, GVSlab, GOSlab>(LayoutLeft{}));
    constexpr auto mma2 = MmaKernelCfg<T, MmaAtom2, MNKTileShape2, decltype(atom_placement_2), FakeASlab, GVSlab, GOSlab, Perf2Cfg, permutation_2_t>{};
    static_assert(size<0>(mma1.MNK_tile_shape) == size<0>(mma2.MNK_tile_shape), "M tile mismatch");
    static_assert(size<1>(mma1.MNK_tile_shape) == size<2>(mma2.MNK_tile_shape), "N/K tile mismatch");
    static_assert(mma1.M == mma2.M, "M mismatch");
    static_assert(mma1.N == mma2.K, "N/K mismatch");
    static_assert(mma1.perf.smempipe / mma1.K_tile_num == mma2.perf.smempipe, "smempipe mismatch");

    // ---- Construct the mask cfg ----
    constexpr auto mask_cfg = make_mask_cfg<T, MNKTileShape1, decltype(mma1.C.mma_FrgThr), CausalMask, OpNegInf>();

    // ---- Construct the softmax cfg ----
    using softmax_cfg_t = decltype(make_SoftmaxCfg<LT, decltype(mma1.C.sTile.shape()), decltype(mma1.C.mma_FrgThr), decltype(mma2.C.sTile.shape()), decltype(mma2.C.mma_FrgThr), static_cast<LT>(1.0), true>());  

    return FlashKernelCfg<LT, decltype(mma1), decltype(mma2), decltype(mask_cfg), softmax_cfg_t, GLSlab>{};
}

template<typename Cfg>
CUTE_HOST_DEVICE auto make_rQ_frg(Cfg const& cfg) {
    using rQ_frg_t = decltype(cfg.mma1.A.make_mma_frg().layout());
    constexpr auto rQ_frg_all_layout = zipped_product(rQ_frg_t{}, Int<cfg.mma1.K_tile_num>{});
    return make_tensor<Cfg::T>(rQ_frg_all_layout);
}


// #define DEBUG 0
#define DEBUG_THREAD (threadIdx.x == 1 && blockIdx.x == 2 && blockIdx.y == 0 && blockIdx.z == 0)

#ifdef DEBUG
#define DEBUG_PRINT_TENSOR(msg) \
    __syncthreads(); \
    store(S_copy_cfg, rS_frg_mma, sS_tile); \
    __syncthreads(); \
    if (DEBUG_THREAD) { \
        print(msg); print_tensor(sS_tile); print("\n"); \
    } \
    __syncthreads()
#else
#define DEBUG_PRINT_TENSOR(msg)
#endif

#ifdef DEBUG
#define DEBUG_PRINT_ROWSUM(msg) \
    __syncthreads(); \
    auto rowsum_frg_copy = make_tensor<LT>(rowsum_frg.layout()); \
    copy(rowsum_frg, rowsum_frg_copy); \
    warp_reduce(cfg.softmax.reduce.warp_reduce, SumOp<LT>{}, rowsum_frg_copy); \
    smart_reduce_cta_all_threads(cfg.softmax.reduce, rowsum_frg_copy, sRow_tile, SumOp<LT>{}); \
    store(R_copy_cfg, rowsum_frg_copy, sRow_tile); \
    __syncthreads(); \
    if (DEBUG_THREAD) { \
        print(msg); print_tensor(sRow_tile); print("\n"); \
    } \
    __syncthreads()
#else
#define DEBUG_PRINT_ROWSUM(msg)
#endif

template <typename Cfg, typename T, typename LT>
__global__ void __launch_bounds__(Cfg::threads, 1) flash_kernel(Cfg const cfg, T* Q_ptr, T* K_ptr, T* V_ptr, T* O_ptr, LT* l_ptr) {
    int tid = threadIdx.x;
    int bid_M = blockIdx.x; int bid_N = blockIdx.y; int bid_P = blockIdx.z;
    auto coord_qk = MmaMNKCoords(cfg.mma1.MNK_tile_shape);
    auto coord_pv = MmaMNKCoords(cfg.mma2.MNK_tile_shape);
    coord_qk.step_M(bid_M); coord_qk.step_P(bid_P);
    coord_pv.step_M(bid_M); coord_pv.step_N(bid_N); coord_pv.step_P(bid_P);
    auto [mask_start, mask_end] = causal_region(cfg, coord_qk);
    coord_qk.step_N(mask_start); coord_pv.step_K(mask_start);
    // ----- Global memory slabs -----
    auto gQ_slab = make_tensor(make_gmem_ptr(Q_ptr), cfg.mma1.A.gSlab);
    auto gK_slab = make_tensor(make_gmem_ptr(K_ptr), cfg.mma1.B.gSlab);
    auto gV_slab = make_tensor(make_gmem_ptr(V_ptr), cfg.mma2.B.gSlab);
    auto gO_slab = make_tensor(make_gmem_ptr(O_ptr), cfg.mma2.C.gSlab);
    auto gl_slab = make_tensor(make_gmem_ptr(l_ptr), cfg.gLSlab);
    // ----- Shared memory pipelines -----
    constexpr static int smempipe_qk = cfg.mma1.perf.smempipe;
    constexpr static int smempipe_pv = cfg.mma2.perf.smempipe;
    auto pipe_qk = SmemPipe<smempipe_qk>();
    auto pipe_pv = SmemPipe<smempipe_pv>();
    extern __shared__ char smem[];
    Allocator<16> alloc(smem);
    T* Q_smem = alloc.allocate<T>(size(cfg.mma1.A.sTile) * (cfg.q_in_reg ? cfg.mma1.K_tile_num : smempipe_qk));
    if constexpr (cfg.q_in_reg) // Q and the rest can share the same smem if Q is in reg
        alloc.reset(smem);
    T* K_smem = alloc.allocate<T>(size(cfg.mma1.B.sTile) * smempipe_qk);
    T* V_smem = alloc.allocate<T>(size(cfg.mma2.B.sTile) * smempipe_pv);
    LT* Reduce_smem = nullptr;
    if constexpr (cfg.softmax.reduce.warp_reduce_size > 1)
        Reduce_smem = alloc.allocate<LT>(size(cfg.softmax.reduce.x.shape));
    auto sRow = make_tensor(make_smem_ptr(Reduce_smem), cfg.softmax.reduce.x.shape);
#ifdef DEBUG
    LT* Debug_S_smem = alloc.allocate<LT>(cfg.mma1.M_tile * cfg.mma1.N_tile);
    LT* Debug_Rowmax_smem = alloc.allocate<LT>(cfg.mma1.M_tile);
#endif
    auto sQ_tile_pipe = [&]() {
        if constexpr (cfg.q_in_reg) {
            auto pipe_q = SmemPipe<cfg.mma1.K_tile_num>();
            return pipe_q.create(Q_smem, cfg.mma1.A.sTile, cfg.mma1.A.tile_copy);
        } else {
            return pipe_qk.create(Q_smem, cfg.mma1.A.sTile, cfg.mma1.A.tile_copy);
        }
    }();
    auto sK_tile_pipe = pipe_qk.create(K_smem, cfg.mma1.B.sTile, cfg.mma1.B.tile_copy);
    auto sV_tile_pipe = pipe_pv.create(V_smem, cfg.mma2.B.sTile, cfg.mma2.B.tile_copy);
    decltype(make_rQ_frg(cfg)) rQ_frg;
    if constexpr (cfg.q_in_reg)
        rQ_frg = make_rQ_frg(cfg);
    auto rS_frg_mma = cfg.mma1.C.make_mma_frg();
    auto rP_frg_mma = rS_frg_mma.compose(cfg.mma1_C_frg__2__mma2_A_frg);
    auto rO_frg_mma = cfg.mma2.C.make_mma_frg();
    clear(rO_frg_mma);
#ifdef DEBUG
    auto sS_tile = make_tensor(make_smem_ptr(Debug_S_smem), Shape<Int<cfg.mma1.M_tile>, Int<cfg.mma1.N_tile>>{});
    auto sRow_tile = make_tensor(make_smem_ptr(Debug_Rowmax_smem), cfg.softmax.reduce.x.shape);
    auto S_copy_cfg = make_smem_FrgCopyCfg<LT, decltype(sS_tile.layout()), decltype(cfg.mma1.C.mma_FrgThr), decltype(rS_frg_mma.layout())>();
    auto R_copy_cfg = make_smem_FrgCopyCfg<LT, decltype(sRow_tile.layout()), decltype(cfg.softmax.reduce.x.frg_thr)>();
#endif
    // ----- Pipeline routines -----
    auto preload_Q = [&]() {
        CUTE_UNROLL
        for (int k_tile = 0; k_tile < cfg.mma1.K_tile_num; k_tile++) {
            auto gQ_tile = coord_qk.slice_A_tile(gQ_slab);
            auto sQ_tile = slice_rest(sQ_tile_pipe, k_tile);
            CTA_copy_tile(cfg.mma1.A.tile_copy, gQ_tile, sQ_tile);
            coord_qk.step_K();
        }
        cp_async_fence(); cp_async_wait<0>(); __syncthreads();
        CUTE_UNROLL
        for (int k_tile = 0; k_tile < cfg.mma1.K_tile_num; k_tile++) {
            auto sQ_tile = slice_rest(sQ_tile_pipe, k_tile);
            auto rQ_frg_tile = slice_rest(rQ_frg, k_tile);
            load(cfg.mma1.A.frg_copy, sQ_tile, rQ_frg_tile);
        }
        __syncthreads();
    };
    auto pipe_fetch_QK = [&](int fetch_end) {
        if (coord_qk.valid_K_tile(cfg.mma1.K) && (int)coord_qk.N_coord() < fetch_end) {
            // if (DEBUG_THREAD) {
            //     print("pipe_fetch_QK: coord_qk: "); print(coord_qk.M_coord()); print(", "); print(coord_qk.N_coord()); print(", "); print(coord_qk.K_coord()); print("\n");
            // }
            if constexpr (!cfg.q_in_reg) {
                if (coord_qk.N_coord() == mask_start || cfg.mma1.K_tile_num > 1) { // Q only needs to be loaded once if K_tile_num == 1
                    auto gQ_tile = coord_qk.slice_A_tile(gQ_slab);
                    pipe_qk.fetch(gQ_tile, sQ_tile_pipe, cfg.mma1.A.tile_copy);
                }
            }
            auto gK_tile = coord_qk.slice_B_tile(gK_slab);
            pipe_qk.fetch(gK_tile, sK_tile_pipe, cfg.mma1.B.tile_copy);
            coord_qk.step_K();
            if (coord_qk.K_coord() == cfg.mma1.K_tile_num) {
                coord_qk.reset_K();
                if (coord_qk.N_coord() == (mask_end - 1))
                    coord_qk.reset_N();
                else
                    coord_qk.step_N();
            }
        }
        pipe_qk.commit();
    };
    auto pipe_fetch_V = [&](int fetch_end) {
        if (coord_pv.valid_K_tile(cfg.mma2.K) && (int)coord_pv.K_coord() < fetch_end) {
            // if (DEBUG_THREAD) {
            //     print("pipe_fetch_V: coord_pv: "); print(coord_pv.M_coord()); print(", "); print(coord_pv.N_coord()); print(", "); print(coord_pv.K_coord()); print("\n");
            // }
            auto gV_tile = coord_pv.slice_B_tile(gV_slab);
            pipe_pv.fetch(gV_tile, sV_tile_pipe, cfg.mma2.B.tile_copy);
            if (coord_pv.K_coord() == (mask_end - 1)) 
                coord_pv.reset_K();
            else
                coord_pv.step_K();
        }
        pipe_pv.commit();
    };
    // ----- Mma routines -----
    auto qk_mma = [&](int fetch_end) {
        clear(rS_frg_mma);
        CUTE_UNROLL
        for (int k_tile = 0; k_tile < cfg.mma1.K_tile_num; k_tile++) {
            pipe_qk.template ready<smempipe_pv>();
            auto sQ_tile = pipe_qk.read(sQ_tile_pipe);
            auto sK_tile = pipe_qk.read(sK_tile_pipe);

            // __syncthreads();
            // if (DEBUG_THREAD) {
            //     print("qk_mma: sQ_tile: "); print_tensor(sQ_tile); print("\n");
            //     print("qk_mma: sK_tile: "); print_tensor(sK_tile); print("\n");
            // }
            // __syncthreads();
            if constexpr (cfg.q_in_reg) {
                auto rQ_frg_tile = slice_rest(rQ_frg, k_tile);
                pipeB_mma(cfg.mma1, rQ_frg_tile, sK_tile, rS_frg_mma);
            } else {
                pipeAB_mma(cfg.mma1, sQ_tile, sK_tile, rS_frg_mma);
            }
            pipe_qk.step();
            pipe_fetch_QK(fetch_end);
        }
    };
    auto pv_mma = [&](int fetch_end) {
        pipe_pv.template ready<smempipe_qk>();
        auto rP_frg = convert_type<T>(rP_frg_mma);
        auto sV_tile = pipe_pv.read(sV_tile_pipe);
        // __syncthreads();
        // if (DEBUG_THREAD) {
        //     print("pv_mma: rP_frg_mma: "); print_tensor(rP_frg_mma); print("\n");
        //     print("pv_mma: rP_frg: "); print_tensor(rP_frg); print("\n");
        //     print("pv_mma: sV_tile: "); print_tensor(sV_tile); print("\n");
        // }
        // __syncthreads();
        pipeB_mma(cfg.mma2, rP_frg, sV_tile, rO_frg_mma);
        // __syncthreads();
        // if (DEBUG_THREAD) {
        //     print("pv_mma: rO_frg_mma: "); print_tensor(rO_frg_mma); print("\n");
        // }
        // __syncthreads();
        pipe_pv.step();
        pipe_fetch_V(fetch_end);
    };
    // ----- Preload Q ----- 
    if constexpr (cfg.q_in_reg) {
        preload_Q();
        coord_qk.reset_K();
    }
    // ----- Prefill pipeline ----- 
    CUTE_UNROLL
    for (; pipe_pv.stage < smempipe_pv; pipe_pv.step()) {
        CUTE_UNROLL
        for (; pipe_qk.stage < (pipe_pv.stage + 1) * cfg.mma1.K_tile_num; pipe_qk.step()) {
            pipe_fetch_QK(mask_end);
        }
        pipe_fetch_V(mask_end);
    }
    // ----- Initialize rowmax and rowsum -----
    auto rowmax_frg = init_tensor<LT, decltype(cfg.softmax.reduce.x.frg_shape)>(-std::numeric_limits<LT>::infinity());
    auto lse_frg = init_tensor<LT, decltype(cfg.softmax.reduce.x.frg_shape)>(0.0f);
    // ----- Main loop 1, masked case -----
    for (int n_tile = mask_start; n_tile < mask_end; n_tile++) {
        auto mstate = make_mask_state(cfg.mask, coord_qk.M_coord() * cfg.mma1.M_tile, n_tile * cfg.mma1.N_tile, (int)cfg.mma1.M, (int)cfg.mma1.N);
        qk_mma(n_tile == mask_end - 1 ? mask_start : mask_end);
        DEBUG_PRINT_TENSOR("loop 1, S_tile:");
        mask_frg(cfg.mask, mstate, rS_frg_mma);
        if (n_tile == mask_start)
            softmax(cfg.softmax, rS_frg_mma, sRow, rowmax_frg, lse_frg);
        else
            softmax_rescale(cfg.softmax, rS_frg_mma, rO_frg_mma, sRow, rowmax_frg, lse_frg);
        DEBUG_PRINT_TENSOR("loop 1, P_tile:");
        pv_mma(n_tile == mask_end - 1 ? mask_start : mask_end);
    }
    // ----- Main loop 2, non-masked case ----- 
    for (int n_tile = 0; n_tile < mask_start; n_tile++) {
        qk_mma(mask_start);
        softmax_rescale(cfg.softmax, rS_frg_mma, rO_frg_mma, sRow, rowmax_frg, lse_frg);
        pv_mma(mask_start);
    }
    pipe_pv.finish();
    // ----- Write O and rowsum to global memory -----
    softmax_epilogue(cfg.softmax, rO_frg_mma, lse_frg, sRow);
    alloc.reset(smem);
    T* O_smem = alloc.allocate<T>(size(cfg.mma2.C.sTile));
    auto sO_tile = make_tensor(make_smem_ptr(O_smem), cfg.mma2.C.sTile);
    __syncthreads();
    copy(rO_frg_mma, slice_rest(sO_tile, cfg.mma2.C.mma_FrgThr, tid));
    __syncthreads();
    auto gO_tile = coord_pv.slice_C_tile(gO_slab);
    auto gL_tile = coord_pv.slice_L_tile(gl_slab);
    if (cfg.softmax.reduce.owns_frg())
        copy(lse_frg, slice_rest(gL_tile, cfg.softmax.reduce.x.frg_thr, tid));
    CTA_copy_tile(cfg.mma2.C.tile_copy, sO_tile, gO_tile);
}

int launch_flash_attention_kernel(auto cfg, auto Q_ptr, auto K_ptr, auto V_ptr, auto O_ptr, auto l_ptr) {
    dim3 blocks(cfg.mma2.M_tile_num, cfg.mma2.N_tile_num, cfg.mma2.P_tile_num);
    int threads = int(size(cfg.mma2.Threads));
    using T = typename decltype(cfg)::T;
    using LT = typename decltype(cfg)::LT;
    int smem_size = cfg.smem_size();
#ifdef DEBUG
    smem_size += cfg.mma1.M_tile * cfg.mma1.N_tile * sizeof(LT);
#endif
    auto kernel = &flash_kernel<decltype(cfg), T, LT>;
    CUDA_CHECK_RETURN(adjust_dynamic_smem_size(kernel, smem_size), "Error setting max dynamic shared memory size");
    kernel<<<blocks, threads, smem_size>>>(cfg, Q_ptr, K_ptr, V_ptr, O_ptr, l_ptr);
    CUDA_CHECK_LAST_ERROR("CUDA kernel launch error");
    return 0;
}


} // namespace vidrial