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} "