From a4bbe167ce03521bf9052d2349f01b2997d67ac7 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 9 Jul 2026 12:00:55 -0600 Subject: [PATCH] Added reference token attention isolation (kv_cache) for Krea2 edit training. Same training cost with significant inference speed up. 2x inference speedup. --- .../diffusion_models/krea2/krea2.py | 11 +++ .../diffusion_models/krea2/src/mmdit.py | 88 +++++++++++++++++-- .../diffusion_models/krea2/src/pipeline.py | 61 ++++++++++++- ui/src/app/jobs/new/SimpleJob.tsx | 8 ++ ui/src/app/jobs/new/options.ts | 9 +- ui/src/docs.tsx | 11 +++ version.py | 2 +- 7 files changed, 177 insertions(+), 13 deletions(-) diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index 14403863..0fa8cc93 100644 --- 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 diff --git a/extensions_built_in/diffusion_models/krea2/src/mmdit.py b/extensions_built_in/diffusion_models/krea2/src/mmdit.py index 014919be..fec307b3 100644 --- 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, :] diff --git a/extensions_built_in/diffusion_models/krea2/src/pipeline.py b/extensions_built_in/diffusion_models/krea2/src/pipeline.py index 6566d2b6..2e901e6a 100644 --- 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: diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 5f9d9bc1..a12809f3 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -330,6 +330,14 @@ export default function SimpleJob({ /> )} + {modelArch?.additionalSections?.includes('model.model_kwargs.kv_cache') && ( + setJobConfig(value, 'config.process[0].model.model_kwargs.kv_cache')} + /> + )} {modelArch?.additionalSections?.includes('model.qie.match_target_res') && ( { diff --git a/ui/src/docs.tsx b/ui/src/docs.tsx index 6a1c0b5a..e178bbb7 100644 --- 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 => { diff --git a/version.py b/version.py index 01668c37..4d445c57 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.10.21" +VERSION = "0.10.22"