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>
101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
from __future__ import annotations
|
|
from typing import Dict, Optional, Tuple, Any
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from . import FLUXSPIKEY as triton_backend
|
|
|
|
|
|
def attend(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
*,
|
|
backend: str = "dense",
|
|
state: Optional[Dict[str, torch.Tensor]] = None,
|
|
neuromorph: Optional[Dict[str, Any]] = None,
|
|
) -> Tuple[torch.Tensor, Dict[str, Any]]:
|
|
if backend == "dense":
|
|
return _attend_dense(q, k, v)
|
|
elif backend == "spike":
|
|
neuromorph_cfg = neuromorph or {}
|
|
if triton_backend.has_triton():
|
|
return triton_backend.attend_spike_triton(
|
|
q, k, v,
|
|
tau=neuromorph_cfg.get("tau", 0.15),
|
|
leak=neuromorph_cfg.get("leak", 0.0),
|
|
event_cap_pct=neuromorph_cfg.get("event_cap_pct", 0.45),
|
|
state=state,
|
|
)
|
|
return _attend_spike(q, k, v, state=state, neuromorph=neuromorph_cfg)
|
|
else:
|
|
raise ValueError(f"Unknown backend: {backend}")
|
|
|
|
|
|
def _attend_dense(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
) -> Tuple[torch.Tensor, Dict[str, Any]]:
|
|
scale = q.shape[-1] ** -0.5
|
|
scores = torch.matmul(q, k.transpose(-2, -1)) * scale
|
|
attn_weights = F.softmax(scores, dim=-1)
|
|
output = torch.matmul(attn_weights, v)
|
|
return output, {"backend": "dense", "attn_weights": attn_weights}
|
|
|
|
|
|
def _attend_spike(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
*,
|
|
state: Optional[Dict[str, torch.Tensor]] = None,
|
|
neuromorph: Optional[Dict[str, Any]] = None,
|
|
) -> Tuple[torch.Tensor, Dict[str, Any]]:
|
|
cfg = neuromorph or {}
|
|
tau = cfg.get("tau", 0.15)
|
|
leak = cfg.get("leak", 0.0)
|
|
event_cap_pct = cfg.get("event_cap_pct", 0.45)
|
|
|
|
batch, heads, seq_len, head_dim = q.shape
|
|
|
|
scale = head_dim ** -0.5
|
|
z = torch.matmul(q, k.transpose(-2, -1)) * scale
|
|
|
|
z_max_per_row = z.max(dim=-1, keepdim=True)[0]
|
|
M_t = (z >= tau * z_max_per_row).float()
|
|
|
|
max_events_per_row = max(1, int(seq_len * event_cap_pct))
|
|
if max_events_per_row < seq_len:
|
|
_, topk_indices = torch.topk(z, k=max_events_per_row, dim=-1, sorted=False)
|
|
M_t_capped = torch.zeros_like(z)
|
|
M_t_capped.scatter_(-1, topk_indices, 1.0)
|
|
M_t = M_t * M_t_capped
|
|
|
|
if state is not None and "S_prev" in state and state["S_prev"].shape == M_t.shape:
|
|
S_t = leak * state["S_prev"] + M_t
|
|
else:
|
|
S_t = M_t
|
|
|
|
S_t_sum = S_t.sum(dim=-1, keepdim=True)
|
|
S_t_sum = torch.where(S_t_sum > 0, S_t_sum, torch.ones_like(S_t_sum))
|
|
attn_weights = (S_t / S_t_sum * M_t).to(v.dtype)
|
|
output = torch.matmul(attn_weights, v)
|
|
|
|
kept = M_t.sum().item()
|
|
total = batch * heads * seq_len * seq_len
|
|
flops_saved = 100.0 * (1.0 - kept / max(1, total))
|
|
|
|
return output, {
|
|
"S_prev": S_t.detach(),
|
|
"M_t": M_t,
|
|
"telemetry": {
|
|
"backend": "spike",
|
|
"kept_events_pct": 100.0 * kept / max(1, total),
|
|
"flops_saved_pct": flops_saved,
|
|
"tau": tau,
|
|
"leak": leak,
|
|
"event_cap_pct": event_cap_pct,
|
|
},
|
|
}
|