Feature: add expert map support for shared experts & EPLB

examples/ops/dispatch_combine/test_dispatch_combine.cpp

@@ -179,10 +179,10 @@ class EpDispatchCombineTestCase {
     RandomInitializeToken();
     if (IsNoDataRank(handle.config)) {
       handle.PrepareInference(GetHipDataType<T>(), nullptr, outTokBuf, weightsBuf, scalesBuf,
-                              nullptr, 0);
+                              nullptr, nullptr, 0);
     } else {
       handle.PrepareInference(GetHipDataType<T>(), inpTokBuf, outTokBuf, weightsBuf, scalesBuf,
-                              tokenIndices, numToken);
+                              tokenIndices, nullptr, numToken);
     }
     // PrintDispatch();
     // PrintDispatchStats();
@@ -465,7 +465,7 @@ class EpDispatchCombineTestCase {
     EpDispatchCombineConfig& config = handle.config;
     if (IsNoDataRank(handle.config)) {
       handle.PrepareInference(GetHipDataType<T>(), inpTokBuf, outTokBuf, weightsBuf, scalesBuf,
-                              nullptr, 0);
+                              nullptr, nullptr, 0);
     }
     // HIP_RUNTIME_CHECK(hipMemcpy(inpTokBuf, outTokBuf,
     //                             config.MaxNumTokensToRecvPerRank() * config.hiddenDim *

examples/ops/dispatch_combine/test_dispatch_combine_internode.py

@@ -42,6 +42,7 @@ def __init__(
         max_tokens,
         kernel_type,
         dtype=torch.bfloat16,
+        with_expert_map=False,
     ):
         self.rank = rank
         self.gpu_per_node = gpu_per_node
@@ -64,6 +65,7 @@ def __init__(
             gpu_per_node=self.gpu_per_node,
             rdma_block_num=64,
         )
+        self.with_expert_map = with_expert_map
 
     def setup(self):
         local_rank = self.rank % self.gpu_per_node
@@ -214,15 +216,21 @@ def gen_test_data(self, use_max_token_num=False):
                 ).to(self.config.data_type)
             )
 
+        total_num_expert = self.config.num_experts_per_rank * self.config.world_size
+        expert_map = torch.randperm(
+            total_num_expert, generator=self.rng, device=self.device
+        ).to(torch.int32)
+
         return (
             num_token,
             all_rank_indices,
             all_rank_input,
             all_rank_weights,
             all_rank_scales,
+            expert_map if self.with_expert_map else None,
         )
 
-    def count_token_num(self, all_rank_indices):
+    def count_token_num(self, all_rank_indices, expert_map):
         # Per-rank counts
         rank_counts = torch.zeros(
             self.config.world_size, dtype=torch.int32, device=self.device
@@ -238,9 +246,11 @@ def count_token_num(self, all_rank_indices):
             src_node = src_rank // self.config.gpu_per_node
 
             # Map expert IDs to rank IDs
-            token_ranks = (
-                indices // self.config.num_experts_per_rank
-            )  # [num_tokens, num_experts_per_token]
+            if expert_map is not None:
+                indices = [expert_map[ids] for ids in indices]
+            else:
+                indices = [ids for ids in indices]
+            token_ranks = [ids // self.config.num_experts_per_rank for ids in indices]
 
             # Deduplicate rank IDs per token
             unique_ranks_per_token = [torch.unique(row) for row in token_ranks]
@@ -265,10 +275,10 @@ def count_token_num(self, all_rank_indices):
                         rank_counts_remote_send[src_rank] += 1
 
         if self.config.rank == 0:
-            print("Rank counts (deduplicated):", rank_counts)
-            # print("Rank counts local nodes:", rank_counts - rank_counts_remote_recv)
-            # print("Rank counts from other nodes:", rank_counts_remote_recv)
-            # print("Rank counts to other nodes:", rank_counts_remote_send)
+            print(self.config.rank, "Rank counts (deduplicated):", rank_counts)
+        # print("Rank counts local nodes:", rank_counts - rank_counts_remote_recv)
+        # print("Rank counts from other nodes:", rank_counts_remote_recv)
+        # print("Rank counts to other nodes:", rank_counts_remote_send)
         return rank_counts, rank_counts_remote_recv, rank_counts_remote_send
 
     def run_test_once(self, op, test_data, error_round, round):
@@ -278,6 +288,7 @@ def run_test_once(self, op, test_data, error_round, round):
             all_rank_input,
             all_rank_weights,
             all_rank_scales,
+            expert_map,
         ) = test_data
         dist.barrier()
 
@@ -293,24 +304,25 @@ def run_test_once(self, op, test_data, error_round, round):
             # None,
             all_rank_scales[self.rank],
             all_rank_indices[self.rank],
+            index_to_exp_id=expert_map,
             block_num=self.config.block_num,
             warp_per_block=16,
         )
         torch.cuda.synchronize()
         dist.barrier()
 
-        rank_counts, _, _ = self.count_token_num(all_rank_indices)
+        rank_counts, _, _ = self.count_token_num(all_rank_indices, expert_map)
 
         src_token_pos = op.get_dispatch_src_token_pos().tolist()
         max_num_token_to_send_per_rank = self.config.max_num_inp_token_per_rank
         recv_token_num = len(src_token_pos)
 
         # Check recv token num
         print(f"rank {self.rank} recv {recv_token_num} tokens")
-        token_num_pass = rank_counts[self.rank] == recv_token_num
+        token_num_pass = rank_counts[self.config.rank] == recv_token_num
         if not token_num_pass:
             print(
-                f"rank {self.rank} expected token num {rank_counts[self.rank]} got {recv_token_num}"
+                f"rank {self.rank} expected token num {rank_counts[self.config.rank]} got {recv_token_num}"
             )
             assert False
 
@@ -353,9 +365,11 @@ def run_test_once(self, op, test_data, error_round, round):
         dist.barrier()
 
         for i in range(all_rank_num_token[self.rank]):
+            indices = all_rank_indices[self.config.rank][i].cpu()
+            if expert_map is not None:
+                indices = expert_map[indices]
             pes = [
-                (idx // self.config.num_experts_per_rank)
-                for idx in all_rank_indices[self.rank][i].cpu().tolist()
+                (idx // self.config.num_experts_per_rank) for idx in indices.tolist()
             ]
             unique_pes = len(set(pes))
             unique_innode_pes = len(
@@ -449,6 +463,7 @@ def run_bench_once(self, op, test_data):
             all_rank_input,
             all_rank_weights,
             all_rank_scales,
+            expert_map,
         ) = test_data
 
         start_event = torch.cuda.Event(enable_timing=True)
@@ -468,6 +483,7 @@ def run_bench_once(self, op, test_data):
             all_rank_weights[self.rank],
             all_rank_scales[self.rank],
             all_rank_indices[self.rank],
+            index_to_exp_id=expert_map,
             block_num=self.config.block_num,
             warp_per_block=16,
         )
@@ -695,6 +711,7 @@ def test_dispatch_combine(
         kernel_type,
         torch.bfloat16,
         # torch.float8_e4m3fnuz,
+        with_expert_map=False,
     )
     test_case.setup()
     if is_bench:

include/mori/ops/dispatch_combine/dispatch_combine.hpp

@@ -69,7 +69,6 @@ inline size_t GetHipDataTypeSize(hipDataType dtype) {
 
 using index_t = int32_t;
 
-#define MAX_EXPERTS_PER_TOKEN (8)
 struct EpDispatchCombineConfig {
   int rank{0};
   int worldSize{0};
@@ -107,25 +106,28 @@ class EpDispatchCombineHandle {
   ~EpDispatchCombineHandle();
 
   void PrepareInference(hipDataType inputType, void* input, void* output, float* weights,
-                        index_t* tokenIndices, index_t numToken) {
+                        index_t* tokenIndices, index_t* indexToExpertId, index_t numToken) {
     this->inputType = inputType;
     this->inpTokenBuf = input;
     this->outTokenBuf = output;
     this->weightsBuf = weights;
     this->tokenIndices = tokenIndices;
     this->curRankNumToken = numToken;
+    this->indexToExpertId = indexToExpertId;
     // printf("handle inputType %s\n", HipDataTypeToString(inputType));
   }
 
   void PrepareInference(hipDataType inputType, void* input, void* output, float* weights,
-                        uint8_t* scales, index_t* tokenIndices, index_t numToken) {
+                        uint8_t* scales, index_t* tokenIndices, index_t* indexToExpertId,
+                        index_t numToken) {
     this->inputType = inputType;
     this->inpTokenBuf = input;
     this->outTokenBuf = output;
     this->weightsBuf = weights;
     this->scalesBuf = scales;
     this->tokenIndices = tokenIndices;
     this->curRankNumToken = numToken;
+    this->indexToExpertId = indexToExpertId;
     // printf("handle inputType %s\n", HipDataTypeToString(inputType));
   }
 
@@ -163,6 +165,7 @@ class EpDispatchCombineHandle {
   EpDispatchCombineConfig config;
   // Routed expert indices for tokens
   index_t* tokenIndices{nullptr};
+  index_t* indexToExpertId{nullptr};
 
   // Kernel input/output buffer
   void* inpTokenBuf{nullptr};
@@ -248,6 +251,7 @@ struct EpDispatchCombineArgs {
   EpDispatchCombineConfig config;
   index_t curRankNumToken{0};
   index_t* tokenIndices{nullptr};
+  index_t* indexToExpertId{nullptr};
   T* inpTokenBuf{nullptr};
   T* outTokenBuf{nullptr};
   float* weightsBuf{nullptr};
@@ -301,6 +305,7 @@ EpDispatchCombineArgs<T> GetEpDispatchCombineArgs(const EpDispatchCombineHandle&
   args.config = handle.config;
   args.curRankNumToken = handle.curRankNumToken;
   args.tokenIndices = handle.tokenIndices;
+  args.indexToExpertId = handle.indexToExpertId;
   args.inpTokenBuf = reinterpret_cast<T*>(handle.inpTokenBuf);
   args.outTokenBuf = reinterpret_cast<T*>(handle.outTokenBuf);
   args.weightsBuf = handle.weightsBuf;

python/mori/ops/dispatch_combine.py

@@ -107,6 +107,7 @@ def dispatch(
         weights: torch.Tensor,
         scales: torch.Tensor,
         indices: torch.Tensor,
+        index_to_exp_id: torch.Tensor = None,
         block_num: int = -1,
         warp_per_block: int = -1,
     ):
@@ -117,6 +118,7 @@ def dispatch(
             weights,
             scales,
             indices,
+            index_to_exp_id,
             block_num,
             warp_per_block,
         )

src/ops/dispatch_combine/internode.hpp

@@ -24,6 +24,7 @@
 #include "mori/core/core.hpp"
 #include "mori/ops/dispatch_combine/dispatch_combine.hpp"
 #include "mori/shmem/shmem.hpp"
+#include "src/ops/dispatch_combine/intranode.hpp"
 
 namespace mori {
 namespace moe {
@@ -91,7 +92,8 @@ __global__ void EpDispatchInterNodeKernel(EpDispatchCombineArgs<T> args) {
     for (int tokenId = globalSubWarpId; tokenId < args.curRankNumToken;
          tokenId += globalSubWarpNum) {
       const int expertOffset = tokenId * numExpertPerToken + laneInSubWarp;
-      index_t destExpert = args.tokenIndices[expertOffset];
+      // index_t destExpert = args.tokenIndices[expertOffset];
+      index_t destExpert = GetExpertIndex(args, expertOffset);
       index_t destPe = destExpert / config.numExpertPerRank;
 
       unsigned long long subWarpMask = ((1ULL << numExpertPerToken) - 1ULL)

src/ops/dispatch_combine/intranode.hpp

@@ -62,6 +62,13 @@ inline __device__ void CrossDeviceBarrierIntraNodeKernel(EpDispatchCombineArgs<T
   __syncthreads();
 }
 
+template <typename T>
+inline __device__ index_t GetExpertIndex(EpDispatchCombineArgs<T> args, index_t id) {
+  index_t expert = args.tokenIndices[id];
+  if (args.indexToExpertId) expert = args.indexToExpertId[expert];
+  return expert;
+}
+
 /* ---------------------------------------------------------------------------------------------- */
 /*                                    EpDispatchIntraNodeKernel                                   */
 /* ---------------------------------------------------------------------------------------------- */
@@ -90,15 +97,18 @@ __global__ void EpDispatchIntraNodeKernel(EpDispatchCombineArgs<T> args) {
     for (int i = globalWarpId; i < args.curRankNumToken * config.numExpertPerToken;
          i += globalWarpNum) {
       index_t srcTokId = i / config.numExpertPerToken;
-      index_t destExpert = args.tokenIndices[i];
+      // index_t destExpert = args.tokenIndices[i];
+      index_t destExpert = GetExpertIndex(args, i);
       index_t destPe = destExpert / config.numExpertPerRank;
       index_t destTokId = 0;
 
       // Deduplicate
       assert(config.numExpertPerToken < warpSize);
       int condition = 0;
       if (laneId < (i % config.numExpertPerToken)) {
-        condition = destPe == (args.tokenIndices[srcTokId * config.numExpertPerToken + laneId] /
+        // condition = destPe == (args.tokenIndices[srcTokId * config.numExpertPerToken + laneId] /
+        //                        config.numExpertPerRank);
+        condition = destPe == (GetExpertIndex(args, srcTokId * config.numExpertPerToken + laneId) /
                                config.numExpertPerRank);
       }
       if (__any(condition)) {
@@ -127,6 +137,7 @@ __global__ void EpDispatchIntraNodeKernel(EpDispatchCombineArgs<T> args) {
               destPe)[destTokId * config.numExpertPerToken + laneId] =
               args.weightsBuf[srcTokId * config.numExpertPerToken + laneId];
         }
+        // Output indicies are used for later computation ops, do not map
         args.shmemOutIndicesMemObj->template GetAs<index_t*>(
             destPe)[destTokId * config.numExpertPerToken + laneId] =
             args.tokenIndices[srcTokId * config.numExpertPerToken + laneId];

src/pybind/mori.cpp

@@ -48,7 +48,8 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<torch::Ten
 LaunchDispatch(mori::moe::EpDispatchCombineHandle& handle, int kernelType,
                const torch::Tensor& input, const std::optional<torch::Tensor>& weights,
                const std::optional<torch::Tensor>& scales, const torch::Tensor& topkIds,
-               int blockNum = -1, int warpPerBlock = -1) {
+               const std::optional<torch::Tensor>& idxToExpIds, int blockNum = -1,
+               int warpPerBlock = -1) {
   assert(input.is_contiguous() && topkIds.is_contiguous());
 
   float* weightPtr = nullptr;
@@ -63,9 +64,16 @@ LaunchDispatch(mori::moe::EpDispatchCombineHandle& handle, int kernelType,
     scalePtr = reinterpret_cast<uint8_t*>(scales->data_ptr());
   }
 
+  mori::moe::index_t* idxToExpIdsPtr = nullptr;
+  if (idxToExpIds.has_value()) {
+    assert(idxToExpIds->is_contiguous() &&
+           idxToExpIds->numel() == (handle.config.numExpertPerRank * handle.config.worldSize));
+    idxToExpIdsPtr = reinterpret_cast<mori::moe::index_t*>(idxToExpIds->data_ptr());
+  }
+
   handle.PrepareInference(mori::ScalarTypeToHipDataType(input.scalar_type()), input.data_ptr(),
                           nullptr, weightPtr, scalePtr, topkIds.data_ptr<mori::moe::index_t>(),
-                          input.size(0));
+                          idxToExpIdsPtr, input.size(0));
   handle.LaunchDispatch((mori::moe::KernelType)kernelType, blockNum, warpPerBlock,
                         at::cuda::getCurrentHIPStream());
 
@@ -120,7 +128,7 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>> LaunchCombine(
   }
 
   handle.PrepareInference(mori::ScalarTypeToHipDataType(input.scalar_type()), input.data_ptr(),
-                          nullptr, weightsPtr, topkIds.data_ptr<mori::moe::index_t>(),
+                          nullptr, weightsPtr, topkIds.data_ptr<mori::moe::index_t>(), nullptr,
                           handle.curRankNumToken);
   handle.LaunchCombine((mori::moe::KernelType)kernelType, blockNum, warpPerBlock,
                        at::cuda::getCurrentHIPStream());

tests/python/ops/bench_dispatch_combine.py

@@ -31,7 +31,7 @@ def __init__(self, config):
         super().__init__(config)
 
     def gen_test_data(self):
-        return super().gen_test_data(use_max_token_num=True)
+        return super().gen_test_data(use_max_token_num=True, with_expert_map=False)
 
     def run_once(self, op, test_data, check_result):
         (
@@ -40,6 +40,7 @@ def run_once(self, op, test_data, check_result):
             all_rank_input,
             all_rank_weights,
             all_rank_scales,
+            _,
         ) = test_data
 
         start_event = torch.cuda.Event(enable_timing=True)
@@ -92,7 +93,7 @@ def run_once(self, op, test_data, check_result):
             None,
             dispatch_indices,
             block_num=80,
-            warp_per_block=16,
+            warp_per_block=4,
         )
         end_event.record()
         self.sync()
@@ -182,7 +183,7 @@ def _bench_dispatch_combine(
     world_size,
     port,
     max_num_inp_token_per_rank=128,
-    data_type=torch.float8_e4m3fnuz,
+    data_type=torch.bfloat16,
     hidden_dim=7168,
     scale_dim=0,
     scale_type_size=0,