claude-opus-4-1 / cuda29819a

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cmath>
#include <limits>
#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_INT32(x) TORCH_CHECK(x.dtype() == torch::kInt32, #x " must be int32")
#define CHECK_DTYPE_F32(x) TORCH_CHECK(x.dtype() == torch::kFloat32, #x " must be float32")

std::tuple<torch::Tensor, torch::Tensor> run(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    torch::Tensor qo_indptr,
    torch::Tensor kv_indptr,
    torch::optional<double> sm_scale_opt = torch::nullopt
) {
    // Input validation
    CHECK_CUDA(q);
    CHECK_CUDA(k);
    CHECK_CUDA(v);
    CHECK_CUDA(qo_indptr);
    CHECK_CUDA(kv_indptr);
    
    CHECK_CONTIGUOUS(q);
    CHECK_CONTIGUOUS(k);
    CHECK_CONTIGUOUS(v);
    CHECK_CONTIGUOUS(qo_indptr);
    CHECK_CONTIGUOUS(kv_indptr);
    
    CHECK_DTYPE_BF16(q);
    CHECK_DTYPE_BF16(k);
    CHECK_DTYPE_BF16(v);
    CHECK_DTYPE_INT32(qo_indptr);
    CHECK_DTYPE_INT32(kv_indptr);
    
    // 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 total_kv = k.size(0);
    const int num_kv_heads = k.size(1);
    
    const int len_indptr = qo_indptr.size(0);
    
    // Verify shape consistency
    TORCH_CHECK(k.size(2) == head_dim, "k head_dim mismatch");
    TORCH_CHECK(v.size(0) == total_kv, "v total_kv mismatch");
    TORCH_CHECK(v.size(1) == num_kv_heads, "v num_kv_heads mismatch");
    TORCH_CHECK(v.size(2) == head_dim, "v head_dim mismatch");
    TORCH_CHECK(kv_indptr.size(0) == len_indptr, "kv_indptr length mismatch");
    
    // Verify constants
    TORCH_CHECK(num_qo_heads == NUM_QO_HEADS, 
                "num_qo_heads must be 32, got " + std::to_string(num_qo_heads));
    TORCH_CHECK(num_kv_heads == NUM_KV_HEADS, 
                "num_kv_heads must be 4, got " + std::to_string(num_kv_heads));
    TORCH_CHECK(head_dim == HEAD_DIM, 
                "head_dim must be 128, got " + std::to_string(head_dim));
    
    // Verify constraints
    if (len_indptr > 0) {
        auto qo_indptr_cpu = qo_indptr.cpu();
        auto kv_indptr_cpu = kv_indptr.cpu();
        
        int32_t last_qo = qo_indptr_cpu[-1].item<int32_t>();
        int32_t last_kv = kv_indptr_cpu[-1].item<int32_t>();
        
        TORCH_CHECK(total_q == last_qo, 
                    "total_q must equal qo_indptr[-1], got " + std::to_string(total_q) + 
                    " vs " + std::to_string(last_qo));
        TORCH_CHECK(total_kv == last_kv, 
                    "total_kv must equal kv_indptr[-1], got " + std::to_string(total_kv) + 
                    " vs " + std::to_string(last_kv));
    }
    
    // Set default sm_scale if not provided
    float sm_scale = sm_scale_opt.has_value() 
        ? static_cast<float>(sm_scale_opt.value())
        : 1.0f / std::sqrt(static_cast<float>(head_dim));
    
    // Allocate output tensors
    auto options_bf16 = torch::TensorOptions()
        .dtype(torch::kBFloat16)
        .device(q.device());
    auto options_f32 = torch::TensorOptions()
        .dtype(torch::kFloat32)
        .device(q.device());
    
    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
    if (total_q == 0 || len_indptr <= 1) {
        return std::make_tuple(output, lse);
    }
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_gqa_ragged_prefill(
        reinterpret_cast<const __nv_bfloat16*>(q.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(k.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(v.data_ptr<at::BFloat16>()),
        qo_indptr.data_ptr<int32_t>(),
        kv_indptr.data_ptr<int32_t>(),
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()),
        lse.data_ptr<float>(),
        sm_scale,
        len_indptr,
        total_q,
        total_kv,
        stream
    );
    
    // Synchronize to ensure kernel completion
    cudaError_t err = cudaStreamSynchronize(stream);
    TORCH_CHECK(err == cudaSuccess, 
                "CUDA kernel execution failed: ", cudaGetErrorString(err));
    
    return std::make_tuple(output, lse);
}

// Python bindings
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, 
          "GQA Ragged Prefill Causal Attention (BF16)",
          py::arg("q"),
          py::arg("k"),
          py::arg("v"),
          py::arg("qo_indptr"),
          py::arg("kv_indptr"),
          py::arg("sm_scale") = py::none());
}