fix: LRU(2) eviction for prepack cache — prevents OOM across 61 layers
Cache was growing unbounded — 61 MoE layers × 2 = 122 prepacked SFB tensors permanently in GPU memory (~1.75 GiB each). With sequential layer execution, only 2 entries are needed at a time (current L1 + L2). Added LRU eviction to keep max 2 entries.
This commit is contained in:
@@ -106,7 +106,11 @@ def _prepack_weight_sf(weight_sf, N, K, tag):
|
||||
)
|
||||
if not hasattr(_prepack_weight_sf, '_cache'):
|
||||
_prepack_weight_sf._cache = {}
|
||||
_prepack_weight_sf._cache_order = [] # LRU order
|
||||
if cache_key in _prepack_weight_sf._cache:
|
||||
# Move to end (most recently used)
|
||||
_prepack_weight_sf._cache_order.remove(cache_key)
|
||||
_prepack_weight_sf._cache_order.append(cache_key)
|
||||
return _prepack_weight_sf._cache[cache_key]
|
||||
|
||||
assert weight_sf.dtype == torch.float8_e4m3fn, weight_sf.dtype
|
||||
@@ -127,6 +131,12 @@ def _prepack_weight_sf(weight_sf, N, K, tag):
|
||||
|
||||
packed = torch.stack(packed, dim=0).contiguous()
|
||||
_prepack_weight_sf._cache[cache_key] = packed
|
||||
_prepack_weight_sf._cache_order.append(cache_key)
|
||||
|
||||
# Evict oldest entries — keep only 2 (current layer's L1 + L2)
|
||||
while len(_prepack_weight_sf._cache) > 2:
|
||||
oldest = _prepack_weight_sf._cache_order.pop(0)
|
||||
del _prepack_weight_sf._cache[oldest]
|
||||
|
||||
if MEGA_MOE_DEBUG:
|
||||
print(f"[PREPACK] {tag}: E={E} N={N} K={K} packed_shape={packed.shape} "
|
||||
|
||||
Reference in New Issue
Block a user