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

2026-07-11 09:42:06 +00:00 · 2026-07-09 12:00:55 -06:00
parent 6233efe1bb
commit a4bbe167ce
7 changed files with 177 additions and 13 deletions
--- a/extensions_built_in/diffusion_models/krea2/krea2.py
+++ b/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

--- a/extensions_built_in/diffusion_models/krea2/src/mmdit.py
+++ b/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, :]
--- a/extensions_built_in/diffusion_models/krea2/src/pipeline.py
+++ b/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:
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/ui/src/app/jobs/new/SimpleJob.tsx
@@ -330,6 +330,14 @@ export default function SimpleJob({
                />
              </FormGroup>
            )}
+            {modelArch?.additionalSections?.includes('model.model_kwargs.kv_cache') && (
+              <Checkbox
+                label="KV Cache"
+                docKey="model.model_kwargs.kv_cache"
+                checked={jobConfig.config.process[0].model.model_kwargs.kv_cache || false}
+                onChange={value => setJobConfig(value, 'config.process[0].model.model_kwargs.kv_cache')}
+              />
+            )}
            {modelArch?.additionalSections?.includes('model.qie.match_target_res') && (
              <Checkbox
                label="Match Target Res"
--- a/ui/src/app/jobs/new/options.ts
+++ b/ui/src/app/jobs/new/options.ts
@@ -33,6 +33,7 @@ type AdditionalSections =
  | 'model.qie.match_target_res'
  | 'model.assistant_lora_path'
  | 'model.unconditional_lora_path'
+  | 'model.model_kwargs.kv_cache'
  | 'ideogram_4_prompt';

 type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental' | 'audio';
@@ -1104,7 +1105,8 @@ export const modelArchs: ModelArch[] = [
      'config.process[0].model.model_kwargs': [
        {
          edit: true,
-          match_target_res: false,
+          match_target_res: true,
+          kv_cache: true,
        },
        {},
      ],
@@ -1142,7 +1144,8 @@ export const modelArchs: ModelArch[] = [
      'config.process[0].model.model_kwargs': [
        {
          edit: true,
-          match_target_res: false,
+          match_target_res: true,
+          kv_cache: true,
        },
        {},
      ],
@@ -1157,6 +1160,7 @@ export const modelArchs: ModelArch[] = [
      'model.layer_offloading',
      'model.assistant_lora_path',
      'model.qie.match_target_res',
+      'model.model_kwargs.kv_cache',
    ],
  },
  {
@@ -1209,6 +1213,7 @@ export const modelArchs: ModelArch[] = [
      'model.low_vram',
      'model.layer_offloading',
      'model.qie.match_target_res',
+      'model.model_kwargs.kv_cache',
    ],
  },
 ].sort((a, b) => {
--- a/ui/src/docs.tsx
+++ b/ui/src/docs.tsx
@@ -340,6 +340,17 @@ const docs: { [key: string]: ConfigDoc } = {
      </>
    ),
  },
+  'model.model_kwargs.kv_cache': {
+    title: 'KV Cache',
+    description: (
+      <>
+        This will enable KV Cache for control images in a model that supports it. LoRAs trained with this on
+        need to also be inferenced with it, and vice versa. This does not speed up or slow down training, but on inference,
+        the control images only need to be processed once for the entire generation, vs being processed for every step.
+        Which leads to a significant speedup on inference.
+      </>
+    ),
+  },
 };

 export const getDoc = (key: string | null | undefined): ConfigDoc | null => {
--- a/version.py
+++ b/version.py
@@ -1 +1 @@
-VERSION = "0.10.21"
+VERSION = "0.10.22"