ULT7RAandCopilot
82bdde669f
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>
506 lines
19 KiB
Python
506 lines
19 KiB
Python
"""
|
||
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}")
|