gpt-o3 / cuda3def09

#include "kernel.h"

#include <torch/extension.h>

/*
 * Python-facing entry point.
 * It validates inputs, allocates the output tensor and
 * dispatches to the CUDA launcher.
 */
torch::Tensor run(torch::Tensor hidden_states,
                  torch::Tensor residual,
                  torch::Tensor weight)
{
    // Ensures the tensors live on the same device / stream
    TORCH_CHECK(hidden_states.is_cuda(), "hidden_states must be CUDA");
    TORCH_CHECK(residual.is_cuda()     , "residual must be CUDA");
    TORCH_CHECK(weight.is_cuda()       , "weight must be CUDA");

    TORCH_CHECK(hidden_states.dtype() == at::kBFloat16,
                "hidden_states must be BF16");
    TORCH_CHECK(residual.dtype() == at::kBFloat16,
                "residual must be BF16");
    TORCH_CHECK(weight.dtype() == at::kBFloat16,
                "weight must be BF16");

    // Output – same shape/dtype/device as hidden_states
    auto output = torch::empty_like(hidden_states);

    // Launch CUDA kernel
    fused_add_rmsnorm_h2048(hidden_states, residual, weight, output);

    return output;
}

/* ------------------------------------------------------------------ */
/*  PyBind11 registration                                             */
/* ------------------------------------------------------------------ */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run,
          "Fused Add + RMSNorm (hidden_size = 2048, BF16, B200-optimised)");
}