From db24f052d731b8882a7c12287f56c5f3f4544496 Mon Sep 17 00:00:00 2001 From: ciaranbor Date: Sat, 13 Dec 2025 16:14:46 +0000 Subject: [PATCH] Remove redundant text kv cache computation --- .../mflux/pipefusion/distributed_denoising.py | 11 +-- .../engines/mflux/pipefusion/kv_cache.py | 96 +++++++++++-------- .../mflux/pipefusion/patched_blocks.py | 62 ++++++------ 3 files changed, 91 insertions(+), 78 deletions(-) diff --git a/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py b/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py index da56a145..3c7c5b44 100644 --- a/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py +++ b/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py @@ -130,15 +130,16 @@ class DistributedDenoising: def _initialize_kv_caches( self, batch_size: int, - text_seq_len: int, num_img_tokens: int, dtype: mx.Dtype, ) -> None: """Initialize KV caches for both sync and async pipelines. + Note: Caches only store IMAGE K/V, not text K/V. Text K/V is always + computed fresh and doesn't need caching (it's the same for all patches). + Args: batch_size: Batch size - text_seq_len: Length of text sequence num_img_tokens: Number of image tokens dtype: Data type for cache tensors """ @@ -146,7 +147,6 @@ class DistributedDenoising: JointPatchKVCache( batch_size=batch_size, num_heads=24, - text_seq_len=text_seq_len, image_seq_len=num_img_tokens, head_dim=128, dtype=dtype, @@ -157,7 +157,7 @@ class DistributedDenoising: PatchKVCache( batch_size=batch_size, num_heads=24, - total_seq_len=text_seq_len + num_img_tokens, + image_seq_len=num_img_tokens, head_dim=128, dtype=dtype, ) @@ -195,7 +195,6 @@ class DistributedDenoising: if self._joint_kv_caches is None: self._initialize_kv_caches( batch_size=batch_size, - text_seq_len=text_seq_len, num_img_tokens=num_img_tokens, dtype=hidden_states.dtype, ) @@ -263,6 +262,7 @@ class DistributedDenoising: hidden_states=hidden_states, text_embeddings=text_embeddings, rotary_embeddings=image_rotary_embeddings, + text_seq_len=text_seq_len, ) # Send to next stage if not last @@ -339,7 +339,6 @@ class DistributedDenoising: if self._joint_kv_caches is None: self._initialize_kv_caches( batch_size=batch_size, - text_seq_len=text_seq_len, num_img_tokens=num_img_tokens, dtype=full_hidden.dtype, ) diff --git a/src/exo/worker/engines/mflux/pipefusion/kv_cache.py b/src/exo/worker/engines/mflux/pipefusion/kv_cache.py index 81e37868..3cfee70a 100644 --- a/src/exo/worker/engines/mflux/pipefusion/kv_cache.py +++ b/src/exo/worker/engines/mflux/pipefusion/kv_cache.py @@ -2,47 +2,37 @@ import mlx.core as mx class JointPatchKVCache: - """KV cache for joint attention where text and image are processed separately. + """KV cache for joint attention - stores only IMAGE K/V (not text). Used for joint transformer blocks (19 double blocks in Flux). - Separates text and image portions: - - Text K/V is always "fresh" (updated each patch since we have full text) - - Image K/V uses stale values for non-current patches + Only caches image K/V since: + - Text K/V is always computed fresh (same for all patches) + - Only image portion needs stale/fresh cache management across patches + + This matches xDiT's approach where encoder K/V is not cached. """ def __init__( self, batch_size: int, num_heads: int, - text_seq_len: int, image_seq_len: int, head_dim: int, dtype: mx.Dtype = mx.float32, ): self.batch_size = batch_size self.num_heads = num_heads - self.text_seq_len = text_seq_len self.image_seq_len = image_seq_len self.head_dim = head_dim - self.total_seq_len = text_seq_len + image_seq_len + # Only store image K/V, not text self.key_cache = mx.zeros( - (batch_size, num_heads, self.total_seq_len, head_dim), dtype=dtype + (batch_size, num_heads, image_seq_len, head_dim), dtype=dtype ) self.value_cache = mx.zeros( - (batch_size, num_heads, self.total_seq_len, head_dim), dtype=dtype + (batch_size, num_heads, image_seq_len, head_dim), dtype=dtype ) - def update_text(self, key: mx.array, value: mx.array) -> None: - """Update text portion (always fresh, not patched). - - Args: - key: Text key tensor [batch, heads, text_seq_len, head_dim] - value: Text value tensor [batch, heads, text_seq_len, head_dim] - """ - self.key_cache[:, :, : self.text_seq_len, :] = key - self.value_cache[:, :, : self.text_seq_len, :] = value - def update_image_patch( self, patch_start: int, patch_end: int, key: mx.array, value: mx.array ) -> None: @@ -54,58 +44,84 @@ class JointPatchKVCache: key: Image patch key tensor [batch, heads, patch_len, head_dim] value: Image patch value tensor [batch, heads, patch_len, head_dim] """ - start = self.text_seq_len + patch_start - end = self.text_seq_len + patch_end - self.key_cache[:, :, start:end, :] = key - self.value_cache[:, :, start:end, :] = value + self.key_cache[:, :, patch_start:patch_end, :] = key + self.value_cache[:, :, patch_start:patch_end, :] = value - def get_full_kv(self) -> tuple[mx.array, mx.array]: - """Return full cached K/V (text + image with fresh/stale mix).""" - return self.key_cache, self.value_cache + def get_full_kv( + self, text_key: mx.array, text_value: mx.array + ) -> tuple[mx.array, mx.array]: + """Return full K/V by concatenating fresh text K/V with cached image K/V. + + Args: + text_key: Fresh text key tensor [batch, heads, text_seq_len, head_dim] + text_value: Fresh text value tensor [batch, heads, text_seq_len, head_dim] + + Returns: + Tuple of (full_key, full_value) with shape [batch, heads, text+image, head_dim] + """ + full_key = mx.concatenate([text_key, self.key_cache], axis=2) + full_value = mx.concatenate([text_value, self.value_cache], axis=2) + return full_key, full_value class PatchKVCache: - """KV cache that stores full sequence K/V with patch-level updates. + """KV cache that stores only IMAGE K/V with patch-level updates. Used for single transformer blocks where text and image tokens are concatenated. - The cache stores K/V for the full sequence [text + image] and allows - updating individual image patch slices while keeping stale values for others. + Only caches image K/V since: + - Text K/V is always computed fresh (same for all patches) + - Only image portion needs stale/fresh cache management across patches + + This matches xDiT's approach where encoder K/V is not cached. """ def __init__( self, batch_size: int, num_heads: int, - total_seq_len: int, + image_seq_len: int, head_dim: int, dtype: mx.Dtype = mx.float32, ): self.batch_size = batch_size self.num_heads = num_heads - self.total_seq_len = total_seq_len + self.image_seq_len = image_seq_len self.head_dim = head_dim + # Only store image K/V, not text self.key_cache = mx.zeros( - (batch_size, num_heads, total_seq_len, head_dim), dtype=dtype + (batch_size, num_heads, image_seq_len, head_dim), dtype=dtype ) self.value_cache = mx.zeros( - (batch_size, num_heads, total_seq_len, head_dim), dtype=dtype + (batch_size, num_heads, image_seq_len, head_dim), dtype=dtype ) - def update( + def update_image_patch( self, patch_start: int, patch_end: int, key: mx.array, value: mx.array ) -> None: - """Update cache with fresh K/V for a patch slice. + """Update cache with fresh K/V for an image patch slice. Args: - patch_start: Start token index in the full sequence - patch_end: End token index in the full sequence + patch_start: Start token index within image portion (0-indexed) + patch_end: End token index within image portion key: Fresh key tensor [batch, heads, patch_seq_len, head_dim] value: Fresh value tensor [batch, heads, patch_seq_len, head_dim] """ self.key_cache[:, :, patch_start:patch_end, :] = key self.value_cache[:, :, patch_start:patch_end, :] = value - def get_full_kv(self) -> tuple[mx.array, mx.array]: - """Return full cached K/V (mix of fresh current patch + stale others).""" - return self.key_cache, self.value_cache + def get_full_kv( + self, text_key: mx.array, text_value: mx.array + ) -> tuple[mx.array, mx.array]: + """Return full K/V by concatenating fresh text K/V with cached image K/V. + + Args: + text_key: Fresh text key tensor [batch, heads, text_seq_len, head_dim] + text_value: Fresh text value tensor [batch, heads, text_seq_len, head_dim] + + Returns: + Tuple of (full_key, full_value) with shape [batch, heads, text+image, head_dim] + """ + full_key = mx.concatenate([text_key, self.key_cache], axis=2) + full_value = mx.concatenate([text_value, self.value_cache], axis=2) + return full_key, full_value diff --git a/src/exo/worker/engines/mflux/pipefusion/patched_blocks.py b/src/exo/worker/engines/mflux/pipefusion/patched_blocks.py index 8650528e..77b41e11 100644 --- a/src/exo/worker/engines/mflux/pipefusion/patched_blocks.py +++ b/src/exo/worker/engines/mflux/pipefusion/patched_blocks.py @@ -104,11 +104,8 @@ class CachedJointAttention: xq=query, xk=patch_key, freqs_cis=patch_rope ) - # 7. Update cache with this patch's K, V (after RoPE) - kv_cache.update_text( - key=patch_key[:, :, :text_seq_len, :], - value=patch_value[:, :, :text_seq_len, :], - ) + # 7. Update cache with this patch's IMAGE K/V only (after RoPE) + # Text K/V is not cached - it's always fresh and the same for all patches kv_cache.update_image_patch( patch_start=patch_start, patch_end=patch_end, @@ -116,8 +113,12 @@ class CachedJointAttention: value=patch_value[:, :, text_seq_len:, :], ) - # 8. Get full K, V from cache (fresh current patch + stale others) - full_key, full_value = kv_cache.get_full_kv() + # 8. Get full K, V by concatenating fresh text K/V with cached image K/V + # Text K/V: fresh (just computed), Image K/V: fresh for current patch, stale for others + full_key, full_value = kv_cache.get_full_kv( + text_key=patch_key[:, :, :text_seq_len, :], + text_value=patch_value[:, :, :text_seq_len, :], + ) # 9. Compute attention: patch query attends to full K, V # Query shape: [B, H, text_seq_len + patch_len, D] @@ -209,25 +210,21 @@ class CachedSingleBlockAttention: # 3. Apply RoPE to Q and K query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=patch_rope) - # 4. Update cache with this patch's K, V (after RoPE) - # Cache stores full [text + image] sequence - # Text portion: indices 0 to text_seq_len - # Image portion: indices text_seq_len to text_seq_len + full_img_len - kv_cache.update( - patch_start=0, - patch_end=text_seq_len, - key=key[:, :, :text_seq_len, :], - value=value[:, :, :text_seq_len, :], - ) - kv_cache.update( - patch_start=text_seq_len + patch_start, - patch_end=text_seq_len + patch_end, + # 4. Update cache with this patch's IMAGE K/V only (after RoPE) + # Text K/V is not cached - it's always fresh and the same for all patches + kv_cache.update_image_patch( + patch_start=patch_start, + patch_end=patch_end, key=key[:, :, text_seq_len:, :], value=value[:, :, text_seq_len:, :], ) - # 5. Get full K, V from cache - full_key, full_value = kv_cache.get_full_kv() + # 5. Get full K, V by concatenating fresh text K/V with cached image K/V + # Text K/V: fresh (just computed), Image K/V: fresh for current patch, stale for others + full_key, full_value = kv_cache.get_full_kv( + text_key=key[:, :, :text_seq_len, :], + text_value=value[:, :, :text_seq_len, :], + ) # 6. Compute attention: patch query attends to full K, V batch_size = norm_hidden.shape[0] @@ -498,11 +495,8 @@ class CachingJointTransformerBlock: xq=query, xk=key, freqs_cis=rotary_embeddings ) - # 6. Store K, V in cache for async pipeline warmstart - self.kv_cache.update_text( - key=key[:, :, :text_seq_len, :], - value=value[:, :, :text_seq_len, :], - ) + # 6. Store only IMAGE K/V in cache for async pipeline warmstart + # Text K/V is not cached - it's always computed fresh self.kv_cache.update_image_patch( patch_start=0, patch_end=num_img_tokens, @@ -578,6 +572,7 @@ class CachingSingleTransformerBlock: hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: mx.array, + text_seq_len: int, ) -> mx.array: """Forward pass that computes attention and populates the KV cache. @@ -585,11 +580,13 @@ class CachingSingleTransformerBlock: hidden_states: Full [text + image] hidden states [B, text_len + img_len, D] text_embeddings: Time + pooled text conditioning rotary_embeddings: Full rotary embeddings for [text + full_image] + text_seq_len: Length of text portion (needed to cache only image K/V) Returns: Output hidden states after block processing """ total_seq_len = hidden_states.shape[1] + num_img_tokens = total_seq_len - text_seq_len batch_size = hidden_states.shape[0] # 0. Establish residual connection @@ -618,12 +615,13 @@ class CachingSingleTransformerBlock: xq=query, xk=key, freqs_cis=rotary_embeddings ) - # 4. Store K, V in cache for async pipeline warmstart - self.kv_cache.update( + # 4. Store only IMAGE K/V in cache for async pipeline warmstart + # Text K/V is not cached - it's always computed fresh + self.kv_cache.update_image_patch( patch_start=0, - patch_end=total_seq_len, - key=key, - value=value, + patch_end=num_img_tokens, + key=key[:, :, text_seq_len:, :], + value=value[:, :, text_seq_len:, :], ) # 5. Compute full attention (standard, not using stale cache)