"""CUDA implementation of the benchmark's simplified block-sparse attention."""
from __future__ import annotations

import os

import torch
import torch.nn as nn


_EXT = None


def _extension():
    global _EXT
    if _EXT is None:
        # This problem targets the RTX PRO 6000 (Blackwell sm_120).
        os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
        from torch.utils.cpp_extension import load_inline

        cpp_source = r"""
        #include <torch/extension.h>
        torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v);
        PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            m.def("forward", &nsa_forward, "Simplified NSA sparse attention (CUDA)");
        }
        """

        cuda_source = r"""
        #include <torch/extension.h>
        #include <ATen/cuda/CUDAContext.h>
        #include <cuda.h>
        #include <cuda_bf16.h>
        #include <float.h>

        namespace {
        constexpr int BLOCK = 64;
        constexpr int TOP_N = 8;
        constexpr int WINDOW = 64;
        constexpr int MAX_BLOCKS = 256;
        constexpr int MAX_KEYS = TOP_N * BLOCK + WINDOW;
        constexpr int THREADS = 256;
        constexpr int WARPS = THREADS / 32;

        __device__ __forceinline__ float warp_sum(float x) {
            #pragma unroll
            for (int offset = 16; offset > 0; offset >>= 1)
                x += __shfl_down_sync(0xffffffff, x, offset);
            return x;
        }

        __device__ __forceinline__ bool better(float av, int ai, float bv, int bi) {
            return av > bv || (av == bv && ai > bi);
        }

        template<int D>
        __global__ void prefix_means_kernel(const __nv_bfloat16* __restrict__ k,
                                             float* __restrict__ prefix,
                                             int H, int S) {
            const int bi = blockIdx.x;
            const int h = blockIdx.y;
            const int b = blockIdx.z;
            const int d = threadIdx.x;
            const int s0 = bi * BLOCK;
            const int s1 = min(s0 + BLOCK, S);
            if (d >= D || s0 >= S) return;
            const int64_t base = ((int64_t)b * H + h) * S * D;
            float sum = 0.0f;
            #pragma unroll 1
            for (int s = s0; s < s1; ++s) {
                sum += __bfloat162float(k[base + (int64_t)s * D + d]);
                prefix[base + (int64_t)s * D + d] = sum / float(s - s0 + 1);
            }
        }

        template<int D>
        __global__ void sparse_attention_kernel(const __nv_bfloat16* __restrict__ q,
                                                 const __nv_bfloat16* __restrict__ k,
                                                 const __nv_bfloat16* __restrict__ v,
                                                 const float* __restrict__ prefix,
                                                 __nv_bfloat16* __restrict__ out,
                                                 int B, int H, int S) {
            const int t = blockIdx.x;
            const int h = blockIdx.y;
            const int b = blockIdx.z;
            const int tid = threadIdx.x;
            const int lane = tid & 31;
            const int warp = tid >> 5;
            const int nblocks = (S + BLOCK - 1) / BLOCK;
            const int current_block = t / BLOCK;
            const int causal_blocks = current_block + 1;
            const int top_count = causal_blocks < TOP_N ? causal_blocks : TOP_N;
            const int64_t base = ((int64_t)b * H + h) * S * D;
            const int64_t qbase = base + (int64_t)t * D;
            const float scale = rsqrtf(float(D));

            __shared__ float block_scores[MAX_BLOCKS];
            __shared__ int selected[MAX_BLOCKS];
            __shared__ float warp_values[WARPS];
            __shared__ int warp_indices[WARPS];
            __shared__ int best_index;
            __shared__ int key_indices[MAX_KEYS];
            __shared__ float key_scores[MAX_KEYS];
            __shared__ int nkeys;
            __shared__ float row_max;
            __shared__ float row_sum;

            const float q0 = __bfloat162float(q[qbase + lane]);
            const float q1 = __bfloat162float(q[qbase + lane + 32]);
            float q2 = 0.0f, q3 = 0.0f;
            if constexpr (D == 128) {
                q2 = __bfloat162float(q[qbase + lane + 64]);
                q3 = __bfloat162float(q[qbase + lane + 96]);
            }

            if (tid < MAX_BLOCKS) {
                block_scores[tid] = -FLT_MAX;
                selected[tid] = 0;
            }
            __syncthreads();

            // One warp scores each block mean; earlier blocks use their full
            // prefix and the query's own block uses the causal prefix at t.
            for (int bi = warp; bi <= current_block; bi += WARPS) {
                const int end = bi == current_block ? t : (bi + 1) * BLOCK - 1;
                const int64_t mbase = base + (int64_t)end * D;
                float dot = q0 * prefix[mbase + lane];
                dot = fmaf(q1, prefix[mbase + lane + 32], dot);
                if constexpr (D == 128) {
                    dot = fmaf(q2, prefix[mbase + lane + 64], dot);
                    dot = fmaf(q3, prefix[mbase + lane + 96], dot);
                }
                dot = warp_sum(dot) * scale;
                if (lane == 0) block_scores[bi] = dot;
            }
            __syncthreads();

            // Parallel argmax repeated eight times gives the exact top-eight
            // block set without materializing or sorting a dense score row.
            for (int rank = 0; rank < top_count; ++rank) {
                float best = -FLT_MAX;
                int best_bi = -1;
                for (int bi = tid; bi < causal_blocks; bi += THREADS) {
                    const float candidate = selected[bi] ? -FLT_MAX : block_scores[bi];
                    if (better(candidate, bi, best, best_bi)) {
                        best = candidate;
                        best_bi = bi;
                    }
                }
                #pragma unroll
                for (int offset = 16; offset > 0; offset >>= 1) {
                    const float other_v = __shfl_down_sync(0xffffffff, best, offset);
                    const int other_i = __shfl_down_sync(0xffffffff, best_bi, offset);
                    if (better(other_v, other_i, best, best_bi)) {
                        best = other_v;
                        best_bi = other_i;
                    }
                }
                if (lane == 0) {
                    warp_values[warp] = best;
                    warp_indices[warp] = best_bi;
                }
                __syncthreads();
                if (warp == 0) {
                    float vbest = lane < WARPS ? warp_values[lane] : -FLT_MAX;
                    int ibest = lane < WARPS ? warp_indices[lane] : -1;
                    #pragma unroll
                    for (int offset = 16; offset > 0; offset >>= 1) {
                        const float other_v = __shfl_down_sync(0xffffffff, vbest, offset);
                        const int other_i = __shfl_down_sync(0xffffffff, ibest, offset);
                        if (better(other_v, other_i, vbest, ibest)) {
                            vbest = other_v;
                            ibest = other_i;
                        }
                    }
                    if (lane == 0) best_index = ibest;
                }
                __syncthreads();
                if (tid == 0) selected[best_index] = 1;
                __syncthreads();
            }

            if (tid == 0) {
                int count = 0;
                for (int bi = 0; bi <= current_block; ++bi) {
                    if (selected[bi]) {
                        const int s0 = bi * BLOCK;
                        const int s1 = min(s0 + BLOCK, t + 1);
                        for (int s = s0; s < s1; ++s) key_indices[count++] = s;
                    }
                }
                const int w0 = max(0, t + 1 - WINDOW);
                for (int s = w0; s <= t; ++s) {
                    if (!selected[s / BLOCK]) key_indices[count++] = s;
                }
                nkeys = count;
            }
            __syncthreads();

            // Compute selected-token logits, eight keys at a time.
            for (int si = warp; si < nkeys; si += WARPS) {
                const int s = key_indices[si];
                const int64_t kbase = base + (int64_t)s * D;
                float dot = q0 * __bfloat162float(k[kbase + lane]);
                dot = fmaf(q1, __bfloat162float(k[kbase + lane + 32]), dot);
                if constexpr (D == 128) {
                    dot = fmaf(q2, __bfloat162float(k[kbase + lane + 64]), dot);
                    dot = fmaf(q3, __bfloat162float(k[kbase + lane + 96]), dot);
                }
                dot = warp_sum(dot) * scale;
                if (lane == 0) key_scores[si] = dot;
            }
            __syncthreads();

            float local_max = -FLT_MAX;
            for (int si = tid; si < nkeys; si += THREADS)
                local_max = fmaxf(local_max, key_scores[si]);
            #pragma unroll
            for (int offset = 16; offset > 0; offset >>= 1)
                local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
            if (lane == 0) warp_values[warp] = local_max;
            __syncthreads();
            if (warp == 0) {
                float m = lane < WARPS ? warp_values[lane] : -FLT_MAX;
                #pragma unroll
                for (int offset = 16; offset > 0; offset >>= 1)
                    m = fmaxf(m, __shfl_down_sync(0xffffffff, m, offset));
                if (lane == 0) row_max = m;
            }
            __syncthreads();

            float local_sum = 0.0f;
            for (int si = tid; si < nkeys; si += THREADS) {
                const float weight = expf(key_scores[si] - row_max);
                key_scores[si] = weight;
                local_sum += weight;
            }
            local_sum = warp_sum(local_sum);
            if (lane == 0) warp_values[warp] = local_sum;
            __syncthreads();
            if (warp == 0) {
                float sum = lane < WARPS ? warp_values[lane] : 0.0f;
                #pragma unroll
                for (int offset = 16; offset > 0; offset >>= 1)
                    sum += __shfl_down_sync(0xffffffff, sum, offset);
                if (lane == 0) row_sum = sum;
            }
            __syncthreads();

            for (int d = tid; d < D; d += THREADS) {
                float acc = 0.0f;
                for (int si = 0; si < nkeys; ++si) {
                    const float weight = key_scores[si] / row_sum;
                    const int s = key_indices[si];
                    acc = fmaf(weight, __bfloat162float(v[base + (int64_t)s * D + d]), acc);
                }
                out[base + (int64_t)t * D + d] = __float2bfloat16_rn(acc);
            }
        }

        template<int D>
        torch::Tensor run_nsa(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
            const int B = static_cast<int>(q.size(0));
            const int H = static_cast<int>(q.size(1));
            const int S = static_cast<int>(q.size(2));
            auto prefix = torch::empty({B, H, S, D}, q.options().dtype(torch::kFloat32));
            auto out = torch::empty_like(q);
            const dim3 pgrid((S + BLOCK - 1) / BLOCK, H, B);
            const dim3 agrid(S, H, B);
            cudaStream_t stream = at::cuda::getCurrentCUDAStream();
            prefix_means_kernel<D><<<pgrid, D, 0, stream>>>(
                reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<at::BFloat16>()),
                prefix.data_ptr<float>(), H, S);
            cudaError_t launch_error = cudaGetLastError();
            TORCH_CHECK(launch_error == cudaSuccess,
                        "prefix-means CUDA launch failed: ", cudaGetErrorString(launch_error));
            sparse_attention_kernel<D><<<agrid, THREADS, 0, stream>>>(
                reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<at::BFloat16>()),
                reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<at::BFloat16>()),
                reinterpret_cast<const __nv_bfloat16*>(v.data_ptr<at::BFloat16>()),
                prefix.data_ptr<float>(),
                reinterpret_cast<__nv_bfloat16*>(out.data_ptr<at::BFloat16>()),
                B, H, S);
            launch_error = cudaGetLastError();
            TORCH_CHECK(launch_error == cudaSuccess,
                        "sparse-attention CUDA launch failed: ", cudaGetErrorString(launch_error));
            return out;
        }
        } // namespace

        torch::Tensor nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
            TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "q, k, v must be CUDA tensors");
            TORCH_CHECK(q.scalar_type() == torch::kBFloat16 &&
                        k.scalar_type() == torch::kBFloat16 &&
                        v.scalar_type() == torch::kBFloat16, "q, k, v must be bf16");
            TORCH_CHECK(q.dim() == 4 && q.sizes() == k.sizes() && q.sizes() == v.sizes(),
                        "q, k, v must have matching (B,H,S,D) shapes");
            TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(),
                        "q, k, v must be contiguous");
            TORCH_CHECK(q.size(3) == 64 || q.size(3) == 128, "D must be 64 or 128");
            TORCH_CHECK(q.size(2) <= MAX_BLOCKS * BLOCK, "S exceeds supported limit");
            if (q.size(3) == 64) return run_nsa<64>(q, k, v);
            return run_nsa<128>(q, k, v);
        }
        """

        _EXT = load_inline(
            name="deepseek_nsa_sparse_cuda_v2",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            with_cuda=True,
            verbose=False,
        )
    return _EXT


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 _extension().forward(q.contiguous(), k.contiguous(), v.contiguous())