# 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