gpt-5 / cuda5eb89c

#include "kernel.h"

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_bf16.h>
#include <vector>
#include <iostream>

using torch::Tensor;

static inline void check_inputs(
    const Tensor& q_nope,
    const Tensor& q_pe,
    const Tensor& ckv_cache,
    const Tensor& kpe_cache,
    const Tensor& kv_indptr,
    const Tensor& kv_indices
) {
    // Dtype checks
    CHECK_DTYPE(q_nope, torch::kBFloat16);
    CHECK_DTYPE(q_pe,   torch::kBFloat16);
    CHECK_DTYPE(ckv_cache, torch::kBFloat16);
    CHECK_DTYPE(kpe_cache, torch::kBFloat16);
    CHECK_DTYPE(kv_indptr, torch::kInt);
    CHECK_DTYPE(kv_indices, torch::kInt);

    // Dim checks
    TORCH_CHECK(q_nope.dim() == 3, "q_nope must be [B, 16, 512]");
    TORCH_CHECK(q_pe.dim()   == 3, "q_pe must be [B, 16, 64]");
    TORCH_CHECK(ckv_cache.dim() == 3, "ckv_cache must be [P, 1, 512]");
    TORCH_CHECK(kpe_cache.dim() == 3, "kpe_cache must be [P, 1, 64]");
    TORCH_CHECK(kv_indptr.dim() == 1, "kv_indptr must be 1D");
    TORCH_CHECK(kv_indices.dim() == 1, "kv_indices must be 1D");

    // Shape checks (fixed constants)
    TORCH_CHECK(q_nope.size(1) == MLA_NUM_QO_HEADS && q_nope.size(2) == MLA_HEAD_DIM_CKV,
                "q_nope must be [B, 16, 512]");
    TORCH_CHECK(q_pe.size(1) == MLA_NUM_QO_HEADS && q_pe.size(2) == MLA_HEAD_DIM_KPE,
                "q_pe must be [B, 16, 64]");
    TORCH_CHECK(ckv_cache.size(1) == MLA_PAGE_SIZE && ckv_cache.size(2) == MLA_HEAD_DIM_CKV,
                "ckv_cache must be [P, 1, 512] with page_size=1");
    TORCH_CHECK(kpe_cache.size(1) == MLA_PAGE_SIZE && kpe_cache.size(2) == MLA_HEAD_DIM_KPE,
                "kpe_cache must be [P, 1, 64] with page_size=1");

    // Constraint checks
    const int64_t B = q_nope.size(0);
    TORCH_CHECK(kv_indptr.size(0) == B + 1, "len_indptr must be batch_size + 1");

    // kv_indices length equals kv_indptr[-1]
    // Make sure to get a CPU scalar to avoid device sync issues
    Tensor last_cpu = kv_indptr.index({kv_indptr.size(0) - 1}).cpu();
    int64_t kv_count = last_cpu.item<int32_t>();
    TORCH_CHECK(kv_indices.size(0) == kv_count,
                "num_kv_indices must equal kv_indptr[-1]");
}

pybind11::dict run(
    Tensor q_nope,       // bf16 [B, 16, 512]
    Tensor q_pe,         // bf16 [B, 16, 64]
    Tensor ckv_cache,    // bf16 [P, 1, 512]
    Tensor kpe_cache,    // bf16 [P, 1, 64]
    Tensor kv_indptr,    // int32 [B+1]
    Tensor kv_indices,   // int32 [kv_indptr[-1]]
    double sm_scale_d
) {
    check_inputs(q_nope, q_pe, ckv_cache, kpe_cache, kv_indptr, kv_indices);

    const bool inputs_on_cpu =
        !q_nope.is_cuda() || !q_pe.is_cuda() || !ckv_cache.is_cuda() ||
        !kpe_cache.is_cuda() || !kv_indptr.is_cuda() || !kv_indices.is_cuda();

    // Decide device
    c10::Device device = inputs_on_cpu ? c10::Device(c10::kCUDA, 0) : q_nope.device();

    // Make contiguous and move to device if needed
    Tensor q_nope_dev     = q_nope.contiguous().to(device);
    Tensor q_pe_dev       = q_pe.contiguous().to(device);
    Tensor ckv_cache_dev  = ckv_cache.contiguous().to(device);
    Tensor kpe_cache_dev  = kpe_cache.contiguous().to(device);
    Tensor kv_indptr_dev  = kv_indptr.contiguous().to(device);
    Tensor kv_indices_dev = kv_indices.contiguous().to(device);

    const int64_t B = q_nope_dev.size(0);
    const int64_t len_indptr = kv_indptr_dev.size(0);

    // Allocate outputs on device
    Tensor output_dev = torch::empty({B, MLA_NUM_QO_HEADS, MLA_HEAD_DIM_CKV},
                                     torch::dtype(torch::kBFloat16).device(device));
    Tensor lse_dev = torch::empty({B, MLA_NUM_QO_HEADS},
                                  torch::dtype(torch::kFloat32).device(device));

    // Launch kernel
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const __nv_bfloat16* q_nope_ptr = reinterpret_cast<const __nv_bfloat16*>(
        q_nope_dev.data_ptr<c10::BFloat16>());
    const __nv_bfloat16* q_pe_ptr = reinterpret_cast<const __nv_bfloat16*>(
        q_pe_dev.data_ptr<c10::BFloat16>());
    const __nv_bfloat16* ckv_cache_ptr = reinterpret_cast<const __nv_bfloat16*>(
        ckv_cache_dev.data_ptr<c10::BFloat16>());
    const __nv_bfloat16* kpe_cache_ptr = reinterpret_cast<const __nv_bfloat16*>(
        kpe_cache_dev.data_ptr<c10::BFloat16>());

    const int32_t* kv_indptr_ptr = kv_indptr_dev.data_ptr<int32_t>();
    const int32_t* kv_indices_ptr = kv_indices_dev.data_ptr<int32_t>();

    __nv_bfloat16* output_ptr = reinterpret_cast<__nv_bfloat16*>(
        output_dev.data_ptr<c10::BFloat16>());
    float* lse_ptr = lse_dev.data_ptr<float>();

    float sm_scale = static_cast<float>(sm_scale_d);

    mla_paged_decode_h16_ckv512_kpe64_ps1_launcher(
        q_nope_ptr, q_pe_ptr,
        ckv_cache_ptr, kpe_cache_ptr,
        kv_indptr_ptr, kv_indices_ptr,
        static_cast<int>(B),
        static_cast<int>(len_indptr),
        sm_scale,
        output_ptr, lse_ptr,
        stream
    );

    // Check for async launch errors
    cudaError_t err_sync = cudaGetLastError();
    TORCH_CHECK(err_sync == cudaSuccess, "CUDA error after kernel launch: ", cudaGetErrorString(err_sync));

    // If original inputs were on CPU, move outputs back to CPU (synchronize stream first)
    Tensor output = output_dev;
    Tensor lse    = lse_dev;
    if (inputs_on_cpu) {
        // Ensure kernel finished before D2H copy
        cudaStreamSynchronize(stream);
        output = output_dev.cpu();
        lse    = lse_dev.cpu();
    }

    pybind11::dict result;
    result["output"] = output;
    result["lse"] = lse;
    return result;
}

// Important: Use TORCH_EXTENSION_NAME so the generated module name matches what
// torch.utils.cpp_extension.load expects at runtime.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.doc() = "Optimized MLA paged decode kernel for h16 ckv512 kpe64 ps1 (B200)";
    m.def(
        "run",
        &run,
        pybind11::arg("q_nope"),
        pybind11::arg("q_pe"),
        pybind11::arg("ckv_cache"),
        pybind11::arg("kpe_cache"),
        pybind11::arg("kv_indptr"),
        pybind11::arg("kv_indices"),
        pybind11::arg("sm_scale"),
        "Run mla_paged_decode_h16_ckv512_kpe64_ps1 (BF16) on the current CUDA device"
    );
}