import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
        out = torch.nn.functional.scaled_dot_product_attention(Q, K, V)
        return out

batch_size = 32
num_heads = 32
sequence_length = 512
embedding_dimension = 1024

def get_inputs():
    Q = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension)
    K = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension)
    V = torch.rand(batch_size, num_heads, sequence_length, embedding_dimension)
    return [Q, K, V]

def get_init_inputs():
    return []

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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor sdpa_forward(torch::Tensor Q, torch::Tensor K, torch::Tensor V);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>

#define LCONST 512
#define DCONST 1024
#define BR 6
#define NTHREADS 256

__inline__ __device__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);
    return v;
}

__device__ __forceinline__ float block_reduce_sum(float val, float* smem) {
    int tid = threadIdx.x;
    smem[tid] = val;
    __syncthreads();
    #pragma unroll
    for (int s = NTHREADS/2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    return smem[0];
}

__device__ __forceinline__ float block_reduce_max(float val, float* smem) {
    int tid = threadIdx.x;
    smem[tid] = val;
    __syncthreads();
    #pragma unroll
    for (int s = NTHREADS/2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        __syncthreads();
    }
    return smem[0];
}

__global__ __launch_bounds__(NTHREADS, 2)
void sdpa_kernel(const float* __restrict__ Q,
                 const float* __restrict__ K,
                 const float* __restrict__ V,
                 float* __restrict__ O,
                 int total_bh) {
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int qblock = blockIdx.x;
    const int bh = blockIdx.y;
    const int qbase = qblock * BR;

    const long head_off = ((long)bh) * LCONST * DCONST;
    const float* Qh = Q + head_off;
    const float* Kh = K + head_off;
    const float* Vh = V + head_off;
    float* Oh = O + head_off;

    __shared__ float qtile[BR * DCONST];
    __shared__ float scores[BR * LCONST];
    __shared__ float warp_sums[BR * 8];
    __shared__ float red[NTHREADS];
    __shared__ float inv_sum[BR];

    // Cache this small block of query rows once.  The last block has two invalid rows.
    for (int idx = tid; idx < BR * DCONST; idx += NTHREADS) {
        int r = idx / DCONST;
        int d = idx - r * DCONST;
        int q = qbase + r;
        qtile[idx] = (q < LCONST) ? Qh[q * DCONST + d] : 0.0f;
    }
    __syncthreads();

    constexpr float scale = 0.03125f; // 1/sqrt(1024)

    // Compute the BR x 512 score tile.  K is reused for all query rows in the block.
    for (int j = 0; j < LCONST; ++j) {
        float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f, acc3 = 0.f, acc4 = 0.f, acc5 = 0.f;
        const float* krow = Kh + j * DCONST;
        #pragma unroll
        for (int d = tid; d < DCONST; d += NTHREADS) {
            float kv = krow[d];
            acc0 = fmaf(qtile[0 * DCONST + d], kv, acc0);
            acc1 = fmaf(qtile[1 * DCONST + d], kv, acc1);
            acc2 = fmaf(qtile[2 * DCONST + d], kv, acc2);
            acc3 = fmaf(qtile[3 * DCONST + d], kv, acc3);
            acc4 = fmaf(qtile[4 * DCONST + d], kv, acc4);
            acc5 = fmaf(qtile[5 * DCONST + d], kv, acc5);
        }
        acc0 = warp_reduce_sum(acc0);
        acc1 = warp_reduce_sum(acc1);
        acc2 = warp_reduce_sum(acc2);
        acc3 = warp_reduce_sum(acc3);
        acc4 = warp_reduce_sum(acc4);
        acc5 = warp_reduce_sum(acc5);
        if (lane == 0) {
            warp_sums[0 * 8 + warp] = acc0;
            warp_sums[1 * 8 + warp] = acc1;
            warp_sums[2 * 8 + warp] = acc2;
            warp_sums[3 * 8 + warp] = acc3;
            warp_sums[4 * 8 + warp] = acc4;
            warp_sums[5 * 8 + warp] = acc5;
        }
        __syncthreads();
        if (tid < BR) {
            float s = 0.f;
            #pragma unroll
            for (int w = 0; w < 8; ++w) s += warp_sums[tid * 8 + w];
            scores[tid * LCONST + j] = s * scale;
        }
        __syncthreads();
    }

    // Softmax each row of the score tile in-place (stored as exp(score-max)); keep reciprocal sum.
    for (int r = 0; r < BR; ++r) {
        if (qbase + r >= LCONST) {
            if (tid == 0) inv_sum[r] = 0.f;
            __syncthreads();
            continue;
        }
        float m = -FLT_MAX;
        for (int j = tid; j < LCONST; j += NTHREADS) m = fmaxf(m, scores[r * LCONST + j]);
        m = block_reduce_max(m, red);
        float s = 0.f;
        for (int j = tid; j < LCONST; j += NTHREADS) {
            float e = __expf(scores[r * LCONST + j] - m);
            scores[r * LCONST + j] = e;
            s += e;
        }
        s = block_reduce_sum(s, red);
        if (tid == 0) inv_sum[r] = 1.0f / s;
        __syncthreads();
    }

    // Multiply by V.  Each V element loaded once per query block and contributes to all BR rows.
    for (int d = tid; d < DCONST; d += NTHREADS) {
        float o0 = 0.f, o1 = 0.f, o2 = 0.f, o3 = 0.f, o4 = 0.f, o5 = 0.f;
        #pragma unroll 1
        for (int j = 0; j < LCONST; ++j) {
            float vv = Vh[j * DCONST + d];
            o0 = fmaf(scores[0 * LCONST + j], vv, o0);
            o1 = fmaf(scores[1 * LCONST + j], vv, o1);
            o2 = fmaf(scores[2 * LCONST + j], vv, o2);
            o3 = fmaf(scores[3 * LCONST + j], vv, o3);
            o4 = fmaf(scores[4 * LCONST + j], vv, o4);
            o5 = fmaf(scores[5 * LCONST + j], vv, o5);
        }
        if (qbase + 0 < LCONST) Oh[(qbase + 0) * DCONST + d] = o0 * inv_sum[0];
        if (qbase + 1 < LCONST) Oh[(qbase + 1) * DCONST + d] = o1 * inv_sum[1];
        if (qbase + 2 < LCONST) Oh[(qbase + 2) * DCONST + d] = o2 * inv_sum[2];
        if (qbase + 3 < LCONST) Oh[(qbase + 3) * DCONST + d] = o3 * inv_sum[3];
        if (qbase + 4 < LCONST) Oh[(qbase + 4) * DCONST + d] = o4 * inv_sum[4];
        if (qbase + 5 < LCONST) Oh[(qbase + 5) * DCONST + d] = o5 * inv_sum[5];
    }
}

torch::Tensor sdpa_forward(torch::Tensor Q, torch::Tensor K, torch::Tensor V) {
    auto O = torch::empty_like(Q);
    const int B = (int)Q.size(0);
    const int H = (int)Q.size(1);
    const int total_bh = B * H;
    dim3 block(NTHREADS);
    dim3 grid((LCONST + BR - 1) / BR, total_bh);
    sdpa_kernel<<<grid, block>>>(Q.data_ptr<float>(), K.data_ptr<float>(), V.data_ptr<float>(), O.data_ptr<float>(), total_bh);
    return O;
}
"""

_ext = load_inline(
    name="sdpa_d1024_l512_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["sdpa_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
        return _ext.sdpa_forward(Q.contiguous(), K.contiguous(), V.contiguous())