483 lines
24 KiB
Python
483 lines
24 KiB
Python
"""Veda sparse attention, the torch half: tiles, plans, predictor, selection, one layer's call.
|
|
|
|
Veda (Han, Yang, Hu, Mei, Jiang, Qi; MIT; arXiv 2605.30325) lets a small distilled predictor say
|
|
which 128-token 3D tiles of the attention map matter, and computes only the top ~10% of them.
|
|
Written here from Veda-on-ComfyUI's core (itself a rewrite of Miowtion's miowtion/veda, MIT), in
|
|
this codebase's style; see VEDA-NOTICE.txt. The released predictors were trained against these
|
|
exact rules (tile order, pooling, budget, own-tile, global rows), so none of them may drift.
|
|
|
|
Pure torch, no ComfyUI import: it runs and is tested on any device. The kernel that attends the
|
|
kept tiles is a callable handed in (veda_kernel.attend on a GPU, `reference_attend` in tests).
|
|
Every tensor made here is built on, or moved to, the device of the tensor it meets.
|
|
"""
|
|
import contextlib
|
|
import functools
|
|
import json
|
|
import math
|
|
import os
|
|
import time
|
|
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
TILE = 128
|
|
FORMAT = "miowtion-veda-predictor-v1"
|
|
_STORAGE = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float8_e4m3fn": torch.float8_e4m3fn}
|
|
FPS = 24
|
|
_CHUNK_BYTES = 256 * 2**20 # one tile-ordered q/k/v/out copy of a head chunk; four are alive at once
|
|
|
|
|
|
class VedaError(ValueError):
|
|
"""The file or layout cannot be used; the message is shown to the person."""
|
|
|
|
|
|
# ── budgets: (ratio, tiles), exactly one set. ratio = share of key tiles KEPT at equal kernel cost ─────────────
|
|
|
|
def sparsity_budget(share):
|
|
"""Skip `share` (0..1) of the key tiles."""
|
|
return (1.0 - float(share), None)
|
|
|
|
|
|
def keeps_all(budget):
|
|
return budget[1] is None and budget[0] >= 1.0
|
|
|
|
|
|
def describe_budget(budget):
|
|
return f"{budget[1]:g} tiles" if budget[1] is not None else f"{100.0 * (1.0 - min(budget[0], 1.0)):g}% sparse"
|
|
|
|
|
|
def _per_row(budget, real_tokens, n_cols):
|
|
if budget[1] is not None:
|
|
return float(budget[1])
|
|
ideal = math.ceil(real_tokens / TILE) # equal cost: a padded tiling never gets a bigger budget
|
|
return budget[0] * ideal * ideal / n_cols
|
|
|
|
|
|
def _split(budget, n_cols):
|
|
"""(k_lo, k_hi, frac) with 1 <= k_lo <= k_hi <= n_cols."""
|
|
if budget >= n_cols:
|
|
return n_cols, n_cols, 0.0
|
|
lo = min(max(1, math.floor(budget)), n_cols)
|
|
return lo, min(lo + 1, n_cols), round(min(max(budget - lo, 0.0), 1.0), 12)
|
|
|
|
|
|
@functools.lru_cache(maxsize=256)
|
|
def _bresenham(n_rows, frac):
|
|
ramp = torch.floor(torch.arange(n_rows + 1, dtype=torch.float64) * frac)
|
|
return ramp[1:] > ramp[:-1]
|
|
|
|
|
|
# ── tiles ──────────────────────────────────────────────────────────────────────────────────────────────────────
|
|
|
|
def shapes():
|
|
"""Every power-of-two (t, h, w) with t * h * w = TILE."""
|
|
return sorted((2 ** i, 2 ** j, TILE >> (i + j)) for i in range(8) for j in range(8 - i))
|
|
|
|
|
|
def _padded(grid, shape):
|
|
return tuple(-(-g // s) * s for g, s in zip(grid, shape))
|
|
|
|
|
|
def _count(grid, shape):
|
|
p = _padded(grid, shape)
|
|
return p[0] * p[1] * p[2] // TILE
|
|
|
|
|
|
def least_padding_shape(grid):
|
|
"""Shape that pads `grid` least (ties: the most cubic, then the first); how reference spans are tiled."""
|
|
fit = [s for s in shapes() if all(a <= b for a, b in zip(s, grid))]
|
|
return min(fit or shapes(), key=lambda s: (_count(grid, s), max(s) / min(s), s))
|
|
|
|
|
|
def span_tiles(start, grid, shape):
|
|
"""[n_tiles, 128] packed row ids of one span, -1 on padding. Box order is (h, w, t) outer to inner, rows
|
|
inside a box are t, h, w row-major, real rows compacted to the front (so a tile's real rows are a prefix)."""
|
|
t, h, w = grid
|
|
tp, hp, wp = _padded(grid, shape)
|
|
cells = torch.full((tp, hp, wp), -1, dtype=torch.long)
|
|
cells[:t, :h, :w] = start + torch.arange(t * h * w).view(t, h, w)
|
|
st, sh, sw = shape
|
|
tiles = cells.view(tp // st, st, hp // sh, sh, wp // sw, sw).permute(2, 4, 0, 1, 3, 5).reshape(-1, TILE)
|
|
return torch.gather(tiles, 1, torch.argsort((tiles < 0).to(torch.int8), dim=1, stable=True))
|
|
|
|
|
|
def build_layout(spans, seq_len, device="cpu"):
|
|
"""One permutation of the packed sequence. `spans` = [(start, grid, shape)] in packed order, the target last;
|
|
every other row is global (text, audio, untiled conditions) and is chunked into tiles after the spans."""
|
|
if not spans:
|
|
raise VedaError("the target span must be tiled")
|
|
covered, tiles, stop = torch.zeros(seq_len, dtype=torch.bool), [], 0
|
|
for start, grid, shape in spans:
|
|
rows = grid[0] * grid[1] * grid[2]
|
|
if start < stop or start + rows > seq_len:
|
|
raise VedaError(f"span at {start} overlaps another or runs past {seq_len}")
|
|
stop = start + rows
|
|
covered[start:stop] = True
|
|
tiles.append(span_tiles(start, grid, shape))
|
|
video = torch.cat(tiles)
|
|
rest = torch.nonzero(~covered).view(-1)
|
|
n_global = -(-rest.numel() // TILE)
|
|
tail = torch.full((n_global * TILE,), -1, dtype=torch.long)
|
|
tail[:rest.numel()] = rest
|
|
perm = torch.cat([video.view(-1), tail])
|
|
count = (perm.view(-1, TILE) >= 0).sum(1).to(torch.int32)
|
|
partial = torch.nonzero((count > 0) & (count < TILE)).view(-1)
|
|
n_video = video.shape[0]
|
|
on = lambda x: x.to(device)
|
|
return SimpleNamespace(
|
|
n_video=n_video, n_ref=sum(t.shape[0] for t in tiles[:-1]), n_global=n_global, n_tiles=n_video + n_global,
|
|
ref_tokens=sum(g[0] * g[1] * g[2] for _, g, _ in spans[:-1]),
|
|
target_tokens=spans[-1][1][0] * spans[-1][1][1] * spans[-1][1][2], seq_len=seq_len,
|
|
widest=int((count > 0).sum()), # an upper bound on the kept tiles of any row (a global row sees every tile): no device read needed
|
|
bres={}, valid_count=on(count), gather_index=on(perm.clamp(min=0)), scatter_index=on(torch.where(perm < 0, seq_len, perm)),
|
|
pad_slots=on(torch.nonzero(perm < 0).view(-1)), partial_video=on(partial[partial < n_video]),
|
|
kv_ok=on(count > 0), slot_valid=on((perm >= 0).to(torch.int32)))
|
|
|
|
|
|
def gather(x, lay, heads):
|
|
"""[S, H, D] -> [N, H', D] in tile order, padding slots zero."""
|
|
out = x[lay.gather_index[:, None], heads[None, :]]
|
|
if lay.pad_slots.numel():
|
|
out.index_fill_(0, lay.pad_slots, 0)
|
|
return out
|
|
|
|
|
|
def scatter_(out, tiled, lay, heads):
|
|
"""Tile-ordered rows back into out[S + 1, H, D]; padding lands in the spare row S."""
|
|
out[lay.scatter_index[:, None], heads[None, :]] = tiled
|
|
|
|
|
|
# ── predictor: pool each video tile to [mean | max | min], score tile pairs ────────────────────────────────────
|
|
|
|
@torch.no_grad()
|
|
def pool(x, lay):
|
|
"""[N, H', D] tile order -> [H', n_video, 3D] fp32 (mean, max, min over each tile's real rows; 0 for empty)."""
|
|
n, (heads, dim) = lay.n_video, x.shape[1:]
|
|
tiles = x[:n * TILE].view(n, TILE, heads, dim)
|
|
mean = tiles.sum(dim=1, dtype=torch.float32) / lay.valid_count[:n].clamp(min=1).to(torch.float32)[:, None, None]
|
|
lo, hi = torch.aminmax(tiles, dim=1)
|
|
part = lay.partial_video
|
|
if part.numel(): # padding rows are 0, a legal value: the partial tiles' extremes are redone masked
|
|
sub = tiles.index_select(0, part)
|
|
live = (torch.arange(TILE, device=x.device)[None, :] < lay.valid_count.index_select(0, part)[:, None])[:, :, None, None]
|
|
hi.index_copy_(0, part, sub.masked_fill(~live, float("-inf")).amax(dim=1))
|
|
lo.index_copy_(0, part, sub.masked_fill(~live, float("inf")).amin(dim=1))
|
|
feats = torch.cat([mean, hi.float(), lo.float()], dim=-1)
|
|
return torch.where(lay.kv_ok[:n, None, None], feats, 0.0).permute(1, 0, 2).contiguous()
|
|
|
|
|
|
@torch.no_grad()
|
|
def logits(fq, fk, pq, pk):
|
|
"""[H', n_q, n_k] fp32 block logits: x_hat = pool @ P + mean, q_hat . k_hat / sqrt(D)."""
|
|
dim = pq.shape[-1]
|
|
qh = torch.bmm(fq, pq.to(device=fq.device, dtype=torch.float32)) + fq[..., :dim]
|
|
kh = torch.bmm(fk, pk.to(device=fk.device, dtype=torch.float32)) + fk[..., :dim]
|
|
return torch.bmm(qh, kh.transpose(1, 2)) / math.sqrt(dim)
|
|
|
|
|
|
# ── selection: top-k per query tile over the video x video quadrant ────────────────────────────────────────────
|
|
|
|
def column_blocks(lay, generated, reference):
|
|
"""[(start, stop, budget, real tokens)]: the reference tiles (when tiled) and the target, each with its own top-k
|
|
(one pooled top-k would starve the references, the target's spatial neighbours always win)."""
|
|
target = (lay.n_ref, lay.n_video, generated, lay.target_tokens)
|
|
return [target] if lay.n_ref == 0 else [(0, lay.n_ref, reference, lay.ref_tokens), target]
|
|
|
|
|
|
@torch.no_grad()
|
|
def select(scores, lay, blocks):
|
|
"""(index [H', n_video, K] int64 key tiles, keep [H', n_video, K] bool). Rules: a query tile always keeps its own
|
|
tile and that uses budget, empty key tiles are never kept, the fractional budget is spread by a Bresenham pattern."""
|
|
heads, nv, dev = scores.shape[0], lay.n_video, scores.device
|
|
rows = torch.arange(nv, device=dev)
|
|
idxs, keeps = [], []
|
|
for start, stop, budget, real in blocks:
|
|
n = stop - start
|
|
if keeps_all(budget):
|
|
idxs.append(torch.arange(start, stop, device=dev).expand(heads, nv, n))
|
|
keeps.append(lay.kv_ok[start:stop].expand(heads, nv, n))
|
|
continue
|
|
s = scores[:, :, start:stop].float().clone()
|
|
s.masked_fill_(~lay.kv_ok[None, None, start:stop], float("-inf"))
|
|
own = ((rows >= start) & (rows < stop))[None, :, None]
|
|
col = (rows - start).clamp(0, n - 1)[None, :, None].expand(heads, nv, 1)
|
|
s.scatter_(-1, col, torch.where(own, float("inf"), s.gather(-1, col)))
|
|
lo, hi, frac = _split(_per_row(budget, real, n), n)
|
|
vals, idx = torch.topk(s, hi, dim=-1, sorted=True)
|
|
extra = lay.bres.get((nv, float(frac), str(dev))) # built on the host once, then kept on the device with the layout
|
|
if extra is None:
|
|
extra = lay.bres[(nv, float(frac), str(dev))] = _bresenham(nv, float(frac)).to(device=dev, dtype=torch.long)
|
|
allowed = lo + extra
|
|
idxs.append(idx + start)
|
|
keeps.append((torch.arange(hi, device=dev)[None, None, :] < allowed[None, :, None]) & (vals > float("-inf")))
|
|
return torch.cat(idxs, -1), torch.cat(keeps, -1)
|
|
|
|
|
|
@torch.no_grad()
|
|
def block_mask(index, keep, lay):
|
|
"""[H', n_tiles, n_tiles] bool, rows are query tiles. Global tiles are dense both ways."""
|
|
n, nv = lay.n_tiles, lay.n_video
|
|
mask = torch.zeros(index.shape[0], n, n, dtype=torch.bool, device=index.device)
|
|
mask[:, :nv].scatter_(2, index, keep)
|
|
mask[:, :, nv:] = lay.kv_ok[nv:]
|
|
mask[:, nv:, :] = lay.kv_ok
|
|
return mask
|
|
|
|
|
|
def tile_index_list(mask, widest=None):
|
|
"""A mask as the per-row list of kept tiles: (index [H', n, widest] int32 ascending, count [H', n] int32). Entries past a row's
|
|
count are never read, so `widest` may be any upper bound (a layout knows one); None reads the exact widest from the device."""
|
|
count = mask.sum(-1, dtype=torch.int32)
|
|
if widest is None:
|
|
widest = int(count.max().item()) if count.numel() else 0
|
|
order = torch.argsort(mask.to(torch.int8), dim=-1, descending=True, stable=True)
|
|
return order[..., :widest].to(torch.int32).contiguous(), count
|
|
|
|
|
|
def reference_attend(q, k, v, mask, lay):
|
|
"""Exact block-sparse attention in fp32 on tile-ordered [N, H', D]: the ground truth and the test stand-in
|
|
for the kernel. O(N^2) memory, never for a real sequence."""
|
|
allowed = mask.repeat_interleave(TILE, 1).repeat_interleave(TILE, 2) & lay.slot_valid.bool()[None, None, :]
|
|
probs = torch.softmax((torch.einsum("qhd,khd->hqk", q.float(), k.float()) * q.shape[-1] ** -0.5)
|
|
.masked_fill(~allowed, float("-inf")), dim=-1).nan_to_num(0.0)
|
|
out = torch.einsum("hqk,khd->qhd", probs, v.float()) * lay.slot_valid.to(torch.float32)[:, None, None]
|
|
return out.to(q.dtype)
|
|
|
|
|
|
# ── what the packed layout holds ────────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _grid_of(position_ids, a, b):
|
|
"""(T, H, W) when rows a:b are T, H, W row-major (checked against the positions), else None."""
|
|
pos = position_ids[a:b]
|
|
axes = [torch.unique(pos[:, i]) for i in range(3)]
|
|
t, h, w = (x.numel() for x in axes)
|
|
if t * h * w != b - a:
|
|
return None
|
|
want = torch.stack(torch.meshgrid(*axes, indexing="ij"), dim=-1).reshape(-1, 3)
|
|
return (t, h, w) if torch.equal(want, pos.to(want.dtype)) else None
|
|
|
|
|
|
def describe_layout(layout):
|
|
"""ComfyUI's H3 PackedLayout -> a hashable (seq_len, (target start, grid), ((ref start, grid), ...), skipped notes).
|
|
References (cond, ref_img) whose rows are not a verified row-major grid stay global, which is always correct."""
|
|
try:
|
|
segments = [(int(a), int(b), kind) for a, b, kind in layout.segments]
|
|
_, lt, lh, lw, _ = layout.signature
|
|
seq_len = int(layout.seq_len)
|
|
except (AttributeError, TypeError, ValueError) as exc:
|
|
raise VedaError(f"not a MiniMax H3 packed layout ({exc})") from exc
|
|
video = [(a, b) for a, b, kind in segments if kind == "video"]
|
|
if len(video) != 1:
|
|
raise VedaError(f"expected one target video segment, found {len(video)}")
|
|
grid = (int(lt), int(lh) // 2, int(lw) // 2)
|
|
if grid[0] * grid[1] * grid[2] != video[0][1] - video[0][0]:
|
|
raise VedaError(f"the video segment does not match the latent grid {grid}")
|
|
pos, refs, skipped = getattr(layout, "position_ids", None), [], []
|
|
for a, b, kind in segments:
|
|
if kind in ("cond", "ref_img"):
|
|
g = None if pos is None else _grid_of(pos, a, b)
|
|
(refs.append((a, g)) if g else skipped.append(f"{kind}[{a}:{b}]"))
|
|
return seq_len, (video[0][0], grid), tuple(refs), tuple(skipped)
|
|
|
|
|
|
# ── tile plans and the predictor file ──────────────────────────────────────────────────────────────────────────
|
|
|
|
def _plan(name, grid, shape_list, head_shape):
|
|
return SimpleNamespace(name=name, grid=tuple(grid), shapes=[tuple(s) for s in shape_list], head_shape=head_shape, groups={})
|
|
|
|
|
|
def _aspect(grid):
|
|
return grid[2] / grid[1]
|
|
|
|
|
|
def describe_grid(grid):
|
|
frames = 17 * ((grid[0] - 2) // 5) + 5 if grid[0] >= 2 else 1 # H3: 17n + 5 frames for a latent length
|
|
return f"{grid[2] * 32}x{grid[1] * 32}, {frames / FPS:.1f} s"
|
|
|
|
|
|
def choose_plan(plans, grid):
|
|
"""(plan, exact, words): the plan with the nearest aspect, then nearest duration, then least padding. The H<->W
|
|
mirrored plans are candidates too. Never fails: a plan of another size still tiles any grid, it is only
|
|
outside what the predictor was trained on, and `words` says so."""
|
|
grid = tuple(int(g) for g in grid)
|
|
mirrored = [_plan(p.name + "_T", (p.grid[0], p.grid[2], p.grid[1]), [(t, w, h) for t, h, w in p.shapes], p.head_shape) for p in plans]
|
|
target = math.log(_aspect(grid))
|
|
best = min(plans + mirrored, key=lambda p: (round(abs(math.log(_aspect(p.grid)) - target), 3), abs(p.grid[0] - grid[0]),
|
|
sum(_count(grid, s) for s in p.shapes)))
|
|
if best.grid == grid:
|
|
return best, True, f"trained for this size ({describe_grid(grid)})"
|
|
return best, False, f"nearest trained size: {describe_grid(best.grid)} (this video: {describe_grid(grid)})"
|
|
|
|
|
|
def head_groups(plan, layer, device="cpu"):
|
|
"""[(shape, heads int64 on `device`)] of one layer: heads that share a tile shape. Cached on the plan, per device."""
|
|
key = (layer, str(device))
|
|
if key not in plan.groups:
|
|
row = torch.tensor(plan.head_shape[layer])
|
|
plan.groups[key] = [(plan.shapes[i], torch.nonzero(row == i).view(-1).to(device)) for i in sorted(set(plan.head_shape[layer]))]
|
|
return plan.groups[key]
|
|
|
|
|
|
def _budget_of(meta, prefix, fallback):
|
|
kind = meta.get(f"{prefix}_budget_kind")
|
|
if kind is None:
|
|
return fallback
|
|
try:
|
|
value = float(meta[f"{prefix}_budget_value"])
|
|
except (KeyError, ValueError) as exc:
|
|
raise VedaError(f"{prefix}_budget_value: {exc}") from exc
|
|
if kind not in ("ratio", "tiles"):
|
|
raise VedaError(f"{prefix}_budget_kind: unknown kind {kind!r}")
|
|
return (value, None) if kind == "ratio" else (None, value)
|
|
|
|
|
|
def read_metadata(path):
|
|
from safetensors import safe_open
|
|
name = os.path.basename(path)
|
|
try:
|
|
with safe_open(path, framework="pt", device="cpu") as f:
|
|
meta = dict(f.metadata() or {})
|
|
except Exception as exc: # noqa: BLE001 - safetensors raises several types
|
|
raise VedaError(f"{name} is not a readable safetensors file ({exc})") from exc
|
|
if meta.get("format") != FORMAT:
|
|
raise VedaError(f"{name} is not a Veda predictor (format {meta.get('format')!r}). Put only Veda predictor files in models/veda.")
|
|
return meta
|
|
|
|
|
|
def load_bundle(path):
|
|
"""The predictor file, read strictly: every tensor must match its declared shape. Projections are held in bf16 on the
|
|
host and follow q's device at call time, so the predictor does not pin VRAM while ComfyUI offloads the model."""
|
|
from safetensors import safe_open
|
|
meta, name = read_metadata(path), os.path.basename(path)
|
|
try:
|
|
layers, heads, dim = int(meta["num_layers"]), int(meta["num_heads"]), int(meta["head_dim"])
|
|
keep = meta.get("keep_ratio")
|
|
generated = _budget_of(meta, "target", None if keep is None else (float(keep), None))
|
|
if generated is None:
|
|
raise KeyError("target_budget_kind")
|
|
reference = _budget_of(meta, "ref", generated)
|
|
plan_json = json.loads(meta["plans"])
|
|
except (KeyError, ValueError) as exc:
|
|
raise VedaError(f"{name}: incomplete metadata ({exc})") from exc
|
|
stored = meta.get("dtype", "float32")
|
|
if stored not in _STORAGE:
|
|
raise VedaError(f"{name}: unknown storage dtype {stored!r}")
|
|
with safe_open(path, framework="pt", device="cpu") as f:
|
|
tensors = {key: f.get_tensor(key) for key in f.keys()}
|
|
shape, proj = (heads, 3 * dim, dim), {"proj_q": [], "proj_k": []}
|
|
for layer in range(layers):
|
|
for kind, out in proj.items():
|
|
key = f"layers.{layer}.{kind}"
|
|
value = tensors.pop(key, None)
|
|
if value is None or tuple(value.shape) != shape:
|
|
raise VedaError(f"{name}: {key} missing or not {list(shape)}")
|
|
if stored == "float8_e4m3fn":
|
|
scale = tensors.pop(key + ".__scale", None)
|
|
if scale is None or tuple(scale.shape) != (heads,):
|
|
raise VedaError(f"{name}: {key} has no per-head fp8 scale")
|
|
value = value.float() * scale.float().view(-1, 1, 1)
|
|
out.append(value.to(torch.bfloat16).contiguous())
|
|
if tensors:
|
|
raise VedaError(f"{name}: unexpected tensors {sorted(tensors)[:4]}")
|
|
plans = [_plan(p["geometry"], p["grid"], [tuple(int(x) for x in s.split("x")) for s in p["shapes"]], [list(r) for r in p["head_shape"]])
|
|
for p in plan_json.values()]
|
|
if not plans:
|
|
raise VedaError(f"{name}: no tile plans")
|
|
for p in plans:
|
|
if len(p.head_shape) != layers or len(p.head_shape[0]) != heads:
|
|
raise VedaError(f"{name}: plan {p.name} does not fit {layers} layers x {heads} heads")
|
|
return SimpleNamespace(path=path, layers=layers, heads=heads, dim=dim, generated=generated, reference=reference,
|
|
plans=plans, proj_q=proj["proj_q"], proj_k=proj["proj_k"], meta=meta)
|
|
|
|
|
|
def describe_bundle(b):
|
|
trained = describe_budget(b.generated) + ("" if b.reference == b.generated else f" / reference {describe_budget(b.reference)}")
|
|
return f"{os.path.basename(b.path)}, {b.layers} layers x {b.heads} heads, trained {trained}"
|
|
|
|
|
|
class PhaseTimer:
|
|
"""Opt-in time per phase, summed over a run and read once at its end. CUDA events on a GPU (the work is asynchronous there);
|
|
plain clock time elsewhere. `with timer("score"): ...`"""
|
|
|
|
def __init__(self, device):
|
|
self.cuda = torch.device(device).type == "cuda"
|
|
self.spans = {}
|
|
|
|
@contextlib.contextmanager
|
|
def __call__(self, name):
|
|
if self.cuda:
|
|
a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
|
|
a.record()
|
|
yield
|
|
b.record()
|
|
self.spans.setdefault(name, []).append((a, b))
|
|
else:
|
|
t = time.perf_counter()
|
|
yield
|
|
self.spans.setdefault(name, []).append(time.perf_counter() - t)
|
|
|
|
def totals(self):
|
|
"""{phase: milliseconds}; reads the device once."""
|
|
if self.cuda:
|
|
torch.cuda.synchronize()
|
|
return {k: sum(a.elapsed_time(b) for a, b in v) for k, v in self.spans.items()}
|
|
return {k: 1000.0 * sum(v) for k, v in self.spans.items()}
|
|
|
|
|
|
_NO_TIMER = lambda name: contextlib.nullcontext()
|
|
|
|
|
|
# ── one layer ───────────────────────────────────────────────────────────────────────────────────────────────────
|
|
|
|
@torch.no_grad()
|
|
def attention(q, k, v, layer, spec, plan, bundle, budgets, attend, cache, stats, timer=_NO_TIMER):
|
|
"""One layer of Veda attention. q, k, v: [S, H, D] (post QK-norm and RoPE, any strides, any device).
|
|
For each group of heads sharing a tile shape, in memory-bounded chunks: gather to tile order, pool, score, select,
|
|
build the block mask, `attend(q_t, k_t, v_t, mask, layout)`, scatter back. Returns [S, H, D] in q's dtype.
|
|
`cache` holds the tile layouts (keyed by geometry, shape and device); `stats` accumulates what was computed."""
|
|
seq_len, heads, dim = q.shape
|
|
dev = q.device
|
|
generated, reference = budgets
|
|
tile_refs = not keeps_all(reference)
|
|
out = q.new_zeros(seq_len + 1, heads, dim)
|
|
pq, pk = bundle.proj_q[layer], bundle.proj_k[layer]
|
|
if dev.type == "cuda" and not pq.is_pinned(): # pinned host memory: the copy below does not stop the host
|
|
try:
|
|
pq = bundle.proj_q[layer] = pq.pin_memory()
|
|
pk = bundle.proj_k[layer] = pk.pin_memory()
|
|
except RuntimeError:
|
|
pass
|
|
pq, pk = pq.to(dev, non_blocking=True), pk.to(dev, non_blocking=True)
|
|
for shape, group in head_groups(plan, layer, dev):
|
|
key = (spec, shape, tile_refs, str(dev))
|
|
if key not in cache:
|
|
spans = [(s, g, least_padding_shape(g)) for s, g in spec[2]] if tile_refs else []
|
|
cache[key] = build_layout(spans + [(spec[1][0], spec[1][1], shape)], seq_len, dev)
|
|
lay = cache[key]
|
|
blocks = column_blocks(lay, generated, reference)
|
|
step = max(1, _CHUNK_BYTES // (lay.n_tiles * TILE * dim * q.element_size()))
|
|
for hs in group.split(step):
|
|
with timer("gather"):
|
|
q_t, k_t = gather(q, lay, hs), gather(k, lay, hs)
|
|
with timer("score"):
|
|
scores = logits(pool(q_t, lay), pool(k_t, lay), pq.index_select(0, hs), pk.index_select(0, hs))
|
|
with timer("select"):
|
|
index, keep = select(scores, lay, blocks)
|
|
del scores
|
|
kept = keep.sum()
|
|
stats["kept"] = stats.get("kept", 0) + kept
|
|
stats["pairs"] = stats.get("pairs", 0) + len(hs) * lay.n_video ** 2
|
|
stats["work"] = stats.get("work", 0) + kept + len(hs) * lay.n_global * (lay.n_tiles + lay.n_video)
|
|
stats["dense"] = stats.get("dense", 0) + len(hs) * (seq_len / TILE) ** 2
|
|
stats["global"] = lay.n_global
|
|
with timer("select"):
|
|
mask = block_mask(index, keep, lay)
|
|
with timer("gather"):
|
|
v_t = gather(v, lay, hs)
|
|
with timer("kernel"): # a kernel that can split its own time takes the timer too; the others are called as before
|
|
o_t = attend(q_t, k_t, v_t, mask, lay, timer=timer) if timer is not _NO_TIMER and getattr(attend, "times_itself", False) \
|
|
else attend(q_t, k_t, v_t, mask, lay)
|
|
with timer("scatter"):
|
|
scatter_(out, o_t, lay, hs)
|
|
return out[:seq_len]
|