gemini-2.5-pro_cuda_425070

#include <torch/extension.h>
#include <vector>
#include <stdexcept>
#include <utility> // For std::pair
#include "kernel.h"

#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

// Helper function to create a mapping from global query index to batch index on the CPU,
// then transfer it to the GPU. This avoids complex lookups inside the kernel.
torch::Tensor create_q_to_batch_idx_map(torch::Tensor qo_indptr, int total_q) {
    int len_indptr = qo_indptr.size(0);
    int batch_size = len_indptr > 0 ? len_indptr - 1 : 0;

    auto qo_indptr_cpu = qo_indptr.to(torch::kCPU);
    auto accessor = qo_indptr_cpu.accessor<int, 1>();

    std::vector<int> q_to_batch_idx_vec(total_q);
    for (int b = 0; b < batch_size; ++b) {
        int q_start = accessor[b];
        int q_end = accessor[b+1];
        for (int i = q_start; i < q_end; ++i) {
            if (i < total_q) {
                q_to_batch_idx_vec[i] = b;
            }
        }
    }

    return torch::tensor(q_to_batch_idx_vec, torch::dtype(torch::kInt32)).to(qo_indptr.device());
}

// C++ entry point wrapped for PyTorch
std::pair<torch::Tensor, torch::Tensor> run(
    torch::Tensor q,
    torch::Tensor k_cache,
    torch::Tensor v_cache,
    torch::Tensor qo_indptr,
    torch::Tensor kv_indptr,
    torch::Tensor kv_indices,
    float sm_scale) {

    // Input validation
    CHECK_INPUT(q);
    CHECK_INPUT(k_cache);
    CHECK_INPUT(v_cache);
    CHECK_INPUT(qo_indptr);
    CHECK_INPUT(kv_indptr);
    CHECK_INPUT(kv_indices);

    TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be BFloat16");
    TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be BFloat16");
    TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be BFloat16");
    TORCH_CHECK(qo_indptr.dtype() == torch::kInt32, "qo_indptr must be Int32");
    TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be Int32");
    TORCH_CHECK(kv_indices.dtype() == torch::kInt32, "kv_indices must be Int32");

    // Shape checks based on specification constants
    const int total_q = q.size(0);
    const int num_qo_heads = q.size(1);
    const int head_dim = q.size(2);
    
    TORCH_CHECK(num_qo_heads == 32, "num_qo_heads must be 32");
    TORCH_CHECK(k_cache.size(2) == 8, "num_kv_heads must be 8");
    TORCH_CHECK(head_dim == 128, "head_dim must be 128");
    TORCH_CHECK(k_cache.size(1) == 1, "page_size must be 1");

    // Allocate output tensors
    auto output = torch::empty_like(q);
    auto lse = torch::empty({total_q, num_qo_heads}, q.options().dtype(torch::kFloat32));

    if (total_q == 0) {
        return {output, lse};
    }

    // Precompute mapping from query index to batch index for efficient kernel lookups
    auto q_to_batch_idx = create_q_to_batch_idx_map(qo_indptr, total_q);

    // Get CUDA stream from PyTorch to ensure proper command ordering
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Call the CUDA kernel launcher
    gqa_paged_prefill_causal_h32_kv8_d128_ps1_launch(
        reinterpret_cast<const nv_bfloat16*>(q.data_ptr()),
        reinterpret_cast<const nv_bfloat16*>(k_cache.data_ptr()),
        reinterpret_cast<const nv_bfloat16*>(v_cache.data_ptr()),
        qo_indptr.data_ptr<int>(),
        kv_indptr.data_ptr<int>(),
        kv_indices.data_ptr<int>(),
        q_to_batch_idx.data_ptr<int>(),
        sm_scale,
        reinterpret_cast<nv_bfloat16*>(output.data_ptr()),
        lse.data_ptr<float>(),
        total_q,
        stream
    );

    return {output, lse};
}

// Pybind11 module definition
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "GQA Paged Prefill (Causal, H=32, KV_H=8, D=128, PS=1) kernel",
          py::arg("q"),
          py::arg("k_cache"),
          py::arg("v_cache"),
          py::arg("qo_indptr"),
          py::arg("kv_indptr"),
          py::arg("kv_indices"),
          py::arg("sm_scale"));
}