chenyu520
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Misc] Lazy import `ep` in `mooncake_ep_buffer.py` (#1014) 

* Lazy import `ep` in `mooncake_ep_buffer.py`

* Fix
This commit is contained in:
Xun Sun 2025-11-04 11:49:15 +08:00 committed by GitHub
parent 7c22adb0d2
commit 1a4c01512d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mooncake-wheel/mooncake/mooncake_ep_buffer.py
View File

@ -2,9 +2,6 @@ import torch
import torch.distributed as dist
from typing import Any, Callable, List, Tuple, Optional, Union
# noinspection PyUnresolvedReferences
from mooncake import ep
class EventOverlap:
"""
@ -15,7 +12,7 @@ class EventOverlap:
extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
"""
def __init__(self, event: Optional[ep.EventHandle] = None,
def __init__(self, event: Optional["ep.EventHandle"] = None,
extra_tensors: Optional[Tuple[torch.Tensor, ...]] = None) -> None:
"""
Initialize the class.
@ -63,6 +60,7 @@ class EventOverlap:
class Buffer:
def __init__(self, group: dist.ProcessGroup, num_ep_buffer_bytes: int = 0):
from mooncake import ep
# Initialize the CPP runtime
self.rank = group.rank()
self.group_size = group.size()
@ -120,6 +118,7 @@ class Buffer:
@staticmethod
def get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int) -> int:
from mooncake import ep
return ep.get_ep_buffer_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)
# noinspection PyTypeChecker