claude-opus-4-1 / cudaa6c279

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

namespace py = pybind11;

// Helper macro for CUDA error checking
#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            throw std::runtime_error(std::string("CUDA error at ") + __FILE__ + ":" + \
                                    std::to_string(__LINE__) + " - " + cudaGetErrorString(err)); \
        } \
    } while(0)

// Main run function
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 = -1.0f
) {
    // Input validation
    TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bfloat16");
    TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bfloat16");
    TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bfloat16");
    TORCH_CHECK(qo_indptr.dtype() == torch::kInt32, "qo_indptr must be int32");
    TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be int32");
    TORCH_CHECK(kv_indices.dtype() == torch::kInt32, "kv_indices must be int32");
    
    TORCH_CHECK(q.is_cuda(), "q must be on CUDA device");
    TORCH_CHECK(k_cache.is_cuda(), "k_cache must be on CUDA device");
    TORCH_CHECK(v_cache.is_cuda(), "v_cache must be on CUDA device");
    TORCH_CHECK(qo_indptr.is_cuda(), "qo_indptr must be on CUDA device");
    TORCH_CHECK(kv_indptr.is_cuda(), "kv_indptr must be on CUDA device");
    TORCH_CHECK(kv_indices.is_cuda(), "kv_indices must be on CUDA device");
    
    TORCH_CHECK(q.is_contiguous(), "q must be contiguous");
    TORCH_CHECK(k_cache.is_contiguous(), "k_cache must be contiguous");
    TORCH_CHECK(v_cache.is_contiguous(), "v_cache must be contiguous");
    TORCH_CHECK(qo_indptr.is_contiguous(), "qo_indptr must be contiguous");
    TORCH_CHECK(kv_indptr.is_contiguous(), "kv_indptr must be contiguous");
    TORCH_CHECK(kv_indices.is_contiguous(), "kv_indices must be contiguous");
    
    // Get dimensions
    const int total_q = q.size(0);
    const int num_qo_heads = q.size(1);
    const int head_dim = q.size(2);
    const int num_pages = k_cache.size(0);
    const int page_size = k_cache.size(1);
    const int num_kv_heads = k_cache.size(2);
    const int batch_size = qo_indptr.size(0) - 1;
    
    // Validate constants
    TORCH_CHECK(num_qo_heads == NUM_QO_HEADS, 
                "num_qo_heads must be " + std::to_string(NUM_QO_HEADS) + ", got " + std::to_string(num_qo_heads));
    TORCH_CHECK(num_kv_heads == NUM_KV_HEADS, 
                "num_kv_heads must be " + std::to_string(NUM_KV_HEADS) + ", got " + std::to_string(num_kv_heads));
    TORCH_CHECK(head_dim == HEAD_DIM, 
                "head_dim must be " + std::to_string(HEAD_DIM) + ", got " + std::to_string(head_dim));
    TORCH_CHECK(page_size == PAGE_SIZE, 
                "page_size must be " + std::to_string(PAGE_SIZE) + ", got " + std::to_string(page_size));
    
    // Validate 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 num_pages mismatch");
    TORCH_CHECK(v_cache.size(1) == page_size, "v_cache page_size mismatch");
    TORCH_CHECK(v_cache.size(2) == num_kv_heads, "v_cache num_kv_heads mismatch");
    TORCH_CHECK(v_cache.size(3) == head_dim, "v_cache head_dim mismatch");
    TORCH_CHECK(kv_indptr.size(0) == qo_indptr.size(0), "kv_indptr and qo_indptr batch size mismatch");
    
    // Set default sm_scale if not provided
    if (sm_scale < 0) {
        sm_scale = 1.0f / std::sqrt(static_cast<float>(head_dim));
    }
    
    // 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}, -INFINITY, options_f32);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_gqa_paged_prefill(
        reinterpret_cast<const __nv_bfloat16*>(q.data_ptr()),
        reinterpret_cast<const __nv_bfloat16*>(k_cache.data_ptr()),
        reinterpret_cast<const __nv_bfloat16*>(v_cache.data_ptr()),
        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()),
        lse.data_ptr<float>(),
        sm_scale,
        batch_size,
        total_q,
        stream
    );
    
    // Synchronize to ensure kernel completion
    CUDA_CHECK(cudaStreamSynchronize(stream));
    
    return std::make_tuple(output, lse);
}

// Python binding
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, 
          "GQA Paged Prefill Causal Attention (optimized for B200)",
          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") = -1.0f);
}