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

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

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

    FP8_AdamW_expand_cuda(params, grads, \
                           exp_avg, scale_exp_avg, expand_exp_avg, sqrtminmax_exp_avg, \
                           exp_avg_sq, scale_exp_avg_sq, expand_exp_avg_sq, sqrtminmax_exp_avg_sq, \
                           beta1, beta2, lr, wd, eps, step, qgroup_size, expand_min);
}
