gpt-5 / cuda8ba217

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <stdexcept>
#include <string>
#include "kernel.h"

namespace py = pybind11;

static void validate_inputs(const torch::Tensor& A, const torch::Tensor& B) {
    if (!A.is_cuda() || !B.is_cuda())
        throw std::invalid_argument("A and B must be CUDA tensors");
    if (A.scalar_type() != at::kHalf || B.scalar_type() != at::kHalf)
        throw std::invalid_argument("A and B must be float16 (Half) tensors");
    if (A.dim() != 2 || B.dim() != 2)
        throw std::invalid_argument("A and B must be 2D tensors");
    if (A.size(1) != CONST_K)
        throw std::invalid_argument("A.shape[1] must be 4096");
    if (B.size(0) != CONST_N || B.size(1) != CONST_K)
        throw std::invalid_argument("B must have shape [28672, 4096]");
    if (A.device().index() != B.device().index())
        throw std::invalid_argument("A and B must be on the same CUDA device");
}

torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    validate_inputs(A, B);

    if (!A.is_contiguous()) A = A.contiguous();
    if (!B.is_contiguous()) B = B.contiguous();

    const int64_t M = A.size(0);
    auto options = A.options();
    torch::Tensor C = torch::empty({M, (int64_t)CONST_N}, options);

    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().stream();

    gemm_n_28672_k_4096(A_ptr, B_ptr, C_ptr, M, stream);

    CUDA_CHECK(cudaGetLastError());

    return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "gemm_n_28672_k_4096 (CUDA, cuBLASLt if available)",
          py::arg("A"), py::arg("B"));
}