#include <torch/extension.h>
#include <torch/torch.h>

void FP8_AdamW_cuda(
    torch::Tensor params, // parameter
    torch::Tensor grads, // gradient
    torch::Tensor exp_avg, // first order momentum
    torch::Tensor scale_exp_avg,
    torch::Tensor exp_avg_sq, // second order momentum
    torch::Tensor scale_exp_avg_sq,
    float beta1, float beta2, float lr, float wd, float eps, int step, int qgroup_size // other parameters
);

void FP8_AdamW(
    torch::Tensor params, // parameter
    torch::Tensor grads, // gradient
    torch::Tensor exp_avg, // first order momentum
    torch::Tensor scale_exp_avg,
    torch::Tensor exp_avg_sq, // second order momentum
    torch::Tensor scale_exp_avg_sq,
    float beta1, float beta2, float lr, float wd, float eps, int step, int qgroup_size) { // other parameters

    FP8_AdamW_cuda(params, grads, exp_avg, scale_exp_avg, exp_avg_sq, scale_exp_avg_sq, \
                   beta1, beta2, lr, wd, eps, step, qgroup_size);
}
