Adding KV-Cache Support to Flux2.Klein Image Editing

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:
 

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: ''

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: ''

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):

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,
     )