claude-opus-4-1 / cuda8eba35

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

// Main run function
torch::Tensor run(
    torch::Tensor hidden_states,
    torch::Tensor residual,
    torch::Tensor weight
) {
    // Validate inputs
    TORCH_CHECK(hidden_states.is_cuda(), "hidden_states must be a CUDA tensor");
    TORCH_CHECK(residual.is_cuda(), "residual must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    
    TORCH_CHECK(hidden_states.dtype() == torch::kBFloat16, "hidden_states must be bfloat16");
    TORCH_CHECK(residual.dtype() == torch::kBFloat16, "residual must be bfloat16");
    TORCH_CHECK(weight.dtype() == torch::kBFloat16, "weight must be bfloat16");
    
    // Check dimensions
    auto hidden_shape = hidden_states.sizes();
    auto residual_shape = residual.sizes();
    auto weight_shape = weight.sizes();
    
    TORCH_CHECK(hidden_shape.size() == 2, "hidden_states must be 2D");
    TORCH_CHECK(residual_shape.size() == 2, "residual must be 2D");
    TORCH_CHECK(weight_shape.size() == 1, "weight must be 1D");
    
    int batch_size = hidden_shape[0];
    int hidden_size = hidden_shape[1];
    
    TORCH_CHECK(hidden_size == HIDDEN_SIZE, "hidden_size must be 2048");
    TORCH_CHECK(residual_shape[0] == batch_size, "residual batch size mismatch");
    TORCH_CHECK(residual_shape[1] == hidden_size, "residual hidden size mismatch");
    TORCH_CHECK(weight_shape[0] == hidden_size, "weight size mismatch");
    
    // Ensure contiguous tensors
    hidden_states = hidden_states.contiguous();
    residual = residual.contiguous();
    weight = weight.contiguous();
    
    // 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_h2048(
        hidden_states.data_ptr(),
        residual.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 bindings
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Fused Add RMSNorm H2048",
          py::arg("hidden_states"),
          py::arg("residual"),
          py::arg("weight"));
}