gpt-o3 / cudac9eefe

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include "kernel.h"

/* --------------------------------------------------------------------- */
/* Basic argument checking                                               */
/* --------------------------------------------------------------------- */
static void check_inputs(const torch::Tensor& hidden,
                         const torch::Tensor& weight)
{
  using namespace at;

  TORCH_CHECK(hidden.is_cuda(),  "hidden_states must reside on CUDA");
  TORCH_CHECK(weight.is_cuda(),  "weight must reside on CUDA");

  TORCH_CHECK(hidden.scalar_type() == kBFloat16,
              "hidden_states must be torch.bfloat16");
  TORCH_CHECK(weight.scalar_type() == kBFloat16,
              "weight must be torch.bfloat16");

  TORCH_CHECK(hidden.dim() == 2 && hidden.size(1) == 1536,
              "hidden_states expected shape [batch_size, 1536]");
  TORCH_CHECK(weight.dim() == 1 && weight.size(0) == 1536,
              "weight expected shape [1536]");
}

/* --------------------------------------------------------------------- */
/* Python-visible entry point                                            */
/* --------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor hidden_states,
                  torch::Tensor weight)
{
  check_inputs(hidden_states, weight);

  /* Make sure memory is contiguous for coalesced accesses              */
  hidden_states = hidden_states.contiguous();
  weight        = weight.contiguous();

  auto output = torch::empty_like(hidden_states);

  const __nv_bfloat16* in_ptr =
      reinterpret_cast<const __nv_bfloat16*>(
          hidden_states.data_ptr<at::BFloat16>());
  const __nv_bfloat16* w_ptr  =
      reinterpret_cast<const __nv_bfloat16*>(
          weight.data_ptr<at::BFloat16>());
  __nv_bfloat16* out_ptr =
      reinterpret_cast<__nv_bfloat16*>(
          output.data_ptr<at::BFloat16>());

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  rmsnorm_h1536_launcher(in_ptr, w_ptr, out_ptr,
                         static_cast<int>(hidden_states.size(0)),
                         stream);

  /* Stream synchronisation is handled by PyTorch on tensor return.     */
  return output;
}

/* --------------------------------------------------------------------- */
/* PyBind11 binding                                                      */
/* --------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
  m.doc() = "Optimised RMSNorm (hidden_size = 1536) for NVIDIA GPUs";
  m.def("run",
        &run,
        pybind11::arg("hidden_states"),
        pybind11::arg("weight"),
        "rmsnorm(hidden_states, weight) → output");
}