electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove redundant text kv cache computation

Browse Source
This commit is contained in:
ciaranbor
2026-01-06 10:51:20 +00:00
parent 7b8382be10
commit db24f052d7
3 changed files with 91 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py
+5 -6
View File
@@ -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,
)
src/exo/worker/engines/mflux/pipefusion/kv_cache.py
+56 -40
View File
@@ -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
src/exo/worker/engines/mflux/pipefusion/patched_blocks.py
+30 -32
View File
@@ -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)