DeepSeek NSA GPT-6 Astra Pro

"""CUDA NSA with FP32 block summaries and fused top-eight sparse attention."""

import os
import shutil
import sys

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


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

_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <math_constants.h>

template<int DIM>
__global__ void key_means(const __nv_bfloat16* __restrict__ keys,
                          float* __restrict__ prefix,
                          float* __restrict__ means, int length, int blocks) {
    constexpr int CHUNKS = 256 / DIM;
    constexpr int ITEMS = 64 / CHUNKS;
    __shared__ float ends[256];
    int dimension = threadIdx.x % DIM;
    int chunk = threadIdx.x / DIM;
    int start = blockIdx.x * 64 + chunk * ITEMS;
    int head = blockIdx.y;
    float values[ITEMS];
    float running = 0.0f;
    #pragma unroll
    for (int item = 0; item < ITEMS; ++item) {
        float value = start + item < length
            ? __bfloat162float(keys[(size_t(head) * length + start + item) * DIM + dimension]) : 0.0f;
        running += value;
        values[item] = running;
    }
    ends[threadIdx.x] = running;
    __syncthreads();
    float offset = 0.0f;
    #pragma unroll
    for (int previous = 0; previous < CHUNKS; ++previous) {
        if (previous < chunk) offset += ends[previous * DIM + dimension];
    }
    #pragma unroll
    for (int item = 0; item < ITEMS; ++item) {
        if (start + item < length) {
            prefix[(size_t(head) * length + start + item) * DIM + dimension]
                = (values[item] + offset) / float(chunk * ITEMS + item + 1);
        }
    }
    if (chunk == CHUNKS - 1) {
        means[(size_t(head) * blocks + blockIdx.x) * DIM + dimension] = (running + offset) * (1.0f / 64.0f);
    }
}

__device__ __forceinline__ void mma_bf16(float* accum, const unsigned* left, const unsigned* right) {
    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};"
        : "+f"(accum[0]), "+f"(accum[1]), "+f"(accum[2]), "+f"(accum[3])
        : "r"(left[0]), "r"(left[1]), "r"(left[2]), "r"(left[3]), "r"(right[0]), "r"(right[1]));
}

__device__ __forceinline__ unsigned round_tf32(float value) {
    unsigned result;
    asm volatile("cvt.rn.tf32.f32 %0, %1;" : "=r"(result) : "f"(value));
    return result;
}

__device__ __forceinline__ void mma_tf32(float* accum, const unsigned* left, const unsigned* right) {
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
        : "+f"(accum[0]), "+f"(accum[1]), "+f"(accum[2]), "+f"(accum[3])
        : "r"(left[0]), "r"(left[1]), "r"(left[2]), "r"(left[3]), "r"(right[0]), "r"(right[1]));
}

template<int DIM, int CAPACITY>
__device__ __forceinline__ void select_masks(
    const __nv_bfloat16* __restrict__ queries,
    const float* __restrict__ prefix,
    const float* __restrict__ means,
    unsigned* top_masks, unsigned* bottom_masks, int length, int blocks, int query_start) {
    int lane = threadIdx.x & 31;
    int row = lane / 4;
    int column = lane % 4;
    int query_top = query_start + row;
    int query_bottom = query_top + 8;
    if (query_start >= length) return;
    int current = query_start / 64;
    size_t head_base = size_t(blockIdx.y) * length * DIM;
    float scores[CAPACITY / 8][4] = {};
    #pragma unroll
    for (int chunk = 0; chunk < DIM / 8; ++chunk) {
        unsigned query_fragment[4];
        int dimension = chunk * 8 + column;
        query_fragment[0] = __float_as_uint(query_top < length ? __bfloat162float(queries[head_base + query_top * DIM + dimension]) : 0.0f);
        query_fragment[1] = __float_as_uint(query_bottom < length ? __bfloat162float(queries[head_base + query_bottom * DIM + dimension]) : 0.0f);
        query_fragment[2] = __float_as_uint(query_top < length ? __bfloat162float(queries[head_base + query_top * DIM + dimension + 4]) : 0.0f);
        query_fragment[3] = __float_as_uint(query_bottom < length ? __bfloat162float(queries[head_base + query_bottom * DIM + dimension + 4]) : 0.0f);
        #pragma unroll
        for (int fragment = 0; fragment < CAPACITY / 8; ++fragment) {
            if (fragment * 8 <= current) {
                int block = fragment * 8 + row;
                unsigned mean_high[2] = {};
                unsigned mean_low[2] = {};
                if (block < blocks) {
                    #pragma unroll
                    for (int element = 0; element < 2; ++element) {
                        float value = means[(size_t(blockIdx.y) * blocks + block) * DIM + dimension + element * 4];
                        mean_high[element] = round_tf32(value);
                        mean_low[element] = round_tf32(value - __uint_as_float(mean_high[element]));
                    }
                }
                mma_tf32(scores[fragment], query_fragment, mean_high);
                mma_tf32(scores[fragment], query_fragment, mean_low);
            }
        }
    }
    float partial_top = 0.0f;
    float partial_bottom = 0.0f;
    #pragma unroll
    for (int dimension = column; dimension < DIM; dimension += 4) {
        if (query_top < length)
            partial_top = fmaf(__bfloat162float(queries[head_base + query_top * DIM + dimension]),
                prefix[head_base + query_top * DIM + dimension], partial_top);
        if (query_bottom < length)
            partial_bottom = fmaf(__bfloat162float(queries[head_base + query_bottom * DIM + dimension]),
                prefix[head_base + query_bottom * DIM + dimension], partial_bottom);
    }
    #pragma unroll
    for (int offset = 1; offset < 4; offset *= 2) {
        partial_top += __shfl_xor_sync(0xffffffff, partial_top, offset, 4);
        partial_bottom += __shfl_xor_sync(0xffffffff, partial_bottom, offset, 4);
    }
    #pragma unroll
    for (int fragment = 0; fragment < CAPACITY / 8; ++fragment) {
        #pragma unroll
        for (int element = 0; element < 2; ++element) {
            int block = fragment * 8 + column * 2 + element;
            if (block == current) {
                scores[fragment][element] = partial_top;
                scores[fragment][element + 2] = partial_bottom;
            }
            if (block > current) scores[fragment][element] = scores[fragment][element + 2] = -CUDART_INF_F;
        }
    }
    unsigned selected_top[4] = {};
    unsigned selected_bottom[4] = {};
    #pragma unroll 1
    for (int rank = 0; rank < 8; ++rank) {
        float best_top = -CUDART_INF_F;
        float best_bottom = -CUDART_INF_F;
        int best_top_index = -1;
        int best_bottom_index = -1;
        #pragma unroll
        for (int fragment = 0; fragment < CAPACITY / 8; ++fragment) {
            #pragma unroll
            for (int element = 0; element < 2; ++element) {
                int index = fragment * 8 + column * 2 + element;
                float top_score = scores[fragment][element];
                float bottom_score = scores[fragment][element + 2];
                if (top_score > best_top || (top_score == best_top && index > best_top_index)) {
                    best_top = top_score;
                    best_top_index = index;
                }
                if (bottom_score > best_bottom || (bottom_score == best_bottom && index > best_bottom_index)) {
                    best_bottom = bottom_score;
                    best_bottom_index = index;
                }
            }
        }
        #pragma unroll
        for (int offset = 1; offset < 4; offset *= 2) {
            float other_top = __shfl_xor_sync(0xffffffff, best_top, offset, 4);
            float other_bottom = __shfl_xor_sync(0xffffffff, best_bottom, offset, 4);
            int other_top_index = __shfl_xor_sync(0xffffffff, best_top_index, offset, 4);
            int other_bottom_index = __shfl_xor_sync(0xffffffff, best_bottom_index, offset, 4);
            if (other_top > best_top || (other_top == best_top && other_top_index > best_top_index)) {
                best_top = other_top;
                best_top_index = other_top_index;
            }
            if (other_bottom > best_bottom || (other_bottom == best_bottom && other_bottom_index > best_bottom_index)) {
                best_bottom = other_bottom;
                best_bottom_index = other_bottom_index;
            }
        }
        #pragma unroll
        for (int fragment = 0; fragment < CAPACITY / 8; ++fragment) {
            #pragma unroll
            for (int element = 0; element < 2; ++element) {
                int index = fragment * 8 + column * 2 + element;
                if (index == best_top_index) scores[fragment][element] = -CUDART_INF_F;
                if (index == best_bottom_index) scores[fragment][element + 2] = -CUDART_INF_F;
            }
        }
        #pragma unroll
        for (int part = 0; part < 4; ++part) {
            if (best_top_index / 32 == part) selected_top[part] |= 1u << (best_top_index % 32);
            if (best_bottom_index / 32 == part) selected_bottom[part] |= 1u << (best_bottom_index % 32);
        }
    }
    #pragma unroll
    for (int part = 0; part < 4; ++part) {
        top_masks[part] = selected_top[part];
        bottom_masks[part] = selected_bottom[part];
    }
}

__device__ __forceinline__ void load_left(unsigned* fragment, const __nv_bfloat16* address) {
    unsigned shared_address = __cvta_generic_to_shared(address);
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
        : "=r"(fragment[0]), "=r"(fragment[1]), "=r"(fragment[2]), "=r"(fragment[3]) : "r"(shared_address));
}

__device__ __forceinline__ void load_right(unsigned* fragment, const __nv_bfloat16* address) {
    unsigned shared_address = __cvta_generic_to_shared(address);
    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
        : "=r"(fragment[0]), "=r"(fragment[1]) : "r"(shared_address));
}

__device__ __forceinline__ void load_key(unsigned* fragment, const __nv_bfloat16* address) {
    unsigned shared_address = __cvta_generic_to_shared(address);
    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
        : "=r"(fragment[0]), "=r"(fragment[1]) : "r"(shared_address));
}

__device__ __forceinline__ void copy_async(__nv_bfloat16* destination, const __nv_bfloat16* source, int bytes) {
    unsigned shared_address = __cvta_generic_to_shared(destination);
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" [REDACTED: IP] "r"(shared_address), "l"(source), "r"(bytes));
}

template<int DIM>
__global__ __launch_bounds__(128) void tensor_attention(
    const __nv_bfloat16* __restrict__ queries,
    const __nv_bfloat16* __restrict__ keys,
    const __nv_bfloat16* __restrict__ values,
    const float* __restrict__ prefix,
    const float* __restrict__ means,
    __nv_bfloat16* __restrict__ output, int length) {
    __shared__ __align__(16) __nv_bfloat16 key_tile[64][DIM + 8];
    __shared__ __align__(16) __nv_bfloat16 value_tile[64][DIM + 8];
    __shared__ __align__(16) __nv_bfloat16 probability_tile[64][72];
    __shared__ unsigned union_words[4][4];
    __shared__ int gathered_blocks[128];
    constexpr int OUTPUT_FRAGMENTS = DIM / 8;
    constexpr float SCALE = (DIM == 64 ? 0.125f : 0.08838834764831844f) * 1.4426950408889634f;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int row = lane / 4;
    int column = (lane % 4) * 2;
    int query_start = blockIdx.x * 64 + warp * 16;
    int query_top = query_start + row;
    int query_bottom = query_top + 8;
    int current = query_start / 64;
    size_t head_base = size_t(blockIdx.y) * length * DIM;
    unsigned top_masks[4] = {};
    unsigned bottom_masks[4] = {};
    if (current >= 8) {
        int blocks = (length + 63) / 64;
        if (current < 16)
            select_masks<DIM, 16>(queries, prefix, means, top_masks, bottom_masks, length, blocks, query_start);
        else if (current < 32)
            select_masks<DIM, 32>(queries, prefix, means, top_masks, bottom_masks, length, blocks, query_start);
        else if (current < 64)
            select_masks<DIM, 64>(queries, prefix, means, top_masks, bottom_masks, length, blocks, query_start);
        else
            select_masks<DIM, 128>(queries, prefix, means, top_masks, bottom_masks, length, blocks, query_start);
    } else {
        top_masks[0] = bottom_masks[0] = (1u << (current + 1)) - 1;
    }
    unsigned active[4];
    #pragma unroll
    for (int part = 0; part < 4; ++part) {
        unsigned mask = top_masks[part] | bottom_masks[part];
        #pragma unroll
        for (int offset = 4; offset < 32; offset *= 2)
            mask |= __shfl_xor_sync(0xffffffff, mask, offset);
        if (current / 32 == part) mask |= 1u << (current % 32);
        if (current > 0 && (current - 1) / 32 == part) mask |= 1u << ((current - 1) % 32);
        if (lane == 0) union_words[warp][part] = mask;
    }
    __syncthreads();
    int block_count = 0;
    #pragma unroll
    for (int part = 0; part < 4; ++part) {
        unsigned mask = 0;
        #pragma unroll
        for (int query_tile = 0; query_tile < 4; ++query_tile)
            mask |= union_words[query_tile][part];
        active[part] = mask;
        block_count += __popc(mask);
    }
    int candidate_block = threadIdx.x;
    int candidate_part = candidate_block / 32;
    unsigned candidate_bit = 1u << (candidate_block % 32);
    if (active[candidate_part] & candidate_bit) {
        int rank = __popc(active[candidate_part] & (candidate_bit - 1));
        #pragma unroll
        for (int earlier = 0; earlier < 4; ++earlier)
            if (earlier < candidate_part) rank += __popc(active[earlier]);
        gathered_blocks[rank] = candidate_block;
    }
    __syncthreads();
    for (int vector = threadIdx.x; vector < 64 * DIM / 8; vector += 128) {
        int [REDACTED credential assignment] / (DIM / 8);
        int dimension = (vector % (DIM / 8)) * 8;
        int key_start = gathered_blocks[0] * 64;
        copy_async(&key_tile[token][dimension], keys + head_base + (key_start + token) * DIM + dimension,
            key_start + token < length ? 16 : 0);
    }
    asm volatile("cp.async.commit_group;");
    unsigned query_fragments[DIM / 16][4];
    #pragma unroll
    for (int chunk = 0; chunk < DIM / 16; ++chunk) {
        int dimension = chunk * 16 + column;
        query_fragments[chunk][0] = query_top < length ? *reinterpret_cast<const unsigned*>(queries + head_base + query_top * DIM + dimension) : 0;
        query_fragments[chunk][1] = query_bottom < length ? *reinterpret_cast<const unsigned*>(queries + head_base + query_bottom * DIM + dimension) : 0;
        query_fragments[chunk][2] = query_top < length ? *reinterpret_cast<const unsigned*>(queries + head_base + query_top * DIM + dimension + 8) : 0;
        query_fragments[chunk][3] = query_bottom < length ? *reinterpret_cast<const unsigned*>(queries + head_base + query_bottom * DIM + dimension + 8) : 0;
    }
    float numerator[OUTPUT_FRAGMENTS][4] = {};
    float maximum_top = -1.0e20f;
    float maximum_bottom = -1.0e20f;
    float denominator_top = 0.0f;
    float denominator_bottom = 0.0f;
    for (int block_rank = 0; block_rank < block_count; ++block_rank) {
        int block = gathered_blocks[block_rank];
        int part = block / 32;
        unsigned bit = 1u << (block % 32);
        int key_start = block * 64;
        bool selected_top = (top_masks[part] & bit) != 0;
        bool selected_bottom = (bottom_masks[part] & bit) != 0;
        bool compute_tile = (union_words[warp][part] & bit) != 0;
        asm volatile("cp.async.wait_group 0;");
        __syncthreads();
        for (int vector = threadIdx.x; vector < 64 * DIM / 8; vector += 128) {
            int [REDACTED credential assignment] / (DIM / 8);
            int dimension = (vector % (DIM / 8)) * 8;
            copy_async(&value_tile[token][dimension], values + head_base + (key_start + token) * DIM + dimension,
                key_start + token < length ? 16 : 0);
        }
        asm volatile("cp.async.commit_group;");
        float scores[8][4] = {};
        float local_maximum_top = -CUDART_INF_F;
        float local_maximum_bottom = -CUDART_INF_F;
        if (compute_tile) {
            #pragma unroll
            for (int chunk = 0; chunk < DIM / 16; ++chunk) {
                #pragma unroll
                for (int key_part = 0; key_part < 8; ++key_part) {
                    int [REDACTED credential assignment] * 8 + lane % 8;
                    unsigned key_fragment[2];
                    load_key(key_fragment, &key_tile[token][chunk * 16 + ((lane / 8) % 2) * 8]);
                    mma_bf16(scores[key_part], query_fragments[chunk], key_fragment);
                }
            }
            #pragma unroll
            for (int key_part = 0; key_part < 8; ++key_part) {
                #pragma unroll
                for (int element = 0; element < 2; ++element) {
                    int [REDACTED credential assignment] + key_part * 8 + column + element;
                    bool valid_top = token <= query_top && query_top < length && (selected_top || token >= query_top - 63);
                    bool valid_bottom = token <= query_bottom && query_bottom < length && (selected_bottom || token >= query_bottom - 63);
                    scores[key_part][element] = valid_top ? scores[key_part][element] * SCALE : -CUDART_INF_F;
                    scores[key_part][element + 2] = valid_bottom ? scores[key_part][element + 2] * SCALE : -CUDART_INF_F;
                    local_maximum_top = fmaxf(local_maximum_top, scores[key_part][element]);
                    local_maximum_bottom = fmaxf(local_maximum_bottom, scores[key_part][element + 2]);
                }
            }
            #pragma unroll
            for (int offset = 1; offset < 4; offset *= 2) {
                local_maximum_top = fmaxf(local_maximum_top, __shfl_xor_sync(0xffffffff, local_maximum_top, offset, 4));
                local_maximum_bottom = fmaxf(local_maximum_bottom, __shfl_xor_sync(0xffffffff, local_maximum_bottom, offset, 4));
            }
        }
        __syncthreads();
        bool has_next = block_rank + 1 < block_count;
        if (has_next) {
            int next_key_start = gathered_blocks[block_rank + 1] * 64;
            for (int vector = threadIdx.x; vector < 64 * DIM / 8; vector += 128) {
                int [REDACTED credential assignment] / (DIM / 8);
                int dimension = (vector % (DIM / 8)) * 8;
                copy_async(&key_tile[token][dimension], keys + head_base + (next_key_start + token) * DIM + dimension,
                    next_key_start + token < length ? 16 : 0);
            }
            asm volatile("cp.async.commit_group;");
        }
        if (compute_tile) {
            float next_maximum_top = fmaxf(maximum_top, local_maximum_top);
            float next_maximum_bottom = fmaxf(maximum_bottom, local_maximum_bottom);
            float rescale_top = exp2f(maximum_top - next_maximum_top);
            float rescale_bottom = exp2f(maximum_bottom - next_maximum_bottom);
            float sum_top = 0.0f;
            float sum_bottom = 0.0f;
            #pragma unroll
            for (int key_part = 0; key_part < 8; ++key_part) {
                float probability_top_first = exp2f(scores[key_part][0] - next_maximum_top);
                float probability_top_second = exp2f(scores[key_part][1] - next_maximum_top);
                float probability_bottom_first = exp2f(scores[key_part][2] - next_maximum_bottom);
                float probability_bottom_second = exp2f(scores[key_part][3] - next_maximum_bottom);
                sum_top += probability_top_first;
                sum_top += probability_top_second;
                sum_bottom += probability_bottom_first;
                sum_bottom += probability_bottom_second;
                int [REDACTED credential assignment] * 8 + column;
                *reinterpret_cast<__nv_bfloat162*>(&probability_tile[warp * 16 + row][token])
                    = __floats2bfloat162_rn(probability_top_first, probability_top_second);
                *reinterpret_cast<__nv_bfloat162*>(&probability_tile[warp * 16 + row + 8][token])
                    = __floats2bfloat162_rn(probability_bottom_first, probability_bottom_second);
            }
            #pragma unroll
            for (int offset = 1; offset < 4; offset *= 2) {
                sum_top += __shfl_xor_sync(0xffffffff, sum_top, offset, 4);
                sum_bottom += __shfl_xor_sync(0xffffffff, sum_bottom, offset, 4);
            }
            #pragma unroll
            for (int fragment = 0; fragment < OUTPUT_FRAGMENTS; ++fragment) {
                numerator[fragment][0] *= rescale_top;
                numerator[fragment][1] *= rescale_top;
                numerator[fragment][2] *= rescale_bottom;
                numerator[fragment][3] *= rescale_bottom;
            }
            denominator_top = denominator_top * rescale_top + sum_top;
            denominator_bottom = denominator_bottom * rescale_bottom + sum_bottom;
            maximum_top = next_maximum_top;
            maximum_bottom = next_maximum_bottom;
        }
        if (has_next) asm volatile("cp.async.wait_group 1;");
        else asm volatile("cp.async.wait_group 0;");
        __syncthreads();
        if (compute_tile) {
            #pragma unroll
            for (int chunk = 0; chunk < 4; ++chunk) {
                unsigned probability_fragment[4];
                load_left(probability_fragment, &probability_tile[warp * 16 + lane % 16][chunk * 16 + (lane / 16) * 8]);
                #pragma unroll
                for (int fragment = 0; fragment < OUTPUT_FRAGMENTS; ++fragment) {
                    unsigned value_fragment[2];
                    load_right(value_fragment, &value_tile[chunk * 16 + lane % 16][fragment * 8]);
                    mma_bf16(numerator[fragment], probability_fragment, value_fragment);
                }
            }
        }
    }
    #pragma unroll
    for (int fragment = 0; fragment < OUTPUT_FRAGMENTS; ++fragment) {
        int dimension = fragment * 8 + column;
        if (query_top < length) {
            auto result = __floats2bfloat162_rn(numerator[fragment][0] / denominator_top, numerator[fragment][1] / denominator_top);
            *reinterpret_cast<__nv_bfloat162*>(output + head_base + query_top * DIM + dimension) = result;
        }
        if (query_bottom < length) {
            auto result = __floats2bfloat162_rn(numerator[fragment][2] / denominator_bottom, numerator[fragment][3] / denominator_bottom);
            *reinterpret_cast<__nv_bfloat162*>(output + head_base + query_bottom * DIM + dimension) = result;
        }
    }
}

torch::Tensor nsa_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "CUDA tensors required");
    TORCH_CHECK(q.device() == k.device() && q.device() == v.device(), "Tensors must share a CUDA device");
    TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && v.scalar_type() == at::kBFloat16, "BF16 tensors required");
    TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), "Contiguous tensors required");
    TORCH_CHECK(q.dim() == 4 && q.sizes() == k.sizes() && q.sizes() == v.sizes(), "Expected matching (B,H,S,D) tensors");
    c10::cuda::CUDAGuard guard(q.device());
    int length = q.size(2);
    int dimension = q.size(3);
    int heads = q.size(0) * q.size(1);
    int blocks = (length + 63) / 64;
    TORCH_CHECK((dimension == 64 || dimension == 128) && blocks <= 128, "Unsupported shape");
    auto output = torch::empty_like(q);
    if (q.numel() == 0) return output;
    torch::Tensor prefix, means;
    float* prefix_data = nullptr;
    float* means_data = nullptr;
    if (blocks > 8) {
        prefix = torch::empty(q.sizes(), q.options().dtype(at::kFloat));
        means = torch::empty({heads, blocks, dimension}, q.options().dtype(at::kFloat));
        prefix_data = prefix.data_ptr<float>();
        means_data = means.data_ptr<float>();
    }
    auto stream = at::cuda::getCurrentCUDAStream();
    const auto* queries = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr());
    const auto* keys = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr());
    const auto* values = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr());
    auto* result = reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
    if (dimension == 64) {
        if (blocks > 8) key_means<64><<<dim3(blocks, heads), 256, 0, stream>>>(keys, prefix_data, means_data, length, blocks);
        tensor_attention<64><<<dim3(blocks, heads), 128, 0, stream>>>(queries, keys, values, prefix_data, means_data, result, length);
    } else {
        if (blocks > 8) key_means<128><<<dim3(blocks, heads), 256, 0, stream>>>(keys, prefix_data, means_data, length, blocks);
        tensor_attention<128><<<dim3(blocks, heads), 128, 0, stream>>>(queries, keys, values, prefix_data, means_data, result, length);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}

"""


if shutil.which("ninja") is None:
    os.environ["PATH"] += os.pathsep + os.path.dirname(sys.executable)


_extension = load_inline(
    name="nsa_sparse_sm120",
    cpp_sources=_CPP,
    cuda_sources=_CUDA,
    functions=["nsa_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo", "-gencode=arch=compute_120,code=sm_120"],
    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 _extension.nsa_cuda(q, k, v)