/*
 * FLUXPATCH_FP16.cpp — TORCH_LIBRARY registration for FLUXNATION (FP16-only)
 *
 * Registers ssb_prepare and ssb_finish as torch.ops.flux_ext.* custom ops.
 * Identical API to the FP8 version — Python code doesn't change.
 */

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

// Forward declarations (in FLUXNATION_FP16.cu)
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
ssb_prepare_impl(
    at::Tensor x, at::Tensor vec, at::Tensor pe,
    at::Tensor mod_w, at::Tensor mod_w_scale, at::Tensor mod_b,
    at::Tensor l1w, at::Tensor l1w_scale, at::Tensor l1b,
    at::Tensor qns, at::Tensor kns);

at::Tensor ssb_finish_impl(
    at::Tensor x, at::Tensor attn, at::Tensor mlp, at::Tensor gate,
    at::Tensor l2w, at::Tensor l2w_scale, at::Tensor l2b);

void ssb_cleanup_impl();

TORCH_LIBRARY(flux_ext, m) {
    m.def("ssb_prepare("
          "Tensor x, Tensor vec, Tensor pe, "
          "Tensor mod_w, Tensor mod_w_scale, Tensor mod_b, "
          "Tensor l1w, Tensor l1w_scale, Tensor l1b, "
          "Tensor qns, Tensor kns"
          ") -> (Tensor, Tensor, Tensor, Tensor, Tensor)");

    m.def("ssb_finish("
          "Tensor x, Tensor attn, Tensor mlp, Tensor gate, "
          "Tensor l2w, Tensor l2w_scale, Tensor l2b"
          ") -> Tensor");

    m.def("ssb_cleanup() -> ()");
}

TORCH_LIBRARY_IMPL(flux_ext, CUDA, m) {
    m.impl("ssb_prepare", &ssb_prepare_impl);
    m.impl("ssb_finish",  &ssb_finish_impl);
    m.impl("ssb_cleanup", &ssb_cleanup_impl);
}

PYBIND11_MODULE(flux_ext, m) {
    m.doc() = "FLUXNATION FP16: Fused CUDA kernels for FLUX.1 SingleStreamBlock (no FP8)";
}