From 211f4070e58f5dde289b7fa170e55a7da6310f43 Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Tue, 4 Nov 2025 12:28:36 -0800 Subject: [PATCH] fix: Lazy import mooncake-ep to fix extra gpu contexts being created (#12641) --- .../srt/layers/moe/token_dispatcher/mooncake.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py index e27128ca7..d3e7831a5 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py @@ -21,13 +21,6 @@ from sglang.srt.utils import get_int_env_var if TYPE_CHECKING: from sglang.srt.single_batch_overlap import CombineOverlapArgs -try: - from mooncake.mooncake_ep_buffer import Buffer - - use_mooncake_ep = True -except ImportError: - use_mooncake_ep = False - from enum import Enum, auto import torch @@ -86,6 +79,9 @@ class EPBuffer: if cls._buffer is not None: return cls._buffer + # Lazy import Buffer to avoid creating CUDA context at module import time + from mooncake.mooncake_ep_buffer import Buffer + cls._hidden_size = hidden_size cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank cls._num_experts = num_experts @@ -122,7 +118,9 @@ class _MooncakeEPDispatcherImpl: return_recv_hook: bool, deepep_mode: DeepEPMode, ): - if not use_mooncake_ep: + try: + from mooncake.mooncake_ep_buffer import Buffer # noqa: F401 + except ImportError: raise ImportError( "Mooncake EP is not installed. Please install Mooncake package at " "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md "