add LRU eviction for mooncacke embedding cache (#25954)
Signed-off-by: Michael Qiu <qiudayu.qdy@antgroup.com> Co-authored-by: Mike_Qiu <qiudayu.qdy@antgroup.com> Co-authored-by: siyu <liusy58@linux.alibaba.com> Co-authored-by: Yuang Chen <1131578721@qq.com>
This commit is contained in:
Mike Qiu
co-authored by
Mike_Qiu
siyu
Yuang Chen
Mike_Qiu
siyu
Yuang Chen
parent
e1164a6dfc
commit
0a1fb0da86
@@ -901,8 +901,9 @@ class MMEncoder:
|
||||
mm_feature, mm_inputs, missing_indices, modality, get_feature_fn
|
||||
)
|
||||
|
||||
# Step 3: Rank 0 prefetches cache-hit embeddings from global cache.
|
||||
prefetch_status = torch.tensor([1], dtype=torch.int32)
|
||||
# Step 3: Rank 0 prefetches cache-hit embeddings and builds fallback_mask.
|
||||
fallback_mask = torch.zeros(num_items, dtype=torch.int32)
|
||||
cached_slices = []
|
||||
|
||||
if self.rank == 0:
|
||||
if hit_indices:
|
||||
@@ -919,32 +920,46 @@ class MMEncoder:
|
||||
await asyncio.sleep(0.005)
|
||||
|
||||
await asyncio.wait_for(_wait_prefetch(), timeout=60.0)
|
||||
|
||||
# Prefetch IO completed; check which items actually loaded.
|
||||
cached_slices = self.mm_global_cache.get_embeddings(hit_hashes)
|
||||
for i, idx in enumerate(hit_indices):
|
||||
if cached_slices[i] is None:
|
||||
fallback_mask[idx] = 1
|
||||
num_partial_fail = int(fallback_mask.sum().item())
|
||||
if num_partial_fail > 0:
|
||||
logger.warning(
|
||||
f"Req {req_id}: {num_partial_fail}/{len(hit_indices)} "
|
||||
f"cache-hit items failed to load (pool full), "
|
||||
f"falling back to ViT"
|
||||
)
|
||||
except (asyncio.TimeoutError, Exception) as e:
|
||||
logger.error(
|
||||
f"Prefetch failed for req {req_id}: {e}. "
|
||||
f"Falling back to ViT for {len(hit_indices)} hit items."
|
||||
)
|
||||
prefetch_status[0] = 0
|
||||
for idx in hit_indices:
|
||||
fallback_mask[idx] = 1
|
||||
|
||||
# Step 4: Broadcast prefetch result to all ranks so they stay in sync.
|
||||
# Step 4: Broadcast fallback_mask to all ranks so they stay in sync.
|
||||
if self.server_args.tp_size > 1:
|
||||
torch.distributed.broadcast(
|
||||
prefetch_status,
|
||||
fallback_mask,
|
||||
src=0,
|
||||
group=self.mm_global_cache.prefetch_tp_group,
|
||||
)
|
||||
|
||||
# Step 5: If prefetch failed, all ranks fallback to ViT for the hit mm items.
|
||||
if prefetch_status.item() == 0 and hit_indices:
|
||||
# Step 5: All ranks run ViT for items that need fallback recomputation.
|
||||
fallback_indices = [i for i in range(num_items) if fallback_mask[i].item() == 1]
|
||||
fallback_slices = None
|
||||
if fallback_indices:
|
||||
logger.info(
|
||||
f"Req {req_id}: Prefetch failed, all ranks running ViT fallback "
|
||||
f"for {len(hit_indices)} mm items."
|
||||
f"Req {req_id}: All ranks running ViT fallback "
|
||||
f"for {len(fallback_indices)} items."
|
||||
)
|
||||
fallback_slices = self._encode_missing(
|
||||
mm_feature, mm_inputs, hit_indices, modality, get_feature_fn
|
||||
mm_feature, mm_inputs, fallback_indices, modality, get_feature_fn
|
||||
)
|
||||
else:
|
||||
fallback_slices = None
|
||||
|
||||
# Step 6: Rank 0 assembles final embedding and prepares for sending.
|
||||
if self.rank == 0:
|
||||
@@ -953,25 +968,37 @@ class MMEncoder:
|
||||
for i, idx in enumerate(missing_indices):
|
||||
final_slices[idx] = new_slices[i]
|
||||
|
||||
# Fill in cache-hit embeddings (from prefetch or fallback)
|
||||
if prefetch_status.item() == 1 and hit_indices:
|
||||
cached_slices = self.mm_global_cache.get_embeddings(
|
||||
[str_mm_hashes[i] for i in hit_indices]
|
||||
)
|
||||
for i, idx in enumerate(hit_indices):
|
||||
final_slices[idx] = cached_slices[i]
|
||||
elif fallback_slices is not None:
|
||||
# Fill in successfully loaded cache-hit embeddings
|
||||
if cached_slices:
|
||||
for i, idx in enumerate(hit_indices):
|
||||
if cached_slices[i] is not None:
|
||||
final_slices[idx] = cached_slices[i]
|
||||
|
||||
# Fill in ViT fallback results for failed items
|
||||
if fallback_slices is not None:
|
||||
for i, idx in enumerate(fallback_indices):
|
||||
final_slices[idx] = fallback_slices[i]
|
||||
|
||||
mm_embedding = torch.cat(final_slices, dim=0)
|
||||
|
||||
# Release embedding cache references now that torch.cat has
|
||||
# copied the data into a new tensor. This allows the cache
|
||||
# entries to be evicted under memory pressure.
|
||||
if cached_slices:
|
||||
loaded_hashes = [
|
||||
str_mm_hashes[idx]
|
||||
for idx in hit_indices
|
||||
if fallback_mask[idx].item() == 0
|
||||
]
|
||||
if loaded_hashes:
|
||||
self.mm_global_cache.release_embeddings(loaded_hashes)
|
||||
|
||||
# Background insert: store newly computed embeddings into global cache.
|
||||
# Includes both original misses and fallback-recomputed hits.
|
||||
all_new_hashes = [str_mm_hashes[i] for i in missing_indices]
|
||||
all_new_slices = list(new_slices)
|
||||
if fallback_slices is not None:
|
||||
all_new_hashes += [str_mm_hashes[i] for i in hit_indices]
|
||||
all_new_hashes += [str_mm_hashes[i] for i in fallback_indices]
|
||||
all_new_slices += list(fallback_slices)
|
||||
|
||||
if all_new_hashes:
|
||||
|
||||
@@ -24,7 +24,8 @@ class ContiguousMemoryAllocator:
|
||||
self.total_size = total_size_bytes
|
||||
# List of (offset, size) for free blocks
|
||||
self.free_blocks = [(0, total_size_bytes)]
|
||||
self.allocated_map = {} # {handle: (offset, size)}
|
||||
self.allocated_map = {} # {offset: size_bytes}
|
||||
self.allocated_size = 0 # Running counter for O(1) get_allocated_size
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def allocate(self, size_bytes: int) -> Optional[int]:
|
||||
@@ -38,11 +39,18 @@ class ContiguousMemoryAllocator:
|
||||
self.free_blocks[i] = (offset + size_bytes, remaining_size)
|
||||
else:
|
||||
self.free_blocks.pop(i)
|
||||
self.allocated_map[offset] = size_bytes
|
||||
self.allocated_size += size_bytes
|
||||
return offset
|
||||
return None
|
||||
|
||||
def free(self, offset: int, size_bytes: int):
|
||||
with self.lock:
|
||||
# Remove from allocated map and update counter
|
||||
if offset in self.allocated_map:
|
||||
self.allocated_size -= self.allocated_map[offset]
|
||||
del self.allocated_map[offset]
|
||||
|
||||
# Return block and merge adjacent free blocks
|
||||
self.free_blocks.append((offset, size_bytes))
|
||||
self.free_blocks.sort()
|
||||
@@ -61,6 +69,16 @@ class ContiguousMemoryAllocator:
|
||||
merged.append((curr_offset, curr_size))
|
||||
self.free_blocks = merged
|
||||
|
||||
def get_allocated_size(self) -> int:
|
||||
"""Return total allocated bytes. O(1) operation."""
|
||||
with self.lock:
|
||||
return self.allocated_size
|
||||
|
||||
def get_free_size(self) -> int:
|
||||
"""Return total free bytes."""
|
||||
with self.lock:
|
||||
return sum(block_size for _, block_size in self.free_blocks)
|
||||
|
||||
|
||||
class EmbeddingPrefetchOperation:
|
||||
"""Groups all missing images of a request for a single batch GET."""
|
||||
@@ -98,12 +116,17 @@ class EmbeddingCacheController:
|
||||
hidden_dims: dict = None,
|
||||
tp_group=None,
|
||||
all_rank_get=False,
|
||||
enable_eviction: bool = True,
|
||||
max_eviction_batch: int = 100,
|
||||
):
|
||||
self.tp_world_size = tp_size
|
||||
self.tp_group = tp_group
|
||||
self.tp_rank = tp_rank
|
||||
self.all_rank_get = all_rank_get
|
||||
self.hidden_dims = hidden_dims or {}
|
||||
self.element_size = torch.float32.itemsize
|
||||
self.enable_eviction = enable_eviction
|
||||
self.max_eviction_batch = max_eviction_batch
|
||||
|
||||
# 1. Mooncake Backend & Pinned Buffer
|
||||
self.mooncake_store = MooncakeEmbeddingStore()
|
||||
@@ -115,10 +138,31 @@ class EmbeddingCacheController:
|
||||
|
||||
# 2. Variable Size Memory Management
|
||||
self.allocator = ContiguousMemoryAllocator(self.total_pool_size_bytes)
|
||||
# {hash: (offset, num_tokens, embedding_dim, size_bytes)}
|
||||
# {hash: (offset, num_tokens, embedding_dim, size_bytes, last_access_time)}
|
||||
self.hash_to_metadata = {}
|
||||
|
||||
# 3. Task Tracking
|
||||
# 3. LRU Tracking
|
||||
# OrderedDict maintains insertion order, used as LRU cache
|
||||
# hash -> access_time
|
||||
self.access_order = {}
|
||||
self.access_lock = threading.Lock()
|
||||
|
||||
# 4. RDMA / read reference counting
|
||||
# {hash: ref_count} — entries with ref_count > 0 cannot be evicted.
|
||||
# Incremented when an RDMA transfer (GET/PUT) is in flight or when
|
||||
# get_embeddings() returns a view into cpu_pool. Decremented after
|
||||
# the RDMA completes or the caller releases the view.
|
||||
self.ref_counts = {}
|
||||
|
||||
# 5. Statistics
|
||||
self.stats = {
|
||||
"total_allocated": 0,
|
||||
"total_evicted": 0,
|
||||
"eviction_count": 0,
|
||||
"allocation_failures": 0,
|
||||
}
|
||||
|
||||
# 6. Task Tracking
|
||||
self.ongoing_prefetch = {} # {req_id: EmbeddingPrefetchOperation}
|
||||
self.prefetch_queue = Queue()
|
||||
self.insert_queue = Queue()
|
||||
@@ -142,6 +186,156 @@ class EmbeddingCacheController:
|
||||
else:
|
||||
self.prefetch_tp_group = None
|
||||
|
||||
def _update_access_time(self, image_hash: str):
|
||||
"""Update LRU access time for a hash."""
|
||||
with self.access_lock:
|
||||
# Move to end (most recently used)
|
||||
if image_hash in self.access_order:
|
||||
del self.access_order[image_hash]
|
||||
self.access_order[image_hash] = time.time()
|
||||
|
||||
def _protect_hash(self, image_hash: str):
|
||||
"""Increment ref count to prevent eviction during RDMA or active read.
|
||||
|
||||
NOTE: Caller must hold self.lock.
|
||||
"""
|
||||
self.ref_counts[image_hash] = self.ref_counts.get(image_hash, 0) + 1
|
||||
|
||||
def _release_hash(self, image_hash: str):
|
||||
"""Decrement ref count after RDMA completes or caller releases a view.
|
||||
|
||||
NOTE: Caller must hold self.lock.
|
||||
"""
|
||||
if image_hash in self.ref_counts:
|
||||
self.ref_counts[image_hash] -= 1
|
||||
if self.ref_counts[image_hash] <= 0:
|
||||
del self.ref_counts[image_hash]
|
||||
|
||||
def _select_eviction_candidates(self, required_bytes: int) -> List[str]:
|
||||
"""Select LRU candidates to free up at least required_bytes.
|
||||
|
||||
NOTE: Caller must hold self.lock before calling this method.
|
||||
"""
|
||||
candidates = []
|
||||
freed_bytes = 0
|
||||
|
||||
with self.access_lock:
|
||||
# Sort by access time (oldest first)
|
||||
# Python dicts are insertion-ordered; the first keys are the oldest.
|
||||
sorted_hashes = list(self.access_order.items())
|
||||
|
||||
for image_hash, _ in sorted_hashes:
|
||||
if image_hash not in self.hash_to_metadata:
|
||||
with self.access_lock:
|
||||
self.access_order.pop(image_hash, None)
|
||||
continue
|
||||
if self.ref_counts.get(image_hash, 0) > 0:
|
||||
continue
|
||||
metadata = self.hash_to_metadata[image_hash]
|
||||
size_bytes = metadata[3] if len(metadata) > 3 else 0
|
||||
candidates.append(image_hash)
|
||||
freed_bytes += size_bytes
|
||||
|
||||
if freed_bytes >= required_bytes:
|
||||
break
|
||||
|
||||
if len(candidates) >= self.max_eviction_batch:
|
||||
break
|
||||
|
||||
return candidates
|
||||
|
||||
def _evict_hashes(self, hashes_to_evict: List[str]) -> int:
|
||||
"""Evict specified hashes and free their memory. Returns freed bytes.
|
||||
|
||||
NOTE: Caller must hold self.lock before calling this method.
|
||||
"""
|
||||
total_freed = 0
|
||||
|
||||
# NOTE: self.lock should be held by the caller (e.g., insert_batch,
|
||||
# prefetch). Do NOT acquire it here to avoid reentrant deadlock.
|
||||
for image_hash in hashes_to_evict:
|
||||
if image_hash not in self.hash_to_metadata:
|
||||
continue
|
||||
|
||||
# Safety check: skip entries with in-flight RDMA or active reads
|
||||
if self.ref_counts.get(image_hash, 0) > 0:
|
||||
logger.warning(
|
||||
f"[Rank {self.tp_rank}] Skipping eviction of {image_hash}: "
|
||||
f"ref_count={self.ref_counts[image_hash]} (in-flight RDMA or active read)"
|
||||
)
|
||||
continue
|
||||
|
||||
offset, num_tokens, dim, size_bytes = self.hash_to_metadata[image_hash][:4]
|
||||
|
||||
# Free memory in allocator
|
||||
self.allocator.free(offset, size_bytes)
|
||||
|
||||
# Remove from metadata and ref counts
|
||||
del self.hash_to_metadata[image_hash]
|
||||
self.ref_counts.pop(image_hash, None)
|
||||
|
||||
# Remove from access order
|
||||
with self.access_lock:
|
||||
self.access_order.pop(image_hash, None)
|
||||
|
||||
total_freed += size_bytes
|
||||
self.stats["total_evicted"] += size_bytes
|
||||
|
||||
if total_freed > 0:
|
||||
self.stats["eviction_count"] += 1
|
||||
|
||||
if total_freed > 0:
|
||||
logger.info(
|
||||
f"[Rank {self.tp_rank}] Evicted {len(hashes_to_evict)} embeddings, "
|
||||
f"freed {total_freed / 1024**2:.2f} MB"
|
||||
)
|
||||
|
||||
return total_freed
|
||||
|
||||
def _allocate_with_eviction(self, size_bytes: int) -> Optional[int]:
|
||||
"""Try to allocate memory, evicting old entries if necessary."""
|
||||
# First try direct allocation
|
||||
offset = self.allocator.allocate(size_bytes)
|
||||
if offset is not None:
|
||||
self.stats["total_allocated"] += size_bytes
|
||||
return offset
|
||||
|
||||
# If failed and eviction is enabled, try eviction
|
||||
if not self.enable_eviction:
|
||||
self.stats["allocation_failures"] += 1
|
||||
return None
|
||||
|
||||
# Select candidates to evict
|
||||
candidates = self._select_eviction_candidates(size_bytes)
|
||||
if not candidates:
|
||||
n_protected = sum(1 for v in self.ref_counts.values() if v > 0)
|
||||
logger.warning(
|
||||
f"[Rank {self.tp_rank}] Cannot allocate {size_bytes / 1024**2:.2f} MB: "
|
||||
f"pool full ({self.allocator.get_allocated_size() / 1024**2:.1f}/"
|
||||
f"{self.total_pool_size_bytes / 1024**2:.1f} MB used), "
|
||||
f"no evictable candidates "
|
||||
f"({len(self.hash_to_metadata)} entries, {n_protected} protected)"
|
||||
)
|
||||
self.stats["allocation_failures"] += 1
|
||||
return None
|
||||
|
||||
# Evict and try again
|
||||
freed = self._evict_hashes(candidates)
|
||||
if freed < size_bytes:
|
||||
logger.warning(
|
||||
f"[Rank {self.tp_rank}] Could not free enough memory: "
|
||||
f"needed {size_bytes / 1024**2:.2f} MB, freed {freed / 1024**2:.2f} MB"
|
||||
)
|
||||
|
||||
# Try allocation again after eviction
|
||||
offset = self.allocator.allocate(size_bytes)
|
||||
if offset is not None:
|
||||
self.stats["total_allocated"] += size_bytes
|
||||
else:
|
||||
self.stats["allocation_failures"] += 1
|
||||
|
||||
return offset
|
||||
|
||||
def prefetch(
|
||||
self,
|
||||
req_id: str,
|
||||
@@ -161,17 +355,25 @@ class EmbeddingCacheController:
|
||||
with self.lock:
|
||||
for h, num_tokens in zip(image_hashes, expected_tokens):
|
||||
if h in self.hash_to_metadata:
|
||||
# Update access time for LRU
|
||||
self._update_access_time(h)
|
||||
logger.debug(
|
||||
f"Req {req_id}: Hash already in local metadata, skipping prefetch."
|
||||
f"Req {req_id}: Hash already in local metadata, skipping prefetch."
|
||||
)
|
||||
continue
|
||||
|
||||
size_bytes = num_tokens * dim * self.element_size
|
||||
offset = self.allocator.allocate(size_bytes)
|
||||
offset = self._allocate_with_eviction(size_bytes)
|
||||
if offset is None:
|
||||
logger.warning(
|
||||
f"Req {req_id}: Failed to allocate {size_bytes / 1024**2:.2f} MB "
|
||||
f"for prefetch, skipping this image."
|
||||
)
|
||||
continue
|
||||
|
||||
self.hash_to_metadata[h] = (offset, num_tokens, dim, size_bytes)
|
||||
self._update_access_time(h)
|
||||
self._protect_hash(h)
|
||||
keys.append(h)
|
||||
ptrs.append(self.cpu_pool.data_ptr() + offset)
|
||||
sizes.append(size_bytes)
|
||||
@@ -204,6 +406,9 @@ class EmbeddingCacheController:
|
||||
with self.lock:
|
||||
for h, tensor in zip(image_hashes, embedding_tensors):
|
||||
if h in self.hash_to_metadata:
|
||||
# Update access time for existing entry
|
||||
self._update_access_time(h)
|
||||
self._protect_hash(h)
|
||||
# Local cache hit: ensure Mooncake has it
|
||||
offset, num_tokens, dim, size_bytes = self.hash_to_metadata[h][:4]
|
||||
|
||||
@@ -218,8 +423,12 @@ class EmbeddingCacheController:
|
||||
# Local cache miss: allocate and copy
|
||||
num_tokens, dim = tensor.shape[0], tensor.shape[1]
|
||||
size_bytes = num_tokens * dim * self.element_size
|
||||
offset = self.allocator.allocate(size_bytes)
|
||||
offset = self._allocate_with_eviction(size_bytes)
|
||||
if offset is None:
|
||||
logger.warning(
|
||||
f"Failed to allocate {size_bytes / 1024**2:.2f} MB for insert, "
|
||||
f"skipping this embedding."
|
||||
)
|
||||
skipped_count += 1
|
||||
continue
|
||||
|
||||
@@ -231,6 +440,8 @@ class EmbeddingCacheController:
|
||||
)
|
||||
target_view.copy_(tensor.cpu())
|
||||
self.hash_to_metadata[h] = (offset, num_tokens, dim, size_bytes)
|
||||
self._update_access_time(h)
|
||||
self._protect_hash(h)
|
||||
|
||||
keys.append(h)
|
||||
ptrs.append(self.cpu_pool.data_ptr() + offset)
|
||||
@@ -258,6 +469,10 @@ class EmbeddingCacheController:
|
||||
f"Mooncake GET Finished: Req {op.req_id}, Successfully fetched {success_count}/{len(op.keys)} images."
|
||||
)
|
||||
op.mark_done(all(results))
|
||||
# Release ref counts now that RDMA GET is complete
|
||||
with self.lock:
|
||||
for h in op.keys:
|
||||
self._release_hash(h)
|
||||
self.prefetch_queue.task_done()
|
||||
processed_any = True
|
||||
except Empty:
|
||||
@@ -269,6 +484,10 @@ class EmbeddingCacheController:
|
||||
logger.info(
|
||||
f"Mooncake PUT Finished: Successfully stored {len(op.keys)} keys in cluster."
|
||||
)
|
||||
# Release ref counts now that RDMA PUT is complete
|
||||
with self.lock:
|
||||
for h in op.keys:
|
||||
self._release_hash(h)
|
||||
self.insert_queue.task_done()
|
||||
processed_any = True
|
||||
except Empty:
|
||||
@@ -306,11 +525,24 @@ class EmbeddingCacheController:
|
||||
return False
|
||||
|
||||
def get_embeddings(self, image_hashes: List[str]) -> List[torch.Tensor]:
|
||||
"""Final reconstruction for model input."""
|
||||
"""Final reconstruction for model input.
|
||||
|
||||
Returns views into the pinned cpu_pool. Callers MUST call
|
||||
release_embeddings() once they no longer need the returned
|
||||
tensors (e.g. after .to(device) or torch.cat) so that the
|
||||
entries can be evicted.
|
||||
"""
|
||||
with self.lock:
|
||||
tensors = []
|
||||
for h in image_hashes:
|
||||
offset, num_tokens, dim, size_bytes = self.hash_to_metadata[h]
|
||||
if h not in self.hash_to_metadata:
|
||||
logger.warning(f"Hash {h} not found in local cache")
|
||||
tensors.append(None)
|
||||
continue
|
||||
# Update access time for LRU
|
||||
self._update_access_time(h)
|
||||
self._protect_hash(h)
|
||||
offset, num_tokens, dim, size_bytes = self.hash_to_metadata[h][:4]
|
||||
tensors.append(
|
||||
self.cpu_pool[offset : offset + size_bytes]
|
||||
.view(torch.float32)
|
||||
@@ -318,6 +550,29 @@ class EmbeddingCacheController:
|
||||
)
|
||||
return tensors
|
||||
|
||||
def release_embeddings(self, image_hashes: List[str]):
|
||||
"""Release reference counts on embeddings after the caller is done.
|
||||
|
||||
Must be called once for every successful get_embeddings() call,
|
||||
after the caller no longer needs the returned tensor views
|
||||
(e.g. after .to(device) or torch.cat has copied the data).
|
||||
"""
|
||||
with self.lock:
|
||||
for h in image_hashes:
|
||||
self._release_hash(h)
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
"""Return cache statistics."""
|
||||
with self.lock:
|
||||
return {
|
||||
**self.stats,
|
||||
"num_cached": len(self.hash_to_metadata),
|
||||
"num_protected": sum(1 for v in self.ref_counts.values() if v > 0),
|
||||
"allocated_mb": self.allocator.get_allocated_size() / 1024**2,
|
||||
"free_mb": self.allocator.get_free_size() / 1024**2,
|
||||
"total_mb": self.total_pool_size_bytes / 1024**2,
|
||||
}
|
||||
|
||||
async def batch_is_exist(self, image_hashes: List[str]) -> List[bool]:
|
||||
with self.lock:
|
||||
local_results = [h in self.hash_to_metadata for h in image_hashes]
|
||||
|
||||
Reference in New Issue
Block a user