From 5dc18df494dad42235f023b0e4be06cc5756f750 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 10:33:53 +0000 Subject: [PATCH] feat: MEGA_MOE_PREPACK_CACHE_MAX env var (default 2) with CUDA graph warning Added big comment block explaining the cache sizing rationale and the CUDA graph trap: default of 2 works for sequential layer execution but will cause use-after-free if CUDA graphs capture multiple layers. Set MEGA_MOE_PREPACK_CACHE_MAX to cover all captured layers in that case. --- src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index c0cd73a8..e415d33b 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -104,6 +104,23 @@ def _prepack_weight_sf(weight_sf, N, K, tag): N, K, ) + # ───────────────────────────────────────────────────────────────────── + # PREPACK CACHE — LRU with configurable max size + # + # Each prepacked SFB tensor is ~1.75 GiB per rank. With 61 MoE layers + # × 2 (L1+L2), an unbounded cache would consume ~214 GiB — well beyond + # B200 capacity. Since vLLM calls layers sequentially, only 2 entries + # are needed at a time (current layer's L1 + L2). + # + # WARNING: If you enable CUDA graphs in the future, multiple layers may + # be captured in a single graph replay, and the prepacked tensors must + # remain alive for the graph's lifetime. In that case, increase + # MEGA_MOE_PREPACK_CACHE_MAX to cover all captured layers, or switch + # to a persistent pre-allocation scheme. The default of 2 will cause + # use-after-free on evicted entries if CUDA graphs span >1 layer. + # ───────────────────────────────────────────────────────────────────── + _max_cache = int(os.environ.get('MEGA_MOE_PREPACK_CACHE_MAX', '2')) + if not hasattr(_prepack_weight_sf, '_cache'): _prepack_weight_sf._cache = {} _prepack_weight_sf._cache_order = [] # LRU order @@ -133,8 +150,8 @@ def _prepack_weight_sf(weight_sf, N, K, tag): _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: + # Evict oldest entries — keep only _max_cache entries + while len(_prepack_weight_sf._cache) > _max_cache: oldest = _prepack_weight_sf._cache_order.pop(0) del _prepack_weight_sf._cache[oldest]