gpt-5 / cuda5c1f52

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_fp16.h>
#include <vector>
#include <stdexcept>
#include <sstream>
#include "kernel.h"

static void check_inputs(const torch::Tensor& A, const torch::Tensor& B) {
  // Shapes: A [M, 4096], B [4096, 4096], dtype float16
  TORCH_CHECK(A.dim() == 2, "A must be 2D [M, 4096]");
  TORCH_CHECK(B.dim() == 2, "B must be 2D [4096, 4096]");
  TORCH_CHECK(A.size(1) == GEMM_K_CONST, "A.shape[1] must be 4096 (K)");
  TORCH_CHECK(B.size(0) == GEMM_N_CONST && B.size(1) == GEMM_K_CONST,
              "B must be [4096, 4096] (N=4096, K=4096)");
  TORCH_CHECK(A.dtype() == torch::kFloat16, "A must be torch.float16");
  TORCH_CHECK(B.dtype() == torch::kFloat16, "B must be torch.float16");
  TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
  TORCH_CHECK(B.is_contiguous(), "B must be contiguous");
}

torch::Tensor run(torch::Tensor A, torch::Tensor B) {
  check_inputs(A, B);
  const int64_t M = A.size(0);

  // Decide device placement
  bool inputs_on_cuda = A.is_cuda() && B.is_cuda();

  torch::Tensor A_cuda = A;
  torch::Tensor B_cuda = B;

  if (!inputs_on_cuda) {
    // Move to CUDA with dtype preserved (float16)
    A_cuda = A.contiguous().to(torch::kCUDA);
    B_cuda = B.contiguous().to(torch::kCUDA);
  } else {
    A_cuda = A.contiguous();
    B_cuda = B.contiguous();
  }

  // Allocate output on CUDA
  auto options = torch::TensorOptions().device(A_cuda.device()).dtype(torch::kFloat16);
  torch::Tensor C_cuda = torch::empty({M, (int64_t)GEMM_N_CONST}, options);

  // Launch kernel on current stream
  auto stream = at::cuda::getCurrentCUDAStream();
  const __half* A_ptr = reinterpret_cast<const __half*>(A_cuda.data_ptr<at::Half>());
  const __half* B_ptr = reinterpret_cast<const __half*>(B_cuda.data_ptr<at::Half>());
  __half* C_ptr = reinterpret_cast<__half*>(C_cuda.data_ptr<at::Half>());

  gemm_n_4096_k_4096_launch(A_ptr, B_ptr, C_ptr, static_cast<int>(M), stream.stream());

  // If inputs were CPU tensors, return result to CPU to match requirement
  if (!inputs_on_cuda) {
    return C_cuda.to(torch::kCPU);
  }
  return C_cuda;
}

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