gemini-2.5-pro / cuda3eed96

#include <torch/extension.h>
#include <pybind11/pybind11.h>
#include <cuda_bf16.h>
#include "kernel.h"
#include <limits> // For std::numeric_limits

// Forward declarations for PyTorch CUDA stream management
namespace at {
namespace cuda {
  cudaStream_t getCurrentCUDAStream();
}
}

namespace py = pybind11;

/**
 * @brief Python-bindable 'run' function for RMS Normalization.
 *
 * This function serves as the interface between Python (PyTorch) and the CUDA C++ implementation.
 * It performs extensive input validation before launching the optimized CUDA kernel.
 *
 * @param hidden_states The input tensor of shape [batch_size, 2048] and dtype bfloat16.
 * @param weight The weight tensor of shape [2048] and dtype bfloat16.
 * @return The output tensor of the same shape and dtype as hidden_states.
 */
torch::Tensor run(
    const torch::Tensor& hidden_states,
    const torch::Tensor& weight
) {
    // --- Input Validation ---
    TORCH_CHECK(hidden_states.device().is_cuda(), "hidden_states must be a CUDA tensor");
    TORCH_CHECK(weight.device().is_cuda(), "weight must be a CUDA tensor");

    TORCH_CHECK(hidden_states.scalar_type() == torch::kBFloat16, "hidden_states must be of type BFloat16");
    TORCH_CHECK(weight.scalar_type() == torch::kBFloat16, "weight must be of type BFloat16");

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

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

    TORCH_CHECK(batch_size_64 <= std::numeric_limits<int>::max(), "batch_size exceeds the maximum representable value for an int");
    const int batch_size = static_cast<int>(batch_size_64);

    TORCH_CHECK(hidden_size == 2048, "hidden_size must be 2048, but got ", hidden_size);
    TORCH_CHECK(weight.size(0) == hidden_size, "weight must have size matching hidden_size (2048)");

    TORCH_CHECK(hidden_states.is_contiguous(), "hidden_states must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");

    // --- Kernel Execution ---
    // Create an output tensor with the same properties as the input
    auto output = torch::empty_like(hidden_states);

    if (batch_size == 0) {
        return output;
    }

    // Get the current CUDA stream from PyTorch
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // PyTorch's at::BFloat16 is layout-compatible with CUDA's __nv_bfloat16.
    // We can safely reinterpret_cast the data pointers.
    cudaError_t err = rmsnorm_h2048_launch(
        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*>(weight.data_ptr<at::BFloat16>()),
        batch_size,
        stream
    );

    TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ", cudaGetErrorString(err));

    return output;
}

// PYBIND11_MODULE is a macro that creates an entry point that will be invoked when
// the Python interpreter imports a C++ extension.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run, "rmsnorm_h2048: RMS Normalization with hidden_size=2048 (CUDA, BFloat16)",
          py::arg("hidden_states"), py::arg("weight"));
}