gpt-o3 / cuda8eed1f

#include "kernel.h"

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

// ---------------------------------------------------------------------------
// Python-visible wrapper
// ---------------------------------------------------------------------------
torch::Tensor run(torch::Tensor hidden_states,
                  torch::Tensor weight)
{
    TORCH_CHECK(hidden_states.is_cuda(),
                "hidden_states must be a CUDA tensor.");
    TORCH_CHECK(weight.is_cuda(),
                "weight must be a CUDA tensor.");
    TORCH_CHECK(hidden_states.dtype() == torch::kBFloat16,
                "hidden_states must have dtype torch.bfloat16.");
    TORCH_CHECK(weight.dtype() == torch::kBFloat16,
                "weight must have dtype torch.bfloat16.");
    TORCH_CHECK(hidden_states.size(1) == 4096,
                "hidden_size must be 4096.");

    at::cuda::CUDAGuard device_guard(hidden_states.device());

    // Ensure contiguous memory for coalesced access
    auto x = hidden_states.contiguous();
    auto w = weight.contiguous();
    auto y = torch::empty_like(x);

    rmsnorm_h4096_cuda_forward(x, w, y);
    return y;
}

// ---------------------------------------------------------------------------
// PyBind11 binding
// ---------------------------------------------------------------------------
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("run",
          &run,
          pybind11::arg("hidden_states"),
          pybind11::arg("weight"),
          "Fast RMS-Norm forward (hidden_size = 4096, BF16, B200)");
}