"""DeepSeek NSA-inspired sparse attention — fused CUDA kernel (SM120).

Bench semantics (same as reference.nsa_attend): per query t, block importance is
the mean of q_t.k_j/sqrt(D) over causal keys in each 64-wide block; the top-8
blocks union the last-64 sliding window form the sparse set; softmax over that
set only.

Speed idea: mean(q.k) over a block == q.(mean k over block), so block scores
cost S/64 dots per query instead of S dots. Two kernels:
  1. kbar: per-block mean of K  -> (B*H*nb, D) fp32.
  2. nsa:  one CUDA block per (b,h,t); loads q, scores nb blocks, single-thread
     top-8 select (exact reference tie-break: score desc, block id desc),
     builds the deduped index set, two-pass online softmax over gathered keys,
     warp-cooperative V accumulate. No forbidden ops anywhere.
"""

from __future__ import annotations

import os

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline  # torch.utils.cpp_extension.load_inline

os.makedirs("/tmp/nsa_cuda_build", exist_ok=True)

_CU = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cfloat>

#define BS 64
#define TOPN 8
#define WIN 64
#define NT 128
#define MAXNB 512
#define MAXSEL 576

template<int D>
__global__ void kbar_kernel(const __nv_bfloat16* __restrict__ K,
                            float* __restrict__ KBAR, int S, int nb) {
  int job = blockIdx.x;
  int bi = job % nb;
  int bh = job / nb;
  int s0 = bi * BS;
  int s1 = s0 + BS < S ? s0 + BS : S;
  int tid = threadIdx.x;
  size_t kbase = (size_t)bh * (size_t)S * (size_t)D;
  size_t kbo = (size_t)job * (size_t)D;
  for (int d = tid; d < D; d += NT) {
    float acc = 0.f;
    for (int j = s0; j < s1; ++j)
      acc += __bfloat162float(K[kbase + (size_t)j * (size_t)D + d]);
    KBAR[kbo + d] = acc / (float)(s1 - s0);
  }
}

template<int D>
__device__ __forceinline__ float dotqb(const float* q, const __nv_bfloat16* p) {
  float dot = 0.f;
  const float4* p4 = reinterpret_cast<const float4*>(p);
  #pragma unroll
  for (int i = 0; i < D / 8; ++i) {
    float4 c = p4[i];
    const __nv_bfloat162* b = reinterpret_cast<const __nv_bfloat162*>(&c);
    #pragma unroll
    for (int k = 0; k < 4; ++k) {
      float2 f = __bfloat1622float2(b[k]);
      int d = i * 8 + k * 2;
      dot += q[d] * f.x + q[d + 1] * f.y;
    }
  }
  return dot;
}

template<int D>
__device__ __forceinline__ float dotqf(const float* q, const float* p) {
  float dot = 0.f;
  const float4* p4 = reinterpret_cast<const float4*>(p);
  #pragma unroll
  for (int i = 0; i < D / 4; ++i) {
    float4 c = p4[i];
    int d = i * 4;
    dot += q[d] * c.x + q[d + 1] * c.y + q[d + 2] * c.z + q[d + 3] * c.w;
  }
  return dot;
}

template<int D>
__global__ void __launch_bounds__(128)
nsa_kernel(const __nv_bfloat16* __restrict__ Q,
           const __nv_bfloat16* __restrict__ K,
           const __nv_bfloat16* __restrict__ V,
           const float* __restrict__ KBAR,
           __nv_bfloat16* __restrict__ O,
           int S, int nb) {
  int t = blockIdx.x % S;
  int bh = blockIdx.x / S;
  int tid = threadIdx.x;
  size_t base = (size_t)bh * (size_t)S * (size_t)D;
  size_t kbar_base = (size_t)bh * (size_t)nb * (size_t)D;
  const float scale = 1.0f / sqrtf((float)D);

  __shared__ float q[128];
  __shared__ float bscore[MAXNB];
  __shared__ float sc[MAXSEL];
  __shared__ int idx[MAXSEL];
  __shared__ float red[NT];
  __shared__ int nsel;

  for (int d = tid; d < D; d += NT)
    q[d] = __bfloat162float(Q[base + (size_t)t * (size_t)D + d]);
  __syncthreads();

  int tb = t >> 6;
  for (int bi = tid; bi < nb; bi += NT) {
    int s0 = bi << 6;
    if (s0 > t) {
      bscore[bi] = -1e9f;
    } else if (bi != tb) {
      size_t kb = kbar_base + (size_t)bi * (size_t)D;
      bscore[bi] = dotqf<D>(q, KBAR + kb) * scale;
    }
  }
  __syncthreads();

  {
    int s0 = tb << 6;
    int cnt = t - s0 + 1;
    float part = 0.f;
    for (int j = s0 + tid; j <= t; j += NT)
      part += dotqb<D>(q, K + base + (size_t)j * (size_t)D);
    red[tid] = part;
    __syncthreads();
    for (int s = NT >> 1; s > 0; s >>= 1) {
      if (tid < s) red[tid] += red[tid + s];
      __syncthreads();
    }
    if (tid == 0) bscore[tb] = red[0] / (float)cnt * scale;
  }
  __syncthreads();

  if (tid == 0) {
    int top[TOPN];
    float tops[TOPN];
    int ntop = 0;
    for (int bi = 0; bi < nb; ++bi) {
      int s0 = bi << 6;
      if (s0 > t) continue;
      float s = bscore[bi];
      int pos = 0;
      while (pos < ntop && (tops[pos] > s || (tops[pos] == s && top[pos] > bi))) ++pos;
      if (pos < TOPN) {
        int up = ntop < TOPN ? ntop : TOPN - 1;
        for (int kk = up; kk > pos; --kk) { tops[kk] = tops[kk-1]; top[kk] = top[kk-1]; }
        tops[pos] = s; top[pos] = bi;
        if (ntop < TOPN) ++ntop;
      }
    }
    for (int i = 0; i < ntop; ++i)
      for (int j = i + 1; j < ntop; ++j)
        if (top[j] < top[i]) { int tmp = top[i]; top[i] = top[j]; top[j] = tmp; }
    int n = 0;
    for (int kk = 0; kk < ntop; ++kk) {
      int s0 = top[kk] << 6;
      int s1 = s0 + BS;
      if (s1 > t + 1) s1 = t + 1;
      for (int j = s0; j < s1; ++j) idx[n++] = j;
    }
    int w0 = t + 1 - WIN;
    if (w0 < 0) w0 = 0;
    for (int j = w0; j <= t; ++j) {
      int bj = j >> 6;
      int found = 0;
      for (int kk = 0; kk < ntop; ++kk) if (top[kk] == bj) { found = 1; break; }
      if (!found) idx[n++] = j;
    }
    nsel = n;
  }
  __syncthreads();
  int myn = nsel;

  for (int i = tid; i < myn; i += NT) {
    int j = idx[i];
    sc[i] = dotqb<D>(q, K + base + (size_t)j * (size_t)D) * scale;
  }
  __syncthreads();

  float mx = -FLT_MAX;
  for (int i = tid; i < myn; i += NT) mx = fmaxf(mx, sc[i]);
  red[tid] = mx;
  __syncthreads();
  for (int s = NT >> 1; s > 0; s >>= 1) {
    if (tid < s) red[tid] = fmaxf(red[tid], red[tid + s]);
    __syncthreads();
  }
  float gmax = red[0];
  __syncthreads();

  float psum = 0.f;
  for (int i = tid; i < myn; i += NT) {
    float e = expf(sc[i] - gmax);
    sc[i] = e;
    psum += e;
  }
  red[tid] = psum;
  __syncthreads();
  for (int s = NT >> 1; s > 0; s >>= 1) {
    if (tid < s) red[tid] += red[tid + s];
    __syncthreads();
  }
  float gsum = red[0];
  __syncthreads();
  float inv = 1.f / gsum;
  for (int i = tid; i < myn; i += NT) sc[i] *= inv;
  __syncthreads();

  if constexpr (D == 64) {
    if (tid < 32) {
      const __nv_bfloat162* V2 = reinterpret_cast<const __nv_bfloat162*>(V);
      __nv_bfloat162* O2 = reinterpret_cast<__nv_bfloat162*>(O);
      size_t base2 = base >> 1;
      float a0 = 0.f, a1 = 0.f;
      for (int i = 0; i < myn; ++i) {
        int j = idx[i];
        float2 f = __bfloat1622float2(V2[base2 + (size_t)j * 32 + tid]);
        float w = sc[i];
        a0 += w * f.x;
        a1 += w * f.y;
      }
      O2[base2 + (size_t)t * 32 + tid] = __floats2bfloat162_rn(a0, a1);
    }
  } else {
    if (tid < D) {
      float a = 0.f;
      for (int i = 0; i < myn; ++i) {
        int j = idx[i];
        a += sc[i] * __bfloat162float(V[base + (size_t)j * (size_t)D + tid]);
      }
      O[base + (size_t)t * (size_t)D + tid] = __float2bfloat16(a);
    }
  }
}

void nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o) {
  const int B = q.size(0), H = q.size(1), S = q.size(2), D = q.size(3);
  const int nb = (S + BS - 1) / BS;
  torch::Tensor kbar = torch::empty({(long)B * H * nb * D},
                                    q.options().dtype(torch::kFloat32));
  cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
  const __nv_bfloat16* Q = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr());
  const __nv_bfloat16* K = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr());
  const __nv_bfloat16* V = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr());
  __nv_bfloat16* O = reinterpret_cast<__nv_bfloat16*>(o.data_ptr());
  float* KB = kbar.data_ptr<float>();
  if (D == 64) {
    kbar_kernel<64><<<(unsigned)B * H * nb, NT, 0, stream>>>(K, KB, S, nb);
    nsa_kernel<64><<<(unsigned)B * H * S, NT, 0, stream>>>(Q, K, V, KB, O, S, nb);
  } else if (D == 128) {
    kbar_kernel<128><<<(unsigned)B * H * nb, NT, 0, stream>>>(K, KB, S, nb);
    nsa_kernel<128><<<(unsigned)B * H * S, NT, 0, stream>>>(Q, K, V, KB, O, S, nb);
  }
}
"""

_CPP = r"""
#include <torch/extension.h>
void nsa_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o);
"""

_mod = load_inline(
    name="nsa_cuda",
    cpp_sources=[_CPP],
    cuda_sources=[_CU],
    functions=["nsa_forward"],
    extra_cuda_cflags=["-O3"],
    build_directory="/tmp/nsa_cuda_build",
)

BLOCK_SIZE = 64
TOP_N_BLOCKS = 8
SLIDING_WINDOW = 64


def _eager_fallback(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    import math

    qf, kf, vf = q.float(), k.float(), v.float()
    B, H, S, D = qf.shape
    scale = 1.0 / math.sqrt(D)
    nb = (S + BLOCK_SIZE - 1) // BLOCK_SIZE
    pad = nb * BLOCK_SIZE - S
    if pad:
        kf = torch.nn.functional.pad(kf, (0, 0, 0, pad))
        vf = torch.nn.functional.pad(vf, (0, 0, 0, pad))
    scores = qf @ kf.transpose(-2, -1) * scale  # (B,H,S,nb*64)
    causal = torch.ones(S, nb * BLOCK_SIZE, dtype=torch.bool, device=q.device)
    causal = torch.tril(causal[:, : nb * BLOCK_SIZE][None, None, :, :].expand(B, H, S, nb * BLOCK_SIZE))
    num = scores.masked_fill(~causal[:, :, :, : nb * BLOCK_SIZE], 0).reshape(B, H, S, nb, BLOCK_SIZE).sum(-1)
    den = causal[:, :, :, : nb * BLOCK_SIZE].reshape(B, H, S, nb, BLOCK_SIZE).sum(-1).clamp_min(1).float()
    bimp = num / den
    s0 = torch.arange(nb, device=q.device) * BLOCK_SIZE
    bimp = bimp.masked_fill((s0[None, None, None, :] > torch.arange(S, device=q.device)[None, None, :, None]), float("-1e9"))
    top = torch.topk(bimp, k=min(TOP_N_BLOCKS, nb), dim=-1).indices  # (B,H,S,k)
    sel = torch.zeros(B, H, S, nb * BLOCK_SIZE, dtype=torch.bool, device=q.device)
    rng = torch.arange(nb * BLOCK_SIZE, device=q.device)
    for kk in range(top.shape[-1]):
        bi = top[..., kk]
        sel |= (rng[None, None, None, :] // BLOCK_SIZE == bi[..., None]) & causal[:, :, :, : nb * BLOCK_SIZE]
    w = torch.arange(S, device=q.device)
    sel |= ((w[None, None, None, :] <= w[None, None, :, None]) & ((w[None, None, :, None] - w[None, None, None, :]) < SLIDING_WINDOW))[:, :, :, : nb * BLOCK_SIZE]
    att = scores.masked_fill(~sel, float("-inf")).softmax(-1)
    out = (att @ vf[:, :, : nb * BLOCK_SIZE, :])[:, :, :S, :]
    return out.to(torch.bfloat16)


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:
        B, H, S, D = q.shape
        nb = (S + BLOCK_SIZE - 1) // BLOCK_SIZE
        if D not in (64, 128) or nb > 512 or q.device.type != "cuda":
            return _eager_fallback(q, k, v)
        q = q.contiguous() if not q.is_contiguous() else q
        k = k.contiguous() if not k.is_contiguous() else k
        v = v.contiguous() if not v.is_contiguous() else v
        o = torch.empty_like(q)
        _mod.nsa_forward(q, k, v, o)
        return o