config.h

#pragma once

#include <cutlass/numeric_types.h>
#include <cutlass/arch/barrier.h>
#include <cute/tensor.hpp>
#include <kerutils/kerutils.cuh>

#include "defines.h"
#include "params.h"

using namespace cute;

namespace sm90::decode::sparse_fp8 {

template<ModelType MODEL_TYPE, int NUM_HEADS>
class KernelTemplate {
public:

static_assert(NUM_HEADS == 64 || NUM_HEADS == 128 || NUM_HEADS == 16);
// todo only support tp8
static constexpr int BLOCK_M = 16;
static constexpr int NUM_M_BLOCKS = NUM_HEADS / BLOCK_M;
static constexpr bool Is_causal = false;
static constexpr int HEAD_DIM_K = MODEL_TYPE == ModelType::V32 ? 576 : 512;
static constexpr int HEAD_DIM_V = 512;
static constexpr int HEAD_DIM_ROPE = 64;
static constexpr int HEAD_DIM_NOPE = HEAD_DIM_K - HEAD_DIM_ROPE;

static constexpr int QUANT_TILE_SIZE = MODEL_TYPE == ModelType::V32 ? 128 : 64;
static constexpr int NUM_SCALES = MODEL_TYPE == ModelType::V32 ? 4 : 8;  // For MODEL1: 7 fp8_e4m3 + 1 padding

static constexpr int NUM_THREADS = 256;
static constexpr int TOPK_BLOCK_SIZE = 64;
using elem_type = cutlass::bfloat16_t;
using MMA_Atom_Arch = std::conditional_t<
    std::is_same_v<elem_type, cutlass::half_t>,
    MMA_Atom<GFX928_16x16x64_F32F16F16F32_NT>,
    MMA_Atom<GFX928_16x16x64_F32BF16BF16F32_NT>
>;
static constexpr int kNWarps = 4;
using ValLayoutMNK = Layout<Shape<_1, _1, _1>>;
using TiledMma = TiledMMA<
    MMA_Atom_Arch,
    Layout<Shape<_1, Int<kNWarps>, _1>>,  // 1x4x1 or 1x8x1 thread group
    ValLayoutMNK>;

using MMA_Atom_Arch_16_16_32 = std::conditional_t<
    std::is_same_v<elem_type, cutlass::half_t>,
    MMA_Atom<GFX928_16x16x32_F32F16F16F32_NN>,
    MMA_Atom<GFX928_16x16x32_F32BF16BF16F32_NN>
>;
using TiledMma_16_16_32 = TiledMMA<
    MMA_Atom_Arch_16_16_32,
    Layout<Shape<_1, Int<kNWarps>, _1>>,  // 1x4x1 or 1x8x1 thread group
    ValLayoutMNK>;

using MMA_Atom_Arch_16x32_NT = std::conditional_t<
    std::is_same_v<elem_type, cutlass::half_t>,
    MMA_Atom<GFX928_16x32x16_F32F16F16F32_NT>,
    MMA_Atom<GFX928_16x32x16_F32BF16BF16F32_NT>
>;
using TiledMma_O = TiledMMA<
    MMA_Atom_Arch_16x32_NT,
    Layout<Shape<_1, Int<kNWarps>, _1>>,  // 1x4x1 or 1x8x1 thread group
    ValLayoutMNK>;

using SmemLayoutAtomK = decltype(composition(
    Swizzle<3, 3, 3>{},
    Layout<Shape<Int<8>, Int<32>>, Stride<Int<32>, _1>>{}));
using SmemLayoutK = decltype(tile_to_shape(
    SmemLayoutAtomK{},
    Shape<Int<TOPK_BLOCK_SIZE>, Int<8 * 32>>{}));

using SmemLayoutAtomV = SmemLayoutAtomK;   
using SmemLayoutV = decltype(tile_to_shape(
    SmemLayoutAtomV{},
    Shape<Int<TOPK_BLOCK_SIZE>, Int<512>>{}));

using SmemLayoutVtransposed = decltype(
    composition(SmemLayoutV{}, make_layout(Shape<Int<512>, Int<TOPK_BLOCK_SIZE>>{}, GenRowMajor{})));
using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{}));

using SmemLayoutAtomP = Layout<Shape<Int<4*16*16>>, Stride<Int<1>>>;
using SmemLayoutP = decltype(tile_to_shape(
    SmemLayoutAtomP{},
    Shape<Int<4*16*16>>{}));
using SmemLayoutRow = Layout<Shape<_128>, Stride<_1>>; 

using Element = cutlass::bfloat16_t;
using ElementAccum = float;
struct SharedMemoryPlan {
    union {
        struct {
            cute::array_aligned<Element, cute::cosize_v<SmemLayoutV>> smem_v;
        };
        struct {
            // cute::array_aligned<typename Kernel_traits::Element, cute::cosize_v<typename Kernel_traits::SmemLayoutV_tmp>> smem_v_tmp;  // Double buffer
            cute::array_aligned<Element, cute::cosize_v<SmemLayoutP>> smem_p;
            cute::array_aligned<ElementAccum, cute::cosize_v<SmemLayoutRow>> smem_row_sum;
            cute::array_aligned<ElementAccum, cute::cosize_v<SmemLayoutRow>> smem_row_max;

        };
        // struct {
        //     cute::array_aligned<typename Kernel_traits::ElementAccum, cute::cosize_v<typename Kernel_traits::SmemLayoutO>> smem_o;
        //     // cute::array_aligned<typename Kernel_traits::Element, cute::cosize_v<typename Kernel_traits::SmemLayoutP>> smem_p;
        //     // cute::array_aligned<typename Kernel_traits::ElementAccum, cute::cosize_v<typename Kernel_traits::SmemLayoutRow>> smem_row_sum;
        //     // cute::array_aligned<typename Kernel_traits::ElementAccum, cute::cosize_v<typename Kernel_traits::SmemLayoutRow>> smem_row_max;
        // };
        // struct {
        //     cute::array_aligned<typename Kernel_traits::Element, cute::cosize_v<typename Kernel_traits::SmemLayoutQ>> smem_q;
        // };
    };

    // array_aligned<bf16, cosize_v<SmemLayoutQ>> q;
    // union {
    //     array_aligned<bf16, cosize_v<SmemLayoutK>> k[NUM_K_BUFS];
    //     array_aligned<bf16, cosize_v<SmemLayoutOBuf>> oBuf;
    //     array_aligned<float, cosize_v<SmemLayoutOAccumBuf>> oAccumBuf;
    // } u;
    // CUTE_ALIGNAS(1024) array_aligned<bf16, cosize_v<SmemLayoutS>> s;
    // bool is_kv_valid[NUM_K_BUFS][TOPK_BLOCK_SIZE];

    // float sM[BLOCK_M], sL[BLOCK_M], sScale[BLOCK_M], sOScale[BLOCK_M];
    // transac_bar_t bar_q, bar_k_local_ready[NUM_K_BUFS], bar_k_remote_ready[NUM_K_BUFS], bar_k_avail[NUM_K_BUFS];
};

// template<
//     typename Shape_Q, typename TMA_Q
// >

// using TiledMMA_QK = decltype(make_tiled_mma(
//     GMMA::MMA_64x64x16_F32BF16BF16_SS<GMMA::Major::K, GMMA::Major::K>{},
//     Layout<Shape<_1, _1, _1>>{}
// ));

// using TiledMMA_QK_rQ = decltype(make_tiled_mma(
//     GMMA::MMA_64x64x16_F32BF16BF16_RS<GMMA::Major::K, GMMA::Major::K>{},
//     Layout<Shape<_1, _1, _1>>{}
// ));

// using TiledMMA_PV_LocalP = decltype(make_tiled_mma(
//     GMMA::MMA_64x256x16_F32BF16BF16_RS<GMMA::Major::K, GMMA::Major::MN>{},
//     Layout<Shape<_1, _1, _1>>{}
// ));

// using TiledMMA_PV_RemoteP = decltype(make_tiled_mma(
//     GMMA::MMA_64x256x16_F32BF16BF16_SS<GMMA::Major::K, GMMA::Major::MN>{},
//     Layout<Shape<_1, _1, _1>>{}
// ));






static __device__ __forceinline__ void
compute_attn_1rowblock_splitkv_sparse_mla_fp8(const SparseAttnDecodeParams &params, const DecodingSchedMeta& sched_meta, int batch_idx);

static __device__ __forceinline__ void
devfunc(const SparseAttnDecodeParams &params);

static void run(const SparseAttnDecodeParams &params);

};

}