mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
feat(attention): trace contiguity and count calls per route
A strided query or key changes what a backend's input prep costs, so the route trace records it and keeps strided and contiguous calls as separate routes. The trace also counts calls per route and logs the counts when a generation ends, so a per-step cost can be divided by the number of calls that paid it.
This commit is contained in:
@@ -64,11 +64,13 @@ def set_layout(layout) -> None:
|
||||
|
||||
|
||||
def end() -> None:
|
||||
from modules.attention import debug
|
||||
current.active = False
|
||||
current.role = None
|
||||
current.model_key = None
|
||||
current.layout = None
|
||||
new_pass(0)
|
||||
debug.end_generation()
|
||||
|
||||
|
||||
def set_role(name: str | None) -> None:
|
||||
|
||||
@@ -6,16 +6,37 @@ from modules.attention import context
|
||||
|
||||
enabled = os.environ.get('SD_ATTN_DEBUG', None) is not None
|
||||
seen: set[tuple] = set()
|
||||
counts: dict[tuple, int] = {}
|
||||
|
||||
|
||||
def observe(name: str, query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None) -> None:
|
||||
"""Log each distinct route once: backend, component role, step, shapes, dtype and mask presence."""
|
||||
signature = (name, context.current.role, tuple(query.shape), tuple(key.shape), str(query.dtype), attn_mask is not None)
|
||||
"""Log each distinct route once: backend, component role, step, shapes, dtype, mask presence and whether the inputs are contiguous; count every call."""
|
||||
contiguous = query.is_contiguous() and key.is_contiguous()
|
||||
signature = (name, context.current.role, tuple(query.shape), tuple(key.shape), str(query.dtype), attn_mask is not None, contiguous)
|
||||
counts[signature] = counts.get(signature, 0) + 1
|
||||
if signature in seen:
|
||||
return
|
||||
seen.add(signature)
|
||||
log.debug(f'Attention route: backend={name} role={context.current.role} step={context.current.step} q={list(query.shape)} k={list(key.shape)} dtype={query.dtype} mask={attn_mask is not None}')
|
||||
log.debug(f'Attention route: backend={name} role={context.current.role} step={context.current.step} q={list(query.shape)} k={list(key.shape)} dtype={query.dtype} mask={attn_mask is not None} contiguous={contiguous}')
|
||||
|
||||
|
||||
def summary() -> list[str]:
|
||||
"""One line per route with its call count since the last generation, busiest first."""
|
||||
lines = []
|
||||
for signature, count in sorted(counts.items(), key=lambda item: -item[1]):
|
||||
name, role, q_shape, k_shape, dtype, masked, contiguous = signature
|
||||
lines.append(f'backend={name} role={role} q={list(q_shape)} k={list(k_shape)} dtype={dtype} mask={masked} contiguous={contiguous} calls={count}')
|
||||
return lines
|
||||
|
||||
|
||||
def end_generation() -> None:
|
||||
"""Log the route counts of the generation that just ended and start the next count."""
|
||||
if enabled and counts:
|
||||
for line in summary():
|
||||
log.debug(f'Attention routes: {line}')
|
||||
counts.clear()
|
||||
|
||||
|
||||
def reset() -> None:
|
||||
seen.clear()
|
||||
counts.clear()
|
||||
|
||||
Reference in New Issue
Block a user