gpt-o3 / cudaa7bbcf

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

#define CHECK_CUDA(x)         TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x)   TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_BFLOAT16(x)     TORCH_CHECK(x.dtype() == torch::kBFloat16, #x " must be bfloat16")

torch::Tensor run(torch::Tensor hidden_states,
                  torch::Tensor residual,
                  torch::Tensor weight)
{
    /* --------------------------- sanity checks --------------------------- */
    CHECK_CUDA(hidden_states);
    CHECK_CUDA(residual);
    CHECK_CUDA(weight);

    CHECK_CONTIGUOUS(hidden_states);
    CHECK_CONTIGUOUS(residual);
    CHECK_CONTIGUOUS(weight);

    CHECK_BFLOAT16(hidden_states);
    CHECK_BFLOAT16(residual);
    CHECK_BFLOAT16(weight);

    TORCH_CHECK(hidden_states.dim() == 2 &&
                hidden_states.size(1) == 4096,
                "hidden_states must be [batch, 4096]");
    TORCH_CHECK(residual.sizes() == hidden_states.sizes(),
                "residual must match hidden_states");
    TORCH_CHECK(weight.dim() == 1 && weight.size(0) == 4096,
                "weight must be [4096]");

    const int batch_size = hidden_states.size(0);
    auto output = torch::empty_like(hidden_states);

    /* Current stream from PyTorch */
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    fused_add_rmsnorm_h4096_launch(
        reinterpret_cast<const __nv_bfloat16*>(hidden_states.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(residual.data_ptr<at::BFloat16>()),
        reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr<at::BFloat16>()),
        reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()),
        batch_size,
        stream);

    return output;
}

/* ------------------------------- bindings ------------------------------- */
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run",
          &run,
          "Fused Add + RMSNorm with hidden_size=4096 (CUDA, B200 optimised)");
}