claude-opus-4-1-20250805 / cuda9a3a58

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <vector>
#include <stdexcept>
#include <sstream>
#include "kernel.h"

// Validation macros
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_DTYPE(x) TORCH_CHECK(x.dtype() == torch::kFloat16, #x " must be float16")
#define CHECK_DIMS(x, d) TORCH_CHECK(x.dim() == d, #x " must be " #d "-dimensional")

// CUDA error checking
inline void checkCudaError(cudaError_t error, const char* msg) {
    if (error != cudaSuccess) {
        std::stringstream ss;
        ss << "CUDA error: " << msg << " - " << cudaGetErrorString(error);
        throw std::runtime_error(ss.str());
    }
}

#define CUDA_CHECK(call, msg) checkCudaError((call), msg)

torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // Validate inputs
    CHECK_CUDA(A);
    CHECK_CUDA(B);
    CHECK_CONTIGUOUS(A);
    CHECK_CONTIGUOUS(B);
    CHECK_DTYPE(A);
    CHECK_DTYPE(B);
    CHECK_DIMS(A, 2);
    CHECK_DIMS(B, 2);
    
    // Get dimensions
    const int64_t M = A.size(0);
    const int64_t K_A = A.size(1);
    const int64_t N = B.size(0);
    const int64_t K_B = B.size(1);
    
    // Validate dimensions
    TORCH_CHECK(K_A == 2048, "A must have K=2048, got K=" + std::to_string(K_A));
    TORCH_CHECK(N == 128, "B must have N=128, got N=" + std::to_string(N));
    TORCH_CHECK(K_B == 2048, "B must have K=2048, got K=" + std::to_string(K_B));
    TORCH_CHECK(M > 0, "M must be positive, got M=" + std::to_string(M));
    
    // Create output tensor
    auto options = torch::TensorOptions()
        .dtype(torch::kFloat16)
        .device(A.device())
        .requires_grad(false);
    
    torch::Tensor C = torch::zeros({M, N}, options);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Get device pointers
    const half* A_ptr = reinterpret_cast<const half*>(A.data_ptr<at::Half>());
    const half* B_ptr = reinterpret_cast<const half*>(B.data_ptr<at::Half>());
    half* C_ptr = reinterpret_cast<half*>(C.data_ptr<at::Half>());
    
    // Launch kernel
    cudaError_t error = launch_gemm(A_ptr, B_ptr, C_ptr, static_cast<int>(M), stream);
    CUDA_CHECK(error, "Kernel launch failed");
    
    // Synchronize to catch any execution errors
    CUDA_CHECK(cudaStreamSynchronize(stream), "Kernel execution failed");
    
    return C;
}

// Python bindings
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Optimized GEMM with N=128, K=2048 for B200 GPU",
          py::arg("A"), py::arg("B"));
}