gpt-o3 / cuda5548d0

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

/*
 * Python-visible entry point.
 * Follows reference signature exactly.
 */
std::vector<torch::Tensor> run(torch::Tensor q,
                               torch::Tensor k_cache,
                               torch::Tensor v_cache,
                               torch::Tensor kv_indptr,
                               torch::Tensor kv_indices,
                               double        sm_scale = 0.08838834764831845 /* 1/sqrt(128) */) {

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

    const int64_t batch_size = q.size(0);

    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({batch_size, kNumQoHeads, kHeadDim}, opts_out);
    torch::Tensor lse    = torch::full({batch_size, kNumQoHeads},
                                       -std::numeric_limits<float>::infinity(),
                                       opts_lse);

    gqa_paged_decode_h32_kv4_d128_ps1_launcher(
        q.contiguous(),
        k_cache.contiguous(),
        v_cache.contiguous(),
        kv_indptr.contiguous(),
        kv_indices.contiguous(),
        static_cast<float>(sm_scale),
        output,
        lse);

    return {output, lse};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run,
          py::arg("q"),
          py::arg("k_cache"),
          py::arg("v_cache"),
          py::arg("kv_indptr"),
          py::arg("kv_indices"),
          py::arg("sm_scale") = 0.08838834764831845 /* 1/sqrt(128) */);
}