gemini-2.5-pro / cuda6f8f8e

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <pybind11/pybind11.h>

#include "kernel.h"

#include <stdexcept>
#include <string>

// Helper function to validate tensor properties
void check_tensor(const torch::Tensor& t, const std::string& name) {
    if (!t.is_cuda()) {
        throw std::runtime_error(name + " must be a CUDA tensor");
    }
    if (!t.is_contiguous()) {
        throw std::runtime_error(name + " must be contiguous");
    }
    if (t.scalar_type() != torch::kFloat16) {
        throw std::runtime_error(name + " must have float16 data type");
    }
    if (t.dim() != 2) {
        throw std::runtime_error(name + " must be a 2D tensor");
    }
}

/**
 * @brief Python-bindable function to execute the GEMM operation C = A * B^T.
 *
 * This function serves as the bridge between Python (PyTorch) and the C++/CUDA
 * backend. It handles tensor validation, memory management, and kernel launch.
 *
 * @param A A PyTorch tensor representing matrix A with shape [M, 7168].
 * @param B A PyTorch tensor representing matrix B with shape [256, 7168].
 * @return A PyTorch tensor representing the output matrix C with shape [M, 256].
 */
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // --- Input Validation ---
    check_tensor(A, "Input A");
    check_tensor(B, "Input B");

    const int N_fixed = 256;
    const int K_fixed = 7168;

    if (B.size(0) != N_fixed || B.size(1) != K_fixed) {
        throw std::runtime_error("Input B must have shape [256, 7168]");
    }
    if (A.size(1) != K_fixed) {
        throw std::runtime_error("Input A must have shape [M, 7168]");
    }

    const int M = A.size(0);

    // --- Output Allocation ---
    auto C = torch::empty({M, N_fixed}, A.options());

    // --- Kernel Execution ---
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();

    // PyTorch's at::Half is bit-compatible with CUDA's half type.
    half* c_ptr = reinterpret_cast<half*>(C.data_ptr<at::Half>());
    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>());

    gemm_n256_k7168_launcher(c_ptr, a_ptr, b_ptr, M, stream);

    return C;
}

// --- Pybind11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "High-performance GEMM (C = A @ B.T) for N=256, K=7168 on B200.",
          pybind11::arg("A"), pybind11::arg("B"));
}