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