From a2e6535b93eef596582e4630b96112e53f85f31c Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 23 Aug 2026 22:16:58 +0100 Subject: [PATCH] feat(krea2): publish the packed token layout The transformer is the only site that knows every segment length, since the pad to a multiple of 256 happens there and the text stream is padded to a fixed 512 whatever the prompt. It now publishes the live runs of the mask, so the selector pins the conditioning dense and drops the padding instead of sparsifying the whole sequence. At 2048x2048 with a short prompt that is 6 of 264 key blocks per query row moved from pinned to dropped. --- modules/attention/sparse/__init__.py | 4 +-- modules/attention/sparse/layout.py | 16 ++++++++++ pipelines/krea2/transformer_krea2.py | 10 +++++++ test/test-attention-sparse.py | 44 ++++++++++++++++++++++++++-- 4 files changed, 70 insertions(+), 4 deletions(-) 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)