gemini-2.5-pro / cuda0ae47c

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

// Macros for concise input validation
#define TORCH_CHECK_DTYPE(x, d) TORCH_CHECK((x).dtype() == (d), #x " dtype must be " #d)
#define TORCH_CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
#define TORCH_CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor")

// Wrapper to check for CUDA errors after kernel launches
#define CUDA_CHECK(call) do { \
    cudaError_t e = call; \
    if (e != cudaSuccess) { \
        TORCH_CHECK(false, "CUDA error in ", #call, ": ", cudaGetErrorString(e)); \
    } \
} while (0)

std::vector<torch::Tensor> run(
    const torch::Tensor& q,
    const torch::Tensor& k_cache,
    const torch::Tensor& v_cache,
    const torch::Tensor& kv_indptr,
    const torch::Tensor& kv_indices,
    const float sm_scale
) {
    // --- Input Validation ---
    TORCH_CHECK_CUDA(q);
    TORCH_CHECK_CUDA(k_cache);
    TORCH_CHECK_CUDA(v_cache);
    TORCH_CHECK_CUDA(kv_indptr);
    TORCH_CHECK_CUDA(kv_indices);

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

    TORCH_CHECK_DTYPE(q, torch::kBFloat16);
    TORCH_CHECK_DTYPE(k_cache, torch::kBFloat16);
    TORCH_CHECK_DTYPE(v_cache, torch::kBFloat16);
    TORCH_CHECK_DTYPE(kv_indptr, torch::kInt32);
    TORCH_CHECK_DTYPE(kv_indices, torch::kInt32);

    TORCH_CHECK(q.dim() == 3, "q must be 3D");
    TORCH_CHECK(k_cache.dim() == 4, "k_cache must be 4D");
    TORCH_CHECK(v_cache.dim() == 4, "v_cache must be 4D");
    TORCH_CHECK(kv_indptr.dim() == 1, "kv_indptr must be 1D");
    TORCH_CHECK(kv_indices.dim() == 1, "kv_indices must be 1D");

    // Extract dimensions from input tensors
    const int batch_size = q.size(0);
    const int num_qo_heads = q.size(1);
    const int head_dim = q.size(2);
    const int page_size = k_cache.size(1);
    const int num_kv_heads = k_cache.size(2);
    const int len_indptr = kv_indptr.size(0);

    // --- Constant Axis Checks from Specification ---
    TORCH_CHECK(num_qo_heads == 32, "num_qo_heads must be 32, but got ", num_qo_heads);
    TORCH_CHECK(num_kv_heads == 4, "num_kv_heads must be 4, but got ", num_kv_heads);
    TORCH_CHECK(head_dim == 128, "head_dim must be 128, but got ", head_dim);
    TORCH_CHECK(page_size == 1, "page_size must be 1, but got ", page_size);

    // --- Constraint Checks from Specification ---
    TORCH_CHECK(len_indptr == batch_size + 1, "len_indptr must be batch_size + 1");
    // The constraint `num_kv_indices == kv_indptr[-1]` is assumed to be held by the caller
    // for performance, as checking it would require a device-to-host synchronization.

    // --- Output Tensor Allocation ---
    auto output_options = torch::TensorOptions().dtype(torch::kBFloat16).device(q.device());
    auto lse_options = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
    
    auto output = torch::zeros({batch_size, num_qo_heads, head_dim}, output_options);
    auto lse = torch::full({batch_size, num_qo_heads}, -std::numeric_limits<float>::infinity(), lse_options);

    // Handle empty batch case
    if (batch_size == 0) {
        return {output, lse};
    }

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

    gqa_paged_decode_h32_kv4_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
    );

    CUDA_CHECK(cudaGetLastError());
    return {output, lse};
}

// --- Pybind11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "GQA Paged Decode Kernel (h32, kv4, 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")
    );
}