gpt-o3 / cudaaf0f3d

#include "kernel.h"

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

/*  Public entry point exposed to Python.
    Accepts:
        A : torch.float16  [M , 2048]   (CUDA)
        B : torch.float16  [5120 , 2048](CUDA)
    Returns:
        C : torch.float16  [M , 5120]   (CUDA)   */
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    /*  Basic argument checking  */
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must reside on the GPU");
    TORCH_CHECK(A.scalar_type() == torch::kFloat16 &&
                B.scalar_type() == torch::kFloat16,
                "Inputs must be float16 / half");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2,
                "Inputs must be 2-D matrices");
    TORCH_CHECK(A.size(1) == GEMM_K,
                "A must have shape [M , 2048]");
    TORCH_CHECK(B.size(0) == GEMM_N && B.size(1) == GEMM_K,
                "B must have shape [5120 , 2048]");

    const int64_t M = A.size(0);

    /*  Allocate output tensor on the same device  */
    auto C = torch::empty({M, GEMM_N},
                          torch::dtype(torch::kFloat16).device(A.device()));

    /*  Extract the current CUDA stream used by PyTorch  */
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();

    /*  Launch the optimised GEMM  */
    gemm_n5120_k2048(
        reinterpret_cast<const __half*>(A.data_ptr<at::Half>()),
        reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
        reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
        static_cast<int>(M),
        stream);

    return C;
}

/*  PyBind11 module declaration  */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Optimised GEMM  (B200, fp16)");
}