gpt-5 / cudabd7484

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

static inline void check_inputs(const torch::Tensor& A, const torch::Tensor& B) {
  TORCH_CHECK(A.dim() == 2, "A must be 2D [M, K]");
  TORCH_CHECK(B.dim() == 2, "B must be 2D [N, K]");
  TORCH_CHECK(B.size(0) == GEMM_N_CONST, "B.size(0) must be 4096 (N constant)");
  TORCH_CHECK(A.size(1) == GEMM_K_CONST, "A.size(1) must be 14336 (K constant)");
  TORCH_CHECK(B.size(1) == GEMM_K_CONST, "B.size(1) must be 14336 (K constant)");
}

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

  // Select device
  at::Device target_device = at::Device(at::kCUDA, 0);
  if (A.is_cuda()) target_device = A.device();
  else if (B.is_cuda()) target_device = B.device();

  // Convert dtypes to half and move to target device
  auto a_opt = torch::TensorOptions().dtype(torch::kFloat16).device(target_device);
  auto b_opt = torch::TensorOptions().dtype(torch::kFloat16).device(target_device);
  auto out_opt = torch::TensorOptions().dtype(torch::kFloat16).device(target_device);

  torch::Tensor A_dev = A.to(a_opt, /*non_blocking=*/true).contiguous();
  torch::Tensor B_dev = B.to(b_opt, /*non_blocking=*/true).contiguous();

  const int64_t M = A.size(0);
  torch::Tensor C_dev = torch::empty({M, (int64_t)GEMM_N_CONST}, out_opt);

  // Launch CUDA kernel on the current stream associated with the chosen device
  c10::cuda::CUDAGuard device_guard(target_device);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  const __half* A_ptr = reinterpret_cast<const __half*>(A_dev.data_ptr<at::Half>());
  const __half* B_ptr = reinterpret_cast<const __half*>(B_dev.data_ptr<at::Half>());
  __half* C_ptr = reinterpret_cast<__half*>(C_dev.data_ptr<at::Half>());

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

  // If both inputs were CPU, return the result on CPU to match the reference behavior
  if (!A.is_cuda() && !B.is_cuda()) {
    return C_dev.to(torch::kCPU);
  }
  return C_dev;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run, "gemm_n_4096_k_14336 (CUDA)");
}