Refactor: moe dispatch combine autotune

.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

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

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

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<uint32_t>(std::stoi(v)) : 0;
+        return v ? static_cast<uint32_t>(std::stoi(v)) : 1;
     }();
     return val != 0;
 }

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<uint8_t *>(gbl_ptr) + offset * sizeof(dtype_t);
-        gbl_ptr     = reinterpret_cast<uint8_t *>(gbl_ptr) + total_bytes;
+    __device__ __forceinline__ Buffer(void *&gbl_ptr, int64_t num_elems, int64_t offset = 0) {
+        total_bytes = num_elems * static_cast<int64_t>(sizeof(dtype_t));
+        ptr = reinterpret_cast<uint8_t *>(gbl_ptr) + offset * static_cast<int64_t>(sizeof(dtype_t));
+        gbl_ptr = reinterpret_cast<uint8_t *>(gbl_ptr) + total_bytes;
     }
 
     __device__ __forceinline__ Buffer advance_also(void *&gbl_ptr) {
@@ -34,47 +34,47 @@ public:
 
     __device__ __forceinline__ dtype_t *buffer() { return reinterpret_cast<dtype_t *>(ptr); }
 
-    __device__ __forceinline__ dtype_t &operator[](int idx) { return buffer()[idx]; }
+    __device__ __forceinline__ dtype_t &operator[](int64_t idx) { return buffer()[idx]; }
 };
 
 template <typename dtype_t, int kNumRanks = 1> 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<int64_t>(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<uint8_t *>(gbl_ptr) + per_channel_bytes * sm_id + num_bytes * offset;
         gbl_ptr = reinterpret_cast<uint8_t *>(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<int64_t>(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<uint8_t *>(gbl_ptrs[i]) + per_channel_bytes * sm_id +
                       num_bytes * offset;
             gbl_ptrs[i] = reinterpret_cast<uint8_t *>(gbl_ptrs[i]) + total_bytes;
         }
     }
 
-    __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<int64_t>(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<dtype_t *>(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<dtype_t *>(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<int64_t>(sizeof(dtype_t));
 
-        int per_channel_bytes = num_bytes * num_ranks;
-        total_bytes           = per_channel_bytes * num_sms * (static_cast<int>(kDecoupled) + 1);
-        send_ptr              = reinterpret_cast<uint8_t *>(gbl_ptr) + per_channel_bytes * sm_id;
-        recv_ptr = reinterpret_cast<uint8_t *>(gbl_ptr) + per_channel_bytes * (sm_id + num_sms);
-        gbl_ptr  = reinterpret_cast<uint8_t *>(gbl_ptr) + total_bytes;
+        int64_t per_channel_bytes = num_bytes * num_ranks;
+        total_bytes = per_channel_bytes * num_sms * (static_cast<int>(kDecoupled) + 1);
+        send_ptr    = reinterpret_cast<uint8_t *>(gbl_ptr) + per_channel_bytes * sm_id;
+        recv_ptr    = reinterpret_cast<uint8_t *>(gbl_ptr) + per_channel_bytes * (sm_id + num_sms);
+        gbl_ptr     = reinterpret_cast<uint8_t *>(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<dtype_t *>(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<dtype_t *>(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<dtype_t *>(send_ptr + num_bytes * idx);
     }

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
 ```

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"

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