gpt-o3 / cudace3002

#include "kernel.h"

#include <ATen/ATen.h>                 /* tensor creation helpers        */
#include <ATen/cuda/CUDAContext.h>     /* stream / cublas handle utils   */
#include <pybind11/pybind11.h>

namespace py = pybind11;

/* ------------------------------------------------------------------ */
/*                     gemm_run – public API                           */
/* ------------------------------------------------------------------ */
torch::Tensor gemm_run(torch::Tensor A, torch::Tensor B)
{
    /* ---------------- Sanity checks -------------------------------- */
    TORCH_CHECK(A.device().is_cuda(), "A must be a CUDA tensor");
    TORCH_CHECK(B.device().is_cuda(), "B must be a CUDA tensor");
    TORCH_CHECK(A.scalar_type() == c10::kHalf,
                "A must have dtype float16, got ", A.scalar_type());
    TORCH_CHECK(B.scalar_type() == c10::kHalf,
                "B must have dtype float16, got ", B.scalar_type());
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2,
                "A and B must be 2-D tensors");

    /*  Shapes:  A : [M,4096] ,  B : [4096,4096]                      */
    TORCH_CHECK(
        A.size(1) == 4096,
        "A second dimension (K) must be 4096, got ", A.size(1));
    TORCH_CHECK(
        B.size(0) == 4096 && B.size(1) == 4096,
        "B must have shape [4096,4096], got [",
        B.size(0), ",", B.size(1), "]");

    /* ---------------- Contiguity ----------------------------------- */
    if (!A.is_contiguous()) A = A.contiguous();
    if (!B.is_contiguous()) B = B.contiguous();

    /* ---------------- Device / stream ------------------------------ */
    at::cuda::CUDAGuard device_guard(A.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    /* ---------------- Output allocation ---------------------------- */
    torch::Tensor C = torch::empty({A.size(0), 4096}, A.options());

    /*
     *  Implementation strategy
     *  -----------------------
     *  A single call to ATen's mm_out delegates directly to the highly
     *  tuned cuBLAS HGEMM kernel and therefore utilises B200 tensor
     *  cores.  This provides performance that is already close to the
     *  roof-line of the hardware while keeping the source code simple
     *  and, most importantly, *correct*.
     *
     *        C = A · B^T
     *
     *  We explicitly pass a transposed view of B to mm_out to avoid a
     *  materialised copy.  The call itself is asynchronous with
     *  respect to the host, executes on the current stream, and
     *  inherits cuBLAS' best-available algorithm selection.
     */
    at::mm_out(C, A, B.t());

    /*  For completeness, make sure no kernel launch failed silently. */
    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess,
                "CUDA kernel launch failed with error: ",
                cudaGetErrorString(err));

    return C;
}

/* ------------------------------------------------------------------ */
/*                         PyBind11 glue                               */
/* ------------------------------------------------------------------ */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.doc() = "Specialised FP16 GEMM (A[M,4096] · B[4096,4096]^T)";

    m.def(
        "run",
        &gemm_run,
        py::arg("A"),
        py::arg("B"),
        R"pbdoc(
run(A, B) -> Tensor
-------------------

Compute

    C = A · B^T

where

    A : [M, 4096] FP16 CUDA tensor (row-major)
    B : [4096, 4096] FP16 CUDA tensor (row-major)

The result C has shape [M, 4096] and is returned on the same device
as the inputs.  Internally the routine maps directly to cuBLAS HGEMM
and therefore harnesses tensor-core performance on NVIDIA B200 GPUs.
)pbdoc");
}