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

@@ -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]