feat(EPv2): extend support for rack-level wide-ep. (#597)

Author: kawhil-amd

python/mori/jit/config.py

@@ -90,6 +90,20 @@ def detect_gpu_arch(rocm_path: str = "/opt/rocm") -> str:
     )
 
 
+def detect_wave_size() -> int:
+    """Return the wavefront width for the current GPU (32 or 64).
+
+    Override with MORI_WAVE_SIZE.  Falls back to 64 when detection fails.
+    """
+    v = os.environ.get("MORI_WAVE_SIZE")
+    if v:
+        return int(v)
+    try:
+        return 32 if detect_gpu_arch().startswith("gfx12") else 64
+    except Exception:
+        return 64
+
+
 def _find_tool(rocm_path: str, name: str) -> str:
     """Locate a ROCm LLVM tool, raising FileNotFoundError if missing."""
     candidates = [

python/mori/ops/dispatch_combine_v2/dispatch_combine_op.py

@@ -50,6 +50,9 @@
 import torch
 
 from mori.tensor_utils import from_gpu_ptr
+from mori.jit.config import detect_wave_size
+
+WAVE = detect_wave_size()
 
 # Where each backend lives. Imported lazily, on selection only.
 _BACKEND_MODULES = {"flydsl": "flydsl_backend", "hip": "hip_backend"}
@@ -128,12 +131,6 @@ def __post_init__(self):
             )
         if self.quant_type != "none":
             self.combine_mode = "scatter"
-        # The dispatch grid barrier resets inside a `range(lane, npes, 64)` loop,
-        # correct only while each lane runs it once (npes <= wavefront).
-        if self.world_size > 64:
-            raise ValueError(
-                f"intranode op supports world_size <= 64, got {self.world_size}"
-            )
         # Token copy moves whole 16 B (vec4) chunks; a non-16 B-aligned per-token
         # size would over-read/write a few dwords past the token.
         if self.token_nbytes % 16 != 0:
@@ -233,6 +230,22 @@ def _resolve_geometry(self):
             if self.combine_warp_num_per_block is None:
                 self.combine_warp_num_per_block = 4
 
+        # Precise world_size check against the resolved geometry.  The combine
+        # xdb barrier polls with `tid < npes`, so every schedule bucket must
+        # have blockDim (= comb_warp * WAVE) >= world_size.
+        if self.schedule:
+            min_comb_warp = min(bucket[4] for bucket in self.schedule)
+        else:
+            min_comb_warp = self.combine_warp_num_per_block
+        max_peers = min_comb_warp * WAVE
+        if self.world_size > max_peers:
+            raise ValueError(
+                f"world_size ({self.world_size}) exceeds the smallest combine "
+                f"blockDim in the schedule ({min_comb_warp} warps × {WAVE}-wide "
+                f"wave = {max_peers} threads); the `tid < npes` barrier requires "
+                f"world_size <= blockDim"
+            )
+
     @property
     def is_scatter(self):
         return self.combine_mode == "scatter"

python/mori/ops/dispatch_combine_v2/intranode_kernels.py

@@ -72,25 +72,9 @@
 
 from . import flydsl_prims as P
 
-# Wavefront size: gfx9 (MI300/MI350) = 64, gfx12 (MI400/gfx1250) = 32. Detected
-# once per process; override with MORI_WAVE_SIZE. get_warp_size(gfx1250) wrongly
-# reports 64, so key off the arch string.
-import os as _os
+from mori.jit.config import detect_wave_size
 
-
-def _detect_wave_size():
-    v = _os.environ.get("MORI_WAVE_SIZE")
-    if v:
-        return int(v)
-    try:
-        from mori.jit.config import detect_gpu_arch
-
-        return 32 if detect_gpu_arch().startswith("gfx12") else 64
-    except Exception:
-        return 64
-
-
-WAVE = _detect_wave_size()
+WAVE = detect_wave_size()
 LANE_MASK = WAVE - 1
 LOG2_WAVE = WAVE.bit_length() - 1
 _BALLOT_INT = T.i64 if WAVE == 64 else T.i32
@@ -348,10 +332,11 @@ def ep_dispatch(
                 P.atomic_add_global(fx.Int64(addr_disp_bar), arith.constant(1))
 
             local_recv_num = fx.Int64(window.lsa_ptr(my_lsa_rank, off_recv_num))
+            if global_warp_id == 0:
+                P.spin_until_eq_i32(fx.Int64(addr_disp_bar), block_num)
+                buffer_store(arith.constant(0), rsrc_disp_bar, 0)
             for dest_pe in range(lane, npes, WAVE):
                 if global_warp_id == 0:
-                    P.spin_until_eq_i32(fx.Int64(addr_disp_bar), block_num)
-                    buffer_store(arith.constant(0), rsrc_disp_bar, 0)
                     signal_value = (
                         buffer_load(rsrc_dest_ctr, dest_pe, vec_width=1, dtype=T.i32())
                         + 1

src/ops/dispatch_combine_v2/ep_intranode_1250x.hpp

@@ -63,7 +63,11 @@ using index_t = int32_t;
 #define MORI_COMB_BARSLEEP 15
 #define MORI_COMB_BARSPREAD 16
 
-#define CUSPLIT_MAX_GPUS 8
+// MORI_EP_WORLD_SIZE is emitted by RenderEpSource before #include-ing this
+// header, so the global arrays below are sized to the exact config.
+#ifndef MORI_EP_WORLD_SIZE
+#define MORI_EP_WORLD_SIZE 8
+#endif
 
 template <typename T>
 __device__ __forceinline__ uint32_t MoriPackTo2(float a, float b) {
@@ -199,15 +203,15 @@ __device__ __forceinline__ gfx1250_TDM_GROUP1 TdmSplitShape(const TdmSplit128& s
   return TdmShape2D(32, sp.rows);
 }
 
-#define CUSPLIT_POOL_SLOTS (CUSPLIT_MAX_GPUS * 32768)
+#define CUSPLIT_POOL_SLOTS (MORI_EP_WORLD_SIZE * 32768)
 #define CUSPLIT_MAX_BLOCKS 512
 #define CUSPLIT_MAX_TOPK 16
 
 __device__ index_t _cusplit_stgIdx[CUSPLIT_POOL_SLOTS * CUSPLIT_MAX_TOPK];
 __device__ float _cusplit_stgWt[CUSPLIT_POOL_SLOTS * CUSPLIT_MAX_TOPK];
 __device__ index_t _cusplit_stgSrc[CUSPLIT_POOL_SLOTS];
-__device__ index_t _cusplit_blkBase[CUSPLIT_MAX_BLOCKS * CUSPLIT_MAX_GPUS];
-__device__ index_t _cusplit_blkCount[CUSPLIT_MAX_BLOCKS * CUSPLIT_MAX_GPUS];
+__device__ index_t _cusplit_blkBase[CUSPLIT_MAX_BLOCKS * MORI_EP_WORLD_SIZE];
+__device__ index_t _cusplit_blkCount[CUSPLIT_MAX_BLOCKS * MORI_EP_WORLD_SIZE];
 // Per-token scale rows, staged like the other meta fields so they ship to a peer as
 // one contiguous run rather than a 224 B transfer per (token, destination) -- the
 // size TDM is worst at. The array is at file scope, which the TU reaches before kCfg
@@ -299,7 +303,7 @@ __device__ void EpDispatch1250xBody(EpArgs args) {
   T* _tdmTile = reinterpret_cast<T*>(_tdmBatchSmem + (size_t)warpId * kSlabBytes);
   const gfx1250_TDM_GROUP1 _tdmG1 = TdmShape<T>(static_cast<int>(hiddenDim));
 
-  constexpr int kMaxNpes = CUSPLIT_MAX_GPUS;
+  constexpr int kMaxNpes = kCfg.worldSize;
   __shared__ index_t s_N[kMaxNpes];
   __shared__ index_t s_base[kMaxNpes];
   __shared__ index_t s_run[kMaxNpes];

src/ops/dispatch_combine_v2/ep_spec.cpp

@@ -162,8 +162,9 @@ std::string RenderEpSource(const EpCfg& cfg, const std::string& entry, const cha
                 "\n#define MORI_EP_SCALE_SLOTS " +
                 std::to_string((long long)cfg.worldSize * EpMaxRecv(cfg)) + "\n";
   }
-  return std::string("// mori jit v2 — generated, do not edit.\n") + scaleDefs + "#include \"" +
-         header +
+  std::string worldDef = "#define MORI_EP_WORLD_SIZE " + std::to_string(cfg.worldSize) + "\n";
+  return std::string("// mori jit v2 — generated, do not edit.\n") + worldDef + scaleDefs +
+         "#include \"" + header +
          "\"\n"
          "using namespace mori::ops::v2;\n"
          "constexpr EpCfg kCfg = " +