gemini-2.5-pro_cuda_977367

#include <torch/extension.h>
#include <pybind11/pybind11.h>
#include <vector>
#include <limits> // for std::numeric_limits
#include "kernel.h"

namespace py = pybind11;

// Helper to check tensor properties
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_BF16(x) TORCH_CHECK(x.scalar_type() == torch::kBFloat16, #x " must be a BFloat16 tensor")
#define CHECK_INT32(x) TORCH_CHECK(x.scalar_type() == torch::kInt, #x " must be an Int32 tensor")

std::vector<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
) {
    // --- Input Validation ---
    CHECK_CUDA(q);
    CHECK_CUDA(k_cache);
    CHECK_CUDA(v_cache);
    CHECK_CUDA(kv_indptr);
    CHECK_CUDA(kv_indices);

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

    CHECK_BF16(q);
    CHECK_BF16(k_cache);
    CHECK_BF16(v_cache);

    CHECK_INT32(kv_indptr);
    CHECK_INT32(kv_indices);

    const int batch_size = q.size(0);
    TORCH_CHECK(q.dim() == 3);
    TORCH_CHECK(q.size(1) == NUM_QO_HEADS);
    TORCH_CHECK(q.size(2) == HEAD_DIM);

    TORCH_CHECK(k_cache.dim() == 4);
    TORCH_CHECK(k_cache.size(1) == 1, "page_size must be 1");
    TORCH_CHECK(k_cache.size(2) == NUM_KV_HEADS);
    TORCH_CHECK(k_cache.size(3) == HEAD_DIM);

    TORCH_CHECK(v_cache.dim() == 4);
    TORCH_CHECK(v_cache.size(1) == 1, "page_size must be 1");
    TORCH_CHECK(v_cache.size(2) == NUM_KV_HEADS);
    TORCH_CHECK(v_cache.size(3) == HEAD_DIM);

    TORCH_CHECK(kv_indptr.dim() == 1);
    TORCH_CHECK(kv_indptr.size(0) == batch_size + 1);

    TORCH_CHECK(kv_indices.dim() == 1);

    // --- Output Allocation and Initialization ---
    // Initialize output to zeros. The kernel will not write to it for empty sequences.
    auto output = torch::zeros_like(q);

    // Initialize LSE to -inf. The kernel will not write to it for empty sequences.
    auto lse_options = torch::TensorOptions()
        .device(q.device())
        .dtype(torch::kFloat32);
    auto lse = torch::full({batch_size, NUM_QO_HEADS}, -std::numeric_limits<float>::infinity(), lse_options);

    // --- Kernel Launch ---
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    gqa_paged_decode_h32_kv8_d128_ps1_launch(
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()),
        lse.data_ptr<float>(),
        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<int>(),
        kv_indices.data_ptr<int>(),
        sm_scale,
        batch_size,
        stream
    );

    return {output, lse};
}

// --- Pybind11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "GQA Paged Decode Kernel (h32_kv8_d128_ps1)",
          py::arg("q"),
          py::arg("k_cache"),
          py::arg("v_cache"),
          py::arg("kv_indptr"),
          py::arg("kv_indices"),
          py::arg("sm_scale")
    );
}