gpt-o3 / cuda73b50f

/*
 *  PyTorch binding for the specialised RMS-Norm kernel (hidden_size = 7168).
 *  The Python-visible entry point is `run(hidden_states, weight)`.
 */

#include "kernel.h"

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

/* -------------------------------------------------------------------------- */
/* Python-visible API                                                         */
/* -------------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor hidden_states,
                  torch::Tensor weight)
{
    TORCH_CHECK(hidden_states.is_cuda(), "hidden_states must reside on GPU");
    TORCH_CHECK(weight.is_cuda(),        "weight must reside on GPU");
    TORCH_CHECK(hidden_states.device() == weight.device(),
                "hidden_states and weight must be on the same GPU");

    /* Ensure execution on the correct device */
    at::cuda::CUDAGuard device_guard(hidden_states.device());

    /* The kernel expects contiguous memory */
    hidden_states = hidden_states.contiguous();
    weight        = weight.contiguous();

    /* Prepare output tensor */
    torch::Tensor output = torch::empty_like(hidden_states);

    /* Launch the CUDA kernel */
    launch_rmsnorm_h7168(hidden_states, weight, output);

    return output;
}

/* -------------------------------------------------------------------------- */
/* PyBind11 registration                                                      */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.doc() = "Optimised RMS-Norm (hidden_size = 7168, BF16)";

    m.def("run",
          &run,
          "Execute RMS-Norm on BF16 tensors",
          pybind11::arg("hidden_states"),
          pybind11::arg("weight"));
}