gemini-2.5-pro / cudad85b77

#include "kernel.h"
#include <torch/extension.h>
#include <vector>
#include <stdexcept>
#include <cmath>

#ifdef _OPENMP
#include <omp.h>
#endif

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_DTYPE(x, t) TORCH_CHECK(x.scalar_type() == t, #x " must have dtype " #t)

// C++ implementation of the 'run' function
std::vector<torch::Tensor> run(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    torch::Tensor qo_indptr,
    torch::Tensor kv_indptr,
    float sm_scale) {
    
    // --- Input Validation ---
    CHECK_CUDA(q); CHECK_CUDA(k); CHECK_CUDA(v);
    CHECK_CUDA(qo_indptr); CHECK_CUDA(kv_indptr);

    CHECK_CONTIGUOUS(q); CHECK_CONTIGUOUS(k); CHECK_CONTIGUOUS(v);
    CHECK_CONTIGUOUS(qo_indptr); CHECK_CONTIGUOUS(kv_indptr);

    CHECK_DTYPE(q, torch::kBFloat16);
    CHECK_DTYPE(k, torch::kBFloat16);
    CHECK_DTYPE(v, torch::kBFloat16);
    CHECK_DTYPE(qo_indptr, torch::kInt32);
    CHECK_DTYPE(kv_indptr, torch::kInt32);

    // --- Get Tensor Properties ---
    const int32_t total_q = q.size(0);
    const int32_t num_qo_heads = q.size(1);
    const int32_t head_dim = q.size(2);
    const int32_t len_indptr = qo_indptr.size(0);
    const int32_t batch_size = len_indptr - 1;

    TORCH_CHECK(num_qo_heads == 32, "num_qo_heads must be 32");
    TORCH_CHECK(head_dim == 128, "head_dim must be 128");
    TORCH_CHECK(k.size(1) == 4, "num_kv_heads must be 4");
    TORCH_CHECK(k.size(2) == 128, "head_dim must be 128");
    TORCH_CHECK(v.size(1) == 4, "num_kv_heads must be 4");
    TORCH_CHECK(v.size(2) == 128, "head_dim must be 128");

    // --- Prepare Outputs ---
    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};
    }
    
    // --- Pre-computation on Host: Create q_to_batch_idx map ---
    // This map avoids a search operation inside the kernel for every query token.
    auto q_to_batch_idx = torch::empty({total_q}, torch::kInt32);
    auto qo_indptr_cpu = qo_indptr.to(torch::kCPU);
    auto qo_indptr_acc = qo_indptr_cpu.accessor<int32_t, 1>();
    auto q_to_batch_idx_acc = q_to_batch_idx.accessor<int32_t, 1>();
    
    #pragma omp parallel for
    for (int b = 0; b < batch_size; ++b) {
        int32_t start = qo_indptr_acc[b];
        int32_t end = qo_indptr_acc[b+1];
        for (int32_t i = start; i < end; ++i) {
            q_to_batch_idx_acc[i] = b;
        }
    }
    auto q_to_batch_idx_gpu = q_to_batch_idx.to(q.device());


    // --- Get CUDA Stream ---
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // --- Launch CUDA Kernel ---
    gqa_ragged_prefill_causal_h32_kv4_d128_kernel_launch(
        q.data_ptr(),
        k.data_ptr(),
        v.data_ptr(),
        qo_indptr.data_ptr<int32_t>(),
        kv_indptr.data_ptr<int32_t>(),
        q_to_batch_idx_gpu.data_ptr<int32_t>(),
        sm_scale,
        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,
        "Grouped-Query Attention for Ragged Tensors (Prefill, Causal)",
        py::arg("q"),
        py::arg("k"),
        py::arg("v"),
        py::arg("qo_indptr"),
        py::arg("kv_indptr"),
        py::arg("sm_scale") = 1.0f / std::sqrt(128.0f)
    );
}