Files

T

ULT7RAandCopilot 82bdde669f 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>

2026-04-05 17:56:53 -06:00

4.7 KiB

Raw Permalink Blame History

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

4.7 KiB Raw Permalink Blame History Unescape Escape