FLUX.2 [klein] 9B-KV: KV-cache optimized variant for accelerated multi-reference editing

TL;DR

FLUX.2-klein-9b-kv is a KV-cache-optimized variant of the 4-step distilled FLUX.2 [klein] 9B image-editing flow model. In standard multi-reference image editing, reference-image tokens are reprocessed through every denoising step; this variant computes them once on step 0, caches the per-layer key-value pairs, and on steps 1–3 attends to the cached KVs without recomputing the reference token representations. Reported speedup is up to 2.5× on multi-reference editing while retaining all capabilities of the standard Klein 9B (sub-second T2I and editing in 4 steps). Distributed as a separate HF checkpoint under the FLUX Non-Commercial License with bf16 / FP8 variants, a reference Diffusers pipeline (Flux2KleinKVPipeline), a reference GitHub implementation, and API availability via bfl.ai.

Key claims

KV-caching is applied to the reference-image tokens in a flow-based image-editing DiT: step 0 runs a full forward pass on the reference tokens and stores per-layer K, V; subsequent steps (1–3) read those K, V from cache rather than recomputing them [§How it works].
Up to 2.5× faster inference for multi-reference editing relative to the standard Klein 9B baseline [§Performance].
All capabilities of the standard FLUX.2 [klein] 9B are retained: sub-second T2I, advanced editing in 4 inference steps, 9B flow transformer with 8B Qwen3 text embedder [§Model description].
The variant is shipped as a separate checkpoint (black-forest-labs/FLUX.2-klein-9b-kv), distinct from FLUX.2-klein-9B, with both bf16 and FP8 forms, a reference Diffusers pipeline Flux2KleinKVPipeline, a reference GitHub implementation, and API availability via bfl.ai [§Usage, §Availability].
The acceleration is structural — it exploits the fact that reference tokens carry information that does not need to be refined across denoising steps, while latent (noisy generation) tokens still go through the full DiT every step [§How it works].
License is unchanged from the rest of the FLUX.2 [klein] 9B family: FLUX Non-Commercial License [§License].

Method

In a standard flow-/diffusion-based image-editing DiT, the network attends jointly over (a) the noisy latent tokens being denoised and (b) reference-image tokens supplied as conditioning. At each denoising step, both sets are projected to Q, K, V and the full attention is computed. Because the reference tokens do not change across steps, their K and V are also unchanged — they are recomputed redundantly.

The KV-cache variant exploits this. On step 0 it runs a full forward pass that includes the reference tokens, capturing per-layer K and V for the reference-token positions into a cache. On steps 1, 2, 3 the forward pass processes only the noisy latent tokens; when each attention layer needs K, V at reference positions, it reads them from the cache rather than re-projecting them through the network. The distilled student already runs in only 4 steps, so caching across the remaining 3 steps captures most of the available redundancy.

The packaging is the standard Klein 9B-family release: a separate HF checkpoint, a reference inference repo on GitHub, a Diffusers Flux2KleinKVPipeline, bf16 and FP8 weights, and bfl.ai API access. A community GGUF quantization (Q2–Q8) and ComfyUI workflows appeared shortly after release.

Results

The model card reports:

Up to 2.5× faster inference on multi-reference editing relative to standard Klein 9B [§Performance, §How it works].
All capabilities of the standard Klein 9B retained — no quality degradation is reported on the card [§Model description].
Particularly beneficial for: editing with multiple reference images, generating variations with the same references, and interactive editing applications where reference images are held constant across many user iterations [§Use cases].

No FID / CLIP / Elo numbers, no plots, no per-task latency breakdowns on the model card itself.

Why it’s interesting

This is the simplest possible “KV-cache transplant” from autoregressive LM inference to flow-/diffusion-based image editing, applied to the sub-population of tokens (reference images) whose representations are invariant across denoising steps. It pairs directly with two threads the wiki has been tracking: it sits on top of step-distillation (the underlying Klein 9B is the 4-step distilled checkpoint covered in FLUX.2 [klein]: Towards Interactive Visual Intelligence), and it contrasts with activation-reuse approaches like TEACache / First Block Caching / DeepCache (under the candidate concept approximate-step-caching, see Run FLUX.1-dev three times faster) — Klein-KV reuses exact KVs for a partition of tokens across all steps, rather than reusing approximate full-network activations across adjacent steps. The use case the originating Slack note flags — applicability to omni — is well-posed: any Luma model that consumes multi-image reference conditioning (omni-style identity / control inputs) where the references stay fixed across denoising or AR-rollout steps has the same redundancy, and the same fix should drop in.