From 047cbc3726ed95e6baf2403a8c2a1f8e7e467da1 Mon Sep 17 00:00:00 2001
From: Binyang Li <binyli@microsoft.com>
Date: Thu, 16 Jul 2026 05:46:05 +0000
Subject: [PATCH] address comments

---
 python/mscclpp/ep/README.md                  |  15 +-
 python/mscclpp/ep/__init__.py                |   8 +-
 python/mscclpp/ep/communicator.py            |  10 +-
 python/mscclpp/ep/high_throughput.py         |  12 +-
 python/mscclpp/ep/types.py                   |  10 +-
 src/ext/ep/bindings.cpp                      |   1 -
 src/ext/ep/include/api.cuh                   |   4 +-
 src/ext/ep/low_latency/combine.cu            |  18 +-
 src/ext/ep/low_latency/dispatch.cu           | 174 ++++---
 test/python/ep/test_low_latency_multirank.py | 473 ++++++++++++-------
 10 files changed, 445 insertions(+), 280 deletions(-)

diff --git a/python/mscclpp/ep/README.md b/python/mscclpp/ep/README.md
index d6c5da1d..968c5e0f 100644
--- a/python/mscclpp/ep/README.md
+++ b/python/mscclpp/ep/README.md
@@ -180,8 +180,6 @@ Use `DispatchLayout` instead of string literals for this field:
 | `DispatchLayout.TOKEN_MAJOR` | HT: `[total_recv_tokens, hidden]`; LL: `[world_size * max_tokens_per_rank, hidden]` |
 | `DispatchLayout.EXPERT_MAJOR` | `[num_local_experts, max_slots_per_expert, hidden]` |
 
-`DispatchLayout.FLAT` is a compatibility alias for `DispatchLayout.TOKEN_MAJOR`.
-
 ## MoECommunicator methods
 
 ```python
@@ -267,7 +265,6 @@ class QuantConfig:
 class DispatchLayout(str, Enum):
     EXPERT_MAJOR = "expert_major"
     TOKEN_MAJOR = "token_major"
-    FLAT = TOKEN_MAJOR
 
 
 @dataclass
@@ -317,11 +314,11 @@ class TokenMajorCombineContext:
 
 
 @dataclass
-class RowMajorCombineContext:
+class HighThroughputCombineContext:
     ...
 
 
-CombineContext = ExpertMajorCombineContext | TokenMajorCombineContext | RowMajorCombineContext
+CombineContext = ExpertMajorCombineContext | TokenMajorCombineContext | HighThroughputCombineContext
 
 
 class DispatchHandle:
@@ -338,8 +335,8 @@ class TokenMajorDispatchHandle(DispatchHandle):
     combine_context: TokenMajorCombineContext
 
 
-class RowMajorDispatchHandle(DispatchHandle):
-    combine_context: RowMajorCombineContext
+class HighThroughputDispatchHandle(DispatchHandle):
+    combine_context: HighThroughputCombineContext
 
 
 @dataclass
@@ -428,7 +425,7 @@ to reverse dispatch and finish combine. `ExpertMajorDispatchHandle` uses
 `ExpertMajorCombineContext` (`topk_ids`, `weights`, source info, layout ranges,
 shape, and capacity). `TokenMajorDispatchHandle` records source-token IDs,
 per-source-rank counts, and the original routing needed for cross-rank combine.
-Row-major handles use the intranode combine context with
+High-throughput handles use the intranode combine context with
 receive-side weights, source indices, prefix matrices, and send-head tensors.
 The MLP should treat the handle as opaque and pass it back to `combine`.
 
@@ -563,7 +560,7 @@ dispatch_out.topk_ids          # [world_size * max_tokens_per_rank, K], int32 lo
 dispatch_out.weights           # [world_size * max_tokens_per_rank, K], float32
 ```
 
-Non-local top-k entries use expert ID `-1`. The valid row count in each
+Non-local top-k entries use expert ID `-1` and weight `0`. The valid row count in each
 source-rank region is returned in `dispatch_out.layout.num_tokens_per_rank`.
 For expert-major output, only the first
 `dispatch_out.layout.num_tokens_per_expert[i]` slots are valid:
diff --git a/python/mscclpp/ep/__init__.py b/python/mscclpp/ep/__init__.py
index e680bb15..948b2af5 100644
--- a/python/mscclpp/ep/__init__.py
+++ b/python/mscclpp/ep/__init__.py
@@ -22,13 +22,13 @@ from .communicator import (  # noqa: F401
     DispatchOutputInfo,
     ExpertMajorDispatchHandle,
     ExpertMajorCombineContext,
+    HighThroughputDispatchHandle,
+    HighThroughputCombineContext,
     MoECommunicator,
     MoECommunicatorConfig,
     MoEMode,
     OperationOverlapConfig,
     QuantConfig,
-    RowMajorDispatchHandle,
-    RowMajorCombineContext,
     TokenMajorDispatchHandle,
     TokenMajorCombineContext,
 )
@@ -46,13 +46,13 @@ __all__ = [
     "DispatchOutputInfo",
     "ExpertMajorDispatchHandle",
     "ExpertMajorCombineContext",
+    "HighThroughputDispatchHandle",
+    "HighThroughputCombineContext",
     "MoECommunicator",
     "MoECommunicatorConfig",
     "MoEMode",
     "OperationOverlapConfig",
     "QuantConfig",
-    "RowMajorDispatchHandle",
-    "RowMajorCombineContext",
     "TokenMajorDispatchHandle",
     "TokenMajorCombineContext",
 ]
diff --git a/python/mscclpp/ep/communicator.py b/python/mscclpp/ep/communicator.py
index c42f77c9..08d7e82a 100644
--- a/python/mscclpp/ep/communicator.py
+++ b/python/mscclpp/ep/communicator.py
@@ -21,11 +21,11 @@ from .types import (
     DispatchOutputInfo,
     ExpertMajorDispatchHandle,
     ExpertMajorCombineContext,
+    HighThroughputDispatchHandle,
+    HighThroughputCombineContext,
     MoECommunicatorConfig,
     OperationOverlapConfig,
     QuantConfig,
-    RowMajorDispatchHandle,
-    RowMajorCombineContext,
     TokenMajorDispatchHandle,
     TokenMajorCombineContext,
 )
@@ -43,13 +43,13 @@ __all__ = [
     "DispatchOutputInfo",
     "ExpertMajorDispatchHandle",
     "ExpertMajorCombineContext",
+    "HighThroughputDispatchHandle",
+    "HighThroughputCombineContext",
     "MoECommunicator",
     "MoECommunicatorConfig",
     "MoEMode",
     "OperationOverlapConfig",
     "QuantConfig",
-    "RowMajorDispatchHandle",
-    "RowMajorCombineContext",
     "TokenMajorDispatchHandle",
     "TokenMajorCombineContext",
 ]
@@ -166,7 +166,7 @@ def _validate_common_config(config: MoECommunicatorConfig) -> None:
 
 def _resolve_output_layout(layout: Optional[DispatchLayout], mode: MoEMode) -> DispatchLayout:
     if layout is None:
-        return DispatchLayout.EXPERT_MAJOR if mode == MoEMode.LOW_LATENCY else DispatchLayout.FLAT
+        return DispatchLayout.EXPERT_MAJOR if mode == MoEMode.LOW_LATENCY else DispatchLayout.TOKEN_MAJOR
     if not isinstance(layout, DispatchLayout):
         raise TypeError("MoECommunicatorConfig.output_layout must be a DispatchLayout")
     return layout
diff --git a/python/mscclpp/ep/high_throughput.py b/python/mscclpp/ep/high_throughput.py
index f281ea96..c48dce96 100644
--- a/python/mscclpp/ep/high_throughput.py
+++ b/python/mscclpp/ep/high_throughput.py
@@ -24,10 +24,10 @@ from .types import (
     DispatchLayoutInfo,
     DispatchOutput,
     DispatchOutputInfo,
+    HighThroughputCombineContext,
+    HighThroughputDispatchHandle,
     MoECommunicatorConfig,
     QuantConfig,
-    RowMajorCombineContext,
-    RowMajorDispatchHandle,
 )
 from .utils import (
     bf16_view as _bf16_view,
@@ -397,7 +397,7 @@ class HighThroughputBackend:
             recv_topk_idx = cache["recv_topk_idx"]
             recv_topk_weights = cache["recv_topk_weights"]
             num_recv_tokens_per_expert_list = cache["num_recv_tokens_per_expert_list"]
-            combine_context = RowMajorCombineContext(
+            combine_context = HighThroughputCombineContext(
                 recv_topk_weights=recv_topk_weights,
                 src_idx=recv_src_idx,
                 rank_prefix_matrix=rank_prefix_matrix,
@@ -430,7 +430,7 @@ class HighThroughputBackend:
                 None,
                 self.expert_alignment,
             )
-            combine_context = RowMajorCombineContext(
+            combine_context = HighThroughputCombineContext(
                 recv_topk_weights=recv_topk_weights,
                 src_idx=recv_src_idx,
                 rank_prefix_matrix=rank_prefix_matrix,
@@ -471,7 +471,7 @@ class HighThroughputBackend:
             topk_ids=recv_topk_idx,
             weights=recv_topk_weights,
         )
-        handle = RowMajorDispatchHandle(output_info=output_info, combine_context=combine_context)
+        handle = HighThroughputDispatchHandle(output_info=output_info, combine_context=combine_context)
         # The torch-free HT runtime orders its work on the caller's CUDA stream
         # (no separate event handle), so there is nothing to attach here.
         handle._event = None  # type: ignore[attr-defined]
@@ -547,7 +547,7 @@ class HighThroughputBackend:
                 raise ValueError("weights shape must match topk_ids")
 
     def _validate_combine_inputs(self, expert_output, handle) -> None:
-        if not isinstance(handle, RowMajorDispatchHandle):
+        if not isinstance(handle, HighThroughputDispatchHandle):
             raise TypeError("handle must be a DispatchHandle returned by dispatch")
         if expert_output.dim() != 2 or not expert_output.is_contiguous():
             raise ValueError("expert_output must be a contiguous [total_recv_tokens, hidden] tensor")
diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py
index 6bda6304..6c53b652 100644
--- a/python/mscclpp/ep/types.py
+++ b/python/mscclpp/ep/types.py
@@ -131,8 +131,8 @@ class TokenMajorCombineContext:
 
 
 @dataclass
-class RowMajorCombineContext:
-    """Combine context for row-major high-throughput dispatch output."""
+class HighThroughputCombineContext:
+    """Combine context for high-throughput dispatch output."""
 
     recv_topk_weights: Optional[torch.Tensor]
     src_idx: torch.Tensor
@@ -141,7 +141,7 @@ class RowMajorCombineContext:
     send_head: torch.Tensor
 
 
-CombineContext = Union[ExpertMajorCombineContext, TokenMajorCombineContext, RowMajorCombineContext]
+CombineContext = Union[ExpertMajorCombineContext, TokenMajorCombineContext, HighThroughputCombineContext]
 
 
 # Opaque dispatch handles returned by dispatch() and consumed by combine().
@@ -165,8 +165,8 @@ class TokenMajorDispatchHandle(DispatchHandle):
 
 
 @dataclass
-class RowMajorDispatchHandle(DispatchHandle):
-    combine_context: RowMajorCombineContext
+class HighThroughputDispatchHandle(DispatchHandle):
+    combine_context: HighThroughputCombineContext
 
 
 # Optional async/overlap configuration.
diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp
index 0e5a68c3..d28411e1 100644
--- a/src/ext/ep/bindings.cpp
+++ b/src/ext/ep/bindings.cpp
@@ -46,7 +46,6 @@ NB_MODULE(mscclpp_ep_cpp, m) {
 
   nb::enum_<mscclpp::ep::DispatchLayout>(m, "DispatchLayout")
       .value("EXPERT_MAJOR", mscclpp::ep::DispatchLayout::EXPERT_MAJOR)
-      .value("FLAT", mscclpp::ep::DispatchLayout::FLAT)
       .value("TOKEN_MAJOR", mscclpp::ep::DispatchLayout::TOKEN_MAJOR);
 
   nb::enum_<mscclpp::ep::low_latency::CombineMode>(m, "CombineMode")
diff --git a/src/ext/ep/include/api.cuh b/src/ext/ep/include/api.cuh
index a4f41b56..e6ae4e44 100644
--- a/src/ext/ep/include/api.cuh
+++ b/src/ext/ep/include/api.cuh
@@ -32,9 +32,7 @@ enum class DispatchLayout {
   /// Token-major rows. Low latency uses
   /// [num_ranks * max_tokens_per_rank, hidden], grouped by source rank; high
   /// throughput uses [num_recv_tokens, hidden].
-  TOKEN_MAJOR,
-  /// Backward-compatible alias for TOKEN_MAJOR.
-  FLAT = TOKEN_MAJOR
+  TOKEN_MAJOR
 };
 
 // ===========================================================================
diff --git a/src/ext/ep/low_latency/combine.cu b/src/ext/ep/low_latency/combine.cu
index a117c93f..cab3c625 100644
--- a/src/ext/ep/low_latency/combine.cu
+++ b/src/ext/ep/low_latency/combine.cu
@@ -47,13 +47,16 @@ MSCCLPP_HOST_DEVICE_INLINE int directSendWorkerCount(int nLocalExperts) {
   return availableWorkers < DirectSendMaxNWorkers ? availableWorkers : DirectSendMaxNWorkers;
 }
 
-template <int Hidden, low_latency::CombineMode Mode>
+template <int Hidden, low_latency::CombineMode Mode, DispatchLayout Layout>
 MSCCLPP_HOST_DEVICE_INLINE size_t combineSharedBytes(int nLocalExperts) {
   constexpr size_t TileBytes = static_cast<size_t>(Hidden) * sizeof(Bf16);
   if constexpr (Mode == low_latency::CombineMode::DIRECT_SEND) {
     return directSendControlBytes(nLocalExperts) +
            static_cast<size_t>(directSendWorkerCount<Hidden>(nLocalExperts)) * directSendWorkerBytes<Hidden>();
   }
+  if constexpr (Layout == DispatchLayout::TOKEN_MAJOR) {
+    return CombineNStages * TileBytes;
+  }
   return CombineNStages * TileBytes;
 }
 
@@ -188,10 +191,9 @@ MSCCLPP_DEVICE_INLINE void sendRankReducedPartials(const void* expertOutput, int
 }
 
 template <int Hidden>
-MSCCLPP_DEVICE_INLINE void sendTokenMajorPartials(const void* expertOutput, const int* srcInfo, int nRanks,
-                                                  int maxTokensPerRank, void* combineRecvBuffer,
-                                                  const TransportView& transport, WorkspaceView& workspaceView,
-                                                  uint8_t* sharedMemory) {
+MSCCLPP_DEVICE_INLINE void sendTokenMajorPartials(const void* expertOutput, const int* srcInfo, int maxTokensPerRank,
+                                                  void* combineRecvBuffer, const TransportView& transport,
+                                                  WorkspaceView& workspaceView, uint8_t* sharedMemory) {
   const int threadId = static_cast<int>(threadIdx.x);
   constexpr size_t HiddenBytes = static_cast<size_t>(Hidden) * sizeof(Bf16);
   constexpr int HiddenInt4 = HiddenBytes / sizeof(int4);
@@ -436,8 +438,8 @@ __global__ __launch_bounds__(CombineNThreads, 1) void combineKernel(
 
   if constexpr (Layout == DispatchLayout::TOKEN_MAJOR) {
     static_assert(Mode == low_latency::CombineMode::RANK_LOCAL_REDUCE);
-    sendTokenMajorPartials<Hidden>(expertOutput, srcInfo, nRanks, maxTokensPerRank, combineRecvBuffer, transport,
-                                   workspaceView, sharedMemory);
+    sendTokenMajorPartials<Hidden>(expertOutput, srcInfo, maxTokensPerRank, combineRecvBuffer, transport, workspaceView,
+                                   sharedMemory);
   } else if constexpr (Mode == low_latency::CombineMode::RANK_LOCAL_REDUCE) {
     sendRankReducedPartials<Hidden, DispatchType, ScaleBlockSize>(
         expertOutput, nExperts, nRanks, nTopk, maxTokensPerRank, combineRecvBuffer, dispatchRecvBuffer, transport,
@@ -476,7 +478,7 @@ inline void combineHiddenMode(void* output, const void* expertOutput, const int6
   }
 
   auto combineFunc = combineKernel<Mode, Hidden, DispatchType, ScaleBlockSize, Layout>;
-  const size_t sharedBytes = combineSharedBytes<Hidden, Mode>(nLocalExperts);
+  const size_t sharedBytes = combineSharedBytes<Hidden, Mode, Layout>(nLocalExperts);
   static thread_local KernelConfigCache kernelConfig;
   const int residentBlocks = configureKernel(combineFunc, CombineNThreads, sharedBytes, comm, kernelConfig);
   EP_HOST_ASSERT(residentBlocks >= numBlocks);
diff --git a/src/ext/ep/low_latency/dispatch.cu b/src/ext/ep/low_latency/dispatch.cu
index 0d7dfa58..659409d2 100644
--- a/src/ext/ep/low_latency/dispatch.cu
+++ b/src/ext/ep/low_latency/dispatch.cu
@@ -413,6 +413,99 @@ MSCCLPP_DEVICE_INLINE bool acquireRecvTask(RecvTask& task, WorkspaceView& worksp
   return task.sourceRank_ >= 0;
 }
 
+template <int Hidden, DispatchDataType DataType, int ScaleBlockSize>
+MSCCLPP_DEVICE_INLINE bool dispatchRecvExpertMajorOutput(
+    void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout,
+    const DispatchPayloadView<DataType>& payloadView, void* sourcePayload, int localExpertIdx, int sourceRank,
+    int sourceTokenIdx, int nLocalExperts, int nRanks, int nTopk, int maxTokensPerRank, WorkspaceView& workspaceView,
+    uint8_t* sharedTile, uint64_t* tmaBarrier, uint32_t& recvTmaPhase) {
+  using OutputType = DispatchElementType<DataType>;
+  constexpr size_t OutputBytes = static_cast<size_t>(Hidden) * sizeof(OutputType);
+  constexpr int NumScales = DataType == DispatchDataType::BF16 ? 0 : Hidden / ScaleBlockSize;
+  const int laneId = get_lane_id();
+  const int nOutputSlotsPerExpert = nRanks * maxTokensPerRank;
+  int outputTokenIdx = -1;
+  int combineInputOffset = -1;
+
+  if (localExpertIdx >= 0) {
+    int expertTokenCount;
+    int outputOffset;
+    unpack2(outputLayout[localExpertIdx * nRanks + sourceRank], expertTokenCount, outputOffset);
+    const int copiedTokenIdx =
+        atomicAdd(workspaceView.dispatchExpertCopiedCounts_ + sourceRank * nLocalExperts + localExpertIdx, 1);
+    EP_DEVICE_ASSERT(copiedTokenIdx < expertTokenCount);
+    if (copiedTokenIdx == expertTokenCount - 1) {
+      workspaceView.dispatchExpertCopiedCounts_[sourceRank * nLocalExperts + localExpertIdx] = 0;
+    }
+    outputTokenIdx = outputOffset + copiedTokenIdx;
+    outputSrcInfo[static_cast<size_t>(localExpertIdx) * nOutputSlotsPerExpert + outputTokenIdx] = sourceTokenIdx;
+    combineInputOffset = localExpertIdx * nOutputSlotsPerExpert + outputTokenIdx;
+  }
+
+  if constexpr (DataType != DispatchDataType::BF16) {
+    const auto* sourceScales = payloadView.scaleFactors(sourcePayload);
+    // Each top-k lane may create a row for a different local expert. All lanes
+    // cooperate to copy the shared payload's scale vector to every such row.
+    for (int topkLane = 0; topkLane < nTopk; ++topkLane) {
+      const int scaleLocalExpertIdx = warpBroadcast(localExpertIdx, topkLane);
+      const int scaleOutputTokenIdx = warpBroadcast(outputTokenIdx, topkLane);
+      if (scaleLocalExpertIdx < 0) continue;
+      for (int scaleIdx = laneId; scaleIdx < NumScales; scaleIdx += WARP_SIZE) {
+        outputScales[(static_cast<size_t>(scaleLocalExpertIdx) * NumScales + scaleIdx) * nOutputSlotsPerExpert +
+                     scaleOutputTokenIdx] = sourceScales[scaleIdx];
+      }
+    }
+  }
+  if (laneId < nTopk) payloadView.topKIndices(sourcePayload)[laneId] = combineInputOffset;
+
+  if (laneId == 0) waitTmaLoad(tmaBarrier, recvTmaPhase);
+  __syncwarp();
+  fenceProxyAsyncSharedCta();
+
+  if (localExpertIdx < 0) return false;
+  auto* outputData = reinterpret_cast<uint8_t*>(output) +
+                     (static_cast<size_t>(localExpertIdx) * nOutputSlotsPerExpert + outputTokenIdx) * OutputBytes;
+  issueTmaStore(outputData, sharedTile, static_cast<uint32_t>(OutputBytes));
+  return true;
+}
+
+template <int Hidden, DispatchDataType DataType, int ScaleBlockSize>
+MSCCLPP_DEVICE_INLINE bool dispatchRecvTokenMajorOutput(void* output, float* outputScales, int* outputSrcInfo,
+                                                        int* outputTopkIdx, float* outputTopkWeights,
+                                                        const DispatchPayloadView<DataType>& payloadView,
+                                                        const void* sourcePayload, int localExpertIdx, int sourceRank,
+                                                        int sourceTokenSlot, int sourceTokenIdx, int nTopk,
+                                                        int maxTokensPerRank, uint8_t* sharedTile, uint64_t* tmaBarrier,
+                                                        uint32_t& recvTmaPhase) {
+  using OutputType = DispatchElementType<DataType>;
+  constexpr size_t OutputBytes = static_cast<size_t>(Hidden) * sizeof(OutputType);
+  constexpr int NumScales = DataType == DispatchDataType::BF16 ? 0 : Hidden / ScaleBlockSize;
+  const int laneId = get_lane_id();
+  const int outputRow = sourceRank * maxTokensPerRank + sourceTokenSlot;
+
+  if (laneId == 0) outputSrcInfo[outputRow] = sourceTokenIdx;
+  if (laneId < nTopk) {
+    outputTopkIdx[static_cast<size_t>(outputRow) * nTopk + laneId] = localExpertIdx;
+    outputTopkWeights[static_cast<size_t>(outputRow) * nTopk + laneId] =
+        localExpertIdx >= 0 ? payloadView.topKValues(sourcePayload)[laneId] : 0.0f;
+  }
+  if constexpr (DataType != DispatchDataType::BF16) {
+    const auto* sourceScales = payloadView.scaleFactors(sourcePayload);
+    for (int scaleIdx = laneId; scaleIdx < NumScales; scaleIdx += WARP_SIZE) {
+      outputScales[static_cast<size_t>(outputRow) * NumScales + scaleIdx] = sourceScales[scaleIdx];
+    }
+  }
+
+  if (laneId == 0) waitTmaLoad(tmaBarrier, recvTmaPhase);
+  __syncwarp();
+  fenceProxyAsyncSharedCta();
+
+  if (laneId != 0) return false;
+  auto* outputData = reinterpret_cast<uint8_t*>(output) + static_cast<size_t>(outputRow) * OutputBytes;
+  issueTmaStore(outputData, sharedTile, static_cast<uint32_t>(OutputBytes));
+  return true;
+}
+
 template <int Hidden, DispatchDataType DataType, int ScaleBlockSize, DispatchLayout Layout>
 MSCCLPP_DEVICE_INLINE void dispatchRecvWorker(void* output, float* outputScales, int* outputSrcInfo, int* outputTopkIdx,
                                               float* outputTopkWeights, int64_t* outputLayout, int nExperts, int rank,
@@ -439,8 +532,6 @@ MSCCLPP_DEVICE_INLINE void dispatchRecvWorker(void* output, float* outputScales,
   const size_t payloadStride = dispatchPayloadStride<DataType>(Hidden, nTopk, ScaleBlockSize);
   constexpr size_t OutputBytes = static_cast<size_t>(Hidden) * sizeof(OutputType);
   constexpr size_t TileBytes = OutputBytes;
-  constexpr int NumScales = DataType == DispatchDataType::BF16 ? 0 : Hidden / ScaleBlockSize;
-  const int nOutputSlotsPerExpert = nRanks * maxTokensPerRank;
   auto* sourcePayloadBase = reinterpret_cast<uint8_t*>(recvBuffer) + dispatchMetadataBytes(nRanks, nExperts) +
                             static_cast<size_t>(sourceRank) * maxTokensPerRank * payloadStride;
   auto* tmaTiles = reinterpret_cast<uint8_t*>(sharedMem) + dispatchSharedControlBytes(nRanks);
@@ -470,79 +561,16 @@ MSCCLPP_DEVICE_INLINE void dispatchRecvWorker(void* output, float* outputScales,
                                    : -1;
     const int sourceTokenIdx = warpBroadcast(
         laneId == 0 ? *payloadView.srcTokenGlobalIdx(sourcePayload) - sourceRank * maxTokensPerRank : 0, 0);
-    int outputTokenIdx = -1;
-    int combineInputOffset = -1;
     if constexpr (Layout == DispatchLayout::EXPERT_MAJOR) {
-      if (localExpertIdx >= 0) {
-        int expertTokenCount;
-        int outputOffset;
-        unpack2(outputLayout[localExpertIdx * nRanks + sourceRank], expertTokenCount, outputOffset);
-        const int copiedTokenIdx =
-            atomicAdd(workspaceView.dispatchExpertCopiedCounts_ + sourceRank * nLocalExperts + localExpertIdx, 1);
-        EP_DEVICE_ASSERT(copiedTokenIdx < expertTokenCount);
-        if (copiedTokenIdx == expertTokenCount - 1) {
-          workspaceView.dispatchExpertCopiedCounts_[sourceRank * nLocalExperts + localExpertIdx] = 0;
-        }
-        outputTokenIdx = outputOffset + copiedTokenIdx;
-        outputSrcInfo[static_cast<size_t>(localExpertIdx) * nOutputSlotsPerExpert + outputTokenIdx] = sourceTokenIdx;
-        combineInputOffset = localExpertIdx * nOutputSlotsPerExpert + outputTokenIdx;
-      }
+      hasPendingStore = dispatchRecvExpertMajorOutput<Hidden, DataType, ScaleBlockSize>(
+          output, outputScales, outputSrcInfo, outputLayout, payloadView, sourcePayload, localExpertIdx, sourceRank,
+          sourceTokenIdx, nLocalExperts, nRanks, nTopk, maxTokensPerRank, workspaceView, sharedTile, tmaBarrier,
+          recvTmaPhase);
     } else {
-      const int outputRow = sourceRank * maxTokensPerRank + sourceTokenSlot;
-      if (laneId == 0) outputSrcInfo[outputRow] = sourceTokenIdx;
-      if (laneId < nTopk) {
-        outputTopkIdx[static_cast<size_t>(outputRow) * nTopk + laneId] = localExpertIdx;
-        outputTopkWeights[static_cast<size_t>(outputRow) * nTopk + laneId] =
-            payloadView.topKValues(sourcePayload)[laneId];
-      }
-      outputTokenIdx = sourceTokenSlot;
-      combineInputOffset = outputRow;
-    }
-    if constexpr (DataType != DispatchDataType::BF16) {
-      const auto* sourceScales = payloadView.scaleFactors(sourcePayload);
-      if constexpr (Layout == DispatchLayout::EXPERT_MAJOR) {
-        // Each top-k lane may create a row for a different local expert. All lanes
-        // cooperate to copy the shared payload's scale vector to every such row.
-        for (int topkLane = 0; topkLane < nTopk; ++topkLane) {
-          const int scaleLocalExpertIdx = warpBroadcast(localExpertIdx, topkLane);
-          const int scaleOutputTokenIdx = warpBroadcast(outputTokenIdx, topkLane);
-          if (scaleLocalExpertIdx < 0) continue;
-          for (int scaleIdx = laneId; scaleIdx < NumScales; scaleIdx += WARP_SIZE) {
-            outputScales[(static_cast<size_t>(scaleLocalExpertIdx) * NumScales + scaleIdx) * nOutputSlotsPerExpert +
-                         scaleOutputTokenIdx] = sourceScales[scaleIdx];
-          }
-        }
-      } else {
-        for (int scaleIdx = laneId; scaleIdx < NumScales; scaleIdx += WARP_SIZE) {
-          outputScales[static_cast<size_t>(combineInputOffset) * NumScales + scaleIdx] = sourceScales[scaleIdx];
-        }
-      }
-    }
-    if constexpr (Layout == DispatchLayout::EXPERT_MAJOR) {
-      if (laneId < nTopk) payloadView.topKIndices(sourcePayload)[laneId] = combineInputOffset;
-    }
-
-    if (laneId == 0) waitTmaLoad(tmaBarrier, recvTmaPhase);
-    __syncwarp();
-    fenceProxyAsyncSharedCta();
-
-    if constexpr (Layout == DispatchLayout::EXPERT_MAJOR) {
-      if (localExpertIdx >= 0) {
-        auto* outputData = reinterpret_cast<uint8_t*>(output) +
-                           (static_cast<size_t>(localExpertIdx) * nOutputSlotsPerExpert + outputTokenIdx) * OutputBytes;
-        issueTmaStore(outputData, sharedTile, static_cast<uint32_t>(OutputBytes));
-        hasPendingStore = true;
-      } else {
-        hasPendingStore = false;
-      }
-    } else {
-      if (laneId == 0) {
-        auto* outputData = reinterpret_cast<uint8_t*>(output) + static_cast<size_t>(combineInputOffset) * OutputBytes;
-        issueTmaStore(outputData, sharedTile, static_cast<uint32_t>(OutputBytes));
-        hasPendingStore = true;
-      } else {
-        hasPendingStore = false;
-      }
+      hasPendingStore = dispatchRecvTokenMajorOutput<Hidden, DataType, ScaleBlockSize>(
+          output, outputScales, outputSrcInfo, outputTopkIdx, outputTopkWeights, payloadView, sourcePayload,
+          localExpertIdx, sourceRank, sourceTokenSlot, sourceTokenIdx, nTopk, maxTokensPerRank, sharedTile, tmaBarrier,
+          recvTmaPhase);
     }
   }
 
diff --git a/test/python/ep/test_low_latency_multirank.py b/test/python/ep/test_low_latency_multirank.py
index 139e10b0..d012f4a8 100644
--- a/test/python/ep/test_low_latency_multirank.py
+++ b/test/python/ep/test_low_latency_multirank.py
@@ -160,6 +160,255 @@ def validate_combine_output(actual, expected, *, exact, group):
     return local_diff.item(), global_diff.item()
 
 
+def validate_expert_major_dispatch(
+    *,
+    rank,
+    num_ranks,
+    hidden,
+    num_local_experts,
+    dispatch_quant,
+    dispatch_out,
+    handle,
+    packed_recv_x,
+    packed_recv_count,
+    packed_recv_layout_range,
+    all_topk_idx,
+    all_x,
+    expected_scales,
+):
+    int_mask = (1 << 32) - 1
+    for local_expert_idx in range(num_local_experts):
+        expert_id = rank * num_local_experts + local_expert_idx
+        recv_count = int(packed_recv_count[local_expert_idx].item())
+        expected_count = int((all_topk_idx == expert_id).sum().item())
+        recv_layout_range = packed_recv_layout_range[local_expert_idx]
+        layout_sum = int((recv_layout_range & int_mask).sum().item())
+        assert (
+            recv_count == expected_count
+        ), f"rank{rank} expert{expert_id}: recv_count={recv_count} != expected={expected_count}"
+        assert (
+            layout_sum == recv_count
+        ), f"rank{rank} expert{expert_id}: layout range sum {layout_sum} != recv_count {recv_count}"
+
+        if recv_count == 0:
+            continue
+        recv_x = packed_recv_x[local_expert_idx, :recv_count]
+        if dispatch_quant is None:
+            recv_x_lo = recv_x[:, :-128]
+            assert torch.equal(
+                recv_x_lo.amin(dim=-1), recv_x_lo.amax(dim=-1)
+            ), f"rank{rank} expert{expert_id}: non-uniform recv block"
+            continue
+
+        assert all_x is not None
+        assert expected_scales is not None
+        assert dispatch_out.quant is not None
+        assert dispatch_out.quant.block_scales is not None
+        for source_rank in range(num_ranks):
+            packed_range = int(recv_layout_range[source_rank].item())
+            source_count = packed_range & int_mask
+            output_offset = packed_range >> 32
+            if source_count == 0:
+                continue
+            source_tokens = handle.combine_context.src_info[
+                local_expert_idx, output_offset : output_offset + source_count
+            ].long()
+            actual_tokens = recv_x[output_offset : output_offset + source_count]
+            actual_scales = dispatch_out.quant.block_scales[
+                local_expert_idx, output_offset : output_offset + source_count
+            ]
+            reference_scales = expected_scales[source_rank, source_tokens]
+            torch.testing.assert_close(actual_scales, reference_scales, rtol=1e-6, atol=1e-7)
+            actual_dequantized = actual_tokens.float().reshape(
+                source_count, hidden // 128, 128
+            ) * actual_scales.unsqueeze(-1)
+            reference_tokens = all_x[source_rank, source_tokens].float().reshape(source_count, hidden // 128, 128)
+            quant_error = (actual_dequantized - reference_tokens).abs()
+            quant_error_bound = reference_scales.unsqueeze(-1) * 16.1 + 1e-6
+            max_scale_error = (quant_error / reference_scales.unsqueeze(-1)).max().item()
+            assert torch.all(quant_error <= quant_error_bound), (
+                f"rank{rank} expert{expert_id}: FP8 payload mismatch from rank {source_rank}, "
+                f"max scale error={max_scale_error}"
+            )
+
+
+def validate_token_major_dispatch(
+    *,
+    rank,
+    num_ranks,
+    num_tokens,
+    hidden,
+    num_topk,
+    num_local_experts,
+    dispatch_quant,
+    dispatch_out,
+    handle,
+    packed_recv_x,
+    packed_recv_count,
+    all_topk_idx,
+    all_topk_weights,
+    all_x,
+    expected_scales,
+):
+    assert all_x is not None
+    assert dispatch_out.topk_ids is not None
+    assert dispatch_out.topk_ids.dtype == torch.int32
+    assert dispatch_out.topk_ids.shape == (num_ranks * num_tokens, num_topk)
+    assert dispatch_out.weights is not None
+    assert dispatch_out.weights.shape == (num_ranks * num_tokens, num_topk)
+    source_token_ids = handle.combine_context.source_token_ids
+    assert source_token_ids.shape == (num_ranks * num_tokens,)
+    local_expert_begin = rank * num_local_experts
+    local_expert_end = local_expert_begin + num_local_experts
+
+    for source_rank in range(num_ranks):
+        recv_count = int(packed_recv_count[source_rank].item())
+        source_routing = all_topk_idx[source_rank]
+        expected_source_tokens = (
+            ((source_routing >= local_expert_begin) & (source_routing < local_expert_end))
+            .any(dim=1)
+            .nonzero()
+            .flatten()
+        )
+        assert recv_count == expected_source_tokens.numel()
+        if recv_count == 0:
+            continue
+
+        row_begin = source_rank * num_tokens
+        row_end = row_begin + recv_count
+        source_tokens = source_token_ids[row_begin:row_end].long()
+        assert torch.equal(torch.sort(source_tokens).values, expected_source_tokens)
+
+        actual_tokens = packed_recv_x[row_begin:row_end]
+        actual_topk_ids = dispatch_out.topk_ids[row_begin:row_end]
+        actual_weights = dispatch_out.weights[row_begin:row_end]
+        expected_global_ids = all_topk_idx[source_rank, source_tokens]
+        local_mask = (expected_global_ids >= local_expert_begin) & (expected_global_ids < local_expert_end)
+        expected_local_ids = torch.where(
+            local_mask, expected_global_ids - local_expert_begin, torch.full_like(expected_global_ids, -1)
+        )
+        expected_weights = torch.where(
+            local_mask, all_topk_weights[source_rank, source_tokens], torch.zeros_like(actual_weights)
+        )
+        assert torch.equal(actual_topk_ids, expected_local_ids.to(torch.int32))
+        torch.testing.assert_close(actual_weights, expected_weights)
+
+        if dispatch_quant is None:
+            assert torch.equal(actual_tokens, all_x[source_rank, source_tokens])
+            continue
+
+        assert expected_scales is not None
+        assert dispatch_out.quant is not None
+        assert dispatch_out.quant.block_scales is not None
+        actual_scales = dispatch_out.quant.block_scales[row_begin:row_end]
+        reference_scales = expected_scales[source_rank, source_tokens]
+        torch.testing.assert_close(actual_scales, reference_scales, rtol=1e-6, atol=1e-7)
+        actual_dequantized = actual_tokens.float().reshape(recv_count, hidden // 128, 128) * actual_scales.unsqueeze(-1)
+        reference_tokens = all_x[source_rank, source_tokens].float().reshape(recv_count, hidden // 128, 128)
+        quant_error = (actual_dequantized - reference_tokens).abs()
+        quant_error_bound = reference_scales.unsqueeze(-1) * 16.1 + 1e-6
+        assert torch.all(quant_error <= quant_error_bound)
+
+
+def reconstruct_expert_major_reference(
+    *,
+    rank,
+    num_ranks,
+    num_tokens,
+    hidden,
+    num_local_experts,
+    all_topk_idx,
+    packed_recv_layout_range,
+    handle,
+    dequantized_x,
+    group,
+):
+    int_mask = (1 << 32) - 1
+    first_expert = all_topk_idx.gather(-1, (all_topk_idx >= 0).to(torch.int32).argmax(dim=-1, keepdim=True)).squeeze(-1)
+    first_expert.masked_fill_(~(all_topk_idx >= 0).any(dim=-1), -1)
+    dispatched_reference_x = torch.zeros((num_ranks, num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
+    for local_expert_idx in range(num_local_experts):
+        expert_id = rank * num_local_experts + local_expert_idx
+        recv_layout_range = packed_recv_layout_range[local_expert_idx]
+        for source_rank in range(num_ranks):
+            packed_range = int(recv_layout_range[source_rank].item())
+            source_count = packed_range & int_mask
+            output_offset = packed_range >> 32
+            if source_count == 0:
+                continue
+            source_tokens = handle.combine_context.src_info[
+                local_expert_idx, output_offset : output_offset + source_count
+            ].long()
+            selected = first_expert[source_rank, source_tokens] == expert_id
+            dispatched_reference_x[source_rank, source_tokens[selected]] = dequantized_x[
+                local_expert_idx, output_offset : output_offset + source_count
+            ][selected]
+    dist.all_reduce(dispatched_reference_x, group=group)
+    return dispatched_reference_x[rank]
+
+
+def reconstruct_token_major_reference(
+    *,
+    rank,
+    num_ranks,
+    num_tokens,
+    hidden,
+    num_local_experts,
+    all_topk_idx,
+    packed_recv_count,
+    handle,
+    dequantized_x,
+    group,
+):
+    source_token_ids = handle.combine_context.source_token_ids
+    destination_ranks = torch.where(
+        all_topk_idx >= 0,
+        all_topk_idx // num_local_experts,
+        torch.full_like(all_topk_idx, num_ranks),
+    )
+    first_destination_rank = destination_ranks.amin(dim=-1)
+    dispatched_reference_x = torch.zeros((num_ranks, num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
+    for source_rank in range(num_ranks):
+        source_count = int(packed_recv_count[source_rank].item())
+        if source_count == 0:
+            continue
+        row_begin = source_rank * num_tokens
+        row_end = row_begin + source_count
+        source_tokens = source_token_ids[row_begin:row_end].long()
+        selected = first_destination_rank[source_rank, source_tokens] == rank
+        dispatched_reference_x[source_rank, source_tokens[selected]] = dequantized_x[row_begin:row_end][selected]
+    dist.all_reduce(dispatched_reference_x, group=group)
+    return dispatched_reference_x[rank]
+
+
+def expected_expert_major_output(reference_x, topk_idx, topk_weights):
+    expected = torch.zeros_like(reference_x, dtype=torch.float32)
+    reference_x_f = reference_x.float()
+    for topk_slot in range(topk_idx.size(1)):
+        weight = (
+            (topk_idx[:, topk_slot] != -1).float()
+            if topk_weights is None
+            else topk_weights[:, topk_slot].masked_fill(topk_idx[:, topk_slot] == -1, 0.0)
+        ).view(-1, 1)
+        expected = torch.addcmul(expected, reference_x_f, weight)
+    return expected.to(torch.bfloat16)
+
+
+def expected_token_major_output(reference_x, topk_idx, topk_weights, num_ranks, num_local_experts):
+    expected = torch.zeros_like(reference_x, dtype=torch.float32)
+    reference_x_f = reference_x.float()
+    for destination_rank in range(num_ranks):
+        rank_partial = torch.zeros_like(reference_x, dtype=torch.float32)
+        for topk_slot in range(topk_idx.size(1)):
+            selected = (topk_idx[:, topk_slot] >= 0) & (topk_idx[:, topk_slot] // num_local_experts == destination_rank)
+            weight = (
+                selected.float() if topk_weights is None else topk_weights[:, topk_slot].masked_fill(~selected, 0.0)
+            ).view(-1, 1)
+            rank_partial = torch.addcmul(rank_partial, reference_x_f, weight)
+        expected += rank_partial.to(torch.bfloat16).float()
+    return expected.to(torch.bfloat16)
+
+
 def main():
     args = parse_args()
     rank, num_ranks, _, group = init_dist()
@@ -296,115 +545,41 @@ def main():
         assert all_x is not None
         expected_scales = fp8_e4m3_block128_scales(all_x)
 
-    int_mask = (1 << 32) - 1
     if output_layout == ep.DispatchLayout.EXPERT_MAJOR:
         assert packed_recv_layout_range is not None
-        for i in range(num_local_experts):
-            expert_id = rank * num_local_experts + i
-            recv_count = int(packed_recv_count[i].item())
-            expected_count = int((all_topk_idx == expert_id).sum().item())
-            recv_layout_range = packed_recv_layout_range[i]
-            layout_sum = int((recv_layout_range & int_mask).sum().item())
-            assert (
-                recv_count == expected_count
-            ), f"rank{rank} expert{expert_id}: recv_count={recv_count} != expected={expected_count}"
-            assert (
-                layout_sum == recv_count
-            ), f"rank{rank} expert{expert_id}: layout range sum {layout_sum} != recv_count {recv_count}"
-
-            if recv_count:
-                recv_x = packed_recv_x[i, :recv_count]
-                if dispatch_quant is None:
-                    recv_x_lo = recv_x[:, :-128]
-                    amin = recv_x_lo.amin(dim=-1)
-                    amax = recv_x_lo.amax(dim=-1)
-                    assert torch.equal(amin, amax), f"rank{rank} expert{expert_id}: non-uniform recv block"
-                else:
-                    assert all_x is not None
-                    assert expected_scales is not None
-                    assert dispatch_out.quant is not None
-                    assert dispatch_out.quant.block_scales is not None
-                    for source_rank in range(num_ranks):
-                        packed_range = int(recv_layout_range[source_rank].item())
-                        source_count = packed_range & int_mask
-                        output_offset = packed_range >> 32
-                        if source_count == 0:
-                            continue
-                        source_tokens = handle.combine_context.src_info[
-                            i, output_offset : output_offset + source_count
-                        ].long()
-                        actual_tokens = recv_x[output_offset : output_offset + source_count]
-                        actual_scales = dispatch_out.quant.block_scales[i, output_offset : output_offset + source_count]
-                        reference_scales = expected_scales[source_rank, source_tokens]
-                        torch.testing.assert_close(actual_scales, reference_scales, rtol=1e-6, atol=1e-7)
-                        actual_dequantized = actual_tokens.float().reshape(
-                            source_count, hidden // 128, 128
-                        ) * actual_scales.unsqueeze(-1)
-                        reference_tokens = (
-                            all_x[source_rank, source_tokens].float().reshape(source_count, hidden // 128, 128)
-                        )
-                        quant_error = (actual_dequantized - reference_tokens).abs()
-                        quant_error_bound = reference_scales.unsqueeze(-1) * 16.1 + 1e-6
-                        max_scale_error = (quant_error / reference_scales.unsqueeze(-1)).max().item()
-                        assert torch.all(quant_error <= quant_error_bound), (
-                            f"rank{rank} expert{expert_id}: FP8 payload mismatch from rank {source_rank}, "
-                            f"max scale error={max_scale_error}"
-                        )
+        validate_expert_major_dispatch(
+            rank=rank,
+            num_ranks=num_ranks,
+            hidden=hidden,
+            num_local_experts=num_local_experts,
+            dispatch_quant=dispatch_quant,
+            dispatch_out=dispatch_out,
+            handle=handle,
+            packed_recv_x=packed_recv_x,
+            packed_recv_count=packed_recv_count,
+            packed_recv_layout_range=packed_recv_layout_range,
+            all_topk_idx=all_topk_idx,
+            all_x=all_x,
+            expected_scales=expected_scales,
+        )
     else:
-        assert all_x is not None
-        assert dispatch_out.topk_ids is not None
-        assert dispatch_out.topk_ids.dtype == torch.int32
-        assert dispatch_out.topk_ids.shape == (num_ranks * num_tokens, num_topk)
-        assert dispatch_out.weights is not None
-        assert dispatch_out.weights.shape == (num_ranks * num_tokens, num_topk)
-        source_token_ids = handle.combine_context.source_token_ids
-        assert source_token_ids.shape == (num_ranks * num_tokens,)
-        local_expert_begin = rank * num_local_experts
-        local_expert_end = local_expert_begin + num_local_experts
-        for source_rank in range(num_ranks):
-            recv_count = int(packed_recv_count[source_rank].item())
-            source_routing = all_topk_idx[source_rank]
-            expected_source_tokens = (
-                ((source_routing >= local_expert_begin) & (source_routing < local_expert_end))
-                .any(dim=1)
-                .nonzero()
-                .flatten()
-            )
-            assert recv_count == expected_source_tokens.numel()
-            if recv_count == 0:
-                continue
-            row_begin = source_rank * num_tokens
-            row_end = row_begin + recv_count
-            source_tokens = source_token_ids[row_begin:row_end].long()
-            assert torch.equal(torch.sort(source_tokens).values, expected_source_tokens)
-
-            actual_tokens = packed_recv_x[row_begin:row_end]
-            actual_topk_ids = dispatch_out.topk_ids[row_begin:row_end]
-            actual_weights = dispatch_out.weights[row_begin:row_end]
-            expected_global_ids = all_topk_idx[source_rank, source_tokens]
-            expected_local_ids = expected_global_ids - local_expert_begin
-            expected_local_ids.masked_fill_(
-                (expected_global_ids < local_expert_begin) | (expected_global_ids >= local_expert_end), -1
-            )
-            assert torch.equal(actual_topk_ids, expected_local_ids.to(torch.int32))
-            torch.testing.assert_close(actual_weights, all_topk_weights[source_rank, source_tokens])
-
-            if dispatch_quant is None:
-                assert torch.equal(actual_tokens, all_x[source_rank, source_tokens])
-            else:
-                assert expected_scales is not None
-                assert dispatch_out.quant is not None
-                assert dispatch_out.quant.block_scales is not None
-                actual_scales = dispatch_out.quant.block_scales[row_begin:row_end]
-                reference_scales = expected_scales[source_rank, source_tokens]
-                torch.testing.assert_close(actual_scales, reference_scales, rtol=1e-6, atol=1e-7)
-                actual_dequantized = actual_tokens.float().reshape(
-                    recv_count, hidden // 128, 128
-                ) * actual_scales.unsqueeze(-1)
-                reference_tokens = all_x[source_rank, source_tokens].float().reshape(recv_count, hidden // 128, 128)
-                quant_error = (actual_dequantized - reference_tokens).abs()
-                quant_error_bound = reference_scales.unsqueeze(-1) * 16.1 + 1e-6
-                assert torch.all(quant_error <= quant_error_bound)
+        validate_token_major_dispatch(
+            rank=rank,
+            num_ranks=num_ranks,
+            num_tokens=num_tokens,
+            hidden=hidden,
+            num_topk=num_topk,
+            num_local_experts=num_local_experts,
+            dispatch_quant=dispatch_quant,
+            dispatch_out=dispatch_out,
+            handle=handle,
+            packed_recv_x=packed_recv_x,
+            packed_recv_count=packed_recv_count,
+            all_topk_idx=all_topk_idx,
+            all_topk_weights=all_topk_weights,
+            all_x=all_x,
+            expected_scales=expected_scales,
+        )
 
     if rank == 0:
         print(f"[dispatch] OK (ranks={num_ranks})", flush=True)
@@ -416,77 +591,43 @@ def main():
     simulated_gemm_x = simulated_gemm_output(dispatch_out)
     reference_x = x
     if dispatch_quant is not None:
-        dispatched_reference_x = torch.zeros((num_ranks, num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
         if output_layout == ep.DispatchLayout.EXPERT_MAJOR:
             assert packed_recv_layout_range is not None
-            first_expert = all_topk_idx.gather(
-                -1, (all_topk_idx >= 0).to(torch.int32).argmax(dim=-1, keepdim=True)
-            ).squeeze(-1)
-            first_expert.masked_fill_(~(all_topk_idx >= 0).any(dim=-1), -1)
-            for i in range(num_local_experts):
-                expert_id = rank * num_local_experts + i
-                recv_layout_range = packed_recv_layout_range[i]
-                for source_rank in range(num_ranks):
-                    packed_range = int(recv_layout_range[source_rank].item())
-                    source_count = packed_range & int_mask
-                    output_offset = packed_range >> 32
-                    if source_count == 0:
-                        continue
-                    source_tokens = handle.combine_context.src_info[
-                        i, output_offset : output_offset + source_count
-                    ].long()
-                    selected = first_expert[source_rank, source_tokens] == expert_id
-                    dispatched_reference_x[source_rank, source_tokens[selected]] = dequantized_x[
-                        i, output_offset : output_offset + source_count
-                    ][selected]
-        else:
-            source_token_ids = handle.combine_context.source_token_ids
-            destination_ranks = torch.where(
-                all_topk_idx >= 0,
-                all_topk_idx // num_local_experts,
-                torch.full_like(all_topk_idx, num_ranks),
+            reference_x = reconstruct_expert_major_reference(
+                rank=rank,
+                num_ranks=num_ranks,
+                num_tokens=num_tokens,
+                hidden=hidden,
+                num_local_experts=num_local_experts,
+                all_topk_idx=all_topk_idx,
+                packed_recv_layout_range=packed_recv_layout_range,
+                handle=handle,
+                dequantized_x=dequantized_x,
+                group=group,
+            )
+        else:
+            reference_x = reconstruct_token_major_reference(
+                rank=rank,
+                num_ranks=num_ranks,
+                num_tokens=num_tokens,
+                hidden=hidden,
+                num_local_experts=num_local_experts,
+                all_topk_idx=all_topk_idx,
+                packed_recv_count=packed_recv_count,
+                handle=handle,
+                dequantized_x=dequantized_x,
+                group=group,
             )
-            first_destination_rank = destination_ranks.amin(dim=-1)
-            for source_rank in range(num_ranks):
-                source_count = int(packed_recv_count[source_rank].item())
-                if source_count == 0:
-                    continue
-                row_begin = source_rank * num_tokens
-                row_end = row_begin + source_count
-                source_tokens = source_token_ids[row_begin:row_end].long()
-                selected = first_destination_rank[source_rank, source_tokens] == rank
-                dispatched_reference_x[source_rank, source_tokens[selected]] = dequantized_x[row_begin:row_end][
-                    selected
-                ]
-        dist.all_reduce(dispatched_reference_x, group=group)
-        reference_x = dispatched_reference_x[rank]
     out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
     combined_x = moe_comm.combine(simulated_gemm_x, handle, out=out)
 
     # Analytical expected: each token i, weighted sum over topk entries that
     # are not -1. Accumulate in the same top-k order as the kernel; multiplying
     # by the pre-summed weights can differ by one BF16 ULP for large token IDs.
-    x_f = reference_x.float()
-    expected_f = torch.zeros_like(x, dtype=torch.float32)
     if output_layout == ep.DispatchLayout.TOKEN_MAJOR:
-        for destination_rank in range(num_ranks):
-            rank_partial = torch.zeros_like(x, dtype=torch.float32)
-            for j in range(num_topk):
-                selected = (topk_idx[:, j] >= 0) & (topk_idx[:, j] // num_local_experts == destination_rank)
-                weight_j = (
-                    selected.float() if topk_weights is None else topk_weights[:, j].masked_fill(~selected, 0.0)
-                ).view(-1, 1)
-                rank_partial = torch.addcmul(rank_partial, x_f, weight_j)
-            expected_f += rank_partial.to(torch.bfloat16).float()
+        expected = expected_token_major_output(reference_x, topk_idx, topk_weights, num_ranks, num_local_experts)
     else:
-        for j in range(num_topk):
-            weight_j = (
-                (topk_idx[:, j] != -1).float()
-                if topk_weights is None
-                else topk_weights[:, j].masked_fill(topk_idx[:, j] == -1, 0.0)
-            ).view(-1, 1)
-            expected_f = torch.addcmul(expected_f, x_f, weight_j)
-    expected = expected_f.to(torch.bfloat16)
+        expected = expected_expert_major_output(reference_x, topk_idx, topk_weights)
     local_diff, _ = validate_combine_output(
         combined_x,
         expected,