claude-opus-4-1 / cuda462ef5

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

// Helper macros for input validation
#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_BFLOAT16(x) TORCH_CHECK(x.scalar_type() == torch::kBFloat16, #x " must be bfloat16")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x); CHECK_BFLOAT16(x)

torch::Tensor run(
    torch::Tensor hidden_states,
    torch::Tensor residual,
    torch::Tensor weight
) {
    // Input validation
    CHECK_INPUT(hidden_states);
    CHECK_INPUT(residual);
    CHECK_INPUT(weight);
    
    // Check dimensions
    TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2D");
    TORCH_CHECK(residual.dim() == 2, "residual must be 2D");
    TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
    
    const int batch_size = hidden_states.size(0);
    const int hidden_size = hidden_states.size(1);
    
    TORCH_CHECK(hidden_size == 4096, 
                "hidden_size must be 4096 but got ", hidden_size);
    TORCH_CHECK(residual.size(0) == batch_size && residual.size(1) == hidden_size,
                "residual shape mismatch: expected [", batch_size, ", ", hidden_size, 
                "] but got [", residual.size(0), ", ", residual.size(1), "]");
    TORCH_CHECK(weight.size(0) == hidden_size, 
                "weight size must be ", hidden_size, " but got ", weight.size(0));
    
    // Allocate output tensor
    auto output = torch::empty_like(hidden_states);
    
    // Get CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch kernel
    launch_fused_add_rmsnorm(
        hidden_states.data_ptr(),
        residual.data_ptr(),
        weight.data_ptr(),
        output.data_ptr(),
        batch_size,
        stream
    );
    
    // Check for kernel errors
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        TORCH_CHECK(false, "CUDA kernel launch failed: ", cudaGetErrorString(err));
    }
    
    return output;
}

// Python bindings
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Fused Add + RMSNorm for hidden_size=4096",
          py::arg("hidden_states"),
          py::arg("residual"),
          py::arg("weight"));
}