gemini-2.5-pro / cuda2c7e9e

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include "kernel.h"

// Define helper macros for concise input tensor validation
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_BF16(x) TORCH_CHECK(x.scalar_type() == torch::kBFloat16, #x " must be a bfloat16 tensor")

// Define constant for the fixed hidden size
constexpr int HIDDEN_SIZE_CONST = 128;

/**
 * @brief C++ and Pybind11 interface for the RMSNorm CUDA kernel.
 *
 * This function acts as the bridge between Python (PyTorch) and the CUDA C++
 * implementation. It handles tensor validation, memory management, and CUDA
 * kernel launching.
 *
 * @param hidden_states The input tensor of shape [batch_size, 128] and dtype bfloat16.
 * @param weight The weight tensor of shape [128] and dtype bfloat16.
 * @return The output tensor of the same shape and dtype as hidden_states.
 */
torch::Tensor run(
    torch::Tensor hidden_states,
    torch::Tensor weight) {

    // --- Input Validation ---
    // Ensure all tensors are on the GPU
    CHECK_CUDA(hidden_states);
    CHECK_CUDA(weight);

    // Ensure all tensors are contiguous in memory for direct pointer access
    CHECK_CONTIGUOUS(hidden_states);
    CHECK_CONTIGUOUS(weight);

    // Check for the correct data type (bfloat16)
    CHECK_BF16(hidden_states);
    CHECK_BF16(weight);

    // Check tensor dimensions
    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 = hidden_states.size(0);
    const int64_t hidden_size = hidden_states.size(1);

    // Check fixed dimension sizes
    TORCH_CHECK(hidden_size == HIDDEN_SIZE_CONST, "hidden_size must be 128");
    TORCH_CHECK(weight.size(0) == HIDDEN_SIZE_CONST, "weight must have size 128");

    // --- Output Tensor Preparation ---
    // Create an output tensor with the same properties (shape, dtype, device) as the input
    auto output = torch::empty_like(hidden_states);

    // --- Kernel Launch ---
    // Get the current CUDA stream from PyTorch's context
    c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();

    rmsnorm_h128_launch(
        output.data_ptr(),
        hidden_states.data_ptr(),
        weight.data_ptr(),
        static_cast<int>(batch_size),
        stream
    );

    return output;
}

// --- Pybind11 Module Definition ---
// This is the entry point that exposes the C++ 'run' function to Python.
// The macro TORCH_EXTENSION_NAME is defined by the build system (e.g., setuptools).
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def(
        "run",                  // Python function name
        &run,                   // C++ function pointer
        "RMSNorm H=128 kernel (CUDA BFloat16)" // Docstring for the Python function
    );
}