claude-opus-4-1 / cuda86c432

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

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

// Helper function to check tensor properties
void check_tensor(const torch::Tensor& t, const std::string& name, 
                  torch::ScalarType dtype, int device_id) {
    if (!t.is_cuda()) {
        throw std::runtime_error(name + " must be a CUDA tensor");
    }
    if (t.device().index() != device_id) {
        throw std::runtime_error(name + " must be on the same CUDA device");
    }
    if (t.scalar_type() != dtype) {
        throw std::runtime_error(name + " must have the correct dtype");
    }
    if (!t.is_contiguous()) {
        throw std::runtime_error(name + " must be contiguous");
    }
}

std::tuple<torch::Tensor, 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
) {
    // Set CUDA device
    int device_id = q.device().index();
    CHECK_CUDA(cudaSetDevice(device_id));
    
    // Ensure tensors are contiguous
    q = q.contiguous();
    k_cache = k_cache.contiguous();
    v_cache = v_cache.contiguous();
    kv_indptr = kv_indptr.contiguous();
    kv_indices = kv_indices.contiguous();
    
    // Check input tensor properties
    check_tensor(q, "q", torch::kBFloat16, device_id);
    check_tensor(k_cache, "k_cache", torch::kBFloat16, device_id);
    check_tensor(v_cache, "v_cache", torch::kBFloat16, device_id);
    check_tensor(kv_indptr, "kv_indptr", torch::kInt32, device_id);
    check_tensor(kv_indices, "kv_indices", torch::kInt32, device_id);
    
    // Get dimensions
    const int batch_size = 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 len_indptr = kv_indptr.size(0);
    const int num_kv_indices = kv_indices.size(0);
    
    // Verify constants match specification
    if (num_qo_heads != NUM_QO_HEADS) {
        throw std::runtime_error("num_qo_heads must be " + std::to_string(NUM_QO_HEADS) + 
                                ", got " + std::to_string(num_qo_heads));
    }
    if (num_kv_heads != NUM_KV_HEADS) {
        throw std::runtime_error("num_kv_heads must be " + std::to_string(NUM_KV_HEADS) + 
                                ", got " + std::to_string(num_kv_heads));
    }
    if (head_dim != HEAD_DIM) {
        throw std::runtime_error("head_dim must be " + std::to_string(HEAD_DIM) + 
                                ", got " + std::to_string(head_dim));
    }
    if (page_size != PAGE_SIZE) {
        throw std::runtime_error("page_size must be " + std::to_string(PAGE_SIZE) + 
                                ", got " + std::to_string(page_size));
    }
    
    // Verify constraints
    if (len_indptr != batch_size + 1) {
        throw std::runtime_error("len_indptr (" + std::to_string(len_indptr) + 
                                ") must equal batch_size + 1 (" + std::to_string(batch_size + 1) + ")");
    }
    
    // Set default sm_scale if not provided or invalid
    if (sm_scale <= 0.0f || std::isnan(sm_scale) || std::isinf(sm_scale)) {
        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({batch_size, num_qo_heads, head_dim}, options_bf16);
    torch::Tensor lse = torch::full({batch_size, num_qo_heads}, 
                                   -std::numeric_limits<float>::infinity(), options_f32);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_gqa_paged_decode(
        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>()),
        reinterpret_cast<const int32_t*>(kv_indptr.data_ptr<int32_t>()),
        reinterpret_cast<const int32_t*>(kv_indices.data_ptr<int32_t>()),
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()),
        reinterpret_cast<float*>(lse.data_ptr<float>()),
        sm_scale,
        batch_size,
        stream
    );
    
    // Check for kernel errors
    CHECK_CUDA(cudaGetLastError());
    
    // Ensure kernel completion for debugging (can be removed for production)
    // CHECK_CUDA(cudaStreamSynchronize(stream));
    
    return std::make_tuple(output, lse);
}

// Python binding
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "GQA Paged Decode 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"));
}