gpt-o3 / cuda9abd34

#include "kernel.h"

#include <torch/extension.h>
#include <vector>
#include <cmath>
#include <tuple>

/* -------------------------------------------------------------------------- */
/*  Convenience: empty BF16 tensor on a device                                */
/* -------------------------------------------------------------------------- */
static inline torch::Tensor bf16_empty(const std::vector<int64_t>& sizes,
                                       const torch::Device& dev)
{
    return torch::empty(
        sizes,
        torch::TensorOptions()
              .dtype(torch::kBFloat16)
              .device(dev));
}

/* -------------------------------------------------------------------------- */
/*  PUBLIC ENTRY POINT (exposed to Python)                                    */
/* -------------------------------------------------------------------------- */
std::tuple<torch::Tensor, torch::Tensor>
run(torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    torch::Tensor qo_indptr,
    torch::Tensor kv_indptr,
    double        sm_scale_d = 1.0 / std::sqrt(128.0))
{
    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(),
                "q, k, v must be CUDA tensors");
    TORCH_CHECK(q.scalar_type() == torch::kBFloat16 &&
                k.scalar_type() == torch::kBFloat16 &&
                v.scalar_type() == torch::kBFloat16,
                "q, k, v must be bfloat16");

    TORCH_CHECK(qo_indptr.scalar_type() == torch::kInt32 &&
                kv_indptr.scalar_type() == torch::kInt32,
                "indptr tensors must be int32");

    /* fixed shapes */
    TORCH_CHECK(q.size(1) == NUM_QO_HEADS && q.size(2) == HEAD_DIM,
                "q wrong second/third dimension");
    TORCH_CHECK(k.size(1) == NUM_KV_HEADS && k.size(2) == HEAD_DIM,
                "k wrong second/third dimension");

    const int64_t total_q    = q.size(0);
    const int64_t total_kv   = k.size(0);
    const int64_t len_indptr = qo_indptr.size(0);

    TORCH_CHECK(qo_indptr[len_indptr - 1].item<int32_t>() == total_q,
                "total_q inconsistent with qo_indptr");
    TORCH_CHECK(kv_indptr[len_indptr - 1].item<int32_t>() == total_kv,
                "total_kv inconsistent with kv_indptr");

    /* allocate outputs */
    const auto device = q.device();
    torch::Tensor output = bf16_empty({total_q, NUM_QO_HEADS, HEAD_DIM}, device);
    torch::Tensor lse    = torch::empty({total_q, NUM_QO_HEADS},
                                        torch::TensorOptions()
                                              .dtype(torch::kFloat32)
                                              .device(device));

    /* launch */
    gqa_ragged_prefill_causal_h32_kv4_d128_launcher(
        q, k, v,
        qo_indptr, kv_indptr,
        static_cast<float>(sm_scale_d),
        output, lse);

    return {output, lse};
}

/* -------------------------------------------------------------------------- */
/*  PYBIND11 MODULE                                                           */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.doc() =
        "Optimised ragged causal prefill attention "
        "(32 QO heads / 4 KV heads / head_dim 128)";
    m.def("run", &run,
          pybind11::arg("q"),
          pybind11::arg("k"),
          pybind11::arg("v"),
          pybind11::arg("qo_indptr"),
          pybind11::arg("kv_indptr"),
          pybind11::arg("sm_scale") = 1.0 / std::sqrt(128.0),
          "Execute the kernel and return (output, lse)");
}