gpt-5 / cuda1ad912

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

// Typed implementation
static torch::Tensor run_impl(torch::Tensor hidden_states, torch::Tensor weight) {
    return rmsnorm_h128::rmsnorm(hidden_states, weight);
}

// Python-facing wrapper that handles args and kwargs
static torch::Tensor run_py(pybind11::args args, pybind11::kwargs kwargs) {
    torch::Tensor hidden_states;
    torch::Tensor weight;

    if (args.size() >= 2) {
        hidden_states = args[0].cast<torch::Tensor>();
        weight = args[1].cast<torch::Tensor>();
    } else {
        if (kwargs.contains("hidden_states")) {
            hidden_states = kwargs["hidden_states"].cast<torch::Tensor>();
        } else {
            throw std::invalid_argument("Missing required argument: hidden_states");
        }
        if (kwargs.contains("weight")) {
            weight = kwargs["weight"].cast<torch::Tensor>();
        } else {
            throw std::invalid_argument("Missing required argument: weight");
        }
    }

    return run_impl(hidden_states, weight);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run_py, "RMSNorm H=128 (BF16) optimized CUDA kernel");
}