gpt-o3 / cudaec73ea

#include "kernel.h"

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

/* shorthand helpers ------------------------------------------------------- */
#define CHECK_CUDA(x)     TORCH_CHECK(x.is_cuda(),       #x " must be CUDA")
#define CHECK_CONTIG(x)   TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_DTYPE(x, t) TORCH_CHECK(x.scalar_type() == t, #x " has wrong dtype")

/* ------------------------------------------------------------------------ */
/*  Python-visible ‘run’ function                                           */
/* ------------------------------------------------------------------------ */
std::tuple<torch::Tensor, torch::Tensor>
run(torch::Tensor q,
    torch::Tensor k_cache,
    torch::Tensor v_cache,
    torch::Tensor kv_indptr,
    torch::Tensor kv_indices,
    float sm_scale)
{
    /* --------- basic checks -------------------------------------- */
    CHECK_CUDA(q);           CHECK_CUDA(k_cache);     CHECK_CUDA(v_cache);
    CHECK_CUDA(kv_indptr);   CHECK_CUDA(kv_indices);

    CHECK_CONTIG(q);         CHECK_CONTIG(k_cache);   CHECK_CONTIG(v_cache);
    CHECK_CONTIG(kv_indptr); CHECK_CONTIG(kv_indices);

    CHECK_DTYPE(q,         at::kBFloat16);
    CHECK_DTYPE(k_cache,   at::kBFloat16);
    CHECK_DTYPE(v_cache,   at::kBFloat16);
    CHECK_DTYPE(kv_indptr, at::kInt);
    CHECK_DTYPE(kv_indices,at::kInt);

    /* geometry */
    const int64_t B      = q.size(0);
    const int64_t n_idx  = kv_indices.size(0);
    const int64_t n_page = k_cache.size(0);      /* not used by kernel */

    /* outputs */
    auto output = torch::empty({B, NUM_QO_HEADS, HEAD_DIM}, q.options());
    auto lse    = torch::empty({B, NUM_QO_HEADS},
                               q.options().dtype(torch::kFloat32));

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

    /* launch */
    gqa_paged_decode_h32_kv8_d128_ps1_launcher(
        reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(k_cache.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(v_cache.data_ptr<at::BFloat16>()),
        kv_indptr.data_ptr<int32_t>(),
        kv_indices.data_ptr<int32_t>(),
        sm_scale,
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()),
        lse.data_ptr<float>(),
        static_cast<int>(B),
        static_cast<int>(n_idx),
        static_cast<int>(n_page),
        stream);

    return {output, lse};
}

/* ------------------------------------------------------------------------ */
/*  PyBind11 module                                                         */
/* ------------------------------------------------------------------------ */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run,
          "Optimised GQA paged-decode kernel (H32 / K8 / D128 / PS1)");
}