A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 19:00 UTC
[Benchmark] Fix OOM during MoE kernel tuning for large models (#31604)
Signed-off-by: Alfred <massif0601@gmail.com>
Alfred committed
Jan 3, 2026 at 06:24 UTC
a0e9ee83c76f28527e5b8dce06df0af41cb44280
1 file changed
+55
-1
benchmarks/kernels/benchmark_moe.py
+55
-1
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import argparse
5
+import gc
6
import json
7
import os
8
import time
@@ -26,6 +27,46 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
27
28
FP8_DTYPE = current_platform.fp8_dtype()
29
30
+# Default interval for clearing Triton JIT cache during tuning
31
+# Set to 0 to disable automatic cache clearing
32
+_CACHE_CLEAR_INTERVAL_ENV = "VLLM_MOE_TUNE_CACHE_CLEAR_INTERVAL"
33
+TRITON_CACHE_CLEAR_INTERVAL = int(os.environ.get(_CACHE_CLEAR_INTERVAL_ENV, "50"))
34
+
35
+
36
+def clear_triton_cache():
37
+ """Clear Triton JIT compilation cache and Python/CUDA memory.
38
+
39
+ This helps prevent OOM during tuning with large models (many experts).
40
+ """
41
+ # Force Python garbage collection
42
+ gc.collect()
43
+
44
+ # Clear CUDA memory cache
45
+ if torch.cuda.is_available():
46
+ torch.cuda.empty_cache()
47
+
48
+ # Try to clear Triton's runtime cache
49
+ try:
50
+ import triton
51
+
52
+ if (
53
+ hasattr(triton, "runtime")
54
+ and hasattr(triton.runtime, "cache")
55
+ and hasattr(triton.runtime.cache, "clear")
56
+ ):
57
+ triton.runtime.cache.clear()
58
+ except ImportError:
59
+ # Triton not installed, skip cache clearing
60
+ pass
61
+ except AttributeError:
62
+ # Triton version doesn't have expected cache API
63
+ pass
64
+ except Exception as e:
65
+ print(f"Warning: Failed to clear Triton cache: {e}")
66
+
67
+ # Additional garbage collection after clearing caches
68
+ gc.collect()
69
+
70
71
def ensure_divisibility(numerator, denominator, text):
72
"""Ensure that numerator is divisible by the denominator."""
@@ -483,7 +524,7 @@ class BenchmarkWorker:
524
need_device_guard = True
525
526
with torch.cuda.device(self.device_id) if need_device_guard else nullcontext():
486
- for config in tqdm(search_space):
527
+ for idx, config in enumerate(tqdm(search_space)):
528
try:
529
kernel_time = benchmark_config(
530
config,
@@ -506,6 +547,19 @@ class BenchmarkWorker:
547
if kernel_time < best_time:
548
best_time = kernel_time
549
best_config = config
550
+
551
+ # Periodically clear Triton JIT cache to prevent OOM
552
+ # This is especially important for large models with many experts
553
+ if (
554
+ TRITON_CACHE_CLEAR_INTERVAL > 0
555
+ and idx > 0
556
+ and idx % TRITON_CACHE_CLEAR_INTERVAL == 0
557
+ ):
558
+ clear_triton_cache()
559
+
560
+ # Final cleanup after tuning completes
561
+ clear_triton_cache()
562
+
563
now = datetime.now()
564
print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}")
565
assert best_config is not None