From 90313f3a9275ee629221b80c74d8e074984d3cb0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 10:31:57 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20LRU(2)=20eviction=20for=20prepack=20cach?= =?UTF-8?q?e=20=E2=80=94=20prevents=20OOM=20across=2061=20layers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 95cb1dee..c0cd73a8 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -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} "