gemini-2.5-pro / cuda4597e0

#include "kernel.h"
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h> // Required for at::cuda::getCurrentCUDAStream

/**
 * @brief Python-bindable 'run' function for the fused add+rmsnorm operation.
 *
 * This function serves as the C++ entry point, callable from Python. It handles
 * tensor validation, memory management, and CUDA kernel launching.
 *
 * @param hidden_states The main input tensor of shape [batch_size, 2048].
 * @param residual The tensor to be added to hidden_states, shape [batch_size, 2048].
 * @param weight The scaling weights for RMSNorm, shape [2048].
 * @return A new tensor containing the result of the operation.
 */
torch::Tensor run(
    const torch::Tensor& hidden_states,
    const torch::Tensor& residual,
    const torch::Tensor& weight) {

    // --- Input Tensor Validation ---
    static constexpr auto HIDDEN_SIZE = 2048;

    // Device checks
    TORCH_CHECK(hidden_states.device().is_cuda(), "hidden_states must be a CUDA tensor");
    TORCH_CHECK(residual.device().is_cuda(), "residual must be a CUDA tensor");
    TORCH_CHECK(weight.device().is_cuda(), "weight must be a CUDA tensor");

    TORCH_CHECK(hidden_states.device() == residual.device() && hidden_states.device() == weight.device(),
                "All tensors must be on the same CUDA device");

    // Dtype checks
    TORCH_CHECK(hidden_states.scalar_type() == torch::kBFloat16, "hidden_states must be BFloat16");
    TORCH_CHECK(residual.scalar_type() == torch::kBFloat16, "residual must be BFloat16");
    TORCH_CHECK(weight.scalar_type() == torch::kBFloat16, "weight must be BFloat16");

    // Shape checks
    TORCH_CHECK(hidden_states.dim() == 2, "hidden_states must be 2-dimensional");
    TORCH_CHECK(residual.dim() == 2, "residual must be 2-dimensional");
    TORCH_CHECK(weight.dim() == 1, "weight must be 1-dimensional");

    const int64_t batch_size = hidden_states.size(0);
    TORCH_CHECK(hidden_states.size(1) == HIDDEN_SIZE, "hidden_states hidden size must be ", HIDDEN_SIZE);
    TORCH_CHECK(residual.size(0) == batch_size, "residual batch size must match hidden_states");
    TORCH_CHECK(residual.size(1) == HIDDEN_SIZE, "residual hidden size must be ", HIDDEN_SIZE);
    TORCH_CHECK(weight.size(0) == HIDDEN_SIZE, "weight size must be ", HIDDEN_SIZE);

    // Contiguity checks for safe pointer access
    TORCH_CHECK(hidden_states.is_contiguous(), "hidden_states must be contiguous");
    TORCH_CHECK(residual.is_contiguous(), "residual must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");

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

    // --- Kernel Execution ---
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Update the weights in constant memory for this run
    fused_add_rmsnorm_h2048_update_weights(
        weight.data_ptr<torch::BFloat16>(),
        stream
    );

    fused_add_rmsnorm_h2048_launch(
        output.data_ptr<torch::BFloat16>(),
        hidden_states.data_ptr<torch::BFloat16>(),
        residual.data_ptr<torch::BFloat16>(),
        batch_size,
        stream
    );

    return output;
}


// --- PYBIND11 Module Definition ---
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "Optimized CUDA kernel for Fused Add + RMSNorm (h=2048, bfloat16).");
}