gpt-o3_cuda_037673

#include "kernel.h"

#include <ATen/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <cmath>
#include <limits>

/* ========================================================================= */
/*  Internal helper – wraps CUDA launcher                                    */
/* ========================================================================= */
static void run_kernel(
        const torch::Tensor& q,
        const torch::Tensor& k_cache,
        const torch::Tensor& v_cache,
        const torch::Tensor& qo_indptr,
        const torch::Tensor& kv_indptr,
        const torch::Tensor& kv_indices,
        double               sm_scale,
        torch::Tensor&       output,
        torch::Tensor&       lse)
{
    /* Basic sanity checks -------------------------------------------------- */
    TORCH_CHECK(q.is_cuda()          && q.scalar_type() == torch::kBFloat16, "q must be bf16 CUDA");
    TORCH_CHECK(k_cache.is_cuda()    && k_cache.scalar_type() == torch::kBFloat16, "k_cache must be bf16 CUDA");
    TORCH_CHECK(v_cache.is_cuda()    && v_cache.scalar_type() == torch::kBFloat16, "v_cache must be bf16 CUDA");
    TORCH_CHECK(qo_indptr.is_cuda()  && qo_indptr.scalar_type() == torch::kInt,     "qo_indptr must be int32 CUDA");
    TORCH_CHECK(kv_indptr.is_cuda()  && kv_indptr.scalar_type() == torch::kInt,     "kv_indptr must be int32 CUDA");
    TORCH_CHECK(kv_indices.is_cuda() && kv_indices.scalar_type() == torch::kInt,    "kv_indices must be int32 CUDA");

    /* Ensure we execute on the right device / stream ---------------------- */
    at::cuda::CUDAGuard device_guard(q.device());
    at::cuda::CUDAStream cuda_stream = at::cuda::getCurrentCUDAStream();

    launch_gqa_paged_prefill_causal_h32_kv4_d128_ps1(
        q, k_cache, v_cache,
        qo_indptr, kv_indptr, kv_indices,
        static_cast<float>(sm_scale),
        output, lse,
        cuda_stream.stream());
}

/* ========================================================================= */
/*  Python binding                                                           */
/* ========================================================================= */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run",
          [](torch::Tensor  q,
             torch::Tensor  k_cache,
             torch::Tensor  v_cache,
             torch::Tensor  qo_indptr,
             torch::Tensor  kv_indptr,
             torch::Tensor  kv_indices,
             double         sm_scale = 1.0 / std::sqrt(static_cast<double>(HEAD_DIM)))
          {
              const int64_t total_q = q.size(0);

              /* Allocate outputs ----------------------------------------- */
              auto opts_out = torch::TensorOptions()
                                  .dtype(torch::kBFloat16)
                                  .device(q.device());
              auto opts_lse = torch::TensorOptions()
                                  .dtype(torch::kFloat32)
                                  .device(q.device());

              torch::Tensor output = torch::zeros({total_q,
                                                   NUM_QO_HEADS,
                                                   HEAD_DIM}, opts_out);

              torch::Tensor lse    = torch::full({total_q,
                                                  NUM_QO_HEADS},
                                                  -std::numeric_limits<float>::infinity(),
                                                  opts_lse);

              /* Invoke CUDA kernel --------------------------------------- */
              run_kernel(q, k_cache, v_cache,
                         qo_indptr, kv_indptr, kv_indices,
                         sm_scale,
                         output, lse);

              return std::make_tuple(output, lse);
          },
          pybind11::arg("q"),
          pybind11::arg("k_cache"),
          pybind11::arg("v_cache"),
          pybind11::arg("qo_indptr"),
          pybind11::arg("kv_indptr"),
          pybind11::arg("kv_indices"),
          pybind11::arg("sm_scale") = 1.0 / std::sqrt(static_cast<double>(HEAD_DIM)),
          R"doc(
B200-optimised GQA paged-prefill causal-attention kernel  
Hard-wired for 32 Qo-heads, 4 KV-heads, d = 128, page = 1.  
Returns  (output, lse)
)doc");
}