#define NUM_WORKERS 16 // This kernel uses 16 workers in parallel per block, to help issue instructions more quickly.using namespace kittens; // this kernel only handles headdim=64 for simplicity. Also n should be a multiple of 256 here.__global__ void attend_ker64(int n, const bf16* __restrict__ __q__, const bf16* __restrict__ __k__, const bf16* __restrict__ __v__, bf16* __o__) {
    auto warpid        = kittens::warpid();    auto block_start   = blockIdx.x*(n*64);    const bf16 *_q = __q__ + block_start, *_k = __k__ + block_start, *_v = __v__ + block_start;          bf16 *_o = __o__ + block_start;
    extern __shared__ alignment_dummy __shm[]; // this is the CUDA shared memory    shared_allocator al((int*)&__shm[0]);        // K and V live in shared memory -- this is about all that will fit.    st_bf_1x4<:st_layout::swizzle> (&k_smem)[NUM_WORKERS] = al.allocate, NUM_WORKERS>();    st_bf_1x4<:st_layout::swizzle> (&v_smem)[NUM_WORKERS] = al.allocate, NUM_WORKERS>();
    // Initialize all of the register tiles.    rt_bf_1x4<> q_reg, k_reg, v_reg; // v_reg need to be swapped into col_l    rt_fl_1x1<> att_block;    rt_bf_1x1<> att_block_mma;    rt_fl_1x4<> o_reg;    rt_fl_1x1<>::col_vec max_vec_last, max_vec; // these are column vectors for the attention block    rt_fl_1x1<>::col_vec norm_vec_last, norm_vec; // these are column vectors for the attention block        int qo_blocks = n / (q_reg.rows*NUM_WORKERS), kv_blocks = n / (q_reg.rows*NUM_WORKERS);
    for(auto q_blk = 0; q_blk < qo_blocks; q_blk++) {
        // each warp loads its own Q tile of 16x64, and then multiplies by 1/sqrt(d)        load(q_reg, _q + (q_blk*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols);        mul(q_reg, q_reg, __float2bfloat16(0.125f)); // temperature adjustment
        // zero flash attention L, M, and O registers.        neg_infty(max_vec); // zero registers for the Q chunk        zero(norm_vec);        zero(o_reg);
        // iterate over k, v for these q's that have been loaded        for(auto kv_idx = 0; kv_idx < kv_blocks; kv_idx++) {
            // each warp loads its own chunk of k, v into shared memory            load(v_smem[warpid], _v + (kv_idx*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols);            load(k_smem[warpid], _k + (kv_idx*NUM_WORKERS + warpid)*q_reg.num_elements, q_reg.cols);            __syncthreads(); // we need to make sure all memory is loaded before we can begin the compute phase
            // now each warp goes through all of the subtiles, loads them, and then does the flash attention internal alg.            for(int subtile = 0; subtile < NUM_WORKERS; subtile++) {
                load(k_reg, k_smem[subtile]); // load k from shared into registers
                zero(att_block); // zero 16x16 attention tile                mma_ABt(att_block, q_reg, k_reg, att_block); // [email protected]
                copy(norm_vec_last, norm_vec);                copy(max_vec_last,  max_vec);
                row_max(max_vec, att_block, max_vec); // accumulate onto the max_vec                sub_row(att_block, att_block, max_vec); // subtract max from attention -- now all <=0                exp(att_block, att_block); // exponentiate the block in-place.
                sub(max_vec_last, max_vec_last, max_vec); // subtract new max from old max to find the new normalization.                exp(max_vec_last, max_vec_last); // exponentiate this vector -- this is what we need to normalize by.                mul(norm_vec, norm_vec, max_vec_last); // and the norm vec is now normalized.
                row_sum(norm_vec, att_block, norm_vec); // accumulate the new attention block onto the now-rescaled norm_vec                div_row(att_block, att_block, norm_vec); // now the attention block is correctly normalized
                mul(norm_vec_last, norm_vec_last, max_vec_last); // normalize the previous norm vec according to the new max                div(norm_vec_last, norm_vec_last, norm_vec); // normalize the previous norm vec according to the new norm
                copy(att_block_mma, att_block); // convert to bf16 for mma_AB
                load(v_reg, v_smem[subtile]); // load v from shared into registers.                rt_bf_1x4<:rt_layout::col> &v_reg_col = swap_layout_inplace(v_reg); // this is a reference and the call has invalidated v_reg
                mul_row(o_reg, o_reg, norm_vec_last); // normalize o_reg in advance of mma_AB'ing onto it                mma_AB(o_reg, att_block_mma, v_reg_col, o_reg); // mfma onto o_reg with the local attention@V matmul.            }            __syncthreads(); // we need to make sure all warps are done before we can start loading the next kv chunk        }
        store(_o + (q_blk*NUM_WORKERS + warpid)*q_reg.num_elements, o_reg, q_reg.cols); // write out o. compiler has an issue with register usage if d is made constexpr q_reg.rows :/    }}

template<int D>__global__  __launch_bounds__((NUM_WORKERS)*kittens::WARP_THREADS, 2)void fwd_attend_ker_dim(int N, const CUtensorMap* tma_q, const CUtensorMap* tma_k, const CUtensorMap* tma_v, CUtensorMap* tma_o) {    extern __shared__ int __shm[]; // this is the CUDA shared memory    tma_swizzle_allocator al((int*)&__shm[0]);
    constexpr int tile_width = fwd_attend_ker_tile_dims::tile_width; // constants    constexpr int qo_height  = fwd_attend_ker_tile_dims::qo_height;    constexpr int kv_height  = fwd_attend_ker_tile_dims::kv_height;
    st_bf          (&q_smem)   [NUM_WARPGROUPS] = al.allocate,          NUM_WARPGROUPS>();    st_bf          (&k_smem)[2][NUM_WORKERS_KV] = al.allocate, 2,       NUM_WORKERS_KV>();    st_bf          (&v_smem)[2][NUM_WORKERS_KV] = al.allocate, 2,       NUM_WORKERS_KV>();
    int tic = 0, toc = 1;     rt_fl<1, kv_height> att_block;    rt_bf<1, kv_height> att_block_mma;    rt_fl<1, qo_height> o_prev;    col_vec1, kv_height>> max_vec_last, max_vec;    col_vec1, kv_height>> norm_vec_last, norm_vec;
    int warpid      = kittens::warpid();    int warpgroupid = warpid/kittens::WARPGROUP_WARPS;
    int kv_blocks = N / (NUM_WORKERS_KV*k_smem[0][0].rows);
    __shared__ uint64_t qsmem_barrier, kvsmem_barrier;//, vsmem_barrier;
    int q_phasebit = 0;    int kv_phasebit = 0;
    if (threadIdx.x == 0) {        tma::init_barrier, NUM_WARPGROUPS>(qsmem_barrier, 1);        tma::init_barrier, NUM_WORKERS_KV*2>(kvsmem_barrier, 1);     }
    if (warpid == 0) {        for (int wg = 0; wg < NUM_WORKERS/kittens::WARPGROUP_WARPS; wg++) { // load q            int tile_idx = (blockIdx.y * NUM_WARPGROUPS * gridDim.x) + (blockIdx.x * NUM_WARPGROUPS) + wg;            tma::load_async((q_smem[wg]), tma_q, qsmem_barrier, tile_idx);         }        for (int w = 0; w < NUM_WORKERS_KV; w++) { // load k, v                  int tile_idx = (blockIdx.y * NUM_WORKERS_KV * kv_blocks) + (0 * NUM_WORKERS_KV) + w;             tma::load_async((k_smem[tic][w]), tma_k, kvsmem_barrier, tile_idx);             tma::load_async((v_smem[tic][w]), tma_v, kvsmem_barrier, tile_idx);         }    }
    neg_infty(max_vec); // zero registers for the Q chunk    zero(norm_vec);    zero(o_prev);    __syncthreads();
    tma::arrive_and_wait(qsmem_barrier, q_phasebit);    q_phasebit ^= 1;
    if constexpr (D == 64) { warpgroup::mul(q_smem[warpgroupid], q_smem[warpgroupid], __float2bfloat16(0.125f)); }     else { warpgroup::mul(q_smem[warpgroupid], q_smem[warpgroupid], __float2bfloat16(0.08838834764f)); }
    for (auto kv_idx = 0; kv_idx < kv_blocks; kv_idx++, tic ^= 1, toc ^= 1) {        tma::arrive_and_wait(kvsmem_barrier, kv_phasebit);        kv_phasebit ^= 1;
        __syncthreads();        if (warpid == 0) {            tma::set_bytes(kvsmem_barrier, 2 * NUM_WORKERS_KV * k_smem[0][0].num_elements * sizeof(bf16));
            if (kv_idx + 1 < kv_blocks) {                for (int w = 0; w < NUM_WORKERS_KV; w++) {                            int tile_idx = (blockIdx.y * NUM_WORKERS_KV * kv_blocks) + ((kv_idx + 1) * NUM_WORKERS_KV) + w;                     tma::load_async((k_smem[toc][w]), tma_k, kvsmem_barrier, tile_idx);                     tma::load_async((v_smem[toc][w]), tma_v, kvsmem_barrier, tile_idx);                }            }        }
        warpgroup::mma_fence(att_block);        warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[tic][0]);        warpgroup::mma_commit_group();
        copy(norm_vec_last, norm_vec);        copy(max_vec_last,  max_vec);
        warpgroup::mma_async_wait();
        row_max(max_vec, att_block, max_vec); // accumulate onto the max_vec        sub_row(att_block, att_block, max_vec);        exp(att_block, att_block);
        sub(max_vec_last, max_vec_last, max_vec);        exp(max_vec_last, max_vec_last);        mul(norm_vec, norm_vec, max_vec_last);
        row_sum(norm_vec, att_block, norm_vec); // accumulate onto the norm_vec        div_row(att_block, att_block, norm_vec);