claude-opus-4-1 / cudaefa2b2

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

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

// Helper function to check tensor properties
void check_tensor(const torch::Tensor& tensor, const std::string& name, 
                  c10::ScalarType expected_dtype, int expected_dims) {
    TORCH_CHECK(tensor.is_cuda(), name + " must be a CUDA tensor");
    TORCH_CHECK(tensor.is_contiguous(), name + " must be contiguous");
    TORCH_CHECK(tensor.scalar_type() == expected_dtype, 
                name + " must have dtype " + c10::toString(expected_dtype));
    TORCH_CHECK(tensor.dim() == expected_dims, 
                name + " must have " + std::to_string(expected_dims) + " dimensions");
}

torch::Tensor run(torch::Tensor hidden_states, torch::Tensor weight) {
    // Check input tensors
    check_tensor(hidden_states, "hidden_states", c10::ScalarType::BFloat16, 2);
    check_tensor(weight, "weight", c10::ScalarType::BFloat16, 1);
    
    // Get dimensions
    const int batch_size = hidden_states.size(0);
    const int hidden_size = hidden_states.size(1);
    
    // Verify hidden_size
    TORCH_CHECK(hidden_size == HIDDEN_SIZE, 
                "hidden_size must be " + std::to_string(HIDDEN_SIZE) + 
                ", got " + std::to_string(hidden_size));
    TORCH_CHECK(weight.size(0) == HIDDEN_SIZE,
                "weight must have size " + std::to_string(HIDDEN_SIZE) + 
                ", got " + std::to_string(weight.size(0)));
    
    // Allocate output tensor
    torch::Tensor output = torch::empty_like(hidden_states);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Get data pointers - use proper casting
    const __nv_bfloat16* hidden_states_ptr = 
        reinterpret_cast<const __nv_bfloat16*>(hidden_states.data_ptr<c10::BFloat16>());
    const __nv_bfloat16* weight_ptr = 
        reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr<c10::BFloat16>());
    __nv_bfloat16* output_ptr = 
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<c10::BFloat16>());
    
    // Launch kernel
    launch_rmsnorm_h4096(
        hidden_states_ptr,
        weight_ptr,
        output_ptr,
        batch_size,
        stream
    );
    
    // Check for kernel launch errors
    CHECK_CUDA(cudaGetLastError());
    
    return output;
}

// Python bindings
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.doc() = "RMSNorm kernel optimized for hidden_size=4096 on B200 GPU";
    m.def("run", &run, "RMSNorm forward pass",
          py::arg("hidden_states"), 
          py::arg("weight"));
}