gpt-o3 / cuda025ec1

#include "kernel.h"

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

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

torch::Tensor gqa_run(torch::Tensor q,
                      torch::Tensor k,
                      torch::Tensor v,
                      torch::Tensor qo_indptr,
                      torch::Tensor kv_indptr,
                      float sm_scale,
                      torch::Tensor output,
                      torch::Tensor lse) {
    CHECK_INPUT(q);
    CHECK_INPUT(k);
    CHECK_INPUT(v);
    CHECK_INPUT(qo_indptr);
    CHECK_INPUT(kv_indptr);
    CHECK_INPUT(output);
    CHECK_INPUT(lse);

    TORCH_CHECK(q.scalar_type()  == at::kBFloat16, "q must be bfloat16");
    TORCH_CHECK(k.scalar_type()  == at::kBFloat16, "k must be bfloat16");
    TORCH_CHECK(v.scalar_type()  == at::kBFloat16, "v must be bfloat16");
    TORCH_CHECK(output.scalar_type() == at::kBFloat16, "output must be bfloat16");
    TORCH_CHECK(lse.scalar_type()    == at::kFloat,    "lse must be float32");

    int total_q  = q.size(0);
    int total_kv = k.size(0);
    int len_indptr = qo_indptr.numel();

    // pointer casts
    const __nv_bfloat16 *q_ptr  = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<at::BFloat16>());
    const __nv_bfloat16 *k_ptr  = reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<at::BFloat16>());
    const __nv_bfloat16 *v_ptr  = reinterpret_cast<const __nv_bfloat16*>(v.data_ptr<at::BFloat16>());
    const int32_t *qo_ptr       = qo_indptr.data_ptr<int32_t>();
    const int32_t *kv_ptr       = kv_indptr.data_ptr<int32_t>();
    __nv_bfloat16 *out_ptr      = reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>());
    float *lse_ptr              = lse.data_ptr<float>();

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    gqa_ragged_prefill_causal_h32_kv8_d128_launcher(
        q_ptr, k_ptr, v_ptr,
        qo_ptr, kv_ptr,
        sm_scale,
        len_indptr,
        total_q,
        total_kv,
        out_ptr,
        lse_ptr,
        stream);

    TORCH_CHECK(cudaGetLastError() == cudaSuccess,
                "Kernel launch failed");

    return output;   // we also changed lse in-place
}

// ------------------------------------------------------------------
// Python binding
// ------------------------------------------------------------------
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run",
          [](torch::Tensor q,
             torch::Tensor k,
             torch::Tensor v,
             torch::Tensor qo_indptr,
             torch::Tensor kv_indptr,
             float sm_scale) {
                auto options_bf16 = torch::dtype(torch::kBFloat16).device(q.device());
                auto options_f32  = torch::dtype(torch::kFloat32).device(q.device());

                int64_t total_q = q.size(0);

                auto output = torch::empty({total_q, 32, 128}, options_bf16);
                auto lse    = torch::empty({total_q, 32},       options_f32);

                gqa_run(q, k, v, qo_indptr, kv_indptr, sm_scale, output, lse);
                return std::make_tuple(output, lse);
          },
          py::arg("q"),
          py::arg("k"),
          py::arg("v"),
          py::arg("qo_indptr"),
          py::arg("kv_indptr"),
          py::arg("sm_scale") = 1.0f / 11.313708498984761f); // 1/sqrt(128)
}