diff --git a/modules/attention/sparse/__init__.py b/modules/attention/sparse/__init__.py index 5d8efc1f1..f8d0cdf9b 100644 --- a/modules/attention/sparse/__init__.py +++ b/modules/attention/sparse/__init__.py @@ -1,8 +1,8 @@ """Block-sparse attention: the selector, the token layout it respects, and the consumers that apply it.""" from modules.attention.sparse.selector import BlockSelection, SparseSpec, block_count, radial_blocks, schedule, select_blocks -from modules.attention.sparse.layout import Span, TokenLayout, block_pins, layout_from_index_kwargs, layout_from_prefix, layout_from_segments +from modules.attention.sparse.layout import Span, TokenLayout, block_pins, layout_from_index_kwargs, layout_from_prefix, layout_from_segments, publish_segments, segments_from_live __all__ = [ 'BlockSelection', 'SparseSpec', 'block_count', 'radial_blocks', 'schedule', 'select_blocks', - 'Span', 'TokenLayout', 'block_pins', 'layout_from_index_kwargs', 'layout_from_prefix', 'layout_from_segments', + 'Span', 'TokenLayout', 'block_pins', 'layout_from_index_kwargs', 'layout_from_prefix', 'layout_from_segments', 'publish_segments', 'segments_from_live', ] diff --git a/modules/attention/sparse/layout.py b/modules/attention/sparse/layout.py index 7362c96d6..c6c3de813 100644 --- a/modules/attention/sparse/layout.py +++ b/modules/attention/sparse/layout.py @@ -83,6 +83,22 @@ def layout_from_segments(segments, length: int | None = None, source: str = 'seg return TokenLayout(spans=tuple(spans), length=length if length is not None else cursor, source=source) +def segments_from_live(live: torch.Tensor, kind: str, pad_kind: str = 'pad') -> list[tuple[str, int]]: + """Run length encode a boolean live mask into ordered (kind, count) pairs, the dead runs labelled as padding.""" + values = live.detach().to('cpu').bool() + if values.numel() == 0: + return [] + changes = (values[1:] != values[:-1]).nonzero().flatten().tolist() + bounds = [0, *[c + 1 for c in changes], values.numel()] + return [(kind if bool(values[bounds[i]]) else pad_kind, bounds[i + 1] - bounds[i]) for i in range(len(bounds) - 1)] + + +def publish_segments(segments, length: int | None = None, source: str = 'segments') -> None: + """Publish a layout from the site that packs the sequence, which is the only place the segment lengths are all known.""" + from modules.attention import context + context.set_layout(layout_from_segments(segments, length=length, source=source)) + + def layout_from_prefix(length: int, prefix: int) -> TokenLayout: """Fallback when nothing published a layout: treat a leading run as conditioning and sparsify the rest.""" return layout_from_segments([('text', prefix), ('image', length - prefix)], length=length, source='prefix') diff --git a/pipelines/krea2/transformer_krea2.py b/pipelines/krea2/transformer_krea2.py index 643d926ba..c09e4be84 100644 --- a/pipelines/krea2/transformer_krea2.py +++ b/pipelines/krea2/transformer_krea2.py @@ -342,6 +342,16 @@ class Krea2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOri else: mask = None + from modules.attention.sparse import layout as sparse_layout # delayed, this module is also importable without the webui + if attention_mask is not None: + # the text stream is padded to a fixed length and the joint sequence to a multiple of 256, so the live + # runs are what the layout must report; a wholly padded key block is dropped rather than pinned + live = attention_mask.any(dim=0) + segments = sparse_layout.segments_from_live(live[:txtlen], "text") + sparse_layout.segments_from_live(live[txtlen:], "image") + else: + segments = [("text", txtlen), ("image", imglen)] + sparse_layout.publish_segments(segments, source="krea2") + freqs = self.posemb(position_ids) for block in self.blocks: diff --git a/test/test-attention-sparse.py b/test/test-attention-sparse.py index d78625da8..ddcf98316 100644 --- a/test/test-attention-sparse.py +++ b/test/test-attention-sparse.py @@ -242,6 +242,22 @@ def test_block_pins_pin_a_boundary_tile(): return True +def test_segments_from_live_splits_interior_padding(): + from modules.attention.sparse import layout as layout_mod + live = torch.zeros(512, dtype=torch.bool, device=device) + live[:40] = True + live[-8:] = True + segments = layout_mod.segments_from_live(live, 'text') + assert segments == [('text', 40), ('pad', 464), ('text', 8)], segments + token_layout = layout_mod.layout_from_segments(segments + [('image', 1024)]) + assert token_layout.length == 1536, token_layout.length + _, drops = layout_mod.block_pins(token_layout, 1536, 1536, 128, 64, device) + # only a key block that is padding all the way through is dropped, so the two straddling blocks survive + assert drops[..., 1:7].all(), 'whole padded key blocks are dropped' + assert not drops[..., 0].any() and not drops[..., 7].any(), 'a straddling block keeps its live tokens' + return True + + def test_block_pins_are_cached_per_geometry(): layout = sparse.layout_from_segments([('text', 128), ('video', 1024)]) first = sparse.block_pins(layout, 1152, 1152, 128, 64, device) @@ -369,8 +385,10 @@ def test_stage_gates(): 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' + mask = torch.zeros(1, 1, 2048, 2048, dtype=torch.bool, device=device) + assert stage(q, k, v, mask, False) is None, 'a masked call needs a backend that composes the two' + assert stage(q, k, v, mask, False, frozenset({'masked_block'})) is not None, 'a composing backend takes the masked call' + assert stage(q, k, v, None, True) is None, 'a causal call is not eligible' 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' assert stage.last_skip == 'below the minimum sequence', stage.last_skip @@ -381,6 +399,26 @@ def test_stage_gates(): return with_context(checks) +def test_published_segments_reach_the_stage(): + from modules.attention.sparse import layout as layout_mod + 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)) + q, k, v = qkv(heads=2, seq=2048) + + def checks(): + layout_mod.publish_segments((('text', 256), ('image', 1536), ('pad', 256)), source='test') + published = ctx.current.layout + assert published.length == 2048 and published.source == 'test', published + assert published.kinds() == ('text', 'image', 'pad'), published.kinds() + selection = stage(q, k, v, None, False) + keep, block_kv = selection.keep, selection.block_kv + assert keep[..., :256 // block_kv].all(), 'conditioning key tiles stay dense' + assert not keep[..., 1792 // block_kv:].any(), 'padding key tiles are dropped' + 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 @@ -440,6 +478,7 @@ def run_all(): test_block_pins_pin_conditioning_and_drop_padding, test_block_pins_pin_a_boundary_tile, test_block_pins_are_cached_per_geometry, + test_segments_from_live_splits_interior_padding, ]: run_test(cat, fn) @@ -461,6 +500,7 @@ def run_all(): test_stage_is_none_when_disabled_or_at_full_budget, test_stage_gates, test_stage_follows_the_step_schedule, + test_published_segments_reach_the_stage, test_stage_uses_a_published_layout_and_falls_back_without_one, ]: run_test(cat, fn)