gpt-5 / cudab704b7

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

namespace py = pybind11;

// Entry point called from Python
// Computes C = A @ B^T
// A: [M, 4096] float16
// B: [6144, 4096] float16
// C: [M, 6144] float16
torch::Tensor run(torch::Tensor A, 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(A.size(1) == K_CONST, "A.shape[1] must be 4096");
  TORCH_CHECK(B.size(0) == N_CONST && B.size(1) == K_CONST, "B must be shaped [6144, 4096]");
  TORCH_CHECK(A.dtype() == torch::kHalf, "A must be float16 (torch.half)");
  TORCH_CHECK(B.dtype() == torch::kHalf, "B must be float16 (torch.half)");

  const int64_t M = A.size(0);
  TORCH_CHECK(M >= 0, "M must be non-negative");

  // Determine target device (prefer device of A, else B, else current)
  int device_index = -1;
  if (A.is_cuda()) device_index = A.get_device();
  else if (B.is_cuda()) device_index = B.get_device();
  else device_index = at::cuda::current_device();

  c10::cuda::CUDAGuard guard(device_index);
  auto device = torch::Device(torch::kCUDA, device_index);

  // Ensure tensors are on device and contiguous
  torch::Tensor A_dev = A;
  if (!(A.is_cuda() && (A.get_device() == device_index))) {
    auto A_opts = A.options().device(device);
    A_dev = A.to(A_opts, /*non_blocking=*/false, /*copy=*/true);
  }
  torch::Tensor B_dev = B;
  if (!(B.is_cuda() && (B.get_device() == device_index))) {
    auto B_opts = B.options().device(device);
    B_dev = B.to(B_opts, /*non_blocking=*/false, /*copy=*/true);
  }

  A_dev = A_dev.contiguous();
  B_dev = B_dev.contiguous();

  // Output tensor on device
  torch::Tensor C_dev = torch::empty({M, (int64_t)N_CONST}, A_dev.options());

  // Raw pointers
  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>());

  // Launch on current CUDA stream
  auto stream = at::cuda::getCurrentCUDAStream();
  launch_gemm_n_6144_k_4096(A_ptr, B_ptr, C_ptr, static_cast<int>(M), stream.stream());

  // If inputs were both on CPU, move result back to CPU; otherwise keep on device
  bool inputs_on_cpu = (!A.is_cuda()) && (!B.is_cuda());
  if (inputs_on_cpu) {
    return C_dev.to(torch::kCPU, /*non_blocking=*/false, /*copy=*/true);
  } else {
    return C_dev;
  }
}

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