gpt-o3 / cuda620cb5

#include "kernel.h"

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

/* quick helpers ------------------------------------------------------------- */
#define CHECK_CUDA(x)      TORCH_CHECK(x.is_cuda(),        #x " must be a CUDA tensor")
#define CHECK_CONTIG(x)    TORCH_CHECK(x.is_contiguous(),  #x " must be contiguous")
#define CHECK_DTYPE(x, t)  TORCH_CHECK(x.scalar_type() == t, #x " must have dtype " #t)

/* -------------------------------------------------------------------------- */
/* Thin C++ wrapper that validates inputs, allocates output and launches CUDA */
/* -------------------------------------------------------------------------- */
torch::Tensor fused_add_rmsnorm_h7168(torch::Tensor hidden_states,
                                      torch::Tensor residual,
                                      torch::Tensor weight)
{
    /* safety checks --------------------------------------------------------- */
    CHECK_CUDA(hidden_states);
    CHECK_CUDA(residual);
    CHECK_CUDA(weight);

    CHECK_CONTIG(hidden_states);
    CHECK_CONTIG(residual);
    CHECK_CONTIG(weight);

    CHECK_DTYPE(hidden_states, at::kBFloat16);
    CHECK_DTYPE(residual,      at::kBFloat16);
    CHECK_DTYPE(weight,        at::kBFloat16);

    TORCH_CHECK(hidden_states.dim() == 2 &&
                hidden_states.size(1) == HIDDEN_SIZE,
                "hidden_states must be [N, 7168]");
    TORCH_CHECK(residual.sizes() == hidden_states.sizes(),
                "residual must have the same shape as hidden_states");
    TORCH_CHECK(weight.numel() == HIDDEN_SIZE,
                "weight must contain 7168 elements");

    /* allocate output ------------------------------------------------------- */
    torch::Tensor output = torch::empty_like(hidden_states);

    /* launch kernel --------------------------------------------------------- */
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    launch_fused_add_rmsnorm_h7168(hidden_states,
                                   residual,
                                   weight,
                                   output,
                                   stream);

    return output;
}

/* -------------------------------------------------------------------------- */
/* Entry point expected by the benchmark harness                              */
/* -------------------------------------------------------------------------- */
torch::Tensor run(torch::Tensor hidden_states,
                  torch::Tensor residual,
                  torch::Tensor weight)
{
    return fused_add_rmsnorm_h7168(std::move(hidden_states),
                                   std::move(residual),
                                   std::move(weight));
}

/* -------------------------------------------------------------------------- */
/* pybind11 binding                                                           */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("run", &run,
          "Fused Add + RMSNorm (hidden=7168, BF16, optimised for Blackwell)");
}