From e39fd6e2443c0607dac67c832c3ad77dd252ede7 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Fri, 2 Oct 2026 09:20:39 +0200 Subject: [PATCH] Run batch-1 Krea 2 forwards on the valid prompt rows (0.11.1) Krea2Pipeline pads every prompt to 512 text tokens and masks the padding. The fused blocks skipped the padded rows with a gather and scatter in every block; a batch-1 forward now drops them once, before the text fusion, and runs without a mask. The trimmed embeddings and positions are reused for the same prompt, so the text-fusion cache stays warm across steps. The padded rows got no attention weight and never reached the output, so the image rows see the same keys. RTX 4060 Ti, 1248x832, 8 steps, through the pipeline with zero-copy offload: a denoising step takes 1.03 s, as with the OrbitQuant 0.10 runner, instead of about 1.09 s. --- README.md | 2 +- pyproject.toml | 2 +- src/orbitquant/__init__.py | 2 +- src/orbitquant/fused/families/krea2.py | 66 ++++++++++++++++++++++++++ tests/test_fused_family_kernels.py | 40 ++++++++++++++++ uv.lock | 2 +- 6 files changed, 110 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 1ae38f7..5c41642 100644 --- a/README.md +++ b/README.md @@ -492,7 +492,7 @@ the same for an `orbitquant-v1` component artifact, and | Family | Transformer class | Fused blocks | | --- | --- | --- | -| Krea 2 | `Krea2Transformer2DModel` | Q\|K\|V\|gate, gated output, SwiGLU; W8A8 down projections | +| Krea 2 | `Krea2Transformer2DModel` | Q\|K\|V\|gate, gated output, SwiGLU; W8A8 down projections; the text fusion runs once per prompt and batch-1 forwards skip the padded prompt rows | | FLUX.2 | `Flux2Transformer2DModel` | double-stream (per-stream Q\|K\|V, gated residuals) and single-stream blocks; reference-image KV cache | | Ideogram 4 | `Ideogram4Transformer2DModel` | sandwich-norm blocks with row-indexed modulation | | Qwen-Image 2.1 | `QwenImage21Transformer2DModel` | block-causal prefill and KV-cached decode | diff --git a/pyproject.toml b/pyproject.toml index 561d981..13e1a5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.11.0" +version = "0.11.1" description = "Calibration-free OrbitQuant for transformer linear projections" readme = "README.md" requires-python = ">=3.11" diff --git a/src/orbitquant/__init__.py b/src/orbitquant/__init__.py index f33e0db..e4c5b9d 100644 --- a/src/orbitquant/__init__.py +++ b/src/orbitquant/__init__.py @@ -24,7 +24,7 @@ from orbitquant.quantizer import register_hf_quantizers from orbitquant.recipes import recipe -__version__ = "0.11.0" +__version__ = "0.11.1" register_hf_quantizers() diff --git a/src/orbitquant/fused/families/krea2.py b/src/orbitquant/fused/families/krea2.py index 2766999..2f3899b 100644 --- a/src/orbitquant/fused/families/krea2.py +++ b/src/orbitquant/fused/families/krea2.py @@ -41,6 +41,72 @@ def prepare(block, kind): def install(model): + _run_on_valid_text_rows(model) + _cache_text_fusion(model) + + +def _run_on_valid_text_rows(model): + """The pipeline pads the prompt to its maximum length and masks the padding. Padded rows + get no attention weight and never reach the output, so a batch-1 forward runs on the valid + text rows only and needs no mask: no per-block gather and scatter of the padded sequence, + and RoPE and the text projection cover fewer rows. The trimmed tensors of the last two + prompts are reused (a CFG pipeline alternates two), which keeps the text-fusion cache warm.""" + run = model.forward + entries = [] + + def same(held, tensors): + return all( + ref() is t and version == t._version + for (ref, version), t in zip(held, tensors, strict=True) + ) + + def trimmed(encoder_hidden_states, mask, position_ids): + tensors = (encoder_hidden_states, mask, position_ids) + for held, out in entries: + if same(held, tensors): + return out + text_len = encoder_hidden_states.shape[1] + rows = mask.reshape(-1).nonzero().squeeze(1) + out = ( + encoder_hidden_states[:, rows], + torch.cat([position_ids[rows], position_ids[text_len:]]), + ) + entries.insert(0, (tuple((weakref.ref(t), t._version) for t in tensors), out)) + del entries[2:] + return out + + def forward( + hidden_states, + encoder_hidden_states, + timestep, + position_ids, + encoder_attention_mask=None, + attention_kwargs=None, + return_dict=True, + ): + if ( + encoder_attention_mask is not None + and hidden_states.shape[0] == 1 + and not torch.is_grad_enabled() + ): + encoder_hidden_states, position_ids = trimmed( + encoder_hidden_states, encoder_attention_mask, position_ids + ) + encoder_attention_mask = None + return run( + hidden_states, + encoder_hidden_states, + timestep, + position_ids, + encoder_attention_mask=encoder_attention_mask, + attention_kwargs=attention_kwargs, + return_dict=return_dict, + ) + + model.forward = forward + + +def _cache_text_fusion(model): """The text fusion stack does not depend on the timestep: run it once per prompt instead of in every denoising step. The last two results are kept (a CFG pipeline alternates two prompts), each for the same embeddings and mask tensor objects at the same versions, so a new diff --git a/tests/test_fused_family_kernels.py b/tests/test_fused_family_kernels.py index f4b15f1..2ff098f 100644 --- a/tests/test_fused_family_kernels.py +++ b/tests/test_fused_family_kernels.py @@ -300,3 +300,43 @@ def test_lloyd_max_activation_codes_match_the_rotated_reference(dim): index = torch.bucketize(unit, codebook.boundaries.to("cuda", torch.float32)) expected = quantizer.codes[index] assert (codes != expected).float().mean().item() < 1e-3 + + +def test_krea2_batch_one_forward_runs_on_the_valid_text_rows(): + _require_cuda() + from orbitquant.fused.families import krea2 + + calls = [] + + class Transformer(torch.nn.Module): + def forward( + self, + hidden_states, + encoder_hidden_states, + timestep, + position_ids, + encoder_attention_mask=None, + attention_kwargs=None, + return_dict=True, + ): + calls.append((encoder_hidden_states, position_ids, encoder_attention_mask)) + return hidden_states + + model = Transformer() + krea2._run_on_valid_text_rows(model) + # The pipeline's layout: [prompt | padding | template suffix]. + text = torch.randn(1, 6, 2, 4, device="cuda") + mask = torch.tensor([[True, True, False, False, True, True]], device="cuda") + positions = torch.cat( + [torch.zeros(6, 3, device="cuda"), torch.arange(9.0, device="cuda").view(3, 3)] + ) + latents = torch.randn(1, 3, 8, device="cuda") + with torch.no_grad(): + for _ in range(2): + model.forward(latents, text, torch.ones(1), positions, encoder_attention_mask=mask) + + (text_1, positions_1, mask_1), (text_2, positions_2, _) = calls + assert mask_1 is None and text_1 is text_2 and positions_1 is positions_2 + valid = torch.tensor([0, 1, 4, 5], device="cuda") + torch.testing.assert_close(text_1, text[:, valid]) + torch.testing.assert_close(positions_1, torch.cat([positions[valid], positions[6:]])) diff --git a/uv.lock b/uv.lock index c2ed311..a4463a5 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.11.0" +version = "0.11.1" source = { editable = "." } dependencies = [ { name = "huggingface-hub" },