diff --git a/README.md b/README.md
index 2d49540d9..885bce96a 100755
--- a/README.md
+++ b/README.md
@@ -17,6 +17,7 @@
 [![Unit Tests](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml/badge.svg)](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml)
 
 # What's new?
+- **`2026/08/28`**: Flux2.Klein text to image and image editing (w/ KV Cache) is now supported.
 - **`2026/07/14`**: Automatic attention tile-size (`block_q`/`block_kv`) search for Wan is now supported.
 - **`2026/06/26`**: 2D ring (USP) attention with a custom splash kernel is now supported for Wan (`ulysses_ring_custom`), splitting context parallelism into an intra-chip Ulysses axis and a cross-chip ring axis.
 - **`2026/04/16`**: Support for Tokamax Ring Attention kernel is now added.
@@ -49,6 +50,7 @@ MaxDiffusion supports
 * Stable Diffusion 2.1 (training and inference)
 * Stable Diffusion XL (training and inference).
 * Flux Dev and Schnell (Training and inference).
+* Flux.2-Klein 4B & 9B (text-to-image and multi-image editing with KV-Cache).
 * Stable Diffusion Lightning (inference).
 * Hyper-SD XL LoRA loading (inference).
 * Load Multiple LoRA (SDXL inference).
@@ -759,6 +761,8 @@ The optimal attention tile sizes (`block_q` / `block_kv`) depend on the sequence
 
   Flux.2-Klein provides ultra-fast 4-step image generation using Qwen3 text embeddings and FLUX.2 transformer blocks.
 
+  #### Text-to-Image Generation:
+
   Flux.2-Klein 4B:
 
   ```bash
@@ -770,6 +774,22 @@ The optimal attention tile sizes (`block_q` / `block_kv`) depend on the sequence
   ```bash
   python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.yml run_name=flux2klein_9b prompt="A detailed vector illustration of a robotic hummingbird"
   ```
+
+  #### Multi-Reference Image Editing:
+
+  Flux.2-Klein supports multi-reference image editing conditioned on up to 4 reference images via the `image_paths` CLI flag.
+
+  Flux.2-Klein 9B Image Editing:
+
+  ```bash
+  python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.yml run_name=flux2klein_9b_image_edit prompt="change the lighting to evening" image_paths="['src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b.png']"
+  ```
+
+  The 9B model also supports KV-Cache for faster inference, and can be toggled with the `use_kv=True` CLI flag:
+
+  ```bash
+  python src/maxdiffusion/generate_flux2klein.py src/maxdiffusion/configs/base_flux2klein_9B.yml run_name=flux2klein_9b_kv_edit prompt="change the lighting to evening" image_paths="['src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b.png']" use_kv=True
+  ```
   ## Fused Attention for GPU:
   Fused Attention for GPU is supported via TransformerEngine. Installation instructions:
 
diff --git a/src/maxdiffusion/configs/base_flux2klein.yml b/src/maxdiffusion/configs/base_flux2klein.yml
index 03758a1b7..c579fc889 100644
--- a/src/maxdiffusion/configs/base_flux2klein.yml
+++ b/src/maxdiffusion/configs/base_flux2klein.yml
@@ -43,6 +43,7 @@ base_shift: 0.5
 max_shift: 1.15
 image_paths: []
 use_base2_exp: True
+use_kv: False
 
 
 unet_checkpoint: ''
diff --git a/src/maxdiffusion/configs/base_flux2klein_9B.yml b/src/maxdiffusion/configs/base_flux2klein_9B.yml
index 2d061421c..48d591258 100644
--- a/src/maxdiffusion/configs/base_flux2klein_9B.yml
+++ b/src/maxdiffusion/configs/base_flux2klein_9B.yml
@@ -43,6 +43,7 @@ base_shift: 0.5
 max_shift: 1.15
 image_paths: []
 use_base2_exp: True
+use_kv: False
 
 
 unet_checkpoint: ''
diff --git a/src/maxdiffusion/generate_flux2klein.py b/src/maxdiffusion/generate_flux2klein.py
index cecdc7b92..71434c643 100644
--- a/src/maxdiffusion/generate_flux2klein.py
+++ b/src/maxdiffusion/generate_flux2klein.py
@@ -223,6 +223,17 @@ def main(argv):
   repo_id = config.pretrained_model_name_or_path
   if not repo_id:
     raise ValueError("pretrained_model_name_or_path must be specified in configuration YAML or CLI.")
+
+  use_kv = config.use_kv
+  if use_kv:
+    if repo_id in ("black-forest-labs/FLUX.2-klein-4B", "black-forest-labs/FLUX.2-klein-4b"):
+      max_logging.log("[WARNING] KV cache not supported for 4B model, ignoring use_kv=True.")
+      pyconfig._config.keys["use_kv"] = False
+    elif repo_id in ("black-forest-labs/FLUX.2-klein-9B", "black-forest-labs/FLUX.2-klein-9b"):
+      repo_id = "black-forest-labs/FLUX.2-klein-9b-kv"
+      pyconfig._config.keys["pretrained_model_name_or_path"] = repo_id
+      max_logging.log(f"[INFO] use_kv=True: switched pretrained_model_name_or_path to KV model variant: {repo_id}")
+
   max_logging.log(f"Target model detected: {repo_id}")
 
   if os.path.exists(repo_id):
diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
index 5a2754b6d..7b2ba0df7 100644
--- a/src/maxdiffusion/models/attention_flax.py
+++ b/src/maxdiffusion/models/attention_flax.py
@@ -273,32 +273,50 @@ def _select_flash_block_sizes(
     flash_block_sizes: BlockSizes,
     dtype: jnp.dtype,
     attention_kernel: str,
+    preserve_asymmetric_block_sizes: bool = False,
 ) -> BlockSizes:
+  """Select Flash/Splash block sizes.
+
+  Existing MaxDiffusion behavior is preserved by default. When
+  preserve_asymmetric_block_sizes=True, explicitly configured block sizes are
+  honored even when Q and KV have different sequence lengths; the existing
+  padding path makes the tensors compatible with those block sizes.
+  """
   query_seq_len = _flash_sequence_length(query)
   key_seq_len = _flash_sequence_length(key)
 
   q_max_block_size = 1024 if dtype == jnp.bfloat16 else 512
+
   if key_seq_len != query_seq_len:
     kv_max_block_size = ((key_seq_len + 127) // 128) * 128
   else:
     kv_max_block_size = q_max_block_size
 
-  # Custom kernels use a lightweight carrier that omits the standard Splash
-  # backward fields. A remapped/local standard kernel still needs a complete
-  # BlockSizes object, including when cross-attention happens to have q_len ==
-  # kv_len.
+  # Preserve the existing Tokamax conversion behavior.
   if flash_block_sizes is not None and not hasattr(flash_block_sizes, "use_fused_bwd_kernel"):
     flash_block_sizes = _coerce_tokamax_block_sizes(flash_block_sizes)
 
-  # Keep configured block sizes for self-attention, but let
-  # cross-attention derive safe KV-aware sizes when q_len != kv_len.
-  if flash_block_sizes and key_seq_len == query_seq_len:
-    if attention_kernel in ["tokamax_flash", "tokamax_ring"]:
+  # Existing self-attention behavior: configured values are returned unchanged.
+  if flash_block_sizes is not None and key_seq_len == query_seq_len:
+    if attention_kernel in ("tokamax_flash", "tokamax_ring"):
+      return _coerce_tokamax_block_sizes(flash_block_sizes)
+    return flash_block_sizes
+
+  # NEW: opt-in behavior required by Klein KV-cache.
+  #
+  # Q and KV may have different sequence lengths, but _pad_data_for_flash()
+  # pads each sequence independently to the configured block size. Therefore
+  # block_q/block_kv do not need to divide the original sequence lengths.
+  if preserve_asymmetric_block_sizes and flash_block_sizes is not None:
+    if attention_kernel in ("tokamax_flash", "tokamax_ring"):
       return _coerce_tokamax_block_sizes(flash_block_sizes)
     return flash_block_sizes
 
-  block_size_q = flash_block_sizes.block_q if flash_block_sizes else q_max_block_size
-  use_tokamax = attention_kernel in ["tokamax_flash", "tokamax_ring"]
+  # Existing MaxDiffusion cross-attention behavior.
+  block_size_q = flash_block_sizes.block_q if flash_block_sizes is not None else q_max_block_size
+
+  use_tokamax = attention_kernel in ("tokamax_flash", "tokamax_ring")
+
   return splash_attention_kernel.BlockSizes(
       block_q=block_size_q,
       block_kv_compute=min(kv_max_block_size, key_seq_len),
@@ -308,7 +326,7 @@ def _select_flash_block_sizes(
       block_kv_dkv_compute=min(kv_max_block_size, query_seq_len),
       block_q_dq=None if use_tokamax else block_size_q,
       block_kv_dq=None if use_tokamax else min(kv_max_block_size, query_seq_len),
-      use_fused_bwd_kernel=True if use_tokamax else False,
+      use_fused_bwd_kernel=use_tokamax,
   )
 
 
@@ -577,6 +595,7 @@ def _tpu_flash_attention(
     use_base2_exp: bool = False,
     use_experimental_scheduler: bool = False,
     is_causal: bool = False,
+    preserve_asymmetric_block_sizes: bool = False,
 ) -> jax.Array:
   """TPU Flash Attention"""
 
@@ -587,7 +606,14 @@ def _tpu_flash_attention(
   attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2])
   if attention_mask is not None and attention_kernel == "tokamax_ring_custom":
     raise NotImplementedError("tokamax_ring_custom does not support attention_mask.")
-  block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, attention_kernel)
+  block_sizes = _select_flash_block_sizes(
+      query,
+      key,
+      flash_block_sizes,
+      dtype,
+      attention_kernel,
+      preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes,
+  )
 
   q_axis_names = nn.logical_to_mesh_axes(axis_names_q)
   kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv)
@@ -832,6 +858,7 @@ def _ulysses_attention(
     use_experimental_scheduler: bool = False,
     use_fixed_m: bool = False,
     ulysses_attention_chunks: int = 1,
+    preserve_asymmetric_block_sizes: bool = False,
 ) -> jax.Array:
   """Ulysses sequence-parallel attention.
 
@@ -862,7 +889,14 @@ def _ulysses_attention(
     )
 
   if not use_custom_kernel:
-    block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, "flash")
+    block_sizes = _select_flash_block_sizes(
+        query,
+        key,
+        flash_block_sizes,
+        dtype,
+        "flash",
+        preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes,
+    )
 
   q_axis_names = nn.logical_to_mesh_axes(axis_names_q)
   kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv)
@@ -1080,6 +1114,7 @@ def _ulysses_ring_attention(
     use_experimental_scheduler: bool = False,
     ulysses_shards: int = -1,
     ulysses_attention_chunks: int = 1,
+    preserve_asymmetric_block_sizes: bool = False,
 ) -> jax.Array:
   """2D context-parallel attention using a private Ulysses x ring mesh.
 
@@ -1124,7 +1159,14 @@ def _ulysses_ring_attention(
   attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2])
   num_heads = query.shape[1]
 
-  block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, "tokamax_ring")
+  block_sizes = _select_flash_block_sizes(
+      query,
+      key,
+      flash_block_sizes,
+      dtype,
+      "tokamax_ring",
+      preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes,
+  )
 
   q_axis_names = nn.logical_to_mesh_axes(axis_names_q)
   kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv)
@@ -1751,6 +1793,7 @@ def ulysses_kernel(q, k, v, context):
       residual_checkpoint_name=context["residual_checkpoint_name"],
       attention_mask=context["attention_mask"],
       ulysses_attention_chunks=context["ulysses_attention_chunks"],
+      preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False),
   )
 
 
@@ -1773,6 +1816,7 @@ def ulysses_ring_kernel(q, k, v, context):
       use_experimental_scheduler=context["use_experimental_scheduler"],
       ulysses_shards=context["ulysses_shards"],
       ulysses_attention_chunks=context["ulysses_attention_chunks"],
+      preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False),
   )
 
 
@@ -1795,6 +1839,7 @@ def flash_kernel(q, k, v, context):
       use_base2_exp=context["use_base2_exp"],
       use_experimental_scheduler=context["use_experimental_scheduler"],
       is_causal=context.get("is_causal", False),
+      preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False),
   )
 
 
@@ -1817,6 +1862,7 @@ def tokamax_flash_kernel(q, k, v, context):
       use_base2_exp=context["use_base2_exp"],
       use_experimental_scheduler=context["use_experimental_scheduler"],
       is_causal=context.get("is_causal", False),
+      preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False),
   )
 
 
@@ -1839,6 +1885,7 @@ def tokamax_ring_kernel(q, k, v, context):
       use_base2_exp=context["use_base2_exp"],
       use_experimental_scheduler=context["use_experimental_scheduler"],
       is_causal=context.get("is_causal", False),
+      preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False),
   )
 
 
@@ -1859,6 +1906,7 @@ def tokamax_ring_custom_kernel(q, k, v, context):
       attention_mask=context["attention_mask"],
       use_base2_exp=context.get("use_base2_exp", True),
       use_experimental_scheduler=context.get("use_experimental_scheduler", False),
+      preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False),
   )
 
 
@@ -1893,6 +1941,7 @@ def _apply_attention(
     ulysses_shards: int = -1,
     ulysses_attention_chunks: int = 1,
     is_causal: bool = False,
+    preserve_asymmetric_block_sizes: bool = False,
 ):
   """Routes to different attention kernels using a module-level registry."""
 
@@ -1959,11 +2008,13 @@ def _apply_attention(
       "use_memory_efficient_attention": use_memory_efficient_attention,
       "dpa_layer": dpa_layer,
       "is_causal": is_causal,
+      "preserve_asymmetric_block_sizes": preserve_asymmetric_block_sizes,
   }
 
   # Module-level Registry lookup
   if effective_attention_kernel in KERNEL_REGISTRY:
-    return KERNEL_REGISTRY[effective_attention_kernel](query, key, value, context)
+    with jax.named_scope(f"kernel_{effective_attention_kernel}"):
+      return KERNEL_REGISTRY[effective_attention_kernel](query, key, value, context)
 
   raise ValueError(f"Unexpected attention kernel {effective_attention_kernel=}.")
 
@@ -2244,6 +2295,7 @@ def apply_attention(
       key: Array,
       value: Array,
       attention_mask: Array = None,
+      preserve_asymmetric_block_sizes: bool = False,
   ):
     return _apply_attention(
         query=query,
@@ -2270,6 +2322,7 @@ def apply_attention(
         use_experimental_scheduler=self.use_experimental_scheduler if hasattr(self, "use_experimental_scheduler") else False,
         ulysses_shards=(self.ulysses_shards if hasattr(self, "ulysses_shards") else -1),
         ulysses_attention_chunks=(self.ulysses_attention_chunks if hasattr(self, "ulysses_attention_chunks") else 1),
+        preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes,
     )
 
 
@@ -2318,7 +2371,14 @@ def setup(self):
       variables = {}
       self.dpa_layer = functools.partial(dpa_layer.apply, variables)
 
-  def apply_attention(self, query: Array, key: Array, value: Array, attention_mask: Array = None):
+  def apply_attention(
+      self,
+      query: Array,
+      key: Array,
+      value: Array,
+      attention_mask: Array = None,
+      preserve_asymmetric_block_sizes: bool = False,
+  ):
     return _apply_attention(
         query=query,
         key=key,
@@ -2343,6 +2403,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
         ulysses_shards=self.ulysses_shards,
         ulysses_attention_chunks=self.ulysses_attention_chunks,
         is_causal=self.is_causal,
+        preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes,
     )
 
 
diff --git a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py
index 3183f462c..2d536ca1c 100644
--- a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py
+++ b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py
@@ -14,7 +14,7 @@
 limitations under the License.
 """
 
-from typing import Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 import jax
 import math
 import jax.numpy as jnp
@@ -1435,58 +1435,118 @@ def __call__(
       hidden_states: jax.Array,
       encoder_hidden_states: Optional[jax.Array] = None,
       image_rotary_emb: Optional[Tuple[jax.Array, jax.Array]] = None,
-  ) -> Tuple[jax.Array, Optional[jax.Array]]:
+      kv_cache: Optional[Tuple[jax.Array, jax.Array]] = None,
+      kv_cache_mode: Optional[str] = None,
+      num_ref_tokens: int = 0,
+  ) -> Tuple[Tuple[jax.Array, Optional[jax.Array]], Optional[Tuple[jax.Array, jax.Array]]]:
     B, L = hidden_states.shape[:2]
     H, D = self.heads, self.dim_head
 
-    qkv_proj = self.i_qkv(hidden_states).reshape(B, L, 3, H, D)
-    query_proj, key_proj, value_proj = jnp.split(qkv_proj, 3, axis=2)
-    query_proj = self.query_norm(query_proj.squeeze(2))
-    key_proj = self.key_norm(key_proj.squeeze(2))
-    value_proj = value_proj.squeeze(2)
-
-    if encoder_hidden_states is not None:
-      B_enc, L_txt = encoder_hidden_states.shape[:2]
-      encoder_qkv_proj = self.e_qkv(encoder_hidden_states).reshape(B_enc, L_txt, 3, H, D)
-      enc_query_proj, enc_key_proj, enc_value_proj = jnp.split(encoder_qkv_proj, 3, axis=2)
-      enc_query_proj = self.encoder_query_norm(enc_query_proj.squeeze(2))
-      enc_key_proj = self.encoder_key_norm(enc_key_proj.squeeze(2))
-      enc_value_proj = enc_value_proj.squeeze(2)
-
-      query_proj = jnp.concatenate((enc_query_proj, query_proj), axis=1)
-      key_proj = jnp.concatenate((enc_key_proj, key_proj), axis=1)
-      value_proj = jnp.concatenate((enc_value_proj, value_proj), axis=1)
+    with jax.named_scope("qkv_projections"):
+      qkv_proj = self.i_qkv(hidden_states).reshape(B, L, 3, H, D)
+      query_proj, key_proj, value_proj = jnp.split(qkv_proj, 3, axis=2)
+      query_proj = self.query_norm(query_proj.squeeze(2))
+      key_proj = self.key_norm(key_proj.squeeze(2))
+      value_proj = value_proj.squeeze(2)
+
+      if encoder_hidden_states is not None:
+        B_enc, L_txt = encoder_hidden_states.shape[:2]
+        encoder_qkv_proj = self.e_qkv(encoder_hidden_states).reshape(B_enc, L_txt, 3, H, D)
+        enc_query_proj, enc_key_proj, enc_value_proj = jnp.split(encoder_qkv_proj, 3, axis=2)
+        enc_query_proj = self.encoder_query_norm(enc_query_proj.squeeze(2))
+        enc_key_proj = self.encoder_key_norm(enc_key_proj.squeeze(2))
+        enc_value_proj = enc_value_proj.squeeze(2)
+
+        query_proj = jnp.concatenate((enc_query_proj, query_proj), axis=1)
+        key_proj = jnp.concatenate((enc_key_proj, key_proj), axis=1)
+        value_proj = jnp.concatenate((enc_value_proj, value_proj), axis=1)
 
     if image_rotary_emb is not None:
-      if not isinstance(image_rotary_emb, (tuple, list)):
-        image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
+      with jax.named_scope("rope_embeddings"):
+        if not isinstance(image_rotary_emb, (tuple, list)):
+          image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
+        else:
+          image_rotary_emb_reordered = image_rotary_emb
+        query_proj = query_proj.swapaxes(1, 2)
+        key_proj = key_proj.swapaxes(1, 2)
+        query_proj, key_proj = apply_rope(query_proj, key_proj, image_rotary_emb_reordered)
+        query_proj = query_proj.swapaxes(1, 2)
+        key_proj = key_proj.swapaxes(1, 2)
+
+    layer_kv = None
+    num_txt_tokens = encoder_hidden_states.shape[1] if encoder_hidden_states is not None else 0
+
+    with jax.named_scope("joint_attention_op"):
+      if kv_cache_mode == "extract" and num_ref_tokens > 0:
+        ref_start = num_txt_tokens
+        ref_end = num_txt_tokens + num_ref_tokens
+        k_ref = key_proj[:, ref_start:ref_end, :, :]
+        v_ref = value_proj[:, ref_start:ref_end, :, :]
+        layer_kv = (k_ref, v_ref)
+
+        q_txt = query_proj[:, :ref_start]
+        q_ref = query_proj[:, ref_start:ref_end]
+        q_img = query_proj[:, ref_end:]
+
+        q_txt_img = jnp.concatenate([q_txt, q_img], axis=1).reshape(B, -1, H * D)
+        k_all = key_proj.reshape(B, -1, H * D)
+        v_all = value_proj.reshape(B, -1, H * D)
+
+        attn_txt_img = self.attention_op.apply_attention(
+            q_txt_img,
+            k_all,
+            v_all,
+            preserve_asymmetric_block_sizes=True,
+        )
+        attn_txt = attn_txt_img[:, :ref_start]
+        attn_img = attn_txt_img[:, ref_start:]
+
+        q_ref_flat = q_ref.reshape(B, -1, H * D)
+        k_ref_flat = k_ref.reshape(B, -1, H * D)
+        v_ref_flat = v_ref.reshape(B, -1, H * D)
+        attn_ref = self.attention_op.apply_attention(q_ref_flat, k_ref_flat, v_ref_flat)
+
+        attn_output = jnp.concatenate([attn_txt, attn_ref, attn_img], axis=1)
+
+      elif kv_cache_mode == "cached" and kv_cache is not None:
+        k_ref, v_ref = kv_cache
+        k_txt = key_proj[:, :num_txt_tokens]
+        k_img = key_proj[:, num_txt_tokens:]
+        v_txt = value_proj[:, :num_txt_tokens]
+        v_img = value_proj[:, num_txt_tokens:]
+
+        k_all = jnp.concatenate([k_txt, k_ref, k_img], axis=1).reshape(B, -1, H * D)
+        v_all = jnp.concatenate([v_txt, v_ref, v_img], axis=1).reshape(B, -1, H * D)
+        q_all = query_proj.reshape(B, -1, H * D)
+
+        attn_output = self.attention_op.apply_attention(
+            q_all,
+            k_all,
+            v_all,
+            preserve_asymmetric_block_sizes=True,
+        )
+
       else:
-        image_rotary_emb_reordered = image_rotary_emb
-      query_proj = query_proj.swapaxes(1, 2)
-      key_proj = key_proj.swapaxes(1, 2)
-      query_proj, key_proj = apply_rope(query_proj, key_proj, image_rotary_emb_reordered)
-      query_proj = query_proj.swapaxes(1, 2)
-      key_proj = key_proj.swapaxes(1, 2)
+        query_proj = query_proj.reshape(B, -1, H * D)
+        key_proj = key_proj.reshape(B, -1, H * D)
+        value_proj = value_proj.reshape(B, -1, H * D)
 
-    query_proj = query_proj.reshape(B, -1, H * D)
-    key_proj = key_proj.reshape(B, -1, H * D)
-    value_proj = value_proj.reshape(B, -1, H * D)
+        if encoder_hidden_states is not None:
+          query_proj = nn.with_logical_constraint(query_proj, ("activation_batch", "activation_length", "activation_heads"))
+          key_proj = nn.with_logical_constraint(key_proj, ("activation_batch", "activation_length", "activation_heads"))
+          value_proj = nn.with_logical_constraint(value_proj, ("activation_batch", "activation_length", "activation_heads"))
 
-    if encoder_hidden_states is not None:
-      query_proj = nn.with_logical_constraint(query_proj, ("activation_batch", "activation_length", "activation_heads"))
-      key_proj = nn.with_logical_constraint(key_proj, ("activation_batch", "activation_length", "activation_heads"))
-      value_proj = nn.with_logical_constraint(value_proj, ("activation_batch", "activation_length", "activation_heads"))
+        attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
 
-    attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
     context_attn_output = None
-
     if encoder_hidden_states is not None:
-      context_attn_output = attn_output[:, : encoder_hidden_states.shape[1]]
-      attn_output = attn_output[:, encoder_hidden_states.shape[1] :]
-      attn_output = self.i_proj(attn_output)
-      context_attn_output = self.e_proj(context_attn_output)
+      with jax.named_scope("attention_out_projections"):
+        context_attn_output = attn_output[:, : encoder_hidden_states.shape[1]]
+        attn_output = attn_output[:, encoder_hidden_states.shape[1] :]
+        attn_output = self.i_proj(attn_output)
+        context_attn_output = self.e_proj(context_attn_output)
 
-    return attn_output, context_attn_output
+    return (attn_output, context_attn_output), layer_kv
 
 
 class NNXFluxSingleAttention(nnx.Module):
@@ -1649,46 +1709,58 @@ def __call__(
       image_rotary_emb: Tuple[jax.Array, jax.Array],
       temb_mod_img: Optional[jax.Array] = None,
       temb_mod_txt: Optional[jax.Array] = None,
-  ) -> Tuple[jax.Array, jax.Array]:
+      kv_cache: Optional[Tuple[jax.Array, jax.Array]] = None,
+      kv_cache_mode: Optional[str] = None,
+      num_ref_tokens: int = 0,
+  ) -> Tuple[jax.Array, jax.Array, Optional[Tuple[jax.Array, jax.Array]]]:
     shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = jnp.split(temb_mod_img, 6, axis=-1)
     c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = jnp.split(temb_mod_txt, 6, axis=-1)
 
-    shift_msa = jnp.expand_dims(shift_msa, axis=1)
-    scale_msa = jnp.expand_dims(scale_msa, axis=1)
-    gate_msa = jnp.expand_dims(gate_msa, axis=1)
-    shift_mlp = jnp.expand_dims(shift_mlp, axis=1)
-    scale_mlp = jnp.expand_dims(scale_mlp, axis=1)
-    gate_mlp = jnp.expand_dims(gate_mlp, axis=1)
-
-    c_shift_msa = jnp.expand_dims(c_shift_msa, axis=1)
-    c_scale_msa = jnp.expand_dims(c_scale_msa, axis=1)
-    c_gate_msa = jnp.expand_dims(c_gate_msa, axis=1)
-    c_shift_mlp = jnp.expand_dims(c_shift_mlp, axis=1)
-    c_scale_mlp = jnp.expand_dims(c_scale_mlp, axis=1)
-    c_gate_mlp = jnp.expand_dims(c_gate_mlp, axis=1)
-
-    norm1_h = self.norm1(hidden_states) * (1.0 + scale_msa) + shift_msa
-    norm1_enc = self.norm1_context(encoder_hidden_states) * (1.0 + c_scale_msa) + c_shift_msa
-
-    attn_img, attn_txt = self.attn(
-        hidden_states=norm1_h,
-        encoder_hidden_states=norm1_enc,
-        image_rotary_emb=image_rotary_emb,
-    )
+    if temb_mod_img.ndim == 2:
+      shift_msa = jnp.expand_dims(shift_msa, axis=1)
+      scale_msa = jnp.expand_dims(scale_msa, axis=1)
+      gate_msa = jnp.expand_dims(gate_msa, axis=1)
+      shift_mlp = jnp.expand_dims(shift_mlp, axis=1)
+      scale_mlp = jnp.expand_dims(scale_mlp, axis=1)
+      gate_mlp = jnp.expand_dims(gate_mlp, axis=1)
+
+    if temb_mod_txt.ndim == 2:
+      c_shift_msa = jnp.expand_dims(c_shift_msa, axis=1)
+      c_scale_msa = jnp.expand_dims(c_scale_msa, axis=1)
+      c_gate_msa = jnp.expand_dims(c_gate_msa, axis=1)
+      c_shift_mlp = jnp.expand_dims(c_shift_mlp, axis=1)
+      c_scale_mlp = jnp.expand_dims(c_scale_mlp, axis=1)
+      c_gate_mlp = jnp.expand_dims(c_gate_mlp, axis=1)
 
-    hidden_states = hidden_states + gate_msa * attn_img
-    encoder_hidden_states = encoder_hidden_states + c_gate_msa * attn_txt
+    with jax.named_scope("norm1_and_modulation"):
+      norm1_h = self.norm1(hidden_states) * (1.0 + scale_msa) + shift_msa
+      norm1_enc = self.norm1_context(encoder_hidden_states) * (1.0 + c_scale_msa) + c_shift_msa
+
+    with jax.named_scope("double_attention"):
+      (attn_img, attn_txt), layer_kv = self.attn(
+          hidden_states=norm1_h,
+          encoder_hidden_states=norm1_enc,
+          image_rotary_emb=image_rotary_emb,
+          kv_cache=kv_cache,
+          kv_cache_mode=kv_cache_mode,
+          num_ref_tokens=num_ref_tokens,
+      )
 
-    norm2_h = self.norm2(hidden_states) * (1.0 + scale_mlp) + shift_mlp
-    norm2_enc = self.norm2_context(encoder_hidden_states) * (1.0 + c_scale_mlp) + c_shift_mlp
+      hidden_states = hidden_states + gate_msa * attn_img
+      encoder_hidden_states = encoder_hidden_states + c_gate_msa * attn_txt
 
-    mlp_output = self.ff(norm2_h)
-    encoder_mlp_output = self.ff_context(norm2_enc)
+    with jax.named_scope("norm2_and_modulation"):
+      norm2_h = self.norm2(hidden_states) * (1.0 + scale_mlp) + shift_mlp
+      norm2_enc = self.norm2_context(encoder_hidden_states) * (1.0 + c_scale_mlp) + c_shift_mlp
 
-    hidden_states = hidden_states + gate_mlp * mlp_output
-    encoder_hidden_states = encoder_hidden_states + c_gate_mlp * encoder_mlp_output
+    with jax.named_scope("double_mlp"):
+      mlp_output = self.ff(norm2_h)
+      encoder_mlp_output = self.ff_context(norm2_enc)
 
-    return encoder_hidden_states, hidden_states
+      hidden_states = hidden_states + gate_mlp * mlp_output
+      encoder_hidden_states = encoder_hidden_states + c_gate_mlp * encoder_mlp_output
+
+    return encoder_hidden_states, hidden_states, layer_kv
 
 
 class NNXFluxSingleTransformerBlock(nnx.Module):
@@ -1769,50 +1841,117 @@ def __call__(
       temb: jax.Array,
       image_rotary_emb: Tuple[jax.Array, jax.Array],
       temb_mod: Optional[jax.Array] = None,
-  ) -> jax.Array:
+      kv_cache: Optional[Tuple[jax.Array, jax.Array]] = None,
+      kv_cache_mode: Optional[str] = None,
+      num_txt_tokens: int = 0,
+      num_ref_tokens: int = 0,
+  ) -> Tuple[jax.Array, Optional[Tuple[jax.Array, jax.Array]]]:
     residual = hidden_states
     shift_msa, scale_msa, gate = jnp.split(temb_mod, 3, axis=-1)
-    shift_msa = jnp.expand_dims(shift_msa, axis=1)
-    scale_msa = jnp.expand_dims(scale_msa, axis=1)
-    gate = jnp.expand_dims(gate, axis=1)
 
-    norm_hidden_states = self.norm(hidden_states)
-    norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
+    if temb_mod.ndim == 2:
+      shift_msa = jnp.expand_dims(shift_msa, axis=1)
+      scale_msa = jnp.expand_dims(scale_msa, axis=1)
+      gate = jnp.expand_dims(gate, axis=1)
 
-    qkv, mlp = jnp.split(self.linear1(norm_hidden_states), [3 * self.dim], axis=-1)
-    qkv = nn.with_logical_constraint(qkv, ("activation_batch", "activation_length", "activation_embed"))
-    mlp = nn.with_logical_constraint(mlp, ("activation_batch", "activation_length", "activation_embed"))
+    with jax.named_scope("norm_and_modulation"):
+      norm_hidden_states = self.norm(hidden_states)
+      norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
 
-    B, L = hidden_states.shape[:2]
-    H, D = self.num_attention_heads, qkv.shape[-1] // (self.num_attention_heads * 3)
-    qkv_proj = qkv.reshape(B, L, 3, H, D).transpose(2, 0, 3, 1, 4)
-    q, k, v = qkv_proj
+    with jax.named_scope("linear1_qkv_and_mlp"):
+      qkv, mlp = jnp.split(self.linear1(norm_hidden_states), [3 * self.dim], axis=-1)
+      qkv = nn.with_logical_constraint(qkv, ("activation_batch", "activation_length", "activation_embed"))
+      mlp = nn.with_logical_constraint(mlp, ("activation_batch", "activation_length", "activation_embed"))
 
-    q = self.attn.query_norm(q)
-    k = self.attn.key_norm(k)
+      B, L = hidden_states.shape[:2]
+      H, D = self.num_attention_heads, qkv.shape[-1] // (self.num_attention_heads * 3)
+      qkv_proj = qkv.reshape(B, L, 3, H, D).transpose(2, 0, 3, 1, 4)
+      q, k, v = qkv_proj
+
+      q = self.attn.query_norm(q)
+      k = self.attn.key_norm(k)
 
     if image_rotary_emb is not None:
-      if isinstance(image_rotary_emb, (tuple, list)):
-        image_rotary_emb_reordered = image_rotary_emb
-      else:
-        image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
-      q, k = apply_rope(q, k, image_rotary_emb_reordered)
+      with jax.named_scope("rope_embeddings"):
+        if isinstance(image_rotary_emb, (tuple, list)):
+          image_rotary_emb_reordered = image_rotary_emb
+        else:
+          image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
+        q, k = apply_rope(q, k, image_rotary_emb_reordered)
+
+    layer_kv = None
+    with jax.named_scope("single_attention_op"):
+      if kv_cache_mode == "extract" and num_ref_tokens > 0:
+        ref_start = num_txt_tokens
+        ref_end = num_txt_tokens + num_ref_tokens
+        k_ref = k[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3)
+        v_ref = v[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3)
+        layer_kv = (k_ref, v_ref)
+
+        q_txt_img = (
+            jnp.concatenate([q[:, :, :ref_start, :], q[:, :, ref_end:, :]], axis=2)
+            .transpose(0, 2, 1, 3)
+            .reshape(B, -1, H * D)
+        )
+        k_all = k.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        v_all = v.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+
+        attn_txt_img = self.attn.attention_op.apply_attention(
+            q_txt_img,
+            k_all,
+            v_all,
+            preserve_asymmetric_block_sizes=True,
+        )
+        attn_txt = attn_txt_img[:, :ref_start]
+        attn_img = attn_txt_img[:, ref_start:]
 
-    q = q.transpose(0, 2, 1, 3).reshape(q.shape[0], q.shape[2], -1)
-    k = k.transpose(0, 2, 1, 3).reshape(k.shape[0], k.shape[2], -1)
-    v = v.transpose(0, 2, 1, 3).reshape(v.shape[0], v.shape[2], -1)
+        q_ref_flat = q[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        k_ref_flat = k[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        v_ref_flat = v[:, :, ref_start:ref_end, :].transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        attn_ref = self.attn.attention_op.apply_attention(q_ref_flat, k_ref_flat, v_ref_flat)
 
-    attn_output = self.attn.attention_op.apply_attention(q, k, v)
+        attn_output = jnp.concatenate([attn_txt, attn_ref, attn_img], axis=1)
 
-    mlp1, mlp2 = jnp.split(mlp, 2, axis=-1)
-    mlp_activated = nnx.silu(mlp1) * mlp2
+      elif kv_cache_mode == "cached" and kv_cache is not None:
+        k_ref, v_ref = kv_cache
+        k_ref_trans = k_ref.transpose(0, 2, 1, 3)
+        v_ref_trans = v_ref.transpose(0, 2, 1, 3)
 
-    attn_mlp = jnp.concatenate([attn_output, mlp_activated], axis=2)
-    attn_mlp = nn.with_logical_constraint(attn_mlp, ("activation_batch", "activation_length", "activation_embed"))
-    hidden_states = self.linear2(attn_mlp)
-    hidden_states = gate * hidden_states
-    hidden_states = residual + hidden_states
-    return hidden_states
+        k_all = (
+            jnp.concatenate([k[:, :, :num_txt_tokens, :], k_ref_trans, k[:, :, num_txt_tokens:, :]], axis=2)
+            .transpose(0, 2, 1, 3)
+            .reshape(B, -1, H * D)
+        )
+        v_all = (
+            jnp.concatenate([v[:, :, :num_txt_tokens, :], v_ref_trans, v[:, :, num_txt_tokens:, :]], axis=2)
+            .transpose(0, 2, 1, 3)
+            .reshape(B, -1, H * D)
+        )
+        q_all = q.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+
+        attn_output = self.attn.attention_op.apply_attention(
+            q_all,
+            k_all,
+            v_all,
+            preserve_asymmetric_block_sizes=True,
+        )
+
+      else:
+        q_flat = q.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        k_flat = k.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        v_flat = v.transpose(0, 2, 1, 3).reshape(B, -1, H * D)
+        attn_output = self.attn.attention_op.apply_attention(q_flat, k_flat, v_flat)
+
+    with jax.named_scope("swiglu_and_linear2"):
+      mlp1, mlp2 = jnp.split(mlp, 2, axis=-1)
+      mlp_activated = nnx.silu(mlp1) * mlp2
+
+      attn_mlp = jnp.concatenate([attn_output, mlp_activated], axis=2)
+      attn_mlp = nn.with_logical_constraint(attn_mlp, ("activation_batch", "activation_length", "activation_embed"))
+      hidden_states = self.linear2(attn_mlp)
+      hidden_states = gate * hidden_states
+      hidden_states = residual + hidden_states
+    return hidden_states, layer_kv
 
 
 class NNXFlux2KleinTransformer2DModel(nnx.Module):
@@ -1978,59 +2117,177 @@ def __call__(
       txt_ids: Optional[jax.Array] = None,
       guidance: Optional[jax.Array] = None,
       return_dict: bool = True,
-  ) -> Union[jax.Array, Transformer2DModelOutput]:
-    hidden_states = self.x_embedder(hidden_states)
-    timestep = timestep * 1000.0
-    if guidance is not None:
-      guidance = guidance * 1000.0
-    temb = self.time_text_embed(timestep, guidance, pooled_projections)
-    temb = temb.astype(hidden_states.dtype)
-
-    temb_silu = nnx.silu(temb)
-    double_stream_mod_img = self.double_stream_modulation_img(temb_silu)
-    double_stream_mod_txt = self.double_stream_modulation_txt(temb_silu)
-    single_stream_mod = self.single_stream_modulation(temb_silu)
-
-    if encoder_hidden_states is not None:
-      encoder_hidden_states = self.context_embedder(encoder_hidden_states)
-
-    if txt_ids.ndim == 3:
-      txt_ids = txt_ids[0]
-    if img_ids.ndim == 3:
-      img_ids = img_ids[0]
-
-    image_rotary_emb = self.pos_embed(img_ids)
-    text_rotary_emb = self.pos_embed(txt_ids)
-    concat_rotary_emb = (
-        jnp.concatenate([text_rotary_emb[0], image_rotary_emb[0]], axis=0),
-        jnp.concatenate([text_rotary_emb[1], image_rotary_emb[1]], axis=0),
-    )
+      kv_cache: Optional[Any] = None,
+      kv_cache_mode: Optional[str] = None,
+      num_ref_tokens: int = 0,
+      ref_fixed_timestep: float = 0.0,
+  ) -> Union[jax.Array, Transformer2DModelOutput, Tuple[Any, Any]]:
+    with jax.named_scope("input_embeddings"):
+      hidden_states = self.x_embedder(hidden_states)
+      timestep_scaled = timestep * 1000.0
+      guidance_scaled = guidance * 1000.0 if guidance is not None else None
+      temb = self.time_text_embed(timestep_scaled, guidance_scaled, pooled_projections)
+      temb = temb.astype(hidden_states.dtype)
+
+      temb_silu = nnx.silu(temb)
+      double_stream_mod_img = self.double_stream_modulation_img(temb_silu)
+      double_stream_mod_txt = self.double_stream_modulation_txt(temb_silu)
+      single_stream_mod = self.single_stream_modulation(temb_silu)
 
-    for double_block in self.double_blocks:
-      encoder_hidden_states, hidden_states = double_block(
-          hidden_states=hidden_states,
-          encoder_hidden_states=encoder_hidden_states,
-          temb=temb,
-          image_rotary_emb=concat_rotary_emb,
-          temb_mod_img=double_stream_mod_img,
-          temb_mod_txt=double_stream_mod_txt,
-      )
+      if encoder_hidden_states is not None:
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
 
-    num_txt_tokens = encoder_hidden_states.shape[1]
-    hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
+      if txt_ids.ndim == 3:
+        txt_ids = txt_ids[0]
+      if img_ids.ndim == 3:
+        img_ids = img_ids[0]
 
-    for single_block in self.single_blocks:
-      hidden_states = single_block(
-          hidden_states=hidden_states,
-          temb=temb,
-          image_rotary_emb=concat_rotary_emb,
-          temb_mod=single_stream_mod,
+      image_rotary_emb = self.pos_embed(img_ids)
+      text_rotary_emb = self.pos_embed(txt_ids)
+      concat_rotary_emb = (
+          jnp.concatenate([text_rotary_emb[0], image_rotary_emb[0]], axis=0),
+          jnp.concatenate([text_rotary_emb[1], image_rotary_emb[1]], axis=0),
       )
 
-    hidden_states = hidden_states[:, num_txt_tokens:, ...]
-    hidden_states = self.norm_out(hidden_states, temb)
-    output = self.proj_out(hidden_states)
+    if kv_cache_mode == "extract" and num_ref_tokens > 0:
+      with jax.named_scope("extract_reference_modulation"):
+        num_img_tokens = hidden_states.shape[1] - num_ref_tokens
+        ref_timestep = jnp.full_like(timestep_scaled, ref_fixed_timestep * 1000.0)
+        ref_temb = self.time_text_embed(ref_timestep, guidance_scaled, pooled_projections).astype(hidden_states.dtype)
+        ref_temb_silu = nnx.silu(ref_temb)
+        ref_double_mod_img = self.double_stream_modulation_img(ref_temb_silu)
+        ref_single_mod = self.single_stream_modulation(ref_temb_silu)
+
+        ref_mod_expanded = jnp.repeat(jnp.expand_dims(ref_double_mod_img, 1), num_ref_tokens, axis=1)
+        img_mod_expanded = jnp.repeat(jnp.expand_dims(double_stream_mod_img, 1), num_img_tokens, axis=1)
+        double_stream_mod_img = jnp.concatenate([ref_mod_expanded, img_mod_expanded], axis=1)
+
+      double_block_caches = []
+      for idx, double_block in enumerate(self.double_blocks):
+        with jax.named_scope(f"double_block_{idx}"):
+          encoder_hidden_states, hidden_states, layer_kv = double_block(
+              hidden_states=hidden_states,
+              encoder_hidden_states=encoder_hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod_img=double_stream_mod_img,
+              temb_mod_txt=double_stream_mod_txt,
+              kv_cache=None,
+              kv_cache_mode="extract",
+              num_ref_tokens=num_ref_tokens,
+          )
+          double_block_caches.append(layer_kv)
+
+      num_txt_tokens = encoder_hidden_states.shape[1]
+      hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+      txt_mod_expanded = jnp.repeat(jnp.expand_dims(single_stream_mod, 1), num_txt_tokens, axis=1)
+      ref_smod_expanded = jnp.repeat(jnp.expand_dims(ref_single_mod, 1), num_ref_tokens, axis=1)
+      img_smod_expanded = jnp.repeat(jnp.expand_dims(single_stream_mod, 1), num_img_tokens, axis=1)
+      single_stream_mod = jnp.concatenate([txt_mod_expanded, ref_smod_expanded, img_smod_expanded], axis=1)
+
+      single_block_caches = []
+      for idx, single_block in enumerate(self.single_blocks):
+        with jax.named_scope(f"single_block_{idx}"):
+          hidden_states, layer_kv = single_block(
+              hidden_states=hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod=single_stream_mod,
+              kv_cache=None,
+              kv_cache_mode="extract",
+              num_txt_tokens=num_txt_tokens,
+              num_ref_tokens=num_ref_tokens,
+          )
+          single_block_caches.append(layer_kv)
+
+      with jax.named_scope("output_norm_and_projection"):
+        hidden_states = hidden_states[:, num_txt_tokens + num_ref_tokens :, ...]
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+      extracted_kv_cache = {
+          "double": tuple(double_block_caches),
+          "single": tuple(single_block_caches),
+          "num_ref_tokens": num_ref_tokens,
+      }
+      if not return_dict:
+        return output, extracted_kv_cache
+      return Transformer2DModelOutput(sample=output), extracted_kv_cache
+
+    elif kv_cache_mode == "cached" and kv_cache is not None:
+      double_caches = kv_cache["double"] if isinstance(kv_cache, dict) else kv_cache[0]
+      single_caches = kv_cache["single"] if isinstance(kv_cache, dict) else kv_cache[1]
+      num_ref = kv_cache.get("num_ref_tokens", 0) if isinstance(kv_cache, dict) else 0
+
+      for idx, double_block in enumerate(self.double_blocks):
+        with jax.named_scope(f"double_block_{idx}"):
+          encoder_hidden_states, hidden_states, _ = double_block(
+              hidden_states=hidden_states,
+              encoder_hidden_states=encoder_hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod_img=double_stream_mod_img,
+              temb_mod_txt=double_stream_mod_txt,
+              kv_cache=double_caches[idx],
+              kv_cache_mode="cached",
+              num_ref_tokens=num_ref,
+          )
+
+      num_txt_tokens = encoder_hidden_states.shape[1]
+      hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+      for idx, single_block in enumerate(self.single_blocks):
+        with jax.named_scope(f"single_block_{idx}"):
+          hidden_states, _ = single_block(
+              hidden_states=hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod=single_stream_mod,
+              kv_cache=single_caches[idx],
+              kv_cache_mode="cached",
+              num_txt_tokens=num_txt_tokens,
+              num_ref_tokens=num_ref,
+          )
+
+      with jax.named_scope("output_norm_and_projection"):
+        hidden_states = hidden_states[:, num_txt_tokens:, ...]
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+      if not return_dict:
+        return (output,)
+      return Transformer2DModelOutput(sample=output)
 
-    if not return_dict:
-      return (output,)
-    return Transformer2DModelOutput(sample=output)
+    else:
+      for idx, double_block in enumerate(self.double_blocks):
+        with jax.named_scope(f"double_block_{idx}"):
+          encoder_hidden_states, hidden_states, _ = double_block(
+              hidden_states=hidden_states,
+              encoder_hidden_states=encoder_hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod_img=double_stream_mod_img,
+              temb_mod_txt=double_stream_mod_txt,
+          )
+
+      num_txt_tokens = encoder_hidden_states.shape[1]
+      hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
+      for idx, single_block in enumerate(self.single_blocks):
+        with jax.named_scope(f"single_block_{idx}"):
+          hidden_states, _ = single_block(
+              hidden_states=hidden_states,
+              temb=temb,
+              image_rotary_emb=concat_rotary_emb,
+              temb_mod=single_stream_mod,
+          )
+
+      with jax.named_scope("output_norm_and_projection"):
+        hidden_states = hidden_states[:, num_txt_tokens:, ...]
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+      if not return_dict:
+        return (output,)
+      return Transformer2DModelOutput(sample=output)
diff --git a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
index 1b8608d98..6b492fb4c 100644
--- a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
+++ b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
@@ -233,6 +233,72 @@ def scan_body(cur_latents, step_idx):
         final_latents, _ = jax.lax.scan(scan_body, latents, steps)
         return final_latents
 
+      @jax.jit(static_argnums=(10, 11))
+      def fused_kv_denoise_loop(
+          t_params,
+          target_latents,
+          ref_latents,
+          target_img_ids,
+          ref_img_ids,
+          prompt_embeds,
+          txt_ids,
+          vec,
+          timesteps,
+          sigmas,
+          guidance=None,
+          num_ref_tokens=0,
+      ):
+        sigmas_padded = jnp.concatenate([sigmas, jnp.array([0.0], dtype=sigmas.dtype)])
+        nnx_merged = nnx.merge(g, t_params, r)
+
+        step0_latents = jnp.concatenate([ref_latents, target_latents], axis=1)
+        step0_img_ids = jnp.concatenate([ref_img_ids, target_img_ids], axis=1)
+        t0_val = timesteps[0]
+        t0_vec = jnp.broadcast_to(t0_val / 1000.0, (target_latents.shape[0],))
+
+        out0, kv_cache = nnx_merged(
+            hidden_states=step0_latents,
+            img_ids=step0_img_ids,
+            encoder_hidden_states=prompt_embeds,
+            txt_ids=txt_ids,
+            pooled_projections=vec,
+            timestep=t0_vec,
+            guidance=guidance,
+            return_dict=True,
+            kv_cache_mode="extract",
+            num_ref_tokens=num_ref_tokens,
+        )
+        dt0 = sigmas_padded[1] - sigmas_padded[0]
+        v0 = out0.sample
+        latents_step1 = target_latents + v0 * dt0
+
+        def scan_body(cur_latents, step_idx):
+          t_val = timesteps[step_idx]
+          t_vec = jnp.broadcast_to(t_val / 1000.0, (cur_latents.shape[0],))
+          model_output = nnx_merged(
+              hidden_states=cur_latents,
+              img_ids=target_img_ids,
+              encoder_hidden_states=prompt_embeds,
+              txt_ids=txt_ids,
+              pooled_projections=vec,
+              timestep=t_vec,
+              guidance=guidance,
+              return_dict=True,
+              kv_cache=kv_cache,
+              kv_cache_mode="cached",
+              num_ref_tokens=num_ref_tokens,
+          )
+          sigma = sigmas_padded[step_idx]
+          sigma_next = sigmas_padded[step_idx + 1]
+          dt = sigma_next - sigma
+          v = model_output.sample
+          next_latents = cur_latents + v * dt
+          return next_latents, None
+
+        steps = jnp.arange(1, timesteps.shape[0])
+        final_latents, _ = jax.lax.scan(scan_body, latents_step1, steps)
+        return final_latents
+
     else:
 
       @jax.jit
@@ -284,9 +350,12 @@ def scan_body(cur_latents, step_idx):
         final_latents, _ = jax.lax.scan(scan_body, latents, steps)
         return final_latents
 
+      fused_kv_denoise_loop = fused_denoise_loop
+
     self._jitted_qwen3_forward = qwen3_forward
     self._jitted_transformer_step = transformer_step
     self._jitted_fused_denoise_loop = fused_denoise_loop
+    self._jitted_fused_kv_denoise_loop = fused_kv_denoise_loop
     self._jitted_vae_encode = vae_encode
     self._jitted_vae_decode = vae_decode
 
@@ -349,6 +418,7 @@ def compile_aot_async(
       images=None,
       image=None,
       num_conditioning_images=0,
+      use_kv=None,
   ):
     """Triggers AOT compilation for Qwen3, Flux Transformer, and VAE concurrently using ThreadPoolExecutor."""
     self._setup_jit_functions()
@@ -417,24 +487,54 @@ def compile_qwen3():
 
     num_steps = self._config.num_inference_steps
     dummy_timesteps = put_data_on_devices(jnp.zeros((num_steps,), dtype=jnp.float32), replicated_sharding)
-    dummy_sigmas = put_data_on_devices(jnp.zeros((num_steps,), dtype=jnp.float32), replicated_sharding)
+    dummy_sigmas = put_data_on_devices(jnp.zeros((num_steps + 1,), dtype=jnp.float32), replicated_sharding)
+
+    use_kv = self._config.use_kv if use_kv is None else use_kv
+    dummy_ref_latents = (
+        put_data_on_devices(jnp.zeros((batch_size, total_ref_tokens, 128), dtype=jnp.float32), data_sharding)
+        if total_ref_tokens > 0
+        else None
+    )
+    dummy_ref_img_ids = (
+        put_data_on_devices(jnp.zeros((batch_size, total_ref_tokens, 4), dtype=jnp.int32), data_sharding)
+        if total_ref_tokens > 0
+        else None
+    )
+    dummy_target_img_ids = put_data_on_devices(jnp.zeros((batch_size, seq_len_img, 4), dtype=jnp.int32), data_sharding)
 
     def compile_transformer():
       t0 = time.perf_counter()
       with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules):
-        self._jitted_fused_denoise_loop.lower(
-            params,
-            dummy_latents,
-            dummy_img_ids,
-            dummy_prompt_embeds,
-            dummy_txt_ids,
-            None,
-            dummy_timesteps,
-            dummy_sigmas,
-            None,
-            seq_len_img,
-        ).compile()
-      max_logging.log(f" -> [AOT COMPILED] Fused Flux Transformer Denoise Scan in {time.perf_counter() - t0:.2f}s")
+        if use_kv and total_ref_tokens > 0 and self._jitted_fused_kv_denoise_loop is not None:
+          self._jitted_fused_kv_denoise_loop.lower(
+              params,
+              dummy_target_latents,
+              dummy_ref_latents,
+              dummy_target_img_ids,
+              dummy_ref_img_ids,
+              dummy_prompt_embeds,
+              dummy_txt_ids,
+              None,
+              dummy_timesteps,
+              dummy_sigmas,
+              None,
+              num_ref_tokens=total_ref_tokens,
+          ).compile()
+          max_logging.log(f" -> [AOT COMPILED] Fused Flux Transformer KV Denoise Scan in {time.perf_counter() - t0:.2f}s")
+        else:
+          self._jitted_fused_denoise_loop.lower(
+              params,
+              dummy_latents,
+              dummy_img_ids,
+              dummy_prompt_embeds,
+              dummy_txt_ids,
+              None,
+              dummy_timesteps,
+              dummy_sigmas,
+              None,
+              seq_len_img,
+          ).compile()
+          max_logging.log(f" -> [AOT COMPILED] Fused Flux Transformer Denoise Scan in {time.perf_counter() - t0:.2f}s")
 
     def compile_vae():
       t0 = time.perf_counter()
@@ -508,6 +608,7 @@ def __call__(
       image: Optional[Union[Any, List[Any]]] = None,
       use_latents: bool = False,
       latents: Optional[Any] = None,
+      use_kv: Optional[bool] = None,
       measure_time: bool = False,
       timing: bool = False,
       warmup: bool = False,
@@ -617,6 +718,8 @@ def put_data_on_devices(x, sharding):
         ref_img_ids_val = prepare_multi_image_ids(norm_ref_latents, scale=10)
         if ref_img_ids_val.shape[0] == 1 and batch_size > 1:
           ref_img_ids_val = jnp.repeat(ref_img_ids_val, batch_size, axis=0)
+        ref_latents_jax = jnp.concatenate(packed_ref_latents, axis=1)
+        num_ref_tokens = ref_latents_jax.shape[1]
         img_ids_val = jnp.concatenate([target_img_ids_val, ref_img_ids_val], axis=1)
         latents_jax = jnp.concatenate([latents_jax] + packed_ref_latents, axis=1)
         max_logging.log(f"  [PIPELINE] Joint latents shape: {latents_jax.shape}, Joint img_ids shape: {img_ids_val.shape}")
@@ -629,6 +732,9 @@ def put_data_on_devices(x, sharding):
           max_logging.log(f" -> [TIMING] Reference Image Encoding (VAE): {trace['vae_encode']:.4f} seconds")
       else:
         img_ids_val = target_img_ids_val
+        packed_ref_latents = []
+        ref_latents_jax = None
+        num_ref_tokens = 0
         trace["vae_encode"] = 0.0
         trace["image_encoding"] = 0.0
 
@@ -734,26 +840,52 @@ def put_data_on_devices(x, sharding):
       replicated_sharding = jax.sharding.NamedSharding(self.mesh, P())
       timesteps_device = put_data_on_devices(scheduler_state.timesteps, replicated_sharding)
       sigmas_device = put_data_on_devices(scheduler_state.sigmas, replicated_sharding)
-
+      use_kv = self._config.use_kv if use_kv is None else use_kv
       do_prof_denoise = profile_target in ("all", "denoise")
       if do_prof_denoise:
         tb_dir = self._config.tensorboard_dir
         jax.profiler.start_trace(os.path.join(tb_dir, "profile_denoise"))
-      with jax.named_scope("fused_flux_denoise_loop"):
-        latents_jax = self._jitted_fused_denoise_loop(
-            params,
-            latents_jax,
-            img_ids_val,
-            prompt_embeds_jax,
-            txt_ids_val,
-            vec_val,
-            timesteps_device,
-            sigmas_device,
-            guidance_vec_val,
-            seq_len_img,
-        )
-        if timing:
-          latents_jax.block_until_ready()
+
+      if use_kv and len(packed_ref_latents) > 0:
+        ref_latents_device = put_data_on_devices(ref_latents_jax, data_sharding)
+        ref_img_ids_device = put_data_on_devices(ref_img_ids_val, data_sharding)
+        target_img_ids_device = put_data_on_devices(target_img_ids_val, data_sharding)
+        target_latents_device = put_data_on_devices(latents_jax[:, :seq_len_img, :], data_sharding)
+
+        with jax.named_scope("fused_flux_kv_denoise_loop"):
+          latents_jax = self._jitted_fused_kv_denoise_loop(
+              params,
+              target_latents_device,
+              ref_latents_device,
+              target_img_ids_device,
+              ref_img_ids_device,
+              prompt_embeds_jax,
+              txt_ids_val,
+              vec_val,
+              timesteps_device,
+              sigmas_device,
+              guidance_vec_val,
+              num_ref_tokens,
+          )
+          if timing:
+            latents_jax.block_until_ready()
+      else:
+        with jax.named_scope("fused_flux_denoise_loop"):
+          latents_jax = self._jitted_fused_denoise_loop(
+              params,
+              latents_jax,
+              img_ids_val,
+              prompt_embeds_jax,
+              txt_ids_val,
+              vec_val,
+              timesteps_device,
+              sigmas_device,
+              guidance_vec_val,
+              seq_len_img,
+          )
+          if timing:
+            latents_jax.block_until_ready()
+
       if do_prof_denoise:
         jax.profiler.stop_trace()
 
diff --git a/src/maxdiffusion/tests/attention_block_sizes_test.py b/src/maxdiffusion/tests/attention_block_sizes_test.py
new file mode 100644
index 000000000..cc33b5e5f
--- /dev/null
+++ b/src/maxdiffusion/tests/attention_block_sizes_test.py
@@ -0,0 +1,127 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import unittest
+import jax.numpy as jnp
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
+from maxdiffusion.models.attention_flax import (
+    _select_flash_block_sizes,
+    _pad_data_for_flash,
+)
+
+
+class AttentionBlockSizesTest(unittest.TestCase):
+  """Unit tests for Flash/Splash block-size selection and padding."""
+
+  def test_symmetric_configured_behavior_unchanged(self):
+    """Test A: User-configured symmetric BlockSizes are returned unchanged."""
+    bs = splash_attention_kernel.BlockSizes(
+        block_q=2048,
+        block_kv=1024,
+        block_kv_compute=512,
+        block_q_dkv=1024,
+        block_kv_dkv=512,
+        block_kv_dkv_compute=256,
+        block_q_dq=512,
+        block_kv_dq=256,
+        use_fused_bwd_kernel=False,
+    )
+    q = jnp.zeros((1, 1, 4096, 128), dtype=jnp.bfloat16)
+    k = jnp.zeros((1, 1, 4096, 128), dtype=jnp.bfloat16)
+
+    result = _select_flash_block_sizes(
+        query=q,
+        key=k,
+        flash_block_sizes=bs,
+        dtype=jnp.bfloat16,
+        attention_kernel="flash",
+        preserve_asymmetric_block_sizes=False,
+    )
+    self.assertIs(result, bs)
+
+  def test_default_asymmetric_behavior_unchanged(self):
+    """Test B: Default asymmetric behavior derives KV-aware BlockSizes matching main."""
+    bs = splash_attention_kernel.BlockSizes(
+        block_q=1024,
+        block_kv=1024,
+        block_kv_compute=1024,
+        block_q_dkv=1024,
+        block_kv_dkv=1024,
+        block_kv_dkv_compute=1024,
+        block_q_dq=1024,
+        block_kv_dq=1024,
+        use_fused_bwd_kernel=False,
+    )
+    q = jnp.zeros((1, 1, 512, 128), dtype=jnp.bfloat16)
+    k = jnp.zeros((1, 1, 4096, 128), dtype=jnp.bfloat16)
+
+    result = _select_flash_block_sizes(
+        query=q,
+        key=k,
+        flash_block_sizes=bs,
+        dtype=jnp.bfloat16,
+        attention_kernel="flash",
+        preserve_asymmetric_block_sizes=False,
+    )
+    self.assertEqual(result.block_q, 1024)
+    self.assertEqual(result.block_kv, 4096)
+    self.assertEqual(result.block_kv_compute, 4096)
+
+  def test_klein_kv_asymmetric_config_preserved(self):
+    """Test C: Klein KV asymmetric config is preserved when preserve_asymmetric_block_sizes=True."""
+    klein_bs = splash_attention_kernel.BlockSizes(
+        block_q=4608,
+        block_kv=1024,
+        block_kv_compute=1024,
+        block_q_dkv=4608,
+        block_kv_dkv=1024,
+        block_kv_dkv_compute=1024,
+        block_q_dq=4608,
+        block_kv_dq=1024,
+        use_fused_bwd_kernel=False,
+    )
+    q = jnp.zeros((1, 1, 4608, 128), dtype=jnp.bfloat16)
+    k = jnp.zeros((1, 1, 8704, 128), dtype=jnp.bfloat16)
+
+    result = _select_flash_block_sizes(
+        query=q,
+        key=k,
+        flash_block_sizes=klein_bs,
+        dtype=jnp.bfloat16,
+        attention_kernel="flash",
+        preserve_asymmetric_block_sizes=True,
+    )
+    self.assertIs(result, klein_bs)
+    self.assertEqual(result.block_q, 4608)
+    self.assertEqual(result.block_kv, 1024)
+    self.assertEqual(result.block_kv_compute, 1024)
+    self.assertNotEqual(8704 % 1024, 0)
+
+  def test_pad_data_for_flash_validity(self):
+    """Test D: _pad_data_for_flash pads 8704 KV tokens to 9216 (multiple of 1024)."""
+    key = jnp.zeros((1, 8704, 8 * 128), dtype=jnp.bfloat16)
+    padded_k, _, original_len = _pad_data_for_flash(
+        key,
+        heads=8,
+        flash_block_size=1024,
+    )
+    self.assertEqual(original_len, 8704)
+    self.assertEqual(padded_k.shape[2], 9216)
+    self.assertEqual(padded_k.shape[2] % 1024, 0)
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxdiffusion/tests/flux2klein_kv_pipeline_e2e_test.py b/src/maxdiffusion/tests/flux2klein_kv_pipeline_e2e_test.py
new file mode 100644
index 000000000..13ccac106
--- /dev/null
+++ b/src/maxdiffusion/tests/flux2klein_kv_pipeline_e2e_test.py
@@ -0,0 +1,466 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import gc
+import json
+import os
+from typing import Optional
+import unittest
+import numpy as np
+import pytest
+from PIL import Image
+from skimage.metrics import structural_similarity as ssim
+import torch
+
+try:
+  from diffusers import Flux2KleinKVPipeline
+except ImportError:
+  Flux2KleinKVPipeline = None
+
+import jax
+import jax.numpy as jnp
+import flax
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+import flax.linen as nn
+from jax.sharding import Mesh
+from transformers import AutoConfig, Qwen2TokenizerFast
+
+from maxdiffusion import max_logging, pyconfig
+from maxdiffusion.max_utils import create_device_mesh, get_flash_block_sizes, device_put_replicated
+from maxdiffusion.models.flux.transformers.transformer_flux_flax import NNXFlux2KleinTransformer2DModel
+from maxdiffusion.models.flux.vae.autoencoder_kl_flux2_nnx import (
+    NNXAutoencoderKLFlux2,
+    load_and_convert_flux2klein_nnx_vae_weights,
+)
+from maxdiffusion.models.flux.util import load_and_convert_flux_klein_nnx_weights
+from maxdiffusion.models.qwen3_flax import FlaxQwen3Model, FlaxQwen3Config
+from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights
+from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler
+from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline
+
+IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+PROMPT = "a vibrant artistic painting combining the dog, car, mountain, and fruit bowl in surreal neon lighting"
+
+
+def compute_psnr(img1: Image.Image, img2: Image.Image) -> float:
+  arr1 = np.array(img1, dtype=np.float64)
+  arr2 = np.array(img2, dtype=np.float64)
+  mse = np.mean((arr1 - arr2) ** 2)
+  if mse == 0:
+    return float("inf")
+  return float(20 * np.log10(255.0 / np.sqrt(mse)))
+
+
+def compute_ssim(img1: Image.Image, img2: Image.Image) -> float:
+  arr1 = np.array(img1.convert("RGB"))
+  arr2 = np.array(img2.convert("RGB"))
+  return float(ssim(arr1, arr2, channel_axis=-1))
+
+
+def get_model_snapshot_dir(model_id: str, env_var: Optional[str] = None) -> str:
+  """Locates the snapshot directory for the given model_id strictly under HF_HOME or explicit env var."""
+  if env_var and os.environ.get(env_var):
+    path = os.environ[env_var]
+    if os.path.exists(path):
+      return path
+
+  try:
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(repo_id=model_id, local_files_only=True)
+  except Exception:
+    pass
+
+  hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+  escaped_id = model_id.replace("/", "--")
+  cache_dir = os.path.join(hf_home, "hub", f"models--{escaped_id}", "snapshots")
+  if not os.path.exists(cache_dir):
+    cache_dir = os.path.join(hf_home, f"models--{escaped_id}", "snapshots")
+
+  if not os.path.exists(cache_dir):
+    raise FileNotFoundError(f"Hugging Face cache directory not found for '{model_id}' under HF_HOME ({hf_home}).")
+
+  snapshots = [s for s in os.listdir(cache_dir) if not s.startswith(".")]
+  if not snapshots:
+    raise FileNotFoundError(f"No snapshot directory found for '{model_id}' in {cache_dir}.")
+
+  return os.path.join(cache_dir, snapshots[0])
+
+
+class TestFlux2KleinKVPipelineE2EBF16Parity(unittest.TestCase):
+  """End-to-end parity test comparing Diffusers Flux2KleinKVPipeline vs MaxDiffusion FlaxFlux2KleinPipeline (use_kv=True) in bfloat16."""
+
+  @classmethod
+  def setUpClass(cls):
+    cls.model_path = get_model_snapshot_dir("black-forest-labs/FLUX.2-klein-9b-kv", "FLUX2_KLEIN_KV_MODEL_PATH")
+    cls.work_dir = "/tmp/flux2klein_kv_e2e"
+    os.makedirs(cls.work_dir, exist_ok=True)
+
+    cls.height = 256
+    cls.width = 256
+    cls.num_inference_steps = 4
+    cls.seed = int(os.getenv("FLUX2_KLEIN_E2E_SEED", "42"))
+
+    # 1. Load 4 real reference images (256x256)
+    ref_dir = os.path.join(THIS_DIR, "images", "flux2klein")
+    cls.ref_images = []
+    if os.path.exists(ref_dir):
+      for i in range(4):
+        p = os.path.join(ref_dir, f"ref_image_{i}.png")
+        if os.path.exists(p):
+          cls.ref_images.append(Image.open(p).convert("RGB").resize((256, 256), Image.Resampling.BICUBIC))
+
+    if len(cls.ref_images) < 4:
+      colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
+      for i, c in enumerate(colors):
+        arr = np.full((256, 256, 3), c, dtype=np.uint8)
+        cls.ref_images.append(Image.fromarray(arr))
+
+    # 2. Generate shared starting noise latents (1, 32, 32, 32)
+    rng = np.random.RandomState(cls.seed)
+    latents_unpacked = rng.randn(1, 32, cls.height // 8, cls.width // 8).astype(np.float32)
+    cls.latents_unpacked_jax = jnp.array(latents_unpacked)
+
+    # Prepare packed latents for PyTorch: (1, 32, H/16, 2, W/16, 2) -> permute(0, 1, 3, 5, 2, 4) -> reshape(1, 128, H/16, W/16)
+    latents_unpacked_pt = torch.from_numpy(latents_unpacked)
+    latents_pt_packed = latents_unpacked_pt.view(1, 32, cls.height // 16, 2, cls.width // 16, 2)
+    latents_pt_packed = latents_pt_packed.permute(0, 1, 3, 5, 2, 4)
+    cls.latents_pt_packed = latents_pt_packed.reshape(1, 128, cls.height // 16, cls.width // 16)
+
+  @pytest.mark.skipif(
+      IN_GITHUB_ACTIONS or Flux2KleinKVPipeline is None,
+      reason="Requires TPU, full FLUX.2-Klein 9B weights, and Flux2KleinKVPipeline support in diffusers",
+  )
+  def test_flux2klein_kv_pipeline_bf16_parity(self):
+    """Executes PyTorch Diffusers Flux2KleinKVPipeline in BF16 and MaxDiffusion FlaxFlux2KleinPipeline in BF16 (use_kv=True) and compares outputs."""
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("FLUX.2-KLEIN-9B KV CACHE END-TO-END BF16 PARITY TEST")
+    max_logging.log("=" * 80)
+    max_logging.log(f"Model Path:          {self.model_path}")
+    max_logging.log(f"Prompt:              '{PROMPT}'")
+    max_logging.log(f"Number of Ref Images:{len(self.ref_images)} (256x256)")
+    max_logging.log(f"Target Resolution:   {self.width}x{self.height}")
+    max_logging.log(f"Inference Steps:     {self.num_inference_steps}")
+
+    # =========================================================================
+    # LEG 1: PyTorch Diffusers Flux2KleinKVPipeline in BF16
+    # =========================================================================
+    max_logging.log("\n" + "-" * 80)
+    max_logging.log("LEG 1: Running PyTorch Diffusers Flux2KleinKVPipeline (bfloat16 on CPU)...")
+    max_logging.log("-" * 80)
+
+    pipe_pt = Flux2KleinKVPipeline.from_pretrained(
+        self.model_path,
+        torch_dtype=torch.bfloat16,
+        local_files_only=True,
+    )
+    pipe_pt.to("cpu")
+
+    with torch.no_grad():
+      pt_out = pipe_pt(
+          prompt=PROMPT,
+          image=self.ref_images,
+          latents=self.latents_pt_packed.to(torch.bfloat16),
+          num_inference_steps=self.num_inference_steps,
+          height=self.height,
+          width=self.width,
+          output_type="pil",
+      ).images[0]
+
+    pt_output_path = os.path.join(self.work_dir, "pt_bf16_output.png")
+    pt_out.save(pt_output_path)
+    max_logging.log(f" -> Saved PyTorch Diffusers BF16 output: {pt_output_path}")
+
+    del pipe_pt
+    gc.collect()
+
+    # =========================================================================
+    # LEG 2: MaxDiffusion FlaxFlux2KleinPipeline in BF16 (use_kv=True)
+    # =========================================================================
+    max_logging.log("\n" + "-" * 80)
+    max_logging.log("LEG 2: Running MaxDiffusion FlaxFlux2KleinPipeline with use_kv=True (bfloat16 on TPU)...")
+    max_logging.log("-" * 80)
+
+    # 1. Device mesh setup
+    active_devices = jax.devices()
+    active_device_count = len(active_devices)
+
+    pyconfig._config = None
+    pyconfig.config = None
+    config_path = os.path.join(THIS_DIR, "..", "configs", "base_flux2klein_9B.yml")
+    args = [
+        None,
+        config_path,
+        "run_name=e2e_kv_parity_test",
+        f"output_dir={self.work_dir}",
+        f"per_device_batch_size={1.0 / active_device_count}",
+        f"height={self.height}",
+        f"width={self.width}",
+        f"num_inference_steps={self.num_inference_steps}",
+        f"seed={self.seed}",
+        "use_kv=True",
+        "weights_dtype=bfloat16",
+        "activations_dtype=bfloat16",
+        "precision=DEFAULT",
+        "attention=tokamax_flash",
+        'flash_block_sizes={"block_q": 512, "block_kv": 512, "block_kv_compute": 512}',
+        "text_encoder_attention=dot_product",
+    ]
+    pyconfig.initialize(args, unittest=True)
+    config = pyconfig.config
+
+    if active_device_count > 1:
+      pyconfig._config.keys["ici_tensor_parallelism"] = active_device_count
+      pyconfig._config.keys["ici_data_parallelism"] = 1
+      pyconfig._config.keys["ici_fsdp_parallelism"] = 1
+      pyconfig._config.keys["ici_context_parallelism"] = 1
+
+    pyconfig._config.keys["flash_block_sizes"] = {
+        "block_q": 512,
+        "block_kv": 512,
+        "block_kv_compute": 512,
+    }
+
+    devices_array = create_device_mesh(config, devices=active_devices)
+    mesh = Mesh(devices_array, config.mesh_axes)
+
+    # 2. Text Encoder & Tokenizer
+    text_encoder_path = os.path.join(self.model_path, "text_encoder")
+    tokenizer_path = os.path.join(self.model_path, "tokenizer")
+    pt_config = AutoConfig.from_pretrained(text_encoder_path)
+    rope_theta = None
+    if hasattr(pt_config, "rope_parameters") and isinstance(pt_config.rope_parameters, dict):
+      rope_theta = pt_config.rope_parameters.get("rope_theta")
+    if rope_theta is None:
+      rope_theta = getattr(pt_config, "rope_theta", None) or getattr(pt_config, "rope_base", None)
+
+    qwen_kwargs = {
+        "vocab_size": pt_config.vocab_size,
+        "hidden_size": pt_config.hidden_size,
+        "intermediate_size": pt_config.intermediate_size,
+        "num_hidden_layers": pt_config.num_hidden_layers,
+        "num_attention_heads": pt_config.num_attention_heads,
+        "num_key_value_heads": pt_config.num_key_value_heads,
+        "max_position_embeddings": pt_config.max_position_embeddings,
+        "rms_norm_eps": pt_config.rms_norm_eps,
+        "dtype": jnp.bfloat16,
+        "attention_kernel": "dot_product",
+        "mesh": mesh,
+        "max_layer_to_run": config.text_encoder_max_layer,
+        "is_causal": True,
+    }
+    if rope_theta is not None:
+      qwen_kwargs["rope_theta"] = rope_theta
+
+    qwen3_config = FlaxQwen3Config(**qwen_kwargs)
+    text_encoder = FlaxQwen3Model(qwen3_config)
+    tokenizer = Qwen2TokenizerFast.from_pretrained(tokenizer_path)
+
+    # 3. NNX Transformer
+    transformer_path = os.path.join(self.model_path, "transformer")
+    transformer_config_json = os.path.join(transformer_path, "config.json")
+    transformer_pt_cfg = {}
+    if os.path.exists(transformer_config_json):
+      with open(transformer_config_json, "r") as f:
+        transformer_pt_cfg = json.load(f)
+
+    num_double_layers = transformer_pt_cfg.get("num_layers", 8)
+    depth = transformer_pt_cfg.get("num_single_layers", 24)
+    num_attention_heads = transformer_pt_cfg.get("num_attention_heads", 32)
+
+    transformer = NNXFlux2KleinTransformer2DModel(
+        rngs=nnx.Rngs(0),
+        in_channels=128,
+        num_layers=num_double_layers,
+        num_single_layers=depth,
+        attention_head_dim=128,
+        num_attention_heads=num_attention_heads,
+        joint_attention_dim=3 * pt_config.hidden_size,
+        pooled_projection_dim=768,
+        guidance_embeds=transformer_pt_cfg.get("guidance_embeds", False),
+        axes_dim=(32, 32, 32, 32),
+        theta=2000.0,
+        mlp_ratio=3.0,
+        attention_kernel=config.attention,
+        flash_min_seq_length=512,
+        flash_block_sizes=get_flash_block_sizes(config),
+        mesh=mesh,
+        dtype=jnp.bfloat16,
+        weights_dtype=jnp.bfloat16,
+        scale_shift_order="scale_shift",
+        use_base2_exp=True,
+    )
+
+    # 4. NNX VAE
+    vae_path = os.path.join(self.model_path, "vae", "diffusion_pytorch_model.safetensors")
+    if not os.path.exists(vae_path):
+      vae_path = os.path.join(self.model_path, "vae")
+    vae = NNXAutoencoderKLFlux2(
+        in_channels=3,
+        out_channels=3,
+        latent_channels=32,
+        block_out_channels=(128, 256, 512, 512),
+        layers_per_block=2,
+        norm_num_groups=32,
+        dtype=jnp.bfloat16,
+        param_dtype=jnp.bfloat16,
+    )
+
+    # 5. Extract mesh shardings for all models
+    abstract_transformer_state = nnx.state(transformer, nnx.Param)
+    abstract_vae_state = nnx.state(vae, nnx.Param)
+
+    def qwen3_init_fn():
+      return text_encoder.init(
+          jax.random.PRNGKey(0), jnp.zeros((1, 512), dtype=jnp.int32), jnp.zeros((1, 512), dtype=jnp.int32)
+      )
+
+    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+      logical_transformer_specs = nnx.get_partition_spec(abstract_transformer_state)
+      logical_vae_specs = nnx.get_partition_spec(abstract_vae_state)
+      abstract_qwen3_vars = jax.eval_shape(qwen3_init_fn)
+      logical_qwen3_specs = nn.get_partition_spec(abstract_qwen3_vars)
+
+      transformer_shardings = nn.logical_to_mesh_sharding(logical_transformer_specs, mesh, config.logical_axis_rules)
+      vae_shardings = nn.logical_to_mesh_sharding(logical_vae_specs, mesh, config.logical_axis_rules)
+      qwen3_shardings = flax.core.freeze(
+          nn.logical_to_mesh_sharding(logical_qwen3_specs, mesh, config.logical_axis_rules)["params"]
+      )
+
+    # 6. Load weights on Host CPU and shard across TPU HBM
+    cpu_device = jax.local_devices(backend="cpu")[0]
+    with jax.default_device(cpu_device):
+      t_params = load_and_convert_flux_klein_nnx_weights(
+          transformer_path,
+          abstract_transformer_state,
+          num_double_layers=num_double_layers,
+          num_single_layers=depth,
+          dtype=jnp.bfloat16,
+      )
+      vae_bn_mean, vae_bn_std = load_and_convert_flux2klein_nnx_vae_weights(vae_path, vae, dtype=jnp.bfloat16)
+      v_params = nnx.state(vae, nnx.Param)
+
+      def unbox_fn(x):
+        import flax.linen.spmd as flax_spmd
+
+        return x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x
+
+      qwen3_params_template = jax.tree_util.tree_map(
+          unbox_fn, abstract_qwen3_vars["params"], is_leaf=lambda k: hasattr(k, "unbox")
+      )
+      qwen3_params_template = flax.core.unfreeze(qwen3_params_template)
+      q_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params_template, qwen3_config)
+      q_params = flax.core.freeze(q_params)
+
+    # Shard onto TPU HBM
+    max_logging.log(" -> Sharding parameters across TPU HBM...")
+    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+      t_params = jax.tree_util.tree_map(device_put_replicated, t_params, transformer_shardings)
+      v_params = jax.tree_util.tree_map(device_put_replicated, v_params, vae_shardings)
+      nnx.update(vae, v_params)
+      q_params = jax.tree_util.tree_map(device_put_replicated, q_params, qwen3_shardings)
+
+    # 7. Scheduler
+    scheduler = FlaxFlowMatchScheduler(
+        num_train_timesteps=1000,
+        shift=1.0,
+        sigma_max=1.0,
+        sigma_min=0.001,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        use_dynamic_shifting=True,
+        time_shift_type="exponential",
+    )
+
+    # 8. Pipeline instantiation & execution
+    pipeline = FlaxFlux2KleinPipeline(
+        transformer=transformer,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        scheduler=scheduler,
+        config=config,
+        mesh=mesh,
+    )
+
+    pipeline.compile_aot_async(
+        params=t_params,
+        vae_params=v_params,
+        qwen3_params=q_params,
+        vae_bn_mean=vae_bn_mean,
+        vae_bn_std=vae_bn_std,
+        batch_size=1,
+        height=self.height,
+        width=self.width,
+        images=self.ref_images,
+        use_kv=True,
+    )
+
+    jax_output_name = "jax_bf16_output.png"
+    pipeline(
+        prompt=PROMPT,
+        params=t_params,
+        vae_params=v_params,
+        qwen3_params=q_params,
+        vae_bn_mean=vae_bn_mean,
+        vae_bn_std=vae_bn_std,
+        transformer_shardings=transformer_shardings,
+        vae_shardings=vae_shardings,
+        qwen3_shardings=qwen3_shardings,
+        height=self.height,
+        width=self.width,
+        num_inference_steps=self.num_inference_steps,
+        batch_size=1,
+        images=self.ref_images,
+        use_latents=True,
+        latents=self.latents_unpacked_jax,
+        use_kv=True,
+        output_dir=self.work_dir,
+        output_name=jax_output_name,
+    )
+
+    jax_output_path = os.path.join(self.work_dir, jax_output_name)
+    self.assertTrue(os.path.exists(jax_output_path), f"JAX output image not found at {jax_output_path}")
+    max_logging.log(f" -> Found MaxDiffusion JAX BF16 output: {jax_output_path}")
+
+    # =========================================================================
+    # LEG 3: Compute Parity Metrics (SSIM & PSNR)
+    # =========================================================================
+    max_logging.log("\n" + "=" * 80)
+    max_logging.log("CROSS-FRAMEWORK BF16 PARITY EVALUATION REPORT")
+    max_logging.log("=" * 80)
+
+    img_pt = Image.open(pt_output_path).convert("RGB")
+    img_jax = Image.open(jax_output_path).convert("RGB")
+
+    score_ssim = compute_ssim(img_jax, img_pt)
+    score_psnr = compute_psnr(img_jax, img_pt)
+
+    max_logging.log(f" -> Structural Similarity (SSIM): {score_ssim:.6f}")
+    max_logging.log(f" -> Peak Signal-to-Noise Ratio (PSNR): {score_psnr:.2f} dB")
+    max_logging.log("=" * 80)
+
+    self.assertGreaterEqual(
+        score_ssim, 0.70, f"End-to-End BF16 SSIM {score_ssim:.6f} is below the required acceptance threshold of 0.70"
+    )
+    max_logging.log("End-to-End BF16 Parity Test PASSED successfully!\n")
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py
index b0c4f3f2f..621300d8e 100644
--- a/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py
+++ b/src/maxdiffusion/tests/generate_flux2klein_smoke_test.py
@@ -23,8 +23,7 @@
 from PIL import Image
 from skimage.metrics import structural_similarity as ssim
 
-from maxdiffusion import pyconfig
-from maxdiffusion import generate_flux2klein
+from maxdiffusion import generate_flux2klein, max_logging, pyconfig
 
 IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -76,7 +75,7 @@ def test_flux2klein_4b_smoke(self):
 
     self.assertEqual(base_image.shape, test_image.shape)
     ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
-    print(f"\n[SMOKE TEST 4B] SSIM Score: {ssim_compare:.6f}")
+    max_logging.log(f"\n[SMOKE TEST 4B] SSIM Score: {ssim_compare:.6f}")
     self.assertGreaterEqual(ssim_compare, 0.8)
 
   @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
@@ -121,7 +120,7 @@ def test_flux2klein_9b_smoke(self):
 
     self.assertEqual(base_image.shape, test_image.shape)
     ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
-    print(f"\n[SMOKE TEST 9B] SSIM Score: {ssim_compare:.6f}")
+    max_logging.log(f"\n[SMOKE TEST 9B] SSIM Score: {ssim_compare:.6f}")
     self.assertGreaterEqual(ssim_compare, 0.8)
 
   @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
@@ -170,7 +169,7 @@ def test_flux2klein_4b_image_edit_smoke(self):
 
     self.assertEqual(base_image.shape, test_image.shape)
     ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
-    print(f"\n[SMOKE TEST 4B IMAGE EDIT] SSIM Score: {ssim_compare:.6f}")
+    max_logging.log(f"\n[SMOKE TEST 4B IMAGE EDIT] SSIM Score: {ssim_compare:.6f}")
     self.assertGreaterEqual(ssim_compare, 0.8)
 
   @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
@@ -219,7 +218,59 @@ def test_flux2klein_9b_image_edit_smoke(self):
 
     self.assertEqual(base_image.shape, test_image.shape)
     ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
-    print(f"\n[SMOKE TEST 9B IMAGE EDIT] SSIM Score: {ssim_compare:.6f}")
+    max_logging.log(f"\n[SMOKE TEST 9B IMAGE EDIT] SSIM Score: {ssim_compare:.6f}")
+    self.assertGreaterEqual(ssim_compare, 0.8)
+
+  @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions (requires TPU HBM)")
+  def test_flux2klein_9b_kv_image_edit_smoke(self):
+    """End-to-end smoke test for Flux.2-klein-9B KV cache image editing at 512x512."""
+    ref_gold_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_9b_kv_image_edit.png")
+    self.assertTrue(os.path.exists(ref_gold_path), f"Golden reference image not found: {ref_gold_path}")
+    base_image = np.array(Image.open(ref_gold_path)).astype(np.uint8)
+
+    input_img_path = os.path.join(THIS_DIR, "images", "flux2klein", "ref_flux2klein_4b.png")
+    self.assertTrue(os.path.exists(input_img_path), f"Input reference image not found: {input_img_path}")
+
+    output_dir = "/tmp/smoke_test_kv_image_edit_9b"
+    os.makedirs(output_dir, exist_ok=True)
+    out_path = os.path.join(output_dir, "flux2klein_generated_image.png")
+    if os.path.exists(out_path):
+      os.remove(out_path)
+
+    pyconfig._config = None
+    pyconfig.config = None
+    args = [
+        None,
+        os.path.join(THIS_DIR, "..", "configs", "base_flux2klein_9B.yml"),
+        "run_name=smoke_test_kv_image_edit_9b",
+        f"output_dir={output_dir}",
+        "jax_cache_dir=/tmp/cache_dir",
+        f"image_paths=['{input_img_path}']",
+        "prompt=change the lighting to evening",
+        "height=512",
+        "width=512",
+        f"per_device_batch_size={1.0 / jax.device_count()}",
+        "seed=42",
+        "use_kv=True",
+        "attention=tokamax_flash",
+        'flash_block_sizes={"block_q": 512, "block_kv": 512, "block_kv_compute": 512}',
+        "weights_dtype=bfloat16",
+        "activations_dtype=bfloat16",
+        "precision=DEFAULT",
+        "num_reps=5",
+        "text_encoder_attention=dot_product",
+    ]
+
+    generate_flux2klein.main(args)
+
+    rep_out_path = os.path.join(output_dir, "rep_1_flux2klein_generated_image.png")
+    final_out_path = rep_out_path if os.path.exists(rep_out_path) else out_path
+    self.assertTrue(os.path.exists(final_out_path), "Smoke test 9B KV image edit failed to produce output image!")
+    test_image = np.array(Image.open(final_out_path)).astype(np.uint8)
+
+    self.assertEqual(base_image.shape, test_image.shape)
+    ssim_compare = ssim(base_image, test_image, channel_axis=-1, data_range=255)
+    max_logging.log(f"\n[SMOKE TEST 9B KV IMAGE EDIT] SSIM Score: {ssim_compare:.6f}")
     self.assertGreaterEqual(ssim_compare, 0.8)
 
 
diff --git a/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_kv_image_edit.png b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_kv_image_edit.png
new file mode 100644
index 000000000..992159ab9
Binary files /dev/null and b/src/maxdiffusion/tests/images/flux2klein/ref_flux2klein_9b_kv_image_edit.png differ