gpt-5-2025-08-07 / cudad50c9b

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <stdexcept>
#include <vector>
#include <string>
#include <limits>
#include <pybind11/pybind11.h>
#include "kernel.h"

namespace py = pybind11;

// Validate shapes/dtypes
static inline void check_shapes(const torch::Tensor& A, const torch::Tensor& B) {
  if (A.scalar_type() != torch::kFloat16 || B.scalar_type() != torch::kFloat16)
    throw std::invalid_argument("A and B must be float16 (torch.float16) tensors");
  if (A.dim() != 2 || B.dim() != 2)
    throw std::invalid_argument("A and B must be 2D tensors");
  if (A.size(1) != GEMM_K)
    throw std::invalid_argument("A shape must be [M, 4096]");
  if (B.size(0) != GEMM_N || B.size(1) != GEMM_K)
    throw std::invalid_argument("B shape must be [2048, 4096]");
}

// Core implementation: C = A @ B.T, where
// A: [M, 4096] float16
// B: [2048, 4096] float16
// C: [M, 2048] float16
static torch::Tensor run_impl(torch::Tensor A_in, torch::Tensor B_in) {
  check_shapes(A_in, B_in);

  // Determine target CUDA device
  int target_device = -1;
  if (A_in.is_cuda()) target_device = A_in.get_device();
  else if (B_in.is_cuda()) target_device = B_in.get_device();
  else target_device = at::cuda::current_device();

  c10::cuda::CUDAGuard device_guard(target_device);
  torch::Device device(torch::kCUDA, target_device);

  // Track whether to return result to CPU (if both inputs were CPU)
  bool return_to_cpu = (!A_in.is_cuda() && !B_in.is_cuda());

  // Move inputs to the target CUDA device if needed and make contiguous
  torch::Tensor A = (A_in.is_cuda() && A_in.get_device() == target_device) ? A_in : A_in.to(device);
  torch::Tensor B = (B_in.is_cuda() && B_in.get_device() == target_device) ? B_in : B_in.to(device);

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

  const int64_t M64 = A.size(0);
  if (M64 < 0 || M64 > std::numeric_limits<int>::max())
    throw std::invalid_argument("M dimension is out of supported int32 range");
  int M = static_cast<int>(M64);

  auto options = A.options().device(device);
  auto C_dev = torch::empty({M64, static_cast<int64_t>(GEMM_N)}, options); // float16 on target GPU device

  // Launch kernel on current CUDA stream
  auto stream = at::cuda::getCurrentCUDAStream();
  const __half* Ap = reinterpret_cast<const __half*>(A.data_ptr<at::Half>());
  const __half* Bp = reinterpret_cast<const __half*>(B.data_ptr<at::Half>());
  __half* Cp = reinterpret_cast<__half*>(C_dev.data_ptr<at::Half>());

  gemm_n2048_k4096_launcher(Ap, Bp, Cp, M, stream.stream());

  // Optionally move result back to CPU if inputs were CPU
  if (return_to_cpu) {
    return C_dev.to(torch::kCPU);
  }
  return C_dev;
}

// Python-exposed entry that accepts both args and kwargs
static torch::Tensor run_py(py::args args, py::kwargs kwargs) {
  torch::Tensor A, B;
  if (args.size() == 2) {
    try {
      A = args[0].cast<torch::Tensor>();
      B = args[1].cast<torch::Tensor>();
    } catch (const std::exception& e) {
      throw std::invalid_argument(std::string("Invalid positional arguments: ") + e.what());
    }
  } else {
    // Expect kwargs: A=..., B=...
    if (!kwargs.contains("A") || !kwargs.contains("B")) {
      throw std::invalid_argument("Expected arguments (A, B) or kwargs A=..., B=...");
    }
    try {
      A = kwargs["A"].cast<torch::Tensor>();
      B = kwargs["B"].cast<torch::Tensor>();
    } catch (const std::exception& e) {
      throw std::invalid_argument(std::string("Invalid keyword arguments: ") + e.what());
    }
  }
  return run_impl(A, B);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.doc() = "Optimized GEMM kernel for N=2048, K=4096 using Tensor Cores (B200-optimized)";
  m.def("run", &run_py, "Run GEMM: C = A @ B.T for A[M,4096], B[2048,4096] (float16)");
}