gpt-o3 / cuda64f9f5

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

/* -------------------------------------------------------------------------- */
/* Python-exposed host function                                               */
/* -------------------------------------------------------------------------- */
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 be BF16");
    TORCH_CHECK(weight.dtype() == torch::kBFloat16,
                "weight must be BF16");
    TORCH_CHECK(hidden_states.dim() == 2 && hidden_states.size(1) == HIDDEN_SIZE,
                "hidden_states shape must be [batch, 128]");
    TORCH_CHECK(weight.numel() == HIDDEN_SIZE,
                "weight size must be 128");

    const int64_t batch_size = hidden_states.size(0);

    /* Ensure contiguous tensors ------------------------------------------------*/
    auto h_in  = hidden_states.contiguous();
    auto w_in  = weight.contiguous();
    auto out   = torch::empty_like(h_in);

    /* Convert tensor pointers to BF16 device pointers ------------------------- */
    const __nv_bfloat16* h_ptr =
        reinterpret_cast<const __nv_bfloat16*>(h_in.data_ptr<torch::BFloat16>());
    const __nv_bfloat16* w_ptr =
        reinterpret_cast<const __nv_bfloat16*>(w_in.data_ptr<torch::BFloat16>());
    __nv_bfloat16* o_ptr =
        reinterpret_cast<__nv_bfloat16*>(out.data_ptr<torch::BFloat16>());

    /* Launch kernel ----------------------------------------------------------- */
    auto stream = at::cuda::getCurrentCUDAStream();
    rmsnorm_h128_launcher(h_ptr, w_ptr, o_ptr, static_cast<int>(batch_size),
                          stream.stream());

    return out;
}

/* -------------------------------------------------------------------------- */
/* PyBind11 binding                                                           */
/* -------------------------------------------------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run,
          "RMSNorm with hidden_size = 128 (BF16, CUDA/B200 optimised)");
}