claude-opus-4-1 / cuda0302e6

#include <torch/extension.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>
#include "kernel.h"

namespace py = pybind11;

// Helper macros for tensor checking
#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_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

std::vector<torch::Tensor> mla_paged_prefill_forward(
    torch::Tensor q_nope,
    torch::Tensor q_pe,
    torch::Tensor ckv_cache,
    torch::Tensor kpe_cache,
    torch::Tensor qo_indptr,
    torch::Tensor kv_indptr,
    torch::Tensor kv_indices,
    float sm_scale
) {
    // Input validation
    CHECK_INPUT(q_nope);
    CHECK_INPUT(q_pe);
    CHECK_INPUT(ckv_cache);
    CHECK_INPUT(kpe_cache);
    CHECK_INPUT(qo_indptr);
    CHECK_INPUT(kv_indptr);
    CHECK_INPUT(kv_indices);
    
    // Verify data types
    TORCH_CHECK(q_nope.dtype() == torch::kBFloat16, "q_nope must be bfloat16");
    TORCH_CHECK(q_pe.dtype() == torch::kBFloat16, "q_pe must be bfloat16");
    TORCH_CHECK(ckv_cache.dtype() == torch::kBFloat16, "ckv_cache must be bfloat16");
    TORCH_CHECK(kpe_cache.dtype() == torch::kBFloat16, "kpe_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");
    
    // Extract dimensions
    const int total_q = q_nope.size(0);
    const int num_qo_heads = q_nope.size(1);
    const int head_dim_ckv = q_nope.size(2);
    const int head_dim_kpe = q_pe.size(2);
    const int page_size = ckv_cache.size(1);
    const int batch_size = qo_indptr.size(0) - 1;
    
    // Validate constants
    TORCH_CHECK(num_qo_heads == NUM_QO_HEADS, 
                "num_qo_heads must be ", NUM_QO_HEADS, ", got ", num_qo_heads);
    TORCH_CHECK(head_dim_ckv == HEAD_DIM_CKV, 
                "head_dim_ckv must be ", HEAD_DIM_CKV, ", got ", head_dim_ckv);
    TORCH_CHECK(head_dim_kpe == HEAD_DIM_KPE, 
                "head_dim_kpe must be ", HEAD_DIM_KPE, ", got ", head_dim_kpe);
    TORCH_CHECK(page_size == PAGE_SIZE, 
                "page_size must be ", PAGE_SIZE, ", got ", page_size);
    
    // Create output tensors
    auto options_bf16 = torch::TensorOptions()
        .dtype(torch::kBFloat16)
        .device(q_nope.device())
        .requires_grad(false);
    
    auto options_f32 = torch::TensorOptions()
        .dtype(torch::kFloat32)
        .device(q_nope.device())
        .requires_grad(false);
    
    torch::Tensor output = torch::zeros({total_q, num_qo_heads, head_dim_ckv}, options_bf16);
    torch::Tensor lse = torch::full({total_q, num_qo_heads}, -INFINITY, options_f32);
    
    // Get current CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch CUDA kernel
    launch_mla_paged_prefill(
        q_nope.data_ptr(),
        q_pe.data_ptr(),
        ckv_cache.data_ptr(),
        kpe_cache.data_ptr(),
        qo_indptr.data_ptr(),
        kv_indptr.data_ptr(),
        kv_indices.data_ptr(),
        output.data_ptr(),
        lse.data_ptr(),
        sm_scale,
        total_q,
        batch_size,
        stream
    );
    
    // Ensure kernel completion
    cudaError_t err = cudaStreamSynchronize(stream);
    if (err != cudaSuccess) {
        TORCH_CHECK(false, "CUDA kernel execution failed: ", cudaGetErrorString(err));
    }
    
    return {output, lse};
}

// Python binding function matching reference implementation signature
std::vector<torch::Tensor> run(
    torch::Tensor q_nope,
    torch::Tensor q_pe,
    torch::Tensor ckv_cache,
    torch::Tensor kpe_cache,
    torch::Tensor qo_indptr,
    torch::Tensor kv_indptr,
    torch::Tensor kv_indices,
    py::object sm_scale_obj
) {
    float sm_scale;
    
    // Handle sm_scale parameter (can be None, float, or scalar tensor)
    if (sm_scale_obj.is_none()) {
        // Default value: 1/sqrt(head_dim_kpe)
        sm_scale = 1.0f / std::sqrt(static_cast<float>(HEAD_DIM_KPE));
    } else {
        try {
            // Try to extract as float
            sm_scale = py::cast<float>(sm_scale_obj);
        } catch (const py::cast_error&) {
            // Try as tensor
            try {
                torch::Tensor sm_scale_tensor = py::cast<torch::Tensor>(sm_scale_obj);
                TORCH_CHECK(sm_scale_tensor.numel() == 1, "sm_scale must be a scalar");
                
                // Move to CPU if needed
                if (sm_scale_tensor.device().is_cuda()) {
                    sm_scale_tensor = sm_scale_tensor.cpu();
                }
                
                // Extract scalar value
                if (sm_scale_tensor.dtype() == torch::kFloat32) {
                    sm_scale = sm_scale_tensor.item<float>();
                } else if (sm_scale_tensor.dtype() == torch::kFloat64) {
                    sm_scale = static_cast<float>(sm_scale_tensor.item<double>());
                } else {
                    sm_scale = sm_scale_tensor.to(torch::kFloat32).item<float>();
                }
            } catch (...) {
                TORCH_CHECK(false, "sm_scale must be a number, scalar tensor, or None");
            }
        }
    }
    
    return mla_paged_prefill_forward(
        q_nope, q_pe, ckv_cache, kpe_cache,
        qo_indptr, kv_indptr, kv_indices, sm_scale
    );
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.doc() = "MLA paged prefill causal attention CUDA kernel optimized for B200 GPU";
    m.def("run", &run,
          "MLA paged prefill causal attention forward pass",
          py::arg("q_nope"),
          py::arg("q_pe"),
          py::arg("ckv_cache"),
          py::arg("kpe_cache"),
          py::arg("qo_indptr"),
          py::arg("kv_indptr"),
          py::arg("kv_indices"),
          py::arg("sm_scale"),
          py::call_guard<py::gil_scoped_release>());
}