claude-opus-4-1 / cudabc88ee

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <vector>
#include <cmath>
#include <stdexcept>
#include "kernel.h"

// Helper macros for input validation
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_DTYPE_BF16(x) TORCH_CHECK(x.dtype() == torch::kBFloat16, #x " must be bfloat16")
#define CHECK_DTYPE_F32(x) TORCH_CHECK(x.dtype() == torch::kFloat32, #x " must be float32")
#define CHECK_DTYPE_I32(x) TORCH_CHECK(x.dtype() == torch::kInt32, #x " must be int32")

std::tuple<torch::Tensor, torch::Tensor> run(
    torch::Tensor q,
    torch::Tensor k_cache,
    torch::Tensor v_cache,
    torch::Tensor qo_indptr,
    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(qo_indptr);
    CHECK_CUDA(kv_indptr);
    CHECK_CUDA(kv_indices);
    
    CHECK_CONTIGUOUS(q);
    CHECK_CONTIGUOUS(k_cache);
    CHECK_CONTIGUOUS(v_cache);
    CHECK_CONTIGUOUS(qo_indptr);
    CHECK_CONTIGUOUS(kv_indptr);
    CHECK_CONTIGUOUS(kv_indices);
    
    CHECK_DTYPE_BF16(q);
    CHECK_DTYPE_BF16(k_cache);
    CHECK_DTYPE_BF16(v_cache);
    CHECK_DTYPE_I32(qo_indptr);
    CHECK_DTYPE_I32(kv_indptr);
    CHECK_DTYPE_I32(kv_indices);
    
    // Get dimensions
    const int64_t total_q = q.size(0);
    const int64_t num_qo_heads = q.size(1);
    const int64_t head_dim = q.size(2);
    
    const int64_t num_pages = k_cache.size(0);
    const int64_t page_size = k_cache.size(1);
    const int64_t num_kv_heads = k_cache.size(2);
    
    const int64_t len_indptr = qo_indptr.size(0);
    const int64_t num_kv_indices = kv_indices.size(0);
    
    // Verify constants
    TORCH_CHECK(num_qo_heads == NUM_QO_HEADS, 
                "num_qo_heads must be 32, got ", num_qo_heads);
    TORCH_CHECK(num_kv_heads == NUM_KV_HEADS, 
                "num_kv_heads must be 4, got ", num_kv_heads);
    TORCH_CHECK(head_dim == HEAD_DIM, 
                "head_dim must be 128, got ", head_dim);
    TORCH_CHECK(page_size == PAGE_SIZE, 
                "page_size must be 1, got ", page_size);
    
    // Verify shape consistency
    TORCH_CHECK(k_cache.size(3) == head_dim, 
                "k_cache head_dim mismatch");
    TORCH_CHECK(v_cache.size(0) == num_pages && 
                v_cache.size(1) == page_size && 
                v_cache.size(2) == num_kv_heads && 
                v_cache.size(3) == head_dim,
                "v_cache shape mismatch");
    TORCH_CHECK(kv_indptr.size(0) == len_indptr,
                "kv_indptr and qo_indptr must have same length");
    
    // Allocate output tensors
    auto options_bf16 = torch::TensorOptions()
        .dtype(torch::kBFloat16)
        .device(q.device())
        .requires_grad(false);
    
    auto options_f32 = torch::TensorOptions()
        .dtype(torch::kFloat32)
        .device(q.device())
        .requires_grad(false);
    
    torch::Tensor output = torch::zeros({total_q, num_qo_heads, head_dim}, options_bf16);
    torch::Tensor lse = torch::full({total_q, num_qo_heads}, 
                                    -std::numeric_limits<float>::infinity(), 
                                    options_f32);
    
    // Handle empty input case
    if (total_q == 0 || len_indptr <= 1) {
        return std::make_tuple(output, lse);
    }
    
    // Verify constraints
    if (len_indptr > 0) {
        // Use accessor for scalar access to avoid warnings
        auto qo_indptr_acc = qo_indptr.accessor<int32_t, 1>();
        auto kv_indptr_acc = kv_indptr.accessor<int32_t, 1>();
        
        int32_t last_qo_val = qo_indptr_acc[len_indptr - 1];
        int32_t last_kv_val = kv_indptr_acc[len_indptr - 1];
        
        TORCH_CHECK(total_q == last_qo_val, 
                    "total_q (", total_q, ") must equal qo_indptr[-1] (", last_qo_val, ")");
        TORCH_CHECK(num_kv_indices == last_kv_val, 
                    "num_kv_indices (", num_kv_indices, ") must equal kv_indptr[-1] (", last_kv_val, ")");
    }
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_gqa_paged_prefill_kernel(
        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>()),
        qo_indptr.data_ptr<int32_t>(),
        kv_indptr.data_ptr<int32_t>(),
        kv_indices.data_ptr<int32_t>(),
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()),
        lse.data_ptr<float>(),
        sm_scale,
        static_cast<int>(total_q),
        static_cast<int>(num_pages),
        static_cast<int>(len_indptr),
        stream
    );
    
    // Synchronize for error checking in debug mode
    #ifdef DEBUG
    cudaError_t err = cudaStreamSynchronize(stream);
    if (err != cudaSuccess) {
        TORCH_CHECK(false, "CUDA kernel execution error: ", cudaGetErrorString(err));
    }
    #endif
    
    return std::make_tuple(output, lse);
}

// Python bindings
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.doc() = "GQA Paged Prefill Causal Attention CUDA implementation optimized for B200";
    
    m.def("run", &run, 
          "GQA Paged Prefill Causal Attention kernel",
          py::arg("q"),
          py::arg("k_cache"),
          py::arg("v_cache"),
          py::arg("qo_indptr"),
          py::arg("kv_indptr"),
          py::arg("kv_indices"),
          py::arg("sm_scale"));
}