gpt-o3_cuda_696722

/*
 *  main.cpp
 *
 *  Thin C++ front-end for the GEMM launcher so that the module can be used
 *  directly from Python.  The exposed `run` function mirrors the reference
 *  implementation:
 *
 *      C = torch.matmul(A , B.T)
 */

#include "kernel.h"

#include <torch/extension.h>

/* --------------------------------------------------------------------- */
/*  Public API exposed to Python                                         */
/* --------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor A, torch::Tensor B,
                  py::kwargs /*unused*/ = {})
{
    TORCH_CHECK(A.device().is_cuda() && B.device().is_cuda(),
                "Input tensors must be on CUDA device");

    const auto M = A.size(0);
    auto C = torch::empty({M, 28672},
                          torch::dtype(at::kHalf).device(A.device()));

    /* Perform the GEMM on the current CUDA stream */
    launch_gemm_n_28672_k_4096(A, B, C);

    return C;
}

/* --------------------------------------------------------------------- */
/*  PyBind11 module                                                      */
/* --------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run,
          "GEMM:  C[M,28672] = A[M,4096] * B[28672,4096]^T  (FP16)",
          py::arg("A"),
          py::arg("B"));
}