gemini-2.5-pro / cuda5808cd

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

// Helper to check common tensor properties.
void check_tensor(const torch::Tensor& tensor, const std::string& name, torch::ScalarType dtype) {
    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() == dtype, name, " must have dtype ", dtype, ", but got ", tensor.scalar_type());
}

/**
 * @brief Python-bindable C++ function for fused Add + RMSNorm.
 *
 * This function serves as the interface between PyTorch and the CUDA kernel.
 * It performs extensive input validation before launching the kernel.
 *
 * @param hidden_states The main input tensor of shape [batch_size, 4096].
 * @param residual The tensor to be added to hidden_states, same shape.
 * @param weight The scaling weight tensor of shape [4096].
 * @param eps A small float for numerical stability in the rsqrt operation.
 * @return A new tensor containing the result of the operation.
 */
torch::Tensor fused_add_rmsnorm_h4096(
    const torch::Tensor& hidden_states,
    const torch::Tensor& residual,
    const torch::Tensor& weight,
    double eps = 1e-5) {

    // --- Input Validation ---
    check_tensor(hidden_states, "hidden_states", torch::kBFloat16);
    check_tensor(residual, "residual", torch::kBFloat16);
    check_tensor(weight, "weight", torch::kBFloat16);

    TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be a 2D tensor");
    TORCH_CHECK(residual.dim() == 2, "residual must be a 2D tensor");
    TORCH_CHECK(weight.dim() == 1, "weight must be a 1D tensor");

    const int64_t batch_size = hidden_states.size(0);
    const int64_t hidden_size = hidden_states.size(1);

    TORCH_CHECK(hidden_size == 4096, "This kernel is specialized for hidden_size=4096, but got ", hidden_size);
    TORCH_CHECK(hidden_states.sizes() == residual.sizes(), "hidden_states and residual must have the same shape");
    TORCH_CHECK(weight.size(0) == hidden_size, "weight must have shape [hidden_size]");

    // --- Kernel Execution ---
    auto output = torch::empty_like(hidden_states);
    auto stream = at::cuda::getCurrentCUDAStream();

    launch_fused_add_rmsnorm(
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(hidden_states.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(residual.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr<at::BFloat16>()),
        static_cast<int>(batch_size),
        static_cast<int>(hidden_size),
        static_cast<float>(eps),
        stream
    );

    C10_CUDA_CHECK(cudaGetLastError());
    return output;
}

// pybind11 module definition to expose the C++ function to Python.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &fused_add_rmsnorm_h4096, "Fused Add + RMSNorm Kernel (CUDA BFloat16, H=4096)",
          py::arg("hidden_states"),
          py::arg("residual"),
          py::arg("weight"),
          py::arg("eps") = 1e-5);
}