gemini-2.5-pro / cuda6c93f0

#include <torch/extension.h>
#include <cstdint>
#include "kernel.h"

// Helper to check tensor properties for robust error handling.
void check_tensor(const torch::Tensor& tensor, const std::string& name, torch::ScalarType dtype, int64_t dims, int64_t last_dim_size) {
    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());
    TORCH_CHECK(tensor.dim() == dims, name, " must be a ", dims, "D tensor, but got ", tensor.dim(), "D");
    if (last_dim_size != -1) {
       TORCH_CHECK(tensor.size(dims - 1) == last_dim_size, name, " last dimension must be ", last_dim_size, ", but got ", tensor.size(dims - 1));
    }
}


/**
 * @brief Python-callable function to run the RMSNorm H512 implementation.
 *
 * This function acts as the bridge between PyTorch and the custom CUDA kernel.
 * It performs extensive input validation before launching the kernel.
 *
 * @param hidden_states The input tensor of shape [batch_size, 512] and dtype bfloat16.
 * @param weight The scaling weight tensor of shape [512] and dtype bfloat16.
 * @return The output tensor of the same shape and dtype as `hidden_states`.
 */
torch::Tensor rmsnorm_h512_run(
    torch::Tensor hidden_states,
    torch::Tensor weight) {

    // --- Input Validation ---
    const auto BFLOAT16 = torch::kBFloat16;
    constexpr int32_t HIDDEN_SIZE = 512;

    check_tensor(hidden_states, "hidden_states", BFLOAT16, 2, HIDDEN_SIZE);
    check_tensor(weight, "weight", BFLOAT16, 1, HIDDEN_SIZE);

    const int batch_size = hidden_states.size(0);
    if (batch_size == 0) {
      return torch::empty_like(hidden_states);
    }

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

    // --- Kernel Launch ---
    // Get the current CUDA stream from PyTorch's dispatcher to ensure proper synchronization.
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    rmsnorm_h512_launcher(
        batch_size,
        hidden_states.data_ptr(),
        weight.data_ptr(),
        output.data_ptr(),
        stream);

    return output;
}

// --- PYBIND11 Module Definition ---
// Exposes the `run` function to Python so it can be called from the benchmark framework.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def(
      "run",
      &rmsnorm_h512_run,
      "RMSNorm H512 forward pass (CUDA)",
      pybind11::arg("hidden_states"),
      pybind11::arg("weight")
  );
}