claude-opus-4-1 / cuda7a69e8

#include <torch/extension.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "kernel.h"

// Helper function to check CUDA errors
#define CHECK_CUDA(x) TORCH_CHECK(x == cudaSuccess, "CUDA error: ", cudaGetErrorString(x))

// Helper function to check tensor properties
void check_input(torch::Tensor tensor, const std::string& name, 
                 c10::ScalarType expected_dtype, 
                 c10::IntArrayRef expected_shape) {
    TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
    TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
    TORCH_CHECK(tensor.dtype() == expected_dtype, 
                name, " must have dtype ", expected_dtype);
    
    if (expected_shape.size() > 0) {
        auto shape = tensor.sizes();
        for (size_t i = 0; i < expected_shape.size(); i++) {
            if (expected_shape[i] >= 0) {
                TORCH_CHECK(shape[i] == expected_shape[i],
                           name, " dimension ", i, " must be ", expected_shape[i],
                           " but got ", shape[i]);
            }
        }
    }
}

torch::Tensor run(torch::Tensor hidden_states, torch::Tensor weight) {
    // Set CUDA device
    c10::cuda::CUDAGuard device_guard(hidden_states.device());
    
    // Check inputs
    TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2-dimensional");
    TORCH_CHECK(weight.dim() == 1, "weight must be 1-dimensional");
    
    int batch_size = hidden_states.size(0);
    int hidden_size = hidden_states.size(1);
    
    TORCH_CHECK(hidden_size == HIDDEN_SIZE, 
                "hidden_size must be ", HIDDEN_SIZE, " but got ", hidden_size);
    TORCH_CHECK(weight.size(0) == HIDDEN_SIZE,
                "weight size must be ", HIDDEN_SIZE, " but got ", weight.size(0));
    
    // Check dtypes
    check_input(hidden_states, "hidden_states", torch::kBFloat16, {-1, HIDDEN_SIZE});
    check_input(weight, "weight", torch::kBFloat16, {HIDDEN_SIZE});
    
    // Allocate output tensor
    auto output = torch::empty_like(hidden_states);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    CHECK_CUDA(launch_rmsnorm_h1536(
        hidden_states.data_ptr(),
        weight.data_ptr(),
        output.data_ptr(),
        batch_size,
        stream
    ));
    
    // Synchronize if needed (PyTorch handles this automatically in most cases)
    // cudaStreamSynchronize(stream);
    
    return output;
}

// Python binding
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "RMSNorm H1536 CUDA kernel",
          py::arg("hidden_states"), py::arg("weight"));
}