claude-opus-4-1 / cudab1507e

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

// Helper macro for CUDA error checking
#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, 
                  torch::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 BFloat16");
    TORCH_CHECK(tensor.dim() == expected_dims,
                name + " must have " + std::to_string(expected_dims) + " dimensions");
}

torch::Tensor run(
    torch::Tensor hidden_states,
    torch::Tensor residual,
    torch::Tensor weight
) {
    // Validate input tensors
    check_tensor(hidden_states, "hidden_states", torch::kBFloat16, 2);
    check_tensor(residual, "residual", torch::kBFloat16, 2);
    check_tensor(weight, "weight", torch::kBFloat16, 1);
    
    // Get dimensions
    const int64_t batch_size = hidden_states.size(0);
    const int64_t hidden_size = hidden_states.size(1);
    
    // Verify dimensions
    TORCH_CHECK(hidden_size == HIDDEN_SIZE,
                "hidden_size must be ", HIDDEN_SIZE, ", got ", hidden_size);
    
    TORCH_CHECK(residual.size(0) == batch_size && residual.size(1) == hidden_size,
                "residual shape mismatch: expected [", batch_size, ", ", hidden_size, 
                "], got [", residual.size(0), ", ", residual.size(1), "]");
    
    TORCH_CHECK(weight.size(0) == hidden_size,
                "weight shape mismatch: expected [", hidden_size, 
                "], got [", weight.size(0), "]");
    
    // Allocate output tensor
    torch::Tensor output = torch::empty({batch_size, hidden_size}, 
                                       torch::TensorOptions()
                                           .dtype(torch::kBFloat16)
                                           .device(hidden_states.device()));
    
    // Get current CUDA stream
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    // Launch the kernel
    CHECK_CUDA(launch_fused_add_rmsnorm(
        hidden_states.data_ptr(),
        residual.data_ptr(),
        weight.data_ptr(),
        output.data_ptr(),
        static_cast<int>(batch_size),
        stream
    ));
    
    // Ensure kernel completion for error checking
    CHECK_CUDA(cudaGetLastError());
    
    return output;
}

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