gemini-2.5-pro_cuda_428669

#include "kernel.h"

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

// Helper to check tensor properties
void check_tensor(const torch::Tensor& tensor, const std::string& name, torch::ScalarType dtype, int64_t last_dim) {
    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");
    if (last_dim > 0) {
        TORCH_CHECK(tensor.size(-1) == last_dim, name, " must have last dimension of size ", last_dim);
    }
}

// C++ entry point, exposed to Python
torch::Tensor run(
    const torch::Tensor& hidden_states,
    const torch::Tensor& residual,
    const torch::Tensor& weight,
    double eps) {
    
    // --- Input Validation ---
    const int64_t hidden_size = 7168;
    check_tensor(hidden_states, "hidden_states", torch::kBFloat16, hidden_size);
    check_tensor(residual, "residual", torch::kBFloat16, hidden_size);
    check_tensor(weight, "weight", torch::kBFloat16, hidden_size);
    
    TORCH_CHECK(hidden_states.sizes() == residual.sizes(), "hidden_states and residual must have the same shape");
    TORCH_CHECK(weight.dim() == 1, "weight must be a 1D tensor");

    // --- Output Allocation ---
    auto output = torch::empty_like(hidden_states);

    // --- Kernel Launch ---
    // Get the current CUDA stream from PyTorch to ensure proper synchronization.
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    
    launch_fused_add_rmsnorm_h7168(
        output,
        hidden_states,
        residual,
        weight,
        static_cast<float>(eps),
        stream
    );

    return output;
}

// --- Pybind11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Fused Add + RMSNorm (BFloat16) for hidden_size=7168",
          py::arg("hidden_states"),
          py::arg("residual"),
          py::arg("weight"),
          py::arg("eps") = 1e-6);
}