address comments

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:

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",
 ]

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

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")

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.