Run batch-1 Krea 2 forwards on the valid prompt rows (0.11.1)

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 |

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"

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()
 

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

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:]]))