FLUXNATION: Fused spike attention + step cache CUDA kernel for FLUX.1
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>
This commit is contained in:
ULT7RA
+100
@@ -0,0 +1,100 @@
|
||||
from __future__ import annotations
|
||||
from typing import Dict, Optional, Tuple, Any
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from . import FLUXSPIKEY as triton_backend
|
||||
|
||||
|
||||
def attend(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
backend: str = "dense",
|
||||
state: Optional[Dict[str, torch.Tensor]] = None,
|
||||
neuromorph: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[torch.Tensor, Dict[str, Any]]:
|
||||
if backend == "dense":
|
||||
return _attend_dense(q, k, v)
|
||||
elif backend == "spike":
|
||||
neuromorph_cfg = neuromorph or {}
|
||||
if triton_backend.has_triton():
|
||||
return triton_backend.attend_spike_triton(
|
||||
q, k, v,
|
||||
tau=neuromorph_cfg.get("tau", 0.15),
|
||||
leak=neuromorph_cfg.get("leak", 0.0),
|
||||
event_cap_pct=neuromorph_cfg.get("event_cap_pct", 0.45),
|
||||
state=state,
|
||||
)
|
||||
return _attend_spike(q, k, v, state=state, neuromorph=neuromorph_cfg)
|
||||
else:
|
||||
raise ValueError(f"Unknown backend: {backend}")
|
||||
|
||||
|
||||
def _attend_dense(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, Dict[str, Any]]:
|
||||
scale = q.shape[-1] ** -0.5
|
||||
scores = torch.matmul(q, k.transpose(-2, -1)) * scale
|
||||
attn_weights = F.softmax(scores, dim=-1)
|
||||
output = torch.matmul(attn_weights, v)
|
||||
return output, {"backend": "dense", "attn_weights": attn_weights}
|
||||
|
||||
|
||||
def _attend_spike(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
state: Optional[Dict[str, torch.Tensor]] = None,
|
||||
neuromorph: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[torch.Tensor, Dict[str, Any]]:
|
||||
cfg = neuromorph or {}
|
||||
tau = cfg.get("tau", 0.15)
|
||||
leak = cfg.get("leak", 0.0)
|
||||
event_cap_pct = cfg.get("event_cap_pct", 0.45)
|
||||
|
||||
batch, heads, seq_len, head_dim = q.shape
|
||||
|
||||
scale = head_dim ** -0.5
|
||||
z = torch.matmul(q, k.transpose(-2, -1)) * scale
|
||||
|
||||
z_max_per_row = z.max(dim=-1, keepdim=True)[0]
|
||||
M_t = (z >= tau * z_max_per_row).float()
|
||||
|
||||
max_events_per_row = max(1, int(seq_len * event_cap_pct))
|
||||
if max_events_per_row < seq_len:
|
||||
_, topk_indices = torch.topk(z, k=max_events_per_row, dim=-1, sorted=False)
|
||||
M_t_capped = torch.zeros_like(z)
|
||||
M_t_capped.scatter_(-1, topk_indices, 1.0)
|
||||
M_t = M_t * M_t_capped
|
||||
|
||||
if state is not None and "S_prev" in state and state["S_prev"].shape == M_t.shape:
|
||||
S_t = leak * state["S_prev"] + M_t
|
||||
else:
|
||||
S_t = M_t
|
||||
|
||||
S_t_sum = S_t.sum(dim=-1, keepdim=True)
|
||||
S_t_sum = torch.where(S_t_sum > 0, S_t_sum, torch.ones_like(S_t_sum))
|
||||
attn_weights = (S_t / S_t_sum * M_t).to(v.dtype)
|
||||
output = torch.matmul(attn_weights, v)
|
||||
|
||||
kept = M_t.sum().item()
|
||||
total = batch * heads * seq_len * seq_len
|
||||
flops_saved = 100.0 * (1.0 - kept / max(1, total))
|
||||
|
||||
return output, {
|
||||
"S_prev": S_t.detach(),
|
||||
"M_t": M_t,
|
||||
"telemetry": {
|
||||
"backend": "spike",
|
||||
"kept_events_pct": 100.0 * kept / max(1, total),
|
||||
"flops_saved_pct": flops_saved,
|
||||
"tau": tau,
|
||||
"leak": leak,
|
||||
"event_cap_pct": event_cap_pct,
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
#ifndef SINGLE_STREAM_BLOCK_FP16_H
|
||||
#define SINGLE_STREAM_BLOCK_FP16_H
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cublas_v2.h>
|
||||
#include <cstdio>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
namespace transformer {
|
||||
|
||||
constexpr int HIDDEN = 3072;
|
||||
constexpr int HEADS = 24;
|
||||
constexpr int HDIM = 128;
|
||||
constexpr int FFDIM = 12288;
|
||||
constexpr int MODDIM = 3 * HIDDEN;
|
||||
constexpr int LIN1DIM = 3 * HIDDEN + FFDIM;
|
||||
constexpr int LIN2DIM = HIDDEN + FFDIM;
|
||||
constexpr float LN_EPS = 1e-6f;
|
||||
constexpr float RMS_EPS = 1e-6f;
|
||||
constexpr float ATT_SC = 0.0883883476f;
|
||||
|
||||
static int const kKernelThreads = 256;
|
||||
static_assert((kKernelThreads % 32) == 0, "kKernelThreads must be a multiple of 32.");
|
||||
|
||||
#define CK(call) do { \
|
||||
cudaError_t err__ = (call); \
|
||||
if (err__ != cudaSuccess) { \
|
||||
throw std::runtime_error(std::string("CUDA ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + cudaGetErrorString(err__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define BK(call) do { \
|
||||
cublasStatus_t st__ = (call); \
|
||||
if (st__ != CUBLAS_STATUS_SUCCESS) { \
|
||||
throw std::runtime_error(std::string("cuBLAS ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + std::to_string((int)st__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static inline size_t ssb_workspace_bytes(int batch, int seq) {
|
||||
size_t rows = static_cast<size_t>(batch) * seq;
|
||||
size_t b = static_cast<size_t>(batch);
|
||||
size_t sz = 0;
|
||||
#define ADDBUF(count, elem_size) \
|
||||
sz = (sz + 255) & ~size_t(255); sz += (count) * (elem_size);
|
||||
ADDBUF(b * HIDDEN, 2); // d_silu_
|
||||
ADDBUF(b * MODDIM, 2); // d_mod_
|
||||
ADDBUF(rows * HIDDEN, 2); // d_xmod_
|
||||
ADDBUF(rows * LIN1DIM, 2); // d_l1o_
|
||||
sz = (sz + 255) & ~size_t(255);
|
||||
sz += 32ull << 20; // 32MB cuBLAS workspace
|
||||
#undef ADDBUF
|
||||
return sz;
|
||||
}
|
||||
|
||||
class SSBEngine {
|
||||
public:
|
||||
SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size);
|
||||
~SSBEngine();
|
||||
|
||||
void prepare(const half* d_x,
|
||||
const half* d_vec,
|
||||
const half* d_pe,
|
||||
const half* mod_w,
|
||||
const half* mod_b,
|
||||
const half* l1w,
|
||||
const half* l1b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
half* d_q,
|
||||
half* d_k,
|
||||
half* d_v,
|
||||
half* d_mlp,
|
||||
half* d_gate,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
void finish(const half* d_x,
|
||||
const half* d_attn,
|
||||
half* d_mlp,
|
||||
const half* d_gate,
|
||||
half* d_out,
|
||||
const half* l2w,
|
||||
const half* l2b,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
int batch() const { return B_; }
|
||||
int seq() const { return S_; }
|
||||
|
||||
private:
|
||||
int B_;
|
||||
int S_;
|
||||
cublasHandle_t cub_;
|
||||
void* workspace_;
|
||||
size_t workspace_bytes_;
|
||||
|
||||
half* d_silu_;
|
||||
half* d_mod_;
|
||||
half* d_xmod_;
|
||||
half* d_l1o_;
|
||||
};
|
||||
|
||||
} // namespace transformer
|
||||
|
||||
#endif // SINGLE_STREAM_BLOCK_FP16_H
|
||||
@@ -0,0 +1,25 @@
|
||||
@echo off
|
||||
setlocal
|
||||
|
||||
REM FLUXNATION FP16 Build — 10-series (SM 61, Pascal)
|
||||
REM Compatible with GTX 1060/1070/1080/1080 Ti
|
||||
|
||||
set "VSVARS=C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat"
|
||||
if not exist "%VSVARS%" (
|
||||
echo [flux_ext] ERROR: VS 2022 vcvars64.bat not found
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
call "%VSVARS%" || exit /b 1
|
||||
|
||||
set DISTUTILS_USE_SDK=1
|
||||
|
||||
echo [flux_ext] Building FLUXNATION FP16 for 10-series (SM 61)...
|
||||
py -3.10 setup_ext.py install
|
||||
if errorlevel 1 (
|
||||
echo [flux_ext] Build FAILED
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
echo [flux_ext] Build succeeded! flux_ext (FP16, SM 61) is now installed.
|
||||
echo [flux_ext] Restart ComfyUI to use the new kernel.
|
||||
@@ -0,0 +1,198 @@
|
||||
/*
|
||||
* FLUXNATION_FP16.cu — PyTorch C++ extension for FLUXNATION (FP16-only)
|
||||
*
|
||||
* Wraps the FP16 SSBEngine as torch.ops custom operators.
|
||||
* All weights converted to FP16 internally — no FP8, no CUTLASS.
|
||||
* Compatible with SM 61+ (Pascal, Turing, Ampere).
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#include "SINGLE_STREAM_BLOCK_FP16.H"
|
||||
|
||||
#include <unordered_map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
|
||||
using namespace transformer;
|
||||
|
||||
// =========================================================================
|
||||
// Global engine cache — keyed by (device, batch, seq_len)
|
||||
// =========================================================================
|
||||
struct EngineEntry {
|
||||
std::unique_ptr<SSBEngine> engine;
|
||||
at::Tensor workspace;
|
||||
};
|
||||
|
||||
static std::unordered_map<int64_t, EngineEntry> g_engine_cache;
|
||||
static std::mutex g_engine_mutex;
|
||||
|
||||
static int64_t make_engine_key(int device, int batch, int seq) {
|
||||
return (int64_t)device * 1000000000LL + (int64_t)batch * 1000000LL + seq;
|
||||
}
|
||||
|
||||
static SSBEngine* get_or_create_engine(int batch, int seq, at::Device device) {
|
||||
int dev_idx = device.index();
|
||||
int64_t key = make_engine_key(dev_idx, batch, seq);
|
||||
|
||||
std::lock_guard<std::mutex> lock(g_engine_mutex);
|
||||
auto it = g_engine_cache.find(key);
|
||||
if (it != g_engine_cache.end()) {
|
||||
return it->second.engine.get();
|
||||
}
|
||||
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
size_t ws_bytes = ssb_workspace_bytes(batch, seq);
|
||||
auto ws_tensor = at::empty({(int64_t)ws_bytes}, at::TensorOptions().dtype(at::kByte).device(device));
|
||||
|
||||
auto engine = std::make_unique<SSBEngine>(batch, seq, ws_tensor.data_ptr(), ws_bytes);
|
||||
auto* ptr = engine.get();
|
||||
|
||||
EngineEntry entry;
|
||||
entry.engine = std::move(engine);
|
||||
entry.workspace = ws_tensor;
|
||||
g_engine_cache[key] = std::move(entry);
|
||||
|
||||
fprintf(stderr, "[flux_ext FP16] Created engine B=%d S=%d ws=%.1fMB dev=%d\n",
|
||||
batch, seq, ws_bytes / (1024.0 * 1024.0), dev_idx);
|
||||
fflush(stderr);
|
||||
|
||||
return ptr;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// torch.ops implementation: ssb_prepare
|
||||
//
|
||||
// Same Python API as FP8 version — scale tensors are accepted but ignored.
|
||||
// All weights converted to FP16 internally.
|
||||
// =========================================================================
|
||||
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
|
||||
ssb_prepare_impl(
|
||||
at::Tensor x, // [B, S, 3072] bf16
|
||||
at::Tensor vec, // [B, 3072] or [B, 1, 3072] fp16
|
||||
at::Tensor pe, // [S, 64, 2] fp16
|
||||
at::Tensor mod_w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor mod_w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor mod_b, // fp16
|
||||
at::Tensor l1w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor l1w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor l1b, // fp16
|
||||
at::Tensor qns, // fp16
|
||||
at::Tensor kns) // fp16
|
||||
{
|
||||
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
|
||||
TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");
|
||||
|
||||
int B = x.size(0);
|
||||
int S = x.size(1);
|
||||
auto device = x.device();
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(device.index());
|
||||
|
||||
// Ensure correct dtypes — x stays BF16 (layernorm kernel reads it as BF16)
|
||||
auto x_bf16 = x.to(at::kBFloat16).contiguous();
|
||||
auto vec_fp16 = vec.dim() == 3 ? vec.squeeze(1).to(at::kHalf).contiguous()
|
||||
: vec.to(at::kHalf).contiguous();
|
||||
auto pe_fp16 = pe.to(at::kHalf).contiguous();
|
||||
|
||||
// Convert ALL weights to FP16 (no FP8 on this target)
|
||||
auto mod_w_fp16 = mod_w.to(at::kHalf).contiguous();
|
||||
auto mod_b_fp16 = mod_b.to(at::kHalf).contiguous();
|
||||
auto l1w_fp16 = l1w.to(at::kHalf).contiguous();
|
||||
auto l1b_fp16 = l1b.to(at::kHalf).contiguous();
|
||||
auto qns_fp16 = qns.to(at::kHalf).contiguous();
|
||||
auto kns_fp16 = kns.to(at::kHalf).contiguous();
|
||||
|
||||
// Allocate output tensors
|
||||
auto opts_fp16 = at::TensorOptions().dtype(at::kHalf).device(device);
|
||||
auto q = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto k = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto v = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto mlp = at::empty({B, S, FFDIM}, opts_fp16);
|
||||
auto gate = at::empty({B, HIDDEN}, opts_fp16);
|
||||
|
||||
SSBEngine* engine = get_or_create_engine(B, S, device);
|
||||
|
||||
engine->prepare(
|
||||
reinterpret_cast<const half*>(x_bf16.data_ptr()),
|
||||
reinterpret_cast<const half*>(vec_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(pe_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(mod_w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(mod_b_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l1w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l1b_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(qns_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(kns_fp16.data_ptr()),
|
||||
reinterpret_cast<half*>(q.data_ptr()),
|
||||
reinterpret_cast<half*>(k.data_ptr()),
|
||||
reinterpret_cast<half*>(v.data_ptr()),
|
||||
reinterpret_cast<half*>(mlp.data_ptr()),
|
||||
reinterpret_cast<half*>(gate.data_ptr()),
|
||||
stream);
|
||||
|
||||
return std::make_tuple(q, k, v, mlp, gate);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// torch.ops implementation: ssb_finish
|
||||
//
|
||||
// FP16 path: converts residual (x) from BF16→FP16, runs all-FP16 engine,
|
||||
// converts result back to BF16 for model compatibility.
|
||||
// =========================================================================
|
||||
at::Tensor ssb_finish_impl(
|
||||
at::Tensor x, // [B, S, 3072] bf16 — original input for residual
|
||||
at::Tensor attn, // [B, S, 3072] fp16 — attention output
|
||||
at::Tensor mlp, // [B, S, 12288] fp16 — MLP intermediate (pre-GELU)
|
||||
at::Tensor gate, // [B, 3072] fp16 — gate values
|
||||
at::Tensor l2w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor l2w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor l2b) // fp16 — L2 bias
|
||||
{
|
||||
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
|
||||
TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");
|
||||
|
||||
int B = x.size(0);
|
||||
int S = x.size(1);
|
||||
auto device = x.device();
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(device.index());
|
||||
|
||||
// Convert residual from BF16 to FP16 for the FP16 engine
|
||||
auto x_fp16 = x.to(at::kHalf).contiguous();
|
||||
auto attn_fp16 = attn.to(at::kHalf).contiguous();
|
||||
auto mlp_c = mlp.contiguous();
|
||||
auto gate_c = gate.contiguous();
|
||||
auto l2w_fp16 = l2w.to(at::kHalf).contiguous();
|
||||
auto l2b_fp16 = l2b.to(at::kHalf).contiguous();
|
||||
|
||||
// Result in FP16 (engine is all-FP16)
|
||||
auto opts_fp16 = at::TensorOptions().dtype(at::kHalf).device(device);
|
||||
auto result_fp16 = at::empty({B, S, HIDDEN}, opts_fp16);
|
||||
|
||||
SSBEngine* engine = get_or_create_engine(B, S, device);
|
||||
|
||||
engine->finish(
|
||||
reinterpret_cast<const half*>(x_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(attn_fp16.data_ptr()),
|
||||
reinterpret_cast<half*>(mlp_c.data_ptr()),
|
||||
reinterpret_cast<const half*>(gate_c.data_ptr()),
|
||||
reinterpret_cast<half*>(result_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l2w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l2b_fp16.data_ptr()),
|
||||
stream);
|
||||
|
||||
// Convert back to BF16 for model compatibility
|
||||
return result_fp16.to(at::kBFloat16);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Cleanup — destroys all cached engines
|
||||
// =========================================================================
|
||||
void ssb_cleanup_impl() {
|
||||
std::lock_guard<std::mutex> lock(g_engine_mutex);
|
||||
g_engine_cache.clear();
|
||||
fprintf(stderr, "[flux_ext FP16] All engines destroyed\n");
|
||||
fflush(stderr);
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
/*
|
||||
* 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)";
|
||||
}
|
||||
@@ -0,0 +1,611 @@
|
||||
/*
|
||||
* flux_opt_fp16.cu — FP16-only SSBEngine for FLUXNATION
|
||||
*
|
||||
* Derived from flux_opt.cu. All FP8/CUTLASS code removed.
|
||||
* Uses cuBLAS FP16 tensor-op GEMMs only.
|
||||
* Compatible with SM 61+ (Pascal, Turing, Ampere).
|
||||
*/
|
||||
|
||||
#include "SINGLE_STREAM_BLOCK_FP16.H"
|
||||
#include <algorithm>
|
||||
#include <cfloat>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
|
||||
namespace transformer {
|
||||
|
||||
// Vectorized load/store helpers — 128-bit (8 half elements) for max bandwidth
|
||||
union VecHalf8 {
|
||||
float4 f4;
|
||||
half h[8];
|
||||
};
|
||||
union VecHalf4 {
|
||||
float2 f2;
|
||||
half h[4];
|
||||
};
|
||||
union VecHalf2 {
|
||||
unsigned int u;
|
||||
half h[2];
|
||||
};
|
||||
union VecBF16_4 {
|
||||
float2 f2;
|
||||
__nv_bfloat16 b[4];
|
||||
};
|
||||
union VecBF16_8 {
|
||||
float4 f4;
|
||||
__nv_bfloat16 b[8];
|
||||
};
|
||||
|
||||
// FP16 GEMM via cuBLAS with tensor-op math
|
||||
static inline void linear_gemm(cublasHandle_t cub,
|
||||
int m,
|
||||
int n,
|
||||
int k,
|
||||
const half* x,
|
||||
const half* w_torch,
|
||||
half* y)
|
||||
{
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
BK(cublasGemmEx(
|
||||
cub,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
n, m, k,
|
||||
&alpha,
|
||||
w_torch, CUDA_R_16F, k,
|
||||
x, CUDA_R_16F, k,
|
||||
&beta,
|
||||
y, CUDA_R_16F, n,
|
||||
CUBLAS_COMPUTE_32F_FAST_16F,
|
||||
CUBLAS_GEMM_DEFAULT_TENSOR_OP));
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float warp_reduce_sum(float v)
|
||||
{
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
v += __shfl_down_sync(0xffffffff, v, offset);
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float warp_reduce_max(float v)
|
||||
{
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
v = fmaxf(v, __shfl_down_sync(0xffffffff, v, offset));
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float block_reduce_sum(float v)
|
||||
{
|
||||
constexpr int kWarps = kKernelThreads / 32;
|
||||
__shared__ float warp_vals[kWarps];
|
||||
int lane = threadIdx.x & 31;
|
||||
int wid = threadIdx.x >> 5;
|
||||
|
||||
v = warp_reduce_sum(v);
|
||||
if (lane == 0) {
|
||||
warp_vals[wid] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (wid == 0) {
|
||||
float x = (lane < kWarps) ? warp_vals[lane] : 0.0f;
|
||||
x = warp_reduce_sum(x);
|
||||
if (lane == 0) {
|
||||
warp_vals[0] = x;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return warp_vals[0];
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float block_reduce_max(float v)
|
||||
{
|
||||
constexpr int kWarps = kKernelThreads / 32;
|
||||
__shared__ float warp_vals[kWarps];
|
||||
int lane = threadIdx.x & 31;
|
||||
int wid = threadIdx.x >> 5;
|
||||
|
||||
v = warp_reduce_max(v);
|
||||
if (lane == 0) {
|
||||
warp_vals[wid] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (wid == 0) {
|
||||
float x = (lane < kWarps) ? warp_vals[lane] : -FLT_MAX;
|
||||
x = warp_reduce_max(x);
|
||||
if (lane == 0) {
|
||||
warp_vals[0] = x;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return warp_vals[0];
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void silu_kernel(const half* __restrict__ in,
|
||||
half* __restrict__ out,
|
||||
int n)
|
||||
{
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < n) {
|
||||
VecHalf8 v;
|
||||
v.f4 = *reinterpret_cast<const float4*>(in + base);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
float x = __half2float(v.h[j]);
|
||||
v.h[j] = __float2half_rn(x / (1.0f + expf(-x)));
|
||||
}
|
||||
*reinterpret_cast<float4*>(out + base) = v.f4;
|
||||
} else {
|
||||
for (int i = base; i < min(base + 8, n); i++) {
|
||||
float x = __half2float(in[i]);
|
||||
out[i] = __float2half_rn(x / (1.0f + expf(-x)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void bias_add_kernel(half* __restrict__ out,
|
||||
const half* __restrict__ bias,
|
||||
int rows,
|
||||
int cols)
|
||||
{
|
||||
int total = rows * cols;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
|
||||
if (base + 3 < total) {
|
||||
int col = base % cols;
|
||||
if (col + 3 < cols) {
|
||||
VecHalf4 ov, bv;
|
||||
ov.f2 = *reinterpret_cast<const float2*>(out + base);
|
||||
bv.f2 = *reinterpret_cast<const float2*>(bias + col);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++)
|
||||
ov.h[j] = __float2half_rn(__half2float(ov.h[j]) + __half2float(bv.h[j]));
|
||||
*reinterpret_cast<float2*>(out + base) = ov.f2;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 4 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int col = idx % cols;
|
||||
float v = __half2float(out[idx]) + __half2float(bias[col]);
|
||||
out[idx] = __float2half_rn(v);
|
||||
}
|
||||
}
|
||||
|
||||
// Input x is bf16 (model runs in bf16). Output is fp16. BF16→FP16 conversion is
|
||||
// software-emulated on SM < 80 but is IO-bound so negligible overhead.
|
||||
__global__ __launch_bounds__(256, 2)
|
||||
void layernorm_mod_kernel(const __nv_bfloat16* __restrict__ x,
|
||||
const half* __restrict__ mod,
|
||||
half* __restrict__ out,
|
||||
int seq)
|
||||
{
|
||||
int row = blockIdx.x;
|
||||
int b = row / seq;
|
||||
|
||||
const __nv_bfloat16* x_row = x + static_cast<long long>(row) * HIDDEN;
|
||||
half* out_row = out + static_cast<long long>(row) * HIDDEN;
|
||||
const half* shift = mod + static_cast<long long>(b) * MODDIM;
|
||||
const half* scale = shift + HIDDEN;
|
||||
|
||||
float sum = 0.0f;
|
||||
float sq = 0.0f;
|
||||
for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
|
||||
VecBF16_4 bv;
|
||||
bv.f2 = *reinterpret_cast<const float2*>(x_row + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float v = __bfloat162float(bv.b[j]);
|
||||
sum += v;
|
||||
sq += v * v;
|
||||
}
|
||||
}
|
||||
float mean = block_reduce_sum(sum) / HIDDEN;
|
||||
float var = block_reduce_sum(sq) / HIDDEN - mean * mean;
|
||||
float inv = rsqrtf(var + LN_EPS);
|
||||
|
||||
for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
|
||||
VecBF16_4 xv;
|
||||
xv.f2 = *reinterpret_cast<const float2*>(x_row + i);
|
||||
VecHalf4 sc_v, sh_v, ov;
|
||||
sc_v.f2 = *reinterpret_cast<const float2*>(scale + i);
|
||||
sh_v.f2 = *reinterpret_cast<const float2*>(shift + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float v = __bfloat162float(xv.b[j]);
|
||||
float nrm = (v - mean) * inv;
|
||||
ov.h[j] = __float2half_rn(nrm * (1.0f + __half2float(sc_v.h[j])) + __half2float(sh_v.h[j]));
|
||||
}
|
||||
*reinterpret_cast<float2*>(out_row + i) = ov.f2;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void split_qkv_kernel(const half* __restrict__ l1o,
|
||||
half* __restrict__ q,
|
||||
half* __restrict__ k,
|
||||
half* __restrict__ v,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = seq * HIDDEN;
|
||||
int batch_total = gridDim.y;
|
||||
int linear = blockIdx.y * total + idx;
|
||||
int max_linear = batch_total * total;
|
||||
if (linear >= max_linear) return;
|
||||
|
||||
int token = linear / HIDDEN;
|
||||
int h = linear % HIDDEN;
|
||||
int b = token / seq;
|
||||
int s = token % seq;
|
||||
int head = h / HDIM;
|
||||
int d = h % HDIM;
|
||||
|
||||
long long base_l1 = static_cast<long long>(token) * LIN1DIM;
|
||||
long long dst = (((static_cast<long long>(b) * HEADS + head) * seq) + s) * HDIM + d;
|
||||
|
||||
q[dst] = l1o[base_l1 + h];
|
||||
k[dst] = l1o[base_l1 + HIDDEN + h];
|
||||
v[dst] = l1o[base_l1 + 2 * HIDDEN + h];
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void copy_mlp_kernel(const half* __restrict__ l1o,
|
||||
half* __restrict__ mlp,
|
||||
int rows)
|
||||
{
|
||||
int total = rows * FFDIM;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int row = base / FFDIM;
|
||||
int col = base % FFDIM;
|
||||
if (col + 7 < FFDIM) {
|
||||
float4 v = *reinterpret_cast<const float4*>(l1o + static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col);
|
||||
*reinterpret_cast<float4*>(mlp + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int row = idx / FFDIM;
|
||||
int col = idx % FFDIM;
|
||||
mlp[idx] = l1o[static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void extract_gate_kernel(const half* __restrict__ mod,
|
||||
half* __restrict__ gate,
|
||||
int batch)
|
||||
{
|
||||
int total = batch * HIDDEN;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int b = base / HIDDEN;
|
||||
int h = base % HIDDEN;
|
||||
if (h + 7 < HIDDEN) {
|
||||
float4 v = *reinterpret_cast<const float4*>(mod + static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h);
|
||||
*reinterpret_cast<float4*>(gate + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int b = idx / HIDDEN;
|
||||
int h = idx % HIDDEN;
|
||||
gate[idx] = mod[static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 2)
|
||||
void rmsnorm_scale_kernel(half* __restrict__ data,
|
||||
const half* __restrict__ scale)
|
||||
{
|
||||
int row = blockIdx.x;
|
||||
half* row_ptr = data + static_cast<long long>(row) * HDIM;
|
||||
|
||||
float sq = 0.0f;
|
||||
for (int i = threadIdx.x * 4; i < HDIM; i += blockDim.x * 4) {
|
||||
VecHalf4 v;
|
||||
v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float val = __half2float(v.h[j]);
|
||||
sq += val * val;
|
||||
}
|
||||
}
|
||||
float inv = rsqrtf(block_reduce_sum(sq) / HDIM + RMS_EPS);
|
||||
|
||||
for (int i = threadIdx.x * 4; i < HDIM; i += blockDim.x * 4) {
|
||||
VecHalf4 v, sc;
|
||||
v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
|
||||
sc.f2 = *reinterpret_cast<const float2*>(scale + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++)
|
||||
v.h[j] = __float2half_rn(__half2float(v.h[j]) * inv * __half2float(sc.h[j]));
|
||||
*reinterpret_cast<float2*>(row_ptr + i) = v.f2;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void rope_kernel(half* __restrict__ q,
|
||||
half* __restrict__ k,
|
||||
const half* __restrict__ pe,
|
||||
int batch,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = batch * HEADS * seq * (HDIM / 2);
|
||||
if (idx >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int pair = idx % (HDIM / 2);
|
||||
int t0 = idx / (HDIM / 2);
|
||||
int s = t0 % seq;
|
||||
int t1 = t0 / seq;
|
||||
int h = t1 % HEADS;
|
||||
int b = t1 / HEADS;
|
||||
|
||||
long long base = (((static_cast<long long>(b) * HEADS + h) * seq) + s) * HDIM + 2 * pair;
|
||||
long long pe_off = (static_cast<long long>(s) * (HDIM / 2) + pair) * 2;
|
||||
|
||||
VecHalf2 pe_v;
|
||||
pe_v.u = *reinterpret_cast<const unsigned int*>(pe + pe_off);
|
||||
float c = __half2float(pe_v.h[0]);
|
||||
float sn = __half2float(pe_v.h[1]);
|
||||
|
||||
VecHalf2 qv;
|
||||
qv.u = *reinterpret_cast<const unsigned int*>(q + base);
|
||||
float q0 = __half2float(qv.h[0]);
|
||||
float q1 = __half2float(qv.h[1]);
|
||||
qv.h[0] = __float2half_rn(q0 * c - q1 * sn);
|
||||
qv.h[1] = __float2half_rn(q0 * sn + q1 * c);
|
||||
*reinterpret_cast<unsigned int*>(q + base) = qv.u;
|
||||
|
||||
VecHalf2 kv;
|
||||
kv.u = *reinterpret_cast<const unsigned int*>(k + base);
|
||||
float k0 = __half2float(kv.h[0]);
|
||||
float k1 = __half2float(kv.h[1]);
|
||||
kv.h[0] = __float2half_rn(k0 * c - k1 * sn);
|
||||
kv.h[1] = __float2half_rn(k0 * sn + k1 * c);
|
||||
*reinterpret_cast<unsigned int*>(k + base) = kv.u;
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void gelu_tanh_kernel(half* __restrict__ data,
|
||||
int n)
|
||||
{
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < n) {
|
||||
VecHalf8 v;
|
||||
v.f4 = *reinterpret_cast<const float4*>(data + base);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
float x = __half2float(v.h[j]);
|
||||
float x3 = x * x * x;
|
||||
v.h[j] = __float2half_rn(0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x3))));
|
||||
}
|
||||
*reinterpret_cast<float4*>(data + base) = v.f4;
|
||||
} else {
|
||||
for (int i = base; i < min(base + 8, n); i++) {
|
||||
float x = __half2float(data[i]);
|
||||
float x3 = x * x * x;
|
||||
data[i] = __float2half_rn(0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x3))));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void concat_final_kernel(const half* __restrict__ attn,
|
||||
const half* __restrict__ mlp,
|
||||
half* __restrict__ out,
|
||||
int rows)
|
||||
{
|
||||
int total = rows * LIN2DIM;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int row = base / LIN2DIM;
|
||||
int col = base % LIN2DIM;
|
||||
if (col + 7 < LIN2DIM) {
|
||||
float4 v;
|
||||
if (col + 7 < HIDDEN) {
|
||||
v = *reinterpret_cast<const float4*>(attn + static_cast<long long>(row) * HIDDEN + col);
|
||||
} else if (col >= HIDDEN) {
|
||||
v = *reinterpret_cast<const float4*>(mlp + static_cast<long long>(row) * FFDIM + (col - HIDDEN));
|
||||
} else {
|
||||
VecHalf8 vh;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
int c = col + j;
|
||||
if (c < HIDDEN)
|
||||
vh.h[j] = attn[static_cast<long long>(row) * HIDDEN + c];
|
||||
else
|
||||
vh.h[j] = mlp[static_cast<long long>(row) * FFDIM + (c - HIDDEN)];
|
||||
}
|
||||
v = vh.f4;
|
||||
}
|
||||
*reinterpret_cast<float4*>(out + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int row = idx / LIN2DIM;
|
||||
int col = idx % LIN2DIM;
|
||||
if (col < HIDDEN)
|
||||
out[idx] = attn[static_cast<long long>(row) * HIDDEN + col];
|
||||
else
|
||||
out[idx] = mlp[static_cast<long long>(row) * FFDIM + (col - HIDDEN)];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void gate_residual_kernel(const half* __restrict__ residual,
|
||||
const half* __restrict__ gate,
|
||||
half* __restrict__ out,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = gridDim.y * seq * HIDDEN;
|
||||
int linear = blockIdx.y * (seq * HIDDEN) + idx;
|
||||
if (linear >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int h = linear % HIDDEN;
|
||||
int token = linear / HIDDEN;
|
||||
int b = token / seq;
|
||||
float gate_v = __half2float(gate[static_cast<long long>(b) * HIDDEN + h]);
|
||||
float res = __half2float(residual[linear]);
|
||||
float v = __half2float(out[linear]);
|
||||
out[linear] = __float2half_rn(res + gate_v * v);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// SSBEngine — FP16-only (no FP8, no CUTLASS)
|
||||
// =========================================================================
|
||||
|
||||
SSBEngine::SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size)
|
||||
: B_(batch),
|
||||
S_(seq),
|
||||
cub_(nullptr),
|
||||
workspace_(nullptr),
|
||||
workspace_bytes_(0),
|
||||
d_silu_(nullptr),
|
||||
d_mod_(nullptr),
|
||||
d_xmod_(nullptr),
|
||||
d_l1o_(nullptr)
|
||||
{
|
||||
BK(cublasCreate(&cub_));
|
||||
BK(cublasSetMathMode(cub_, CUBLAS_TENSOR_OP_MATH));
|
||||
|
||||
size_t min_needed = ssb_workspace_bytes(B_, S_);
|
||||
if (workspace_buf_size < min_needed) {
|
||||
throw std::runtime_error("Workspace buffer too small: need " +
|
||||
std::to_string(min_needed) + " got " + std::to_string(workspace_buf_size));
|
||||
}
|
||||
|
||||
char* p = static_cast<char*>(workspace_buf);
|
||||
size_t off = 0;
|
||||
size_t rows = static_cast<size_t>(B_) * S_;
|
||||
|
||||
#define CARVE(ptr, type, count) \
|
||||
off = (off + 255) & ~size_t(255); \
|
||||
ptr = reinterpret_cast<type*>(p + off); \
|
||||
off += static_cast<size_t>(count) * sizeof(type);
|
||||
|
||||
CARVE(d_silu_, half, B_ * HIDDEN)
|
||||
CARVE(d_mod_, half, static_cast<size_t>(B_) * MODDIM)
|
||||
CARVE(d_xmod_, half, rows * HIDDEN)
|
||||
CARVE(d_l1o_, half, rows * LIN1DIM)
|
||||
|
||||
#undef CARVE
|
||||
|
||||
off = (off + 255) & ~size_t(255);
|
||||
workspace_ = p + off;
|
||||
workspace_bytes_ = workspace_buf_size - off;
|
||||
}
|
||||
|
||||
SSBEngine::~SSBEngine()
|
||||
{
|
||||
if (cub_) cublasDestroy(cub_);
|
||||
}
|
||||
|
||||
void SSBEngine::prepare(const half* d_x,
|
||||
const half* d_vec,
|
||||
const half* d_pe,
|
||||
const half* mod_w,
|
||||
const half* mod_b,
|
||||
const half* l1w,
|
||||
const half* l1b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
half* d_q,
|
||||
half* d_k,
|
||||
half* d_v,
|
||||
half* d_mlp,
|
||||
half* d_gate,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
BK(cublasSetStream(cub_, stream));
|
||||
|
||||
int rows = B_ * S_;
|
||||
int threads = kKernelThreads;
|
||||
|
||||
// SiLU on vec
|
||||
silu_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_vec, d_silu_, B_ * HIDDEN);
|
||||
|
||||
// Modulation GEMM: FP16→FP16
|
||||
linear_gemm(cub_, B_, MODDIM, HIDDEN, d_silu_, mod_w, d_mod_);
|
||||
int out_n = B_ * MODDIM;
|
||||
bias_add_kernel<<<(out_n + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(d_mod_, mod_b, B_, MODDIM);
|
||||
|
||||
// LayerNorm + Modulation (reads x as BF16, outputs FP16)
|
||||
layernorm_mod_kernel<<<rows, threads, 0, stream>>>(reinterpret_cast<const __nv_bfloat16*>(d_x), d_mod_, d_xmod_, S_);
|
||||
|
||||
// Linear1 GEMM: FP16→FP16
|
||||
linear_gemm(cub_, rows, LIN1DIM, HIDDEN, d_xmod_, l1w, d_l1o_);
|
||||
int l1_n = rows * LIN1DIM;
|
||||
bias_add_kernel<<<(l1_n + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(d_l1o_, l1b, rows, LIN1DIM);
|
||||
|
||||
// Split Q/K/V from Linear1 output
|
||||
dim3 split_grid((S_ * HIDDEN + threads - 1) / threads, B_);
|
||||
split_qkv_kernel<<<split_grid, threads, 0, stream>>>(d_l1o_, d_q, d_k, d_v, S_);
|
||||
|
||||
// Copy MLP portion
|
||||
copy_mlp_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_l1o_, d_mlp, rows);
|
||||
|
||||
// Extract gate from modulation output
|
||||
extract_gate_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mod_, d_gate, B_);
|
||||
|
||||
// RMSNorm + scale on Q and K
|
||||
int qk_rows = B_ * HEADS * S_;
|
||||
rmsnorm_scale_kernel<<<qk_rows, threads, 0, stream>>>(d_q, qns);
|
||||
rmsnorm_scale_kernel<<<qk_rows, threads, 0, stream>>>(d_k, kns);
|
||||
|
||||
// RoPE
|
||||
int rope_total = B_ * HEADS * S_ * (HDIM / 2);
|
||||
rope_kernel<<<(rope_total + threads - 1) / threads, threads, 0, stream>>>(d_q, d_k, d_pe, B_, S_);
|
||||
|
||||
CK(cudaGetLastError());
|
||||
}
|
||||
|
||||
void SSBEngine::finish(const half* d_x,
|
||||
const half* d_attn,
|
||||
half* d_mlp,
|
||||
const half* d_gate,
|
||||
half* d_out,
|
||||
const half* l2w,
|
||||
const half* l2b,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
BK(cublasSetStream(cub_, stream));
|
||||
|
||||
int rows = B_ * S_;
|
||||
int threads = kKernelThreads;
|
||||
|
||||
// GELU on MLP
|
||||
gelu_tanh_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mlp, rows * FFDIM);
|
||||
|
||||
// Concat attention + MLP → d_l1o_ (reused as temp)
|
||||
concat_final_kernel<<<(rows * LIN2DIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_attn, d_mlp, d_l1o_, rows);
|
||||
|
||||
// Linear2 GEMM: FP16→FP16
|
||||
linear_gemm(cub_, rows, HIDDEN, LIN2DIM, d_l1o_, l2w, d_out);
|
||||
|
||||
// Bias add
|
||||
int out_blocks = (rows * HIDDEN + threads * 4 - 1) / (threads * 4);
|
||||
bias_add_kernel<<<out_blocks, threads, 0, stream>>>(d_out, l2b, rows, HIDDEN);
|
||||
|
||||
// Gate + residual: out = x + gate * out
|
||||
dim3 gate_grid((S_ * HIDDEN + threads - 1) / threads, B_);
|
||||
gate_residual_kernel<<<gate_grid, threads, 0, stream>>>(d_x, d_gate, d_out, S_);
|
||||
|
||||
CK(cudaGetLastError());
|
||||
}
|
||||
|
||||
} // namespace transformer
|
||||
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
Build script for flux_ext FP16 PyTorch C++/CUDA extension — 10-series (SM 61, Pascal).
|
||||
|
||||
Usage:
|
||||
py -3.10 setup_ext.py install
|
||||
|
||||
Compiles the FP16-only FLUXNATION kernel (no FP8, no CUTLASS).
|
||||
Compatible with GTX 1060, 1070, 1080, 1080 Ti.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["CUDA_HOME"] = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8"
|
||||
os.environ["CUDA_PATH"] = os.environ["CUDA_HOME"]
|
||||
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
|
||||
|
||||
ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
setup(
|
||||
name="flux_ext",
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
name="flux_ext",
|
||||
sources=[
|
||||
os.path.join(ROOT, "csrc", "FLUXPATCH_FP16.cpp"),
|
||||
os.path.join(ROOT, "csrc", "FLUXNATION_FP16.cu"),
|
||||
os.path.join(ROOT, "flux_opt_fp16.cu"),
|
||||
],
|
||||
include_dirs=[
|
||||
ROOT, # SINGLE_STREAM_BLOCK_FP16.H
|
||||
],
|
||||
extra_compile_args={
|
||||
"cxx": ["/O2", "/std:c++17", "/EHsc", "/MD"],
|
||||
"nvcc": [
|
||||
"-O3",
|
||||
"-std=c++17",
|
||||
"-arch=sm_61",
|
||||
"-DFLUX_TORCH_EXT=1",
|
||||
"--allow-unsupported-compiler",
|
||||
"-diag-suppress=20013",
|
||||
"-diag-suppress=20015",
|
||||
"-Xcompiler", "/MD",
|
||||
"-Xcompiler", "/EHsc",
|
||||
"-Xcompiler", "/Zc:preprocessor",
|
||||
],
|
||||
},
|
||||
libraries=["cublas"],
|
||||
)
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension},
|
||||
)
|
||||
@@ -0,0 +1,106 @@
|
||||
#ifndef SINGLE_STREAM_BLOCK_FP16_H
|
||||
#define SINGLE_STREAM_BLOCK_FP16_H
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cublas_v2.h>
|
||||
#include <cstdio>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
namespace transformer {
|
||||
|
||||
constexpr int HIDDEN = 3072;
|
||||
constexpr int HEADS = 24;
|
||||
constexpr int HDIM = 128;
|
||||
constexpr int FFDIM = 12288;
|
||||
constexpr int MODDIM = 3 * HIDDEN;
|
||||
constexpr int LIN1DIM = 3 * HIDDEN + FFDIM;
|
||||
constexpr int LIN2DIM = HIDDEN + FFDIM;
|
||||
constexpr float LN_EPS = 1e-6f;
|
||||
constexpr float RMS_EPS = 1e-6f;
|
||||
constexpr float ATT_SC = 0.0883883476f;
|
||||
|
||||
static int const kKernelThreads = 256;
|
||||
static_assert((kKernelThreads % 32) == 0, "kKernelThreads must be a multiple of 32.");
|
||||
|
||||
#define CK(call) do { \
|
||||
cudaError_t err__ = (call); \
|
||||
if (err__ != cudaSuccess) { \
|
||||
throw std::runtime_error(std::string("CUDA ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + cudaGetErrorString(err__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define BK(call) do { \
|
||||
cublasStatus_t st__ = (call); \
|
||||
if (st__ != CUBLAS_STATUS_SUCCESS) { \
|
||||
throw std::runtime_error(std::string("cuBLAS ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + std::to_string((int)st__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static inline size_t ssb_workspace_bytes(int batch, int seq) {
|
||||
size_t rows = static_cast<size_t>(batch) * seq;
|
||||
size_t b = static_cast<size_t>(batch);
|
||||
size_t sz = 0;
|
||||
#define ADDBUF(count, elem_size) \
|
||||
sz = (sz + 255) & ~size_t(255); sz += (count) * (elem_size);
|
||||
ADDBUF(b * HIDDEN, 2); // d_silu_
|
||||
ADDBUF(b * MODDIM, 2); // d_mod_
|
||||
ADDBUF(rows * HIDDEN, 2); // d_xmod_
|
||||
ADDBUF(rows * LIN1DIM, 2); // d_l1o_
|
||||
sz = (sz + 255) & ~size_t(255);
|
||||
sz += 32ull << 20; // 32MB cuBLAS workspace
|
||||
#undef ADDBUF
|
||||
return sz;
|
||||
}
|
||||
|
||||
class SSBEngine {
|
||||
public:
|
||||
SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size);
|
||||
~SSBEngine();
|
||||
|
||||
void prepare(const half* d_x,
|
||||
const half* d_vec,
|
||||
const half* d_pe,
|
||||
const half* mod_w,
|
||||
const half* mod_b,
|
||||
const half* l1w,
|
||||
const half* l1b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
half* d_q,
|
||||
half* d_k,
|
||||
half* d_v,
|
||||
half* d_mlp,
|
||||
half* d_gate,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
void finish(const half* d_x,
|
||||
const half* d_attn,
|
||||
half* d_mlp,
|
||||
const half* d_gate,
|
||||
half* d_out,
|
||||
const half* l2w,
|
||||
const half* l2b,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
int batch() const { return B_; }
|
||||
int seq() const { return S_; }
|
||||
|
||||
private:
|
||||
int B_;
|
||||
int S_;
|
||||
cublasHandle_t cub_;
|
||||
void* workspace_;
|
||||
size_t workspace_bytes_;
|
||||
|
||||
half* d_silu_;
|
||||
half* d_mod_;
|
||||
half* d_xmod_;
|
||||
half* d_l1o_;
|
||||
};
|
||||
|
||||
} // namespace transformer
|
||||
|
||||
#endif // SINGLE_STREAM_BLOCK_FP16_H
|
||||
@@ -0,0 +1,25 @@
|
||||
@echo off
|
||||
setlocal
|
||||
|
||||
REM FLUXNATION FP16 Build — 20-series (SM 75, Turing)
|
||||
REM Compatible with RTX 2060/2070/2080/2080 Ti
|
||||
|
||||
set "VSVARS=C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat"
|
||||
if not exist "%VSVARS%" (
|
||||
echo [flux_ext] ERROR: VS 2022 vcvars64.bat not found
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
call "%VSVARS%" || exit /b 1
|
||||
|
||||
set DISTUTILS_USE_SDK=1
|
||||
|
||||
echo [flux_ext] Building FLUXNATION FP16 for 20-series (SM 75)...
|
||||
py -3.10 setup_ext.py install
|
||||
if errorlevel 1 (
|
||||
echo [flux_ext] Build FAILED
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
echo [flux_ext] Build succeeded! flux_ext (FP16, SM 75) is now installed.
|
||||
echo [flux_ext] Restart ComfyUI to use the new kernel.
|
||||
@@ -0,0 +1,198 @@
|
||||
/*
|
||||
* FLUXNATION_FP16.cu — PyTorch C++ extension for FLUXNATION (FP16-only)
|
||||
*
|
||||
* Wraps the FP16 SSBEngine as torch.ops custom operators.
|
||||
* All weights converted to FP16 internally — no FP8, no CUTLASS.
|
||||
* Compatible with SM 61+ (Pascal, Turing, Ampere).
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#include "SINGLE_STREAM_BLOCK_FP16.H"
|
||||
|
||||
#include <unordered_map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
|
||||
using namespace transformer;
|
||||
|
||||
// =========================================================================
|
||||
// Global engine cache — keyed by (device, batch, seq_len)
|
||||
// =========================================================================
|
||||
struct EngineEntry {
|
||||
std::unique_ptr<SSBEngine> engine;
|
||||
at::Tensor workspace;
|
||||
};
|
||||
|
||||
static std::unordered_map<int64_t, EngineEntry> g_engine_cache;
|
||||
static std::mutex g_engine_mutex;
|
||||
|
||||
static int64_t make_engine_key(int device, int batch, int seq) {
|
||||
return (int64_t)device * 1000000000LL + (int64_t)batch * 1000000LL + seq;
|
||||
}
|
||||
|
||||
static SSBEngine* get_or_create_engine(int batch, int seq, at::Device device) {
|
||||
int dev_idx = device.index();
|
||||
int64_t key = make_engine_key(dev_idx, batch, seq);
|
||||
|
||||
std::lock_guard<std::mutex> lock(g_engine_mutex);
|
||||
auto it = g_engine_cache.find(key);
|
||||
if (it != g_engine_cache.end()) {
|
||||
return it->second.engine.get();
|
||||
}
|
||||
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
size_t ws_bytes = ssb_workspace_bytes(batch, seq);
|
||||
auto ws_tensor = at::empty({(int64_t)ws_bytes}, at::TensorOptions().dtype(at::kByte).device(device));
|
||||
|
||||
auto engine = std::make_unique<SSBEngine>(batch, seq, ws_tensor.data_ptr(), ws_bytes);
|
||||
auto* ptr = engine.get();
|
||||
|
||||
EngineEntry entry;
|
||||
entry.engine = std::move(engine);
|
||||
entry.workspace = ws_tensor;
|
||||
g_engine_cache[key] = std::move(entry);
|
||||
|
||||
fprintf(stderr, "[flux_ext FP16] Created engine B=%d S=%d ws=%.1fMB dev=%d\n",
|
||||
batch, seq, ws_bytes / (1024.0 * 1024.0), dev_idx);
|
||||
fflush(stderr);
|
||||
|
||||
return ptr;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// torch.ops implementation: ssb_prepare
|
||||
//
|
||||
// Same Python API as FP8 version — scale tensors are accepted but ignored.
|
||||
// All weights converted to FP16 internally.
|
||||
// =========================================================================
|
||||
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
|
||||
ssb_prepare_impl(
|
||||
at::Tensor x, // [B, S, 3072] bf16
|
||||
at::Tensor vec, // [B, 3072] or [B, 1, 3072] fp16
|
||||
at::Tensor pe, // [S, 64, 2] fp16
|
||||
at::Tensor mod_w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor mod_w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor mod_b, // fp16
|
||||
at::Tensor l1w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor l1w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor l1b, // fp16
|
||||
at::Tensor qns, // fp16
|
||||
at::Tensor kns) // fp16
|
||||
{
|
||||
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
|
||||
TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");
|
||||
|
||||
int B = x.size(0);
|
||||
int S = x.size(1);
|
||||
auto device = x.device();
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(device.index());
|
||||
|
||||
// Ensure correct dtypes — x stays BF16 (layernorm kernel reads it as BF16)
|
||||
auto x_bf16 = x.to(at::kBFloat16).contiguous();
|
||||
auto vec_fp16 = vec.dim() == 3 ? vec.squeeze(1).to(at::kHalf).contiguous()
|
||||
: vec.to(at::kHalf).contiguous();
|
||||
auto pe_fp16 = pe.to(at::kHalf).contiguous();
|
||||
|
||||
// Convert ALL weights to FP16 (no FP8 on this target)
|
||||
auto mod_w_fp16 = mod_w.to(at::kHalf).contiguous();
|
||||
auto mod_b_fp16 = mod_b.to(at::kHalf).contiguous();
|
||||
auto l1w_fp16 = l1w.to(at::kHalf).contiguous();
|
||||
auto l1b_fp16 = l1b.to(at::kHalf).contiguous();
|
||||
auto qns_fp16 = qns.to(at::kHalf).contiguous();
|
||||
auto kns_fp16 = kns.to(at::kHalf).contiguous();
|
||||
|
||||
// Allocate output tensors
|
||||
auto opts_fp16 = at::TensorOptions().dtype(at::kHalf).device(device);
|
||||
auto q = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto k = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto v = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto mlp = at::empty({B, S, FFDIM}, opts_fp16);
|
||||
auto gate = at::empty({B, HIDDEN}, opts_fp16);
|
||||
|
||||
SSBEngine* engine = get_or_create_engine(B, S, device);
|
||||
|
||||
engine->prepare(
|
||||
reinterpret_cast<const half*>(x_bf16.data_ptr()),
|
||||
reinterpret_cast<const half*>(vec_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(pe_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(mod_w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(mod_b_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l1w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l1b_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(qns_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(kns_fp16.data_ptr()),
|
||||
reinterpret_cast<half*>(q.data_ptr()),
|
||||
reinterpret_cast<half*>(k.data_ptr()),
|
||||
reinterpret_cast<half*>(v.data_ptr()),
|
||||
reinterpret_cast<half*>(mlp.data_ptr()),
|
||||
reinterpret_cast<half*>(gate.data_ptr()),
|
||||
stream);
|
||||
|
||||
return std::make_tuple(q, k, v, mlp, gate);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// torch.ops implementation: ssb_finish
|
||||
//
|
||||
// FP16 path: converts residual (x) from BF16→FP16, runs all-FP16 engine,
|
||||
// converts result back to BF16 for model compatibility.
|
||||
// =========================================================================
|
||||
at::Tensor ssb_finish_impl(
|
||||
at::Tensor x, // [B, S, 3072] bf16 — original input for residual
|
||||
at::Tensor attn, // [B, S, 3072] fp16 — attention output
|
||||
at::Tensor mlp, // [B, S, 12288] fp16 — MLP intermediate (pre-GELU)
|
||||
at::Tensor gate, // [B, 3072] fp16 — gate values
|
||||
at::Tensor l2w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor l2w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor l2b) // fp16 — L2 bias
|
||||
{
|
||||
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
|
||||
TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");
|
||||
|
||||
int B = x.size(0);
|
||||
int S = x.size(1);
|
||||
auto device = x.device();
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(device.index());
|
||||
|
||||
// Convert residual from BF16 to FP16 for the FP16 engine
|
||||
auto x_fp16 = x.to(at::kHalf).contiguous();
|
||||
auto attn_fp16 = attn.to(at::kHalf).contiguous();
|
||||
auto mlp_c = mlp.contiguous();
|
||||
auto gate_c = gate.contiguous();
|
||||
auto l2w_fp16 = l2w.to(at::kHalf).contiguous();
|
||||
auto l2b_fp16 = l2b.to(at::kHalf).contiguous();
|
||||
|
||||
// Result in FP16 (engine is all-FP16)
|
||||
auto opts_fp16 = at::TensorOptions().dtype(at::kHalf).device(device);
|
||||
auto result_fp16 = at::empty({B, S, HIDDEN}, opts_fp16);
|
||||
|
||||
SSBEngine* engine = get_or_create_engine(B, S, device);
|
||||
|
||||
engine->finish(
|
||||
reinterpret_cast<const half*>(x_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(attn_fp16.data_ptr()),
|
||||
reinterpret_cast<half*>(mlp_c.data_ptr()),
|
||||
reinterpret_cast<const half*>(gate_c.data_ptr()),
|
||||
reinterpret_cast<half*>(result_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l2w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l2b_fp16.data_ptr()),
|
||||
stream);
|
||||
|
||||
// Convert back to BF16 for model compatibility
|
||||
return result_fp16.to(at::kBFloat16);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Cleanup — destroys all cached engines
|
||||
// =========================================================================
|
||||
void ssb_cleanup_impl() {
|
||||
std::lock_guard<std::mutex> lock(g_engine_mutex);
|
||||
g_engine_cache.clear();
|
||||
fprintf(stderr, "[flux_ext FP16] All engines destroyed\n");
|
||||
fflush(stderr);
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
/*
|
||||
* 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)";
|
||||
}
|
||||
@@ -0,0 +1,611 @@
|
||||
/*
|
||||
* flux_opt_fp16.cu — FP16-only SSBEngine for FLUXNATION
|
||||
*
|
||||
* Derived from flux_opt.cu. All FP8/CUTLASS code removed.
|
||||
* Uses cuBLAS FP16 tensor-op GEMMs only.
|
||||
* Compatible with SM 61+ (Pascal, Turing, Ampere).
|
||||
*/
|
||||
|
||||
#include "SINGLE_STREAM_BLOCK_FP16.H"
|
||||
#include <algorithm>
|
||||
#include <cfloat>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
|
||||
namespace transformer {
|
||||
|
||||
// Vectorized load/store helpers — 128-bit (8 half elements) for max bandwidth
|
||||
union VecHalf8 {
|
||||
float4 f4;
|
||||
half h[8];
|
||||
};
|
||||
union VecHalf4 {
|
||||
float2 f2;
|
||||
half h[4];
|
||||
};
|
||||
union VecHalf2 {
|
||||
unsigned int u;
|
||||
half h[2];
|
||||
};
|
||||
union VecBF16_4 {
|
||||
float2 f2;
|
||||
__nv_bfloat16 b[4];
|
||||
};
|
||||
union VecBF16_8 {
|
||||
float4 f4;
|
||||
__nv_bfloat16 b[8];
|
||||
};
|
||||
|
||||
// FP16 GEMM via cuBLAS with tensor-op math
|
||||
static inline void linear_gemm(cublasHandle_t cub,
|
||||
int m,
|
||||
int n,
|
||||
int k,
|
||||
const half* x,
|
||||
const half* w_torch,
|
||||
half* y)
|
||||
{
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
BK(cublasGemmEx(
|
||||
cub,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
n, m, k,
|
||||
&alpha,
|
||||
w_torch, CUDA_R_16F, k,
|
||||
x, CUDA_R_16F, k,
|
||||
&beta,
|
||||
y, CUDA_R_16F, n,
|
||||
CUBLAS_COMPUTE_32F_FAST_16F,
|
||||
CUBLAS_GEMM_DEFAULT_TENSOR_OP));
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float warp_reduce_sum(float v)
|
||||
{
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
v += __shfl_down_sync(0xffffffff, v, offset);
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float warp_reduce_max(float v)
|
||||
{
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
v = fmaxf(v, __shfl_down_sync(0xffffffff, v, offset));
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float block_reduce_sum(float v)
|
||||
{
|
||||
constexpr int kWarps = kKernelThreads / 32;
|
||||
__shared__ float warp_vals[kWarps];
|
||||
int lane = threadIdx.x & 31;
|
||||
int wid = threadIdx.x >> 5;
|
||||
|
||||
v = warp_reduce_sum(v);
|
||||
if (lane == 0) {
|
||||
warp_vals[wid] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (wid == 0) {
|
||||
float x = (lane < kWarps) ? warp_vals[lane] : 0.0f;
|
||||
x = warp_reduce_sum(x);
|
||||
if (lane == 0) {
|
||||
warp_vals[0] = x;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return warp_vals[0];
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float block_reduce_max(float v)
|
||||
{
|
||||
constexpr int kWarps = kKernelThreads / 32;
|
||||
__shared__ float warp_vals[kWarps];
|
||||
int lane = threadIdx.x & 31;
|
||||
int wid = threadIdx.x >> 5;
|
||||
|
||||
v = warp_reduce_max(v);
|
||||
if (lane == 0) {
|
||||
warp_vals[wid] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (wid == 0) {
|
||||
float x = (lane < kWarps) ? warp_vals[lane] : -FLT_MAX;
|
||||
x = warp_reduce_max(x);
|
||||
if (lane == 0) {
|
||||
warp_vals[0] = x;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return warp_vals[0];
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void silu_kernel(const half* __restrict__ in,
|
||||
half* __restrict__ out,
|
||||
int n)
|
||||
{
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < n) {
|
||||
VecHalf8 v;
|
||||
v.f4 = *reinterpret_cast<const float4*>(in + base);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
float x = __half2float(v.h[j]);
|
||||
v.h[j] = __float2half_rn(x / (1.0f + expf(-x)));
|
||||
}
|
||||
*reinterpret_cast<float4*>(out + base) = v.f4;
|
||||
} else {
|
||||
for (int i = base; i < min(base + 8, n); i++) {
|
||||
float x = __half2float(in[i]);
|
||||
out[i] = __float2half_rn(x / (1.0f + expf(-x)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void bias_add_kernel(half* __restrict__ out,
|
||||
const half* __restrict__ bias,
|
||||
int rows,
|
||||
int cols)
|
||||
{
|
||||
int total = rows * cols;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
|
||||
if (base + 3 < total) {
|
||||
int col = base % cols;
|
||||
if (col + 3 < cols) {
|
||||
VecHalf4 ov, bv;
|
||||
ov.f2 = *reinterpret_cast<const float2*>(out + base);
|
||||
bv.f2 = *reinterpret_cast<const float2*>(bias + col);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++)
|
||||
ov.h[j] = __float2half_rn(__half2float(ov.h[j]) + __half2float(bv.h[j]));
|
||||
*reinterpret_cast<float2*>(out + base) = ov.f2;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 4 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int col = idx % cols;
|
||||
float v = __half2float(out[idx]) + __half2float(bias[col]);
|
||||
out[idx] = __float2half_rn(v);
|
||||
}
|
||||
}
|
||||
|
||||
// Input x is bf16 (model runs in bf16). Output is fp16. BF16→FP16 conversion is
|
||||
// software-emulated on SM < 80 but is IO-bound so negligible overhead.
|
||||
__global__ __launch_bounds__(256, 2)
|
||||
void layernorm_mod_kernel(const __nv_bfloat16* __restrict__ x,
|
||||
const half* __restrict__ mod,
|
||||
half* __restrict__ out,
|
||||
int seq)
|
||||
{
|
||||
int row = blockIdx.x;
|
||||
int b = row / seq;
|
||||
|
||||
const __nv_bfloat16* x_row = x + static_cast<long long>(row) * HIDDEN;
|
||||
half* out_row = out + static_cast<long long>(row) * HIDDEN;
|
||||
const half* shift = mod + static_cast<long long>(b) * MODDIM;
|
||||
const half* scale = shift + HIDDEN;
|
||||
|
||||
float sum = 0.0f;
|
||||
float sq = 0.0f;
|
||||
for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
|
||||
VecBF16_4 bv;
|
||||
bv.f2 = *reinterpret_cast<const float2*>(x_row + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float v = __bfloat162float(bv.b[j]);
|
||||
sum += v;
|
||||
sq += v * v;
|
||||
}
|
||||
}
|
||||
float mean = block_reduce_sum(sum) / HIDDEN;
|
||||
float var = block_reduce_sum(sq) / HIDDEN - mean * mean;
|
||||
float inv = rsqrtf(var + LN_EPS);
|
||||
|
||||
for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
|
||||
VecBF16_4 xv;
|
||||
xv.f2 = *reinterpret_cast<const float2*>(x_row + i);
|
||||
VecHalf4 sc_v, sh_v, ov;
|
||||
sc_v.f2 = *reinterpret_cast<const float2*>(scale + i);
|
||||
sh_v.f2 = *reinterpret_cast<const float2*>(shift + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float v = __bfloat162float(xv.b[j]);
|
||||
float nrm = (v - mean) * inv;
|
||||
ov.h[j] = __float2half_rn(nrm * (1.0f + __half2float(sc_v.h[j])) + __half2float(sh_v.h[j]));
|
||||
}
|
||||
*reinterpret_cast<float2*>(out_row + i) = ov.f2;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void split_qkv_kernel(const half* __restrict__ l1o,
|
||||
half* __restrict__ q,
|
||||
half* __restrict__ k,
|
||||
half* __restrict__ v,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = seq * HIDDEN;
|
||||
int batch_total = gridDim.y;
|
||||
int linear = blockIdx.y * total + idx;
|
||||
int max_linear = batch_total * total;
|
||||
if (linear >= max_linear) return;
|
||||
|
||||
int token = linear / HIDDEN;
|
||||
int h = linear % HIDDEN;
|
||||
int b = token / seq;
|
||||
int s = token % seq;
|
||||
int head = h / HDIM;
|
||||
int d = h % HDIM;
|
||||
|
||||
long long base_l1 = static_cast<long long>(token) * LIN1DIM;
|
||||
long long dst = (((static_cast<long long>(b) * HEADS + head) * seq) + s) * HDIM + d;
|
||||
|
||||
q[dst] = l1o[base_l1 + h];
|
||||
k[dst] = l1o[base_l1 + HIDDEN + h];
|
||||
v[dst] = l1o[base_l1 + 2 * HIDDEN + h];
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void copy_mlp_kernel(const half* __restrict__ l1o,
|
||||
half* __restrict__ mlp,
|
||||
int rows)
|
||||
{
|
||||
int total = rows * FFDIM;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int row = base / FFDIM;
|
||||
int col = base % FFDIM;
|
||||
if (col + 7 < FFDIM) {
|
||||
float4 v = *reinterpret_cast<const float4*>(l1o + static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col);
|
||||
*reinterpret_cast<float4*>(mlp + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int row = idx / FFDIM;
|
||||
int col = idx % FFDIM;
|
||||
mlp[idx] = l1o[static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void extract_gate_kernel(const half* __restrict__ mod,
|
||||
half* __restrict__ gate,
|
||||
int batch)
|
||||
{
|
||||
int total = batch * HIDDEN;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int b = base / HIDDEN;
|
||||
int h = base % HIDDEN;
|
||||
if (h + 7 < HIDDEN) {
|
||||
float4 v = *reinterpret_cast<const float4*>(mod + static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h);
|
||||
*reinterpret_cast<float4*>(gate + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int b = idx / HIDDEN;
|
||||
int h = idx % HIDDEN;
|
||||
gate[idx] = mod[static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 2)
|
||||
void rmsnorm_scale_kernel(half* __restrict__ data,
|
||||
const half* __restrict__ scale)
|
||||
{
|
||||
int row = blockIdx.x;
|
||||
half* row_ptr = data + static_cast<long long>(row) * HDIM;
|
||||
|
||||
float sq = 0.0f;
|
||||
for (int i = threadIdx.x * 4; i < HDIM; i += blockDim.x * 4) {
|
||||
VecHalf4 v;
|
||||
v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float val = __half2float(v.h[j]);
|
||||
sq += val * val;
|
||||
}
|
||||
}
|
||||
float inv = rsqrtf(block_reduce_sum(sq) / HDIM + RMS_EPS);
|
||||
|
||||
for (int i = threadIdx.x * 4; i < HDIM; i += blockDim.x * 4) {
|
||||
VecHalf4 v, sc;
|
||||
v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
|
||||
sc.f2 = *reinterpret_cast<const float2*>(scale + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++)
|
||||
v.h[j] = __float2half_rn(__half2float(v.h[j]) * inv * __half2float(sc.h[j]));
|
||||
*reinterpret_cast<float2*>(row_ptr + i) = v.f2;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void rope_kernel(half* __restrict__ q,
|
||||
half* __restrict__ k,
|
||||
const half* __restrict__ pe,
|
||||
int batch,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = batch * HEADS * seq * (HDIM / 2);
|
||||
if (idx >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int pair = idx % (HDIM / 2);
|
||||
int t0 = idx / (HDIM / 2);
|
||||
int s = t0 % seq;
|
||||
int t1 = t0 / seq;
|
||||
int h = t1 % HEADS;
|
||||
int b = t1 / HEADS;
|
||||
|
||||
long long base = (((static_cast<long long>(b) * HEADS + h) * seq) + s) * HDIM + 2 * pair;
|
||||
long long pe_off = (static_cast<long long>(s) * (HDIM / 2) + pair) * 2;
|
||||
|
||||
VecHalf2 pe_v;
|
||||
pe_v.u = *reinterpret_cast<const unsigned int*>(pe + pe_off);
|
||||
float c = __half2float(pe_v.h[0]);
|
||||
float sn = __half2float(pe_v.h[1]);
|
||||
|
||||
VecHalf2 qv;
|
||||
qv.u = *reinterpret_cast<const unsigned int*>(q + base);
|
||||
float q0 = __half2float(qv.h[0]);
|
||||
float q1 = __half2float(qv.h[1]);
|
||||
qv.h[0] = __float2half_rn(q0 * c - q1 * sn);
|
||||
qv.h[1] = __float2half_rn(q0 * sn + q1 * c);
|
||||
*reinterpret_cast<unsigned int*>(q + base) = qv.u;
|
||||
|
||||
VecHalf2 kv;
|
||||
kv.u = *reinterpret_cast<const unsigned int*>(k + base);
|
||||
float k0 = __half2float(kv.h[0]);
|
||||
float k1 = __half2float(kv.h[1]);
|
||||
kv.h[0] = __float2half_rn(k0 * c - k1 * sn);
|
||||
kv.h[1] = __float2half_rn(k0 * sn + k1 * c);
|
||||
*reinterpret_cast<unsigned int*>(k + base) = kv.u;
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void gelu_tanh_kernel(half* __restrict__ data,
|
||||
int n)
|
||||
{
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < n) {
|
||||
VecHalf8 v;
|
||||
v.f4 = *reinterpret_cast<const float4*>(data + base);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
float x = __half2float(v.h[j]);
|
||||
float x3 = x * x * x;
|
||||
v.h[j] = __float2half_rn(0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x3))));
|
||||
}
|
||||
*reinterpret_cast<float4*>(data + base) = v.f4;
|
||||
} else {
|
||||
for (int i = base; i < min(base + 8, n); i++) {
|
||||
float x = __half2float(data[i]);
|
||||
float x3 = x * x * x;
|
||||
data[i] = __float2half_rn(0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x3))));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void concat_final_kernel(const half* __restrict__ attn,
|
||||
const half* __restrict__ mlp,
|
||||
half* __restrict__ out,
|
||||
int rows)
|
||||
{
|
||||
int total = rows * LIN2DIM;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int row = base / LIN2DIM;
|
||||
int col = base % LIN2DIM;
|
||||
if (col + 7 < LIN2DIM) {
|
||||
float4 v;
|
||||
if (col + 7 < HIDDEN) {
|
||||
v = *reinterpret_cast<const float4*>(attn + static_cast<long long>(row) * HIDDEN + col);
|
||||
} else if (col >= HIDDEN) {
|
||||
v = *reinterpret_cast<const float4*>(mlp + static_cast<long long>(row) * FFDIM + (col - HIDDEN));
|
||||
} else {
|
||||
VecHalf8 vh;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
int c = col + j;
|
||||
if (c < HIDDEN)
|
||||
vh.h[j] = attn[static_cast<long long>(row) * HIDDEN + c];
|
||||
else
|
||||
vh.h[j] = mlp[static_cast<long long>(row) * FFDIM + (c - HIDDEN)];
|
||||
}
|
||||
v = vh.f4;
|
||||
}
|
||||
*reinterpret_cast<float4*>(out + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int row = idx / LIN2DIM;
|
||||
int col = idx % LIN2DIM;
|
||||
if (col < HIDDEN)
|
||||
out[idx] = attn[static_cast<long long>(row) * HIDDEN + col];
|
||||
else
|
||||
out[idx] = mlp[static_cast<long long>(row) * FFDIM + (col - HIDDEN)];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void gate_residual_kernel(const half* __restrict__ residual,
|
||||
const half* __restrict__ gate,
|
||||
half* __restrict__ out,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = gridDim.y * seq * HIDDEN;
|
||||
int linear = blockIdx.y * (seq * HIDDEN) + idx;
|
||||
if (linear >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int h = linear % HIDDEN;
|
||||
int token = linear / HIDDEN;
|
||||
int b = token / seq;
|
||||
float gate_v = __half2float(gate[static_cast<long long>(b) * HIDDEN + h]);
|
||||
float res = __half2float(residual[linear]);
|
||||
float v = __half2float(out[linear]);
|
||||
out[linear] = __float2half_rn(res + gate_v * v);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// SSBEngine — FP16-only (no FP8, no CUTLASS)
|
||||
// =========================================================================
|
||||
|
||||
SSBEngine::SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size)
|
||||
: B_(batch),
|
||||
S_(seq),
|
||||
cub_(nullptr),
|
||||
workspace_(nullptr),
|
||||
workspace_bytes_(0),
|
||||
d_silu_(nullptr),
|
||||
d_mod_(nullptr),
|
||||
d_xmod_(nullptr),
|
||||
d_l1o_(nullptr)
|
||||
{
|
||||
BK(cublasCreate(&cub_));
|
||||
BK(cublasSetMathMode(cub_, CUBLAS_TENSOR_OP_MATH));
|
||||
|
||||
size_t min_needed = ssb_workspace_bytes(B_, S_);
|
||||
if (workspace_buf_size < min_needed) {
|
||||
throw std::runtime_error("Workspace buffer too small: need " +
|
||||
std::to_string(min_needed) + " got " + std::to_string(workspace_buf_size));
|
||||
}
|
||||
|
||||
char* p = static_cast<char*>(workspace_buf);
|
||||
size_t off = 0;
|
||||
size_t rows = static_cast<size_t>(B_) * S_;
|
||||
|
||||
#define CARVE(ptr, type, count) \
|
||||
off = (off + 255) & ~size_t(255); \
|
||||
ptr = reinterpret_cast<type*>(p + off); \
|
||||
off += static_cast<size_t>(count) * sizeof(type);
|
||||
|
||||
CARVE(d_silu_, half, B_ * HIDDEN)
|
||||
CARVE(d_mod_, half, static_cast<size_t>(B_) * MODDIM)
|
||||
CARVE(d_xmod_, half, rows * HIDDEN)
|
||||
CARVE(d_l1o_, half, rows * LIN1DIM)
|
||||
|
||||
#undef CARVE
|
||||
|
||||
off = (off + 255) & ~size_t(255);
|
||||
workspace_ = p + off;
|
||||
workspace_bytes_ = workspace_buf_size - off;
|
||||
}
|
||||
|
||||
SSBEngine::~SSBEngine()
|
||||
{
|
||||
if (cub_) cublasDestroy(cub_);
|
||||
}
|
||||
|
||||
void SSBEngine::prepare(const half* d_x,
|
||||
const half* d_vec,
|
||||
const half* d_pe,
|
||||
const half* mod_w,
|
||||
const half* mod_b,
|
||||
const half* l1w,
|
||||
const half* l1b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
half* d_q,
|
||||
half* d_k,
|
||||
half* d_v,
|
||||
half* d_mlp,
|
||||
half* d_gate,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
BK(cublasSetStream(cub_, stream));
|
||||
|
||||
int rows = B_ * S_;
|
||||
int threads = kKernelThreads;
|
||||
|
||||
// SiLU on vec
|
||||
silu_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_vec, d_silu_, B_ * HIDDEN);
|
||||
|
||||
// Modulation GEMM: FP16→FP16
|
||||
linear_gemm(cub_, B_, MODDIM, HIDDEN, d_silu_, mod_w, d_mod_);
|
||||
int out_n = B_ * MODDIM;
|
||||
bias_add_kernel<<<(out_n + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(d_mod_, mod_b, B_, MODDIM);
|
||||
|
||||
// LayerNorm + Modulation (reads x as BF16, outputs FP16)
|
||||
layernorm_mod_kernel<<<rows, threads, 0, stream>>>(reinterpret_cast<const __nv_bfloat16*>(d_x), d_mod_, d_xmod_, S_);
|
||||
|
||||
// Linear1 GEMM: FP16→FP16
|
||||
linear_gemm(cub_, rows, LIN1DIM, HIDDEN, d_xmod_, l1w, d_l1o_);
|
||||
int l1_n = rows * LIN1DIM;
|
||||
bias_add_kernel<<<(l1_n + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(d_l1o_, l1b, rows, LIN1DIM);
|
||||
|
||||
// Split Q/K/V from Linear1 output
|
||||
dim3 split_grid((S_ * HIDDEN + threads - 1) / threads, B_);
|
||||
split_qkv_kernel<<<split_grid, threads, 0, stream>>>(d_l1o_, d_q, d_k, d_v, S_);
|
||||
|
||||
// Copy MLP portion
|
||||
copy_mlp_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_l1o_, d_mlp, rows);
|
||||
|
||||
// Extract gate from modulation output
|
||||
extract_gate_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mod_, d_gate, B_);
|
||||
|
||||
// RMSNorm + scale on Q and K
|
||||
int qk_rows = B_ * HEADS * S_;
|
||||
rmsnorm_scale_kernel<<<qk_rows, threads, 0, stream>>>(d_q, qns);
|
||||
rmsnorm_scale_kernel<<<qk_rows, threads, 0, stream>>>(d_k, kns);
|
||||
|
||||
// RoPE
|
||||
int rope_total = B_ * HEADS * S_ * (HDIM / 2);
|
||||
rope_kernel<<<(rope_total + threads - 1) / threads, threads, 0, stream>>>(d_q, d_k, d_pe, B_, S_);
|
||||
|
||||
CK(cudaGetLastError());
|
||||
}
|
||||
|
||||
void SSBEngine::finish(const half* d_x,
|
||||
const half* d_attn,
|
||||
half* d_mlp,
|
||||
const half* d_gate,
|
||||
half* d_out,
|
||||
const half* l2w,
|
||||
const half* l2b,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
BK(cublasSetStream(cub_, stream));
|
||||
|
||||
int rows = B_ * S_;
|
||||
int threads = kKernelThreads;
|
||||
|
||||
// GELU on MLP
|
||||
gelu_tanh_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mlp, rows * FFDIM);
|
||||
|
||||
// Concat attention + MLP → d_l1o_ (reused as temp)
|
||||
concat_final_kernel<<<(rows * LIN2DIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_attn, d_mlp, d_l1o_, rows);
|
||||
|
||||
// Linear2 GEMM: FP16→FP16
|
||||
linear_gemm(cub_, rows, HIDDEN, LIN2DIM, d_l1o_, l2w, d_out);
|
||||
|
||||
// Bias add
|
||||
int out_blocks = (rows * HIDDEN + threads * 4 - 1) / (threads * 4);
|
||||
bias_add_kernel<<<out_blocks, threads, 0, stream>>>(d_out, l2b, rows, HIDDEN);
|
||||
|
||||
// Gate + residual: out = x + gate * out
|
||||
dim3 gate_grid((S_ * HIDDEN + threads - 1) / threads, B_);
|
||||
gate_residual_kernel<<<gate_grid, threads, 0, stream>>>(d_x, d_gate, d_out, S_);
|
||||
|
||||
CK(cudaGetLastError());
|
||||
}
|
||||
|
||||
} // namespace transformer
|
||||
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
Build script for flux_ext FP16 PyTorch C++/CUDA extension — 20-series (SM 75, Turing).
|
||||
|
||||
Usage:
|
||||
py -3.10 setup_ext.py install
|
||||
|
||||
Compiles the FP16-only FLUXNATION kernel (no FP8, no CUTLASS).
|
||||
Compatible with RTX 2060, 2070, 2080, 2080 Ti.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["CUDA_HOME"] = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8"
|
||||
os.environ["CUDA_PATH"] = os.environ["CUDA_HOME"]
|
||||
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
|
||||
|
||||
ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
setup(
|
||||
name="flux_ext",
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
name="flux_ext",
|
||||
sources=[
|
||||
os.path.join(ROOT, "csrc", "FLUXPATCH_FP16.cpp"),
|
||||
os.path.join(ROOT, "csrc", "FLUXNATION_FP16.cu"),
|
||||
os.path.join(ROOT, "flux_opt_fp16.cu"),
|
||||
],
|
||||
include_dirs=[
|
||||
ROOT, # SINGLE_STREAM_BLOCK_FP16.H
|
||||
],
|
||||
extra_compile_args={
|
||||
"cxx": ["/O2", "/std:c++17", "/EHsc", "/MD"],
|
||||
"nvcc": [
|
||||
"-O3",
|
||||
"-std=c++17",
|
||||
"-arch=sm_75",
|
||||
"-DFLUX_TORCH_EXT=1",
|
||||
"--allow-unsupported-compiler",
|
||||
"-diag-suppress=20013",
|
||||
"-diag-suppress=20015",
|
||||
"-Xcompiler", "/MD",
|
||||
"-Xcompiler", "/EHsc",
|
||||
"-Xcompiler", "/Zc:preprocessor",
|
||||
],
|
||||
},
|
||||
libraries=["cublas"],
|
||||
)
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension},
|
||||
)
|
||||
@@ -0,0 +1,106 @@
|
||||
#ifndef SINGLE_STREAM_BLOCK_FP16_H
|
||||
#define SINGLE_STREAM_BLOCK_FP16_H
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cublas_v2.h>
|
||||
#include <cstdio>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
namespace transformer {
|
||||
|
||||
constexpr int HIDDEN = 3072;
|
||||
constexpr int HEADS = 24;
|
||||
constexpr int HDIM = 128;
|
||||
constexpr int FFDIM = 12288;
|
||||
constexpr int MODDIM = 3 * HIDDEN;
|
||||
constexpr int LIN1DIM = 3 * HIDDEN + FFDIM;
|
||||
constexpr int LIN2DIM = HIDDEN + FFDIM;
|
||||
constexpr float LN_EPS = 1e-6f;
|
||||
constexpr float RMS_EPS = 1e-6f;
|
||||
constexpr float ATT_SC = 0.0883883476f;
|
||||
|
||||
static int const kKernelThreads = 256;
|
||||
static_assert((kKernelThreads % 32) == 0, "kKernelThreads must be a multiple of 32.");
|
||||
|
||||
#define CK(call) do { \
|
||||
cudaError_t err__ = (call); \
|
||||
if (err__ != cudaSuccess) { \
|
||||
throw std::runtime_error(std::string("CUDA ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + cudaGetErrorString(err__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define BK(call) do { \
|
||||
cublasStatus_t st__ = (call); \
|
||||
if (st__ != CUBLAS_STATUS_SUCCESS) { \
|
||||
throw std::runtime_error(std::string("cuBLAS ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + std::to_string((int)st__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static inline size_t ssb_workspace_bytes(int batch, int seq) {
|
||||
size_t rows = static_cast<size_t>(batch) * seq;
|
||||
size_t b = static_cast<size_t>(batch);
|
||||
size_t sz = 0;
|
||||
#define ADDBUF(count, elem_size) \
|
||||
sz = (sz + 255) & ~size_t(255); sz += (count) * (elem_size);
|
||||
ADDBUF(b * HIDDEN, 2); // d_silu_
|
||||
ADDBUF(b * MODDIM, 2); // d_mod_
|
||||
ADDBUF(rows * HIDDEN, 2); // d_xmod_
|
||||
ADDBUF(rows * LIN1DIM, 2); // d_l1o_
|
||||
sz = (sz + 255) & ~size_t(255);
|
||||
sz += 32ull << 20; // 32MB cuBLAS workspace
|
||||
#undef ADDBUF
|
||||
return sz;
|
||||
}
|
||||
|
||||
class SSBEngine {
|
||||
public:
|
||||
SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size);
|
||||
~SSBEngine();
|
||||
|
||||
void prepare(const half* d_x,
|
||||
const half* d_vec,
|
||||
const half* d_pe,
|
||||
const half* mod_w,
|
||||
const half* mod_b,
|
||||
const half* l1w,
|
||||
const half* l1b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
half* d_q,
|
||||
half* d_k,
|
||||
half* d_v,
|
||||
half* d_mlp,
|
||||
half* d_gate,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
void finish(const half* d_x,
|
||||
const half* d_attn,
|
||||
half* d_mlp,
|
||||
const half* d_gate,
|
||||
half* d_out,
|
||||
const half* l2w,
|
||||
const half* l2b,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
int batch() const { return B_; }
|
||||
int seq() const { return S_; }
|
||||
|
||||
private:
|
||||
int B_;
|
||||
int S_;
|
||||
cublasHandle_t cub_;
|
||||
void* workspace_;
|
||||
size_t workspace_bytes_;
|
||||
|
||||
half* d_silu_;
|
||||
half* d_mod_;
|
||||
half* d_xmod_;
|
||||
half* d_l1o_;
|
||||
};
|
||||
|
||||
} // namespace transformer
|
||||
|
||||
#endif // SINGLE_STREAM_BLOCK_FP16_H
|
||||
@@ -0,0 +1,25 @@
|
||||
@echo off
|
||||
setlocal
|
||||
|
||||
REM FLUXNATION FP16 Build — 30-series (SM 86, Ampere)
|
||||
REM Compatible with RTX 3060/3070/3080/3090
|
||||
|
||||
set "VSVARS=C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat"
|
||||
if not exist "%VSVARS%" (
|
||||
echo [flux_ext] ERROR: VS 2022 vcvars64.bat not found
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
call "%VSVARS%" || exit /b 1
|
||||
|
||||
set DISTUTILS_USE_SDK=1
|
||||
|
||||
echo [flux_ext] Building FLUXNATION FP16 for 30-series (SM 86)...
|
||||
py -3.10 setup_ext.py install
|
||||
if errorlevel 1 (
|
||||
echo [flux_ext] Build FAILED
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
echo [flux_ext] Build succeeded! flux_ext (FP16, SM 86) is now installed.
|
||||
echo [flux_ext] Restart ComfyUI to use the new kernel.
|
||||
@@ -0,0 +1,198 @@
|
||||
/*
|
||||
* FLUXNATION_FP16.cu — PyTorch C++ extension for FLUXNATION (FP16-only)
|
||||
*
|
||||
* Wraps the FP16 SSBEngine as torch.ops custom operators.
|
||||
* All weights converted to FP16 internally — no FP8, no CUTLASS.
|
||||
* Compatible with SM 61+ (Pascal, Turing, Ampere).
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#include "SINGLE_STREAM_BLOCK_FP16.H"
|
||||
|
||||
#include <unordered_map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
|
||||
using namespace transformer;
|
||||
|
||||
// =========================================================================
|
||||
// Global engine cache — keyed by (device, batch, seq_len)
|
||||
// =========================================================================
|
||||
struct EngineEntry {
|
||||
std::unique_ptr<SSBEngine> engine;
|
||||
at::Tensor workspace;
|
||||
};
|
||||
|
||||
static std::unordered_map<int64_t, EngineEntry> g_engine_cache;
|
||||
static std::mutex g_engine_mutex;
|
||||
|
||||
static int64_t make_engine_key(int device, int batch, int seq) {
|
||||
return (int64_t)device * 1000000000LL + (int64_t)batch * 1000000LL + seq;
|
||||
}
|
||||
|
||||
static SSBEngine* get_or_create_engine(int batch, int seq, at::Device device) {
|
||||
int dev_idx = device.index();
|
||||
int64_t key = make_engine_key(dev_idx, batch, seq);
|
||||
|
||||
std::lock_guard<std::mutex> lock(g_engine_mutex);
|
||||
auto it = g_engine_cache.find(key);
|
||||
if (it != g_engine_cache.end()) {
|
||||
return it->second.engine.get();
|
||||
}
|
||||
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
size_t ws_bytes = ssb_workspace_bytes(batch, seq);
|
||||
auto ws_tensor = at::empty({(int64_t)ws_bytes}, at::TensorOptions().dtype(at::kByte).device(device));
|
||||
|
||||
auto engine = std::make_unique<SSBEngine>(batch, seq, ws_tensor.data_ptr(), ws_bytes);
|
||||
auto* ptr = engine.get();
|
||||
|
||||
EngineEntry entry;
|
||||
entry.engine = std::move(engine);
|
||||
entry.workspace = ws_tensor;
|
||||
g_engine_cache[key] = std::move(entry);
|
||||
|
||||
fprintf(stderr, "[flux_ext FP16] Created engine B=%d S=%d ws=%.1fMB dev=%d\n",
|
||||
batch, seq, ws_bytes / (1024.0 * 1024.0), dev_idx);
|
||||
fflush(stderr);
|
||||
|
||||
return ptr;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// torch.ops implementation: ssb_prepare
|
||||
//
|
||||
// Same Python API as FP8 version — scale tensors are accepted but ignored.
|
||||
// All weights converted to FP16 internally.
|
||||
// =========================================================================
|
||||
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
|
||||
ssb_prepare_impl(
|
||||
at::Tensor x, // [B, S, 3072] bf16
|
||||
at::Tensor vec, // [B, 3072] or [B, 1, 3072] fp16
|
||||
at::Tensor pe, // [S, 64, 2] fp16
|
||||
at::Tensor mod_w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor mod_w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor mod_b, // fp16
|
||||
at::Tensor l1w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor l1w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor l1b, // fp16
|
||||
at::Tensor qns, // fp16
|
||||
at::Tensor kns) // fp16
|
||||
{
|
||||
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
|
||||
TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");
|
||||
|
||||
int B = x.size(0);
|
||||
int S = x.size(1);
|
||||
auto device = x.device();
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(device.index());
|
||||
|
||||
// Ensure correct dtypes — x stays BF16 (layernorm kernel reads it as BF16)
|
||||
auto x_bf16 = x.to(at::kBFloat16).contiguous();
|
||||
auto vec_fp16 = vec.dim() == 3 ? vec.squeeze(1).to(at::kHalf).contiguous()
|
||||
: vec.to(at::kHalf).contiguous();
|
||||
auto pe_fp16 = pe.to(at::kHalf).contiguous();
|
||||
|
||||
// Convert ALL weights to FP16 (no FP8 on this target)
|
||||
auto mod_w_fp16 = mod_w.to(at::kHalf).contiguous();
|
||||
auto mod_b_fp16 = mod_b.to(at::kHalf).contiguous();
|
||||
auto l1w_fp16 = l1w.to(at::kHalf).contiguous();
|
||||
auto l1b_fp16 = l1b.to(at::kHalf).contiguous();
|
||||
auto qns_fp16 = qns.to(at::kHalf).contiguous();
|
||||
auto kns_fp16 = kns.to(at::kHalf).contiguous();
|
||||
|
||||
// Allocate output tensors
|
||||
auto opts_fp16 = at::TensorOptions().dtype(at::kHalf).device(device);
|
||||
auto q = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto k = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto v = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto mlp = at::empty({B, S, FFDIM}, opts_fp16);
|
||||
auto gate = at::empty({B, HIDDEN}, opts_fp16);
|
||||
|
||||
SSBEngine* engine = get_or_create_engine(B, S, device);
|
||||
|
||||
engine->prepare(
|
||||
reinterpret_cast<const half*>(x_bf16.data_ptr()),
|
||||
reinterpret_cast<const half*>(vec_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(pe_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(mod_w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(mod_b_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l1w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l1b_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(qns_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(kns_fp16.data_ptr()),
|
||||
reinterpret_cast<half*>(q.data_ptr()),
|
||||
reinterpret_cast<half*>(k.data_ptr()),
|
||||
reinterpret_cast<half*>(v.data_ptr()),
|
||||
reinterpret_cast<half*>(mlp.data_ptr()),
|
||||
reinterpret_cast<half*>(gate.data_ptr()),
|
||||
stream);
|
||||
|
||||
return std::make_tuple(q, k, v, mlp, gate);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// torch.ops implementation: ssb_finish
|
||||
//
|
||||
// FP16 path: converts residual (x) from BF16→FP16, runs all-FP16 engine,
|
||||
// converts result back to BF16 for model compatibility.
|
||||
// =========================================================================
|
||||
at::Tensor ssb_finish_impl(
|
||||
at::Tensor x, // [B, S, 3072] bf16 — original input for residual
|
||||
at::Tensor attn, // [B, S, 3072] fp16 — attention output
|
||||
at::Tensor mlp, // [B, S, 12288] fp16 — MLP intermediate (pre-GELU)
|
||||
at::Tensor gate, // [B, 3072] fp16 — gate values
|
||||
at::Tensor l2w, // weight (any dtype → converted to FP16)
|
||||
at::Tensor l2w_scale, // IGNORED in FP16 mode
|
||||
at::Tensor l2b) // fp16 — L2 bias
|
||||
{
|
||||
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
|
||||
TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");
|
||||
|
||||
int B = x.size(0);
|
||||
int S = x.size(1);
|
||||
auto device = x.device();
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(device.index());
|
||||
|
||||
// Convert residual from BF16 to FP16 for the FP16 engine
|
||||
auto x_fp16 = x.to(at::kHalf).contiguous();
|
||||
auto attn_fp16 = attn.to(at::kHalf).contiguous();
|
||||
auto mlp_c = mlp.contiguous();
|
||||
auto gate_c = gate.contiguous();
|
||||
auto l2w_fp16 = l2w.to(at::kHalf).contiguous();
|
||||
auto l2b_fp16 = l2b.to(at::kHalf).contiguous();
|
||||
|
||||
// Result in FP16 (engine is all-FP16)
|
||||
auto opts_fp16 = at::TensorOptions().dtype(at::kHalf).device(device);
|
||||
auto result_fp16 = at::empty({B, S, HIDDEN}, opts_fp16);
|
||||
|
||||
SSBEngine* engine = get_or_create_engine(B, S, device);
|
||||
|
||||
engine->finish(
|
||||
reinterpret_cast<const half*>(x_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(attn_fp16.data_ptr()),
|
||||
reinterpret_cast<half*>(mlp_c.data_ptr()),
|
||||
reinterpret_cast<const half*>(gate_c.data_ptr()),
|
||||
reinterpret_cast<half*>(result_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l2w_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(l2b_fp16.data_ptr()),
|
||||
stream);
|
||||
|
||||
// Convert back to BF16 for model compatibility
|
||||
return result_fp16.to(at::kBFloat16);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Cleanup — destroys all cached engines
|
||||
// =========================================================================
|
||||
void ssb_cleanup_impl() {
|
||||
std::lock_guard<std::mutex> lock(g_engine_mutex);
|
||||
g_engine_cache.clear();
|
||||
fprintf(stderr, "[flux_ext FP16] All engines destroyed\n");
|
||||
fflush(stderr);
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
/*
|
||||
* 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)";
|
||||
}
|
||||
@@ -0,0 +1,611 @@
|
||||
/*
|
||||
* flux_opt_fp16.cu — FP16-only SSBEngine for FLUXNATION
|
||||
*
|
||||
* Derived from flux_opt.cu. All FP8/CUTLASS code removed.
|
||||
* Uses cuBLAS FP16 tensor-op GEMMs only.
|
||||
* Compatible with SM 61+ (Pascal, Turing, Ampere).
|
||||
*/
|
||||
|
||||
#include "SINGLE_STREAM_BLOCK_FP16.H"
|
||||
#include <algorithm>
|
||||
#include <cfloat>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
|
||||
namespace transformer {
|
||||
|
||||
// Vectorized load/store helpers — 128-bit (8 half elements) for max bandwidth
|
||||
union VecHalf8 {
|
||||
float4 f4;
|
||||
half h[8];
|
||||
};
|
||||
union VecHalf4 {
|
||||
float2 f2;
|
||||
half h[4];
|
||||
};
|
||||
union VecHalf2 {
|
||||
unsigned int u;
|
||||
half h[2];
|
||||
};
|
||||
union VecBF16_4 {
|
||||
float2 f2;
|
||||
__nv_bfloat16 b[4];
|
||||
};
|
||||
union VecBF16_8 {
|
||||
float4 f4;
|
||||
__nv_bfloat16 b[8];
|
||||
};
|
||||
|
||||
// FP16 GEMM via cuBLAS with tensor-op math
|
||||
static inline void linear_gemm(cublasHandle_t cub,
|
||||
int m,
|
||||
int n,
|
||||
int k,
|
||||
const half* x,
|
||||
const half* w_torch,
|
||||
half* y)
|
||||
{
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
BK(cublasGemmEx(
|
||||
cub,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
n, m, k,
|
||||
&alpha,
|
||||
w_torch, CUDA_R_16F, k,
|
||||
x, CUDA_R_16F, k,
|
||||
&beta,
|
||||
y, CUDA_R_16F, n,
|
||||
CUBLAS_COMPUTE_32F_FAST_16F,
|
||||
CUBLAS_GEMM_DEFAULT_TENSOR_OP));
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float warp_reduce_sum(float v)
|
||||
{
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
v += __shfl_down_sync(0xffffffff, v, offset);
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float warp_reduce_max(float v)
|
||||
{
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
v = fmaxf(v, __shfl_down_sync(0xffffffff, v, offset));
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float block_reduce_sum(float v)
|
||||
{
|
||||
constexpr int kWarps = kKernelThreads / 32;
|
||||
__shared__ float warp_vals[kWarps];
|
||||
int lane = threadIdx.x & 31;
|
||||
int wid = threadIdx.x >> 5;
|
||||
|
||||
v = warp_reduce_sum(v);
|
||||
if (lane == 0) {
|
||||
warp_vals[wid] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (wid == 0) {
|
||||
float x = (lane < kWarps) ? warp_vals[lane] : 0.0f;
|
||||
x = warp_reduce_sum(x);
|
||||
if (lane == 0) {
|
||||
warp_vals[0] = x;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return warp_vals[0];
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float block_reduce_max(float v)
|
||||
{
|
||||
constexpr int kWarps = kKernelThreads / 32;
|
||||
__shared__ float warp_vals[kWarps];
|
||||
int lane = threadIdx.x & 31;
|
||||
int wid = threadIdx.x >> 5;
|
||||
|
||||
v = warp_reduce_max(v);
|
||||
if (lane == 0) {
|
||||
warp_vals[wid] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (wid == 0) {
|
||||
float x = (lane < kWarps) ? warp_vals[lane] : -FLT_MAX;
|
||||
x = warp_reduce_max(x);
|
||||
if (lane == 0) {
|
||||
warp_vals[0] = x;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return warp_vals[0];
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void silu_kernel(const half* __restrict__ in,
|
||||
half* __restrict__ out,
|
||||
int n)
|
||||
{
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < n) {
|
||||
VecHalf8 v;
|
||||
v.f4 = *reinterpret_cast<const float4*>(in + base);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
float x = __half2float(v.h[j]);
|
||||
v.h[j] = __float2half_rn(x / (1.0f + expf(-x)));
|
||||
}
|
||||
*reinterpret_cast<float4*>(out + base) = v.f4;
|
||||
} else {
|
||||
for (int i = base; i < min(base + 8, n); i++) {
|
||||
float x = __half2float(in[i]);
|
||||
out[i] = __float2half_rn(x / (1.0f + expf(-x)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void bias_add_kernel(half* __restrict__ out,
|
||||
const half* __restrict__ bias,
|
||||
int rows,
|
||||
int cols)
|
||||
{
|
||||
int total = rows * cols;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
|
||||
if (base + 3 < total) {
|
||||
int col = base % cols;
|
||||
if (col + 3 < cols) {
|
||||
VecHalf4 ov, bv;
|
||||
ov.f2 = *reinterpret_cast<const float2*>(out + base);
|
||||
bv.f2 = *reinterpret_cast<const float2*>(bias + col);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++)
|
||||
ov.h[j] = __float2half_rn(__half2float(ov.h[j]) + __half2float(bv.h[j]));
|
||||
*reinterpret_cast<float2*>(out + base) = ov.f2;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 4 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int col = idx % cols;
|
||||
float v = __half2float(out[idx]) + __half2float(bias[col]);
|
||||
out[idx] = __float2half_rn(v);
|
||||
}
|
||||
}
|
||||
|
||||
// Input x is bf16 (model runs in bf16). Output is fp16. BF16→FP16 conversion is
|
||||
// software-emulated on SM < 80 but is IO-bound so negligible overhead.
|
||||
__global__ __launch_bounds__(256, 2)
|
||||
void layernorm_mod_kernel(const __nv_bfloat16* __restrict__ x,
|
||||
const half* __restrict__ mod,
|
||||
half* __restrict__ out,
|
||||
int seq)
|
||||
{
|
||||
int row = blockIdx.x;
|
||||
int b = row / seq;
|
||||
|
||||
const __nv_bfloat16* x_row = x + static_cast<long long>(row) * HIDDEN;
|
||||
half* out_row = out + static_cast<long long>(row) * HIDDEN;
|
||||
const half* shift = mod + static_cast<long long>(b) * MODDIM;
|
||||
const half* scale = shift + HIDDEN;
|
||||
|
||||
float sum = 0.0f;
|
||||
float sq = 0.0f;
|
||||
for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
|
||||
VecBF16_4 bv;
|
||||
bv.f2 = *reinterpret_cast<const float2*>(x_row + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float v = __bfloat162float(bv.b[j]);
|
||||
sum += v;
|
||||
sq += v * v;
|
||||
}
|
||||
}
|
||||
float mean = block_reduce_sum(sum) / HIDDEN;
|
||||
float var = block_reduce_sum(sq) / HIDDEN - mean * mean;
|
||||
float inv = rsqrtf(var + LN_EPS);
|
||||
|
||||
for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
|
||||
VecBF16_4 xv;
|
||||
xv.f2 = *reinterpret_cast<const float2*>(x_row + i);
|
||||
VecHalf4 sc_v, sh_v, ov;
|
||||
sc_v.f2 = *reinterpret_cast<const float2*>(scale + i);
|
||||
sh_v.f2 = *reinterpret_cast<const float2*>(shift + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float v = __bfloat162float(xv.b[j]);
|
||||
float nrm = (v - mean) * inv;
|
||||
ov.h[j] = __float2half_rn(nrm * (1.0f + __half2float(sc_v.h[j])) + __half2float(sh_v.h[j]));
|
||||
}
|
||||
*reinterpret_cast<float2*>(out_row + i) = ov.f2;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void split_qkv_kernel(const half* __restrict__ l1o,
|
||||
half* __restrict__ q,
|
||||
half* __restrict__ k,
|
||||
half* __restrict__ v,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = seq * HIDDEN;
|
||||
int batch_total = gridDim.y;
|
||||
int linear = blockIdx.y * total + idx;
|
||||
int max_linear = batch_total * total;
|
||||
if (linear >= max_linear) return;
|
||||
|
||||
int token = linear / HIDDEN;
|
||||
int h = linear % HIDDEN;
|
||||
int b = token / seq;
|
||||
int s = token % seq;
|
||||
int head = h / HDIM;
|
||||
int d = h % HDIM;
|
||||
|
||||
long long base_l1 = static_cast<long long>(token) * LIN1DIM;
|
||||
long long dst = (((static_cast<long long>(b) * HEADS + head) * seq) + s) * HDIM + d;
|
||||
|
||||
q[dst] = l1o[base_l1 + h];
|
||||
k[dst] = l1o[base_l1 + HIDDEN + h];
|
||||
v[dst] = l1o[base_l1 + 2 * HIDDEN + h];
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void copy_mlp_kernel(const half* __restrict__ l1o,
|
||||
half* __restrict__ mlp,
|
||||
int rows)
|
||||
{
|
||||
int total = rows * FFDIM;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int row = base / FFDIM;
|
||||
int col = base % FFDIM;
|
||||
if (col + 7 < FFDIM) {
|
||||
float4 v = *reinterpret_cast<const float4*>(l1o + static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col);
|
||||
*reinterpret_cast<float4*>(mlp + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int row = idx / FFDIM;
|
||||
int col = idx % FFDIM;
|
||||
mlp[idx] = l1o[static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void extract_gate_kernel(const half* __restrict__ mod,
|
||||
half* __restrict__ gate,
|
||||
int batch)
|
||||
{
|
||||
int total = batch * HIDDEN;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int b = base / HIDDEN;
|
||||
int h = base % HIDDEN;
|
||||
if (h + 7 < HIDDEN) {
|
||||
float4 v = *reinterpret_cast<const float4*>(mod + static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h);
|
||||
*reinterpret_cast<float4*>(gate + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int b = idx / HIDDEN;
|
||||
int h = idx % HIDDEN;
|
||||
gate[idx] = mod[static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 2)
|
||||
void rmsnorm_scale_kernel(half* __restrict__ data,
|
||||
const half* __restrict__ scale)
|
||||
{
|
||||
int row = blockIdx.x;
|
||||
half* row_ptr = data + static_cast<long long>(row) * HDIM;
|
||||
|
||||
float sq = 0.0f;
|
||||
for (int i = threadIdx.x * 4; i < HDIM; i += blockDim.x * 4) {
|
||||
VecHalf4 v;
|
||||
v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++) {
|
||||
float val = __half2float(v.h[j]);
|
||||
sq += val * val;
|
||||
}
|
||||
}
|
||||
float inv = rsqrtf(block_reduce_sum(sq) / HDIM + RMS_EPS);
|
||||
|
||||
for (int i = threadIdx.x * 4; i < HDIM; i += blockDim.x * 4) {
|
||||
VecHalf4 v, sc;
|
||||
v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
|
||||
sc.f2 = *reinterpret_cast<const float2*>(scale + i);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; j++)
|
||||
v.h[j] = __float2half_rn(__half2float(v.h[j]) * inv * __half2float(sc.h[j]));
|
||||
*reinterpret_cast<float2*>(row_ptr + i) = v.f2;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void rope_kernel(half* __restrict__ q,
|
||||
half* __restrict__ k,
|
||||
const half* __restrict__ pe,
|
||||
int batch,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = batch * HEADS * seq * (HDIM / 2);
|
||||
if (idx >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int pair = idx % (HDIM / 2);
|
||||
int t0 = idx / (HDIM / 2);
|
||||
int s = t0 % seq;
|
||||
int t1 = t0 / seq;
|
||||
int h = t1 % HEADS;
|
||||
int b = t1 / HEADS;
|
||||
|
||||
long long base = (((static_cast<long long>(b) * HEADS + h) * seq) + s) * HDIM + 2 * pair;
|
||||
long long pe_off = (static_cast<long long>(s) * (HDIM / 2) + pair) * 2;
|
||||
|
||||
VecHalf2 pe_v;
|
||||
pe_v.u = *reinterpret_cast<const unsigned int*>(pe + pe_off);
|
||||
float c = __half2float(pe_v.h[0]);
|
||||
float sn = __half2float(pe_v.h[1]);
|
||||
|
||||
VecHalf2 qv;
|
||||
qv.u = *reinterpret_cast<const unsigned int*>(q + base);
|
||||
float q0 = __half2float(qv.h[0]);
|
||||
float q1 = __half2float(qv.h[1]);
|
||||
qv.h[0] = __float2half_rn(q0 * c - q1 * sn);
|
||||
qv.h[1] = __float2half_rn(q0 * sn + q1 * c);
|
||||
*reinterpret_cast<unsigned int*>(q + base) = qv.u;
|
||||
|
||||
VecHalf2 kv;
|
||||
kv.u = *reinterpret_cast<const unsigned int*>(k + base);
|
||||
float k0 = __half2float(kv.h[0]);
|
||||
float k1 = __half2float(kv.h[1]);
|
||||
kv.h[0] = __float2half_rn(k0 * c - k1 * sn);
|
||||
kv.h[1] = __float2half_rn(k0 * sn + k1 * c);
|
||||
*reinterpret_cast<unsigned int*>(k + base) = kv.u;
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 4)
|
||||
void gelu_tanh_kernel(half* __restrict__ data,
|
||||
int n)
|
||||
{
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < n) {
|
||||
VecHalf8 v;
|
||||
v.f4 = *reinterpret_cast<const float4*>(data + base);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
float x = __half2float(v.h[j]);
|
||||
float x3 = x * x * x;
|
||||
v.h[j] = __float2half_rn(0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x3))));
|
||||
}
|
||||
*reinterpret_cast<float4*>(data + base) = v.f4;
|
||||
} else {
|
||||
for (int i = base; i < min(base + 8, n); i++) {
|
||||
float x = __half2float(data[i]);
|
||||
float x3 = x * x * x;
|
||||
data[i] = __float2half_rn(0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x3))));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(256, 6)
|
||||
void concat_final_kernel(const half* __restrict__ attn,
|
||||
const half* __restrict__ mlp,
|
||||
half* __restrict__ out,
|
||||
int rows)
|
||||
{
|
||||
int total = rows * LIN2DIM;
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base + 7 < total) {
|
||||
int row = base / LIN2DIM;
|
||||
int col = base % LIN2DIM;
|
||||
if (col + 7 < LIN2DIM) {
|
||||
float4 v;
|
||||
if (col + 7 < HIDDEN) {
|
||||
v = *reinterpret_cast<const float4*>(attn + static_cast<long long>(row) * HIDDEN + col);
|
||||
} else if (col >= HIDDEN) {
|
||||
v = *reinterpret_cast<const float4*>(mlp + static_cast<long long>(row) * FFDIM + (col - HIDDEN));
|
||||
} else {
|
||||
VecHalf8 vh;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j++) {
|
||||
int c = col + j;
|
||||
if (c < HIDDEN)
|
||||
vh.h[j] = attn[static_cast<long long>(row) * HIDDEN + c];
|
||||
else
|
||||
vh.h[j] = mlp[static_cast<long long>(row) * FFDIM + (c - HIDDEN)];
|
||||
}
|
||||
v = vh.f4;
|
||||
}
|
||||
*reinterpret_cast<float4*>(out + base) = v;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 8 && base + j < total; j++) {
|
||||
int idx = base + j;
|
||||
int row = idx / LIN2DIM;
|
||||
int col = idx % LIN2DIM;
|
||||
if (col < HIDDEN)
|
||||
out[idx] = attn[static_cast<long long>(row) * HIDDEN + col];
|
||||
else
|
||||
out[idx] = mlp[static_cast<long long>(row) * FFDIM + (col - HIDDEN)];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void gate_residual_kernel(const half* __restrict__ residual,
|
||||
const half* __restrict__ gate,
|
||||
half* __restrict__ out,
|
||||
int seq)
|
||||
{
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = gridDim.y * seq * HIDDEN;
|
||||
int linear = blockIdx.y * (seq * HIDDEN) + idx;
|
||||
if (linear >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int h = linear % HIDDEN;
|
||||
int token = linear / HIDDEN;
|
||||
int b = token / seq;
|
||||
float gate_v = __half2float(gate[static_cast<long long>(b) * HIDDEN + h]);
|
||||
float res = __half2float(residual[linear]);
|
||||
float v = __half2float(out[linear]);
|
||||
out[linear] = __float2half_rn(res + gate_v * v);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// SSBEngine — FP16-only (no FP8, no CUTLASS)
|
||||
// =========================================================================
|
||||
|
||||
SSBEngine::SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size)
|
||||
: B_(batch),
|
||||
S_(seq),
|
||||
cub_(nullptr),
|
||||
workspace_(nullptr),
|
||||
workspace_bytes_(0),
|
||||
d_silu_(nullptr),
|
||||
d_mod_(nullptr),
|
||||
d_xmod_(nullptr),
|
||||
d_l1o_(nullptr)
|
||||
{
|
||||
BK(cublasCreate(&cub_));
|
||||
BK(cublasSetMathMode(cub_, CUBLAS_TENSOR_OP_MATH));
|
||||
|
||||
size_t min_needed = ssb_workspace_bytes(B_, S_);
|
||||
if (workspace_buf_size < min_needed) {
|
||||
throw std::runtime_error("Workspace buffer too small: need " +
|
||||
std::to_string(min_needed) + " got " + std::to_string(workspace_buf_size));
|
||||
}
|
||||
|
||||
char* p = static_cast<char*>(workspace_buf);
|
||||
size_t off = 0;
|
||||
size_t rows = static_cast<size_t>(B_) * S_;
|
||||
|
||||
#define CARVE(ptr, type, count) \
|
||||
off = (off + 255) & ~size_t(255); \
|
||||
ptr = reinterpret_cast<type*>(p + off); \
|
||||
off += static_cast<size_t>(count) * sizeof(type);
|
||||
|
||||
CARVE(d_silu_, half, B_ * HIDDEN)
|
||||
CARVE(d_mod_, half, static_cast<size_t>(B_) * MODDIM)
|
||||
CARVE(d_xmod_, half, rows * HIDDEN)
|
||||
CARVE(d_l1o_, half, rows * LIN1DIM)
|
||||
|
||||
#undef CARVE
|
||||
|
||||
off = (off + 255) & ~size_t(255);
|
||||
workspace_ = p + off;
|
||||
workspace_bytes_ = workspace_buf_size - off;
|
||||
}
|
||||
|
||||
SSBEngine::~SSBEngine()
|
||||
{
|
||||
if (cub_) cublasDestroy(cub_);
|
||||
}
|
||||
|
||||
void SSBEngine::prepare(const half* d_x,
|
||||
const half* d_vec,
|
||||
const half* d_pe,
|
||||
const half* mod_w,
|
||||
const half* mod_b,
|
||||
const half* l1w,
|
||||
const half* l1b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
half* d_q,
|
||||
half* d_k,
|
||||
half* d_v,
|
||||
half* d_mlp,
|
||||
half* d_gate,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
BK(cublasSetStream(cub_, stream));
|
||||
|
||||
int rows = B_ * S_;
|
||||
int threads = kKernelThreads;
|
||||
|
||||
// SiLU on vec
|
||||
silu_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_vec, d_silu_, B_ * HIDDEN);
|
||||
|
||||
// Modulation GEMM: FP16→FP16
|
||||
linear_gemm(cub_, B_, MODDIM, HIDDEN, d_silu_, mod_w, d_mod_);
|
||||
int out_n = B_ * MODDIM;
|
||||
bias_add_kernel<<<(out_n + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(d_mod_, mod_b, B_, MODDIM);
|
||||
|
||||
// LayerNorm + Modulation (reads x as BF16, outputs FP16)
|
||||
layernorm_mod_kernel<<<rows, threads, 0, stream>>>(reinterpret_cast<const __nv_bfloat16*>(d_x), d_mod_, d_xmod_, S_);
|
||||
|
||||
// Linear1 GEMM: FP16→FP16
|
||||
linear_gemm(cub_, rows, LIN1DIM, HIDDEN, d_xmod_, l1w, d_l1o_);
|
||||
int l1_n = rows * LIN1DIM;
|
||||
bias_add_kernel<<<(l1_n + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(d_l1o_, l1b, rows, LIN1DIM);
|
||||
|
||||
// Split Q/K/V from Linear1 output
|
||||
dim3 split_grid((S_ * HIDDEN + threads - 1) / threads, B_);
|
||||
split_qkv_kernel<<<split_grid, threads, 0, stream>>>(d_l1o_, d_q, d_k, d_v, S_);
|
||||
|
||||
// Copy MLP portion
|
||||
copy_mlp_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_l1o_, d_mlp, rows);
|
||||
|
||||
// Extract gate from modulation output
|
||||
extract_gate_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mod_, d_gate, B_);
|
||||
|
||||
// RMSNorm + scale on Q and K
|
||||
int qk_rows = B_ * HEADS * S_;
|
||||
rmsnorm_scale_kernel<<<qk_rows, threads, 0, stream>>>(d_q, qns);
|
||||
rmsnorm_scale_kernel<<<qk_rows, threads, 0, stream>>>(d_k, kns);
|
||||
|
||||
// RoPE
|
||||
int rope_total = B_ * HEADS * S_ * (HDIM / 2);
|
||||
rope_kernel<<<(rope_total + threads - 1) / threads, threads, 0, stream>>>(d_q, d_k, d_pe, B_, S_);
|
||||
|
||||
CK(cudaGetLastError());
|
||||
}
|
||||
|
||||
void SSBEngine::finish(const half* d_x,
|
||||
const half* d_attn,
|
||||
half* d_mlp,
|
||||
const half* d_gate,
|
||||
half* d_out,
|
||||
const half* l2w,
|
||||
const half* l2b,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
BK(cublasSetStream(cub_, stream));
|
||||
|
||||
int rows = B_ * S_;
|
||||
int threads = kKernelThreads;
|
||||
|
||||
// GELU on MLP
|
||||
gelu_tanh_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mlp, rows * FFDIM);
|
||||
|
||||
// Concat attention + MLP → d_l1o_ (reused as temp)
|
||||
concat_final_kernel<<<(rows * LIN2DIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_attn, d_mlp, d_l1o_, rows);
|
||||
|
||||
// Linear2 GEMM: FP16→FP16
|
||||
linear_gemm(cub_, rows, HIDDEN, LIN2DIM, d_l1o_, l2w, d_out);
|
||||
|
||||
// Bias add
|
||||
int out_blocks = (rows * HIDDEN + threads * 4 - 1) / (threads * 4);
|
||||
bias_add_kernel<<<out_blocks, threads, 0, stream>>>(d_out, l2b, rows, HIDDEN);
|
||||
|
||||
// Gate + residual: out = x + gate * out
|
||||
dim3 gate_grid((S_ * HIDDEN + threads - 1) / threads, B_);
|
||||
gate_residual_kernel<<<gate_grid, threads, 0, stream>>>(d_x, d_gate, d_out, S_);
|
||||
|
||||
CK(cudaGetLastError());
|
||||
}
|
||||
|
||||
} // namespace transformer
|
||||
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
Build script for flux_ext FP16 PyTorch C++/CUDA extension — 30-series (SM 86, Ampere).
|
||||
|
||||
Usage:
|
||||
py -3.10 setup_ext.py install
|
||||
|
||||
Compiles the FP16-only FLUXNATION kernel (no FP8, no CUTLASS).
|
||||
Compatible with RTX 3060, 3070, 3080, 3090.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["CUDA_HOME"] = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8"
|
||||
os.environ["CUDA_PATH"] = os.environ["CUDA_HOME"]
|
||||
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
|
||||
|
||||
ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
setup(
|
||||
name="flux_ext",
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
name="flux_ext",
|
||||
sources=[
|
||||
os.path.join(ROOT, "csrc", "FLUXPATCH_FP16.cpp"),
|
||||
os.path.join(ROOT, "csrc", "FLUXNATION_FP16.cu"),
|
||||
os.path.join(ROOT, "flux_opt_fp16.cu"),
|
||||
],
|
||||
include_dirs=[
|
||||
ROOT, # SINGLE_STREAM_BLOCK_FP16.H
|
||||
],
|
||||
extra_compile_args={
|
||||
"cxx": ["/O2", "/std:c++17", "/EHsc", "/MD"],
|
||||
"nvcc": [
|
||||
"-O3",
|
||||
"-std=c++17",
|
||||
"-arch=sm_86",
|
||||
"-DFLUX_TORCH_EXT=1",
|
||||
"--allow-unsupported-compiler",
|
||||
"-diag-suppress=20013",
|
||||
"-diag-suppress=20015",
|
||||
"-Xcompiler", "/MD",
|
||||
"-Xcompiler", "/EHsc",
|
||||
"-Xcompiler", "/Zc:preprocessor",
|
||||
],
|
||||
},
|
||||
libraries=["cublas"],
|
||||
)
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension},
|
||||
)
|
||||
+505
@@ -0,0 +1,505 @@
|
||||
"""
|
||||
Monkey-patch ComfyUI's FLUX SingleStreamBlock to use FLUXNATION via torch.ops.
|
||||
|
||||
Architecture: torch.ops.flux_ext.ssb_prepare + torch.ops.flux_ext.ssb_finish
|
||||
- Zero ctypes overhead (direct tensor passing)
|
||||
- Fused gate+residual+bias in CUDA (no Python-side float32 intermediates)
|
||||
- Weight caching (extracted once per module, not every forward call)
|
||||
- Compatible with torch.compile (via register_fake)
|
||||
|
||||
Spike Attention Integration:
|
||||
- After switch_pct% of diffusion steps, routes attention through neuromorphic
|
||||
spike backend (block-sparse Triton kernel) instead of dense flash attention
|
||||
- Saves ~90% attention FLOPs with rollback safety
|
||||
- Enable via SPIKE_CONFIG dict or environment variable FLUX_SPIKE=1
|
||||
"""
|
||||
|
||||
import math
|
||||
import os
|
||||
import torch
|
||||
from comfy.ldm.flux.math import attention as flux_attention
|
||||
|
||||
# ── Spike attention import (LAZY — loaded on first use, not at startup) ──
|
||||
_SPIKE_AVAILABLE = None # None = not checked yet, True/False = result
|
||||
_spike_attend = None
|
||||
_spike_has_triton = None
|
||||
|
||||
|
||||
def _lazy_load_spike():
|
||||
"""Load spike attention modules on first use."""
|
||||
global _SPIKE_AVAILABLE, _spike_attend, _spike_has_triton
|
||||
if _SPIKE_AVAILABLE is not None:
|
||||
return _SPIKE_AVAILABLE
|
||||
from . import FLUXATTENHUT as _FLUXATTENHUT_mod
|
||||
from . import FLUXSPIKEY as _FLUXSPIKEY_mod
|
||||
_spike_attend = _FLUXATTENHUT_mod.attend
|
||||
_spike_has_triton = _FLUXSPIKEY_mod.has_triton
|
||||
_SPIKE_AVAILABLE = True
|
||||
triton_ok = _spike_has_triton() if _spike_has_triton else False
|
||||
print(f"[FLUXNATION] Spike loaded (triton={'YES' if triton_ok else 'NO'})")
|
||||
return _SPIKE_AVAILABLE
|
||||
|
||||
_FP8_DTYPE = getattr(torch, "float8_e4m3fn", None)
|
||||
_FP8_MAX = 448.0
|
||||
_FP32_MAX = 3.4028235e+38
|
||||
|
||||
# ── Spike attention configuration ──
|
||||
# Override via environment: FLUX_SPIKE=1 FLUX_SPIKE_SWITCH=0.35 FLUX_SPIKE_CAP=0.35
|
||||
SPIKE_CONFIG = {
|
||||
"enabled": os.environ.get("FLUX_SPIKE", "0") == "1",
|
||||
"switch_pct": float(os.environ.get("FLUX_SPIKE_SWITCH", "0.35")), # activate after 35%
|
||||
"tau": float(os.environ.get("FLUX_SPIKE_TAU", "0.05")), # low threshold for block scoring
|
||||
"leak": 0.0, # NO accumulator across SSBs
|
||||
"event_cap_pct": float(os.environ.get("FLUX_SPIKE_CAP", "0.35")), # keep 35% of blocks
|
||||
"cache_steps": os.environ.get("FLUX_SPIKE_CACHE", "1") == "1", # step caching ON by default
|
||||
}
|
||||
|
||||
|
||||
class SpikeController:
|
||||
"""Tracks diffusion progress and manages spike attention + step caching.
|
||||
|
||||
Universal: works with ANY sampler (Euler=1 eval/step, Heun=2, DPM++=varies).
|
||||
Progress is based on actual forward calls, not sampler step count.
|
||||
Self-calibrates from the previous generation's actual call count.
|
||||
|
||||
Three attention modes per model eval:
|
||||
1. DENSE — full flash attention (early evals + last 2 evals)
|
||||
2. SPIKE — block-sparse Triton attention (compute evals in spike zone)
|
||||
3. CACHED — replay previous eval's output (alternating with spike)
|
||||
|
||||
Each SSB block (38 per eval) is cached independently by block index.
|
||||
"""
|
||||
|
||||
BLOCKS_PER_STEP = 38 # SingleStreamBlock count in FLUX.1
|
||||
|
||||
def __init__(self, config=None):
|
||||
cfg = config or SPIKE_CONFIG
|
||||
self.enabled = cfg.get("enabled", False)
|
||||
self.switch_pct = cfg.get("switch_pct", 0.35)
|
||||
self.tau = cfg.get("tau", 0.05)
|
||||
self.leak = 0.0
|
||||
self.event_cap_pct = cfg.get("event_cap_pct", 0.35)
|
||||
self.cache_enabled = cfg.get("cache_steps", True)
|
||||
|
||||
# Per-generation state
|
||||
self.total_steps = 0
|
||||
self.call_count = 0
|
||||
self._total_calls = 0 # expected total forward calls this gen
|
||||
self._prev_gen_calls = 0 # actual calls from previous gen
|
||||
self._logged_activation = False
|
||||
self._logged_cache = False
|
||||
self._spike_loaded = False
|
||||
|
||||
# SSB step cache
|
||||
self._attn_cache = {}
|
||||
self._last_compute_step = -1
|
||||
|
||||
def _ensure_spike_loaded(self):
|
||||
if self._spike_loaded:
|
||||
return
|
||||
self._spike_loaded = True
|
||||
if self.enabled:
|
||||
_lazy_load_spike()
|
||||
|
||||
# Eval multipliers for ALL ComfyUI samplers (from k_diffusion model() call count)
|
||||
SAMPLER_EVALS = {
|
||||
"euler": 1, "euler_cfg_pp": 1,
|
||||
"euler_ancestral": 1, "euler_ancestral_cfg_pp": 1,
|
||||
"heun": 2, "heunpp2": 4,
|
||||
"exp_heun_2_x0": 2, "exp_heun_2_x0_sde": 2,
|
||||
"dpm_2": 2, "dpm_2_ancestral": 2,
|
||||
"lms": 2, "dpm_fast": 1, "dpm_adaptive": 1,
|
||||
"dpmpp_2s_ancestral": 2, "dpmpp_2s_ancestral_cfg_pp": 2,
|
||||
"dpmpp_sde": 2, "dpmpp_sde_gpu": 2,
|
||||
"dpmpp_2m": 1, "dpmpp_2m_cfg_pp": 1,
|
||||
"dpmpp_2m_sde": 1, "dpmpp_2m_sde_gpu": 1,
|
||||
"dpmpp_2m_sde_heun": 1, "dpmpp_2m_sde_heun_gpu": 1,
|
||||
"dpmpp_3m_sde": 1, "dpmpp_3m_sde_gpu": 1,
|
||||
"ddpm": 1, "ddim": 1, "lcm": 1,
|
||||
"ipndm": 1, "ipndm_v": 1, "deis": 1,
|
||||
"res_multistep": 1, "res_multistep_cfg_pp": 1,
|
||||
"res_multistep_ancestral": 1, "res_multistep_ancestral_cfg_pp": 1,
|
||||
"gradient_estimation": 1, "gradient_estimation_cfg_pp": 1,
|
||||
"er_sde": 1, "seeds_2": 2, "seeds_3": 3,
|
||||
"sa_solver": 2, "sa_solver_pece": 2,
|
||||
"uni_pc": 1, "uni_pc_bh2": 1,
|
||||
}
|
||||
|
||||
def begin_generation(self, total_steps: int, sampler_name: str = ""):
|
||||
# Capture actual call count from previous gen for self-calibration
|
||||
if self.call_count > 0:
|
||||
self._prev_gen_calls = self.call_count
|
||||
|
||||
self.total_steps = total_steps
|
||||
|
||||
# Detect eval multiplier from sampler name
|
||||
evals_per_step = self.SAMPLER_EVALS.get(sampler_name, 0)
|
||||
|
||||
if evals_per_step > 0:
|
||||
# Known sampler — use exact multiplier
|
||||
self._total_calls = total_steps * self.BLOCKS_PER_STEP * evals_per_step
|
||||
elif self._prev_gen_calls > 0:
|
||||
# Unknown sampler — use previous gen's actual count
|
||||
self._total_calls = self._prev_gen_calls
|
||||
else:
|
||||
# First gen, unknown sampler — assume 1 eval/step
|
||||
self._total_calls = total_steps * self.BLOCKS_PER_STEP
|
||||
|
||||
self._sampler_name = sampler_name
|
||||
|
||||
self.call_count = 0
|
||||
self._logged_activation = False
|
||||
self._logged_cache = False
|
||||
self._attn_cache.clear()
|
||||
self._last_compute_step = -1
|
||||
|
||||
@property
|
||||
def total_evals(self):
|
||||
"""Total expected model evaluations this generation."""
|
||||
return self._total_calls // max(1, self.BLOCKS_PER_STEP)
|
||||
|
||||
@property
|
||||
def current_eval(self):
|
||||
"""Which model evaluation we're on (0-indexed)."""
|
||||
if self.BLOCKS_PER_STEP <= 0:
|
||||
return 0
|
||||
return self.call_count // self.BLOCKS_PER_STEP
|
||||
|
||||
@property
|
||||
def block_in_step(self):
|
||||
"""Which SSB block within the current eval (0..37)."""
|
||||
return self.call_count % self.BLOCKS_PER_STEP
|
||||
|
||||
@property
|
||||
def progress(self):
|
||||
"""Generation progress 0.0→1.0 based on actual forward calls."""
|
||||
if self._total_calls <= 0:
|
||||
return 0.0
|
||||
return self.call_count / self._total_calls
|
||||
|
||||
def _in_spike_zone(self):
|
||||
"""Check if we're in the spike-eligible zone."""
|
||||
if not self.enabled:
|
||||
return False
|
||||
if self.progress < self.switch_pct:
|
||||
return False
|
||||
# Last 2 model evaluations always dense for final quality
|
||||
remaining_evals = self.total_evals - self.current_eval
|
||||
if remaining_evals <= 2:
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_attention_mode(self):
|
||||
"""Determine attention mode for current forward call.
|
||||
|
||||
Returns: "dense", "spike", or "cached"
|
||||
"""
|
||||
if not self._in_spike_zone():
|
||||
return "dense"
|
||||
|
||||
# Lazy-load spike on first need
|
||||
self._ensure_spike_loaded()
|
||||
if not self.enabled:
|
||||
return "dense"
|
||||
|
||||
# Eval index since entering spike zone
|
||||
first_spike_eval = math.ceil(self.switch_pct * self.total_evals)
|
||||
evals_in_zone = self.current_eval - first_spike_eval
|
||||
|
||||
# Guard: if negative (boundary rounding), this IS the first eval — always compute
|
||||
if evals_in_zone <= 0:
|
||||
return "spike"
|
||||
|
||||
# Alternating: even evals in zone = compute (spike), odd = cache
|
||||
if not self.cache_enabled:
|
||||
return "spike"
|
||||
|
||||
if evals_in_zone % 2 == 0:
|
||||
return "spike"
|
||||
else:
|
||||
return "cached"
|
||||
|
||||
def attend(self, q, k, v):
|
||||
"""Run spike attention. q/k/v are [B, heads, S, head_dim] fp16."""
|
||||
neuromorph_cfg = {
|
||||
"tau": self.tau,
|
||||
"leak": 0.0,
|
||||
"event_cap_pct": self.event_cap_pct,
|
||||
}
|
||||
|
||||
output, new_state = _spike_attend(
|
||||
q, k, v,
|
||||
backend="spike",
|
||||
state=None,
|
||||
neuromorph=neuromorph_cfg,
|
||||
)
|
||||
|
||||
if not self._logged_activation:
|
||||
telem = new_state.get("telemetry", {})
|
||||
flops = telem.get("flops_saved_pct", 0)
|
||||
kept = telem.get("kept_events_pct", 0)
|
||||
evals_per_step = self.total_evals / max(1, self.total_steps)
|
||||
print(f"[flux_ext] SPIKE active @ eval {self.current_eval}/{self.total_evals} "
|
||||
f"({self.progress*100:.0f}%) — {flops:.1f}% FLOPs saved, "
|
||||
f"{kept:.1f}% events kept"
|
||||
f"{' + step caching' if self.cache_enabled else ''}"
|
||||
f" ({evals_per_step:.0f} eval/step)")
|
||||
self._logged_activation = True
|
||||
|
||||
return output
|
||||
|
||||
def cache_attn(self, block_idx, attn_tensor):
|
||||
"""Store attention output for replay on next step."""
|
||||
self._attn_cache[block_idx] = attn_tensor.detach()
|
||||
self._last_compute_step = self.current_eval
|
||||
|
||||
def get_cached_attn(self, block_idx):
|
||||
"""Retrieve cached attention output."""
|
||||
cached = self._attn_cache.get(block_idx)
|
||||
if cached is not None and not self._logged_cache:
|
||||
print(f"[flux_ext] CACHE replay active @ eval {self.current_eval}/{self.total_evals}")
|
||||
self._logged_cache = True
|
||||
return cached
|
||||
|
||||
def tick(self):
|
||||
"""Call once per SSB forward() invocation."""
|
||||
self.call_count += 1
|
||||
|
||||
|
||||
# Global controller instance
|
||||
_spike_controller = SpikeController()
|
||||
|
||||
# (output_name, path_to_weight, path_to_linear_module)
|
||||
_WEIGHT_SPECS = (
|
||||
("mod_w", ("modulation", "lin", "weight"), ("modulation", "lin")),
|
||||
("l1w", ("linear1", "weight"), ("linear1",)),
|
||||
("l2w", ("linear2", "weight"), ("linear2",)),
|
||||
)
|
||||
|
||||
_weight_cache = {}
|
||||
_pe_cache = {}
|
||||
|
||||
|
||||
def _get_attr_path(obj, path):
|
||||
cur = obj
|
||||
for part in path:
|
||||
cur = getattr(cur, part)
|
||||
return cur
|
||||
|
||||
|
||||
def _normalize_scale(scale, device):
|
||||
return scale.to(dtype=torch.float32, device=device).reshape(1).contiguous()
|
||||
|
||||
|
||||
def _to_f16(tensor, device):
|
||||
if tensor.dtype == torch.float16 and tensor.is_contiguous() and tensor.device == device:
|
||||
return tensor
|
||||
return tensor.to(dtype=torch.float16, device=device).contiguous()
|
||||
|
||||
|
||||
def _quantize_weight_to_fp8(linear_mod, weight, device):
|
||||
w = _to_f16(weight, device)
|
||||
amax = torch.amax(torch.abs(w)).to(dtype=torch.float32)
|
||||
amax = torch.nan_to_num(amax, nan=0.0, posinf=_FP32_MAX, neginf=0.0)
|
||||
scale = torch.where(amax < 1e-12, torch.ones_like(amax), amax / _FP8_MAX).reshape(1).contiguous()
|
||||
inv_scale = (1.0 / scale).to(dtype=w.dtype)
|
||||
w_fp8 = torch.clamp(w * inv_scale, min=-_FP8_MAX, max=_FP8_MAX).to(dtype=_FP8_DTYPE).contiguous()
|
||||
with torch.no_grad():
|
||||
linear_mod.weight = torch.nn.Parameter(w_fp8, requires_grad=False)
|
||||
linear_mod.scale_weight = _normalize_scale(scale, device).reshape(())
|
||||
return linear_mod.weight.contiguous(), _normalize_scale(linear_mod.scale_weight, device)
|
||||
|
||||
|
||||
def _cache_weights(module, device):
|
||||
"""Extract and cache all weight tensors ONCE per module. Not per forward call."""
|
||||
mod_id = id(module)
|
||||
if mod_id in _weight_cache:
|
||||
return _weight_cache[mod_id]
|
||||
|
||||
materialized = {}
|
||||
|
||||
for name, weight_path, linear_path in _WEIGHT_SPECS:
|
||||
tensor = _get_attr_path(module, weight_path)
|
||||
if not tensor.is_cuda or tensor.device != device:
|
||||
tensor = tensor.to(device=device)
|
||||
linear_mod = _get_attr_path(module, linear_path)
|
||||
|
||||
if _FP8_DTYPE is not None and tensor.dtype == _FP8_DTYPE:
|
||||
materialized[name] = tensor if tensor.is_contiguous() else tensor.contiguous()
|
||||
scale = getattr(linear_mod, "scale_weight", None)
|
||||
if scale is not None:
|
||||
materialized[f"{name}_scale"] = _normalize_scale(scale, device)
|
||||
else:
|
||||
materialized[f"{name}_scale"] = torch.ones(1, dtype=torch.float32, device=device)
|
||||
else:
|
||||
materialized[name], materialized[f"{name}_scale"] = _quantize_weight_to_fp8(
|
||||
linear_mod, tensor, device
|
||||
)
|
||||
|
||||
for name, attr_path in (
|
||||
("mod_b", ("modulation", "lin", "bias")),
|
||||
("l1b", ("linear1", "bias")),
|
||||
("l2b", ("linear2", "bias")),
|
||||
("qns", ("norm", "query_norm", "weight")),
|
||||
("kns", ("norm", "key_norm", "weight")),
|
||||
):
|
||||
tensor = _get_attr_path(module, attr_path)
|
||||
if tensor is None:
|
||||
materialized[name] = torch.zeros(1, dtype=torch.float16, device=device)
|
||||
continue
|
||||
if not tensor.is_cuda or tensor.device != device:
|
||||
tensor = tensor.to(device=device)
|
||||
materialized[name] = _to_f16(tensor, device)
|
||||
|
||||
_weight_cache[mod_id] = materialized
|
||||
return materialized
|
||||
|
||||
|
||||
def _convert_pe(pe):
|
||||
"""Convert PE tensor from [1,1,S,64,2,1] to [S,64,2] fp16."""
|
||||
pe_key = (pe.data_ptr(), pe.shape)
|
||||
if pe_key in _pe_cache:
|
||||
return _pe_cache[pe_key]
|
||||
|
||||
pe_cos = pe[0, 0, :, :, 0, 0]
|
||||
pe_sin = pe[0, 0, :, :, 1, 0]
|
||||
result = torch.stack([pe_cos, pe_sin], dim=-1).to(dtype=torch.float16, device=pe.device).contiguous()
|
||||
|
||||
_pe_cache[pe_key] = result
|
||||
return result
|
||||
|
||||
|
||||
def patched_forward(self, x, vec, pe, attn_mask=None, modulation_dims=None, transformer_options={}):
|
||||
b, s, h = x.shape
|
||||
if not (x.is_cuda and h == 3072 and self.num_heads == 24
|
||||
and modulation_dims is None and pe is not None):
|
||||
raise RuntimeError(
|
||||
f"[flux_ext] Unsupported block: cuda={x.is_cuda} h={h} heads={self.num_heads}"
|
||||
)
|
||||
|
||||
device = x.device
|
||||
weights = _cache_weights(self, device)
|
||||
|
||||
# Convert inputs (no-op if already correct dtype/contiguity)
|
||||
pe_h = _convert_pe(pe)
|
||||
|
||||
# ── PREPARE: one torch.ops call replaces 19 ctypes args ──
|
||||
q, k, v, mlp, gate = torch.ops.flux_ext.ssb_prepare(
|
||||
x, vec, pe_h,
|
||||
weights["mod_w"], weights["mod_w_scale"], weights["mod_b"],
|
||||
weights["l1w"], weights["l1w_scale"], weights["l1b"],
|
||||
weights["qns"], weights["kns"],
|
||||
)
|
||||
|
||||
# ── ATTENTION: dense / spike / cached ──
|
||||
attn_mode = _spike_controller.get_attention_mode()
|
||||
block_idx = _spike_controller.block_in_step
|
||||
|
||||
if attn_mode == "cached":
|
||||
attn = _spike_controller.get_cached_attn(block_idx)
|
||||
elif attn_mode == "spike":
|
||||
# Spike: block-sparse Triton attention
|
||||
attn_out = _spike_controller.attend(q, k, v)
|
||||
attn = attn_out.transpose(1, 2).reshape(b, s, h)
|
||||
_spike_controller.cache_attn(block_idx, attn)
|
||||
else:
|
||||
# Dense: ComfyUI's flash attention
|
||||
attn = flux_attention(q, k, v, pe=None, mask=attn_mask,
|
||||
transformer_options=transformer_options)
|
||||
|
||||
# Tick AFTER using block_in_step so indexing is correct
|
||||
_spike_controller.tick()
|
||||
|
||||
if not attn.is_contiguous():
|
||||
attn = attn.contiguous()
|
||||
if attn.dtype != torch.float16:
|
||||
attn = attn.to(dtype=torch.float16).contiguous()
|
||||
|
||||
# ── FINISH: L2 GEMM + fused gate+residual+bias in ONE call ──
|
||||
result = torch.ops.flux_ext.ssb_finish(
|
||||
x, attn, mlp, gate,
|
||||
weights["l2w"], weights["l2w_scale"], weights["l2b"],
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def apply_patch():
|
||||
try:
|
||||
from comfy.ldm.flux.layers import SingleStreamBlock
|
||||
|
||||
if not hasattr(SingleStreamBlock, "_original_forward"):
|
||||
SingleStreamBlock._original_forward = SingleStreamBlock.forward
|
||||
|
||||
# Try torch.compile on the dense inner path
|
||||
use_compile = os.environ.get("FLUX_COMPILE", "0") == "1"
|
||||
if use_compile:
|
||||
try:
|
||||
compiled_forward = torch.compile(patched_forward, mode="reduce-overhead", dynamic=True)
|
||||
SingleStreamBlock.forward = compiled_forward
|
||||
print("[flux_ext] torch.compile enabled (reduce-overhead, dynamic)")
|
||||
except Exception as e:
|
||||
print(f"[flux_ext] torch.compile failed ({e}), using uncompiled")
|
||||
SingleStreamBlock.forward = patched_forward
|
||||
else:
|
||||
SingleStreamBlock.forward = patched_forward
|
||||
|
||||
mode = "torch.ops"
|
||||
if _spike_controller.enabled:
|
||||
mode += f" + SPIKE (switch={_spike_controller.switch_pct*100:.0f}%"
|
||||
mode += f", cap={_spike_controller.event_cap_pct*100:.0f}%"
|
||||
if _spike_controller.cache_enabled:
|
||||
mode += " + step-cache"
|
||||
mode += ")"
|
||||
if use_compile:
|
||||
mode += " + torch.compile"
|
||||
print(f"[flux_ext] Patched SingleStreamBlock.forward ({mode})")
|
||||
|
||||
# Hook into ComfyUI's sampling to track generation start
|
||||
_hook_sampling_start()
|
||||
|
||||
except ImportError as exc:
|
||||
print(f"[flux_ext] Could not patch: {exc}")
|
||||
except Exception as exc:
|
||||
print(f"[flux_ext] Patch error: {exc}")
|
||||
|
||||
|
||||
def _hook_sampling_start():
|
||||
"""Hook into ComfyUI's sample function to call begin_generation() at each prompt."""
|
||||
try:
|
||||
import comfy.sample as sample_mod
|
||||
|
||||
# Store the REAL original only once
|
||||
if not hasattr(sample_mod, "_flux_ext_original_sample"):
|
||||
sample_mod._flux_ext_original_sample = sample_mod.sample
|
||||
|
||||
real_original = sample_mod._flux_ext_original_sample
|
||||
|
||||
def hooked_sample(*args, **kwargs):
|
||||
steps = 20
|
||||
sampler_name = ""
|
||||
if len(args) > 2 and isinstance(args[2], int):
|
||||
steps = args[2]
|
||||
elif "steps" in kwargs:
|
||||
steps = kwargs["steps"]
|
||||
# comfy.sample.sample(model, noise, steps, cfg, sampler_name, ...)
|
||||
if len(args) > 4:
|
||||
sampler_name = str(args[4])
|
||||
elif "sampler_name" in kwargs:
|
||||
sampler_name = str(kwargs["sampler_name"])
|
||||
|
||||
_spike_controller.begin_generation(steps, sampler_name)
|
||||
_pe_cache.clear()
|
||||
evals = _spike_controller.total_evals
|
||||
source = f"sampler={sampler_name}" if sampler_name else "estimated"
|
||||
print(f"[flux_ext] Generation start: {steps} steps × {evals//max(1,steps)} evals/step = "
|
||||
f"{evals} evals ({source}), "
|
||||
f"spike={'ON' if _spike_controller.enabled else 'OFF'}")
|
||||
|
||||
return real_original(*args, **kwargs)
|
||||
|
||||
sample_mod.sample = hooked_sample
|
||||
print("[flux_ext] Sampling hook installed on comfy.sample.sample")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[flux_ext] Could not hook sampling start: {e}")
|
||||
+195
@@ -0,0 +1,195 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
import os
|
||||
import warnings
|
||||
os.environ['TRITON_DISABLE_LINE_INFO'] = '1'
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("ignore", category=UserWarning, module="triton")
|
||||
import triton
|
||||
import triton.language as tl
|
||||
TRITON_AVAILABLE = True
|
||||
except Exception:
|
||||
TRITON_AVAILABLE = False
|
||||
|
||||
class _TritonStub:
|
||||
def jit(self, fn=None, **kwargs):
|
||||
if fn is None:
|
||||
return lambda inner_fn: inner_fn
|
||||
return fn
|
||||
|
||||
class _TLStub:
|
||||
def __getattr__(self, name):
|
||||
raise RuntimeError("Triton is required for spike attention.")
|
||||
|
||||
triton = _TritonStub()
|
||||
tl = _TLStub()
|
||||
|
||||
BLOCK_SIZE = 32
|
||||
|
||||
|
||||
def _ceil_div(x: int, y: int) -> int:
|
||||
return (x + y - 1) // y
|
||||
|
||||
|
||||
if TRITON_AVAILABLE:
|
||||
|
||||
@triton.jit
|
||||
def _blocksparse_attention_kernel(
|
||||
Q, K, V, Out,
|
||||
BlockIndices, BlockCounts,
|
||||
seq_len, head_dim, num_heads,
|
||||
stride_q_batch, stride_q_head, stride_q_seq, stride_q_dim,
|
||||
stride_k_batch, stride_k_head, stride_k_seq, stride_k_dim,
|
||||
stride_v_batch, stride_v_head, stride_v_seq, stride_v_dim,
|
||||
stride_o_batch, stride_o_head, stride_o_seq, stride_o_dim,
|
||||
stride_idx_batch, stride_idx_head, stride_idx_qblock,
|
||||
stride_counts_batch, stride_counts_head, stride_counts_qblock,
|
||||
scale,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
MAX_BLOCKS: tl.constexpr,
|
||||
):
|
||||
pid_bh = tl.program_id(0)
|
||||
pid_qb = tl.program_id(1)
|
||||
head = pid_bh % num_heads
|
||||
batch = pid_bh // num_heads
|
||||
q_start = pid_qb * BLOCK_SIZE
|
||||
rows = q_start + tl.arange(0, BLOCK_SIZE)
|
||||
cols = tl.arange(0, HEAD_DIM)
|
||||
|
||||
q_ptr = Q + batch * stride_q_batch + head * stride_q_head + rows[:, None] * stride_q_seq + cols[None, :] * stride_q_dim
|
||||
q_mask = rows[:, None] < seq_len
|
||||
q = tl.where(q_mask, tl.load(q_ptr, mask=q_mask, other=0.0), 0.0).to(tl.float32)
|
||||
|
||||
out_acc = tl.zeros((BLOCK_SIZE, HEAD_DIM), dtype=tl.float32)
|
||||
m_i = tl.full((BLOCK_SIZE,), float("-inf"), dtype=tl.float32)
|
||||
l_i = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
|
||||
|
||||
block_count = tl.load(BlockCounts + batch * stride_counts_batch + head * stride_counts_head + pid_qb * stride_counts_qblock, mask=True, other=0).to(tl.int32)
|
||||
idx_ptr_base = BlockIndices + batch * stride_idx_batch + head * stride_idx_head + pid_qb * stride_idx_qblock
|
||||
|
||||
for block_pos in range(MAX_BLOCKS):
|
||||
blk_in_range = block_pos < block_count
|
||||
block_id = tl.load(idx_ptr_base + block_pos, mask=blk_in_range, other=-1)
|
||||
valid = blk_in_range & (block_id >= 0)
|
||||
k_rows = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
kv_mask = (k_rows >= 0) & (k_rows < seq_len) & valid
|
||||
|
||||
k_ptr = K + batch * stride_k_batch + head * stride_k_head + k_rows[None, :] * stride_k_seq + cols[:, None] * stride_k_dim
|
||||
v_ptr = V + batch * stride_v_batch + head * stride_v_head + k_rows[:, None] * stride_v_seq + cols[None, :] * stride_v_dim
|
||||
|
||||
k_block = tl.load(k_ptr, mask=kv_mask[None, :], other=0.0).to(tl.float32)
|
||||
v_block = tl.load(v_ptr, mask=kv_mask[:, None], other=0.0).to(tl.float32)
|
||||
|
||||
scores = tl.dot(q, k_block) * scale
|
||||
scores = tl.where(kv_mask[None, :], scores, float("-inf"))
|
||||
scores = tl.where(q_mask, scores, float("-inf"))
|
||||
|
||||
m_new = tl.maximum(m_i, tl.max(scores, axis=1))
|
||||
exp_scores = tl.where(kv_mask[None, :], tl.exp(scores - m_new[:, None]), 0.0)
|
||||
l_new = l_i * tl.exp(m_i - m_new) + tl.sum(exp_scores, axis=1)
|
||||
out_acc = (out_acc * (l_i * tl.exp(m_i - m_new))[:, None] + tl.dot(exp_scores, v_block)) / tl.maximum(l_new, 1e-6)[:, None]
|
||||
m_i = m_new
|
||||
l_i = l_new
|
||||
|
||||
out_ptr = Out + batch * stride_o_batch + head * stride_o_head + rows[:, None] * stride_o_seq + cols[None, :] * stride_o_dim
|
||||
tl.store(out_ptr, out_acc, mask=q_mask)
|
||||
|
||||
def _launch_block_sparse_attention(q, k, v, block_indices, block_counts):
|
||||
batch, heads, seq_len, head_dim = q.shape
|
||||
num_blocks = _ceil_div(seq_len, BLOCK_SIZE)
|
||||
max_blocks = block_indices.shape[-1]
|
||||
out_fp32 = torch.zeros(batch, heads, seq_len, head_dim, dtype=torch.float32, device=q.device)
|
||||
q_c, k_c, v_c = q.contiguous(), k.contiguous(), v.contiguous()
|
||||
bi_c, bc_c = block_indices.contiguous(), block_counts.contiguous()
|
||||
_blocksparse_attention_kernel[(batch * heads, num_blocks)](
|
||||
Q=q_c, K=k_c, V=v_c, Out=out_fp32,
|
||||
BlockIndices=bi_c, BlockCounts=bc_c,
|
||||
seq_len=seq_len, head_dim=head_dim, num_heads=heads,
|
||||
stride_q_batch=q_c.stride(0), stride_q_head=q_c.stride(1), stride_q_seq=q_c.stride(2), stride_q_dim=q_c.stride(3),
|
||||
stride_k_batch=k_c.stride(0), stride_k_head=k_c.stride(1), stride_k_seq=k_c.stride(2), stride_k_dim=k_c.stride(3),
|
||||
stride_v_batch=v_c.stride(0), stride_v_head=v_c.stride(1), stride_v_seq=v_c.stride(2), stride_v_dim=v_c.stride(3),
|
||||
stride_o_batch=out_fp32.stride(0), stride_o_head=out_fp32.stride(1), stride_o_seq=out_fp32.stride(2), stride_o_dim=out_fp32.stride(3),
|
||||
stride_idx_batch=bi_c.stride(0), stride_idx_head=bi_c.stride(1), stride_idx_qblock=bi_c.stride(2),
|
||||
stride_counts_batch=bc_c.stride(0), stride_counts_head=bc_c.stride(1), stride_counts_qblock=bc_c.stride(2),
|
||||
scale=1.0 / math.sqrt(head_dim),
|
||||
BLOCK_SIZE=BLOCK_SIZE, HEAD_DIM=head_dim, MAX_BLOCKS=max_blocks,
|
||||
)
|
||||
return out_fp32 if q.dtype == torch.float32 else out_fp32.to(q.dtype)
|
||||
|
||||
else:
|
||||
def _launch_block_sparse_attention(q, k, v, block_indices, block_counts):
|
||||
raise RuntimeError("Triton unavailable — cannot run block-sparse attention.")
|
||||
|
||||
|
||||
def has_triton() -> bool:
|
||||
return TRITON_AVAILABLE and torch.cuda.is_available()
|
||||
|
||||
|
||||
def attend_spike_triton(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
tau: float,
|
||||
leak: float,
|
||||
event_cap_pct: float,
|
||||
state: Optional[Dict[str, torch.Tensor]] = None,
|
||||
debug: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[torch.Tensor, Dict[str, Any]]:
|
||||
batch, heads, seq_len, head_dim = q.shape
|
||||
|
||||
if seq_len == 0:
|
||||
return torch.empty_like(q), {"telemetry": {"spike_events": 0}}
|
||||
|
||||
scale = head_dim ** -0.5
|
||||
q_fp32 = q.to(torch.float32)
|
||||
k_fp32 = k.to(torch.float32)
|
||||
|
||||
num_blocks = _ceil_div(seq_len, BLOCK_SIZE)
|
||||
target_blocks = max(1, int(num_blocks * (int(seq_len * event_cap_pct) / max(1, seq_len))))
|
||||
block_cap = max(1, min(num_blocks, target_blocks))
|
||||
|
||||
pad = num_blocks * BLOCK_SIZE - seq_len
|
||||
if pad > 0:
|
||||
q_fp32 = torch.nn.functional.pad(q_fp32, (0, 0, 0, pad))
|
||||
k_fp32 = torch.nn.functional.pad(k_fp32, (0, 0, 0, pad))
|
||||
|
||||
q_blocks = q_fp32.view(batch, heads, num_blocks, BLOCK_SIZE, head_dim)
|
||||
k_blocks = k_fp32.view(batch, heads, num_blocks, BLOCK_SIZE, head_dim)
|
||||
q_rep = q_blocks.mean(dim=3)
|
||||
k_rep = k_blocks.mean(dim=3)
|
||||
|
||||
block_scores = torch.matmul(q_rep, k_rep.transpose(-2, -1)) * scale
|
||||
row_max = block_scores.max(dim=-1, keepdim=True)[0]
|
||||
block_scores = torch.where(block_scores >= tau * row_max, block_scores, torch.full_like(block_scores, float('-inf')))
|
||||
|
||||
top_vals, top_idx = torch.topk(block_scores, k=block_cap, dim=-1)
|
||||
valid = torch.isfinite(top_vals)
|
||||
block_counts = valid.sum(dim=-1).to(torch.int32)
|
||||
block_indices = torch.where(valid, top_idx, top_idx.new_full(top_idx.shape, -1))
|
||||
|
||||
output = _launch_block_sparse_attention(q, k, v, block_indices, block_counts)
|
||||
|
||||
avg_blocks = block_counts.float().mean().item()
|
||||
kept_pct = 100.0 * (avg_blocks / max(1.0, float(num_blocks)))
|
||||
tok_pct = (avg_blocks * BLOCK_SIZE) / max(1.0, float(seq_len))
|
||||
|
||||
return output, {
|
||||
"block_indices": block_indices.detach(),
|
||||
"block_counts": block_counts.detach(),
|
||||
"telemetry": {
|
||||
"backend": "block-topk",
|
||||
"blocks_avg": avg_blocks,
|
||||
"blocks_pct": kept_pct,
|
||||
"block_cap": block_cap,
|
||||
"flops_saved_pct": 100.0 * (1.0 - tok_pct),
|
||||
"kept_events_pct": tok_pct * 100.0,
|
||||
},
|
||||
}
|
||||
+141
@@ -0,0 +1,141 @@
|
||||
#ifndef FLUX_STREAM_H
|
||||
#define FLUX_STREAM_H
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cublas_v2.h>
|
||||
#include <cstdio>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
namespace transformer {
|
||||
|
||||
constexpr int HIDDEN = 3072;
|
||||
constexpr int HEADS = 24;
|
||||
constexpr int HDIM = 128;
|
||||
constexpr int FFDIM = 12288;
|
||||
constexpr int MODDIM = 3 * HIDDEN;
|
||||
constexpr int LIN1DIM = 3 * HIDDEN + FFDIM;
|
||||
constexpr int LIN2DIM = HIDDEN + FFDIM;
|
||||
constexpr float LN_EPS = 1e-6f;
|
||||
constexpr float RMS_EPS = 1e-6f;
|
||||
constexpr float ATT_SC = 0.0883883476f;
|
||||
constexpr int FP8_MOD_W = 1;
|
||||
constexpr int FP8_L1_W = 2;
|
||||
constexpr int FP8_L2_W = 4;
|
||||
#define CK(call) do { \
|
||||
cudaError_t err__ = (call); \
|
||||
if (err__ != cudaSuccess) { \
|
||||
throw std::runtime_error(std::string("CUDA ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + cudaGetErrorString(err__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define BK(call) do { \
|
||||
cublasStatus_t st__ = (call); \
|
||||
if (st__ != CUBLAS_STATUS_SUCCESS) { \
|
||||
throw std::runtime_error(std::string("cuBLAS ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + std::to_string((int)st__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
// Returns minimum workspace bytes needed for ssb_create
|
||||
static inline size_t ssb_workspace_bytes(int batch, int seq) {
|
||||
size_t rows = static_cast<size_t>(batch) * seq;
|
||||
size_t b = static_cast<size_t>(batch);
|
||||
size_t sz = 0;
|
||||
// Each sub-buffer 256-byte aligned
|
||||
#define ADDBUF(count, elem_size) \
|
||||
sz = (sz + 255) & ~size_t(255); sz += (count) * (elem_size);
|
||||
ADDBUF(b * HIDDEN, 2); // d_silu_
|
||||
ADDBUF(b * MODDIM, 2); // d_mod_
|
||||
ADDBUF(rows * HIDDEN, 2); // d_xmod_
|
||||
ADDBUF(rows * LIN1DIM, 2); // d_l1o_
|
||||
ADDBUF(rows * LIN2DIM, 1); // d_act_fp8_
|
||||
ADDBUF(1, 4); // d_act_amax_
|
||||
ADDBUF(1, 4); // d_act_scale_
|
||||
ADDBUF(1, 4); // d_out_scale_
|
||||
sz = (sz + 255) & ~size_t(255);
|
||||
sz += 32ull << 20; // 32MB CUTLASS/cuBLAS workspace
|
||||
#undef ADDBUF
|
||||
return sz;
|
||||
}
|
||||
|
||||
class SSBEngine {
|
||||
public:
|
||||
// workspace_buf: Python-allocated GPU memory (torch tensor .data_ptr())
|
||||
// workspace_buf_size: must be >= ssb_workspace_bytes(batch, seq)
|
||||
SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size);
|
||||
~SSBEngine();
|
||||
|
||||
void prepare(const half* d_x,
|
||||
const half* d_vec,
|
||||
const half* d_pe,
|
||||
const void* mod_w,
|
||||
const float* mod_w_scale,
|
||||
const half* mod_b,
|
||||
const void* l1w,
|
||||
const float* l1w_scale,
|
||||
const half* l1b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
int fp8_mask,
|
||||
half* d_q,
|
||||
half* d_k,
|
||||
half* d_v,
|
||||
half* d_mlp,
|
||||
half* d_gate,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
void finish(const half* d_x,
|
||||
const half* d_attn,
|
||||
half* d_mlp,
|
||||
const half* d_gate,
|
||||
half* d_out,
|
||||
const void* l2w,
|
||||
const float* l2w_scale,
|
||||
const half* l2b,
|
||||
int fp8_mask,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
void forward(const half* d_x,
|
||||
const half* d_vec,
|
||||
half* d_out,
|
||||
const half* d_pe,
|
||||
const void* mod_w,
|
||||
const float* mod_w_scale,
|
||||
const half* mod_b,
|
||||
const void* l1w,
|
||||
const float* l1w_scale,
|
||||
const half* l1b,
|
||||
const void* l2w,
|
||||
const float* l2w_scale,
|
||||
const half* l2b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
int fp8_mask,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
int batch() const { return B_; }
|
||||
int seq() const { return S_; }
|
||||
|
||||
private:
|
||||
int B_;
|
||||
int S_;
|
||||
cublasHandle_t cub_;
|
||||
void* workspace_;
|
||||
size_t workspace_bytes_;
|
||||
|
||||
half* d_silu_;
|
||||
half* d_mod_;
|
||||
half* d_xmod_;
|
||||
half* d_l1o_;
|
||||
__nv_fp8_e4m3* d_act_fp8_;
|
||||
unsigned int* d_act_amax_;
|
||||
float* d_act_scale_;
|
||||
float* d_out_scale_;
|
||||
};
|
||||
|
||||
} // namespace transformer
|
||||
|
||||
#endif // FLUX_STREAM_H
|
||||
@@ -0,0 +1,164 @@
|
||||
# FLUXNATION
|
||||
|
||||
**Fused CUDA kernel for FLUX.1 image generation — RTX 4090 optimized.**
|
||||
|
||||
~9.3 seconds for 20 steps at 1024×1024. No other kernel does what this does.
|
||||
|
||||
---
|
||||
|
||||
## What it does
|
||||
|
||||
Three technologies fused into one ComfyUI custom node:
|
||||
|
||||
- **FP8 fused CUDA kernel** — replaces ComfyUI's SingleStreamBlock forward pass with a single `torch.ops` call. Modulation, L1 GEMM, QKV split, RoPE, L2 GEMM, gate+residual+bias all fused. Zero ctypes overhead.
|
||||
- **Spike attention** — neuromorphic block-sparse Triton attention. After step 10/20, scores attention blocks by dot-product similarity and only computes the top 45%. Saves ~90% attention FLOPs on spike steps.
|
||||
- **Step caching** — alternates between computing spike attention and replaying the cached output. Zero attention compute on every other spike step.
|
||||
|
||||
No fallbacks. If it breaks, it crashes — no silent quality degradation.
|
||||
|
||||
---
|
||||
|
||||
## Requirements
|
||||
|
||||
- Windows (tested) / Linux (should work)
|
||||
- RTX 4090 (SM 8.9) — other Ampere/Ada cards may work but untested
|
||||
- Python 3.10
|
||||
- CUDA 12.x
|
||||
- PyTorch 2.x with CUDA
|
||||
- ComfyUI
|
||||
- Visual Studio 2022 (for building the extension)
|
||||
|
||||
---
|
||||
|
||||
## Installation
|
||||
|
||||
### Step 1 — Clone and copy the custom node
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ULT7RA/cuda-kernels.git
|
||||
```
|
||||
|
||||
Copy `FLUXNATION/custom_nodes/FLUXNATION/` into your ComfyUI custom nodes folder:
|
||||
|
||||
```
|
||||
ComfyUI/custom_nodes/FLUXNATION/
|
||||
```
|
||||
|
||||
### Step 2 — Build the CUDA extension
|
||||
|
||||
From the `FLUXNATION/` folder:
|
||||
|
||||
```bat
|
||||
build_ext.bat
|
||||
```
|
||||
|
||||
This requires:
|
||||
- Visual Studio 2022 (with C++ workload)
|
||||
- CUDA Toolkit 12.x
|
||||
- Python 3.10
|
||||
|
||||
The build installs `flux_ext.cp310-win_amd64.pyd` into your Python site-packages automatically.
|
||||
|
||||
### Step 3 — Launch ComfyUI with spike enabled
|
||||
|
||||
Set these environment variables before launching:
|
||||
|
||||
```bat
|
||||
set FLUX_SPIKE=1
|
||||
set FLUX_SPIKE_SWITCH=0.50
|
||||
set FLUX_SPIKE_CAP=0.45
|
||||
set FLUX_SPIKE_TAU=0.05
|
||||
set FLUX_SPIKE_CACHE=1
|
||||
py -3.10 -s ComfyUI\main.py --windows-standalone-build --fast --dont-upcast-attention --disable-cuda-malloc --disable-mmap
|
||||
```
|
||||
|
||||
Or use the included `run_nvidia_gpu_spike.bat`.
|
||||
|
||||
---
|
||||
|
||||
## ComfyUI Nodes
|
||||
|
||||
Four nodes appear under the **FLUXNATION** category:
|
||||
|
||||
| Node | Purpose |
|
||||
|------|---------|
|
||||
| `ULT7RAFLUX 👑` | Master node — all controls in one. Connect between Load Checkpoint and KSampler. |
|
||||
| `FLUXNATION 🔥` | Pass-through — logs kernel status |
|
||||
| `FLUXATTENHUT ⚡` | Configure spike params (switch %, cap %, tau) |
|
||||
| `FLUXSPIKEY 🧠` | Toggle spike and step cache on/off |
|
||||
|
||||
**Basic workflow:**
|
||||
```
|
||||
Load Checkpoint → ULT7RAFLUX → KSampler
|
||||
```
|
||||
|
||||
The kernel activates at ComfyUI startup regardless of whether the node is in your graph. The nodes give you live control without restarting.
|
||||
|
||||
---
|
||||
|
||||
## Environment Variables
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `FLUX_SPIKE` | `0` | Set to `1` to enable spike attention |
|
||||
| `FLUX_SPIKE_SWITCH` | `0.35` | Fraction of steps before spike activates (`0.50` = step 10/20) |
|
||||
| `FLUX_SPIKE_CAP` | `0.35` | Fraction of attention blocks to keep (`0.45` = keep top 45%) |
|
||||
| `FLUX_SPIKE_TAU` | `0.05` | Block scoring threshold — lower keeps more blocks |
|
||||
| `FLUX_SPIKE_CACHE` | `1` | Set to `1` to enable step caching |
|
||||
| `FLUX_EXT_DISABLE` | `0` | Set to `1` to run stock ComfyUI (benchmarking) |
|
||||
|
||||
---
|
||||
|
||||
## How spike attention works
|
||||
|
||||
FLUX.1 has 38 SingleStreamBlocks per step. At step 10/20:
|
||||
|
||||
```
|
||||
Steps 0–9: DENSE — full flash attention (image structure forms here)
|
||||
Step 10: SPIKE — block-sparse Triton, fills cache for blocks 0–37
|
||||
Step 11: CACHED — replay step 10 output, zero attention compute
|
||||
Step 12: SPIKE — refill cache
|
||||
Step 13: CACHED
|
||||
...
|
||||
Steps 18–19: DENSE — refinement tail
|
||||
```
|
||||
|
||||
Block scoring uses real dot-product: `q_rep · k_rep^T` across heads. Top 45% of blocks by score get computed, rest are skipped. Quality is maintained because image structure is already established by step 10.
|
||||
|
||||
---
|
||||
|
||||
## Performance
|
||||
|
||||
RTX 4090 — FLUX.1-dev — 20 steps — 1024×1024:
|
||||
|
||||
| Mode | Time |
|
||||
|------|------|
|
||||
| Stock ComfyUI | ~13–14s |
|
||||
| FLUXNATION (kernel only) | ~11s |
|
||||
| FLUXNATION + spike + cache | **~9.3s** |
|
||||
|
||||
---
|
||||
|
||||
## File structure
|
||||
|
||||
```
|
||||
FLUXNATION/
|
||||
├── README.md
|
||||
├── setup_ext.py — builds flux_ext CUDA extension
|
||||
├── build_ext.bat — one-click build (Windows)
|
||||
├── csrc/
|
||||
│ ├── flux_ext.cpp — torch.ops registration
|
||||
│ └── flux_ext_kernels.cu — FP8 GEMM + fused gate kernel
|
||||
└── custom_nodes/
|
||||
└── FLUXNATION/
|
||||
├── __init__.py — ComfyUI loader + node definitions
|
||||
├── FLUXPATCH.py — monkey-patch + spike controller
|
||||
├── spike_attention.py — attend() router
|
||||
└── spike_triton.py — Triton block-sparse kernel
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,141 @@
|
||||
#ifndef SINGLE_STREAM_BLOCK_H
|
||||
#define SINGLE_STREAM_BLOCK_H
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cublas_v2.h>
|
||||
#include <cstdio>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
namespace transformer {
|
||||
|
||||
constexpr int HIDDEN = 3072;
|
||||
constexpr int HEADS = 24;
|
||||
constexpr int HDIM = 128;
|
||||
constexpr int FFDIM = 12288;
|
||||
constexpr int MODDIM = 3 * HIDDEN;
|
||||
constexpr int LIN1DIM = 3 * HIDDEN + FFDIM;
|
||||
constexpr int LIN2DIM = HIDDEN + FFDIM;
|
||||
constexpr float LN_EPS = 1e-6f;
|
||||
constexpr float RMS_EPS = 1e-6f;
|
||||
constexpr float ATT_SC = 0.0883883476f;
|
||||
constexpr int FP8_MOD_W = 1;
|
||||
constexpr int FP8_L1_W = 2;
|
||||
constexpr int FP8_L2_W = 4;
|
||||
#define CK(call) do { \
|
||||
cudaError_t err__ = (call); \
|
||||
if (err__ != cudaSuccess) { \
|
||||
throw std::runtime_error(std::string("CUDA ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + cudaGetErrorString(err__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define BK(call) do { \
|
||||
cublasStatus_t st__ = (call); \
|
||||
if (st__ != CUBLAS_STATUS_SUCCESS) { \
|
||||
throw std::runtime_error(std::string("cuBLAS ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + std::to_string((int)st__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
// Returns minimum workspace bytes needed for ssb_create
|
||||
static inline size_t ssb_workspace_bytes(int batch, int seq) {
|
||||
size_t rows = static_cast<size_t>(batch) * seq;
|
||||
size_t b = static_cast<size_t>(batch);
|
||||
size_t sz = 0;
|
||||
// Each sub-buffer 256-byte aligned
|
||||
#define ADDBUF(count, elem_size) \
|
||||
sz = (sz + 255) & ~size_t(255); sz += (count) * (elem_size);
|
||||
ADDBUF(b * HIDDEN, 2); // d_silu_
|
||||
ADDBUF(b * MODDIM, 2); // d_mod_
|
||||
ADDBUF(rows * HIDDEN, 2); // d_xmod_
|
||||
ADDBUF(rows * LIN1DIM, 2); // d_l1o_
|
||||
ADDBUF(rows * LIN2DIM, 1); // d_act_fp8_
|
||||
ADDBUF(1, 4); // d_act_amax_
|
||||
ADDBUF(1, 4); // d_act_scale_
|
||||
ADDBUF(1, 4); // d_out_scale_
|
||||
sz = (sz + 255) & ~size_t(255);
|
||||
sz += 32ull << 20; // 32MB CUTLASS/cuBLAS workspace
|
||||
#undef ADDBUF
|
||||
return sz;
|
||||
}
|
||||
|
||||
class SSBEngine {
|
||||
public:
|
||||
// workspace_buf: Python-allocated GPU memory (torch tensor .data_ptr())
|
||||
// workspace_buf_size: must be >= ssb_workspace_bytes(batch, seq)
|
||||
SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size);
|
||||
~SSBEngine();
|
||||
|
||||
void prepare(const half* d_x,
|
||||
const half* d_vec,
|
||||
const half* d_pe,
|
||||
const void* mod_w,
|
||||
const float* mod_w_scale,
|
||||
const half* mod_b,
|
||||
const void* l1w,
|
||||
const float* l1w_scale,
|
||||
const half* l1b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
int fp8_mask,
|
||||
half* d_q,
|
||||
half* d_k,
|
||||
half* d_v,
|
||||
half* d_mlp,
|
||||
half* d_gate,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
void finish(const half* d_x,
|
||||
const half* d_attn,
|
||||
half* d_mlp,
|
||||
const half* d_gate,
|
||||
half* d_out,
|
||||
const void* l2w,
|
||||
const float* l2w_scale,
|
||||
const half* l2b,
|
||||
int fp8_mask,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
void forward(const half* d_x,
|
||||
const half* d_vec,
|
||||
half* d_out,
|
||||
const half* d_pe,
|
||||
const void* mod_w,
|
||||
const float* mod_w_scale,
|
||||
const half* mod_b,
|
||||
const void* l1w,
|
||||
const float* l1w_scale,
|
||||
const half* l1b,
|
||||
const void* l2w,
|
||||
const float* l2w_scale,
|
||||
const half* l2b,
|
||||
const half* qns,
|
||||
const half* kns,
|
||||
int fp8_mask,
|
||||
cudaStream_t stream = 0);
|
||||
|
||||
int batch() const { return B_; }
|
||||
int seq() const { return S_; }
|
||||
|
||||
private:
|
||||
int B_;
|
||||
int S_;
|
||||
cublasHandle_t cub_;
|
||||
void* workspace_;
|
||||
size_t workspace_bytes_;
|
||||
|
||||
half* d_silu_;
|
||||
half* d_mod_;
|
||||
half* d_xmod_;
|
||||
half* d_l1o_;
|
||||
__nv_fp8_e4m3* d_act_fp8_;
|
||||
unsigned int* d_act_amax_;
|
||||
float* d_act_scale_;
|
||||
float* d_out_scale_;
|
||||
};
|
||||
|
||||
} // namespace transformer
|
||||
|
||||
#endif // SINGLE_STREAM_BLOCK_H
|
||||
+334
@@ -0,0 +1,334 @@
|
||||
"""
|
||||
FLUXNATION — Fused CUDA kernel for FLUX.1 SingleStreamBlock.
|
||||
|
||||
Architecture: PyTorch C++ Extension via torch.ops (no ctypes)
|
||||
- torch.ops.flux_ext.ssb_prepare: modulation + L1 GEMM + split QKV + RoPE
|
||||
- torch.ops.flux_ext.ssb_finish: L2 GEMM + fused gate+residual+bias
|
||||
|
||||
Build: run build_ext.bat (or: py -3.10 setup_ext.py install)
|
||||
"""
|
||||
|
||||
class FLUXNATION:
|
||||
"""Main FLUXNATION kernel node. Pass-through MODEL — kernel is already active."""
|
||||
CATEGORY = "FLUXNATION"
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "apply"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"model": ("MODEL",)}}
|
||||
|
||||
def apply(self, model):
|
||||
from .FLUXPATCH import _spike_controller
|
||||
spike = "SPIKE ON" if _spike_controller.enabled else "SPIKE OFF"
|
||||
print(f"[FLUXNATION] FP8 fused CUDA kernel ACTIVE | {spike}")
|
||||
return (model,)
|
||||
|
||||
|
||||
class FLUXATTENHUT_Node:
|
||||
"""Configure spike attention parameters."""
|
||||
CATEGORY = "FLUXNATION"
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "configure"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"switch_pct": ("FLOAT", {"default": 0.50, "min": 0.1, "max": 0.9, "step": 0.05,
|
||||
"tooltip": "Fraction of steps before spike activates (0.50 = step 10/20)"}),
|
||||
"cap_pct": ("FLOAT", {"default": 0.45, "min": 0.1, "max": 0.9, "step": 0.05,
|
||||
"tooltip": "Fraction of attention blocks to keep (0.45 = 45%)"}),
|
||||
"tau": ("FLOAT", {"default": 0.05, "min": 0.01, "max": 0.5, "step": 0.01,
|
||||
"tooltip": "Block scoring threshold — lower = more blocks kept"}),
|
||||
}
|
||||
}
|
||||
|
||||
def configure(self, model, switch_pct, cap_pct, tau):
|
||||
from .FLUXPATCH import _spike_controller
|
||||
_spike_controller.switch_pct = switch_pct
|
||||
_spike_controller.event_cap_pct = cap_pct
|
||||
_spike_controller.tau = tau
|
||||
print(f"[FLUXATTENHUT] switch={switch_pct*100:.0f}% cap={cap_pct*100:.0f}% tau={tau}")
|
||||
return (model,)
|
||||
|
||||
|
||||
class FLUXSPIKEY_Node:
|
||||
"""Enable or disable spike attention + step caching."""
|
||||
CATEGORY = "FLUXNATION"
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "configure"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"spike": ("BOOLEAN", {"default": True, "tooltip": "Enable block-sparse spike attention"}),
|
||||
"step_cache": ("BOOLEAN", {"default": True, "tooltip": "Cache attention output on alternating spike steps"}),
|
||||
}
|
||||
}
|
||||
|
||||
def configure(self, model, spike, step_cache):
|
||||
from .FLUXPATCH import _spike_controller
|
||||
_spike_controller.enabled = spike
|
||||
_spike_controller.cache_enabled = step_cache
|
||||
print(f"[FLUXSPIKEY] spike={'ON' if spike else 'OFF'} cache={'ON' if step_cache else 'OFF'}")
|
||||
return (model,)
|
||||
|
||||
|
||||
class ULT7RAFLUX:
|
||||
"""Master status node — shows kernel state, spike config, and timing info."""
|
||||
CATEGORY = "FLUXNATION"
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "status"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"spike": ("BOOLEAN", {"default": True}),
|
||||
"step_cache": ("BOOLEAN", {"default": True}),
|
||||
"switch_pct": ("FLOAT", {"default": 0.50, "min": 0.1, "max": 0.9, "step": 0.05}),
|
||||
"cap_pct": ("FLOAT", {"default": 0.45, "min": 0.1, "max": 0.9, "step": 0.05}),
|
||||
"tau": ("FLOAT", {"default": 0.05, "min": 0.01, "max": 0.5, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
def status(self, model, spike, step_cache, switch_pct, cap_pct, tau):
|
||||
from .FLUXPATCH import _spike_controller
|
||||
_spike_controller.enabled = spike
|
||||
_spike_controller.cache_enabled = step_cache
|
||||
_spike_controller.switch_pct = switch_pct
|
||||
_spike_controller.event_cap_pct = cap_pct
|
||||
_spike_controller.tau = tau
|
||||
print(
|
||||
f"[ULT7RAFLUX] FP8 fused kernel | "
|
||||
f"spike={'ON' if spike else 'OFF'} | "
|
||||
f"cache={'ON' if step_cache else 'OFF'} | "
|
||||
f"switch={switch_pct*100:.0f}% cap={cap_pct*100:.0f}% tau={tau}"
|
||||
)
|
||||
return (model,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FLUXNATION": FLUXNATION,
|
||||
"FLUXATTENHUT": FLUXATTENHUT_Node,
|
||||
"FLUXSPIKEY": FLUXSPIKEY_Node,
|
||||
"ULT7RAFLUX": ULT7RAFLUX,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FLUXNATION": "FLUXNATION",
|
||||
"FLUXATTENHUT": "FLUXATTENHUT",
|
||||
"FLUXSPIKEY": "FLUXSPIKEY",
|
||||
"ULT7RAFLUX": "ULT7RAFLUX",
|
||||
}
|
||||
|
||||
_EXT_LOADED = False
|
||||
|
||||
|
||||
def _patch_safetensor_loading():
|
||||
"""Fix Windows mmap access violation in safetensors loading.
|
||||
|
||||
PyTorch 2.5.x + safetensors safe_open() on Windows triggers access violations
|
||||
when memory-mapping large files (>10GB). This patches ComfyUI's load_torch_file
|
||||
to use safetensors.torch.load_file() which loads the full file without mmap.
|
||||
"""
|
||||
import sys
|
||||
if sys.platform != "win32":
|
||||
return
|
||||
|
||||
try:
|
||||
import comfy.utils as comfy_utils
|
||||
import safetensors.torch
|
||||
import torch
|
||||
|
||||
_original_load = comfy_utils.load_torch_file
|
||||
|
||||
def _read_safetensor_streaming(path, target_device='cpu'):
|
||||
"""Stream-load a safetensor file without mmap or full-file buffering.
|
||||
|
||||
Reads the header to get tensor metadata, then reads each tensor
|
||||
individually from disk into a pre-allocated torch tensor.
|
||||
"""
|
||||
import struct, json, numpy as np
|
||||
|
||||
dtype_map = {
|
||||
'F16': (torch.float16, 2),
|
||||
'BF16': (torch.bfloat16, 2),
|
||||
'F32': (torch.float32, 4),
|
||||
'F64': (torch.float64, 8),
|
||||
'I8': (torch.int8, 1),
|
||||
'I16': (torch.int16, 2),
|
||||
'I32': (torch.int32, 4),
|
||||
'I64': (torch.int64, 8),
|
||||
'U8': (torch.uint8, 1),
|
||||
'BOOL': (torch.bool, 1),
|
||||
'F8_E4M3': (getattr(torch, 'float8_e4m3fn', torch.uint8), 1),
|
||||
'F8_E5M2': (getattr(torch, 'float8_e5m2', torch.uint8), 1),
|
||||
}
|
||||
np_dtype_map = {
|
||||
'F16': np.float16, 'BF16': np.void, 'F32': np.float32,
|
||||
'F64': np.float64, 'I8': np.int8, 'I16': np.int16,
|
||||
'I32': np.int32, 'I64': np.int64, 'U8': np.uint8,
|
||||
'BOOL': np.bool_,
|
||||
}
|
||||
|
||||
with open(path, 'rb') as f:
|
||||
header_size = struct.unpack('<Q', f.read(8))[0]
|
||||
header_bytes = f.read(header_size)
|
||||
data_offset = 8 + header_size
|
||||
|
||||
header = json.loads(header_bytes)
|
||||
metadata = header.pop('__metadata__', {})
|
||||
|
||||
sd = {}
|
||||
for name, info in header.items():
|
||||
dtype_str = info['dtype']
|
||||
shape = info['shape']
|
||||
offsets = info['data_offsets']
|
||||
start, end = offsets[0], offsets[1]
|
||||
nbytes = end - start
|
||||
|
||||
torch_dtype, elem_size = dtype_map.get(dtype_str, (torch.float32, 4))
|
||||
|
||||
f.seek(data_offset + start)
|
||||
raw = f.read(nbytes)
|
||||
|
||||
if dtype_str == 'BF16':
|
||||
# numpy doesn't support bfloat16 — read as uint16 then view
|
||||
arr = np.frombuffer(raw, dtype=np.uint16).reshape(shape)
|
||||
t = torch.from_numpy(arr.copy()).view(torch.bfloat16)
|
||||
elif dtype_str in ('F8_E4M3', 'F8_E5M2'):
|
||||
arr = np.frombuffer(raw, dtype=np.uint8).reshape(shape)
|
||||
t = torch.from_numpy(arr.copy()).view(torch_dtype)
|
||||
else:
|
||||
np_dt = np_dtype_map.get(dtype_str, np.float32)
|
||||
arr = np.frombuffer(raw, dtype=np_dt).reshape(shape)
|
||||
t = torch.from_numpy(arr.copy())
|
||||
|
||||
if target_device != 'cpu':
|
||||
t = t.to(device=target_device)
|
||||
sd[name] = t
|
||||
|
||||
return sd, metadata
|
||||
|
||||
def _safe_load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
|
||||
if not ckpt.lower().endswith((".safetensors", ".sft")):
|
||||
return _original_load(ckpt, safe_load=safe_load, device=device,
|
||||
return_metadata=return_metadata)
|
||||
|
||||
import os
|
||||
size_gb = os.path.getsize(ckpt) / (1024**3)
|
||||
|
||||
# Only use streaming loader for large files (>1GB) to avoid mmap crash
|
||||
if size_gb > 1.0:
|
||||
print(f"[FLUXNATION] Loading {os.path.basename(ckpt)} ({size_gb:.1f} GB) via streaming loader")
|
||||
try:
|
||||
import time
|
||||
t0 = time.time()
|
||||
sd, metadata = _read_safetensor_streaming(ckpt, device.type)
|
||||
elapsed = time.time() - t0
|
||||
print(f"[FLUXNATION] Loaded {len(sd)} tensors in {elapsed:.1f}s (streaming, no-mmap)")
|
||||
if return_metadata:
|
||||
return sd, metadata
|
||||
return sd
|
||||
except Exception as e:
|
||||
print(f"[FLUXNATION] Streaming load failed ({e}), falling back to original")
|
||||
return _original_load(ckpt, safe_load=safe_load, device=device,
|
||||
return_metadata=return_metadata)
|
||||
else:
|
||||
return _original_load(ckpt, safe_load=safe_load, device=device,
|
||||
return_metadata=return_metadata)
|
||||
|
||||
comfy_utils.load_torch_file = _safe_load_torch_file
|
||||
print("[FLUXNATION] Patched safetensor loading (non-mmap, Windows crash fix)")
|
||||
except Exception as e:
|
||||
print(f"[FLUXNATION] Could not patch safetensor loading: {e}")
|
||||
|
||||
|
||||
def _try_load_extension():
|
||||
"""Try to import the compiled flux_ext extension."""
|
||||
global _EXT_LOADED
|
||||
try:
|
||||
import flux_ext # noqa: F401 — triggers TORCH_LIBRARY registration
|
||||
_EXT_LOADED = True
|
||||
print("[FLUXNATION] Loaded PyTorch C++ extension successfully")
|
||||
except ImportError:
|
||||
import torch
|
||||
import os
|
||||
root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||||
for candidate in [
|
||||
os.path.join(root, "flux_ext.pyd"),
|
||||
os.path.join(root, "flux_ext.so"),
|
||||
]:
|
||||
if os.path.exists(candidate):
|
||||
try:
|
||||
torch.ops.load_library(candidate)
|
||||
_EXT_LOADED = True
|
||||
print(f"[FLUXNATION] Loaded from {candidate}")
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"[FLUXNATION] Failed to load {candidate}: {e}")
|
||||
|
||||
if not _EXT_LOADED:
|
||||
print("[FLUXNATION] Extension not found. Run build_ext.bat to compile.")
|
||||
return False
|
||||
|
||||
_register_fake_kernels()
|
||||
return True
|
||||
|
||||
|
||||
def _register_fake_kernels():
|
||||
"""Register Meta/FakeTensor implementations for torch.compile."""
|
||||
import torch
|
||||
|
||||
if not hasattr(torch, 'library') or not hasattr(torch.library, 'register_fake'):
|
||||
print("[FLUXNATION] torch.library.register_fake not available, skipping")
|
||||
return
|
||||
|
||||
try:
|
||||
@torch.library.register_fake("flux_ext::ssb_prepare")
|
||||
def _fake_prepare(x, vec, pe, mod_w, mod_w_scale, mod_b,
|
||||
l1w, l1w_scale, l1b, qns, kns):
|
||||
B, S = x.size(0), x.size(1)
|
||||
q = torch.empty([B, 24, S, 128], dtype=torch.float16, device=x.device)
|
||||
k = torch.empty([B, 24, S, 128], dtype=torch.float16, device=x.device)
|
||||
v = torch.empty([B, 24, S, 128], dtype=torch.float16, device=x.device)
|
||||
mlp = torch.empty([B, S, 12288], dtype=torch.float16, device=x.device)
|
||||
gate = torch.empty([B, 3072], dtype=torch.float16, device=x.device)
|
||||
return (q, k, v, mlp, gate)
|
||||
|
||||
@torch.library.register_fake("flux_ext::ssb_finish")
|
||||
def _fake_finish(x, attn, mlp, gate, l2w, l2w_scale, l2b):
|
||||
return torch.empty_like(x, dtype=torch.bfloat16)
|
||||
|
||||
@torch.library.register_fake("flux_ext::ssb_cleanup")
|
||||
def _fake_cleanup():
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FLUXNATION] Warning: could not register fake kernels: {e}")
|
||||
|
||||
|
||||
# ── Initialization: patch safetensor loading + load extension ──
|
||||
import os as _os
|
||||
_patch_safetensor_loading()
|
||||
|
||||
if _os.environ.get("FLUX_EXT_DISABLE", "0") == "1":
|
||||
print("[FLUXNATION] DISABLED by FLUX_EXT_DISABLE env var (stock mode)")
|
||||
elif _try_load_extension():
|
||||
from .FLUXPATCH import apply_patch
|
||||
apply_patch()
|
||||
else:
|
||||
print("[FLUXNATION] DISABLED — extension not found. Run build_ext.bat.")
|
||||
@@ -0,0 +1,23 @@
|
||||
@echo off
|
||||
setlocal
|
||||
|
||||
REM Use VS 2022 (CUDA 12.8 supports MSVC 14.3x from VS 2022, not MSVC 14.5x from VS 2025)
|
||||
set "VSVARS=C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat"
|
||||
if not exist "%VSVARS%" (
|
||||
echo [flux_ext] ERROR: VS 2022 vcvars64.bat not found
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
call "%VSVARS%" || exit /b 1
|
||||
|
||||
set DISTUTILS_USE_SDK=1
|
||||
|
||||
echo [flux_ext] Building PyTorch C++ extension with VS 2022 + CUDA 12.8...
|
||||
py -3.10 setup_ext.py install
|
||||
if errorlevel 1 (
|
||||
echo [flux_ext] Build FAILED
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
echo [flux_ext] Build succeeded! flux_ext is now installed.
|
||||
echo [flux_ext] Restart ComfyUI to use the new torch.ops path.
|
||||
@@ -0,0 +1,271 @@
|
||||
/*
|
||||
* flux_ext_kernels.cu — PyTorch C++ extension for FLUXNATION
|
||||
*
|
||||
* Wraps SSBEngine (from flux_opt.cu) as torch.ops custom operators.
|
||||
* Adds fused gate+residual+bias kernel to eliminate Python-side overhead.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
// SSBEngine class + constants from existing code
|
||||
#include "SINGLE_STREAM_BLOCK.H"
|
||||
|
||||
#include <unordered_map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
|
||||
using namespace transformer;
|
||||
|
||||
// =========================================================================
|
||||
// Vectorized load/store types (same as flux_opt.cu but local to this TU)
|
||||
// =========================================================================
|
||||
union VecBF16_8_ext {
|
||||
float4 f4;
|
||||
__nv_bfloat16 b[8];
|
||||
};
|
||||
union VecHalf8_ext {
|
||||
float4 f4;
|
||||
half h[8];
|
||||
};
|
||||
|
||||
// =========================================================================
|
||||
// NEW: Fused gate + residual + L2-bias kernel
|
||||
//
|
||||
// result[b,s,h] = x[b,s,h] + gate[b,h] * (l2_out[b,s,h] + l2b[h])
|
||||
//
|
||||
// All reads/writes vectorized 8-wide. Eliminates 5 intermediate float32
|
||||
// tensors that the Python path previously created per block.
|
||||
// =========================================================================
|
||||
__global__ void gate_residual_bias_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ d_l2out, // [rows, HIDDEN] L2 GEMM output (bf16)
|
||||
const __nv_bfloat16* __restrict__ d_x, // [rows, HIDDEN] original input / residual (bf16)
|
||||
const half* __restrict__ d_gate, // [B, HIDDEN] gate values (fp16)
|
||||
const half* __restrict__ l2b, // [HIDDEN] L2 bias (fp16)
|
||||
__nv_bfloat16* __restrict__ d_result, // [rows, HIDDEN] output (bf16)
|
||||
int S,
|
||||
int total_elems)
|
||||
{
|
||||
int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
|
||||
if (base >= total_elems) return;
|
||||
|
||||
int row = base / HIDDEN;
|
||||
int col = base % HIDDEN;
|
||||
int batch = row / S;
|
||||
|
||||
// Vectorized loads
|
||||
VecBF16_8_ext out_v = *reinterpret_cast<const VecBF16_8_ext*>(&d_l2out[base]);
|
||||
VecBF16_8_ext x_v = *reinterpret_cast<const VecBF16_8_ext*>(&d_x[base]);
|
||||
VecHalf8_ext gate_v = *reinterpret_cast<const VecHalf8_ext*>(&d_gate[batch * HIDDEN + col]);
|
||||
VecHalf8_ext bias_v = *reinterpret_cast<const VecHalf8_ext*>(&l2b[col]);
|
||||
|
||||
VecBF16_8_ext result_v;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i++) {
|
||||
float l2_f = __bfloat162float(out_v.b[i]);
|
||||
float x_f = __bfloat162float(x_v.b[i]);
|
||||
float g_f = __half2float(gate_v.h[i]);
|
||||
float b_f = __half2float(bias_v.h[i]);
|
||||
result_v.b[i] = __float2bfloat16(x_f + g_f * (l2_f + b_f));
|
||||
}
|
||||
*reinterpret_cast<VecBF16_8_ext*>(&d_result[base]) = result_v;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Global engine cache — keyed by (device, batch, seq_len)
|
||||
// =========================================================================
|
||||
struct EngineEntry {
|
||||
std::unique_ptr<SSBEngine> engine;
|
||||
at::Tensor workspace; // prevents deallocation
|
||||
};
|
||||
|
||||
static std::unordered_map<int64_t, EngineEntry> g_engine_cache;
|
||||
static std::mutex g_engine_mutex;
|
||||
|
||||
static int64_t make_engine_key(int device, int batch, int seq) {
|
||||
return (int64_t)device * 1000000000LL + (int64_t)batch * 1000000LL + seq;
|
||||
}
|
||||
|
||||
static SSBEngine* get_or_create_engine(int batch, int seq, at::Device device) {
|
||||
int dev_idx = device.index();
|
||||
int64_t key = make_engine_key(dev_idx, batch, seq);
|
||||
|
||||
std::lock_guard<std::mutex> lock(g_engine_mutex);
|
||||
auto it = g_engine_cache.find(key);
|
||||
if (it != g_engine_cache.end()) {
|
||||
return it->second.engine.get();
|
||||
}
|
||||
|
||||
// Create new engine
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
size_t ws_bytes = ssb_workspace_bytes(batch, seq);
|
||||
auto ws_tensor = at::empty({(int64_t)ws_bytes}, at::TensorOptions().dtype(at::kByte).device(device));
|
||||
|
||||
auto engine = std::make_unique<SSBEngine>(batch, seq, ws_tensor.data_ptr(), ws_bytes);
|
||||
auto* ptr = engine.get();
|
||||
|
||||
EngineEntry entry;
|
||||
entry.engine = std::move(engine);
|
||||
entry.workspace = ws_tensor;
|
||||
g_engine_cache[key] = std::move(entry);
|
||||
|
||||
fprintf(stderr, "[flux_ext] Created engine B=%d S=%d ws=%.1fMB dev=%d\n",
|
||||
batch, seq, ws_bytes / (1024.0 * 1024.0), dev_idx);
|
||||
fflush(stderr);
|
||||
|
||||
return ptr;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// torch.ops implementation: ssb_prepare
|
||||
//
|
||||
// Runs: silu → mod GEMM → layernorm_mod → linear1 GEMM → split_qkv →
|
||||
// copy_mlp → extract_gate → rmsnorm → rope
|
||||
//
|
||||
// Returns: (q, k, v, mlp, gate)
|
||||
// =========================================================================
|
||||
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
|
||||
ssb_prepare_impl(
|
||||
at::Tensor x, // [B, S, 3072] bf16
|
||||
at::Tensor vec, // [B, 3072] or [B, 1, 3072] fp16
|
||||
at::Tensor pe, // [S, 64, 2] fp16 (pre-converted)
|
||||
at::Tensor mod_w, // FP8 [MODDIM, HIDDEN]
|
||||
at::Tensor mod_w_scale, // float32 [1]
|
||||
at::Tensor mod_b, // fp16 [MODDIM]
|
||||
at::Tensor l1w, // FP8 [LIN1DIM, HIDDEN]
|
||||
at::Tensor l1w_scale, // float32 [1]
|
||||
at::Tensor l1b, // fp16 [LIN1DIM]
|
||||
at::Tensor qns, // fp16 [128]
|
||||
at::Tensor kns) // fp16 [128]
|
||||
{
|
||||
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
|
||||
TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");
|
||||
|
||||
int B = x.size(0);
|
||||
int S = x.size(1);
|
||||
auto device = x.device();
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(device.index());
|
||||
|
||||
// Ensure correct dtypes and contiguity
|
||||
auto x_bf16 = x.to(at::kBFloat16).contiguous();
|
||||
auto vec_fp16 = vec.dim() == 3 ? vec.squeeze(1).to(at::kHalf).contiguous()
|
||||
: vec.to(at::kHalf).contiguous();
|
||||
auto pe_fp16 = pe.to(at::kHalf).contiguous();
|
||||
|
||||
// Allocate output tensors
|
||||
auto opts_fp16 = at::TensorOptions().dtype(at::kHalf).device(device);
|
||||
auto q = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto k = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto v = at::empty({B, HEADS, S, HDIM}, opts_fp16);
|
||||
auto mlp = at::empty({B, S, FFDIM}, opts_fp16);
|
||||
auto gate = at::empty({B, HIDDEN}, opts_fp16);
|
||||
|
||||
// Get or create engine (manages workspace internally)
|
||||
SSBEngine* engine = get_or_create_engine(B, S, device);
|
||||
|
||||
// Call prepare — all pointer casts match existing kernel expectations
|
||||
engine->prepare(
|
||||
reinterpret_cast<const half*>(x_bf16.data_ptr()),
|
||||
reinterpret_cast<const half*>(vec_fp16.data_ptr()),
|
||||
reinterpret_cast<const half*>(pe_fp16.data_ptr()),
|
||||
mod_w.data_ptr(), // FP8 weights (void*)
|
||||
mod_w_scale.defined() ? reinterpret_cast<const float*>(mod_w_scale.data_ptr()) : nullptr,
|
||||
reinterpret_cast<const half*>(mod_b.data_ptr()),
|
||||
l1w.data_ptr(), // FP8 weights (void*)
|
||||
l1w_scale.defined() ? reinterpret_cast<const float*>(l1w_scale.data_ptr()) : nullptr,
|
||||
reinterpret_cast<const half*>(l1b.data_ptr()),
|
||||
reinterpret_cast<const half*>(qns.data_ptr()),
|
||||
reinterpret_cast<const half*>(kns.data_ptr()),
|
||||
7, // fp8_mask = all FP8
|
||||
reinterpret_cast<half*>(q.data_ptr()),
|
||||
reinterpret_cast<half*>(k.data_ptr()),
|
||||
reinterpret_cast<half*>(v.data_ptr()),
|
||||
reinterpret_cast<half*>(mlp.data_ptr()),
|
||||
reinterpret_cast<half*>(gate.data_ptr()),
|
||||
stream);
|
||||
|
||||
return std::make_tuple(q, k, v, mlp, gate);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// torch.ops implementation: ssb_finish
|
||||
//
|
||||
// Runs: gelu → concat → L2 GEMM → FUSED gate+residual+bias
|
||||
//
|
||||
// Returns: result [B, S, 3072] bf16
|
||||
// =========================================================================
|
||||
at::Tensor ssb_finish_impl(
|
||||
at::Tensor x, // [B, S, 3072] bf16 — original input for residual
|
||||
at::Tensor attn, // [B, S, 3072] fp16 — attention output
|
||||
at::Tensor mlp, // [B, S, 12288] fp16 — MLP intermediate (pre-GELU)
|
||||
at::Tensor gate, // [B, 3072] fp16 — gate values
|
||||
at::Tensor l2w, // FP8 [HIDDEN, LIN2DIM]
|
||||
at::Tensor l2w_scale, // float32 [1]
|
||||
at::Tensor l2b) // fp16 [3072] — L2 bias
|
||||
{
|
||||
TORCH_CHECK(x.is_cuda(), "x must be CUDA");
|
||||
TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");
|
||||
|
||||
int B = x.size(0);
|
||||
int S = x.size(1);
|
||||
int rows = B * S;
|
||||
auto device = x.device();
|
||||
at::cuda::CUDAGuard guard(device);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(device.index());
|
||||
|
||||
auto x_bf16 = x.to(at::kBFloat16).contiguous();
|
||||
auto attn_fp16 = attn.to(at::kHalf).contiguous();
|
||||
// mlp is already fp16 from prepare — ensure contiguous
|
||||
auto mlp_c = mlp.contiguous();
|
||||
auto gate_c = gate.contiguous();
|
||||
auto l2b_c = l2b.to(at::kHalf).contiguous();
|
||||
|
||||
// Temp buffer for L2 GEMM output (bf16) — ATen caching allocator makes this free
|
||||
auto l2_out = at::empty({B, S, HIDDEN}, at::TensorOptions().dtype(at::kBFloat16).device(device));
|
||||
|
||||
SSBEngine* engine = get_or_create_engine(B, S, device);
|
||||
|
||||
// Call existing finish() — does gelu, concat, L2 GEMM
|
||||
// NOTE: finish() FP8 path writes bf16 to d_out, does NOT add bias/gate/residual
|
||||
engine->finish(
|
||||
reinterpret_cast<const half*>(x_bf16.data_ptr()), // d_x (bf16 in half*)
|
||||
reinterpret_cast<const half*>(attn_fp16.data_ptr()), // d_attn
|
||||
reinterpret_cast<half*>(mlp_c.data_ptr()), // d_mlp (in/out)
|
||||
reinterpret_cast<const half*>(gate_c.data_ptr()), // d_gate
|
||||
reinterpret_cast<half*>(l2_out.data_ptr()), // d_out (bf16)
|
||||
l2w.data_ptr(), // FP8 weights
|
||||
l2w_scale.defined() ? reinterpret_cast<const float*>(l2w_scale.data_ptr()) : nullptr,
|
||||
reinterpret_cast<const half*>(l2b_c.data_ptr()), // l2b (unused by finish FP8 path)
|
||||
7, // fp8_mask = all FP8
|
||||
stream);
|
||||
|
||||
// FUSED: gate + residual + L2 bias in one kernel
|
||||
// result = x + gate * (l2_out + l2b)
|
||||
auto result = at::empty_like(x_bf16);
|
||||
int total_elems = rows * HIDDEN;
|
||||
int threads = 256;
|
||||
int blocks = (total_elems + threads * 8 - 1) / (threads * 8);
|
||||
|
||||
gate_residual_bias_bf16_kernel<<<blocks, threads, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(l2_out.data_ptr()),
|
||||
reinterpret_cast<const __nv_bfloat16*>(x_bf16.data_ptr()),
|
||||
reinterpret_cast<const half*>(gate_c.data_ptr()),
|
||||
reinterpret_cast<const half*>(l2b_c.data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16*>(result.data_ptr()),
|
||||
S,
|
||||
total_elems);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Cleanup function — destroys all cached engines
|
||||
// =========================================================================
|
||||
void ssb_cleanup_impl() {
|
||||
std::lock_guard<std::mutex> lock(g_engine_mutex);
|
||||
g_engine_cache.clear();
|
||||
fprintf(stderr, "[flux_ext] All engines destroyed\n");
|
||||
fflush(stderr);
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
/*
|
||||
* 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";
|
||||
}
|
||||
+1418
File diff suppressed because it is too large
Load Diff
+24
@@ -0,0 +1,24 @@
|
||||
"""
|
||||
ComfyUI Manager install script.
|
||||
Runs build_ext.bat to compile the CUDA extension after clone.
|
||||
Requires: CUDA Toolkit 12.x, Visual Studio 2022 (Windows) or gcc/nvcc (Linux).
|
||||
"""
|
||||
import subprocess
|
||||
import sys
|
||||
import os
|
||||
|
||||
node_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
if sys.platform == "win32":
|
||||
build_script = os.path.join(node_dir, "build_ext.bat")
|
||||
if os.path.exists(build_script):
|
||||
print("[FLUXNATION] Building CUDA extension...")
|
||||
result = subprocess.run([build_script], cwd=node_dir, shell=True)
|
||||
if result.returncode == 0:
|
||||
print("[FLUXNATION] Build succeeded.")
|
||||
else:
|
||||
print("[FLUXNATION] Build failed. Install CUDA Toolkit 12.x and Visual Studio 2022, then run build_ext.bat manually.")
|
||||
else:
|
||||
print("[FLUXNATION] build_ext.bat not found.")
|
||||
else:
|
||||
print("[FLUXNATION] Linux build: run 'python setup_ext.py install' from the node directory.")
|
||||
@@ -0,0 +1,9 @@
|
||||
set CUDA_LAUNCH_BLOCKING=0
|
||||
set TORCH_CUDNN_V8_API_ENABLED=1
|
||||
set FLUX_SPIKE=1
|
||||
set FLUX_SPIKE_SWITCH=0.50
|
||||
set FLUX_SPIKE_CAP=0.45
|
||||
set FLUX_SPIKE_TAU=0.05
|
||||
set FLUX_SPIKE_CACHE=1
|
||||
py -3.10 -s ComfyUI\main.py --windows-standalone-build --fast --dont-upcast-attention --disable-cuda-malloc --disable-mmap
|
||||
pause
|
||||
@@ -0,0 +1,15 @@
|
||||
@echo off
|
||||
:: FLUXNATION Low-VRAM Profile
|
||||
:: For GPUs with 8-12GB VRAM (RTX 3060, 3070, etc.)
|
||||
:: Aggressive spike + step caching to minimize attention memory
|
||||
:: Quality trade-off is real but generation is POSSIBLE where it wasn't before
|
||||
|
||||
set CUDA_LAUNCH_BLOCKING=0
|
||||
set TORCH_CUDNN_V8_API_ENABLED=1
|
||||
set FLUX_SPIKE=1
|
||||
set FLUX_SPIKE_SWITCH=0.15
|
||||
set FLUX_SPIKE_CAP=0.20
|
||||
set FLUX_SPIKE_TAU=0.03
|
||||
set FLUX_SPIKE_CACHE=1
|
||||
py -3.10 -s ComfyUI\main.py --windows-standalone-build --fast --dont-upcast-attention --disable-cuda-malloc --lowvram
|
||||
pause
|
||||
@@ -0,0 +1,55 @@
|
||||
"""
|
||||
Build script for flux_ext PyTorch C++/CUDA extension.
|
||||
|
||||
Usage:
|
||||
py -3.10 setup_ext.py install
|
||||
|
||||
This compiles the CUTLASS FP8 kernels (from flux_opt.cu) into a proper
|
||||
PyTorch extension, eliminating all ctypes overhead.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Use CUDA 12.8 (compatible with PyTorch 2.5.1 which was built with CUDA 12.1)
|
||||
os.environ["CUDA_HOME"] = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8"
|
||||
os.environ["CUDA_PATH"] = os.environ["CUDA_HOME"]
|
||||
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
|
||||
|
||||
ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
setup(
|
||||
name="flux_ext",
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
name="flux_ext",
|
||||
sources=[
|
||||
os.path.join(ROOT, "csrc", "FLUXPATCH.cpp"),
|
||||
os.path.join(ROOT, "csrc", "FLUXNATION.cu"),
|
||||
os.path.join(ROOT, "flux_opt.cu"),
|
||||
],
|
||||
include_dirs=[
|
||||
ROOT, # SINGLE_STREAM_BLOCK.H
|
||||
os.path.join(ROOT, "cutlass", "include"), # CUTLASS headers
|
||||
],
|
||||
extra_compile_args={
|
||||
"cxx": ["/O2", "/std:c++17", "/EHsc", "/MD"],
|
||||
"nvcc": [
|
||||
"-O3",
|
||||
"-std=c++17",
|
||||
"-arch=sm_89",
|
||||
"-DFLUX_TORCH_EXT=1",
|
||||
"--allow-unsupported-compiler",
|
||||
"-diag-suppress=20013",
|
||||
"-diag-suppress=20015",
|
||||
"-Xcompiler", "/MD",
|
||||
"-Xcompiler", "/EHsc",
|
||||
"-Xcompiler", "/Zc:preprocessor",
|
||||
],
|
||||
},
|
||||
libraries=["cublas"],
|
||||
)
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension},
|
||||
)
|
||||
Reference in New Issue
Block a user