gpt-o3 / cuda5a050d

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include "kernel.h"

/* -------------------------------------------------------------------------- */
/*                       Python-visible entry point                           */
/* -------------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(),
                "Input tensors must be on CUDA device");
    TORCH_CHECK(A.dtype() == torch::kFloat16 &&
                B.dtype() == torch::kFloat16,
                "Only fp16 tensors are supported");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2,
                "Inputs must be 2-D matrices");
    TORCH_CHECK(A.size(1) == 4096,
                "A must have shape (M,4096)");
    TORCH_CHECK(B.size(0) == 2048 && B.size(1) == 4096,
                "B must have shape (2048,4096)");

    const int64_t M = A.size(0);

    /* Output tensor ---------------------------------------------------- */
    auto options = torch::TensorOptions()
                     .dtype(torch::kFloat16)
                     .device(torch::kCUDA, A.device().index());
    torch::Tensor C = torch::empty({M, 2048}, options);

    /* Raw device pointers ---------------------------------------------- */
    const __half *A_ptr = reinterpret_cast<const __half*>(A.data_ptr<at::Half>());
    const __half *B_ptr = reinterpret_cast<const __half*>(B.data_ptr<at::Half>());
    __half       *C_ptr = reinterpret_cast<__half*>(C.data_ptr<at::Half>());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    gemm_n2048_k4096_launcher(A_ptr, B_ptr, C_ptr,
                              static_cast<int>(M), stream);

    return C;
}

/* ------------------------------ PyBind11 ---------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("run", &run,
          "GEMM  (A[M,4096] · B[2048,4096]^T → C[M,2048])  "
          "optimised for NVIDIA B200");
}