mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
feat(attention): sparse attention settings and router stage
Sparse attention is a stage over the chain rather than a member of it: one switch, and the router hands the selection to whichever active backend advertises that it consumes a block mask, currently flex. A backend declares that through a capability set, so the quantized kernel joins later without touching the router. The stage gates on the component role, self attention, a minimum sequence length defaulting to the measured 8192 token crossover, and the absence of a token mask or causal flag, which flex cannot combine with a block only mask. Budgets follow a precomputed per step schedule with at most two distinct values. Enabling the feature with no capable backend in the chain warns and leaves attention dense rather than doing nothing quietly. The modular pre-forward hook now receives kwargs and publishes whatever token layout the pipeline passes by name, so a packed sequence gets its conditioning pinned without any model specific code. Without a layout the whole sequence is sparsified and that is logged once per length.
This commit is contained in:
@@ -9,7 +9,10 @@ def prepare(platform: Platform, original): # pylint: disable=unused-argument
|
||||
def causal_mask(b, h, q_idx, kv_idx): # pylint: disable=unused-argument
|
||||
return q_idx >= kv_idx
|
||||
|
||||
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
|
||||
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=None): # pylint: disable=unused-argument
|
||||
if selection is not None:
|
||||
from modules.attention.sparse import flex as sparse_flex
|
||||
return sparse_flex.attend(query, key, value, selection, scale=scale, enable_gqa=enable_gqa)
|
||||
score_mod = None
|
||||
block_mask = None
|
||||
if attn_mask is not None:
|
||||
@@ -36,4 +39,5 @@ def prepare(platform: Platform, original): # pylint: disable=unused-argument
|
||||
backend = AttentionBackend(
|
||||
name='flex', label='Flex attention', priority=20, prepare=prepare,
|
||||
constraints=Constraints(min_ndim=4, same_device=True), # flex_attention takes 4d tensors on one device and compiles on cpu
|
||||
caps=frozenset({'block_mask'}),
|
||||
)
|
||||
|
||||
@@ -13,6 +13,7 @@ class GenerationContext:
|
||||
forwards: int = 0
|
||||
model_key: tuple[str, str | None] | None = None # pipeline class and denoiser class, telemetry only
|
||||
step_buffer: torch.Tensor | None = None # the step as a device scalar updated in place, so compiled readers keep their graph
|
||||
layout: object | None = None # TokenLayout published by whoever knows the packing, None until something does
|
||||
|
||||
|
||||
current = GenerationContext()
|
||||
@@ -30,6 +31,7 @@ def begin(pipe, steps: int = 0) -> None:
|
||||
from modules import devices
|
||||
current.active = True
|
||||
current.role = 'transformer'
|
||||
current.layout = None
|
||||
current.model_key = (pipe.__class__.__name__, denoiser_name(pipe)) if pipe is not None else None
|
||||
device = devices.device if devices.device is not None else torch.device('cpu')
|
||||
if current.step_buffer is None or current.step_buffer.device != device:
|
||||
@@ -56,10 +58,16 @@ def tick(step: int | None = None) -> None:
|
||||
current.forwards = current.step + 1
|
||||
|
||||
|
||||
def set_layout(layout) -> None:
|
||||
"""Publish what the packed sequence holds; callers that know the packing set this per forward."""
|
||||
current.layout = layout
|
||||
|
||||
|
||||
def end() -> None:
|
||||
current.active = False
|
||||
current.role = None
|
||||
current.model_key = None
|
||||
current.layout = None
|
||||
new_pass(0)
|
||||
|
||||
|
||||
|
||||
@@ -64,6 +64,7 @@ class AttentionBackend:
|
||||
terminal: bool = False # serves every call the entries decline, in place of the original sdpa
|
||||
platforms: frozenset[str] | None = None # devices backends the implementation exists for, None for all
|
||||
options: tuple[str, ...] = () # settings the prepared call captures; a change to one rebuilds the chain
|
||||
caps: frozenset[str] = frozenset() # what the call can consume beyond plain sdpa arguments, currently 'block_mask'
|
||||
|
||||
def available_on(self, platform: Platform) -> bool:
|
||||
return self.platforms is None or platform.backend in self.platforms
|
||||
@@ -94,5 +95,8 @@ class Registry:
|
||||
def options(self) -> list[str]:
|
||||
return sorted({name for backend in self.backends.values() for name in backend.options})
|
||||
|
||||
def with_cap(self, cap: str) -> list[AttentionBackend]:
|
||||
return [backend for backend in self.ordered() if cap in backend.caps]
|
||||
|
||||
|
||||
registry = Registry()
|
||||
|
||||
@@ -59,7 +59,7 @@ def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registr
|
||||
return Plan(entries=tuple(entries), terminal=terminal, original=original, platform=platform, labels=tuple(labels))
|
||||
|
||||
|
||||
def make_router(plan: Plan, observer: Callable | None = None) -> AttentionCall:
|
||||
def make_router(plan: Plan, observer: Callable | None = None, stage: Callable | None = None) -> AttentionCall:
|
||||
entries = plan.entries
|
||||
terminal = plan.terminal.call if plan.terminal is not None else None
|
||||
terminal_name = plan.terminal.backend.name if plan.terminal is not None else 'sdpa'
|
||||
@@ -69,10 +69,16 @@ def make_router(plan: Plan, observer: Callable | None = None) -> AttentionCall:
|
||||
def sdpa_router(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs):
|
||||
for entry in entries:
|
||||
if entry.backend.constraints.accepts(query, key, value, attn_mask):
|
||||
if stage is not None and 'block_mask' in entry.backend.caps:
|
||||
selection = stage(query, key, value, attn_mask, is_causal)
|
||||
if selection is not None:
|
||||
if observer is not None:
|
||||
observer(f'{entry.backend.name}+sparse', query, key, attn_mask)
|
||||
return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=selection)
|
||||
if observer is not None:
|
||||
observer(entry.backend.name, query, key, attn_mask)
|
||||
return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa)
|
||||
if observer is not None:
|
||||
if observer is not None: # pylint: disable=duplicate-code
|
||||
observer(terminal_name, query, key, attn_mask)
|
||||
if terminal is not None:
|
||||
return terminal(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, **kwargs)
|
||||
@@ -89,21 +95,43 @@ def install_router(labels, platform: Platform, original: AttentionCall, reg: Reg
|
||||
plan = build_plan(labels, platform, original, reg)
|
||||
debug.reset()
|
||||
observer = debug.observe if debug.enabled else None
|
||||
torch.nn.functional.scaled_dot_product_attention = make_router(plan, observer) if (plan.entries or plan.terminal is not None) else original
|
||||
stage = build_sparse_stage(plan)
|
||||
torch.nn.functional.scaled_dot_product_attention = make_router(plan, observer, stage) if (plan.entries or plan.terminal is not None) else original
|
||||
current_plan = plan
|
||||
torch_info.set(attention='>'.join(plan.chain()))
|
||||
log.debug(f'Torch attention: chain={">".join(plan.chain())} overrides={list(labels)} backend={platform.backend}')
|
||||
log.debug(f'Torch attention: chain={">".join(plan.chain())} overrides={list(labels)} backend={platform.backend} sparse={stage is not None}')
|
||||
return plan
|
||||
|
||||
|
||||
def build_sparse_stage(plan: Plan):
|
||||
"""Sparse attention is a stage over the chain rather than a chain member, so it needs a backend in the chain that consumes a block mask."""
|
||||
from modules.attention.sparse import stage as sparse_stage
|
||||
try:
|
||||
options = sparse_stage.read_options()
|
||||
except Exception:
|
||||
return None
|
||||
if not options.enabled:
|
||||
return None
|
||||
capable = [entry.backend.name for entry in plan.entries if 'block_mask' in entry.backend.caps]
|
||||
if not capable:
|
||||
names = [backend.label for backend in default_registry.with_cap('block_mask')]
|
||||
log.warning(f'Sparse attention: enabled but no active backend consumes a block mask, enable one of {names} in sdp overrides; attention stays dense')
|
||||
return None
|
||||
built = sparse_stage.make_stage(options)
|
||||
if built is not None:
|
||||
log.info(f'Sparse attention: backend={capable[0]} budget={options.budget:.0%} gate={options.gate} schedule={options.schedule_steps}x+{options.schedule_bump:.0%}')
|
||||
return built
|
||||
|
||||
|
||||
def get_plan() -> Plan | None:
|
||||
return current_plan
|
||||
|
||||
|
||||
def reapply_options(reg: Registry | None = None) -> list[str]:
|
||||
"""Settings whose change rebuilds the chain: the override set, the torch kernel flags, and every option a backend captures."""
|
||||
"""Settings whose change rebuilds the chain: the override set, the torch kernel flags, every option a backend captures, and the sparse stage."""
|
||||
from modules.attention.sparse import stage as sparse_stage
|
||||
reg = reg if reg is not None else default_registry
|
||||
return ['sdp_options', 'sdp_overrides', *reg.options()]
|
||||
return ['sdp_options', 'sdp_overrides', *reg.options(), *sparse_stage.OPTION_NAMES]
|
||||
|
||||
|
||||
def reapply() -> None:
|
||||
@@ -117,12 +145,22 @@ def reapply() -> None:
|
||||
|
||||
|
||||
def report() -> dict:
|
||||
"""The active chain and generation context, for the api and the debug log."""
|
||||
"""The active chain, sparse stage and generation context, for the api and the debug log."""
|
||||
from modules.attention.sparse import stage as sparse_stage
|
||||
plan = current_plan
|
||||
state = context.current
|
||||
options = sparse_stage.read_options()
|
||||
layout = state.layout
|
||||
return {
|
||||
'chain': plan.chain() if plan is not None else ['sdpa'],
|
||||
'overrides': list(plan.labels) if plan is not None else [],
|
||||
'sparse': {
|
||||
'enabled': options.enabled,
|
||||
'budget': options.budget,
|
||||
'gate': options.gate,
|
||||
'capable': [entry.backend.name for entry in plan.entries if 'block_mask' in entry.backend.caps] if plan is not None else [],
|
||||
'layout': {'source': layout.source, 'kinds': list(layout.kinds()), 'length': layout.length} if layout is not None else None,
|
||||
},
|
||||
'backend': plan.platform.backend if plan is not None else None,
|
||||
'context': {'active': state.active, 'role': state.role, 'step': state.step, 'steps': state.steps, 'model': state.model_key},
|
||||
}
|
||||
|
||||
@@ -53,7 +53,7 @@ def runs(indices: torch.Tensor) -> list[tuple[int, int]]:
|
||||
return [(int(values[bounds[i]].item()), int(values[bounds[i + 1] - 1].item()) + 1) for i in range(len(bounds) - 1)]
|
||||
|
||||
|
||||
def layout_from_index_kwargs(kwargs: dict, length: int) -> TokenLayout | None:
|
||||
def layout_from_index_kwargs(kwargs: dict, length: int | None = None) -> TokenLayout | None:
|
||||
"""Read a layout off the *_indices tensors a pipeline passes its transformer by name."""
|
||||
spans: list[Span] = []
|
||||
for name, value in kwargs.items():
|
||||
@@ -68,7 +68,7 @@ def layout_from_index_kwargs(kwargs: dict, length: int) -> TokenLayout | None:
|
||||
if not spans:
|
||||
return None
|
||||
spans.sort(key=lambda s: s.start)
|
||||
return TokenLayout(spans=tuple(spans), length=length, source='indices')
|
||||
return TokenLayout(spans=tuple(spans), length=length if length is not None else spans[-1].end, source='indices')
|
||||
|
||||
|
||||
def layout_from_segments(segments, length: int | None = None, source: str = 'segments') -> TokenLayout:
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
"""The router stage that turns settings plus a published layout into a per call block selection."""
|
||||
from dataclasses import dataclass
|
||||
from modules.logger import log
|
||||
from modules.attention import context
|
||||
from modules.attention.sparse import layout as layout_mod
|
||||
from modules.attention.sparse.selector import SparseSpec, block_count, schedule, select_blocks
|
||||
|
||||
|
||||
# measured on a 3090: below roughly this length a 30 percent budget caps under 1.25x per block,
|
||||
# so the selector cannot pay for itself; see docs/sparse-attention-tracker.md
|
||||
AUTO_MIN_TOKENS = 8192
|
||||
|
||||
# settings the stage reads, so a change to any of them rebuilds the chain
|
||||
OPTION_NAMES = ('sparse_attention_enabled', 'sparse_attention_budget', 'sparse_attention_min_tokens', 'sparse_attention_schedule_steps', 'sparse_attention_schedule_bump', 'sparse_attention_head_shared')
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StageOptions:
|
||||
enabled: bool = False
|
||||
budget: float = 0.30
|
||||
min_tokens: int = 0 # 0 selects AUTO_MIN_TOKENS
|
||||
schedule_steps: int = 0
|
||||
schedule_bump: float = 0.0
|
||||
head_shared: bool = False
|
||||
|
||||
@property
|
||||
def gate(self) -> int:
|
||||
return self.min_tokens if self.min_tokens > 0 else AUTO_MIN_TOKENS
|
||||
|
||||
|
||||
def read_options() -> StageOptions:
|
||||
from modules import shared
|
||||
opts = shared.opts
|
||||
return StageOptions(
|
||||
enabled=bool(getattr(opts, 'sparse_attention_enabled', False)),
|
||||
budget=float(getattr(opts, 'sparse_attention_budget', 30)) / 100.0,
|
||||
min_tokens=int(getattr(opts, 'sparse_attention_min_tokens', 0)),
|
||||
schedule_steps=int(getattr(opts, 'sparse_attention_schedule_steps', 0)),
|
||||
schedule_bump=float(getattr(opts, 'sparse_attention_schedule_bump', 0)) / 100.0,
|
||||
head_shared=bool(getattr(opts, 'sparse_attention_head_shared', False)),
|
||||
)
|
||||
|
||||
|
||||
def resolve_layout(seq: int, reported: set) -> layout_mod.TokenLayout:
|
||||
"""The published layout when there is one, otherwise sparsify the whole sequence and say so once."""
|
||||
published = context.current.layout
|
||||
if isinstance(published, layout_mod.TokenLayout) and published.length == seq:
|
||||
return published
|
||||
if seq not in reported:
|
||||
reported.add(seq)
|
||||
detail = 'none published' if published is None else f'published length {getattr(published, "length", None)} does not match {seq}'
|
||||
log.info(f'Sparse attention: no token layout ({detail}), sparsifying the whole sequence at tokens={seq}')
|
||||
return layout_mod.layout_from_prefix(seq, 0)
|
||||
|
||||
|
||||
def make_stage(options: StageOptions):
|
||||
"""Return the per call selector, or None when the feature is off."""
|
||||
if not options.enabled or options.budget >= 1.0:
|
||||
return None
|
||||
reported: set = set()
|
||||
cache: dict = {}
|
||||
|
||||
def budget_for_step() -> float:
|
||||
state = context.current
|
||||
if options.schedule_steps <= 0 or options.schedule_bump <= 0 or state.steps <= 0:
|
||||
return options.budget
|
||||
key = (state.steps, options.budget, options.schedule_bump, options.schedule_steps)
|
||||
table = cache.get(key)
|
||||
if table is None:
|
||||
table = schedule(state.steps, options.budget, options.schedule_bump, options.schedule_steps)
|
||||
cache.clear()
|
||||
cache[key] = table
|
||||
return table[min(state.step, len(table) - 1)] if table else options.budget
|
||||
|
||||
def stage(query, key, value, attn_mask, is_causal): # pylint: disable=unused-argument
|
||||
state = context.current
|
||||
if state.role != 'transformer' or not state.active:
|
||||
return None
|
||||
if attn_mask is not None or is_causal: # flex would need a mask_mod to combine these; the quantized kernel composes them in R2
|
||||
return None
|
||||
if query.device.type == 'cpu' or query.dim() != 4:
|
||||
return None
|
||||
seq_q, seq_kv = query.shape[-2], key.shape[-2]
|
||||
if seq_q != seq_kv or seq_q < options.gate: # cross attention is short and already cheap
|
||||
return None
|
||||
budget = budget_for_step()
|
||||
if budget >= 1.0:
|
||||
return None
|
||||
spec = SparseSpec(budget=budget, head_shared=options.head_shared)
|
||||
token_layout = resolve_layout(seq_q, reported)
|
||||
nq, nk = block_count(seq_q, spec.block_q), block_count(seq_kv, spec.block_kv)
|
||||
pins, drops = layout_mod.block_pins(token_layout, seq_q, seq_kv, spec.block_q, spec.block_kv, query.device)
|
||||
if pins.shape[-2:] != (nq, nk):
|
||||
return None
|
||||
return select_blocks(query, key, spec, pins=pins, drops=drops)
|
||||
|
||||
stage.options = options # pylint: disable=attribute-defined-outside-init
|
||||
return stage
|
||||
@@ -75,6 +75,15 @@ class InterruptLogFilter(logging.Filter):
|
||||
return 'Interrupted...' not in record.msg
|
||||
|
||||
|
||||
def publish_layout(kwargs):
|
||||
"""Hand the attention router whatever the pipeline says about its packed sequence, keyed on the *_indices tensors rather than the model."""
|
||||
try:
|
||||
from modules.attention.sparse import layout as sparse_layout
|
||||
attention_context.set_layout(sparse_layout.layout_from_index_kwargs(kwargs or {}))
|
||||
except Exception as e:
|
||||
log.debug(f'Pipeline: token layout {e}')
|
||||
|
||||
|
||||
def install_state_hook(pipe):
|
||||
runner_log = logging.getLogger('diffusers.modular_pipelines.modular_pipeline')
|
||||
if not any(isinstance(f, InterruptLogFilter) for f in runner_log.filters):
|
||||
@@ -91,9 +100,10 @@ def install_state_hook(pipe):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _pre_transformer_hook(module, args): # pylint: disable=unused-argument
|
||||
def _pre_transformer_hook(module, args, kwargs): # pylint: disable=unused-argument
|
||||
new_phase = set_phase('Generate', module)
|
||||
attention_context.set_role('transformer')
|
||||
publish_layout(kwargs)
|
||||
if new_phase:
|
||||
sd_offload.offload_ondemand(pipe, exclude=['transformer', 'transformer_ref'], reason='generate', force=hasattr(pipe, 'sdnext_force_offload'))
|
||||
if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0:
|
||||
@@ -138,7 +148,8 @@ def install_state_hook(pipe):
|
||||
if module is not None:
|
||||
target = getattr(module, 'model', module) # conditioning calls the inner model directly
|
||||
if isinstance(target, torch.nn.Module) and getattr(target, 'sdnext_state_hook', None) is None:
|
||||
target.sdnext_state_hook = target.register_forward_pre_hook(_pre_transformer_hook)
|
||||
# with_kwargs, because the blocks call the transformer entirely by keyword and the token layout rides in those kwargs
|
||||
target.sdnext_state_hook = target.register_forward_pre_hook(_pre_transformer_hook, with_kwargs=True)
|
||||
|
||||
for name in ('text_encoder', 'text_encoder_2'):
|
||||
module = getattr(pipe, name, None)
|
||||
|
||||
@@ -262,6 +262,14 @@ def create_settings(cmd_opts):
|
||||
"sdnq_attention_pv_matmul_type": OptionInfo("disabled", "SDNQ Attention PV MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}),
|
||||
"sdnq_attention_hadamard_group_size": OptionInfo(256, "SDNQ Attention Hadamard Group Size", gr.Slider, {"minimum": 4, "maximum": 1024, "step": 1}),
|
||||
|
||||
"sparse_attention_sep": OptionInfo("<h2>Sparse Attention</h2>", "", gr.HTML),
|
||||
"sparse_attention_enabled": OptionInfo(False, "Sparse Attention", gr.Checkbox),
|
||||
"sparse_attention_budget": OptionInfo(30, "Sparse Attention KV budget", gr.Slider, {"minimum": 5, "maximum": 100, "step": 5}),
|
||||
"sparse_attention_min_tokens": OptionInfo(0, "Sparse Attention minimum sequence", gr.Slider, {"minimum": 0, "maximum": 65536, "step": 1024}),
|
||||
"sparse_attention_schedule_steps": OptionInfo(0, "Sparse Attention dense steps", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
|
||||
"sparse_attention_schedule_bump": OptionInfo(30, "Sparse Attention dense step bonus", gr.Slider, {"minimum": 0, "maximum": 70, "step": 5}),
|
||||
"sparse_attention_head_shared": OptionInfo(False, "Sparse Attention share selection across heads", gr.Checkbox),
|
||||
|
||||
"hf_attention_sep": OptionInfo("<h2>Attention Dispatcher</h2>", "", gr.HTML),
|
||||
"hf_attention": OptionInfo('', "Attention dispatcher kernel", gr.Textbox),
|
||||
}))
|
||||
|
||||
@@ -15,6 +15,8 @@ Covers:
|
||||
- the flex consumer: a full-keep selection through flex_attention reproduces dense sdpa, and a
|
||||
selection with dropped tiles reproduces sdpa given the same tiles masked out
|
||||
- the density matched radial control and the step schedule
|
||||
- the router stage: the gates it applies (component role, mask, causal, cross attention, minimum
|
||||
sequence), the per step budget schedule, and layout resolution with and without a publisher
|
||||
|
||||
The flex rows need a cuda device and compile the flex kernel; they skip on cpu.
|
||||
|
||||
@@ -329,6 +331,90 @@ def test_flex_handles_a_ragged_tail():
|
||||
return True
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Router stage
|
||||
# ============================================================
|
||||
|
||||
def stage_options(**kwargs):
|
||||
from modules.attention.sparse import stage as stage_mod
|
||||
base = dict(enabled=True, budget=0.30, min_tokens=1024)
|
||||
base.update(kwargs)
|
||||
return stage_mod.StageOptions(**base)
|
||||
|
||||
|
||||
def with_context(fn):
|
||||
from modules.attention import context as ctx
|
||||
ctx.begin(None, steps=10)
|
||||
try:
|
||||
return fn()
|
||||
finally:
|
||||
ctx.end()
|
||||
|
||||
|
||||
def test_stage_is_none_when_disabled_or_at_full_budget():
|
||||
from modules.attention.sparse import stage as stage_mod
|
||||
assert stage_mod.make_stage(stage_options(enabled=False)) is None
|
||||
assert stage_mod.make_stage(stage_options(budget=1.0)) is None
|
||||
assert stage_mod.make_stage(stage_options()) is not None
|
||||
return True
|
||||
|
||||
|
||||
def test_stage_gates():
|
||||
from modules.attention.sparse import stage as stage_mod
|
||||
stage = stage_mod.make_stage(stage_options())
|
||||
q, k, v = qkv(heads=2, seq=2048)
|
||||
short_q, short_k, short_v = qkv(heads=2, seq=512)
|
||||
cross_k, cross_v = randn(1, 2, 77, 64), randn(1, 2, 77, 64)
|
||||
|
||||
def checks():
|
||||
from modules.attention import context as ctx
|
||||
assert stage(q, k, v, None, False) is not None, 'an eligible call must be selected'
|
||||
assert stage(q, k, v, torch.zeros(1, 1, 2048, 2048, dtype=torch.bool, device=device), False) is None, 'a masked call is not eligible yet'
|
||||
assert stage(q, k, v, None, True) is None, 'a causal call is not eligible yet'
|
||||
assert stage(q, cross_k, cross_v, None, False) is None, 'cross attention is not eligible'
|
||||
assert stage(short_q, short_k, short_v, None, False) is None, 'below the gate attention stays dense'
|
||||
ctx.set_role('vae')
|
||||
assert stage(q, k, v, None, False) is None, 'only the denoiser is sparsified'
|
||||
ctx.set_role('transformer')
|
||||
return True
|
||||
return with_context(checks)
|
||||
|
||||
|
||||
def test_stage_follows_the_step_schedule():
|
||||
from modules.attention.sparse import stage as stage_mod
|
||||
from modules.attention import context as ctx
|
||||
stage = stage_mod.make_stage(stage_options(budget=0.30, schedule_steps=2, schedule_bump=0.40))
|
||||
q, k, v = qkv(heads=2, seq=2048)
|
||||
|
||||
def checks():
|
||||
densities = []
|
||||
for step in range(10):
|
||||
ctx.set_step(step)
|
||||
selection = stage(q, k, v, None, False)
|
||||
densities.append(selection.budget)
|
||||
assert densities[0] == densities[1] > densities[5], densities
|
||||
assert densities[-1] == densities[-2] > densities[5], densities
|
||||
assert len(set(densities)) == 2, set(densities)
|
||||
return True
|
||||
return with_context(checks)
|
||||
|
||||
|
||||
def test_stage_uses_a_published_layout_and_falls_back_without_one():
|
||||
from modules.attention.sparse import stage as stage_mod
|
||||
from modules.attention import context as ctx
|
||||
stage = stage_mod.make_stage(stage_options(budget=0.20))
|
||||
q, k, v = qkv(heads=2, seq=2048)
|
||||
|
||||
def checks():
|
||||
loose = stage(q, k, v, None, False)
|
||||
ctx.set_layout(sparse.layout_from_segments([('text', 256), ('video', 1792)]))
|
||||
pinned = stage(q, k, v, None, False)
|
||||
assert pinned.density > loose.density, f'pinning conditioning must keep more tiles: {pinned.density} vs {loose.density}'
|
||||
assert bool(pinned.keep[..., 0:4].all()), 'the pinned text columns must survive'
|
||||
return True
|
||||
return with_context(checks)
|
||||
|
||||
|
||||
def run_all():
|
||||
log.warning(f'=== selector (device={device}) ===')
|
||||
cat = category('selector')
|
||||
@@ -368,6 +454,16 @@ def run_all():
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== stage ===')
|
||||
cat = category('stage')
|
||||
for fn in [
|
||||
test_stage_is_none_when_disabled_or_at_full_budget,
|
||||
test_stage_gates,
|
||||
test_stage_follows_the_step_schedule,
|
||||
test_stage_uses_a_published_layout_and_falls_back_without_one,
|
||||
]:
|
||||
run_test(cat, fn)
|
||||
|
||||
log.warning('=== Results ===')
|
||||
total_passed = total_failed = total_skipped = 0
|
||||
for cat_name, info in results.items():
|
||||
|
||||
@@ -1470,7 +1470,13 @@
|
||||
{"id":"","label":"SDNQ Attention MatMul type","localized":"","hint":"Precision the query-key matmul is computed in, the first of the two matmuls in attention.<br><br><b>enabled</b> selects int8, and <b>int8</b> and <b>uint8</b> reach the same kernel.<br><b>float16</b> and <b>float8_e4m3fn</b> take the floating point path. fp8 needs a GPU with fp8 tensor cores and fails on hardware without them rather than falling back.<br><b>disabled</b> leaves queries and keys in the model's own precision, which also idles <b><i>SDNQ Attention use Smooth K</i></b> and <b><i>SDNQ Attention use Hadamard</i></b>.<br><br>Default enabled.","ui":"settings_cuda"},
|
||||
{"id":"","label":"SDNQ Attention PV MatMul type","localized":"","hint":"Precision the probability-value matmul is computed in, the second of the two matmuls in attention. Choices match <b><i>SDNQ Attention MatMul type</i></b>.<br><br>Quantizing this one as well takes out the floating point work the first setting leaves behind, and it is the more delicate of the two: its inputs are already normalized probabilities, and the small ones among them carry the fine detail.<br><b>disabled</b> keeps this matmul in the model's own precision.<br><br>Default disabled.","ui":"settings_cuda"},
|
||||
{"id":"","label":"SDNQ Attention Hadamard Group Size","localized":"","hint":"Width of the Hadamard rotation in channels. Wider groups mix more channels together and spread outliers further.<br><br>Clamped to the head dimension of the running model, rounded down to a power of two that divides it. On a model with 64 or 128 channels per head the upper part of this range resolves to that head dimension rather than to the number shown. Rotation is skipped below 4.<br>Applies while <b><i>SDNQ Attention use Hadamard</i></b> is enabled.<br><br>Default 256.","ui":"settings_cuda"},
|
||||
{"id":"","label":"SDNQ Attention Quantize FP32","localized":"","hint":"Upcasts queries, keys and values to fp32 for the quantization step, meaning the mean subtraction, scale and rounding that produce the low precision operands. The matmuls themselves are unaffected, and the kernel applies the scales in fp32 either way.<br>Turned off, that arithmetic runs in the model's own precision. bf16 carries eight mantissa bits, so a scale derived in it is coarser than one derived in fp32, and <b><i>SDNQ Attention use Smooth K</i></b> loses the most from it, since a mean across the whole sequence is exactly the kind of sum that wants the extra bits.<br><br>Whether the upcast costs anything depends on how the GPU runs fp32 vector work against fp16 and bf16. NVIDIA and AMD run them at the same rate here, so there is nothing to save; Intel runs fp32 slower and takes a noticeable hit.<br><br>Enabled by default.","ui":"settings_cuda"}
|
||||
{"id":"","label":"SDNQ Attention Quantize FP32","localized":"","hint":"Upcasts queries, keys and values to fp32 for the quantization step, meaning the mean subtraction, scale and rounding that produce the low precision operands. The matmuls themselves are unaffected, and the kernel applies the scales in fp32 either way.<br>Turned off, that arithmetic runs in the model's own precision. bf16 carries eight mantissa bits, so a scale derived in it is coarser than one derived in fp32, and <b><i>SDNQ Attention use Smooth K</i></b> loses the most from it, since a mean across the whole sequence is exactly the kind of sum that wants the extra bits.<br><br>Whether the upcast costs anything depends on how the GPU runs fp32 vector work against fp16 and bf16. NVIDIA and AMD run them at the same rate here, so there is nothing to save; Intel runs fp32 slower and takes a noticeable hit.<br><br>Enabled by default.","ui":"settings_cuda"},
|
||||
{"id":"","label":"Sparse Attention","localized":"","hint":"Attention is computed on a subset of key tiles instead of all of them, which lowers cost on long sequences. Requires an attention backend that accepts a block mask, currently <b>Flex attention</b> in <b><i>SDP overrides</i></b>. Below the minimum sequence the setting stays inactive and attention is dense.","ui":"settings_cuda"},
|
||||
{"id":"","label":"Sparse Attention KV budget","localized":"","hint":"Percentage of the eligible key tiles each query tile keeps. Lower budgets are faster and coarser. Text, conditioning and audio tokens are always kept, as are the tiles on the diagonal, so the budget applies only to the bulk image or video tokens.","ui":"settings_cuda"},
|
||||
{"id":"","label":"Sparse Attention minimum sequence","localized":"","hint":"Shortest sequence that is sparsified. Below it attention stays dense, because the selection costs more than it saves. <b>0</b> uses the built in threshold.","ui":"settings_cuda"},
|
||||
{"id":"","label":"Sparse Attention dense steps","localized":"","hint":"Number of steps at the start and end of sampling that receive a larger budget, where composition and detail are set. <b>0</b> applies one budget to every step.","ui":"settings_cuda"},
|
||||
{"id":"","label":"Sparse Attention dense step bonus","localized":"","hint":"Percentage points added to the budget on the dense steps.","ui":"settings_cuda"},
|
||||
{"id":"","label":"Sparse Attention share selection across heads","localized":"","hint":"One selection is computed for all attention heads rather than one per head. Cheaper to select and coarser in what it keeps.","ui":"settings_cuda"}
|
||||
],
|
||||
"t": [
|
||||
{"id":"txt2img_nav","label":"T2I","localized":"","hint":"Create image from text<br>Legacy interface that mimics original text-to-image interface and behavior"},
|
||||
|
||||
Reference in New Issue
Block a user