from __future__ import annotations

import os
try:
    import ninja
    os.environ["PATH"] = f"{ninja.BIN_DIR}:{os.environ.get('PATH', '')}"
except ImportError:
    pass
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

OP_TYPE = "deepseek_nsa"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000"]

B, H, S, D = 1, 16, 1024, 64
BLOCK_SIZE = 64
TOP_N_BLOCKS = 8
SLIDING_WINDOW = 64

cuda_source = r"""
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <torch/extension.h>
#include <ATen/cuda/EmptyTensor.h>
#include <cstdint>

namespace {

__device__ __forceinline__ uint32_t smem_u32addr(const void* ptr) {
    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
}

__device__ __forceinline__ void cp_async_ca_16(uint32_t smem_addr, const void* gmem_ptr) {
    asm volatile(
        "cp.async.ca.shared.global [%0], [%1], 16;\n"
        :
        : "r"(smem_addr), "l"(gmem_ptr)
    );
}

__device__ __forceinline__ void cp_async_commit() {
    asm volatile("cp.async.commit_group;\n" ::);
}

template <int N>
__device__ __forceinline__ void cp_async_wait() {
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}

__device__ __forceinline__ void ldmatrix_x4(uint32_t &r0, uint32_t &r1, uint32_t &r2, uint32_t &r3, uint32_t addr) {
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
        : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
        : "r"(addr)
    );
}

__device__ __forceinline__ void ldmatrix_x4_trans(uint32_t &r0, uint32_t &r1, uint32_t &r2, uint32_t &r3, uint32_t addr) {
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n"
        : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
        : "r"(addr)
    );
}

__device__ __forceinline__ void mma_m16n8k16_bf16(
    float &d0, float &d1, float &d2, float &d3,
    uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
    uint32_t b0, uint32_t b1
) {
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
        "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)
    );
}

template <int D>
struct CTASharedStorage {
    __nv_bfloat16 k_tile[64][D + 8];
    __nv_bfloat16 v_tile[64][D + 8];
    uint64_t warp_mask0[4];
    uint64_t warp_mask1[4];
};

union SPFrag {
    float s[8][4];
    uint32_t p[4][4];
};

template <int D, int CTA_LOAD_ITERS, int UINT4_PER_ROW>
__device__ __forceinline__ void issue_cta_tile_cp_async(
    __nv_bfloat16 tile[64][D + 8],
    const __nv_bfloat16* __restrict__ gmem_base,
    int tid_cta
) {
    #pragma unroll
    for (int i = 0; i < CTA_LOAD_ITERS; ++i) {
        int idx = tid_cta + i * 128;
        int r = idx / UINT4_PER_ROW;
        int c = (idx % UINT4_PER_ROW) * 8;
        uint32_t dst = smem_u32addr(&tile[r][c]);
        cp_async_ca_16(dst, gmem_base + r * D + c);
    }
}

template <int D, int WARPS_PER_CTA = 4, int MIN_BLOCKS = 2>
__global__ __launch_bounds__(WARPS_PER_CTA * 32, MIN_BLOCKS)
void nsa_cta_coop_kernel(
    const __nv_bfloat16* __restrict__ Q,
    const __nv_bfloat16* __restrict__ K,
    const __nv_bfloat16* __restrict__ V,
    float* __restrict__ K_mean,
    uint32_t* __restrict__ km_flag,
    uint32_t gen,
    __nv_bfloat16* __restrict__ O,
    int S,
    int n_full_blocks,
    float scale
) {
    constexpr int TK = D / 16;          // 4 (D=64) or 8 (D=128)
    constexpr int TN = D / 8;           // 8 (D=64) or 16 (D=128)
    constexpr int UINT4_PER_ROW = D / 8; // 8 (D=64) or 16 (D=128)
    constexpr int CTA_LOAD_ITERS = (64 * UINT4_PER_ROW) / 128; // 4 (D=64) or 8 (D=128)

    __shared__ CTASharedStorage<D> smem;

    const int tid_cta = threadIdx.x;
    const int w = tid_cta >> 5;
    const int lane = tid_cta & 31;
    const int groupID = lane >> 2; // 0..7
    const int tid = lane & 3;      // 0..3

    const int bi_curr = blockIdx.x;
    const int s0 = bi_curr * 64;
    const int t_warp_base = s0 + w * 16;

    const int bh = blockIdx.y;
    const int64_t bh_offset = static_cast<int64_t>(bh) * S * D;
    const __nv_bfloat16* q_head = Q + bh_offset;
    const __nv_bfloat16* k_head = K + bh_offset;
    const __nv_bfloat16* v_head = V + bh_offset;
    __nv_bfloat16* o_head = O + bh_offset;

    const int q_top_local = w * 16 + groupID;
    const int q_bot_local = w * 16 + groupID + 8;
    const int t_top = s0 + q_top_local;
    const int t_bot = s0 + q_bot_local;
    const bool valid_q_top = (t_top < S);
    const bool valid_q_bot = (t_bot < S);

    const int sub_r = lane & 15;
    const int sub_c = (lane >> 4) << 3;

    // 1. Cooperative CTA load of Q[s0 .. s0+63] into smem.k_tile
    #pragma unroll
    for (int i = 0; i < CTA_LOAD_ITERS; ++i) {
        int idx = tid_cta + i * 128;
        int r = idx / UINT4_PER_ROW;
        int c = (idx % UINT4_PER_ROW) * 8;
        int t_r = s0 + r;
        uint4 val = (t_r < S)
            ? __ldg(reinterpret_cast<const uint4*>(q_head + static_cast<int64_t>(t_r) * D + c))
            : make_uint4(0, 0, 0, 0);
        *reinterpret_cast<uint4*>(&smem.k_tile[r][c]) = val;
    }
    __syncthreads();

    uint32_t q_frag[TK][4];
    #pragma unroll
    for (int tk = 0; tk < TK; ++tk) {
        uint32_t addr = smem_u32addr(&smem.k_tile[w * 16 + sub_r][tk * 16 + sub_c]);
        ldmatrix_x4(q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], addr);
    }
    __syncthreads();

    // Online softmax state for q_top and q_bot
    float m_top = -1e30f, l_top = 0.0f;
    float m_bot = -1e30f, l_bot = 0.0f;
    float o_frag[TN][4];
    #pragma unroll
    for (int tn = 0; tn < TN; ++tn) {
        o_frag[tn][0] = 0.0f;
        o_frag[tn][1] = 0.0f;
        o_frag[tn][2] = 0.0f;
        o_frag[tn][3] = 0.0f;
    }

    float sum_curr_top = 0.0f;
    float sum_curr_bot = 0.0f;

    // Step 1: Cooperative CTA load of K[bi_curr] and V[bi_curr]
    if (s0 + 64 <= S) {
        issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.k_tile, k_head + static_cast<int64_t>(s0) * D, tid_cta);
        issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.v_tile, v_head + static_cast<int64_t>(s0) * D, tid_cta);
        cp_async_commit();
        cp_async_wait<0>();
    } else {
        #pragma unroll
        for (int i = 0; i < CTA_LOAD_ITERS; ++i) {
            int idx = tid_cta + i * 128;
            int r = idx / UINT4_PER_ROW;
            int c = (idx % UINT4_PER_ROW) * 8;
            int k_idx = s0 + r;
            uint4 kv = (k_idx < S) ? __ldg(reinterpret_cast<const uint4*>(k_head + static_cast<int64_t>(k_idx) * D + c)) : make_uint4(0, 0, 0, 0);
            uint4 vv = (k_idx < S) ? __ldg(reinterpret_cast<const uint4*>(v_head + static_cast<int64_t>(k_idx) * D + c)) : make_uint4(0, 0, 0, 0);
            *reinterpret_cast<uint4*>(&smem.k_tile[r][c]) = kv;
            *reinterpret_cast<uint4*>(&smem.v_tile[r][c]) = vv;
        }
    }
    __syncthreads();

    if (t_warp_base < S) {
        const int max_local_k = min(w * 16 + 15, S - 1 - s0);
        for (int k_chunk = 0; k_chunk <= w; ++k_chunk) {
            const int k_base = k_chunk * 16;
            if (k_base > max_local_k) break;

            float s_0[4] = {0.f, 0.f, 0.f, 0.f};
            float s_1[4] = {0.f, 0.f, 0.f, 0.f};
            #pragma unroll
            for (int tk = 0; tk < TK; ++tk) {
                uint32_t r0, r1, r2, r3;
                uint32_t addr = smem_u32addr(&smem.k_tile[k_base + sub_r][tk * 16 + sub_c]);
                ldmatrix_x4(r0, r1, r2, r3, addr);
                mma_m16n8k16_bf16(s_0[0], s_0[1], s_0[2], s_0[3],
                                  q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r0, r2);
                mma_m16n8k16_bf16(s_1[0], s_1[1], s_1[2], s_1[3],
                                  q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r1, r3);
            }

            const int kc0 = k_base + tid * 2 + 0;
            const int kc1 = k_base + tid * 2 + 1;
            const int kc2 = k_base + 8 + tid * 2 + 0;
            const int kc3 = k_base + 8 + tid * 2 + 1;

            bool vt0 = valid_q_top && (kc0 <= q_top_local);
            bool vt1 = valid_q_top && (kc1 <= q_top_local);
            bool vt2 = valid_q_top && (kc2 <= q_top_local);
            bool vt3 = valid_q_top && (kc3 <= q_top_local);

            bool vb0 = valid_q_bot && (kc0 <= q_bot_local);
            bool vb1 = valid_q_bot && (kc1 <= q_bot_local);
            bool vb2 = valid_q_bot && (kc2 <= q_bot_local);
            bool vb3 = valid_q_bot && (kc3 <= q_bot_local);

            float st0 = vt0 ? (s_0[0] * scale) : -1e30f;
            float st1 = vt1 ? (s_0[1] * scale) : -1e30f;
            float st2 = vt2 ? (s_1[0] * scale) : -1e30f;
            float st3 = vt3 ? (s_1[1] * scale) : -1e30f;

            float sb0 = vb0 ? (s_0[2] * scale) : -1e30f;
            float sb1 = vb1 ? (s_0[3] * scale) : -1e30f;
            float sb2 = vb2 ? (s_1[2] * scale) : -1e30f;
            float sb3 = vb3 ? (s_1[3] * scale) : -1e30f;

            if (vt0) sum_curr_top += st0;
            if (vt1) sum_curr_top += st1;
            if (vt2) sum_curr_top += st2;
            if (vt3) sum_curr_top += st3;

            if (vb0) sum_curr_bot += sb0;
            if (vb1) sum_curr_bot += sb1;
            if (vb2) sum_curr_bot += sb2;
            if (vb3) sum_curr_bot += sb3;

            float m_chunk_top = fmaxf(fmaxf(st0, st1), fmaxf(st2, st3));
            m_chunk_top = fmaxf(m_chunk_top, __shfl_xor_sync(0xffffffff, m_chunk_top, 1));
            m_chunk_top = fmaxf(m_chunk_top, __shfl_xor_sync(0xffffffff, m_chunk_top, 2));

            float m_chunk_bot = fmaxf(fmaxf(sb0, sb1), fmaxf(sb2, sb3));
            m_chunk_bot = fmaxf(m_chunk_bot, __shfl_xor_sync(0xffffffff, m_chunk_bot, 1));
            m_chunk_bot = fmaxf(m_chunk_bot, __shfl_xor_sync(0xffffffff, m_chunk_bot, 2));

            float m_new_top = fmaxf(m_top, m_chunk_top);
            float exp_old_top = (l_top > 0.0f) ? __expf(m_top - m_new_top) : 0.0f;
            float pt0 = vt0 ? __expf(st0 - m_new_top) : 0.0f;
            float pt1 = vt1 ? __expf(st1 - m_new_top) : 0.0f;
            float pt2 = vt2 ? __expf(st2 - m_new_top) : 0.0f;
            float pt3 = vt3 ? __expf(st3 - m_new_top) : 0.0f;
            float p_sum_top = pt0 + pt1 + pt2 + pt3;
            p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 1);
            p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 2);
            l_top = l_top * exp_old_top + p_sum_top;
            m_top = m_new_top;

            float m_new_bot = fmaxf(m_bot, m_chunk_bot);
            float exp_old_bot = (l_bot > 0.0f) ? __expf(m_bot - m_new_bot) : 0.0f;
            float pb0 = vb0 ? __expf(sb0 - m_new_bot) : 0.0f;
            float pb1 = vb1 ? __expf(sb1 - m_new_bot) : 0.0f;
            float pb2 = vb2 ? __expf(sb2 - m_new_bot) : 0.0f;
            float pb3 = vb3 ? __expf(sb3 - m_new_bot) : 0.0f;
            float p_sum_bot = pb0 + pb1 + pb2 + pb3;
            p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 1);
            p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 2);
            l_bot = l_bot * exp_old_bot + p_sum_bot;
            m_bot = m_new_bot;

            #pragma unroll
            for (int tn = 0; tn < TN; ++tn) {
                o_frag[tn][0] *= exp_old_top;
                o_frag[tn][1] *= exp_old_top;
                o_frag[tn][2] *= exp_old_bot;
                o_frag[tn][3] *= exp_old_bot;
            }

            uint32_t p_frag[4];
            *reinterpret_cast<__nv_bfloat162*>(&p_frag[0]) = __floats2bfloat162_rn(pt0, pt1);
            *reinterpret_cast<__nv_bfloat162*>(&p_frag[1]) = __floats2bfloat162_rn(pb0, pb1);
            *reinterpret_cast<__nv_bfloat162*>(&p_frag[2]) = __floats2bfloat162_rn(pt2, pt3);
            *reinterpret_cast<__nv_bfloat162*>(&p_frag[3]) = __floats2bfloat162_rn(pb2, pb3);

            #pragma unroll
            for (int tn2 = 0; tn2 < TN / 2; ++tn2) {
                uint32_t r0, r1, r2, r3;
                uint32_t addr = smem_u32addr(&smem.v_tile[k_base + sub_r][tn2 * 16 + sub_c]);
                ldmatrix_x4_trans(r0, r1, r2, r3, addr);
                mma_m16n8k16_bf16(o_frag[2 * tn2 + 0][0], o_frag[2 * tn2 + 0][1], o_frag[2 * tn2 + 0][2], o_frag[2 * tn2 + 0][3],
                                  p_frag[0], p_frag[1], p_frag[2], p_frag[3], r0, r1);
                mma_m16n8k16_bf16(o_frag[2 * tn2 + 1][0], o_frag[2 * tn2 + 1][1], o_frag[2 * tn2 + 1][2], o_frag[2 * tn2 + 1][3],
                                  p_frag[0], p_frag[1], p_frag[2], p_frag[3], r2, r3);
            }
        }
    }

    // Fused K_mean computation by Warp 0 (D=64) or Warps 0..1 (D=128) while Warp 3 finishes Step 1!
    if (tid_cta < (D / 2) && bi_curr < ((S - 1) >> 6)) {
        float* out_ptr = K_mean + (static_cast<int64_t>(bh) * n_full_blocks + bi_curr) * D;
        const float scale_inv64 = scale * (1.0f / 64.0f);
        const int idx2 = tid_cta;
        const int tk_km = idx2 >> 3;
        const int sub_tid = (idx2 >> 1) & 3;
        const int half_km = idx2 & 1;
        const int col_km = tk_km * 16 + half_km * 8 + sub_tid * 2;
        float2 s0 = make_float2(0.f, 0.f);
        float2 s1 = make_float2(0.f, 0.f);
        float2 s2 = make_float2(0.f, 0.f);
        float2 s3 = make_float2(0.f, 0.f);
        #pragma unroll 4
        for (int r = 0; r < 64; r += 4) {
            float2 v0 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&smem.k_tile[r + 0][col_km]));
            float2 v1 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&smem.k_tile[r + 1][col_km]));
            float2 v2 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&smem.k_tile[r + 2][col_km]));
            float2 v3 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&smem.k_tile[r + 3][col_km]));
            s0.x += v0.x; s0.y += v0.y;
            s1.x += v1.x; s1.y += v1.y;
            s2.x += v2.x; s2.y += v2.y;
            s3.x += v3.x; s3.y += v3.y;
        }
        float2 sum_km = make_float2(
            ((s0.x + s1.x) + (s2.x + s3.x)) * scale_inv64,
            ((s0.y + s1.y) + (s2.y + s3.y)) * scale_inv64
        );
        reinterpret_cast<float2*>(out_ptr)[idx2] = sum_km;
        if constexpr (D == 64) {
            __syncwarp();
        } else {
            asm volatile("bar.sync 1, 64;\n" ::: "memory");
        }
        if (tid_cta == 0) {
            uint32_t* flag_ptr = km_flag + bh * n_full_blocks + bi_curr;
            asm volatile("st.global.release.gpu.b32 [%0], %1;\n" :: "l"(flag_ptr), "r"(gen) : "memory");
        }
    }

    // If bi_curr >= 1, wait for K_mean[0..bi_curr-1] (if bi_curr >= 8) and prefetch K[bi_curr - 1], V[bi_curr - 1]!
    if (bi_curr >= 1) {
        if (bi_curr >= 8 && tid_cta < bi_curr) {
            const uint32_t* flag_head = km_flag + bh * n_full_blocks;
            uint32_t flg;
            do {
                asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(flg) : "l"(flag_head + tid_cta));
            } while (flg != gen);
        }
        __syncthreads();
        const int64_t prev_base = static_cast<int64_t>(bi_curr - 1) * 64;
        issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.k_tile, k_head + prev_base * D, tid_cta);
        issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.v_tile, v_head + prev_base * D, tid_cta);
        cp_async_commit();
    }

    // Step 2: Determine selected blocks (mask0_top/bot, mask1_top/bot) using online Top-8 in registers
    uint64_t mask0_top = 0ULL, mask1_top = 0ULL;
    uint64_t mask0_bot = 0ULL, mask1_bot = 0ULL;

    if (t_warp_base < S) {
        if (bi_curr < 8) {
            uint64_t m = (1ULL << bi_curr) - 1ULL;
            mask0_top = valid_q_top ? m : 0ULL;
            mask0_bot = valid_q_bot ? m : 0ULL;
        } else {
            sum_curr_top += __shfl_xor_sync(0xffffffff, sum_curr_top, 1);
            sum_curr_top += __shfl_xor_sync(0xffffffff, sum_curr_top, 2);
            sum_curr_bot += __shfl_xor_sync(0xffffffff, sum_curr_bot, 1);
            sum_curr_bot += __shfl_xor_sync(0xffffffff, sum_curr_bot, 2);

            float imp_curr_top = sum_curr_top / static_cast<float>(q_top_local + 1);
            float imp_curr_bot = sum_curr_bot / static_cast<float>(q_bot_local + 1);

            float ts0 = -1e30f, ts1 = -1e30f;
            int   tb0 = 2 * tid - 8, tb1 = 2 * tid + 1 - 8;
            float min_ts = -1e30f;
            int   min_tb = -8;

            float bs0 = -1e30f, bs1 = -1e30f;
            int   bb0 = 2 * tid - 8, bb1 = 2 * tid + 1 - 8;
            float min_bs = -1e30f;
            int   min_bb = -8;

            const float* km_head = K_mean + static_cast<int64_t>(bh) * n_full_blocks * D;
            for (int bi = 0; bi <= bi_curr; ++bi) {
                float st, sb;
                if (bi < bi_curr) {
                    const float* km_row = km_head + bi * D;
                    float dot_top = 0.0f;
                    float dot_bot = 0.0f;
                    #pragma unroll
                    for (int tk = 0; tk < TK; ++tk) {
                        float2 ft0 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&q_frag[tk][0]));
                        float2 fb0 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&q_frag[tk][1]));
                        float2 ft1 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&q_frag[tk][2]));
                        float2 fb1 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>(&q_frag[tk][3]));

                        float4 k01;
                        asm volatile("ld.global.ca.v4.f32 {%0, %1, %2, %3}, [%4];\n"
                                     : "=f"(k01.x), "=f"(k01.y), "=f"(k01.z), "=f"(k01.w)
                                     : "l"(km_row + tk * 16 + tid * 4));
                        dot_top = __fmaf_rn(ft0.x, k01.x, dot_top);
                        dot_top = __fmaf_rn(ft0.y, k01.y, dot_top);
                        dot_top = __fmaf_rn(ft1.x, k01.z, dot_top);
                        dot_top = __fmaf_rn(ft1.y, k01.w, dot_top);

                        dot_bot = __fmaf_rn(fb0.x, k01.x, dot_bot);
                        dot_bot = __fmaf_rn(fb0.y, k01.y, dot_bot);
                        dot_bot = __fmaf_rn(fb1.x, k01.z, dot_bot);
                        dot_bot = __fmaf_rn(fb1.y, k01.w, dot_bot);
                    }
                    dot_top += __shfl_xor_sync(0xffffffff, dot_top, 1);
                    dot_top += __shfl_xor_sync(0xffffffff, dot_top, 2);
                    dot_bot += __shfl_xor_sync(0xffffffff, dot_bot, 1);
                    dot_bot += __shfl_xor_sync(0xffffffff, dot_bot, 2);
                    st = dot_top;
                    sb = dot_bot;
                } else {
                    st = imp_curr_top;
                    sb = imp_curr_bot;
                }

                bool upd_t = (st >= min_ts);
                bool upd_b = (sb >= min_bs);
                if (__any_sync(0xffffffff, upd_t || upd_b)) {
                    if (upd_t) {
                        if (tb0 == min_tb) { ts0 = st; tb0 = bi; }
                        if (tb1 == min_tb) { ts1 = st; tb1 = bi; }
                    }
                    if (upd_b) {
                        if (bb0 == min_bb) { bs0 = sb; bb0 = bi; }
                        if (bb1 == min_bb) { bs1 = sb; bb1 = bi; }
                    }
                    bool t1_worse = (ts1 < ts0) || (ts1 == ts0 && tb1 < tb0);
                    float cur_ts = t1_worse ? ts1 : ts0;
                    int   cur_tb = t1_worse ? tb1 : tb0;

                    bool b1_worse = (bs1 < bs0) || (bs1 == bs0 && bb1 < bb0);
                    float cur_bs = b1_worse ? bs1 : bs0;
                    int   cur_bb = b1_worse ? bb1 : bb0;

                    #pragma unroll
                    for (int mask = 1; mask <= 2; mask <<= 1) {
                        float oth_ts = __shfl_xor_sync(0xffffffff, cur_ts, mask);
                        int   oth_tb = __shfl_xor_sync(0xffffffff, cur_tb, mask);
                        if (oth_ts < cur_ts || (oth_ts == cur_ts && oth_tb < cur_tb)) {
                            cur_ts = oth_ts;
                            cur_tb = oth_tb;
                        }
                        float oth_bs = __shfl_xor_sync(0xffffffff, cur_bs, mask);
                        int   oth_bb = __shfl_xor_sync(0xffffffff, cur_bb, mask);
                        if (oth_bs < cur_bs || (oth_bs == cur_bs && oth_bb < cur_bb)) {
                            cur_bs = oth_bs;
                            cur_bb = oth_bb;
                        }
                    }
                    min_ts = cur_ts;
                    min_tb = cur_tb;
                    min_bs = cur_bs;
                    min_bb = cur_bb;
                }
            }

            if (valid_q_top) {
                if (tb0 >= 0 && tb0 < bi_curr) {
                    if (tb0 < 64) mask0_top |= (1ULL << tb0);
                    else          mask1_top |= (1ULL << (tb0 - 64));
                }
                if (tb1 >= 0 && tb1 < bi_curr) {
                    if (tb1 < 64) mask0_top |= (1ULL << tb1);
                    else          mask1_top |= (1ULL << (tb1 - 64));
                }
            }
            if (valid_q_bot) {
                if (bb0 >= 0 && bb0 < bi_curr) {
                    if (bb0 < 64) mask0_bot |= (1ULL << bb0);
                    else          mask1_bot |= (1ULL << (bb0 - 64));
                }
                if (bb1 >= 0 && bb1 < bi_curr) {
                    if (bb1 < 64) mask0_bot |= (1ULL << bb1);
                    else          mask1_bot |= (1ULL << (bb1 - 64));
                }
            }
            #pragma unroll
            for (int mask = 1; mask <= 2; mask <<= 1) {
                mask0_top |= __shfl_xor_sync(0xffffffff, mask0_top, mask);
                mask1_top |= __shfl_xor_sync(0xffffffff, mask1_top, mask);
                mask0_bot |= __shfl_xor_sync(0xffffffff, mask0_bot, mask);
                mask1_bot |= __shfl_xor_sync(0xffffffff, mask1_bot, mask);
            }
        }
    }

    // Step 4: Handle block bi_curr - 1 (already prefetched into smem.k_tile and smem.v_tile!)
    if (bi_curr >= 1) {
        const int prev_bi = bi_curr - 1;
        cp_async_wait<0>();
        __syncthreads();

        if (t_warp_base < S) {
            bool has_prev_top = (prev_bi < 64) ? ((mask0_top >> prev_bi) & 1ULL) : ((mask1_top >> (prev_bi - 64)) & 1ULL);
            bool has_prev_bot = (prev_bi < 64) ? ((mask0_bot >> prev_bi) & 1ULL) : ((mask1_bot >> (prev_bi - 64)) & 1ULL);

            if (prev_bi < 64) {
                mask0_top &= ~(1ULL << prev_bi);
                mask0_bot &= ~(1ULL << prev_bi);
            } else {
                mask1_top &= ~(1ULL << (prev_bi - 64));
                mask1_bot &= ~(1ULL << (prev_bi - 64));
            }

            const int min_k_top = has_prev_top ? 0 : (q_top_local + 1);
            const int min_k_bot = has_prev_bot ? 0 : (q_bot_local + 1);

            int min_k_warp = min(valid_q_top ? min_k_top : 64, valid_q_bot ? min_k_bot : 64);
            #pragma unroll
            for (int mask = 16; mask >= 4; mask >>= 1) {
                min_k_warp = min(min_k_warp, __shfl_xor_sync(0xffffffff, min_k_warp, mask));
            }

            for (int k_chunk = (min_k_warp >> 4); k_chunk < 4; ++k_chunk) {
                const int k_base = k_chunk * 16;

                float s_0[4] = {0.f, 0.f, 0.f, 0.f};
                float s_1[4] = {0.f, 0.f, 0.f, 0.f};
                #pragma unroll
                for (int tk = 0; tk < TK; ++tk) {
                    uint32_t r0, r1, r2, r3;
                    uint32_t addr = smem_u32addr(&smem.k_tile[k_base + sub_r][tk * 16 + sub_c]);
                    ldmatrix_x4(r0, r1, r2, r3, addr);
                    mma_m16n8k16_bf16(s_0[0], s_0[1], s_0[2], s_0[3],
                                      q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r0, r2);
                    mma_m16n8k16_bf16(s_1[0], s_1[1], s_1[2], s_1[3],
                                      q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r1, r3);
                }

                const int kc0 = k_base + tid * 2 + 0;
                const int kc1 = k_base + tid * 2 + 1;
                const int kc2 = k_base + 8 + tid * 2 + 0;
                const int kc3 = k_base + 8 + tid * 2 + 1;

                bool vt0 = valid_q_top && (kc0 >= min_k_top);
                bool vt1 = valid_q_top && (kc1 >= min_k_top);
                bool vt2 = valid_q_top && (kc2 >= min_k_top);
                bool vt3 = valid_q_top && (kc3 >= min_k_top);

                bool vb0 = valid_q_bot && (kc0 >= min_k_bot);
                bool vb1 = valid_q_bot && (kc1 >= min_k_bot);
                bool vb2 = valid_q_bot && (kc2 >= min_k_bot);
                bool vb3 = valid_q_bot && (kc3 >= min_k_bot);

                float st0 = vt0 ? (s_0[0] * scale) : -1e30f;
                float st1 = vt1 ? (s_0[1] * scale) : -1e30f;
                float st2 = vt2 ? (s_1[0] * scale) : -1e30f;
                float st3 = vt3 ? (s_1[1] * scale) : -1e30f;

                float sb0 = vb0 ? (s_0[2] * scale) : -1e30f;
                float sb1 = vb1 ? (s_0[3] * scale) : -1e30f;
                float sb2 = vb2 ? (s_1[2] * scale) : -1e30f;
                float sb3 = vb3 ? (s_1[3] * scale) : -1e30f;

                float m_chunk_top = fmaxf(fmaxf(st0, st1), fmaxf(st2, st3));
                m_chunk_top = fmaxf(m_chunk_top, __shfl_xor_sync(0xffffffff, m_chunk_top, 1));
                m_chunk_top = fmaxf(m_chunk_top, __shfl_xor_sync(0xffffffff, m_chunk_top, 2));

                float m_chunk_bot = fmaxf(fmaxf(sb0, sb1), fmaxf(sb2, sb3));
                m_chunk_bot = fmaxf(m_chunk_bot, __shfl_xor_sync(0xffffffff, m_chunk_bot, 1));
                m_chunk_bot = fmaxf(m_chunk_bot, __shfl_xor_sync(0xffffffff, m_chunk_bot, 2));

                float m_new_top = fmaxf(m_top, m_chunk_top);
                float exp_old_top = (m_new_top > m_top) ? __expf(m_top - m_new_top) : 1.0f;
                float pt0 = vt0 ? __expf(st0 - m_new_top) : 0.0f;
                float pt1 = vt1 ? __expf(st1 - m_new_top) : 0.0f;
                float pt2 = vt2 ? __expf(st2 - m_new_top) : 0.0f;
                float pt3 = vt3 ? __expf(st3 - m_new_top) : 0.0f;
                float p_sum_top = pt0 + pt1 + pt2 + pt3;
                p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 1);
                p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 2);
                l_top = l_top * exp_old_top + p_sum_top;
                m_top = m_new_top;

                float m_new_bot = fmaxf(m_bot, m_chunk_bot);
                float exp_old_bot = (m_new_bot > m_bot) ? __expf(m_bot - m_new_bot) : 1.0f;
                float pb0 = vb0 ? __expf(sb0 - m_new_bot) : 0.0f;
                float pb1 = vb1 ? __expf(sb1 - m_new_bot) : 0.0f;
                float pb2 = vb2 ? __expf(sb2 - m_new_bot) : 0.0f;
                float pb3 = vb3 ? __expf(sb3 - m_new_bot) : 0.0f;
                float p_sum_bot = pb0 + pb1 + pb2 + pb3;
                p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 1);
                p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 2);
                l_bot = l_bot * exp_old_bot + p_sum_bot;
                m_bot = m_new_bot;

                #pragma unroll
                for (int tn = 0; tn < TN; ++tn) {
                    o_frag[tn][0] *= exp_old_top;
                    o_frag[tn][1] *= exp_old_top;
                    o_frag[tn][2] *= exp_old_bot;
                    o_frag[tn][3] *= exp_old_bot;
                }

                uint32_t p_frag[4];
                *reinterpret_cast<__nv_bfloat162*>(&p_frag[0]) = __floats2bfloat162_rn(pt0, pt1);
                *reinterpret_cast<__nv_bfloat162*>(&p_frag[1]) = __floats2bfloat162_rn(pb0, pb1);
                *reinterpret_cast<__nv_bfloat162*>(&p_frag[2]) = __floats2bfloat162_rn(pt2, pt3);
                *reinterpret_cast<__nv_bfloat162*>(&p_frag[3]) = __floats2bfloat162_rn(pb2, pb3);

                #pragma unroll
                for (int tn2 = 0; tn2 < TN / 2; ++tn2) {
                    uint32_t r0, r1, r2, r3;
                    uint32_t addr = smem_u32addr(&smem.v_tile[k_base + sub_r][tn2 * 16 + sub_c]);
                    ldmatrix_x4_trans(r0, r1, r2, r3, addr);
                    mma_m16n8k16_bf16(o_frag[2 * tn2 + 0][0], o_frag[2 * tn2 + 0][1], o_frag[2 * tn2 + 0][2], o_frag[2 * tn2 + 0][3],
                                      p_frag[0], p_frag[1], p_frag[2], p_frag[3], r0, r1);
                    mma_m16n8k16_bf16(o_frag[2 * tn2 + 1][0], o_frag[2 * tn2 + 1][1], o_frag[2 * tn2 + 1][2], o_frag[2 * tn2 + 1][3],
                                      p_frag[0], p_frag[1], p_frag[2], p_frag[3], r2, r3);
                }
            }
        }
    }

    // Step 3: Cooperative CTA union across all 4 warps with software-pipelined cp.async.ca!
    uint64_t w_mask0 = mask0_top | mask0_bot;
    uint64_t w_mask1 = mask1_top | mask1_bot;
    #pragma unroll
    for (int mask = 16; mask >= 4; mask >>= 1) {
        w_mask0 |= __shfl_xor_sync(0xffffffff, w_mask0, mask);
        w_mask1 |= __shfl_xor_sync(0xffffffff, w_mask1, mask);
    }

    if (lane == 0) {
        smem.warp_mask0[w] = w_mask0;
        smem.warp_mask1[w] = w_mask1;
    }
    __syncthreads();

    uint64_t cta_mask0 = smem.warp_mask0[0] | smem.warp_mask0[1] | smem.warp_mask0[2] | smem.warp_mask0[3];
    uint64_t cta_mask1 = smem.warp_mask1[0] | smem.warp_mask1[1] | smem.warp_mask1[2] | smem.warp_mask1[3];

    if (cta_mask0 != 0ULL || cta_mask1 != 0ULL) {
        int bi;
        bool active_top, active_bot, warp_active;
        if (cta_mask0 != 0ULL) {
            bi = __ffsll(cta_mask0) - 1;
            cta_mask0 &= (cta_mask0 - 1ULL);
            active_top = (mask0_top >> bi) & 1ULL;
            active_bot = (mask0_bot >> bi) & 1ULL;
            warp_active = (w_mask0 >> bi) & 1ULL;
        } else {
            int bit = __ffsll(cta_mask1) - 1;
            bi = 64 + bit;
            cta_mask1 &= (cta_mask1 - 1ULL);
            active_top = (mask1_top >> bit) & 1ULL;
            active_bot = (mask1_bot >> bit) & 1ULL;
            warp_active = (w_mask1 >> bit) & 1ULL;
        }

        // Prologue: issue K[bi_0] and V[bi_0]
        issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.k_tile, k_head + static_cast<int64_t>(bi) * 64 * D, tid_cta);
        cp_async_commit();
        issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.v_tile, v_head + static_cast<int64_t>(bi) * 64 * D, tid_cta);
        cp_async_commit();
        // Wait for K[bi_0] (group 1), leaving V[bi_0] (group 0) in flight!
        cp_async_wait<1>();
        __syncthreads();

        while (true) {
            SPFrag sp;

            if (warp_active) {
                // Compute S = Q @ K[bi].T while V[bi] is in flight!
                #pragma unroll
                for (int kc = 0; kc < 4; ++kc) {
                    const int k_base = kc * 16;
                    float s_0[4] = {0.f, 0.f, 0.f, 0.f};
                    float s_1[4] = {0.f, 0.f, 0.f, 0.f};
                    #pragma unroll
                    for (int tk = 0; tk < TK; ++tk) {
                        uint32_t r0, r1, r2, r3;
                        uint32_t addr = smem_u32addr(&smem.k_tile[k_base + sub_r][tk * 16 + sub_c]);
                        ldmatrix_x4(r0, r1, r2, r3, addr);
                        mma_m16n8k16_bf16(s_0[0], s_0[1], s_0[2], s_0[3],
                                          q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r0, r2);
                        mma_m16n8k16_bf16(s_1[0], s_1[1], s_1[2], s_1[3],
                                          q_frag[tk][0], q_frag[tk][1], q_frag[tk][2], q_frag[tk][3], r1, r3);
                    }

                    sp.s[2 * kc + 0][0] = s_0[0] * scale;
                    sp.s[2 * kc + 0][1] = s_0[1] * scale;
                    sp.s[2 * kc + 0][2] = s_0[2] * scale;
                    sp.s[2 * kc + 0][3] = s_0[3] * scale;

                    sp.s[2 * kc + 1][0] = s_1[0] * scale;
                    sp.s[2 * kc + 1][1] = s_1[1] * scale;
                    sp.s[2 * kc + 1][2] = s_1[2] * scale;
                    sp.s[2 * kc + 1][3] = s_1[3] * scale;
                }

                float m_blk_top = -1e30f;
                float m_blk_bot = -1e30f;
                #pragma unroll
                for (int tn = 0; tn < 8; ++tn) {
                    m_blk_top = fmaxf(m_blk_top, fmaxf(sp.s[tn][0], sp.s[tn][1]));
                    m_blk_bot = fmaxf(m_blk_bot, fmaxf(sp.s[tn][2], sp.s[tn][3]));
                }
                m_blk_top = fmaxf(m_blk_top, __shfl_xor_sync(0xffffffff, m_blk_top, 1));
                m_blk_top = fmaxf(m_blk_top, __shfl_xor_sync(0xffffffff, m_blk_top, 2));
                m_blk_bot = fmaxf(m_blk_bot, __shfl_xor_sync(0xffffffff, m_blk_bot, 1));
                m_blk_bot = fmaxf(m_blk_bot, __shfl_xor_sync(0xffffffff, m_blk_bot, 2));

                float m_new_top = active_top ? fmaxf(m_top, m_blk_top) : m_top;
                float exp_old_top = active_top ? __expf(m_top - m_new_top) : 1.0f;

                float m_new_bot = active_bot ? fmaxf(m_bot, m_blk_bot) : m_bot;
                float exp_old_bot = active_bot ? __expf(m_bot - m_new_bot) : 1.0f;

                float p_sum_top = 0.0f;
                float p_sum_bot = 0.0f;
                #pragma unroll
                for (int kc = 0; kc < 4; ++kc) {
                    float pt0 = active_top ? __expf(sp.s[2 * kc + 0][0] - m_new_top) : 0.0f;
                    float pt1 = active_top ? __expf(sp.s[2 * kc + 0][1] - m_new_top) : 0.0f;
                    float pt2 = active_top ? __expf(sp.s[2 * kc + 1][0] - m_new_top) : 0.0f;
                    float pt3 = active_top ? __expf(sp.s[2 * kc + 1][1] - m_new_top) : 0.0f;
                    p_sum_top += pt0 + pt1 + pt2 + pt3;

                    float pb0 = active_bot ? __expf(sp.s[2 * kc + 0][2] - m_new_bot) : 0.0f;
                    float pb1 = active_bot ? __expf(sp.s[2 * kc + 0][3] - m_new_bot) : 0.0f;
                    float pb2 = active_bot ? __expf(sp.s[2 * kc + 1][2] - m_new_bot) : 0.0f;
                    float pb3 = active_bot ? __expf(sp.s[2 * kc + 1][3] - m_new_bot) : 0.0f;
                    p_sum_bot += pb0 + pb1 + pb2 + pb3;

                    uint32_t pf0, pf1, pf2, pf3;
                    *reinterpret_cast<__nv_bfloat162*>(&pf0) = __floats2bfloat162_rn(pt0, pt1);
                    *reinterpret_cast<__nv_bfloat162*>(&pf1) = __floats2bfloat162_rn(pb0, pb1);
                    *reinterpret_cast<__nv_bfloat162*>(&pf2) = __floats2bfloat162_rn(pt2, pt3);
                    *reinterpret_cast<__nv_bfloat162*>(&pf3) = __floats2bfloat162_rn(pb2, pb3);
                    sp.p[kc][0] = pf0;
                    sp.p[kc][1] = pf1;
                    sp.p[kc][2] = pf2;
                    sp.p[kc][3] = pf3;
                }

                p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 1);
                p_sum_top += __shfl_xor_sync(0xffffffff, p_sum_top, 2);
                l_top = l_top * exp_old_top + p_sum_top;
                m_top = m_new_top;

                p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 1);
                p_sum_bot += __shfl_xor_sync(0xffffffff, p_sum_bot, 2);
                l_bot = l_bot * exp_old_bot + p_sum_bot;
                m_bot = m_new_bot;

                #pragma unroll
                for (int tn = 0; tn < TN; ++tn) {
                    o_frag[tn][0] *= exp_old_top;
                    o_frag[tn][1] *= exp_old_top;
                    o_frag[tn][2] *= exp_old_bot;
                    o_frag[tn][3] *= exp_old_bot;
                }
            }

            // Wait for V[bi] to arrive and ensure all warps are done reading smem.k_tile
            cp_async_wait<0>();
            __syncthreads();

            const bool has_next = (cta_mask0 != 0ULL || cta_mask1 != 0ULL);
            int next_bi = 0;
            bool next_active_top = false, next_active_bot = false, next_warp_active = false;
            if (has_next) {
                if (cta_mask0 != 0ULL) {
                    next_bi = __ffsll(cta_mask0) - 1;
                    cta_mask0 &= (cta_mask0 - 1ULL);
                    next_active_top = (mask0_top >> next_bi) & 1ULL;
                    next_active_bot = (mask0_bot >> next_bi) & 1ULL;
                    next_warp_active = (w_mask0 >> next_bi) & 1ULL;
                } else {
                    int bit = __ffsll(cta_mask1) - 1;
                    next_bi = 64 + bit;
                    cta_mask1 &= (cta_mask1 - 1ULL);
                    next_active_top = (mask1_top >> bit) & 1ULL;
                    next_active_bot = (mask1_bot >> bit) & 1ULL;
                    next_warp_active = (w_mask1 >> bit) & 1ULL;
                }
                // Issue K[next_bi] into smem.k_tile WHILE warps compute P @ V[bi]!
                issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.k_tile, k_head + static_cast<int64_t>(next_bi) * 64 * D, tid_cta);
                cp_async_commit();
            }

            if (warp_active) {
                #pragma unroll
                for (int kc = 0; kc < 4; ++kc) {
                    const int k_base = kc * 16;
                    #pragma unroll
                    for (int tn2 = 0; tn2 < TN / 2; ++tn2) {
                        uint32_t r0, r1, r2, r3;
                        uint32_t addr = smem_u32addr(&smem.v_tile[k_base + sub_r][tn2 * 16 + sub_c]);
                        ldmatrix_x4_trans(r0, r1, r2, r3, addr);
                        mma_m16n8k16_bf16(o_frag[2 * tn2 + 0][0], o_frag[2 * tn2 + 0][1], o_frag[2 * tn2 + 0][2], o_frag[2 * tn2 + 0][3],
                                          sp.p[kc][0], sp.p[kc][1], sp.p[kc][2], sp.p[kc][3], r0, r1);
                        mma_m16n8k16_bf16(o_frag[2 * tn2 + 1][0], o_frag[2 * tn2 + 1][1], o_frag[2 * tn2 + 1][2], o_frag[2 * tn2 + 1][3],
                                          sp.p[kc][0], sp.p[kc][1], sp.p[kc][2], sp.p[kc][3], r2, r3);
                    }
                }
            }

            if (!has_next) break;

            // Wait for K[next_bi] and ensure all warps are done reading smem.v_tile
            cp_async_wait<0>();
            __syncthreads();

            // Issue V[next_bi] into smem.v_tile WHILE warps compute Q @ K[next_bi].T!
            issue_cta_tile_cp_async<D, CTA_LOAD_ITERS, UINT4_PER_ROW>(smem.v_tile, v_head + static_cast<int64_t>(next_bi) * 64 * D, tid_cta);
            cp_async_commit();

            bi = next_bi;
            active_top = next_active_top;
            active_bot = next_active_bot;
            warp_active = next_warp_active;
        }
    }

    // Step 5: Normalize o_frag by l_top / l_bot, stage via smem.k_tile for coalesced uint4 global writes
    __syncthreads();
    if (t_warp_base < S) {
        float inv_l_top = (l_top > 0.0f) ? (1.0f / l_top) : 0.0f;
        float inv_l_bot = (l_bot > 0.0f) ? (1.0f / l_bot) : 0.0f;

        #pragma unroll
        for (int tn = 0; tn < TN; ++tn) {
            __nv_bfloat162 b_top = __floats2bfloat162_rn(o_frag[tn][0] * inv_l_top, o_frag[tn][1] * inv_l_top);
            __nv_bfloat162 b_bot = __floats2bfloat162_rn(o_frag[tn][2] * inv_l_bot, o_frag[tn][3] * inv_l_bot);
            *reinterpret_cast<__nv_bfloat162*>(&smem.k_tile[w * 16 + groupID][tn * 8 + tid * 2]) = b_top;
            *reinterpret_cast<__nv_bfloat162*>(&smem.k_tile[w * 16 + groupID + 8][tn * 8 + tid * 2]) = b_bot;
        }
    }
    __syncthreads();

    #pragma unroll
    for (int i = 0; i < CTA_LOAD_ITERS; ++i) {
        int idx = tid_cta + i * 128;
        int r = idx / UINT4_PER_ROW;
        int c = (idx % UINT4_PER_ROW) * 8;
        int t_r = s0 + r;
        if (t_r < S) {
            *reinterpret_cast<uint4*>(o_head + static_cast<int64_t>(t_r) * D + c) =
                *reinterpret_cast<const uint4*>(&smem.k_tile[r][c]);
        }
    }
}

struct KMeanBuffers {
    float* km_ptr = nullptr;
    uint32_t* flag_ptr = nullptr;
    uint32_t gen = 0;
};

KMeanBuffers& get_km_buffers(size_t needed_blocks, int D) {
    static KMeanBuffers bufs;
    static size_t cap_floats = 0;
    static size_t cap_flags = 0;
    size_t needed_floats = needed_blocks * D;
    if (needed_floats > cap_floats || needed_blocks > cap_flags) {
        if (bufs.km_ptr) cudaFree(bufs.km_ptr);
        if (bufs.flag_ptr) cudaFree(bufs.flag_ptr);
        cap_flags = std::max(needed_blocks, static_cast<size_t>(32 * 128));
        cap_floats = std::max(needed_floats, static_cast<size_t>(32 * 128 * 128));
        cudaMalloc(&bufs.km_ptr, cap_floats * sizeof(float));
        cudaMalloc(&bufs.flag_ptr, cap_flags * sizeof(uint32_t));
        cudaMemset(bufs.flag_ptr, 0, cap_flags * sizeof(uint32_t));
    }
    ++bufs.gen;
    return bufs;
}

} // namespace

torch::Tensor nsa_forward_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
    const int B = q.size(0);
    const int H = q.size(1);
    const int S = q.size(2);
    const int D = q.size(3);
    const int BH = B * H;
    const int n_full_blocks = S / 64;
    const int n_blocks = (S + 63) / 64;
    const float scale = 1.0f / std::sqrt(static_cast<float>(D));

    auto o = at::detail::empty_cuda(q.sizes(), at::ScalarType::BFloat16, q.device(), std::nullopt);
    auto& bufs = get_km_buffers(static_cast<size_t>(BH) * std::max(n_full_blocks, 1), D);

    const auto* q_ptr = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr());
    const auto* k_ptr = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr());
    const auto* v_ptr = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr());
    auto* o_ptr = reinterpret_cast<__nv_bfloat16*>(o.data_ptr());

    dim3 grid_attn(n_blocks, BH);
    dim3 block_attn(128);

    if (D == 64) {
        nsa_cta_coop_kernel<64, 4, 3><<<grid_attn, block_attn>>>(
            q_ptr, k_ptr, v_ptr, bufs.km_ptr, bufs.flag_ptr, bufs.gen, o_ptr, S, n_full_blocks, scale
        );
    } else if (D == 128) {
        nsa_cta_coop_kernel<128, 4, 2><<<grid_attn, block_attn>>>(
            q_ptr, k_ptr, v_ptr, bufs.km_ptr, bufs.flag_ptr, bufs.gen, o_ptr, S, n_full_blocks, scale
        );
    }
    return o;
}
"""

cpp_source = "torch::Tensor nsa_forward_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v);"

_ext = load_inline(
    name="nsa_cuda_ext_v10",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["nsa_forward_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_120a"],
    verbose=False,
)


class Model(nn.Module):
    def __init__(self, B: int, H: int, S: int, D: int):
        super().__init__()
        self.B, self.H, self.S, self.D = B, H, S, D
        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16))

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        return _ext.nsa_forward_cuda(q, k, v)


def get_init_inputs():
    return [B, H, S, D]


def get_inputs():
    q = torch.randn(B, H, S, D, dtype=torch.bfloat16)
    k = torch.randn(B, H, S, D, dtype=torch.bfloat16)
    v = torch.randn(B, H, S, D, dtype=torch.bfloat16)
    return [q, k, v]