Added reference token attention isolation (kv_cache) for Krea2 edit training. Same training cost with significant inference speed up. 2x inference speedup.

extensions_built_in/diffusion_models/krea2/krea2.py

@@ -208,6 +208,16 @@ class Krea2Model(BaseModel):
         self.has_multiple_control_images = self.is_edit
         # Reference images keep their own aspect/size (not resized to the target).
         self.use_raw_control_images = self.is_edit
+        # model_kwargs.kv_cache = true: train with an asymmetric attention mask
+        # where the clean reference tokens attend only to each other (never to
+        # text / noisy tokens). Their hidden states then depend only on the
+        # refs + t=0 modulation, so at inference their per-layer K/V can be
+        # computed once and reused across all denoising steps
+        # (OminiControl2-style conditioning feature reuse). Off by default:
+        # the base model was trained fully bidirectional, so a LoRA must be
+        # trained with kv_cache enabled for kv-cached inference (the ComfyUI
+        # node / hub pipeline kv_cache toggles) to work properly.
+        self.kv_cache = bool(self.model_config.model_kwargs.get("kv_cache", False))
 
     @staticmethod
     def get_train_scheduler():
@@ -642,6 +652,7 @@ class Krea2Model(BaseModel):
             context,
             text_mask,
             ref_latents=ref_latents,
+            isolate_refs=self.kv_cache,
         )
         return pred
 

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -210,7 +210,13 @@ class Attention(torch.nn.Module):
         self.wo = torch.nn.Linear(dim, dim, bias=bias)
 
     def forward(
-        self, qkv: Tensor, freqs: Tensor | None = None, mask: Tensor | None = None
+        self,
+        qkv: Tensor,
+        freqs: Tensor | None = None,
+        mask: Tensor | None = None,
+        ref_span: tuple[int, int] | None = None,
+        kv_capture: list | None = None,
+        kv_cache: tuple[Tensor, Tensor] | None = None,
     ) -> Tensor:
         q, k, v, gate = self.wq(qkv), self.wk(qkv), self.wv(qkv), self.gate(qkv)
 
@@ -223,6 +229,20 @@ class Attention(torch.nn.Module):
         q, k, v = self.qknorm(q, k, v)
         if freqs is not None:
             q, k = ropeapply(q, k, freqs)
+        if kv_capture is not None and ref_span is not None:
+            # Stash this block's post-RoPE ref K/V so later denoising steps can
+            # run without the ref tokens in the sequence (clone: drop the view
+            # into the full-sequence K/V so only the ref span stays alive).
+            kv_capture.append(
+                (
+                    k[:, :, ref_span[0] : ref_span[1]].clone(),
+                    v[:, :, ref_span[0] : ref_span[1]].clone(),
+                )
+            )
+        if kv_cache is not None:
+            # Cached ref K/V are already RoPE'd at their original positions.
+            k = torch.cat((k, kv_cache[0]), dim=2)
+            v = torch.cat((v, kv_cache[1]), dim=2)
         out = self.wo(attention(q, k, v, mask=mask, gqa=self.gqa) * F.sigmoid(gate))
 
         return out
@@ -326,8 +346,16 @@ class SingleStreamBlock(nn.Module):
         self.mlp = SwiGLU(features, multiplier, bias)
 
     def forward(
-        self, x: Tensor, vec: Tensor, freqs: Tensor, mask: Tensor | None = None
+        self,
+        x: Tensor,
+        vec: Tensor,
+        freqs: Tensor,
+        mask: Tensor | None = None,
+        ref_span: tuple[int, int] | None = None,
+        kv_capture: list | None = None,
+        kv_cache: tuple[Tensor, Tensor] | None = None,
     ) -> Tensor:
+        attn_kwargs = dict(ref_span=ref_span, kv_capture=kv_capture, kv_cache=kv_cache)
         # ``vec`` is the (B, 1, 6*features) modulation input, or a tuple
         # ``(vec, refvec, split)`` for reference-image conditioning: tokens
         # ``[:split]`` (text + noisy image) are modulated with ``vec`` while
@@ -351,13 +379,15 @@ class SingleStreamBlock(nn.Module):
             def gate(h, g):
                 return torch.cat((m[g] * h[:, :split], r[g] * h[:, split:]), dim=1)
 
-            x = x + gate(self.attn(mod(self.prenorm(x), 0, 1), freqs, mask), 2)
+            x = x + gate(
+                self.attn(mod(self.prenorm(x), 0, 1), freqs, mask, **attn_kwargs), 2
+            )
             x = x + gate(self.mlp(mod(self.postnorm(x), 3, 4)), 5)
             return x
 
         prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
         x = x + pregate * self.attn(
-            (1 + prescale) * self.prenorm(x) + preshift, freqs, mask
+            (1 + prescale) * self.prenorm(x) + preshift, freqs, mask, **attn_kwargs
         )
         x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
 
@@ -445,6 +475,9 @@ class SingleStreamDiT(nn.Module):
         pos: Tensor,
         mask: Tensor | None = None,
         reflen: int = 0,
+        isolate_refs: bool = False,
+        ref_kv_capture: list | None = None,
+        ref_kv_cache: tuple[list, Tensor] | None = None,
     ) -> Tensor:
         img = self.first(img)
         t = self.tmlp(temb(t, self.config.tdim, device=img.device, dtype=img.dtype))
@@ -483,11 +516,46 @@ class SingleStreamDiT(nn.Module):
             )
             blockvec = (tvec, self.tproj(t0), txtlen + imglen - reflen)
 
+        padmask = mask  # (B, L) key-padding mask, incl. the 256-alignment pad
         mask = _mask(mask)
 
+        if reflen > 0 and isolate_refs:
+            # Asymmetric attention (OminiControl2-style "feature reuse"): ref
+            # queries attend only to ref keys, while text + noisy queries still
+            # see everything. Combined with the t=0 modulation above, ref hidden
+            # states become independent of t and of the noisy tokens, so their
+            # per-layer K/V can be computed once and cached across denoising
+            # steps at inference. Changes attention flow vs the base model, so
+            # it needs to be trained in.
+            split = txtlen + imglen - reflen
+            is_ref = torch.zeros(
+                combined.shape[1], dtype=torch.bool, device=combined.device
+            )
+            is_ref[split : split + reflen] = True
+            mask = mask & (~is_ref[:, None] | is_ref[None, :])
+
+        # Ref K/V caching (inference-only; requires isolate_refs so the cached
+        # features are step-invariant). Capture mode: this pass has the refs in
+        # the sequence and records each block's post-RoPE ref K/V. Reuse mode:
+        # the refs are dropped from the sequence (reflen == 0) and the cached
+        # K/V are appended as extra attention keys instead.
+        ref_span = None
+        if ref_kv_capture is not None and reflen > 0:
+            assert isolate_refs, "ref K/V capture requires isolate_refs"
+            split = txtlen + imglen - reflen
+            ref_span = (split, split + reflen)
+
+        blockcaches = [None] * len(self.blocks)
+        if ref_kv_cache is not None:
+            blockcaches, refmask = ref_kv_cache
+            # live queries may attend a cached ref key wherever that ref token
+            # is real (refmask right-pads samples with fewer ref tokens)
+            extra = padmask.unsqueeze(1).unsqueeze(3) & refmask.unsqueeze(1).unsqueeze(2)
+            mask = torch.cat((mask, extra), dim=3)
+
         freqs = self.posemb(pos)
 
-        for block in self.blocks:
+        for block, blockkv in zip(self.blocks, blockcaches):
             if self.gradient_checkpointing and torch.is_grad_enabled():
                 combined = checkpoint(
                     block,
@@ -498,7 +566,15 @@ class SingleStreamDiT(nn.Module):
                     use_reentrant=False,
                 )
             else:
-                combined = block(combined, blockvec, freqs, mask)
+                combined = block(
+                    combined,
+                    blockvec,
+                    freqs,
+                    mask,
+                    ref_span=ref_span,
+                    kv_capture=ref_kv_capture,
+                    kv_cache=blockkv,
+                )
 
         final = self.last(combined, t)
         output = final[:, txtlen : txtlen + imglen - reflen, :]

extensions_built_in/diffusion_models/krea2/src/pipeline.py

@@ -151,6 +151,8 @@ def predict_velocity(
     context: torch.Tensor,  # (B, Lt, n*d) flattened stacked Qwen3-VL features
     text_mask: torch.Tensor,  # (B, Lt) 1 for real text tokens
     ref_latents: Optional[List[List[torch.Tensor]]] = None,  # per-sample (C, h, w) refs
+    isolate_refs: bool = False,
+    ref_kv_cache: Optional[dict] = None,
 ) -> torch.Tensor:
     """Run the MMDiT on the packed [text | image | refs] sequence.
 
@@ -159,13 +161,29 @@ def predict_velocity(
     flattened ``(B, Lt, n*d)`` and is restored to ``(B, Lt, n, d)`` for the MMDiT.
     ``ref_latents`` (optional) are clean reference latents appended after the
     image tokens and conditioned at t=0 ("index_timestep_zero"); the prediction
-    only ever covers the noisy target tokens. Returns the velocity
-    ``noise - clean`` reshaped back to ``(B, C, h, w)``. No time flip / negation:
-    Krea's convention matches toolkit's.
+    only ever covers the noisy target tokens. ``isolate_refs`` restricts ref
+    tokens to attending only among themselves (see ``SingleStreamDiT.forward``),
+    making their per-layer K/V cacheable across steps. ``ref_kv_cache`` is a
+    ``{"kv": None, "mask": None}`` dict enabling that cache (inference only,
+    requires ``isolate_refs``): while ``kv`` is unset the call runs the refs
+    normally and fills the dict with each block's ref K/V; once filled, the ref
+    tokens are dropped from the sequence and the cached K/V are injected as
+    extra attention keys -- identical math, no ref recompute per step. Returns
+    the velocity ``noise - clean`` reshaped back to ``(B, C, h, w)``. No time
+    flip / negation: Krea's convention matches toolkit's.
     """
     patch = model.config.patch
     b, c, h, w = latents.shape
 
+    if ref_kv_cache is not None and not isolate_refs:
+        raise ValueError(
+            "ref_kv_cache requires isolate_refs: cached ref K/V are only "
+            "step-invariant when ref tokens attend solely to each other"
+        )
+    reuse_ref_kv = ref_kv_cache is not None and ref_kv_cache.get("kv") is not None
+    if reuse_ref_kv:
+        ref_latents = None  # the refs are consumed from the cache instead
+
     # Restore the stacked-layer axis flattened in pad_text_features: F -> (n, d).
     n = model.config.txtlayers
     context = context.reshape(
@@ -175,6 +193,7 @@ def predict_velocity(
     img_tokens, pos, mask = prepare(latents, context.shape[1], patch, text_mask)
 
     reflen = 0
+    ref_mask = None
     if ref_latents is not None and any(len(r) > 0 for r in ref_latents):
         ref_tokens, ref_pos, ref_mask = pack_ref_latents(
             ref_latents, patch, img_tokens.device, img_tokens.dtype
@@ -184,7 +203,27 @@ def predict_velocity(
         pos = torch.cat((pos, ref_pos), dim=1)
         mask = torch.cat((mask, ref_mask), dim=1)
 
-    out = model(img=img_tokens, context=context, t=t, pos=pos, mask=mask, reflen=reflen)
+    capture = None
+    if ref_kv_cache is not None and not reuse_ref_kv and reflen > 0:
+        capture = []
+
+    out = model(
+        img=img_tokens,
+        context=context,
+        t=t,
+        pos=pos,
+        mask=mask,
+        reflen=reflen,
+        isolate_refs=isolate_refs,
+        ref_kv_capture=capture,
+        ref_kv_cache=(ref_kv_cache["kv"], ref_kv_cache["mask"])
+        if reuse_ref_kv
+        else None,
+    )
+
+    if capture is not None:
+        ref_kv_cache["kv"] = capture
+        ref_kv_cache["mask"] = ref_mask
 
     # (B, imglen, c*p*p) -> (B, c, h, w)
     velocity = rearrange(
@@ -307,6 +346,16 @@ class Krea2Pipeline:
         x2 = (maxres // align) ** 2
         ts = timesteps(gh * gw, num_inference_steps, x1, x2, y1=y1, y2=y2, mu=mu)
 
+        # With the kv_cache model kwarg (isolated ref attention) the ref K/V
+        # are step-invariant, so the very first model call doubles as the
+        # precompute pass: it runs with the refs in the sequence and fills this
+        # cache; every later call (including step 1's uncond pass) drops the
+        # ref tokens and reuses it.
+        isolate = model.kv_cache
+        ref_cache = None
+        if isolate and ref_latents is not None and any(len(r) > 0 for r in ref_latents):
+            ref_cache = {"kv": None, "mask": None}
+
         # Euler integration of the flow ODE (with optional CFG).
         for tcurr, tprev in zip(ts[:-1], ts[1:]):
             t = torch.full((latents.shape[0],), tcurr, dtype=dtype, device=device)
@@ -317,6 +366,8 @@ class Krea2Pipeline:
                 cond_feats,
                 cond_mask,
                 ref_latents=ref_latents,
+                isolate_refs=isolate,
+                ref_kv_cache=ref_cache,
             )
             if do_cfg:
                 v_uncond = predict_velocity(
@@ -326,6 +377,8 @@ class Krea2Pipeline:
                     uncond_feats,
                     uncond_mask,
                     ref_latents=ref_latents,
+                    isolate_refs=isolate,
+                    ref_kv_cache=ref_cache,
                 )
                 v = v_cond + guidance_scale * (v_cond - v_uncond)
             else: