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

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:

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:

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