[DeepEP] Gate DeepEP MNNVL on fabric support (#28942)
This commit is contained in:
@@ -68,6 +68,15 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _is_mnnvl_fabric_supported() -> bool:
|
||||
if not is_flashinfer_available():
|
||||
return False
|
||||
|
||||
from flashinfer.comm.mnnvl import is_mnnvl_fabric_supported
|
||||
|
||||
return is_mnnvl_fabric_supported(torch.cuda.current_device())
|
||||
|
||||
|
||||
def _deepep_precompile_tp_barrier() -> None:
|
||||
# DeepEP's all-to-all operation has a much shorter timeout compared to torch.distributed,
|
||||
# so if different ranks compile at different speeds, it may quickly trigger a timeout.
|
||||
@@ -234,11 +243,11 @@ class DeepEPBuffer:
|
||||
f"Consider using --deepep-config to change the behavior."
|
||||
)
|
||||
|
||||
use_mnnvl_fabric = _is_mnnvl_fabric_supported()
|
||||
buffer_kwargs = dict(
|
||||
low_latency_mode=deepep_mode.enable_low_latency(),
|
||||
num_qps_per_rank=num_qps_per_rank,
|
||||
# TODO can be false when unneeded
|
||||
allow_mnnvl=True,
|
||||
allow_mnnvl=use_mnnvl_fabric,
|
||||
)
|
||||
# Use CU_MEM_HANDLE_TYPE_FABRIC on hardware that advertises MNNVL fabric
|
||||
# support, so cross-pod GB200/GB300 EP groups use
|
||||
@@ -251,11 +260,8 @@ class DeepEPBuffer:
|
||||
# auto-enables fabric in C++ when supported, so we skip it:
|
||||
# https://github.com/fzyzcjy/DeepEP/blob/814e508537c6ffc775d59f6f1b9ba43f3a65968c/csrc/deep_ep.cpp#L52
|
||||
is_cu12 = get_cuda_version()[0] == 12
|
||||
if not is_cu12 and is_flashinfer_available():
|
||||
from flashinfer.comm.mnnvl import is_mnnvl_fabric_supported
|
||||
|
||||
if is_mnnvl_fabric_supported(torch.cuda.current_device()):
|
||||
buffer_kwargs["use_fabric"] = True
|
||||
if not is_cu12 and use_mnnvl_fabric:
|
||||
buffer_kwargs["use_fabric"] = True
|
||||
|
||||
cls._buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes, **buffer_kwargs)
|
||||
return cls._buffer
|
||||
|
||||
Reference in New Issue
Block a user