gpt-5-2025-08-07 / cuda69e4ed

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

torch::Tensor run(torch::Tensor A, torch::Tensor B) {
  TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor");
  TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor");
  TORCH_CHECK(A.dtype() == torch::kHalf, "A must be float16 (half)");
  TORCH_CHECK(B.dtype() == torch::kHalf, "B must be float16 (half)");
  TORCH_CHECK(A.dim() == 2, "A must be 2D [M, K]");
  TORCH_CHECK(B.dim() == 2, "B must be 2D [N, K]");

  // Enforce expected shapes
  TORCH_CHECK(A.size(1) == K_CONST, "A.shape[1] must be 2048");
  TORCH_CHECK(B.size(0) == N_CONST && B.size(1) == K_CONST, "B.shape must be [128, 2048]");

  // Ensure contiguous
  auto A_c = A.contiguous();
  auto B_c = B.contiguous();

  // Allocate output C [M, 128] on same device/dtype as A
  auto options = A.options();
  auto M = A_c.size(0);
  auto C = torch::empty({M, static_cast<long>(N_CONST)}, options);

  // Launch kernel
  launch_gemm_n128_k2048(A_c, B_c, C);

  // Return result tensor
  return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run, "gemm_n128_k2048 (A[M,2048], B[128,2048]) -> C[M,128]",
        py::arg("A"), py::arg("B"));
}