claude-opus-4-1-20250805 / cudaa1d4a7

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

// Main run function exposed to Python
torch::Tensor run(torch::Tensor A, torch::Tensor B) {
    // Validate input tensors
    TORCH_CHECK(A.dim() == 2, "A must be a 2D tensor");
    TORCH_CHECK(B.dim() == 2, "B must be a 2D tensor");
    TORCH_CHECK(A.scalar_type() == torch::ScalarType::Half, "A must be float16");
    TORCH_CHECK(B.scalar_type() == torch::ScalarType::Half, "B must be float16");
    TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor");
    TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor");
    TORCH_CHECK(A.device() == B.device(), "A and B must be on the same device");
    
    // 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 match specification
    TORCH_CHECK(K_A == K_FIXED, "A must have K=4096 columns, got ", K_A);
    TORCH_CHECK(K_B == K_FIXED, "B must have K=4096 columns, got ", K_B);
    TORCH_CHECK(N == N_FIXED, "B must have N=2048 rows, got ", N);
    
    // Ensure contiguous memory layout for optimal memory access
    A = A.contiguous();
    B = B.contiguous();
    
    // Create output tensor on the same device as inputs
    auto options = torch::TensorOptions()
        .dtype(torch::kHalf)
        .device(A.device())
        .requires_grad(false);
    torch::Tensor C = torch::zeros({M, static_cast<int64_t>(N_FIXED)}, options);
    
    // Get CUDA stream from PyTorch
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Get raw pointers to tensor data
    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 CUDA kernel
    launch_gemm(A_ptr, B_ptr, C_ptr, static_cast<int>(M), stream);
    
    // Check for any CUDA errors
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        TORCH_CHECK(false, "CUDA kernel error: ", cudaGetErrorString(err));
    }
    
    return C;
}

// Python module definition
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Optimized GEMM kernel for B200 GPU",
          py::arg("A"), py::arg("B"));
}