diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index f8f20091b..e4da1d058 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -59,7 +59,7 @@ jobs: contains(github.event.pull_request.labels.*.name, 'ci:gpu') runs-on: turbo-torch-ut-v2602 env: - UCCL_COMMIT: 69a21407d261a8ff0f01f83bb0f5b35da059876d # [rdma] Make CQ Ex flags configurable for NIC compatibility. (#689) + UCCL_COMMIT: 12fdebd8d9a40f12364239f51d74fc773c41a4a7 # [EP] LL clean & Zero-token LL rank issues (#993) UCCL_PY_VER: 3.12 steps: - name: Print job info diff --git a/README.md b/README.md index 66e123e55..8a2cbf8f8 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,8 @@ Note: JAX support is under active development. Optim support is planned but not - Python >= 3.10 - PyTorch >= 2.6.0 (with ROCm support) - [AITER](https://github.com/ROCm/aiter) (required for some operators, e.g. FlashAttention / FP8): `pip3 install "amd-aiter @ git+https://github.com/ROCm/aiter.git@v0.1.14.post1"` +- [MORI](https://github.com/ROCm/mori) (optional, required for the **MORI** expert-parallel backend): `BUILD_UMBP=OFF pip3 install "amd_mori @ git+https://github.com/ROCm/mori.git@v1.2.0"` +- [UCCL](https://github.com/uccl-project/uccl) (optional, required for the **UCCL** expert-parallel backend): `pip3 install "uccl @ git+https://github.com/uccl-project/uccl.git@34bdf4f49c8e33d618b983afc34fc5bec8686cfa"` - rocSHMEM (optional, required for **experimental DeepEP**). Please refer to our [DeepEP Installation Guide](primus_turbo/pytorch/deep_ep/README.md) for instructions. #### Hardware diff --git a/benchmark/ops/training/deep_ep/test_intranode.py b/benchmark/ops/training/deep_ep/test_intranode.py index 515c364b0..986480563 100644 --- a/benchmark/ops/training/deep_ep/test_intranode.py +++ b/benchmark/ops/training/deep_ep/test_intranode.py @@ -17,15 +17,28 @@ import torch import torch.distributed as dist -from .utils import ( - bench, - calc_diff, - get_deep_ep_backend, - init_dist, - inplace_unique, - per_token_cast_back, - per_token_cast_to_fp8, -) +try: + # Package-relative import (run as a module). + from .utils import ( + bench, + calc_diff, + get_deep_ep_backend, + init_dist, + inplace_unique, + per_token_cast_back, + per_token_cast_to_fp8, + ) +except ImportError: + # Script-style import (run directly with this dir on sys.path). + from utils import ( + bench, + calc_diff, + get_deep_ep_backend, + init_dist, + inplace_unique, + per_token_cast_back, + per_token_cast_to_fp8, + ) # noinspection PyShadowingNames diff --git a/csrc/include/primus_turbo/deep_ep/configs.h b/csrc/include/primus_turbo/deep_ep/configs.h index 3fc9ea2fd..e89318e3f 100644 --- a/csrc/include/primus_turbo/deep_ep/configs.h +++ b/csrc/include/primus_turbo/deep_ep/configs.h @@ -91,11 +91,11 @@ inline static bool is_enable_cheap_fence() { // When set to 1, EP dispatch/combine kernels run on the caller's current // CUDA stream instead of the Buffer's internal comm stream. -// Default: 0. +// Default: 1. inline static bool is_ep_force_current_stream() { static uint32_t val = []() { const char *v = std::getenv("PRIMUS_TURBO_EP_FORCE_CURRENT_STREAM"); - return v ? static_cast(std::stoi(v)) : 0; + return v ? static_cast(std::stoi(v)) : 1; }(); return val != 0; } diff --git a/csrc/kernels/deep_ep/buffer.cuh b/csrc/kernels/deep_ep/buffer.cuh index c06217c46..45f401888 100644 --- a/csrc/kernels/deep_ep/buffer.cuh +++ b/csrc/kernels/deep_ep/buffer.cuh @@ -17,14 +17,14 @@ private: uint8_t *ptr; public: - int total_bytes; + int64_t total_bytes; __device__ __forceinline__ Buffer() : ptr(nullptr), total_bytes(0) {} - __device__ __forceinline__ Buffer(void *&gbl_ptr, int num_elems, int offset = 0) { - total_bytes = num_elems * sizeof(dtype_t); - ptr = reinterpret_cast(gbl_ptr) + offset * sizeof(dtype_t); - gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; + __device__ __forceinline__ Buffer(void *&gbl_ptr, int64_t num_elems, int64_t offset = 0) { + total_bytes = num_elems * static_cast(sizeof(dtype_t)); + ptr = reinterpret_cast(gbl_ptr) + offset * static_cast(sizeof(dtype_t)); + gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; } __device__ __forceinline__ Buffer advance_also(void *&gbl_ptr) { @@ -34,36 +34,36 @@ public: __device__ __forceinline__ dtype_t *buffer() { return reinterpret_cast(ptr); } - __device__ __forceinline__ dtype_t &operator[](int idx) { return buffer()[idx]; } + __device__ __forceinline__ dtype_t &operator[](int64_t idx) { return buffer()[idx]; } }; template struct AsymBuffer { private: uint8_t *ptrs[kNumRanks]; - int num_bytes; + int64_t num_bytes; public: - int total_bytes; + int64_t total_bytes; - __device__ __forceinline__ AsymBuffer(void *&gbl_ptr, int num_elems, int num_ranks, - int sm_id = 0, int num_sms = 1, int offset = 0) { + __device__ __forceinline__ AsymBuffer(void *&gbl_ptr, int64_t num_elems, int num_ranks, + int sm_id = 0, int num_sms = 1, int64_t offset = 0) { PRIMUS_TURBO_STATIC_CHECK(kNumRanks == 1, ""); - num_bytes = num_elems * sizeof(dtype_t); + num_bytes = num_elems * static_cast(sizeof(dtype_t)); - int per_channel_bytes = num_bytes * num_ranks; - total_bytes = per_channel_bytes * num_sms; + int64_t per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms; ptrs[0] = reinterpret_cast(gbl_ptr) + per_channel_bytes * sm_id + num_bytes * offset; gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; } - __device__ __forceinline__ AsymBuffer(void **gbl_ptrs, int num_elems, int num_ranks, - int sm_id = 0, int num_sms = 1, int offset = 0) { + __device__ __forceinline__ AsymBuffer(void **gbl_ptrs, int64_t num_elems, int num_ranks, + int sm_id = 0, int num_sms = 1, int64_t offset = 0) { PRIMUS_TURBO_STATIC_CHECK(kNumRanks > 1, ""); - num_bytes = num_elems * sizeof(dtype_t); + num_bytes = num_elems * static_cast(sizeof(dtype_t)); - int per_channel_bytes = num_bytes * num_ranks; - total_bytes = per_channel_bytes * num_sms; + int64_t per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms; for (int i = 0; i < kNumRanks; ++i) { ptrs[i] = reinterpret_cast(gbl_ptrs[i]) + per_channel_bytes * sm_id + num_bytes * offset; @@ -71,10 +71,10 @@ public: } } - __device__ __forceinline__ void advance(int shift) { + __device__ __forceinline__ void advance(int64_t shift) { #pragma unroll for (int i = 0; i < kNumRanks; ++i) - ptrs[i] = ptrs[i] + shift * sizeof(dtype_t); + ptrs[i] = ptrs[i] + shift * static_cast(sizeof(dtype_t)); } __device__ __forceinline__ AsymBuffer advance_also(void *&gbl_ptr) { @@ -89,13 +89,13 @@ public: return *this; } - __device__ __forceinline__ dtype_t *buffer(int idx = 0) { + __device__ __forceinline__ dtype_t *buffer(int64_t idx = 0) { PRIMUS_TURBO_STATIC_CHECK(kNumRanks == 1, "`buffer` is only available for single rank case"); return reinterpret_cast(ptrs[0] + num_bytes * idx); } - __device__ __forceinline__ dtype_t *buffer_by(int rank_idx, int idx = 0) { + __device__ __forceinline__ dtype_t *buffer_by(int rank_idx, int64_t idx = 0) { PRIMUS_TURBO_STATIC_CHECK(kNumRanks > 1, "`buffer` is only available for single rank case"); return reinterpret_cast(ptrs[rank_idx] + num_bytes * idx); } @@ -106,35 +106,35 @@ private: // NOTES: for non-decoupled case, `recv_ptr` is not used uint8_t *send_ptr; uint8_t *recv_ptr; - int num_bytes; + int64_t num_bytes; public: - int total_bytes; + int64_t total_bytes; - __device__ __forceinline__ SymBuffer(void *&gbl_ptr, int num_elems, int num_ranks, + __device__ __forceinline__ SymBuffer(void *&gbl_ptr, int64_t num_elems, int num_ranks, int sm_id = 0, int num_sms = 1) { - num_bytes = num_elems * sizeof(dtype_t); + num_bytes = num_elems * static_cast(sizeof(dtype_t)); - int per_channel_bytes = num_bytes * num_ranks; - total_bytes = per_channel_bytes * num_sms * (static_cast(kDecoupled) + 1); - send_ptr = reinterpret_cast(gbl_ptr) + per_channel_bytes * sm_id; - recv_ptr = reinterpret_cast(gbl_ptr) + per_channel_bytes * (sm_id + num_sms); - gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; + int64_t per_channel_bytes = num_bytes * num_ranks; + total_bytes = per_channel_bytes * num_sms * (static_cast(kDecoupled) + 1); + send_ptr = reinterpret_cast(gbl_ptr) + per_channel_bytes * sm_id; + recv_ptr = reinterpret_cast(gbl_ptr) + per_channel_bytes * (sm_id + num_sms); + gbl_ptr = reinterpret_cast(gbl_ptr) + total_bytes; } - __device__ __forceinline__ dtype_t *send_buffer(int idx = 0) { + __device__ __forceinline__ dtype_t *send_buffer(int64_t idx = 0) { PRIMUS_TURBO_STATIC_CHECK(kDecoupled, "`send_buffer` is only available for non-decoupled case"); return reinterpret_cast(send_ptr + num_bytes * idx); } - __device__ __forceinline__ dtype_t *recv_buffer(int idx = 0) { + __device__ __forceinline__ dtype_t *recv_buffer(int64_t idx = 0) { PRIMUS_TURBO_STATIC_CHECK(kDecoupled, "`recv_buffer` is only available for non-decoupled case"); return reinterpret_cast(recv_ptr + num_bytes * idx); } - __device__ __forceinline__ dtype_t *buffer(int idx = 0) { + __device__ __forceinline__ dtype_t *buffer(int64_t idx = 0) { PRIMUS_TURBO_STATIC_CHECK(not kDecoupled, "`buffer` is only available for decoupled case"); return reinterpret_cast(send_ptr + num_bytes * idx); } diff --git a/docs/examples.md b/docs/examples.md index 8835f992a..afcc884d3 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -597,5 +597,5 @@ You can also control backend selection and AutoTune via environment variables: export PRIMUS_TURBO_AUTO_TUNE=1 export PRIMUS_TURBO_GEMM_BACKEND=HIPBLASLT export PRIMUS_TURBO_GROUPED_GEMM_BACKEND=CK -export PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=DEEP_EP +export PRIMUS_TURBO_EP_BACKEND=DEEP_EP ``` diff --git a/primus_turbo/common/constants.py b/primus_turbo/common/constants.py index 8b87e5218..46214b10f 100644 --- a/primus_turbo/common/constants.py +++ b/primus_turbo/common/constants.py @@ -24,9 +24,9 @@ # Default: None (auto-select) ENV_GROUPED_GEMM_BACKEND = "PRIMUS_TURBO_GROUPED_GEMM_BACKEND" -# MoE dispatch/combine EP backend (TURBO, DEEP_EP, or custom names like UCCL_EP). +# MoE dispatch/combine EP backend (TURBO, DEEP_EP, MORI, UCCL). # Default: TURBO -ENV_MOE_DISPATCH_COMBINE_BACKEND = "PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND" +ENV_EP_BACKEND = "PRIMUS_TURBO_EP_BACKEND" # Enable auto-tuning across registered kernel backends ("1" to enable). # Default: "0" (disabled) @@ -39,3 +39,7 @@ # When set to "1", EP dispatch/combine kernels run on the caller's current CUDA stream. # Default: "0" ENV_EP_FORCE_CURRENT_STREAM = "PRIMUS_TURBO_EP_FORCE_CURRENT_STREAM" + +# Number of RDMA queue pairs per PE used by the Mori EP backend. +# Default: "2" +ENV_MORI_NUM_QP_PER_PE = "PRIMUS_TURBO_MORI_NUM_QP_PER_PE" diff --git a/primus_turbo/common/mori_utils.py b/primus_turbo/common/mori_utils.py new file mode 100644 index 000000000..a837a8f3e --- /dev/null +++ b/primus_turbo/common/mori_utils.py @@ -0,0 +1,93 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Helpers for the optional ``mori`` dependency. + +mori is imported lazily (only when the MORI EP backend runs). Primus-Turbo +requires the amd_mori release below; it is not on PyPI, so install from git tag. +``BUILD_UMBP=OFF`` skips mori's gRPC-dependent umbp component (unused here). +""" + +import importlib.metadata +from typing import NoReturn + +from primus_turbo.common.logger import logger + +# Required mori release. Keep in sync with MORI_VERSION in the ci / benchmark +# / release workflows. +MORI_VERSION = "1.2.0" +MORI_GIT_TAG = "v1.2.0" +_MORI_DIST_NAME = "amd_mori" +_MORI_GIT_URL = "https://github.com/ROCm/mori.git" + +_MORI_PIP_INSTALL = f'BUILD_UMBP=OFF pip install "amd_mori @ git+{_MORI_GIT_URL}@{MORI_GIT_TAG}"' + +MORI_INSTALL_HINT = ( + f"Primus-Turbo requires amd_mori=={MORI_VERSION} for the MORI EP backend. Install it with:\n" + f" {_MORI_PIP_INSTALL}" +) + +_version_checked = False + + +def _installed_mori_version(): + try: + return importlib.metadata.version(_MORI_DIST_NAME) + except importlib.metadata.PackageNotFoundError: + return None + + +def _versions_match(installed: str, expected: str) -> bool: + # Ignore any local/dev suffix (e.g. a source build's "1.2.1.dev6+g1234567"). + try: + from packaging.version import InvalidVersion, Version + + try: + return Version(installed).public == Version(expected).public + except InvalidVersion: + return False + except ImportError: + return installed.split("+")[0] == expected + + +def check_mori_version_once(): + """Warn once if the installed mori version differs from the pin.""" + global _version_checked + if _version_checked: + return + _version_checked = True + + installed = _installed_mori_version() + if installed and not _versions_match(installed, MORI_VERSION): + logger.warning( + "mori version mismatch: installed=%s, expected=%s; behavior/perf may differ. " + "To match, run:\n %s", + installed, + MORI_VERSION, + _MORI_PIP_INSTALL, + once=True, + ) + + +def raise_mori_missing(exc: Exception) -> NoReturn: + logger.error(MORI_INSTALL_HINT, once=True) + raise ImportError(MORI_INSTALL_HINT) from exc + + +_mori_module = None + + +def get_mori(): + """Import and return the ``mori`` module, lazily and with a clear error.""" + global _mori_module + if _mori_module is None: + try: + import mori + except ImportError as exc: + raise_mori_missing(exc) + check_mori_version_once() + _mori_module = mori + return _mori_module diff --git a/primus_turbo/common/uccl_utils.py b/primus_turbo/common/uccl_utils.py new file mode 100644 index 000000000..ff3075228 --- /dev/null +++ b/primus_turbo/common/uccl_utils.py @@ -0,0 +1,92 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Helpers for the optional ``uccl`` dependency. + +uccl is imported lazily (only when the UCCL EP backend runs). Primus-Turbo +requires the uccl release below; it is not on PyPI, so install from git tag. +""" + +import importlib.metadata +from typing import NoReturn + +from primus_turbo.common.logger import logger + +# Required uccl release. uccl has no usable release tag, so pin a main commit. +# Keep in sync with UCCL_VERSION/UCCL_GIT_TAG in the ci / benchmark / release workflows. +UCCL_VERSION = "0.1.1" +UCCL_GIT_TAG = "12fdebd8d9a40f12364239f51d74fc773c41a4a7" # [EP] LL clean & Zero-token LL rank issues (#993) +_UCCL_DIST_NAME = "uccl" +_UCCL_GIT_URL = "https://github.com/uccl-project/uccl.git" + +_UCCL_PIP_INSTALL = f'pip install "uccl @ git+{_UCCL_GIT_URL}@{UCCL_GIT_TAG}"' + +UCCL_INSTALL_HINT = ( + f"Primus-Turbo requires uccl=={UCCL_VERSION} for the UCCL EP backend. Install it with:\n" + f" {_UCCL_PIP_INSTALL}" +) + +_version_checked = False + + +def _installed_uccl_version(): + try: + return importlib.metadata.version(_UCCL_DIST_NAME) + except importlib.metadata.PackageNotFoundError: + return None + + +def _versions_match(installed: str, expected: str) -> bool: + # Ignore any local/dev suffix (e.g. a source build's "+g1234567"). + try: + from packaging.version import InvalidVersion, Version + + try: + return Version(installed).public == Version(expected).public + except InvalidVersion: + return False + except ImportError: + return installed.split("+")[0] == expected + + +def check_uccl_version_once(): + """Warn once if the installed uccl version differs from the pin.""" + global _version_checked + if _version_checked: + return + _version_checked = True + + installed = _installed_uccl_version() + if installed and not _versions_match(installed, UCCL_VERSION): + logger.warning( + "uccl version mismatch: installed=%s, expected=%s; behavior/perf may differ. " + "To match, run:\n %s", + installed, + UCCL_VERSION, + _UCCL_PIP_INSTALL, + once=True, + ) + + +def raise_uccl_missing(exc: Exception) -> NoReturn: + logger.error(UCCL_INSTALL_HINT, once=True) + raise ImportError(UCCL_INSTALL_HINT) from exc + + +_uccl_module = None + + +def get_uccl(): + """Import and return the ``uccl`` module, lazily and with a clear error.""" + global _uccl_module + if _uccl_module is None: + try: + import uccl + except ImportError as exc: + raise_uccl_missing(exc) + check_uccl_version_once() + _uccl_module = uccl + return _uccl_module diff --git a/primus_turbo/pytorch/core/backend.py b/primus_turbo/pytorch/core/backend.py index d56aa0bab..7ec16053c 100644 --- a/primus_turbo/pytorch/core/backend.py +++ b/primus_turbo/pytorch/core/backend.py @@ -15,9 +15,9 @@ from primus_turbo.common.constants import ( ENV_AUTO_TUNE, + ENV_EP_BACKEND, ENV_GEMM_BACKEND, ENV_GROUPED_GEMM_BACKEND, - ENV_MOE_DISPATCH_COMBINE_BACKEND, ) from primus_turbo.common.logger import logger from primus_turbo.triton.utils.origami import origami_clear_caches @@ -64,6 +64,8 @@ class BackendType(Enum): DEEP_EP = auto() TURBO = auto() FLYDSL = auto() + MORI = auto() + UCCL = auto() class GlobalBackendManager: @@ -80,7 +82,7 @@ class GlobalBackendManager: _gemm_backend: Dict[PrecisionType, Optional[BackendType]] = None _grouped_gemm_backend: Dict[PrecisionType, Optional[BackendType]] = None - _moe_dispatch_combine_backend: Dict[PrecisionType, Optional[BackendType]] = None + _ep_backend: Dict[PrecisionType, Optional[BackendType]] = None _auto_tune: Optional[bool] = None _env_cache: Dict[str, Dict["PrecisionType", "BackendType"]] = {} @@ -221,16 +223,16 @@ def get_grouped_gemm_backend(cls, precision: PrecisionType) -> Optional[BackendT return None @classmethod - def get_moe_dispatch_combine_backend(cls, precision: PrecisionType) -> Optional[BackendType]: - """Get the MoE dispatch combine backend configuration. Returns None if not set. + def get_ep_backend(cls, precision: PrecisionType) -> Optional[BackendType]: + """Get the EP backend configuration. Returns None if not set. - If the environment variable contains a value that is not a valid ``BackendType`` - (e.g. a custom EP backend name like ``UCCL_EP``), this method returns ``None`` so - the EP-specific backend registry in ``moe_dispatch_combine_impl`` can handle it. + If the environment variable contains a value that is not a valid ``BackendType``, + this method returns ``None`` so the EP-specific backend registry in + ``moe_dispatch_combine_impl`` can handle it. """ - if cls._moe_dispatch_combine_backend is not None: - return cls._moe_dispatch_combine_backend[precision] - env_value = os.environ.get(ENV_MOE_DISPATCH_COMBINE_BACKEND, None) + if cls._ep_backend is not None: + return cls._ep_backend[precision] + env_value = os.environ.get(ENV_EP_BACKEND, None) if env_value is not None and env_value.strip(): try: backend = cls._extract_backend_from_env(env_value).get(precision, None) @@ -240,7 +242,7 @@ def get_moe_dispatch_combine_backend(cls, precision: PrecisionType) -> Optional[ if backend is None: logger.warning( f"Precision {precision.name} not found in the environment variable " - f"{ENV_MOE_DISPATCH_COMBINE_BACKEND}. Using default backend.", + f"{ENV_EP_BACKEND}. Using default backend.", once=True, ) diff --git a/primus_turbo/pytorch/kernels/moe/_backend/__init__.py b/primus_turbo/pytorch/kernels/moe/_backend/__init__.py new file mode 100644 index 000000000..9937f6bf8 --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/__init__.py @@ -0,0 +1,34 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""EP backend implementations; submodules touch optional deps only when imported.""" + +from ._config import EPBufferConfig +from .base import ( + EPBackend, + _broadcast_from_rank0_float, + _broadcast_from_rank0_int, + _DeepEPLikeBackend, + bench, + detect_group_topology, +) +from .deep_ep import DeepEPBackend +from .mori import MoriEPBackend +from .turbo import TurboEPBackend +from .uccl_ep import UCCLEPBackend + +__all__ = [ + "EPBackend", + "EPBufferConfig", + "_DeepEPLikeBackend", + "_broadcast_from_rank0_float", + "_broadcast_from_rank0_int", + "bench", + "detect_group_topology", + "TurboEPBackend", + "DeepEPBackend", + "UCCLEPBackend", + "MoriEPBackend", +] diff --git a/primus_turbo/pytorch/kernels/moe/_backend/_config.py b/primus_turbo/pytorch/kernels/moe/_backend/_config.py new file mode 100644 index 000000000..499274c65 --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/_config.py @@ -0,0 +1,22 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +from dataclasses import dataclass +from typing import Any + + +@dataclass +class EPBufferConfig: + """Configuration for EP communication buffer initialization. + + Attributes: + num_sms: Number of SMs used by high-throughput kernels. + dispatch_config: Dispatch config; ``None`` means use backend default. + combine_config: Combine config; ``None`` means use backend default. + """ + + num_sms: int = 64 + dispatch_config: Any = None + combine_config: Any = None diff --git a/primus_turbo/pytorch/kernels/moe/_backend/base.py b/primus_turbo/pytorch/kernels/moe/_backend/base.py new file mode 100644 index 000000000..2f31cdbaf --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/base.py @@ -0,0 +1,920 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""``EPBackend`` Protocol + ``_DeepEPLikeBackend`` shared base for DeepEP-style backends.""" + +import functools +import inspect +import os +import socket +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Protocol, + Sequence, + Tuple, + Union, + runtime_checkable, +) + +import torch +import torch.distributed as dist + +from primus_turbo.common.logger import logger + +from ._config import EPBufferConfig + +_TOPOLOGY_CACHE: Dict[int, Tuple[int, int]] = {} + + +def call_once(method): + flag_attr = "_call_once__" + method.__qualname__.replace(".", "_") + + @functools.wraps(method) + def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + if getattr(self, flag_attr, False): + return None + result = method(self, *args, **kwargs) + setattr(self, flag_attr, True) + return result + + return wrapper # type: ignore[return-value] + + +def detect_group_topology(group: dist.ProcessGroup) -> Tuple[int, int]: + """Return ``(node_idx, num_nodes)`` for ``group``, cached per ``id(group)``.""" + cache_key = id(group) + cached = _TOPOLOGY_CACHE.get(cache_key) + if cached is not None: + return cached + + node_token = ( + os.environ.get("NODE_RANK") + or os.environ.get("GROUP_RANK") + or os.environ.get("SLURM_NODEID") + or socket.gethostname() + ) + + world = dist.get_world_size(group) + node_tokens = [None] * world + dist.all_gather_object(node_tokens, node_token, group=group) + + token_to_idx: Dict[str, int] = {} + for token in node_tokens: + if token not in token_to_idx: + token_to_idx[token] = len(token_to_idx) + + result = (token_to_idx[node_token], len(token_to_idx)) + _TOPOLOGY_CACHE[cache_key] = result + return result + + +def inplace_unique(x: torch.Tensor, num_slots: int) -> None: + assert x.dim() == 2 + mask = x < 0 + x_padded = x.masked_fill(mask, num_slots) + bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device) + bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded)) + bin_count = bin_count[:, :num_slots] + sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True) + sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1) + sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values + x[:, :].fill_(-1) + valid_len = min(num_slots, x.size(1)) + x[:, :valid_len] = sorted_bin_idx[:, :valid_len] + + +@torch.no_grad() +def bench( + fn, + *, + num_tests: int, + num_warmup: int, + l2_flush_mb: int = 256, + sync_fn: Optional[Callable[[], None]] = None, +) -> float: + """Time ``fn()`` per iteration with CUDA events; return mean seconds. + + Uses ``torch.cuda.Event`` (HIP events, no CUPTI) rather than the Kineto + profiler, which races with UCCL's RDMA workers on multi-node ROCm. An + LLC-sized buffer (``l2_flush_mb``) is zero-filled before each call to evict + ``fn``'s data from Infinity Cache; one event pair per iteration keeps that + flush outside the timed window. ``sync_fn`` realigns ranks before each call. + """ + flush_buf = torch.empty( + max(1, l2_flush_mb) * 1024 * 1024 // 4, + dtype=torch.float32, + device="cuda", + ) + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + + for _ in range(num_warmup): + if sync_fn is not None: + sync_fn() + fn() + torch.cuda.synchronize() + + for i in range(num_tests): + # Realign ranks before the timed call (outside the timed window). + if sync_fn is not None: + sync_fn() + flush_buf.zero_() + starts[i].record() + fn() + ends[i].record() + + torch.cuda.synchronize() + total_ms = sum(s.elapsed_time(e) for s, e in zip(starts, ends)) + return (total_ms / num_tests) * 1e-3 + + +def bench_dispatch_combine( + dispatch_fn: Callable[[], Any], + combine_fn: Callable[[Any], Any], + *, + num_tests: int, + num_warmup: int, + sync_fn: Optional[Callable[[], None]] = None, + l2_flush_mb: int = 256, + flush_buf: Optional[torch.Tensor] = None, +) -> Tuple[float, float]: + """Time a dispatch->combine pair; return ``(dispatch_s, combine_s)``. + + ``dispatch_fn()`` returns whatever ``combine_fn(recv)`` needs. They must run + as a pair: EP backends reuse a one-sided comm buffer that a dispatch-only run + leaves in a state the combine faults on. ``sync_fn`` realigns ranks before + each timed iteration, outside the timed window. + """ + if flush_buf is None: + flush_buf = torch.empty(max(1, l2_flush_mb) * 1024 * 1024 // 4, dtype=torch.float32, device="cuda") + + for _ in range(num_warmup): + if sync_fn: + sync_fn() + combine_fn(dispatch_fn()) + torch.cuda.synchronize() + + # 3 marks per iter: before dispatch | between | after combine. + new_events = lambda: [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] + pre, mid, post = new_events(), new_events(), new_events() + for i in range(num_tests): + if sync_fn: + sync_fn() + flush_buf.zero_() + pre[i].record() + recv = dispatch_fn() + mid[i].record() + combine_fn(recv) + post[i].record() + torch.cuda.synchronize() + + disp_ms = sum(a.elapsed_time(b) for a, b in zip(pre, mid)) / num_tests + comb_ms = sum(a.elapsed_time(b) for a, b in zip(mid, post)) / num_tests + return disp_ms * 1e-3, comb_ms * 1e-3 + + +def sweep_configs( + candidates: Sequence[Any], + dispatch_fn: Callable[[Any], Any], + combine_fn: Callable[[Any, Any], Any], + *, + group: dist.ProcessGroup, + num_tests: int, + num_warmup: int, + on_skip: Optional[Callable[[Any, BaseException], None]] = None, +) -> Tuple[Optional[Any], Optional[Any], float, float]: + """Benchmark each candidate's dispatch->combine pair; return the best + dispatch and best combine candidate (chosen independently) and their times. + + ``dispatch_fn(cand) -> recv`` and ``combine_fn(cand, recv)`` are the backend + launch closures; a candidate whose launch raises is skipped (``on_skip``). + """ + + def _sync() -> None: + torch.cuda.synchronize() + dist.barrier(group=group) + + flush_buf = torch.empty(256 * 1024 * 1024 // 4, dtype=torch.float32, device="cuda") + best_d, best_d_t = None, float("inf") + best_c, best_c_t = None, float("inf") + for cand in candidates: + dist.barrier(group=group) # realign in case a prior candidate raised on some ranks + try: + d_t, c_t = bench_dispatch_combine( + lambda c=cand: dispatch_fn(c), + lambda recv, c=cand: combine_fn(c, recv), + num_tests=num_tests, + num_warmup=num_warmup, + sync_fn=_sync, + flush_buf=flush_buf, + ) + except Exception as exc: # noqa: BLE001 - skip bad candidate, keep sweeping + if on_skip: + on_skip(cand, exc) + continue + if 0 < d_t < best_d_t: + best_d, best_d_t = cand, d_t + if 0 < c_t < best_c_t: + best_c, best_c_t = cand, c_t + return best_d, best_c, best_d_t, best_c_t + + +# ========================================================================= +# EPBackend Protocol +# ========================================================================= + + +@runtime_checkable +class EPBackend(Protocol): + """Structural interface for Expert-Parallel communication backends.""" + + @staticmethod + def is_available() -> bool: + """Return True if this backend's dependencies are importable.""" + ... + + @staticmethod + def supports_cuda_graph() -> bool: + """Return True if dispatch/combine are safe inside ``torch.cuda.graph``. + + Override to ``False`` on backends that rely on synchronous HIP APIs + or per-call device->host plumbing (they would deadlock under capture). + """ + ... + + # Capability methods are provided by ``_EPCapabilities`` (can_release / is_feasible). + + def is_initialized(self) -> bool: + """Return True if the backend is initialized.""" + ... + + def setup_env(self, **overrides: Optional[str]) -> None: + """Configure backend RDMA/network env vars before init. + + Resolution per var: explicit override > existing env > matching ``NCCL_*`` + > backend default. Backends without network state implement as no-op. + """ + ... + + def init_buffer( + self, + group: dist.ProcessGroup, + hidden_size: int, + num_experts: int, + num_topk: int, + seqlen: int, + fp8_dispatch: bool, + config: "EPBufferConfig", + ) -> None: + """(Re-)create the communication buffer if needed. + + Args: + group: EP process group. + hidden_size: Per-token hidden dimension (element count). + num_experts: Total number of experts. + num_topk: Number of topk experts to dispatch. + seqlen: Per-rank maximum dispatch tokens. + fp8_dispatch: Whether the dispatch payload is FP8 (``x`` is a + ``(fp8_tensor, scales)`` tuple). + config: Backend-specific tuning config. + """ + ... + + def dispatch( + self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + handle: Optional[tuple] = None, + topk_idx: Optional[torch.Tensor] = None, + token_weights: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + num_worst_tokens: int = 0, + config: Optional[Any] = None, + ) -> Tuple[ + Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[Union[List[int], torch.Tensor]], + Optional[tuple], + ]: + """Dispatch tokens to experts. + + Args: + x: Input tensor, or ``(fp8_tensor, scales)`` tuple for FP8 path. + handle: Cached dispatch handle; ``None`` for a primary call. + topk_idx: ``[num_tokens, num_topk]`` expert indices (primary only). + token_weights: ``[num_tokens, num_topk]`` expert weights. + num_experts: Total number of experts (primary only). + async_finish: If True, return before the kernel finishes. + allocate_on_comm_stream: Allocate outputs on the comm stream. + num_worst_tokens: Worst-case receive token count (for padding). + config: Backend-specific dispatch config. + + Returns: + ``(recv_x, recv_topk_idx, recv_topk_weights, tokens_per_expert, handle)``. + """ + ... + + def combine( + self, + x: torch.Tensor, + handle: tuple, + topk_weights: Optional[torch.Tensor] = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + config: Optional[Any] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Combine expert outputs back to the original token order. + + Args: + x: Expert-side tensor to be combined. + handle: Handle returned by the paired ``dispatch`` call. + topk_weights: Optional per-token expert weights. + async_finish: If True, return before the kernel finishes. + allocate_on_comm_stream: Allocate outputs on the comm stream. + config: Backend-specific combine config. + + Returns: + ``(combined_x, combined_topk_weights)``. + """ + ... + + def release_buffer(self) -> None: + """Release the communication buffer.""" + ... + + +def _apply_env_with_nccl_fallback( + mappings: Sequence[Tuple[str, str, Optional[str]]], + backend_name: str, +) -> None: + """Set each ``backend_env`` with the first available source: + explicit value > existing env > ``nccl_env`` > leave unset. + """ + for backend_env, nccl_env, explicit in mappings: + if explicit is not None: + os.environ[backend_env] = str(explicit) + elif backend_env not in os.environ: + nccl_val = os.environ.get(nccl_env) + if nccl_val is not None: + os.environ[backend_env] = nccl_val + if backend_env in os.environ: + logger.info( + f"[{backend_name} Network Settings] {backend_env}: {os.environ[backend_env]}", + rank=0, + ) + + +# Per-NIC inflight-env defaults; avoid hangs on AINIC (ionic_*) and Thor-2 (bnxt_re_*). +_UCCL_NIC_INFLIGHT_DEFAULTS: Dict[str, Dict[str, str]] = { + "ionic": { + "UCCL_IB_MAX_INFLIGHT_NORMAL": "1", + "UCCL_IB_MAX_INFLIGHT_LOW_LATENCY": "1", + "UCCL_IB_MAX_INFLIGHT_BYTES": "4194304", # 4MB + }, + "bnxt": { + "UCCL_IB_MAX_INFLIGHT_NORMAL": "1", + "UCCL_IB_MAX_INFLIGHT_LOW_LATENCY": "1", + "UCCL_IB_MAX_INFLIGHT_BYTES": "1572864", # 1.5MB + }, +} + + +def apply_uccl_network_env( + *, + ib_gid_index: Optional[str] = None, + ib_hca: Optional[str] = None, + socket_ifname: Optional[str] = None, + ib_tc: Optional[str] = None, + ib_sl: Optional[str] = None, + ib_max_inflight_normal: Optional[str] = None, + ib_max_inflight_low_latency: Optional[str] = None, + ib_max_inflight_bytes: Optional[str] = None, +) -> None: + """Set UCCL_* RDMA env vars, falling back to ``NCCL_*`` equivalents. + + Shared by the UCCL backend and the uccl-backed deep_ep wrapper. ``None`` + kwargs inherit from ``NCCL_*`` if set, else per-NIC defaults. + """ + from .nic_detect import detect_nic_type + + _apply_env_with_nccl_fallback( + [ + ("UCCL_IB_GID_INDEX", "NCCL_IB_GID_INDEX", ib_gid_index), + ("UCCL_IB_HCA", "NCCL_IB_HCA", ib_hca), + ("UCCL_SOCKET_IFNAME", "NCCL_SOCKET_IFNAME", socket_ifname), + ("UCCL_IB_TC", "NCCL_IB_TC", ib_tc), + ("UCCL_IB_SL", "NCCL_IB_SL", ib_sl), + ], + backend_name="UCCL", + ) + + # NIC-specific defaults for inflight envs (run after HCA is resolved). + nic_defaults = _UCCL_NIC_INFLIGHT_DEFAULTS.get(detect_nic_type() or "", {}) + inflight_envs = [ + ("UCCL_IB_MAX_INFLIGHT_NORMAL", ib_max_inflight_normal), + ("UCCL_IB_MAX_INFLIGHT_LOW_LATENCY", ib_max_inflight_low_latency), + ("UCCL_IB_MAX_INFLIGHT_BYTES", ib_max_inflight_bytes), + ] + for env_name, explicit in inflight_envs: + if explicit is not None: + os.environ[env_name] = str(explicit) + elif env_name not in os.environ and env_name in nic_defaults: + os.environ[env_name] = nic_defaults[env_name] + if env_name in os.environ: + logger.info( + f"[UCCL Network Settings] {env_name}: {os.environ[env_name]}", + rank=0, + ) + + +def _broadcast_from_rank0_int(values: Sequence[int], group: dist.ProcessGroup) -> List[int]: + """Return rank-0's ``values`` on every rank (via ``all_gather``).""" + t = torch.tensor(values, dtype=torch.int32, device="cuda") + gathered = [torch.zeros_like(t) for _ in range(dist.get_world_size(group))] + dist.all_gather(gathered, t, group=group) + return gathered[0].tolist() + + +def _broadcast_from_rank0_float(value: float, group: dist.ProcessGroup) -> float: + """Float counterpart of :func:`_broadcast_from_rank0_int`.""" + t = torch.tensor([value], dtype=torch.float64, device="cuda") + gathered = [torch.zeros_like(t) for _ in range(dist.get_world_size(group))] + dist.all_gather(gathered, t, group=group) + return float(gathered[0].item()) + + +# ========================================================================= +# Backend capability defaults (shared by every EP backend) +# ========================================================================= + + +class _EPCapabilities: + """Backend capability defaults; each backend overrides the one that differs.""" + + @classmethod + def can_release(cls, *, will_reinit: bool) -> bool: # noqa: ARG003 + """Whether release_buffer() is safe now; ``will_reinit`` = re-init'd afterwards. + + Default always safe. uccl: never (destroy corrupts HIP). mori: only when not reinit. + """ + return True + + @classmethod + def is_feasible(cls, group: dist.ProcessGroup) -> bool: # noqa: ARG003 + """Can this backend run on ``group``'s topology (e.g. inter-node)?""" + return True + + +# ========================================================================= +# _DeepEPLikeBackend — shared implementation for DeepEP-compatible backends +# ========================================================================= + + +class _DeepEPLikeBackend(_EPCapabilities): + """Shared base class for backends that follow the DeepEP Buffer protocol.""" + + def __init__(self) -> None: + self._buffer = None + + # ------------------------------------------------------------------ + # Subclass hooks + # ------------------------------------------------------------------ + + def is_initialized(self) -> bool: + """Return True if the backend is initialized.""" + return self._buffer is not None + + @staticmethod + def is_available() -> bool: + """Return True if backend dependencies are importable.""" + raise NotImplementedError + + @staticmethod + def supports_cuda_graph() -> bool: + """DeepEP-style backends run on the caller's stream and are graph-safe.""" + return True + + @staticmethod + def _get_module(): + """Return the backend Python module exposing ``Buffer``/``Config``/events.""" + raise NotImplementedError + + def _make_buffer_kwargs(self, group: dist.ProcessGroup) -> dict: + """Extra kwargs for ``BufferClass(...)``; pass ``is_intranode`` if it requires one.""" + BufferClass = self._get_module().Buffer + try: + param = inspect.signature(BufferClass).parameters.get("is_intranode") + except (TypeError, ValueError): + param = None + if param is not None and param.default is False: + return {"is_intranode": group.size() <= 8} + return {} + + def setup_env(self, **overrides: Optional[str]) -> None: # noqa: ARG002 + """No backend-specific env vars to configure by default. + + Subclasses that need RDMA / socket env wiring (UCCL, ...) override + this to translate NCCL_* vars into their own namespace. + """ + return + + # ------------------------------------------------------------------ + # EPBackend interface + # ------------------------------------------------------------------ + + def init_buffer( + self, + group: dist.ProcessGroup, + hidden_size: int, + num_experts: int, # noqa: ARG002 - DeepEP sizes by hidden_bytes only. + num_topk: int, # noqa: ARG002 - DeepEP sizes by hidden_bytes only. + seqlen: int, # noqa: ARG002 - DeepEP sizes by hidden_bytes only. + fp8_dispatch: bool, + config: EPBufferConfig, + ) -> None: + # Best-effort NCCL-fallback network env; no-op without RDMA, idempotent. + self.setup_env() + + mod = self._get_module() + BufferClass = mod.Buffer + + BufferClass.set_num_sms(config.num_sms) + + dispatch_config = config.dispatch_config or BufferClass.get_dispatch_config(group.size()) + combine_config = config.combine_config or BufferClass.get_combine_config(group.size()) + + element_size = 1 if fp8_dispatch else 2 + hidden_bytes = hidden_size * max(element_size, 2) + + num_nvl_bytes, num_rdma_bytes = 0, 0 + for cfg in (dispatch_config, combine_config): + num_nvl_bytes = max( + cfg.get_nvl_buffer_size_hint(hidden_bytes, group.size()), + num_nvl_bytes, + ) + try: + num_rdma_bytes = max( + cfg.get_rdma_buffer_size_hint(hidden_bytes, group.size()), + num_rdma_bytes, + ) + except (RuntimeError, AttributeError): + pass + + buf_kwargs = self._make_buffer_kwargs(group) + # Required: GC-based teardown leaks UCCL proxies and hangs. + buf_kwargs.setdefault("explicitly_destroy", True) + + if ( + self._buffer is None + # or not isinstance(self._buffer, BufferClass) + or self._buffer.group != group + or self._buffer.num_nvl_bytes < num_nvl_bytes + or self._buffer.num_rdma_bytes < num_rdma_bytes + ): + if self._buffer is not None: + assert isinstance(self._buffer, BufferClass) + # Tear down old buffer to clear UCCL proxy/IPC state before realloc. + self.release_buffer() + self._buffer = BufferClass(group, num_nvl_bytes, num_rdma_bytes, **buf_kwargs) + logger.info( + f"[{self.__class__.__name__} init] world_size={group.size()} " + f"rank={group.rank()} hidden={hidden_size} " + f"fp8_dispatch={fp8_dispatch} num_nvl_bytes={num_nvl_bytes} " + f"num_rdma_bytes={num_rdma_bytes} ", + rank=0, + ) + + # ----- helpers ------------------------------------------------------ + + def _get_event_classes(self): + mod = self._get_module() + EventOverlapClass = mod.utils.EventOverlap if hasattr(mod, "utils") else mod.EventOverlap + EventHandleClass = mod.utils.EventHandle if hasattr(mod, "utils") else mod.EventHandle + return EventOverlapClass, EventHandleClass + + # ----- dispatch / combine ------------------------------------------- + + def dispatch( + self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + handle: Optional[tuple] = None, + topk_idx: Optional[torch.Tensor] = None, + token_weights: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + num_worst_tokens: int = 0, + config: Optional[Any] = None, + ): + assert self.is_initialized(), "Backend is not initialized" + + EventOverlapClass, EventHandleClass = self._get_event_classes() + previous_event = EventOverlapClass(EventHandleClass()) if async_finish else None + + if handle is None: + assert topk_idx is not None + assert token_weights is not None + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + event, + ) = self._buffer.get_dispatch_layout( + topk_idx, + num_experts, + previous_event=previous_event, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + + ( + recv_x, + recv_token_indices, + recv_token_probs, + tokens_per_expert, + handle, + after_event, + ) = self._buffer.dispatch( + x, + topk_idx=topk_idx, + topk_weights=token_weights, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=event, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + num_worst_tokens=num_worst_tokens, + config=config, + ) + else: + recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle, after_event = ( + self._buffer.dispatch( + x, + handle=handle, + previous_event=previous_event, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + config=config, + ) + ) + + if async_finish: + after_event.current_stream_wait() + + return recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle + + def combine( + self, + x: torch.Tensor, + handle: tuple, + topk_weights: Optional[torch.Tensor] = None, + async_finish: bool = False, + allocate_on_comm_stream: bool = False, + config: Optional[Any] = None, + ): + assert self.is_initialized(), "Backend is not initialized" + EventOverlapClass, EventHandleClass = self._get_event_classes() + previous_event = EventOverlapClass(EventHandleClass()) if async_finish else None + + combined_x, combined_topk_weights, after_event = self._buffer.combine( + x, + handle=handle, + topk_weights=None if topk_weights is None else topk_weights.float(), + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + previous_event=previous_event, + config=config, + ) + + if async_finish: + after_event.current_stream_wait() + + return combined_x, combined_topk_weights + + def release_buffer(self) -> None: + """Explicitly destroy the buffer (paired with ``explicitly_destroy=True``).""" + buf = self._buffer + if buf is None: + return + # Detach first so a partial destroy() can't be reused by init_buffer. + self._buffer = None + try: + # Drain comm-stream work before C++ destroy()'s NVL barrier. + torch.cuda.synchronize() + if getattr(buf, "explicitly_destroy", False): + buf.destroy() + except Exception as exc: + logger.warning( + f"{type(self).__name__}.release_buffer: buffer.destroy() " + f"raised {type(exc).__name__}: {exc!r}; resources may leak." + ) + + # ----- autotune ------------------------------------------------------ + + @torch.no_grad() + def tune_configs( + self, + group: dist.ProcessGroup, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + num_experts: int, + *, + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + num_sms: int = 32, + num_tests: int = 20, + num_topk: Optional[int] = None, + uniform_dispatch: bool = True, + ) -> Tuple[Any, Any, float, float]: + """Sweep ``(nvl_chunk_size, rdma_chunk_size)`` candidates and pick the + fastest dispatch / combine ``Config`` for this backend. + + Args: + group: The EP process group. + x: Dispatch input (Tensor or ``(fp8_tensor, scales)`` tuple). + num_experts: Total number of experts. + topk_idx: ``[num_tokens, num_topk]`` expert indices. Required when + ``uniform_dispatch=False``; only used to derive ``num_topk`` + in uniform mode. + topk_weights: ``[num_tokens, num_topk]`` expert weights. Required + when ``uniform_dispatch=False``; ignored in uniform mode. + num_sms: SM count to use for tuning. + num_tests: Timed iterations per candidate. + num_topk: Number of experts per token (uniform mode fallback + when ``topk_idx`` is not supplied). + uniform_dispatch: If True, resample a near-uniform ``topk_idx`` / + ``topk_weights`` internally for more robust tuning. + + Returns: + ``(best_dispatch_config, best_combine_config, best_dispatch_s, best_combine_s)`` + with times identical on all ranks (rank 0's pick, broadcast). + """ + mod = self._get_module() + BufferClass = mod.Buffer + BufferClass.set_num_sms(num_sms) + ConfigClass = mod.Config + ep_size = group.size() + _, num_nodes = detect_group_topology(group) + + x_tensor = x if isinstance(x, torch.Tensor) else x[0] + hidden_size = int(x_tensor.size(1)) + seqlen = int(x_tensor.size(0)) + fp8_dispatch = isinstance(x, tuple) + # Byte count used for RDMA bandwidth accounting. + hidden_size * max(1 if fp8_dispatch else 2, 2) + + # Resample (topk_idx, topk_weights) when uniform_dispatch is on. + x_inp = x if isinstance(x, torch.Tensor) else x[0] + if uniform_dispatch: + if num_topk is None: + if topk_idx is None: + raise ValueError( + "tune_configs(uniform_dispatch=True): need either " + "``num_topk`` or a shape-carrying ``topk_idx`` to " + "sample a uniform distribution." + ) + num_topk = int(topk_idx.size(1)) + if num_topk > num_experts: + raise ValueError( + "tune_configs(uniform_dispatch=True): ``num_topk`` " + f"({num_topk}) cannot exceed ``num_experts`` ({num_experts})." + ) + num_tokens = int(x_inp.size(0)) + device = x_inp.device + scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device=device).abs() + 1 + topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)[1] + topk_weights = torch.randn((num_tokens, num_topk), dtype=torch.float32, device=device) + else: + if topk_idx is None or topk_weights is None: + raise ValueError( + "tune_configs(uniform_dispatch=False): ``topk_idx`` and " + "``topk_weights`` are both required." + ) + + # Per-(EP size) NVL/RDMA buffer sizes (chunks), from the upstream UCCL-EP sweep. + rdma_buffer_size, nvl_buffer_size = 512, (720 if ep_size in (144, 160) else 512) + if ep_size == 24: + nvl_buffer_size = 540 + + # Sweep ranges; intranode pins rdma_chunk since RDMA is unused. + nvl_chunk_range = range(1, 8, 1) + if num_nodes <= 1: + rdma_chunk_range = (16,) + else: + rdma_chunk_range = range(12 if num_nodes == 2 else 8, 33, 4) + + worst_nvl_chunk = max(nvl_chunk_range) + worst_rdma_chunk = max( + rdma_chunk_range, + key=lambda c: -(-rdma_buffer_size // c) * c, + ) + worst_config = ConfigClass( + num_sms, + worst_nvl_chunk, + nvl_buffer_size, + worst_rdma_chunk, + rdma_buffer_size, + ) + self.init_buffer( + group, + hidden_size=hidden_size, + num_experts=num_experts, + num_topk=num_topk, + seqlen=seqlen, + fp8_dispatch=fp8_dispatch, + config=EPBufferConfig( + num_sms=num_sms, + dispatch_config=worst_config, + combine_config=worst_config, + ), + ) + + # Seed one real dispatch so later tune runs can use the cached path. + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + _, + ) = self._buffer.get_dispatch_layout(topk_idx, num_experts) + + topk_weights_f = topk_weights.float() if topk_weights is not None else None + seed_args = { + "x": x, + "num_tokens_per_rank": num_tokens_per_rank, + "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, + "is_token_in_rank": is_token_in_rank, + "num_tokens_per_expert": num_tokens_per_expert, + "topk_idx": topk_idx, + "topk_weights": topk_weights_f, + } + # Seed one dispatch for the layout ``handle`` reused below; recv_x dropped. + _, _, _, _, handle, _ = self._buffer.dispatch(**seed_args) + + num_warmup = max(1, num_tests // 5) + + # Sweep dispatch+combine over the (nvl_chunk, rdma_chunk) grid. + candidates = [ + (nvl_chunk_size, rdma_chunk_size) + for nvl_chunk_size in nvl_chunk_range + for rdma_chunk_size in rdma_chunk_range + ] + + def _cfg(cand): + nvl, rdma = cand + return ConfigClass(num_sms, nvl, nvl_buffer_size, rdma, rdma_buffer_size) + + def _dispatch(cand): + # recv_x is element [0] of the handle-path dispatch return tuple. + return self._buffer.dispatch(x=x, handle=handle, config=_cfg(cand))[0] + + def _combine(cand, recv): + self._buffer.combine(x=recv, handle=handle, config=_cfg(cand)) + + def _on_skip(cand, exc): + if group.rank() == 0: + logger.debug(f"[tuning] skip invalid cfg nvl/rdma={cand}: {exc!r}") + + best_d, best_c, best_dispatch_time, best_combine_time = sweep_configs( + candidates, + _dispatch, + _combine, + group=group, + num_tests=num_tests, + num_warmup=num_warmup, + on_skip=_on_skip, + ) + if best_d is None or best_dispatch_time == float("inf"): + raise RuntimeError("tune_configs: no valid dispatch config found.") + if best_c is None or best_combine_time == float("inf"): + raise RuntimeError("tune_configs: no valid combine config found.") + + # Agree on rank-0's winners + times so every rank uses identical configs. + d_win = _broadcast_from_rank0_int(list(best_d), group) + c_win = _broadcast_from_rank0_int(list(best_c), group) + best_dispatch_time = _broadcast_from_rank0_float(best_dispatch_time, group) + best_combine_time = _broadcast_from_rank0_float(best_combine_time, group) + + dispatch_config = ConfigClass(num_sms, d_win[0], nvl_buffer_size, d_win[1], rdma_buffer_size) + combine_config = ConfigClass(num_sms, c_win[0], nvl_buffer_size, c_win[1], rdma_buffer_size) + + if group.rank() == 0: + logger.debug( + f"[tuning] Best dispatch: SMs {num_sms}, NVL chunk {d_win[0]}, " + f"RDMA chunk {d_win[1]}, {best_dispatch_time * 1e6:.2f} us; " + f"Best combine: NVL chunk {c_win[0]}, RDMA chunk {c_win[1]}, " + f"{best_combine_time * 1e6:.2f} us" + ) + + return dispatch_config, combine_config, best_dispatch_time, best_combine_time diff --git a/primus_turbo/pytorch/kernels/moe/_backend/deep_ep.py b/primus_turbo/pytorch/kernels/moe/_backend/deep_ep.py new file mode 100644 index 000000000..04ea5738c --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/deep_ep.py @@ -0,0 +1,49 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""External ``deep_ep`` package backend (optional).""" + +from typing import Optional + +from .base import _DeepEPLikeBackend, apply_uccl_network_env, call_once + + +class DeepEPBackend(_DeepEPLikeBackend): + """External ``deep_ep`` package backend (optional).""" + + @classmethod + def is_available(cls) -> bool: + try: + import deep_ep # noqa: F401 + except ImportError: + return False + # uccl-backed deep_ep is owned by the UCCL backend. + return not cls._is_uccl_backed() + + @staticmethod + def _get_module(): + import deep_ep + + return deep_ep + + @classmethod + def _is_uccl_backed(cls) -> bool: + # uccl wrapper re-exports Config from uccl.ep; on import error assume uccl (never free). + try: + config = getattr(cls._get_module(), "Config", None) + except Exception: # noqa: BLE001 - import quirks must not break capability checks + return True + return getattr(config, "__module__", "").startswith("uccl") + + @classmethod + def can_release(cls, *, will_reinit: bool) -> bool: # noqa: ARG003 + # uccl destroy() corrupts the HIP context: never release. + return not cls._is_uccl_backed() + + @call_once + def setup_env(self, **overrides: Optional[str]) -> None: + # The uccl-backed wrapper needs the UCCL_* RDMA env like the UCCL backend. + if self._is_uccl_backed(): + apply_uccl_network_env(**overrides) diff --git a/primus_turbo/pytorch/kernels/moe/_backend/expert_token_count.py b/primus_turbo/pytorch/kernels/moe/_backend/expert_token_count.py new file mode 100644 index 000000000..14adf3ad1 --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/expert_token_count.py @@ -0,0 +1,164 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""Per-expert token counting (Mori dispatch post-processing).""" + +from typing import List, Optional, Tuple + +import torch +import triton +import triton.language as tl + + +def _compute_expert_token_info_configs() -> List[triton.Config]: + """Autotune space for :func:`compute_expert_token_info_kernel`.""" + configs: List[triton.Config] = [] + for block_m in (64, 128, 256, 512, 1024): + for num_warps in (1, 2, 4, 8): + # Wave=64, cap threads/CTA to 1024 (HW limit on CDNA). + if num_warps * 64 > 1024: + continue + # Keep at least one element per thread in the M dimension. + if block_m < num_warps * 64: + continue + for num_stages in (1, 2): + configs.append( + triton.Config( + {"BLOCK_SIZE_M": block_m}, + num_warps=num_warps, + num_stages=num_stages, + ) + ) + return configs + + +@triton.autotune( + configs=_compute_expert_token_info_configs(), + key=["num_tokens", "num_topk"], + reset_to_zero=["num_recv_tokens_per_expert_ptr"], +) +@triton.jit +def compute_expert_token_info_kernel( + recv_topk_idx_ptr, + deepep_topk_idx_ptr, + num_recv_tokens_per_expert_ptr, + total_recv_ptr, + num_tokens: tl.int32, + num_topk: tl.int32, + num_local_experts: tl.int32, + expert_base: tl.int32, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + HAS_TOTAL_RECV: tl.constexpr, +): + """Count per-expert received tokens and emit a DeepEP-format topk_idx.""" + pid = tl.program_id(0) + row_offs = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + col_offs = tl.arange(0, BLOCK_SIZE_K) + bin_offs = tl.arange(0, BLOCK_SIZE_N) + + if HAS_TOTAL_RECV: + total_recv = tl.load(total_recv_ptr) + row_limit = tl.minimum(total_recv, num_tokens) + else: + row_limit = num_tokens + + in_footprint = row_offs < num_tokens + in_valid_rows = row_offs < row_limit + col_mask = col_offs < num_topk + + footprint_mask = in_footprint[:, None] & col_mask[None, :] + load_mask = in_valid_rows[:, None] & col_mask[None, :] + + safe_row = tl.where(in_footprint, row_offs, 0) + expert_offs = safe_row[:, None] * num_topk + col_offs[None, :] + + topk_experts = tl.load( + recv_topk_idx_ptr + expert_offs, + mask=load_mask, + other=0, + ) + + local_experts = topk_experts - expert_base + valid = load_mask & (local_experts >= 0) & (local_experts < num_local_experts) + # Clamp invalid lanes to a legal bin; the histogram mask excludes them. + safe_experts = tl.where(valid, local_experts, 0) + + # ``tl.histogram`` requires a flat 1D input; reshape the 2D tile. + flat_experts = tl.reshape(safe_experts, [BLOCK_SIZE_M * BLOCK_SIZE_K]) + flat_mask = tl.reshape(valid, [BLOCK_SIZE_M * BLOCK_SIZE_K]) + local_counts = tl.histogram(flat_experts, BLOCK_SIZE_N, mask=flat_mask) + + # Flush CTA-local accumulator to global memory with one atomic per bin. + bin_mask = (bin_offs < num_local_experts) & (local_counts > 0) + tl.atomic_add( + num_recv_tokens_per_expert_ptr + bin_offs, + local_counts, + sem="relaxed", + scope="gpu", + mask=bin_mask, + ) + + deepep_experts = tl.where(valid, local_experts, -1).to(topk_experts.dtype) + tl.store(deepep_topk_idx_ptr + expert_offs, deepep_experts, mask=footprint_mask) + + +def compute_expert_token_info( + recv_topk_idx: torch.Tensor, + num_local_experts: int, + *, + expert_base: int = 0, + total_recv: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Count per-expert tokens and return recv_topk_idx in DeepEP layout.""" + assert recv_topk_idx.is_cuda, "recv_topk_idx must be a CUDA tensor" + assert recv_topk_idx.dim() == 2, "recv_topk_idx must be 2D [num_tokens, num_topk]" + + device = recv_topk_idx.device + num_tokens, num_topk = recv_topk_idx.shape + + num_recv_tokens_per_expert = torch.zeros(num_local_experts, dtype=recv_topk_idx.dtype, device=device) + deepep_like_topk_idx = torch.empty_like(recv_topk_idx) + + if num_tokens == 0 or num_topk == 0 or num_local_experts == 0: + # Nothing to count; fill with the DeepEP padding sentinel. + if deepep_like_topk_idx.numel() > 0: + deepep_like_topk_idx.fill_(-1) + return num_recv_tokens_per_expert, deepep_like_topk_idx + + has_total_recv = total_recv is not None + if has_total_recv: + assert total_recv.is_cuda, "total_recv must be a CUDA tensor" + assert total_recv.numel() >= 1, "total_recv must have at least 1 element" + # Triton expects an int32 pointer; coerce if needed. + if total_recv.dtype != torch.int32: + total_recv = total_recv.to(torch.int32) + total_recv_ptr = total_recv + else: + # Dummy pointer; unused when HAS_TOTAL_RECV=False. + total_recv_ptr = recv_topk_idx + + # ``tl.histogram`` requires ``num_bins`` to be a power of 2 on AMD Triton. + block_size_n = triton.next_power_of_2(num_local_experts) + + # ``BLOCK_SIZE_M`` is picked by the autotuner; the grid size depends on it. + def grid(META): + return (triton.cdiv(num_tokens, META["BLOCK_SIZE_M"]),) # noqa: E731 + + compute_expert_token_info_kernel[grid]( + recv_topk_idx, + deepep_like_topk_idx, + num_recv_tokens_per_expert, + total_recv_ptr, + num_tokens, + num_topk, + num_local_experts, + expert_base, + BLOCK_SIZE_K=triton.next_power_of_2(num_topk), + BLOCK_SIZE_N=block_size_n, + HAS_TOTAL_RECV=has_total_recv, + ) + return num_recv_tokens_per_expert, deepep_like_topk_idx diff --git a/primus_turbo/pytorch/kernels/moe/_backend/mori.py b/primus_turbo/pytorch/kernels/moe/_backend/mori.py new file mode 100644 index 000000000..fbf282efb --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/mori.py @@ -0,0 +1,710 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""ROCm Mori EP backend (optional).""" + +import dataclasses +import gc +import os +from dataclasses import dataclass +from enum import Enum +from functools import lru_cache +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.distributed as dist + +from primus_turbo.common.constants import ENV_MORI_NUM_QP_PER_PE +from primus_turbo.common.logger import logger +from primus_turbo.common.mori_utils import get_mori + +from ._config import EPBufferConfig +from .base import ( + _apply_env_with_nccl_fallback, + _broadcast_from_rank0_float, + _broadcast_from_rank0_int, + _EPCapabilities, + call_once, + detect_group_topology, + sweep_configs, +) +from .expert_token_count import compute_expert_token_info + +# ========================================================================== +# Mori EP backend helpers +# ========================================================================== + + +class _MoriEpMode(Enum): + """Mori EP deployment mode (normal dispatch only).""" + + INTRA_NODE = "intra_node" + INTER_NODE = "inter_node" + + +@dataclass(frozen=True) +class _MoriDispatchCfg: + """Mori kernel launch config for a given mode / token budget.""" + + kernel_type: Any # mori.ops.EpDispatchCombineKernelType + warp_num_per_block: int + block_num: int + rdma_block_num: int + + +def _get_mori_dispatch_configs() -> Dict[_MoriEpMode, _MoriDispatchCfg]: + """Return per-mode Mori kernel launch configs.""" + import mori.ops + + # Defaults aligned with mori's official gfx950 tuning (used when autotune off). + return { + _MoriEpMode.INTRA_NODE: _MoriDispatchCfg( + kernel_type=mori.ops.EpDispatchCombineKernelType.IntraNode, + warp_num_per_block=16, + block_num=256, + rdma_block_num=0, + ), + _MoriEpMode.INTER_NODE: _MoriDispatchCfg( + kernel_type=mori.ops.EpDispatchCombineKernelType.InterNodeV1, + warp_num_per_block=8, + block_num=128, + rdma_block_num=64, + ), + } + + +def _mori_block_num_candidates(num_sms: int, sm_count: int) -> List[int]: + """block_num sweep set: powers of two up to sm_count, plus sm_count and num_sms. + + Mirrors mori's bench; excludes over-subscription (>sm_count, hang risk). + block_num does not change op buffer size, so the sweep is memory-free. + """ + cap = int(sm_count) if sm_count and sm_count > 0 else int(num_sms or 64) + candidates = {cap} + p = 32 + while p <= cap: + candidates.add(p) + p *= 2 + if num_sms and num_sms > 0: + candidates.add(min(int(num_sms), cap)) + return sorted(v for v in candidates if 0 < v <= cap) + + +# mori's native env: AUTO makes the op auto-apply its shipped tuning JSON per-call. +_MORI_LAUNCH_CONFIG_MODE_ENV = "MORI_EP_LAUNCH_CONFIG_MODE" + + +def _mori_auto_mode() -> bool: + """True when mori auto-tunes from its shipped JSON DB (opt-in via env=AUTO). + + Default MANUAL: our benchmark sweep beats the JSON for bf16 (mori's dispatch + JSON is fp8-only, so bf16 dispatch falls back to a slower hard-coded cfg). + """ + return os.environ.get(_MORI_LAUNCH_CONFIG_MODE_ENV, "MANUAL").upper() == "AUTO" + + +_MORI_SHMEM_PG_NAME = "mori" + + +def _align_mori_hip_with_torch() -> None: + """Make Mori's JIT reuse torch's libamdhip64.so (fixes HIP error 500 on ROCm 7.2+).""" + try: + import ctypes + + import mori.jit.hip_driver as _hd + except ImportError: + return + + torch_hip = os.path.join(os.path.dirname(torch.__file__), "lib", "libamdhip64.so") + if not os.path.isfile(torch_hip): + return + + try: + _hd._hip = ctypes.CDLL(torch_hip) + logger.info(f"[MORI init] Aligned Mori HIP with torch ({torch_hip}).", rank=0, once=True) + except OSError as e: # pragma: no cover - best-effort. + logger.warning(f"[MORI init] Failed to preload {torch_hip}: {e}.", once=True) + + +def _register_and_init_mori_shmem(group: dist.ProcessGroup) -> None: + """Register ``group`` with the Mori SHMEM runtime (idempotent).""" + import mori.shmem + + assert dist.is_initialized(), "torch.distributed must be initialized before Mori SHMEM init." + + try: + torch._C._distributed_c10d._register_process_group(_MORI_SHMEM_PG_NAME, group) + except Exception as e: # noqa: BLE001 - mori binds raise a mix of exception types. + if "already registered" not in str(e): + raise + logger.info( + f"[MORI init] Process group already registered under " + f"'{_MORI_SHMEM_PG_NAME}'; reusing existing SHMEM binding " + f"({e}).", + rank=0, + ) + return + mori.shmem.shmem_torch_process_group_init(_MORI_SHMEM_PG_NAME) + + +def _extract_mori_launch_override( + config: Optional[Any], +) -> Tuple[int, int, int]: + """Return (block_num, warp_per_block, rdma_block_num); -1 means op default.""" + if isinstance(config, _MoriDispatchCfg): + return config.block_num, config.warp_num_per_block, config.rdma_block_num + return -1, -1, -1 + + +def _resolve_mori_dispatch_cfg( + group: dist.ProcessGroup, + config: Optional[EPBufferConfig] = None, +) -> _MoriDispatchCfg: + """Resolve launch cfg: topology default, overridden by num_sms and dispatch_config.""" + _, num_nodes = detect_group_topology(group) + mode = _MoriEpMode.INTRA_NODE if num_nodes <= 1 else _MoriEpMode.INTER_NODE + base = _get_mori_dispatch_configs()[mode] + + kernel_type = base.kernel_type + warp_num_per_block = base.warp_num_per_block + block_num = base.block_num + rdma_block_num = base.rdma_block_num + + num_sms_override = config.num_sms if config is not None else 0 + full_override = ( + config.dispatch_config + if config is not None and isinstance(config.dispatch_config, _MoriDispatchCfg) + else None + ) + + if num_sms_override > 0: + block_num = int(num_sms_override) + + if full_override is not None: + kernel_type = full_override.kernel_type + warp_num_per_block = full_override.warp_num_per_block + block_num = full_override.block_num + rdma_block_num = full_override.rdma_block_num + + return _MoriDispatchCfg( + kernel_type=kernel_type, + warp_num_per_block=warp_num_per_block, + block_num=block_num, + rdma_block_num=rdma_block_num, + ) + + +@lru_cache(maxsize=8) +def _build_mori_op( + rank: int, + world_size: int, + hidden: int, + params_dtype: torch.dtype, + num_local_experts: int, + num_topk: int, + num_max_dispatch_tokens_per_rank: int, + kernel_type: Any, + warp_num_per_block: int, + block_num: int, + rdma_block_num: int, +): + """Build and cache a :class:`mori.ops.EpDispatchCombineOp`.""" + import mori.ops + + # mori reads MORI_EP_LAUNCH_CONFIG_MODE itself (native default MANUAL); set it + # to AUTO to make the op auto-apply its shipped tuning JSON per-call. + num_qp_per_pe = int(os.environ.get(ENV_MORI_NUM_QP_PER_PE, "2")) + + common_kwargs = dict( + data_type=params_dtype, + rank=rank, + world_size=world_size, + hidden_dim=hidden, + scale_dim=0, + scale_type_size=0, + max_token_type_size=max(2, params_dtype.itemsize), + max_num_inp_token_per_rank=num_max_dispatch_tokens_per_rank, + num_experts_per_rank=num_local_experts, + num_experts_per_token=num_topk, + warp_num_per_block=warp_num_per_block, + block_num=block_num, + max_total_recv_tokens=0, + use_external_inp_buf=True, + kernel_type=kernel_type, + gpu_per_node=torch.cuda.device_count(), + rdma_block_num=rdma_block_num, + num_qp_per_pe=num_qp_per_pe, + quant_type="none", + ) + + def _check_mori_compatibility(kwargs: dict) -> dict: + """Drop kwargs not accepted by the installed Mori config.""" + valid = {f.name for f in dataclasses.fields(mori.ops.EpDispatchCombineConfig)} + cleaned = dict(kwargs) + for key in list(cleaned.keys()): + if key not in valid: + logger.warning( + f"[MORI compat] Dropping incompatible EpDispatchCombineConfig " f"argument '{key}'.", + once=True, + ) + del cleaned[key] + return cleaned + + common_kwargs = _check_mori_compatibility(common_kwargs) + + logger.info( + f"[MORI init] world_size={world_size} rank={rank} hidden={hidden} " + f"dtype={params_dtype} max_inp_tokens={num_max_dispatch_tokens_per_rank} " + f"num_local_experts={num_local_experts} num_topk={num_topk} " + f"kernel={kernel_type} block_num={block_num} " + f"warp_num_per_block={warp_num_per_block} rdma_block_num={rdma_block_num}", + rank=0, + ) + + config = mori.ops.EpDispatchCombineConfig(**common_kwargs) + return mori.ops.EpDispatchCombineOp(config) + + +class MoriEPBackend(_EPCapabilities): + """ROCm Mori EP backend for MoE dispatch/combine (normal mode).""" + + def __init__(self) -> None: + self._group: Optional[dist.ProcessGroup] = None + self._op = None + self._hidden_size: int = 0 + self._num_local_experts: int = 0 + self._num_topk: int = 0 + self._seqlen: int = 0 + self._fp8_dispatch: bool = False + self._params_dtype: Optional[torch.dtype] = None + # Last resolved kernel cfg; used to detect cfg-only rebuilds. + self._kernel_cfg: Optional[_MoriDispatchCfg] = None + + # ------------------------------------------------------------------ + # EPBackend protocol + # ------------------------------------------------------------------ + + def is_initialized(self) -> bool: + """Return True if the backend is initialized.""" + return self._op is not None + + @staticmethod + def is_available() -> bool: + try: + get_mori() + return True + except ImportError: + return False + + @staticmethod + def supports_cuda_graph() -> bool: + """not supported""" + return False + + @classmethod + def can_release(cls, *, will_reinit: bool) -> bool: + # inter-node free+rebuild hangs: keep if reinit'd, free a dead loser. + return not will_reinit + + @staticmethod + def _derive_params_dtype(fp8_dispatch: bool) -> torch.dtype: + """Pick the Mori ``data_type`` argument from ``fp8_dispatch``.""" + assert not fp8_dispatch, "Not implemented" + return torch.bfloat16 + + @call_once + def setup_env( + self, + *, + ib_gid_index: Optional[str] = None, + ib_hca: Optional[str] = None, + socket_ifname: Optional[str] = None, + ib_tc: Optional[str] = None, + ib_sl: Optional[str] = None, + ) -> None: + """Set MORI_* RDMA env from NCCL_* fallbacks; None kwargs stay unset (auto-detect).""" + _apply_env_with_nccl_fallback( + [ + ("MORI_IB_GID_INDEX", "NCCL_IB_GID_INDEX", ib_gid_index), + ("MORI_RDMA_DEVICES", "NCCL_IB_HCA", ib_hca), + ("MORI_SOCKET_IFNAME", "NCCL_SOCKET_IFNAME", socket_ifname), + ("MORI_RDMA_TC", "NCCL_IB_TC", ib_tc), + ("MORI_RDMA_SL", "NCCL_IB_SL", ib_sl), + ], + backend_name="MORI", + ) + + def init_buffer( + self, + group: dist.ProcessGroup, + hidden_size: int, + num_experts: int, + num_topk: int, + seqlen: int, + fp8_dispatch: bool, + config: EPBufferConfig, + ) -> None: + """Register Mori SHMEM and (re)build the op on shape/kernel-type change. + + block/warp/rdma are per-call, so the op is reused (not rebuilt) on cfg + change — the autotune winner reuses the tuned op (inter-node rebuild hangs). + """ + assert not fp8_dispatch, "not implemented" + + # Lazy mori import (install hint + version check). + get_mori() + + # Must run before any Mori JIT/SHMEM activity (fixes HIP error 500 on ROCm 7.2+). + _align_mori_hip_with_torch() + + # Apply MORI_* fallbacks from NCCL_* before SHMEM init. + self.setup_env() + + if self._group is not group: + self._group = group + _register_and_init_mori_shmem(group) + logger.info("[MoriEPBackend init] Mori SHMEM initialized.", rank=0) + + num_local_experts = num_experts // group.size() + params_dtype = self._derive_params_dtype(fp8_dispatch) + kernel_cfg = _resolve_mori_dispatch_cfg(group, config) + + # Rebuild on shape/kernel-type change only; block/warp/rdma are per-call. + needs_rebuild = ( + self._op is None + or self._hidden_size != hidden_size + or self._num_local_experts != num_local_experts + or self._num_topk != num_topk + or self._seqlen != seqlen + or self._fp8_dispatch != fp8_dispatch + or self._params_dtype != params_dtype + or self._kernel_cfg is None + or self._kernel_cfg.kernel_type != kernel_cfg.kernel_type + ) + + self._hidden_size = hidden_size + self._num_local_experts = num_local_experts + self._num_topk = num_topk + self._seqlen = seqlen + self._fp8_dispatch = fp8_dispatch + self._params_dtype = params_dtype + self._kernel_cfg = kernel_cfg + + if needs_rebuild: + self._op = _build_mori_op( + rank=group.rank(), + world_size=group.size(), + hidden=hidden_size, + params_dtype=params_dtype, + num_local_experts=num_local_experts, + num_topk=num_topk, + num_max_dispatch_tokens_per_rank=seqlen, + kernel_type=kernel_cfg.kernel_type, + warp_num_per_block=kernel_cfg.warp_num_per_block, + block_num=kernel_cfg.block_num, + rdma_block_num=kernel_cfg.rdma_block_num, + ) + + def dispatch( + self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + handle: Optional[tuple] = None, + topk_idx: Optional[torch.Tensor] = None, + token_weights: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, # noqa: ARG002 - sized via init_buffer. + async_finish: bool = False, # noqa: ARG002 - Mori runs sync on caller stream. + allocate_on_comm_stream: bool = False, # noqa: ARG002 - Mori runs sync on caller stream. + num_worst_tokens: int = 0, + config: Optional[Any] = None, + ): + assert self.is_initialized(), "Backend is not initialized" + + # Graph-friendly caller: keep counts on device (no D2H sync). + keep_tokens_per_expert_on_device = num_worst_tokens > 0 + + if handle is None: + assert topk_idx is not None + topk_idx_i32 = topk_idx.to(torch.int32) + else: + assert topk_idx is None and token_weights is None + (topk_idx_i32, token_weights, _) = handle + + # Autotuner launch override; -1 keeps the op-level default. + block_num, warp_per_block, rdma_block_num = _extract_mori_launch_override(config) + + # FP8 scales arg; BF16-only path for now. + scale = None + recv_x, recv_topk_weights, _recv_x_scales, recv_topk_idx, total_recv = self._op.dispatch( + x, + token_weights, + scale, + topk_idx_i32, + block_num=block_num, + rdma_block_num=rdma_block_num, + warp_per_block=warp_per_block, + ) + expert_base = self._group.rank() * self._num_local_experts + num_recv_tokens_per_expert, deepep_like_recv_topk_idx = compute_expert_token_info( + recv_topk_idx, + self._num_local_experts, + expert_base=expert_base, + total_recv=total_recv, + ) + if not keep_tokens_per_expert_on_device: + num_recv_tokens_per_expert = num_recv_tokens_per_expert.tolist() + + # Keep token_weights in the handle so backward can reach it. + handle = (topk_idx_i32, token_weights, recv_topk_idx) + return (recv_x, deepep_like_recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert, handle) + + def combine( + self, + x: torch.Tensor, + handle: tuple, + topk_weights: Optional[torch.Tensor] = None, + async_finish: bool = False, # noqa: ARG002 - API compat only. + allocate_on_comm_stream: bool = False, # noqa: ARG002 - API compat only. + config: Optional[Any] = None, + ): + assert self.is_initialized(), "Backend is not initialized" + + (_, _, recv_topk_idx) = handle + block_num, warp_per_block, rdma_block_num = _extract_mori_launch_override(config) + combined_x, combined_topk_weights = self._op.combine( + x.contiguous(), + topk_weights, + recv_topk_idx, + block_num=block_num, + rdma_block_num=rdma_block_num, + warp_per_block=warp_per_block, + ) + return combined_x, combined_topk_weights + + @torch.no_grad() + def tune_configs( + self, + group: dist.ProcessGroup, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + num_experts: int, + *, + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + num_sms: int = 32, + num_tests: int = 20, + num_topk: Optional[int] = None, + uniform_dispatch: bool = True, + ) -> Tuple[_MoriDispatchCfg, _MoriDispatchCfg, float, float]: + """Sweep (block_num, warp, rdma) on one op via per-call overrides (mirrors mori's bench). + + Intra pins rdma=0; inter sweeps it. Best dispatch/combine tracked + independently. Returns (dispatch_cfg, combine_cfg, dispatch_s, combine_s). + """ + # Lazy mori import (install hint + version check). + get_mori() + + if isinstance(x, tuple): + raise NotImplementedError( + "MoriEPBackend.tune_configs: FP8 dispatch (tuple ``x``) is not supported." + ) + + # Topology-derived kernel pinning: matches the live ``init_buffer`` path. + _, num_nodes = detect_group_topology(group) + mode = _MoriEpMode.INTRA_NODE if num_nodes <= 1 else _MoriEpMode.INTER_NODE + base_cfg = _get_mori_dispatch_configs()[mode] + kernel_type = base_cfg.kernel_type + + x_tensor = x + hidden_size = int(x_tensor.size(1)) + seqlen = int(x_tensor.size(0)) + device = x_tensor.device + + if uniform_dispatch: + if num_topk is None: + if topk_idx is None: + raise ValueError( + "tune_configs(uniform_dispatch=True): need either " + "``num_topk`` or a shape-carrying ``topk_idx`` to " + "sample a uniform distribution." + ) + num_topk = int(topk_idx.size(1)) + if num_topk > num_experts: + raise ValueError( + "tune_configs(uniform_dispatch=True): ``num_topk`` " + f"({num_topk}) cannot exceed ``num_experts`` ({num_experts})." + ) + num_tokens = int(x_tensor.size(0)) + scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device=device).abs() + 1 + topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)[1] + topk_weights = torch.randn((num_tokens, num_topk), dtype=torch.float32, device=device) + else: + if topk_idx is None or topk_weights is None: + raise ValueError( + "tune_configs(uniform_dispatch=False): ``topk_idx`` and " + "``topk_weights`` are both required." + ) + if num_topk is None: + num_topk = int(topk_idx.size(1)) + + topk_idx_i32 = topk_idx.to(torch.int32) + topk_weights_f = topk_weights.float() if topk_weights is not None else None + + sm_count = torch.cuda.get_device_properties(device).multi_processor_count + + if _mori_auto_mode(): + # mori self-tunes from its JSON per-call; just measure one representative config. + warp_list = [base_cfg.warp_num_per_block] + block_list = [base_cfg.block_num] + candidates = [(base_cfg.block_num, base_cfg.warp_num_per_block, base_cfg.rdma_block_num)] + else: + block_list = _mori_block_num_candidates(num_sms, sm_count) + # Intra pins rdma=0; inter sweeps rdma ~ {block/2, block*2/3} (fewer warps to bound cost). + if num_nodes <= 1: + warp_list = [4, 5, 6, 8, 10, 12, 14, 15, 16] + + def _rdma_candidates(block_num: int) -> List[int]: + return [0] + + else: + warp_list = [4, 8, 16] + + def _rdma_candidates(block_num: int) -> List[int]: + cands = sorted( + {v for v in (max(block_num // 2, 1), block_num * 2 // 3) if 1 <= v < block_num} + ) + return cands or [max(block_num // 2, 1)] + + candidates = [ + (block_num, warp_per_block, rdma_block_num) + for block_num in block_list + for warp_per_block in warp_list + for rdma_block_num in _rdma_candidates(block_num) + ] + + # Size the op for the largest candidate so per-call overrides fit. + worst_cfg = _MoriDispatchCfg( + kernel_type=kernel_type, + warp_num_per_block=max(warp_list), + block_num=max(block_list), + rdma_block_num=base_cfg.rdma_block_num, + ) + self.init_buffer( + group, + hidden_size=hidden_size, + num_experts=num_experts, + num_topk=num_topk, + seqlen=seqlen, + fp8_dispatch=False, + config=EPBufferConfig( + num_sms=num_sms, + dispatch_config=worst_cfg, + combine_config=worst_cfg, + ), + ) + + num_warmup = max(1, num_tests // 5) + is_rank0 = group.rank() == 0 + if is_rank0: + mode = "AUTO (mori tuning DB)" if _mori_auto_mode() else "MANUAL sweep" + logger.info( + f"[Mori tuning] mode={mode} sm_count={sm_count} kernel={kernel_type} " + f"block candidates={block_list} warp candidates={warp_list} " + f"rdma swept={num_nodes > 1 and not _mori_auto_mode()} total={len(candidates)}", + rank=0, + ) + + # Paired dispatch->combine sweep; _combine matches the live contiguous() prep. + def _dispatch(cand): + block_num, warp_per_block, rdma_block_num = cand + recv_x, recv_weights, _, recv_idx, _ = self._op.dispatch( + x_tensor, + topk_weights_f, + None, + topk_idx_i32, + block_num=block_num, + rdma_block_num=rdma_block_num, + warp_per_block=warp_per_block, + ) + return (recv_x, recv_weights, recv_idx, block_num, warp_per_block, rdma_block_num) + + def _combine(cand, recv): + recv_x, recv_weights, recv_idx, block_num, warp_per_block, rdma_block_num = recv + self._op.combine( + recv_x.contiguous(), + recv_weights, + recv_idx, + block_num=block_num, + rdma_block_num=rdma_block_num, + warp_per_block=warp_per_block, + ) + + def _on_skip(cand, exc): + if is_rank0: + logger.debug(f"[Mori tuning] skip invalid cfg {cand}: {exc!r}") + + best_d, best_c, best_d_time, best_c_time = sweep_configs( + candidates, + _dispatch, + _combine, + group=group, + num_tests=num_tests, + num_warmup=num_warmup, + on_skip=_on_skip, + ) + if best_d is None or best_d_time == float("inf"): + raise RuntimeError("MoriEPBackend.tune_configs: no valid dispatch config in sweep.") + if best_c is None or best_c_time == float("inf"): + raise RuntimeError("MoriEPBackend.tune_configs: no valid combine config in sweep.") + + # Broadcast rank 0's pick so every rank agrees on the winners. + disp_winner = _broadcast_from_rank0_int(list(best_d), group) # [block, warp, rdma] + comb_winner = _broadcast_from_rank0_int(list(best_c), group) + best_d_time = _broadcast_from_rank0_float(best_d_time, group) + best_c_time = _broadcast_from_rank0_float(best_c_time, group) + + final_d = _MoriDispatchCfg( + kernel_type=kernel_type, + warp_num_per_block=int(disp_winner[1]), + block_num=int(disp_winner[0]), + rdma_block_num=int(disp_winner[2]), + ) + final_c = _MoriDispatchCfg( + kernel_type=kernel_type, + warp_num_per_block=int(comb_winner[1]), + block_num=int(comb_winner[0]), + rdma_block_num=int(comb_winner[2]), + ) + + if is_rank0: + logger.info( + f"[Mori tuning] best dispatch block_num={final_d.block_num} " + f"warp={final_d.warp_num_per_block} rdma={final_d.rdma_block_num} " + f"t={best_d_time * 1e6:.2f} us; best combine block_num={final_c.block_num} " + f"warp={final_c.warp_num_per_block} rdma={final_c.rdma_block_num} " + f"t={best_c_time * 1e6:.2f} us", + rank=0, + ) + + return final_d, final_c, best_d_time, best_c_time + + def release_buffer(self) -> None: + """Drop the op and free its Mori SHMEM (keeps the global SHMEM binding). + + cache_clear() is required: the lru_cache ref otherwise pins the op so its + dtor (ShmemFree) never runs and autotune leaks the tuning op. + """ + # Drain in-flight kernels touching the buffers before they are freed. + if torch.cuda.is_available() and torch.cuda.is_initialized(): + torch.cuda.synchronize() + self._op = None + self._hidden_size = 0 + self._num_local_experts = 0 + self._num_topk = 0 + self._seqlen = 0 + self._fp8_dispatch = False + self._params_dtype = None + self._kernel_cfg = None + # Drop the cache's ref so the handle dtor runs (ShmemFree). + _build_mori_op.cache_clear() + gc.collect() diff --git a/primus_turbo/pytorch/kernels/moe/_backend/nic_detect.py b/primus_turbo/pytorch/kernels/moe/_backend/nic_detect.py new file mode 100644 index 000000000..53841b8dc --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/nic_detect.py @@ -0,0 +1,81 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +from __future__ import annotations + +import os +from typing import List, Optional, Tuple + +__all__ = ["detect_nic_type"] + +# Ordered ``(device_name_prefix, canonical_family)`` map. First match wins, +# so put more specific prefixes first if a generic one could shadow them. +_NIC_NAME_PREFIXES: Tuple[Tuple[str, str], ...] = ( + ("ionic", "ionic"), # AMD Pollara AINIC + ("bnxt_re", "bnxt"), # Broadcom Thor / Thor-2 + ("mlx5", "mlx5"), # Mellanox ConnectX +) + +_SYSFS_IB_PATH = "/sys/class/infiniband" + + +def _classify(name: str) -> Optional[str]: + """Map an IB device name to its canonical NIC family, or ``None``.""" + for prefix, family in _NIC_NAME_PREFIXES: + if name.startswith(prefix): + return family + return None + + +def _read_hca_env() -> List[str]: + """Parse ``UCCL_IB_HCA`` / ``NCCL_IB_HCA`` into bare device-name tokens.""" + value = ( + os.environ.get("NCCL_IB_HCA") + or os.environ.get("UCCL_IB_HCA") + or os.environ.get("MORI_RDMA_DEVICES") + or "" + ) + seen: set = set() + tokens: List[str] = [] + for raw in value.split(","): + tok = raw.strip().lstrip("^").split(":", 1)[0] + if tok and tok not in seen: + seen.add(tok) + tokens.append(tok) + return tokens + + +def _list_local_ib_devices() -> List[str]: + """Return sorted IB device names from sysfs, or ``[]`` if unreadable.""" + try: + return sorted(os.listdir(_SYSFS_IB_PATH)) + except OSError: + return [] + + +def detect_nic_type() -> Optional[str]: + """Return canonical NIC family (``"ionic"`` / ``"bnxt"`` / ``"mlx5"``) or ``None``. + + Resolution order: + 1. Enumerate ``/sys/class/infiniband``; if the HCA env is set, prefer + devices it matches (NCCL substring rule). Fall through to the + unfiltered list when nothing matches so we still classify + *something* on mixed-NIC hosts. + 2. If sysfs is empty/unreadable, classify the env tokens directly. + """ + env_tokens = _read_hca_env() + local = _list_local_ib_devices() + + if local and env_tokens: + preferred = [n for n in local if any(t in n for t in env_tokens)] + candidates = preferred or local + else: + candidates = local or env_tokens + + for name in candidates: + family = _classify(name) + if family is not None: + return family + return None diff --git a/primus_turbo/pytorch/kernels/moe/_backend/turbo.py b/primus_turbo/pytorch/kernels/moe/_backend/turbo.py new file mode 100644 index 000000000..70196b18e --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/turbo.py @@ -0,0 +1,30 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""In-tree Primus-Turbo DeepEP backend.""" + +import torch.distributed as dist + +from .base import _DeepEPLikeBackend, detect_group_topology + + +class TurboEPBackend(_DeepEPLikeBackend): + """In-tree Primus-Turbo DeepEP backend (always available).""" + + @staticmethod + def is_available() -> bool: + return True + + @classmethod + def is_feasible(cls, group: dist.ProcessGroup) -> bool: + # rocSHMEM internode is unsupported on AINIC (SIGABRTs); intra-node only. + _, num_nodes = detect_group_topology(group) + return num_nodes <= 1 + + @staticmethod + def _get_module(): + import primus_turbo.pytorch.deep_ep as turbo_ep + + return turbo_ep diff --git a/primus_turbo/pytorch/kernels/moe/_backend/uccl_ep.py b/primus_turbo/pytorch/kernels/moe/_backend/uccl_ep.py new file mode 100644 index 000000000..7cec54d70 --- /dev/null +++ b/primus_turbo/pytorch/kernels/moe/_backend/uccl_ep.py @@ -0,0 +1,43 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""UCCL-EP backend (optional).""" + +from typing import Optional + +from primus_turbo.common.uccl_utils import get_uccl + +from .base import _DeepEPLikeBackend, apply_uccl_network_env, call_once + + +class UCCLEPBackend(_DeepEPLikeBackend): + """UCCL-EP backend via uccl's deep_ep-compatible wrapper ``Buffer``.""" + + @staticmethod + def is_available() -> bool: + # Needs uccl and its deep_ep wrapper (raw uccl.ep lacks the high-level API). + try: + import deep_ep # noqa: F401 + import uccl.ep # noqa: F401 + except ImportError: + return False + # Only the uccl-backed wrapper (Config re-exported from uccl.ep) is ours. + return getattr(getattr(deep_ep, "Config", None), "__module__", "").startswith("uccl") + + @classmethod + def can_release(cls, *, will_reinit: bool) -> bool: # noqa: ARG003 + # destroy() corrupts the HIP context: never release. + return False + + @staticmethod + def _get_module(): + get_uccl() + import deep_ep + + return deep_ep + + @call_once + def setup_env(self, **overrides: Optional[str]) -> None: + apply_uccl_network_env(**overrides) diff --git a/primus_turbo/pytorch/kernels/moe/moe_dispatch_combine_impl.py b/primus_turbo/pytorch/kernels/moe/moe_dispatch_combine_impl.py index 1faa7d8fa..512c61801 100644 --- a/primus_turbo/pytorch/kernels/moe/moe_dispatch_combine_impl.py +++ b/primus_turbo/pytorch/kernels/moe/moe_dispatch_combine_impl.py @@ -5,30 +5,31 @@ # See LICENSE for license information. ############################################################################### -import inspect -import os +import dataclasses +import time import warnings -from dataclasses import dataclass -from typing import ( - Any, - Dict, - List, - Optional, - Protocol, - Tuple, - Type, - Union, - runtime_checkable, -) +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, Union import torch import torch.distributed as dist -from primus_turbo.common.constants import ENV_MOE_DISPATCH_COMBINE_BACKEND +from primus_turbo.common.logger import logger from primus_turbo.pytorch.core.backend import ( BackendType, GlobalBackendManager, PrecisionType, + TuneCache, +) +from primus_turbo.pytorch.kernels.moe._backend import ( + DeepEPBackend, + EPBackend, + EPBufferConfig, + MoriEPBackend, + TurboEPBackend, + UCCLEPBackend, + detect_group_topology, ) # ========================================================================= @@ -36,382 +37,96 @@ # ========================================================================= -@dataclass -class EPBufferConfig: - """Configuration for EP communication buffer initialization. - - Attributes: - num_sms: Number of SMs to use in high-throughput kernels. - dispatch_config: Optional user-provided dispatch config (from offline - benchmarking). When ``None``, the backend's default for the current - ``ep_size`` is used (``Buffer.get_dispatch_config(ep_size)``). - combine_config: Optional user-provided combine config. Same fallback - behaviour as *dispatch_config*. - """ +# Per-backend default buffer configuration. Keys must match the names used +# in ``_BACKEND_REGISTRY`` so lookups can be done by backend name. +_DEFAULT_BUFFER_CONFIG_PER_BACKEND: Dict[str, EPBufferConfig] = { + "TURBO": EPBufferConfig( + num_sms=64, + dispatch_config=None, + combine_config=None, + ), + "DEEP_EP": EPBufferConfig( + num_sms=64, + dispatch_config=None, + combine_config=None, + ), + "UCCL": EPBufferConfig( + num_sms=64, + dispatch_config=None, + combine_config=None, + ), + "MORI": EPBufferConfig(num_sms=64, dispatch_config=None, combine_config=None), +} - num_sms: int = 32 - dispatch_config: Any = None - combine_config: Any = None +# Deep copy of the defaults so mutations via ``set_buffer_global_config`` +# don't leak into ``_DEFAULT_BUFFER_CONFIG_PER_BACKEND``. +_buffer_config_per_backend: Dict[str, EPBufferConfig] = { + name: dataclasses.replace(cfg) for name, cfg in _DEFAULT_BUFFER_CONFIG_PER_BACKEND.items() +} -_DEFAULT_BUFFER_CONFIG = EPBufferConfig( - num_sms=32, - dispatch_config=None, - combine_config=None, -) -_buffer_config: EPBufferConfig = _DEFAULT_BUFFER_CONFIG +def _get_buffer_config(backend_name: str) -> EPBufferConfig: + """Return the runtime buffer config for ``backend_name``, or a default.""" + cfg = _buffer_config_per_backend.get(backend_name) + if cfg is not None: + return cfg + default = _DEFAULT_BUFFER_CONFIG_PER_BACKEND.get(backend_name) + return dataclasses.replace(default) if default is not None else EPBufferConfig() def set_buffer_global_config( num_use_cu: int = 32, autotune_config: Optional[tuple] = None, + backend: Optional[str] = None, ) -> None: - """Store the SM count and optional per-operation configs. - - This is typically called once by the token dispatcher during ``__init__``. + """Set the EP buffer config. - If a backend has already allocated a buffer against the previous config, - we *warn* on any actual change: the live buffer is sized to the old config - and may be referenced by a captured CUDA graph, so silently swapping in a - new ``num_sms`` would either re-trigger reallocation on the next + Args: + num_use_cu: SM count for high-throughput kernels. + autotune_config: Optional ``(dispatch_config, combine_config)`` from + offline tuning. + backend: Target backend name (matches ``_BACKEND_REGISTRY``). When + ``None``, the same config is broadcast to every registered + backend (backward-compat with the pre-multi-backend API). + + If a backend has already allocated a buffer against its previous config, we + *warn* on any actual change: the live buffer is sized to the old config and + may be referenced by a captured CUDA graph, so silently swapping in a new + ``num_sms`` would either re-trigger reallocation on the next ``_ensure_buffer`` (invalidating the captured pointer) or be ignored entirely. Equal configs are a no-op. - - Args: - num_use_cu: Number of SMs (compute units) for high-throughput kernels. - autotune_config: Legacy parameter — a ``(dispatch_config, combine_config)`` - tuple obtained from offline benchmarking. ``None`` means use the - backend's built-in defaults for the current EP group size. """ - global _buffer_config dispatch_cfg, combine_cfg = autotune_config if autotune_config is not None else (None, None) - new_config = EPBufferConfig( + new_cfg = EPBufferConfig( num_sms=num_use_cu, dispatch_config=dispatch_cfg, combine_config=combine_cfg, ) - if new_config != _buffer_config: - for name, backend in _backend_instances.items(): - if getattr(backend, "_buffer", None) is not None: - warnings.warn( - f"set_buffer_global_config called with a different config after " - f"backend '{name}' was already initialized " - f"(old={_buffer_config!r}, new={new_config!r}). Previously-captured " - "CUDA graphs continue to reference the old buffer; the new config " - "only takes effect when the buffer is next reallocated.", - stacklevel=2, - ) - break - _buffer_config = new_config - - -# ========================================================================= -# EPBackend Protocol -# ========================================================================= - - -@runtime_checkable -class EPBackend(Protocol): - """Structural (``typing.Protocol``) interface for Expert-Parallel - communication backends. - - Each backend encapsulates a specific EP library (e.g. in-tree Turbo DeepEP, - external ``deep_ep``, UCCL-EP, ...) and owns its own buffer lifecycle. - Adding a new backend is a single-class change plus one - ``register_ep_backend()`` call — any class that structurally conforms to - this Protocol is accepted; no explicit inheritance is required. - """ - - @staticmethod - def is_available() -> bool: - """Return True if this backend's dependencies are importable.""" - ... - - def init_buffer( - self, - group: dist.ProcessGroup, - hidden_bytes: int, - config: EPBufferConfig, - ) -> None: - """(Re-)create the communication buffer if needed.""" - ... - - def dispatch( - self, - x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], - handle: Optional[tuple] = None, - topk_idx: Optional[torch.Tensor] = None, - token_weights: Optional[torch.Tensor] = None, - num_experts: Optional[int] = None, - async_finish: bool = False, - allocate_on_comm_stream: bool = False, - num_worst_tokens: int = 0, - ) -> Tuple[ - Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[Union[List[int], torch.Tensor]], - Optional[tuple], - ]: - """Execute dispatch (layout + send) and return - ``(recv_x, recv_topk_idx, recv_topk_weights, tokens_per_expert, handle)``. - """ - ... - - def combine( - self, - x: torch.Tensor, - handle: tuple, - topk_weights: Optional[torch.Tensor] = None, - async_finish: bool = False, - allocate_on_comm_stream: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """Execute combine and return ``(combined_x, combined_topk_weights)``.""" - ... - - -# ========================================================================= -# _DeepEPLikeBackend — shared implementation for DeepEP-compatible backends -# ========================================================================= - - -class _DeepEPLikeBackend: - """Shared logic for all backends that follow the DeepEP Buffer protocol - (``get_dispatch_layout`` / ``dispatch`` / ``combine`` / ``set_num_sms`` / - ``get_dispatch_config`` / ``get_combine_config``). - - This is a plain implementation base class — it does **not** inherit from - ``EPBackend``. Conformance to the ``EPBackend`` Protocol is checked - structurally by the type system. - - Subclasses only need to override ``is_available``, ``_get_module``, and - optionally ``_make_buffer_kwargs`` to supply backend-specific constructor - arguments. - """ - - def __init__(self) -> None: - self._buffer = None - - # ------------------------------------------------------------------ - # Subclass hooks - # ------------------------------------------------------------------ - - @staticmethod - def is_available() -> bool: - """Return True if this backend's dependencies are importable.""" - raise NotImplementedError - - @staticmethod - def _get_module(): - """Return the Python module that exposes ``Buffer``, ``Config``, - ``EventHandle``, ``EventOverlap`` (or a compatible ``utils`` sub-module). - """ - raise NotImplementedError - - def _make_buffer_kwargs(self, group: dist.ProcessGroup) -> dict: - """Extra keyword arguments forwarded to ``BufferClass(group, nvl, rdma, **kwargs)``.""" - return {} - - # ------------------------------------------------------------------ - # EPBackend interface - # ------------------------------------------------------------------ - - def init_buffer( - self, - group: dist.ProcessGroup, - hidden_bytes: int, - config: EPBufferConfig, - ) -> None: - mod = self._get_module() - BufferClass = mod.Buffer - - BufferClass.set_num_sms(config.num_sms) - - dispatch_config = config.dispatch_config or BufferClass.get_dispatch_config(group.size()) - combine_config = config.combine_config or BufferClass.get_combine_config(group.size()) - - num_nvl_bytes, num_rdma_bytes = 0, 0 - for cfg in (dispatch_config, combine_config): - num_nvl_bytes = max( - cfg.get_nvl_buffer_size_hint(hidden_bytes, group.size()), - num_nvl_bytes, - ) - try: - num_rdma_bytes = max( - cfg.get_rdma_buffer_size_hint(hidden_bytes, group.size()), - num_rdma_bytes, - ) - except (RuntimeError, AttributeError): - pass - - buf_kwargs = self._make_buffer_kwargs(group) - - if ( - self._buffer is None - or not isinstance(self._buffer, BufferClass) - or self._buffer.group != group - or self._buffer.num_nvl_bytes < num_nvl_bytes - or self._buffer.num_rdma_bytes < num_rdma_bytes - ): - self._buffer = BufferClass(group, num_nvl_bytes, num_rdma_bytes, **buf_kwargs) - - # ----- helpers ------------------------------------------------------ - - def _get_event_classes(self): - mod = self._get_module() - EventOverlapClass = mod.utils.EventOverlap if hasattr(mod, "utils") else mod.EventOverlap - EventHandleClass = mod.utils.EventHandle if hasattr(mod, "utils") else mod.EventHandle - return EventOverlapClass, EventHandleClass - - # ----- dispatch / combine ------------------------------------------- - - def dispatch( - self, - x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], - handle: Optional[tuple] = None, - topk_idx: Optional[torch.Tensor] = None, - token_weights: Optional[torch.Tensor] = None, - num_experts: Optional[int] = None, - async_finish: bool = False, - allocate_on_comm_stream: bool = False, - num_worst_tokens: int = 0, - ): - EventOverlapClass, EventHandleClass = self._get_event_classes() - buffer = self._buffer - assert buffer is not None, "init_buffer() must be called before dispatch()" - - previous_event = EventOverlapClass(EventHandleClass()) if async_finish else None - if handle is None: - assert topk_idx is not None - assert token_weights is not None - ( - num_tokens_per_rank, - num_tokens_per_rdma_rank, - num_tokens_per_expert, - is_token_in_rank, - event, - ) = buffer.get_dispatch_layout( - topk_idx, - num_experts, - previous_event=previous_event, - async_finish=async_finish, - allocate_on_comm_stream=allocate_on_comm_stream, - ) - - ( - recv_x, - recv_token_indices, - recv_token_probs, - tokens_per_expert, - handle, - after_event, - ) = buffer.dispatch( - x, - topk_idx=topk_idx, - topk_weights=token_weights, - num_tokens_per_rank=num_tokens_per_rank, - num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, - is_token_in_rank=is_token_in_rank, - num_tokens_per_expert=num_tokens_per_expert, - previous_event=event, - async_finish=async_finish, - allocate_on_comm_stream=allocate_on_comm_stream, - num_worst_tokens=num_worst_tokens, + def _warn_if_live_buffer(name: str) -> None: + old_cfg = _buffer_config_per_backend.get(name) + if old_cfg is None or new_cfg == old_cfg: + return + inst = _backend_instances.get(name) + if inst is not None and getattr(inst, "_buffer", None) is not None: + warnings.warn( + f"set_buffer_global_config called with a different config after " + f"backend '{name}' was already initialized " + f"(old={old_cfg!r}, new={new_cfg!r}). Previously-captured CUDA " + "graphs continue to reference the old buffer; the new config only " + "takes effect when the buffer is next reallocated.", + stacklevel=2, ) - else: - recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle, after_event = ( - buffer.dispatch( - x, - handle=handle, - previous_event=previous_event, - async_finish=async_finish, - allocate_on_comm_stream=allocate_on_comm_stream, - ) - ) - - if async_finish: - after_event.current_stream_wait() - - return recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle - - def combine( - self, - x: torch.Tensor, - handle: tuple, - topk_weights: Optional[torch.Tensor] = None, - async_finish: bool = False, - allocate_on_comm_stream: bool = False, - ): - EventOverlapClass, EventHandleClass = self._get_event_classes() - buffer = self._buffer - assert buffer is not None, "init_buffer() must be called before combine()" - - previous_event = EventOverlapClass(EventHandleClass()) if async_finish else None - - combined_x, combined_topk_weights, after_event = buffer.combine( - x, - handle=handle, - topk_weights=None if topk_weights is None else topk_weights.float(), - async_finish=async_finish, - allocate_on_comm_stream=allocate_on_comm_stream, - previous_event=previous_event, - ) - - if async_finish: - after_event.current_stream_wait() - - return combined_x, combined_topk_weights - - -# ========================================================================= -# Concrete backends -# ========================================================================= - - -class TurboEPBackend(_DeepEPLikeBackend): - """In-tree Primus-Turbo DeepEP backend (always available).""" - - @staticmethod - def is_available() -> bool: - return True - - @staticmethod - def _get_module(): - import primus_turbo.pytorch.deep_ep as turbo_ep - - return turbo_ep - -class DeepEPBackend(_DeepEPLikeBackend): - """External ``deep_ep`` package backend (optional).""" - - @staticmethod - def is_available() -> bool: - try: - import deep_ep # noqa: F401 - - return True - except ImportError: - return False - - @staticmethod - def _get_module(): - import deep_ep - - return deep_ep - - def _make_buffer_kwargs(self, group: dist.ProcessGroup) -> dict: - BufferClass = self._get_module().Buffer - try: - param = inspect.signature(BufferClass).parameters.get("is_intranode") - except (TypeError, ValueError): - param = None - if param is not None and param.default is False: - # uccl-ep special handle - return {"is_intranode": group.size() <= 8} - return {} + if backend is None: + targets = set(_buffer_config_per_backend.keys()) | set(_BACKEND_REGISTRY.keys()) + for name in targets: + _warn_if_live_buffer(name) + _buffer_config_per_backend[name] = dataclasses.replace(new_cfg) + else: + _warn_if_live_buffer(backend) + _buffer_config_per_backend[backend] = new_cfg # ========================================================================= @@ -421,6 +136,8 @@ def _make_buffer_kwargs(self, group: dist.ProcessGroup) -> dict: _BACKEND_REGISTRY: Dict[str, Type[EPBackend]] = { "TURBO": TurboEPBackend, "DEEP_EP": DeepEPBackend, + "MORI": MoriEPBackend, + "UCCL": UCCLEPBackend, } _backend_instances: Dict[str, EPBackend] = {} @@ -447,7 +164,7 @@ def clear_backend_instances(): def register_ep_backend(name: str, cls: Type[EPBackend]) -> None: - """Register a new EP backend class (e.g. ``UCCL_EP``).""" + """Register a new EP backend class (e.g. ``UCCL``).""" _BACKEND_REGISTRY[name] = cls @@ -466,33 +183,449 @@ def _get_backend_instance(name: str) -> EPBackend: return _backend_instances[name] +def _release_other_backends(active_name: str) -> None: + """Release other EP backends' buffers (one live backend per process; skips uccl).""" + for name, inst in _backend_instances.items(): + if name == active_name or not inst.can_release(will_reinit=True): + continue + if not inst.is_initialized(): + continue + try: + inst.release_buffer() + logger.info( + f"EP backend switch: released '{name}' buffer before using " + f"'{active_name}' (one live EP backend per process).", + rank=0, + ) + except Exception as exc: # noqa: BLE001 - best-effort teardown + logger.warning(f"EP backend switch: release_buffer('{name}') failed: {exc!r}") + + +# ========================================================================= +# Autotuner — finds best (backend, dispatch_config, combine_config) per case +# ========================================================================= + + +@dataclass(frozen=True) +class _EPAutoTuneKey: + """Shape-based signature used to cache autotune results.""" + + num_tokens: int + hidden: int + num_topk: int + num_experts: int + ep_size: int + dtype: torch.dtype + use_fp8: bool + + +@dataclass +class EPAutoTuneResult: + """Result of :class:`MoEDispatchCombineAutoTuner.tune`. + + Attributes: + backend_name: Winning backend registry name. + dispatch_config: Best dispatch ``Config``. + combine_config: Best combine ``Config``. + dispatch_time_us: Dispatch latency of the winner (us). + combine_time_us: Combine latency of the winner (us). + per_backend: ``backend_name -> (d_time_us, c_time_us)`` of every + candidate that completed tuning. + tune_duration_s: Wall-clock seconds spent tuning (rank 0 clock). + """ + + backend_name: str + dispatch_config: Any + combine_config: Any + dispatch_time_us: float + combine_time_us: float + per_backend: Dict[str, Tuple[float, float]] = field(default_factory=dict) + tune_duration_s: float = 0.0 + + +class MoEDispatchCombineAutoTuner: + """Picks the best ``(backend, dispatch_config, combine_config)`` per shape. + + The shape-keyed result is bound to the dispatch ``handle`` so the paired + combine / cached dispatch / backward reuse it. Enable via + ``PRIMUS_TURBO_AUTO_TUNE=1`` or :meth:`get_or_tune`. + """ + + # Shape-based cache shared across layers with identical input shape. + _cache: TuneCache = TuneCache(capacity=1024) + + # ``handle -> tune result`` LRU. Handles are tuples (not weakref-able), so + # we keep a strong ref alongside the result and rely on ``stored_handle + # is handle`` to defend against ``id()`` reuse. Cap is small because a + # dispatch/combine pair turns over quickly; the LRU evicts stale entries + # before they pile up. ``clear()`` empties this cache as well. + _HANDLE_CACHE_MAX: int = 128 + _handle_cache: "OrderedDict[int, Tuple[Any, EPAutoTuneResult]]" = OrderedDict() + + # Fallback for calls made before any handle mapping is registered. + _current_result: Optional[EPAutoTuneResult] = None + + @staticmethod + def make_key( + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: Optional[torch.Tensor], + num_experts: int, + ep_size: int, + num_topk: Optional[int] = None, + ) -> _EPAutoTuneKey: + inp = x if isinstance(x, torch.Tensor) else x[0] + num_tokens, hidden = int(inp.size(0)), int(inp.size(1)) + if num_topk is not None: + resolved_num_topk = int(num_topk) + elif topk_idx is not None: + resolved_num_topk = int(topk_idx.size(1)) + else: + resolved_num_topk = 0 + return _EPAutoTuneKey( + num_tokens=num_tokens, + hidden=hidden, + num_topk=resolved_num_topk, + num_experts=int(num_experts), + ep_size=int(ep_size), + dtype=inp.dtype, + use_fp8=isinstance(x, tuple), + ) + + # ------------------------------------------------------------------ + # Handle <-> result binding (dispatch/combine pair) + # ------------------------------------------------------------------ + + @classmethod + def register_handle(cls, handle: Any, result: EPAutoTuneResult) -> None: + """Bind ``result`` to ``handle`` for paired-call reuse. + + Args: + handle: Handle returned by the backend's ``dispatch`` call. + result: Tune result to associate with the handle. + """ + if handle is None: + return + hid = id(handle) + cache = cls._handle_cache + if hid in cache: + cache.move_to_end(hid) + cache[hid] = (handle, result) + if len(cache) > cls._HANDLE_CACHE_MAX: + cache.popitem(last=False) + + @classmethod + def lookup_handle(cls, handle: Any) -> Optional[EPAutoTuneResult]: + """Return the result bound to ``handle``, or ``None`` if unknown. + + ``stored_handle is handle`` guards against ``id()`` collisions after + the original handle was GC'd and its id was reused by a new one. + """ + if handle is None: + return None + entry = cls._handle_cache.get(id(handle)) + if entry is None: + return None + stored_handle, res = entry + if stored_handle is not handle: + cls._handle_cache.pop(id(handle), None) + return None + cls._handle_cache.move_to_end(id(handle)) + return res + + @classmethod + def discard_handle(cls, handle: Any) -> None: + """Drop the cache entry for ``handle`` once the caller is done with it. + + Call this after the paired combine (and any backward that reuses the + handle) so long-running jobs don't pin tune-result tensors via the + LRU. + """ + if handle is None: + return + cls._handle_cache.pop(id(handle), None) + + @classmethod + def handle_cache_size(cls) -> int: + """Return current entry count in the handle cache (tests / metrics).""" + return len(cls._handle_cache) + + @classmethod + def _candidate_backend_names(cls) -> List[str]: + names: List[str] = [] + for name, impl in _BACKEND_REGISTRY.items(): + try: + if impl.is_available(): + names.append(name) + except Exception: + continue + return names + + @classmethod + @torch.no_grad() + def tune( + cls, + group: dist.ProcessGroup, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + num_experts: int, + *, + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + num_sms: int = 32, + candidate_backends: Optional[Sequence[str]] = None, + num_tests: int = 20, + num_topk: Optional[int] = None, + uniform_dispatch: bool = True, + ) -> EPAutoTuneResult: + """Tune across candidate backends + configs and return the best. + + Args: + group: EP process group. + x: Dispatch input (Tensor or ``(fp8_tensor, scales)`` tuple). + num_experts: Total number of experts. + topk_idx: ``[num_tokens, num_topk]`` expert indices (optional in + uniform mode). + topk_weights: ``[num_tokens, num_topk]`` expert weights (optional + in uniform mode). + num_sms: SM count to use for tuning. + candidate_backends: Optional explicit backend name list. + num_tests: Timed iterations per candidate. + num_topk: Number of experts per token (uniform-mode fallback). + uniform_dispatch: If True, resample topk internally for robustness. + + Returns: + Best :class:`EPAutoTuneResult`. Not cached; use :meth:`get_or_tune`. + """ + names = list(candidate_backends) if candidate_backends is not None else cls._candidate_backend_names() + if not names: + raise RuntimeError("MoE autotune: no EP backends are available.") + + # Drop backends infeasible on this topology (e.g. TURBO inter-node SIGABRTs). + feasible: List[str] = [] + for name in names: + impl = _BACKEND_REGISTRY.get(name) + if impl is not None and not impl.is_feasible(group): + logger.info( + f"MoE autotune: skipping backend '{name}'; not feasible on this topology.", + once=True, + rank=0, + ) + continue + feasible.append(name) + names = feasible + if not names: + raise RuntimeError("MoE autotune: no feasible EP backend for this topology.") + + tune_start = time.perf_counter() + + # Detect multi-node topology once; reused to guard backend candidates + # whose worst-case buffer sizing is unsafe in inter-node EP. + _, num_nodes = detect_group_topology(group) + + best: Optional[EPAutoTuneResult] = None + best_total = float("inf") + per_backend: Dict[str, Tuple[float, float]] = {} + + for name in names: + if name not in _BACKEND_REGISTRY: + logger.info( + f"MoE autotune: skipping backend '{name}'; not in " + f"backend registry ({list(_BACKEND_REGISTRY.keys())}).", + once=True, + rank=0, + ) + continue + try: + available = _BACKEND_REGISTRY[name].is_available() + except Exception as exc: + logger.info( + f"MoE autotune: skipping backend '{name}'; " + f"is_available() raised {type(exc).__name__}: {exc!r}.", + once=True, + rank=0, + ) + continue + if not available: + logger.info( + f"MoE autotune: skipping backend '{name}'; " + f"is_available()=False (dependencies missing or backend " + f"disabled).", + once=True, + rank=0, + ) + continue + backend = _get_backend_instance(name) + _release_other_backends(name) + try: + d_cfg, c_cfg, d_t, c_t = backend.tune_configs( + group=group, + x=x, + num_experts=num_experts, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_sms=num_sms, + num_tests=num_tests, + num_topk=num_topk, + uniform_dispatch=uniform_dispatch, + ) + except NotImplementedError: + # Backend opts out of config tuning. Skip quietly. + logger.info( + f"MoE autotune: backend '{name}' does not support " f"tune_configs; skipping.", + once=True, + rank=0, + ) + continue + except Exception as exc: # pragma: no cover - defensive + # repr() so pybind-backed errors aren't silently blank. + logger.warning( + f"MoE autotune: backend '{name}' failed: \n" f"{type(exc).__name__}: {exc!r}", + exc_info=True, + ) + continue + finally: + # Free the tuning buffer (winner will re-init); uccl/mori keep it for reuse. + if backend.can_release(will_reinit=True): + try: + backend.release_buffer() + except Exception as exc: # noqa: BLE001 - best-effort cleanup + logger.debug(f"MoE autotune: release_buffer('{name}') failed: {exc!r}") + + per_backend[name] = (round(d_t * 1e6, 3), round(c_t * 1e6, 3)) + total = d_t + c_t + if total < best_total: + best_total = total + best = EPAutoTuneResult( + backend_name=name, + dispatch_config=d_cfg, + combine_config=c_cfg, + dispatch_time_us=round(d_t * 1e6, 3), + combine_time_us=round(c_t * 1e6, 3), + ) + + if best is None: + raise RuntimeError( + f"MoE autotune: all candidate backends failed ({names}). " + f"Check logs and EP buffer configuration." + ) + + # Free non-winner buffers so they don't waste HBM (skips uccl, can't free). + for name, inst in list(_backend_instances.items()): + if name == best.backend_name: + continue + if not inst.is_initialized(): + continue + if not inst.can_release(will_reinit=False): # dead loser, never re-init'd + continue + try: + inst.release_buffer() + logger.info( + f"MoE autotune: freed non-winner '{name}' buffer " f"(winner={best.backend_name}).", + rank=0, + ) + except Exception as exc: # noqa: BLE001 - best-effort cleanup + logger.debug(f"MoE autotune: post-sweep release_buffer('{name}') failed: {exc!r}") + + best.per_backend = per_backend + best.tune_duration_s = time.perf_counter() - tune_start + return best + + @classmethod + @torch.no_grad() + def get_or_tune( + cls, + group: dist.ProcessGroup, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + num_experts: int, + *, + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + num_sms: int = 32, + num_topk: Optional[int] = None, + uniform_dispatch: bool = True, + **tune_kwargs: Any, + ) -> EPAutoTuneResult: + """Return a cached tune result for this shape or run :meth:`tune` now. + + Args: + group: EP process group. + x: Dispatch input (Tensor or ``(fp8_tensor, scales)`` tuple). + num_experts: Total number of experts. + topk_idx: Optional ``[num_tokens, num_topk]`` expert indices. + topk_weights: Optional ``[num_tokens, num_topk]`` expert weights. + num_sms: SM count to use for tuning. + num_topk: Number of experts per token. Also used in the shape + cache key, so the same signature reuses the cached winner. + uniform_dispatch: If True, resample topk internally for robustness. + **tune_kwargs: Forwarded to :meth:`tune`. + """ + key = cls.make_key(x, topk_idx, num_experts, group.size(), num_topk=num_topk) + cached = cls._cache.get(key) + if cached is not None: + cls._current_result = cached + return cached + + result = cls.tune( + group=group, + x=x, + num_experts=num_experts, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_sms=num_sms, + num_topk=num_topk, + uniform_dispatch=uniform_dispatch, + **tune_kwargs, + ) + cls._cache.put(key, result) + cls._current_result = result + logger.info( + f"[MoE AutoTune] key={key} -> backend={result.backend_name} " + f"dispatch={result.dispatch_time_us:.1f}us combine={result.combine_time_us:.1f}us " + f"tune_time={result.tune_duration_s:.2f}s " + f"per_backend={result.per_backend}", + rank=0, + ) + return result + + @classmethod + def current(cls) -> Optional[EPAutoTuneResult]: + """Return the most recent tune result, or ``None`` if never tuned.""" + return cls._current_result + + @classmethod + def set_current(cls, result: Optional[EPAutoTuneResult]) -> None: + """Override the current runtime result (mainly for testing).""" + cls._current_result = result + + @classmethod + def clear(cls) -> None: + """Drop every cache (shape, handle, current).""" + cls._cache.clear() + cls._handle_cache.clear() + cls._current_result = None + + # ========================================================================= # Backend selection # ========================================================================= +_DEFAULT_BACKEND_NAME = "TURBO" + +# Explicit ``BackendType -> registry name`` map so renaming the enum doesn't +# silently break EP backend selection. Keys must mirror ``_BACKEND_REGISTRY``. _BACKEND_TYPE_TO_NAME: Dict[BackendType, str] = { BackendType.TURBO: "TURBO", BackendType.DEEP_EP: "DEEP_EP", + BackendType.MORI: "MORI", + BackendType.UCCL: "UCCL", } -def _resolve_backend_name() -> str: - """Determine which EP backend to use. - - Priority (high -> low): - 1. ``GlobalBackendManager`` code-level setting (via ``set_moe_dispatch_combine_backend``) - 2. ``PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND`` env var (supports names beyond ``BackendType``) - 3. Default: ``TURBO`` - """ - user_backend = GlobalBackendManager.get_moe_dispatch_combine_backend(PrecisionType.BF16_FP16_FP32) - if user_backend is not None: - return _BACKEND_TYPE_TO_NAME.get(user_backend, user_backend.name) - - env_val = os.environ.get(ENV_MOE_DISPATCH_COMBINE_BACKEND) - if env_val is not None: - return env_val.strip().upper() - - return "TURBO" +def _get_backend_name() -> str: + """Return the user-selected backend name, or ``TURBO`` by default.""" + bt = GlobalBackendManager.get_ep_backend(PrecisionType.BF16_FP16_FP32) + if bt is None: + return _DEFAULT_BACKEND_NAME + return _BACKEND_TYPE_TO_NAME.get(bt, _DEFAULT_BACKEND_NAME) # ========================================================================= @@ -500,37 +633,92 @@ def _resolve_backend_name() -> str: # ========================================================================= -def get_hidden_bytes( +def _maybe_tune_config( x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], -) -> int: - """Calculate the number of hidden bytes for a tensor. + group: dist.ProcessGroup, + handle: Optional[tuple], + topk_idx: Optional[torch.Tensor], + token_weights: Optional[torch.Tensor], + num_experts: Optional[int], +) -> Tuple[str, EPBufferConfig, Optional[EPAutoTuneResult]]: + """Tune the config for the given shape or get default config.""" + backend_name = _get_backend_name() + user_cfg = _get_buffer_config(backend_name) + + # Autotune off: use the user's backend + per-backend config. + if not GlobalBackendManager.auto_tune_enabled(): + return backend_name, user_cfg, None + + # Candidates: pinned backend (tune only that one), else all available + # (infeasible backends are dropped in tune() via is_feasible). + pinned_bt = GlobalBackendManager.get_ep_backend(PrecisionType.BF16_FP16_FP32) + pinned_candidates = ( + [_BACKEND_TYPE_TO_NAME[pinned_bt]] + if pinned_bt is not None and pinned_bt in _BACKEND_TYPE_TO_NAME + else None + ) - Uses at least 2 bytes (bf16 size) so buffers work for both fp8 and bf16 - without reallocation. - """ - inp = x if isinstance(x, torch.Tensor) else x[0] - return inp.size(1) * max(inp.element_size(), 2) + result: Optional[EPAutoTuneResult] = None + if handle is not None: + result = MoEDispatchCombineAutoTuner.lookup_handle(handle) + else: + if topk_idx is not None and token_weights is not None and num_experts is not None: + result = MoEDispatchCombineAutoTuner.get_or_tune( + group=group, + x=x, + num_experts=int(num_experts), + topk_idx=topk_idx, + topk_weights=token_weights, + num_sms=user_cfg.num_sms, + candidate_backends=pinned_candidates, + ) -def _ensure_buffer( - group: dist.ProcessGroup, - hidden_bytes: int, - backend: EPBackend, -) -> None: - """Make sure the backend's buffer is initialized.""" - if _buffer_config is None: - raise RuntimeError( - "set_buffer_global_config() must be called before dispatch/combine. " - "This is typically done by the token dispatcher during __init__." + if result is None: + result = MoEDispatchCombineAutoTuner.current() + + if result is not None: + winner_cfg = _get_buffer_config(result.backend_name) + cfg = EPBufferConfig( + num_sms=winner_cfg.num_sms, + dispatch_config=result.dispatch_config, + combine_config=result.combine_config, ) - backend.init_buffer(group, hidden_bytes, _buffer_config) + return result.backend_name, cfg, result + + # Nothing cached yet - fall back to user/default. + return backend_name, user_cfg, None # ========================================================================= -# Public API — used by ``moe_dispatch_combine.py`` +# Public API - used by ``moe_dispatch_combine.py`` # ========================================================================= +def _warn_if_graph_capture_incompatible(backend_name: str, op: str) -> None: + """Warn when a graph-incompatible backend runs under active capture. + + Triggered iff ``torch.cuda.is_current_stream_capturing()`` is True and the + backend's ``supports_cuda_graph()`` returns False. Does not raise; + kernels will likely deadlock so the warning makes the cause obvious. + """ + if not torch.cuda.is_current_stream_capturing(): + return + cls = _BACKEND_REGISTRY.get(backend_name) + if cls is None: + return + supports = getattr(cls, "supports_cuda_graph", None) + if supports is None or supports(): + return + logger.warning( + f"EP backend '{backend_name}' does not support CUDA graph capture " + f"(detected during moe_{op}); dispatch/combine may hang. Switch to " + f"a graph-capable backend (e.g. PRIMUS_TURBO_EP_BACKEND=TURBO) or " + f"disable capture for this layer.", + once=True, + ) + + def moe_dispatch_impl( x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], group: dist.ProcessGroup, @@ -542,10 +730,42 @@ def moe_dispatch_impl( allocate_on_comm_stream: bool = False, num_worst_tokens: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, tuple]: - name = _resolve_backend_name() + """Dispatch tokens to experts via the selected EP backend. + + Args: + x: Input tensor, or ``(fp8_tensor, scales)`` tuple for FP8 path. + group: EP process group. + handle: Cached dispatch handle; ``None`` for a primary call. + topk_idx: ``[num_tokens, num_topk]`` expert indices (primary only). + token_weights: ``[num_tokens, num_topk]`` expert weights. + num_experts: Total number of experts (primary only). + async_finish: If True, return before the kernel finishes. + allocate_on_comm_stream: Allocate outputs on the comm stream. + num_worst_tokens: Worst-case receive token count (for padding). + + Returns: + ``(recv_x, recv_topk_idx, recv_topk_weights, tokens_per_expert, handle)``. + """ + name, cfg, result = _maybe_tune_config(x, group, handle, topk_idx, token_weights, num_experts) + _warn_if_graph_capture_incompatible(name, "dispatch") backend = _get_backend_instance(name) - _ensure_buffer(group, get_hidden_bytes(x), backend) - return backend.dispatch( + if num_experts is None or topk_idx is None: + assert backend.is_initialized(), "Backend is not initialized" + else: + _release_other_backends(name) + # Unpack FP8 ``(tensor, scales)`` tuple to get shape metadata. + fp8_dispatch = isinstance(x, tuple) + x_inp = x[0] if fp8_dispatch else x + backend.init_buffer( + group, + hidden_size=x_inp.size(1), + num_experts=num_experts, + num_topk=topk_idx.size(1), + seqlen=x_inp.size(0), + fp8_dispatch=fp8_dispatch, + config=cfg, + ) + outputs = backend.dispatch( x, handle=handle, topk_idx=topk_idx, @@ -554,8 +774,16 @@ def moe_dispatch_impl( async_finish=async_finish, allocate_on_comm_stream=allocate_on_comm_stream, num_worst_tokens=num_worst_tokens, + config=cfg.dispatch_config, ) + # Bind the autotune result to the new handle for paired-call reuse. + if result is not None: + new_handle = outputs[-1] + MoEDispatchCombineAutoTuner.register_handle(new_handle, result) + + return outputs + def moe_combine_impl( x: torch.Tensor, @@ -565,13 +793,28 @@ def moe_combine_impl( async_finish: bool = False, allocate_on_comm_stream: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - name = _resolve_backend_name() + """Combine expert outputs back to original token order via the selected backend. + + Args: + x: Expert-side tensor to be combined. + group: EP process group. + handle: Handle returned by the paired ``moe_dispatch_impl`` call. + topk_weights: Optional per-token expert weights. + async_finish: If True, return before the kernel finishes. + allocate_on_comm_stream: Allocate outputs on the comm stream. + + Returns: + ``(combined_x, combined_topk_weights)``. + """ + name, cfg, _ = _maybe_tune_config(x, group, handle, None, None, None) + _warn_if_graph_capture_incompatible(name, "combine") backend = _get_backend_instance(name) - _ensure_buffer(group, get_hidden_bytes(x), backend) + assert backend.is_initialized(), "Backend is not initialized" return backend.combine( x, handle=handle, topk_weights=topk_weights, async_finish=async_finish, allocate_on_comm_stream=allocate_on_comm_stream, + config=cfg.combine_config, ) diff --git a/primus_turbo/pytorch/ops/moe/moe_dispatch_combine.py b/primus_turbo/pytorch/ops/moe/moe_dispatch_combine.py index 3df0ae77d..791e8cd5f 100644 --- a/primus_turbo/pytorch/ops/moe/moe_dispatch_combine.py +++ b/primus_turbo/pytorch/ops/moe/moe_dispatch_combine.py @@ -52,7 +52,8 @@ def forward( ctx.async_finish = async_finish ctx.allocate_on_comm_stream = allocate_on_comm_stream - tokens_per_expert = torch.tensor(tokens_per_expert) + if not isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = torch.tensor(tokens_per_expert) return (recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle) diff --git a/tests/pytorch/core/test_global_backend_manager.py b/tests/pytorch/core/test_global_backend_manager.py index 61834a0ef..89d3997dc 100644 --- a/tests/pytorch/core/test_global_backend_manager.py +++ b/tests/pytorch/core/test_global_backend_manager.py @@ -21,7 +21,7 @@ def clean_backend_state(monkeypatch): for key in ( "PRIMUS_TURBO_GEMM_BACKEND", "PRIMUS_TURBO_GROUPED_GEMM_BACKEND", - "PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", + "PRIMUS_TURBO_EP_BACKEND", "PRIMUS_TURBO_AUTO_TUNE", ): monkeypatch.delenv(key, raising=False) @@ -48,9 +48,9 @@ def test_grouped_gemm_backend_env(self, monkeypatch): monkeypatch.setenv("PRIMUS_TURBO_GROUPED_GEMM_BACKEND", "hipblaslt") assert GlobalBackendManager.get_grouped_gemm_backend(PrecisionType.FP8) == BackendType.HIPBLASLT - def test_moe_dispatch_combine_backend_env(self, monkeypatch): - monkeypatch.setenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", "triton") - assert GlobalBackendManager.get_moe_dispatch_combine_backend(PrecisionType.FP8) == BackendType.TRITON + def test_ep_backend_env(self, monkeypatch): + monkeypatch.setenv("PRIMUS_TURBO_EP_BACKEND", "triton") + assert GlobalBackendManager.get_ep_backend(PrecisionType.FP8) == BackendType.TRITON def test_auto_tune_env_enabled(self, monkeypatch): monkeypatch.setenv("PRIMUS_TURBO_AUTO_TUNE", "1") @@ -76,7 +76,7 @@ def test_gemm_backend_invalid_precision_raises(self, monkeypatch): def test_returns_none_when_env_not_set(self): assert GlobalBackendManager.get_gemm_backend(PrecisionType.FP8) is None assert GlobalBackendManager.get_grouped_gemm_backend(PrecisionType.FP8) is None - assert GlobalBackendManager.get_moe_dispatch_combine_backend(PrecisionType.FP8) is None + assert GlobalBackendManager.get_ep_backend(PrecisionType.FP8) is None assert GlobalBackendManager.auto_tune_enabled() is False diff --git a/tests/pytorch/modules/test_token_dispatcher.py b/tests/pytorch/modules/test_token_dispatcher.py index 719b9ff9a..b34078543 100644 --- a/tests/pytorch/modules/test_token_dispatcher.py +++ b/tests/pytorch/modules/test_token_dispatcher.py @@ -29,14 +29,45 @@ def _get_backends(): """Return available backend names.""" + all_backends = ["TURBO"] try: import deep_ep # noqa: F401 - # TODO: add backend deepep, temporal disable it since rocm7.2 deepep bug. + # TODO: add backend deepep, temporarily disable it since rocm7.2 deepep bug. # return ["TURBO", "DEEP_EP"] - return ["TURBO"] + # all_backends.append("DEEP_EP") except ImportError: - return ["TURBO"] + pass + + try: + import mori # noqa: F401 + + all_backends.append("MORI") + except ImportError: + pass + return all_backends + + +def _get_cuda_graph_backends(): + """Filter ``_get_backends`` down to ones that self-declare graph-safe. + + Done at parametrize time so graph-unsafe backends (e.g. Mori) never enter + the worker process: calling ``self.skipTest`` mid-test under + ``MultiProcContinuousTest`` re-raises as ``RuntimeError`` and fails the run. + """ + from primus_turbo.pytorch.kernels.moe.moe_dispatch_combine_impl import ( + _BACKEND_REGISTRY, + ) + + out = [] + for name in _get_backends(): + cls = _BACKEND_REGISTRY.get(name) + if cls is None: + continue + supports = getattr(cls, "supports_cuda_graph", None) + if supports is None or supports(): + out.append(name) + return out def _build_tp_groups(tp_size): @@ -152,6 +183,10 @@ class TestTokenDispatcher(MultiProcContinuousTest): # -2 tells MultiProcContinuousTest to use torch.cuda.device_count() world_size = -2 + @classmethod + def backend_str(cls) -> str: + return "nccl" + @property def device(self) -> torch.device: return torch.device("cuda", self.rank) @@ -178,10 +213,10 @@ def _bind_device(self) -> None: @parametrize("expert_capacity_factor", [None, 0.5]) def test_basic(self, backend, deepep_use_cuda_num_tokens_per_expert, expert_capacity_factor): self._bind_device() - with patch.dict(os.environ, {"PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND": backend}): + with patch.dict(os.environ, {"PRIMUS_TURBO_EP_BACKEND": backend}): _run_dispatch_combine( self.rank, - dist.group.WORLD, + self.pg, deepep_use_cuda_num_tokens_per_expert=deepep_use_cuda_num_tokens_per_expert, expert_capacity_factor=expert_capacity_factor, ) @@ -194,10 +229,10 @@ def test_basic(self, backend, deepep_use_cuda_num_tokens_per_expert, expert_capa @parametrize("permute_max_token_num", [0, NUM_TOKENS * 8 * ROUTER_TOPK]) def test_worst_tokens(self, backend, permute_max_token_num): self._bind_device() - with patch.dict(os.environ, {"PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND": backend}): + with patch.dict(os.environ, {"PRIMUS_TURBO_EP_BACKEND": backend}): _run_dispatch_combine( self.rank, - dist.group.WORLD, + self.pg, deepep_use_cuda_num_tokens_per_expert=True, deepep_num_worst_tokens=NUM_TOKENS * 8, permute_max_token_num=permute_max_token_num, @@ -211,7 +246,7 @@ def test_worst_tokens(self, backend, permute_max_token_num): @parametrize("pad_multiple", [128, 256]) def test_pad_multiple(self, backend, pad_multiple): self._bind_device() - with patch.dict(os.environ, {"PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND": backend}): + with patch.dict(os.environ, {"PRIMUS_TURBO_EP_BACKEND": backend}): _run_dispatch_combine( self.rank, dist.group.WORLD, @@ -223,20 +258,51 @@ def test_pad_multiple(self, backend, pad_multiple): # Autotune env var (PRIMUS_TURBO_AUTO_TUNE=1) # ------------------------------------------------------------------ - @parametrize("backend", _get_backends()) - def test_autotune(self, backend): + def test_autotune(self): self._bind_device() - with patch.dict( - os.environ, - { - "PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND": backend, - "PRIMUS_TURBO_AUTO_TUNE": "1", - }, - ): - _run_dispatch_combine( - self.rank, - dist.group.WORLD, - ) + # ``patch.dict`` snapshots ``os.environ`` and restores it on exit, so + # the inner ``pop`` is reverted automatically. + with patch.dict(os.environ, {"PRIMUS_TURBO_AUTO_TUNE": "1"}): + os.environ.pop("PRIMUS_TURBO_EP_BACKEND", None) + _run_dispatch_combine(self.rank, self.pg) + + # ------------------------------------------------------------------ + # Autotune without an explicit backend override: actually sweep configs + # and (when multiple backends exist) select the fastest backend. + # ------------------------------------------------------------------ + + def test_autotune_sweep(self): + """Exercise the real autotune path added to ``_DeepEPLikeBackend``. + + With ``PRIMUS_TURBO_AUTO_TUNE=1`` and **no** backend pinned, the first + ``moe_dispatch`` for a given shape triggers a full (dispatch + + combine) ``nvl_chunk_size`` sweep, binds the result to the freshly + returned ``handle``, and the paired ``moe_combine`` simply looks the + result up by ``id(handle)``. A second run on the same shape hits the + shape cache without re-measuring. + """ + from primus_turbo.pytorch.kernels.moe.moe_dispatch_combine_impl import ( + MoEDispatchCombineAutoTuner, + ) + + self._bind_device() + with patch.dict(os.environ, {"PRIMUS_TURBO_AUTO_TUNE": "1"}): + os.environ.pop("PRIMUS_TURBO_EP_BACKEND", None) + MoEDispatchCombineAutoTuner.clear() + + # First run triggers autotune and registers the handle mapping. + _run_dispatch_combine(self.rank, self.pg) + first = MoEDispatchCombineAutoTuner.current() + assert first is not None, "autotune should have populated a result" + assert ( + MoEDispatchCombineAutoTuner.handle_cache_size() > 0 + ), "moe_dispatch should have bound the tuned result to at least one handle" + + # Second run hits the shape cache (no extra measurements) and + # yields the same backend selection. + _run_dispatch_combine(self.rank, self.pg) + second = MoEDispatchCombineAutoTuner.current() + assert second is first or second.backend_name == first.backend_name # ------------------------------------------------------------------ # tp_size > 1 coverage @@ -247,7 +313,7 @@ def test_tp_size_2(self, backend): if dist.get_world_size() < 2: self.skipTest("requires world_size >= 2") self._bind_device() - with patch.dict(os.environ, {"PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND": backend}): + with patch.dict(os.environ, {"PRIMUS_TURBO_EP_BACKEND": backend}): _run_dispatch_combine( self.rank, None, @@ -270,7 +336,7 @@ def test_tp_size_2_with_routing_map(self, backend): # for every token (pre-TP-expansion expert id space). routing_map = torch.zeros((num_tokens, num_experts), dtype=torch.bool, device="cuda") routing_map[:, :router_topk] = True - with patch.dict(os.environ, {"PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND": backend}): + with patch.dict(os.environ, {"PRIMUS_TURBO_EP_BACKEND": backend}): _run_dispatch_combine( self.rank, None, @@ -294,7 +360,7 @@ def test_tp_size_2_with_token_indices(self, backend): # is expected to TP-expand them internally. probs = torch.ones((num_tokens, num_experts), dtype=torch.float32, device="cuda") / router_topk _, token_indices = torch.topk(probs, router_topk, dim=-1) - with patch.dict(os.environ, {"PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND": backend}): + with patch.dict(os.environ, {"PRIMUS_TURBO_EP_BACKEND": backend}): _run_dispatch_combine( self.rank, None, @@ -309,23 +375,6 @@ def test_tp_size_2_with_token_indices(self, backend): # ---------------------------------------------------------------------- # CUDA graph compatibility tests -# -# The C++ helper ``is_ep_force_current_stream()`` caches the value of -# ``PRIMUS_TURBO_EP_FORCE_CURRENT_STREAM`` in a ``static`` local on first -# call, so toggling the env var inside a worker process after the first -# ``Buffer`` has been constructed is a no-op. Because -# ``MultiProcContinuousTest`` reuses the same worker processes across every -# test method in a class, we cannot rely on ``patch.dict(os.environ, ...)`` -# inside a test body to enable current-stream mode. -# -# Instead we put the CUDA-graph tests in a dedicated ``MultiProcContinuousTest`` -# subclass. Each subclass spawns its own worker pool via -# ``torch.multiprocessing.Process`` with the ``spawn`` start method, which -# copies the parent's ``os.environ`` at ``process.start()`` time. By setting -# the env var in ``setUpClass`` and eagerly spawning the workers before -# restoring ``os.environ``, the workers for this class inherit the flag and -# their static cache is initialized to ``1`` on the very first Buffer -# construction, while other test classes' worker pools remain unaffected. # ---------------------------------------------------------------------- @@ -333,6 +382,10 @@ def test_tp_size_2_with_token_indices(self, backend): class TestTokenDispatcherCudaGraph(MultiProcContinuousTest): world_size = -2 + @classmethod + def backend_str(cls) -> str: + return "nccl" + @property def device(self) -> torch.device: return torch.device("cuda", self.rank) @@ -345,14 +398,14 @@ def setUpClass(cls): os.environ["PRIMUS_TURBO_EP_FORCE_CURRENT_STREAM"] = "1" super().setUpClass() - @parametrize("backend", _get_backends()) + @parametrize("backend", _get_cuda_graph_backends()) def test_cuda_graph(self, backend): from primus_turbo.pytorch.kernels.moe import moe_dispatch_combine_impl self._bind_device() with patch.dict( os.environ, - {"PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND": backend}, + {"PRIMUS_TURBO_EP_BACKEND": backend}, ): num_worst_tokens = NUM_TOKENS * 8 permute_max_token_num = NUM_TOKENS * 8 * ROUTER_TOPK @@ -366,7 +419,8 @@ def test_cuda_graph(self, backend): dispatcher = turbo.modules.DeepEPTokenDispatcher( NUM_EXPERTS, ROUTER_TOPK, - dist.group.WORLD, + self.pg, + permute_fusion=True, deepep_use_cuda_num_tokens_per_expert=True, deepep_num_worst_tokens=num_worst_tokens, permute_max_token_num=permute_max_token_num,