claude-opus-4-1 / cudafbec80

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

// Helper function to check CUDA errors
#define CUDA_CHECK(call) \
    do { \
        cudaError_t error = call; \
        if (error != cudaSuccess) { \
            throw std::runtime_error(std::string("CUDA error at ") + __FILE__ + ":" + \
                std::to_string(__LINE__) + " - " + cudaGetErrorString(error)); \
        } \
    } while(0)

torch::Tensor run(
    torch::Tensor hidden_states,
    torch::Tensor weight
) {
    // Input validation
    TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2D tensor");
    TORCH_CHECK(weight.dim() == 1, "weight must be 1D tensor");
    TORCH_CHECK(hidden_states.size(1) == 2048, "hidden_size must be 2048");
    TORCH_CHECK(weight.size(0) == 2048, "weight size must be 2048");
    TORCH_CHECK(hidden_states.scalar_type() == torch::ScalarType::BFloat16, 
                "hidden_states must be BFloat16");
    TORCH_CHECK(weight.scalar_type() == torch::ScalarType::BFloat16, 
                "weight must be BFloat16");
    TORCH_CHECK(hidden_states.is_cuda(), "hidden_states must be on CUDA device");
    TORCH_CHECK(weight.is_cuda(), "weight must be on CUDA device");
    TORCH_CHECK(hidden_states.is_contiguous(), "hidden_states must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");
    
    // Get dimensions
    const int batch_size = hidden_states.size(0);
    
    // Create output tensor with the same properties as input
    auto output = torch::empty_like(hidden_states);
    
    // Get CUDA stream from PyTorch
    c10::cuda::CUDAStream torch_stream = c10::cuda::getCurrentCUDAStream();
    cudaStream_t stream = torch_stream.stream();
    
    // Get data pointers - use correct casting for bfloat16
    const __nv_bfloat16* hidden_states_ptr = reinterpret_cast<const __nv_bfloat16*>(
        hidden_states.data_ptr<at::BFloat16>()
    );
    const __nv_bfloat16* weight_ptr = reinterpret_cast<const __nv_bfloat16*>(
        weight.data_ptr<at::BFloat16>()
    );
    __nv_bfloat16* output_ptr = reinterpret_cast<__nv_bfloat16*>(
        output.data_ptr<at::BFloat16>()
    );
    
    // Launch kernel
    launch_rmsnorm_h2048(
        hidden_states_ptr,
        weight_ptr,
        output_ptr,
        batch_size,
        stream
    );
    
    // Check for errors after kernel launch
    CUDA_CHECK(cudaGetLastError());
    
    // PyTorch manages stream synchronization, so no need for explicit sync
    
    return output;
}

// Python bindings
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "RMSNorm kernel for hidden_size=2048",
          py::arg("hidden_states"),
          py::arg("weight"));
}