gpt-5-2025-08-07_cuda_353791

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

namespace py = pybind11;

static inline torch::Tensor to_contig_if_needed(const torch::Tensor& t) {
  return t.is_contiguous() ? t : t.contiguous();
}

// Entry point: run(A, B) -> C
// - A: [M, 2048] float16
// - B: [5120, 2048] float16
// - Returns: C: [M, 5120] float16
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
  // Validate shapes/dtypes
  TORCH_CHECK(A.dim() == 2, "A must be 2D [M, 2048]");
  TORCH_CHECK(B.dim() == 2, "B must be 2D [5120, 2048]");
  const int64_t M = A.size(0);
  TORCH_CHECK(A.size(1) == GEMM_K_CONST, "A.shape[1] must be 2048");
  TORCH_CHECK(B.size(0) == GEMM_N_CONST, "B.shape[0] must be 5120");
  TORCH_CHECK(B.size(1) == GEMM_K_CONST, "B.shape[1] must be 2048");

  // Type checks and conversions
  if (A.scalar_type() != at::kHalf) A = A.to(at::kHalf);
  if (B.scalar_type() != at::kHalf) B = B.to(at::kHalf);

  // Move to GPU if needed
  bool inputs_on_cpu = (!A.is_cuda() || !B.is_cuda());
  torch::Device device = (A.is_cuda() ? A.device()
                          : (B.is_cuda() ? B.device()
                                         : torch::Device(torch::kCUDA, 0)));
  if (!A.is_cuda()) A = A.to(device, /*non_blocking=*/false);
  if (!B.is_cuda()) B = B.to(device, /*non_blocking=*/false);

  TORCH_CHECK(A.device().is_cuda() && B.device().is_cuda(),
              "Both A and B must be CUDA tensors for this kernel.");
  TORCH_CHECK(A.get_device() == B.get_device(),
              "A and B must be on the same CUDA device.");

  // Ensure contiguous
  A = to_contig_if_needed(A);
  B = to_contig_if_needed(B);

  // Allocate output on same device
  auto C = torch::empty({M, (int64_t)GEMM_N_CONST},
                        A.options().dtype(at::kHalf).device(A.device()));

  // Launch on current stream
  c10::cuda::CUDAGuard device_guard(A.device());
  c10::cuda::CUDAStream stream_obj = at::cuda::getCurrentCUDAStream();
  cudaStream_t stream = stream_obj.stream();

  launch_gemm_n5120_k2048_kernel(
      reinterpret_cast<const __half*>(A.data_ptr<at::Half>())),
      reinterpret_cast<const __half*>(B.data_ptr<at::Half>()),
      reinterpret_cast<__half*>(C.data_ptr<at::Half>()),
      M,
      stream);

  // If inputs were CPU, move result back to CPU to match expected behavior
  if (inputs_on_cpu) {
    cudaError_t err = cudaStreamSynchronize(stream);
    if (err != cudaSuccess) {
      throw std::runtime_error(std::string("CUDA stream sync failed: ") + cudaGetErrorString(err));
    }
    return C.cpu();
  }

  return C;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("run", &run, "gemm_n5120_k2048 (FP16 TensorCore-optimized)",
        py::arg("A"), py::arg("B"));
}
]]>