ULT7RAandCopilot
82bdde669f
Fused FP8 CUDA kernel for FLUX.1 SingleStreamBlock with spike attention and step caching. Replaces ComfyUI's entire SingleStreamBlock forward pass with a single torch.ops call — modulation, dual GEMMs, QKV, RoPE, attention, gating and residual all fused. 30%+ faster than stock ComfyUI on RTX 4090. Spike attention: block-sparse Triton kernel that scores attention blocks by dot-product similarity and keeps only the top 45% after image structure forms. Step caching replays cached attention output on alternating spike steps for zero attention compute every other step. Works with all 38 ComfyUI samplers at any step count. FP16 ports included for 10-series, 20-series, and 30-series GPUs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
59 lines
2.2 KiB
C++
59 lines
2.2 KiB
C++
/*
|
|
* flux_ext.cpp — TORCH_LIBRARY registration for FLUXNATION
|
|
*
|
|
* Registers ssb_prepare and ssb_finish as torch.ops.flux_ext.* custom ops.
|
|
* This eliminates ctypes overhead and enables torch.compile compatibility.
|
|
*/
|
|
|
|
#include <torch/extension.h>
|
|
#include <ATen/ATen.h>
|
|
|
|
// Forward declarations of CUDA implementations (in flux_ext_kernels.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();
|
|
|
|
// =========================================================================
|
|
// Op schema definitions
|
|
// =========================================================================
|
|
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() -> ()");
|
|
}
|
|
|
|
// =========================================================================
|
|
// CUDA backend implementations
|
|
// =========================================================================
|
|
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);
|
|
}
|
|
|
|
// =========================================================================
|
|
// Python module init (required for import flux_ext to work)
|
|
// =========================================================================
|
|
PYBIND11_MODULE(flux_ext, m) {
|
|
m.doc() = "FLUXNATION: Fused CUDA kernels for FLUX.1 SingleStreamBlock";
|
|
}
|