claude-opus-4-1 / cuda8b7225

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

#include "kernel.h"

// Error checking macro
#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_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

// Main run function
torch::Tensor run(
    torch::Tensor hidden_states,
    torch::Tensor weight
) {
    // Validate inputs
    CHECK_INPUT(hidden_states);
    CHECK_INPUT(weight);
    
    // Check dimensions
    TORCH_CHECK(hidden_states.dim() == 2, 
                "hidden_states must be 2D tensor, got ", hidden_states.dim(), "D");
    TORCH_CHECK(weight.dim() == 1, 
                "weight must be 1D tensor, got ", weight.dim(), "D");
    
    const int64_t batch_size = hidden_states.size(0);
    const int64_t hidden_size = hidden_states.size(1);
    
    TORCH_CHECK(hidden_size == HIDDEN_SIZE, 
                "hidden_size must be 128, got ", hidden_size);
    TORCH_CHECK(weight.size(0) == HIDDEN_SIZE, 
                "weight size must be 128, got ", weight.size(0));
    
    // Check data types
    TORCH_CHECK(hidden_states.scalar_type() == torch::kBFloat16, 
                "hidden_states must be bfloat16, got ", hidden_states.scalar_type());
    TORCH_CHECK(weight.scalar_type() == torch::kBFloat16, 
                "weight must be bfloat16, got ", weight.scalar_type());
    
    // Ensure same device
    TORCH_CHECK(hidden_states.device() == weight.device(),
                "hidden_states and weight must be on the same device");
    
    // Set device guard
    c10::cuda::CUDAGuard device_guard(hidden_states.device());
    
    // Allocate output tensor
    auto output = torch::empty_like(hidden_states);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_rmsnorm_h128(
        hidden_states.data_ptr(),
        weight.data_ptr(),
        output.data_ptr(),
        static_cast<int>(batch_size),
        stream
    );
    
    // Check for errors
    auto error = cudaGetLastError();
    TORCH_CHECK(error == cudaSuccess, 
                "CUDA kernel launch failed: ", cudaGetErrorString(error));
    
    return output;
}

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