gpt-o3 / cudaa1fa19

#include "kernel.h"

#include <torch/extension.h>
#include <pybind11/pybind11.h>

/*
 * Python entry point:
 *
 *   output = run(hidden_states : bf16[batch,512],
 *                weight        : bf16[512],
 *                **kwargs)                       # kwargs ignored
 */
torch::Tensor run(torch::Tensor hidden_states,
                  torch::Tensor weight,
                  pybind11::kwargs /*kwargs*/ = {})
{
    TORCH_CHECK(hidden_states.device().is_cuda(),
                "hidden_states must be on a CUDA device.");
    TORCH_CHECK(weight.device().is_cuda(),
                "weight must be on a CUDA device.");

    /* Ensure contiguous layout (becomes a no-op if already contiguous) */
    auto hidden_c = hidden_states.contiguous();
    auto weight_c = weight.contiguous();

    /* Allocate output tensor */
    auto output = torch::empty_like(hidden_c);

    /* Launch CUDA kernel */
    rmsnorm_h512_cuda(hidden_c, weight_c, output);

    return output;
}

/* ------------------------------ pybind11 ---------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("run",
          &run,
          pybind11::arg("hidden_states"),
          pybind11::arg("weight"),
          pybind11::kw_only(),
          "Optimised RMSNorm kernel (hidden_size = 512, bf16, B200)");
}