diff --git a/python/mscclpp/ep/README.md b/python/mscclpp/ep/README.md index 4ea79cb2..968c5e0f 100644 --- a/python/mscclpp/ep/README.md +++ b/python/mscclpp/ep/README.md @@ -167,7 +167,7 @@ The selected mode determines the default dispatch output layout: | Mode | Default layout | |---|---| -| `ht` | `DispatchLayout.FLAT` | +| `ht` | `DispatchLayout.TOKEN_MAJOR` | | `ll` | `DispatchLayout.EXPERT_MAJOR` | `output_layout` may still be kept as an advanced override if a backend supports @@ -177,7 +177,7 @@ Use `DispatchLayout` instead of string literals for this field: | Layout enum | Tensor shape | |---|---| -| `DispatchLayout.FLAT` | HT: `[total_recv_tokens, hidden]`; LL: `[num_local_experts * max_slots_per_expert, hidden]` | +| `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]` | ## MoECommunicator methods @@ -238,11 +238,7 @@ dispatch_out, handle = moe_comm.dispatch( output_buffer=output_buffer, ) -expert_output = mlp( - dispatch_out.tokens, - dispatch_out.layout, - dispatch_out.quant, -) +expert_output = mlp(dispatch_out) output = moe_comm.combine(expert_output, handle) ``` @@ -250,7 +246,7 @@ output = moe_comm.combine(expert_output, handle) `dispatch_out` is for the local MLP. `handle` is for `combine`. The MLP should not need to inspect the opaque handle. -`DispatchOutput.layout` carries both the layout kind (`FLAT` or `EXPERT_MAJOR`) +`DispatchOutput.layout` carries both the layout kind (`TOKEN_MAJOR` or `EXPERT_MAJOR`) and layout-specific metadata. Expert-grouped layouts populate `num_tokens_per_expert`; future layouts that do not expose per-expert grouping @@ -267,8 +263,8 @@ class QuantConfig: class DispatchLayout(str, Enum): - FLAT = "flat" EXPERT_MAJOR = "expert_major" + TOKEN_MAJOR = "token_major" @dataclass @@ -276,6 +272,7 @@ class DispatchLayoutInfo: kind: DispatchLayout num_tokens_per_expert: Optional[torch.Tensor | list[int]] = None offsets: Optional[torch.Tensor] = None + num_tokens_per_rank: Optional[torch.Tensor | list[int]] = None @dataclass @@ -289,6 +286,8 @@ class DispatchOutput: tokens: torch.Tensor quant: Optional[QuantConfig] layout: DispatchLayoutInfo + topk_ids: Optional[torch.Tensor] = None + weights: Optional[torch.Tensor] = None @dataclass @@ -304,11 +303,22 @@ class ExpertMajorCombineContext: @dataclass -class RowMajorCombineContext: +class TokenMajorCombineContext: + topk_ids: torch.Tensor + num_experts: int + num_tokens: int + hidden_size: int + source_token_ids: torch.Tensor + num_tokens_per_rank: torch.Tensor + num_max_dispatch_tokens_per_rank: int + + +@dataclass +class HighThroughputCombineContext: ... -CombineContext = ExpertMajorCombineContext | RowMajorCombineContext +CombineContext = ExpertMajorCombineContext | TokenMajorCombineContext | HighThroughputCombineContext class DispatchHandle: @@ -321,8 +331,12 @@ class ExpertMajorDispatchHandle(DispatchHandle): combine_context: ExpertMajorCombineContext -class RowMajorDispatchHandle(DispatchHandle): - combine_context: RowMajorCombineContext +class TokenMajorDispatchHandle(DispatchHandle): + combine_context: TokenMajorCombineContext + + +class HighThroughputDispatchHandle(DispatchHandle): + combine_context: HighThroughputCombineContext @dataclass @@ -409,7 +423,9 @@ overlap is operation-level only. Each concrete `DispatchHandle` stores a layout-specific `combine_context` used to reverse dispatch and finish combine. `ExpertMajorDispatchHandle` uses `ExpertMajorCombineContext` (`topk_ids`, `weights`, source info, layout ranges, -shape, and capacity). Row-major handles use the intranode combine context with +shape, and capacity). `TokenMajorDispatchHandle` records source-token IDs, +per-source-rank counts, and the original routing needed for cross-rank combine. +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`. @@ -492,6 +508,16 @@ For padded expert-major LL layout: output_buffer: [num_local_experts, world_size * max_tokens_per_rank, hidden] ``` +For token-major LL layout: + +```text +output_buffer: [world_size * max_tokens_per_rank, hidden] +``` + +The token-major rows are grouped into fixed source-rank regions. For source rank +`r`, only the first `dispatch_out.layout.num_tokens_per_rank[r]` rows in region +`[r * max_tokens_per_rank : (r + 1) * max_tokens_per_rank]` are valid. + The dtype must match the dispatch output dtype. For BF16 dispatch it is BF16. For FP8 dispatch it is FP8 and the returned `DispatchOutput.quant` carries the matching format and scale tensor. @@ -505,38 +531,18 @@ buffer instead of allocating it internally. `dispatch` should return MLP-ready tokens. The MLP should not run another token-major to expert-major permutation unless it uses a custom adapter. -### Normal / high-throughput flat layout +### Normal / high-throughput token-major layout -HT uses `DispatchLayout.FLAT`, a flat expert-major layout: +HT uses `DispatchLayout.TOKEN_MAJOR`: ```python dispatch_out.tokens # [total_recv_tokens, H] ``` -Rows are grouped by local expert id: - -```text -expert0 tokens -expert1 tokens -expert2 tokens -... -``` - -`dispatch_out.layout.num_tokens_per_expert` is ordered by local expert id: - -```python -num_tokens_per_expert[i] = valid token count for local expert i -``` - -For flat layout, `dispatch_out.layout.offsets` may be provided or derived by cumulative sum: - -```python -offsets = cumsum([0] + num_tokens_per_expert) -tokens[offsets[i] : offsets[i + 1]] -``` - -This layout is efficient for Triton or grouped GEMM kernels because it avoids -padding. +Each row represents one `(source token, destination rank)` and is accompanied by +`dispatch_out.topk_ids`, `dispatch_out.weights`, and source-token metadata. A +token routed to multiple experts on the same destination rank is transferred +only once. ### Low-latency output layouts @@ -546,14 +552,18 @@ LL defaults to `DispatchLayout.EXPERT_MAJOR`, a padded expert-major tensor: dispatch_out.tokens # [num_local_experts, max_slots_per_expert, H] ``` -LL can also return `DispatchLayout.FLAT`, which is the same contiguous -local-expert-major storage viewed as 2D: +LL can also return `DispatchLayout.TOKEN_MAJOR`: ```python -dispatch_out.tokens # [num_local_experts * max_slots_per_expert, H] +dispatch_out.tokens # [world_size * max_tokens_per_rank, H] +dispatch_out.topk_ids # [world_size * max_tokens_per_rank, K], int32 local expert IDs +dispatch_out.weights # [world_size * max_tokens_per_rank, K], float32 ``` -For expert `i`, only the first `dispatch_out.layout.num_tokens_per_expert[i]` slots are valid: +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 expert_major_tokens = dispatch_out.tokens.view(num_local_experts, max_slots_per_expert, H) @@ -572,8 +582,8 @@ dimension replaced by the scale dimension. Examples: ```text -flat tokens: HT [total_recv_tokens, H]; LL [num_local_experts * max_slots, H] -flat FP8 scales: HT [total_recv_tokens, H / 128]; LL [num_local_experts, max_slots, H / 128] +token-major tokens: HT [total_recv_tokens, H]; LL [world_size * max_tokens_per_rank, H] +token-major scales: LL [world_size * max_tokens_per_rank, H / 128] expert-major tokens: [num_local_experts, max_slots, H] expert-major scales: [num_local_experts, max_slots, H / 128] @@ -583,12 +593,15 @@ expert-major scales: [num_local_experts, max_slots, H / 128] The MLP consumes `dispatch_out`, not the original token-major input. -For flat expert-major output: +For token-major output, the local MLP consumes each token once, runs the local +experts selected by `topk_ids`, applies `weights`, and returns one pre-reduced +rank partial in the same row: ```python -expert_output = triton_mlp( +rank_partial = token_major_mlp( dispatch_out.tokens, - dispatch_out.layout, + dispatch_out.topk_ids, + dispatch_out.weights, dispatch_out.quant, ) ``` @@ -603,9 +616,10 @@ expert_output = expert_major_mlp( ) ``` -The MLP must preserve the dispatch output layout and row/slot order. It may -apply expert-specific GEMMs, but it must not compact or reorder tokens unless it -also produces compatible metadata for combine. +The MLP must preserve the dispatch output layout and row/slot order. For +token-major output, combine assumes each row is already weighted and reduced +across all local experts. `CombineMode.DIRECT_SEND` is therefore available only +for expert-major output. ## Combine API diff --git a/python/mscclpp/ep/__init__.py b/python/mscclpp/ep/__init__.py index 2719de2f..948b2af5 100644 --- a/python/mscclpp/ep/__init__.py +++ b/python/mscclpp/ep/__init__.py @@ -22,13 +22,15 @@ from .communicator import ( # noqa: F401 DispatchOutputInfo, ExpertMajorDispatchHandle, ExpertMajorCombineContext, + HighThroughputDispatchHandle, + HighThroughputCombineContext, MoECommunicator, MoECommunicatorConfig, MoEMode, OperationOverlapConfig, QuantConfig, - RowMajorDispatchHandle, - RowMajorCombineContext, + TokenMajorDispatchHandle, + TokenMajorCombineContext, ) __all__ = [ @@ -44,11 +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 fca7f2c4..08d7e82a 100644 --- a/python/mscclpp/ep/communicator.py +++ b/python/mscclpp/ep/communicator.py @@ -21,11 +21,13 @@ from .types import ( DispatchOutputInfo, ExpertMajorDispatchHandle, ExpertMajorCombineContext, + HighThroughputDispatchHandle, + HighThroughputCombineContext, MoECommunicatorConfig, OperationOverlapConfig, QuantConfig, - RowMajorDispatchHandle, - RowMajorCombineContext, + TokenMajorDispatchHandle, + TokenMajorCombineContext, ) __all__ = [ @@ -41,21 +43,23 @@ __all__ = [ "DispatchOutputInfo", "ExpertMajorDispatchHandle", "ExpertMajorCombineContext", + "HighThroughputDispatchHandle", + "HighThroughputCombineContext", "MoECommunicator", "MoECommunicatorConfig", "MoEMode", "OperationOverlapConfig", "QuantConfig", - "RowMajorDispatchHandle", - "RowMajorCombineContext", + "TokenMajorDispatchHandle", + "TokenMajorCombineContext", ] class MoECommunicator: """High-level MoE communicator for dispatch/combine. - ``mode=MoEMode.LOW_LATENCY`` selects the LL backend (EXPERT_MAJOR); - ``mode=MoEMode.HIGH_THROUGHPUT`` selects the HT backend (FLAT). + ``mode=MoEMode.LOW_LATENCY`` selects the LL backend (EXPERT_MAJOR by default); + ``mode=MoEMode.HIGH_THROUGHPUT`` selects the HT backend (TOKEN_MAJOR). """ def __init__(self, config: Optional[MoECommunicatorConfig] = None, **kwargs) -> None: @@ -162,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 9f025d86..c48dce96 100644 --- a/python/mscclpp/ep/high_throughput.py +++ b/python/mscclpp/ep/high_throughput.py @@ -24,15 +24,14 @@ from .types import ( DispatchLayoutInfo, DispatchOutput, DispatchOutputInfo, + HighThroughputCombineContext, + HighThroughputDispatchHandle, MoECommunicatorConfig, QuantConfig, - RowMajorCombineContext, - RowMajorDispatchHandle, ) from .utils import ( bf16_view as _bf16_view, current_stream_ptr as _stream_ptr, - exclusive_cumsum, ptr as _ptr, resolve_expert_placement, ) @@ -289,8 +288,8 @@ class HighThroughputBackend: self.num_sms = config.num_sms self.enable_overlap = config.enable_overlap - if self.output_layout != DispatchLayout.FLAT: - raise NotImplementedError("HT mode currently supports only DispatchLayout.FLAT") + if self.output_layout != DispatchLayout.TOKEN_MAJOR: + raise NotImplementedError("HT mode currently supports only DispatchLayout.TOKEN_MAJOR") self.num_local_experts, self.local_expert_start = resolve_expert_placement( num_experts=self.num_experts, @@ -351,10 +350,13 @@ class HighThroughputBackend: previous_handle: Optional[DispatchHandle], ) -> tuple[DispatchOutput, DispatchHandle]: self._validate_dispatch_inputs(input, topk_ids, weights, quant) + implicit_weights = weights is None if weights is None: weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device) cache = getattr(previous_handle, "_dispatch_cache", None) if previous_handle is not None else None + if cache is not None and not self._cache_matches(cache, input, topk_ids, weights, implicit_weights): + cache = None if cache is not None: num_tokens_per_rank = cache["num_tokens_per_rank"] num_tokens_per_expert = cache["num_tokens_per_expert"] @@ -370,9 +372,9 @@ class HighThroughputBackend: ( recv_x, _recv_x_scales, - _recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, + _runtime_recv_topk_idx, + _runtime_recv_topk_weights, + _runtime_num_recv_tokens_per_expert_list, rank_prefix_matrix, _channel_prefix_matrix, recv_channel_prefix_matrix, @@ -391,7 +393,11 @@ class HighThroughputBackend: cache["channel_prefix_matrix"], self.expert_alignment, ) - combine_context = RowMajorCombineContext( + del _runtime_recv_topk_idx, _runtime_recv_topk_weights, _runtime_num_recv_tokens_per_expert_list + 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 = HighThroughputCombineContext( recv_topk_weights=recv_topk_weights, src_idx=recv_src_idx, rank_prefix_matrix=rank_prefix_matrix, @@ -403,7 +409,7 @@ class HighThroughputBackend: ( recv_x, _recv_x_scales, - _recv_topk_idx, + recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, rank_prefix_matrix, @@ -424,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, @@ -438,13 +444,23 @@ class HighThroughputBackend: "rank_prefix_matrix": rank_prefix_matrix, "channel_prefix_matrix": channel_prefix_matrix, "num_recv_tokens": int(recv_x.size(0)), + "recv_topk_idx": recv_topk_idx, + "recv_topk_weights": recv_topk_weights, + "num_recv_tokens_per_expert_list": num_recv_tokens_per_expert_list, + "backend_id": id(self), + "num_tokens": int(input.size(0)), + "device": input.device, + "topk_ids_ptr": topk_ids.data_ptr(), + "topk_ids_version": topk_ids._version, + "implicit_weights": implicit_weights, + "weights_ptr": 0 if implicit_weights else weights.data_ptr(), + "weights_version": 0 if implicit_weights else weights._version, } output_info = DispatchOutputInfo( layout=DispatchLayoutInfo( - kind=DispatchLayout.FLAT, + kind=self.output_layout, num_tokens_per_expert=num_recv_tokens_per_expert_list, - offsets=exclusive_cumsum(num_recv_tokens_per_expert_list), ), quant=None, ) @@ -452,14 +468,28 @@ class HighThroughputBackend: tokens=recv_x, quant=output_info.quant, layout=output_info.layout, + 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] handle._dispatch_cache = dispatch_cache # type: ignore[attr-defined] return dispatch_out, handle + def _cache_matches(self, cache, input, topk_ids, weights, implicit_weights) -> bool: + return ( + cache.get("backend_id") == id(self) + and cache.get("num_tokens") == int(input.size(0)) + and cache.get("device") == input.device + and cache.get("topk_ids_ptr") == topk_ids.data_ptr() + and cache.get("topk_ids_version") == topk_ids._version + and cache.get("implicit_weights") == implicit_weights + and (implicit_weights or cache.get("weights_ptr") == weights.data_ptr()) + and (implicit_weights or cache.get("weights_version") == weights._version) + ) + def combine( self, expert_output: torch.Tensor, @@ -517,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/low_latency.py b/python/mscclpp/ep/low_latency.py index 19ed3674..56a35f55 100644 --- a/python/mscclpp/ep/low_latency.py +++ b/python/mscclpp/ep/low_latency.py @@ -18,6 +18,8 @@ from .types import ( ExpertMajorCombineContext, MoECommunicatorConfig, QuantConfig, + TokenMajorCombineContext, + TokenMajorDispatchHandle, ) from .utils import cuda_stream_ptr, resolve_expert_placement @@ -90,14 +92,16 @@ class LowLatencyBackend: self.combine_mode = config.low_latency_combine_mode self.enable_overlap = config.enable_overlap - if self.output_layout != DispatchLayout.EXPERT_MAJOR: - raise NotImplementedError("low-latency mode currently supports only DispatchLayout.EXPERT_MAJOR") + if self.output_layout not in (DispatchLayout.EXPERT_MAJOR, DispatchLayout.TOKEN_MAJOR): + raise NotImplementedError("low-latency mode supports EXPERT_MAJOR and TOKEN_MAJOR output") if self.num_experts % self.world_size != 0: raise ValueError("low-latency mode requires num_experts divisible by world_size") if not self.world_size + 2 <= self.num_blocks <= 130: raise ValueError("low_latency_num_blocks must be between world_size + 2 and 130") if not isinstance(self.combine_mode, CombineMode): raise TypeError("low_latency_combine_mode must be a CombineMode") + if self.output_layout == DispatchLayout.TOKEN_MAJOR and self.combine_mode != CombineMode.RANK_LOCAL_REDUCE: + raise ValueError("TOKEN_MAJOR output requires RANK_LOCAL_REDUCE combine") self.num_local_experts, self.local_expert_start = resolve_expert_placement( num_experts=self.num_experts, @@ -113,6 +117,8 @@ class LowLatencyBackend: self._dispatch_scales: Optional[torch.Tensor] = None self._dispatch_src_info: Optional[torch.Tensor] = None + self._dispatch_topk_ids: Optional[torch.Tensor] = None + self._dispatch_weights: Optional[torch.Tensor] = None self._dispatch_layout_range: Optional[torch.Tensor] = None self._dispatch_count: Optional[torch.Tensor] = None @@ -148,7 +154,9 @@ class LowLatencyBackend: del previous_handle self._validate_dispatch_inputs(input, topk_ids, weights, quant, output_buffer) - out_buf, scales, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer) + out_buf, scales, src_info, recv_topk_ids, recv_weights, layout_range, count = self._get_dispatch_output_tensors( + output_buffer + ) self._runtime.cpp_runtime.dispatch( input.data_ptr(), topk_ids.data_ptr(), @@ -156,13 +164,16 @@ class LowLatencyBackend: out_buf.data_ptr(), 0 if scales is None else scales.data_ptr(), src_info.data_ptr(), - layout_range.data_ptr(), + 0 if recv_topk_ids is None else recv_topk_ids.data_ptr(), + 0 if recv_weights is None else recv_weights.data_ptr(), + 0 if layout_range is None else layout_range.data_ptr(), count.data_ptr(), input.size(0), self.hidden_size, self.topk, self.max_tokens_per_rank, self.num_experts, + self.output_layout, self.dispatch_data_type, self.num_blocks, cuda_stream_ptr(stream), @@ -175,28 +186,46 @@ class LowLatencyBackend: block_scales=scales, ) ) - output_info = DispatchOutputInfo( - layout=DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count), - quant=output_quant, - ) + if self.output_layout == DispatchLayout.EXPERT_MAJOR: + layout_info = DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count) + else: + layout_info = DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_rank=count) + output_info = DispatchOutputInfo(layout=layout_info, quant=output_quant) dispatch_out = DispatchOutput( tokens=out_buf, quant=output_info.quant, layout=output_info.layout, + topk_ids=recv_topk_ids, + weights=recv_weights, ) - handle = ExpertMajorDispatchHandle( - output_info=output_info, - combine_context=ExpertMajorCombineContext( - topk_ids=topk_ids, - weights=weights, - num_experts=self.num_experts, - num_tokens=input.size(0), - hidden_size=self.hidden_size, - src_info=src_info, - layout_range=layout_range, - num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank, - ), - ) + if self.output_layout == DispatchLayout.EXPERT_MAJOR: + assert layout_range is not None + handle: DispatchHandle = ExpertMajorDispatchHandle( + output_info=output_info, + combine_context=ExpertMajorCombineContext( + topk_ids=topk_ids, + weights=weights, + num_experts=self.num_experts, + num_tokens=input.size(0), + hidden_size=self.hidden_size, + src_info=src_info, + layout_range=layout_range, + num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank, + ), + ) + else: + handle = TokenMajorDispatchHandle( + output_info=output_info, + combine_context=TokenMajorCombineContext( + topk_ids=topk_ids, + num_experts=self.num_experts, + num_tokens=input.size(0), + hidden_size=self.hidden_size, + source_token_ids=src_info, + num_tokens_per_rank=count, + num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank, + ), + ) return dispatch_out, handle def combine( @@ -208,23 +237,33 @@ class LowLatencyBackend: stream: Optional[torch.cuda.Stream], ) -> torch.Tensor: self._validate_combine_inputs(expert_output, handle, out) - if not isinstance(handle, ExpertMajorDispatchHandle): - raise ValueError("DispatchHandle does not contain expert-major combine context") - context = handle.combine_context + if isinstance(handle, ExpertMajorDispatchHandle): + context = handle.combine_context + topk_weights = context.weights + src_info = context.src_info + layout_range = context.layout_range + elif isinstance(handle, TokenMajorDispatchHandle): + context = handle.combine_context + topk_weights = None + src_info = context.source_token_ids + layout_range = None + else: + raise ValueError("DispatchHandle does not contain low-latency combine context") if out is None: out = torch.empty((context.num_tokens, self.hidden_size), dtype=torch.bfloat16, device=expert_output.device) self._runtime.cpp_runtime.combine( expert_output.data_ptr(), context.topk_ids.data_ptr(), - 0 if context.weights is None else context.weights.data_ptr(), - context.src_info.data_ptr(), - context.layout_range.data_ptr(), + 0 if topk_weights is None else topk_weights.data_ptr(), + src_info.data_ptr(), + 0 if layout_range is None else layout_range.data_ptr(), out.data_ptr(), context.num_tokens, self.hidden_size, self.topk, context.num_max_dispatch_tokens_per_rank, context.num_experts, + self.output_layout, self.dispatch_data_type, self.combine_mode, self.num_blocks - 2, @@ -236,24 +275,42 @@ class LowLatencyBackend: device = output_buffer.device slots_per_expert = self.world_size * self.max_tokens_per_rank if self._dispatch_src_info is None or self._dispatch_src_info.device != device: - self._dispatch_src_info = torch.empty( - (self.num_local_experts, slots_per_expert), dtype=torch.int32, device=device - ) - self._dispatch_layout_range = torch.empty( - (self.num_local_experts, self.world_size), dtype=torch.int64, device=device - ) - self._dispatch_count = torch.empty((self.num_local_experts,), dtype=torch.int32, device=device) self._dispatch_scales = None - if self.dispatch_data_type == DispatchDataType.FP8_E4M3: - num_scales = self.hidden_size // 128 - scale_storage = torch.empty( - (self.num_local_experts, num_scales, slots_per_expert), dtype=torch.float32, device=device + self._dispatch_topk_ids = None + self._dispatch_weights = None + if self.output_layout == DispatchLayout.EXPERT_MAJOR: + self._dispatch_src_info = torch.empty( + (self.num_local_experts, slots_per_expert), dtype=torch.int32, device=device ) - self._dispatch_scales = scale_storage.transpose(1, 2) + self._dispatch_layout_range = torch.empty( + (self.num_local_experts, self.world_size), dtype=torch.int64, device=device + ) + self._dispatch_count = torch.empty((self.num_local_experts,), dtype=torch.int32, device=device) + if self.dispatch_data_type == DispatchDataType.FP8_E4M3: + num_scales = self.hidden_size // 128 + scale_storage = torch.empty( + (self.num_local_experts, num_scales, slots_per_expert), dtype=torch.float32, device=device + ) + self._dispatch_scales = scale_storage.transpose(1, 2) + else: + token_capacity = self.world_size * self.max_tokens_per_rank + self._dispatch_src_info = torch.empty((token_capacity,), dtype=torch.int32, device=device) + self._dispatch_topk_ids = torch.empty((token_capacity, self.topk), dtype=torch.int32, device=device) + self._dispatch_weights = torch.empty((token_capacity, self.topk), dtype=torch.float32, device=device) + self._dispatch_layout_range = None + self._dispatch_count = torch.empty((self.world_size,), dtype=torch.int32, device=device) + if self.dispatch_data_type == DispatchDataType.FP8_E4M3: + self._dispatch_scales = torch.empty( + (token_capacity, self.hidden_size // 128), dtype=torch.float32, device=device + ) + assert self._dispatch_src_info is not None + assert self._dispatch_count is not None return ( output_buffer, self._dispatch_scales, self._dispatch_src_info, + self._dispatch_topk_ids, + self._dispatch_weights, self._dispatch_layout_range, self._dispatch_count, ) @@ -290,7 +347,7 @@ class LowLatencyBackend: if self.output_layout == DispatchLayout.EXPERT_MAJOR: expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size) else: - expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size) + expected_shape = (self.world_size * self.max_tokens_per_rank, self.hidden_size) if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous(): raise ValueError(f"output_buffer must be a contiguous {self.output_layout} tensor") expected_dtype = torch.float8_e4m3fn if self.dispatch_data_type == DispatchDataType.FP8_E4M3 else torch.bfloat16 @@ -300,11 +357,13 @@ class LowLatencyBackend: raise ValueError(f"output_buffer shape must be {expected_shape}") def _validate_combine_inputs(self, expert_output, handle, out) -> None: - if not isinstance(handle, ExpertMajorDispatchHandle): - raise ValueError("DispatchHandle does not contain expert-major combine context") + if not isinstance(handle, (ExpertMajorDispatchHandle, TokenMajorDispatchHandle)): + raise ValueError("DispatchHandle does not contain low-latency combine context") context = handle.combine_context if context.num_experts != self.num_experts or context.hidden_size != self.hidden_size: raise ValueError("DispatchHandle does not belong to this MoECommunicator configuration") + if handle.output_info.layout.kind != self.output_layout: + raise ValueError("DispatchHandle output layout does not match this MoECommunicator") output_quant = handle.output_info.quant handle_data_type = DispatchDataType.BF16 if output_quant is None else output_quant.format if handle_data_type != self.dispatch_data_type: @@ -313,7 +372,7 @@ class LowLatencyBackend: if handle.output_info.layout.kind == DispatchLayout.EXPERT_MAJOR: expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size) else: - expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size) + expected_shape = (self.world_size * self.max_tokens_per_rank, self.hidden_size) if expert_output.dim() != len(expected_shape) or not expert_output.is_contiguous(): raise ValueError("expert_output must keep dispatch output's contiguous layout") if tuple(expert_output.shape) != expected_shape: diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py index 35fde230..6c53b652 100644 --- a/python/mscclpp/ep/types.py +++ b/python/mscclpp/ep/types.py @@ -78,6 +78,7 @@ class DispatchLayoutInfo: kind: DispatchLayout num_tokens_per_expert: Optional[Union[torch.Tensor, List[int]]] = None offsets: Optional[torch.Tensor] = None + num_tokens_per_rank: Optional[Union[torch.Tensor, List[int]]] = None @dataclass @@ -95,6 +96,8 @@ class DispatchOutput: tokens: torch.Tensor quant: Optional[QuantConfig] layout: DispatchLayoutInfo + topk_ids: Optional[torch.Tensor] = None + weights: Optional[torch.Tensor] = None # Combine-side context. These objects are layout-specific and opaque to the MLP. @@ -115,8 +118,21 @@ class ExpertMajorCombineContext: @dataclass -class RowMajorCombineContext: - """Combine context for row-major high-throughput dispatch output.""" +class TokenMajorCombineContext: + """Combine context for token-major, rank-local pre-reduced output.""" + + topk_ids: torch.Tensor + num_experts: int + num_tokens: int + hidden_size: int + source_token_ids: torch.Tensor + num_tokens_per_rank: torch.Tensor + num_max_dispatch_tokens_per_rank: int + + +@dataclass +class HighThroughputCombineContext: + """Combine context for high-throughput dispatch output.""" recv_topk_weights: Optional[torch.Tensor] src_idx: torch.Tensor @@ -125,7 +141,7 @@ class RowMajorCombineContext: send_head: torch.Tensor -CombineContext = Union[ExpertMajorCombineContext, RowMajorCombineContext] +CombineContext = Union[ExpertMajorCombineContext, TokenMajorCombineContext, HighThroughputCombineContext] # Opaque dispatch handles returned by dispatch() and consumed by combine(). @@ -144,8 +160,13 @@ class ExpertMajorDispatchHandle(DispatchHandle): @dataclass -class RowMajorDispatchHandle(DispatchHandle): - combine_context: RowMajorCombineContext +class TokenMajorDispatchHandle(DispatchHandle): + combine_context: TokenMajorCombineContext + + +@dataclass +class HighThroughputDispatchHandle(DispatchHandle): + combine_context: HighThroughputCombineContext # Optional async/overlap configuration. diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index d76c1dd4..3f8d05b2 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -30,6 +30,13 @@ The optimized LL backend is available when all participating ranks belong to one detected GPU IPC domain. That domain may span hosts when CUDA fabric handles and the required fabric services are available. +LL dispatch supports two user-visible layouts: + +- `EXPERT_MAJOR`: one row per `(token, local expert)`. +- `TOKEN_MAJOR`: one row per `(token, destination rank)`, plus local top-k expert + IDs, routing weights, source-token IDs, and per-source-rank counts. The caller + must produce one pre-weighted local partial per row before combine. + ### High throughput HT follows the same direct-mapping resource model: diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index 24403f86..d28411e1 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -46,7 +46,7 @@ NB_MODULE(mscclpp_ep_cpp, m) { nb::enum_(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_(m, "CombineMode") .value("RANK_LOCAL_REDUCE", mscclpp::ep::low_latency::CombineMode::RANK_LOCAL_REDUCE) @@ -64,37 +64,41 @@ NB_MODULE(mscclpp_ep_cpp, m) { .def( "dispatch", [](mscclpp::ep::MoERuntime& self, uintptr_t inputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, - uintptr_t outputPtr, uintptr_t outputScalesPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputLayoutRangePtr, - uintptr_t outputCountPtr, int numTokens, int hidden, int numTopk, int maxTokensPerRank, int numExperts, + uintptr_t outputPtr, uintptr_t outputScalesPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputTopkIdxPtr, + uintptr_t outputTopkWeightsPtr, uintptr_t outputLayoutRangePtr, uintptr_t outputCountPtr, int numTokens, + int hidden, int numTopk, int maxTokensPerRank, int numExperts, mscclpp::ep::DispatchLayout dispatchLayout, mscclpp::ep::low_latency::DispatchDataType dispatchDataType, int numBlocks, uintptr_t streamPtr) { - self.dispatch( - ptr(outputPtr), reinterpret_cast(ptr(outputScalesPtr)), - reinterpret_cast(ptr(outputSrcInfoPtr)), reinterpret_cast(ptr(outputLayoutRangePtr)), - reinterpret_cast(ptr(outputCountPtr)), ptr(inputPtr), reinterpret_cast(ptr(topkIdxPtr)), - reinterpret_cast(ptr(topkWeightsPtr)), numTokens, hidden, numTopk, maxTokensPerRank, numExperts, - dispatchDataType, numBlocks, stream(streamPtr)); + self.dispatch(ptr(outputPtr), reinterpret_cast(ptr(outputScalesPtr)), + reinterpret_cast(ptr(outputSrcInfoPtr)), reinterpret_cast(ptr(outputTopkIdxPtr)), + reinterpret_cast(ptr(outputTopkWeightsPtr)), + reinterpret_cast(ptr(outputLayoutRangePtr)), + reinterpret_cast(ptr(outputCountPtr)), ptr(inputPtr), + reinterpret_cast(ptr(topkIdxPtr)), reinterpret_cast(ptr(topkWeightsPtr)), + numTokens, hidden, numTopk, maxTokensPerRank, numExperts, dispatchLayout, dispatchDataType, + numBlocks, stream(streamPtr)); }, nb::arg("input_ptr"), nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("output_ptr"), - nb::arg("output_scales_ptr"), nb::arg("output_src_info_ptr"), nb::arg("output_layout_range_ptr"), - nb::arg("output_count_ptr"), nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), - nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("dispatch_data_type"), nb::arg("num_blocks"), + nb::arg("output_scales_ptr"), nb::arg("output_src_info_ptr"), nb::arg("output_topk_idx_ptr"), + nb::arg("output_topk_weights_ptr"), nb::arg("output_layout_range_ptr"), nb::arg("output_count_ptr"), + nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), + nb::arg("num_experts"), nb::arg("dispatch_layout"), nb::arg("dispatch_data_type"), nb::arg("num_blocks"), nb::arg("stream_ptr")) .def( "combine", [](mscclpp::ep::MoERuntime& self, uintptr_t expertOutputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, uintptr_t srcInfoPtr, uintptr_t layoutRangePtr, uintptr_t outputPtr, int numTokens, int hidden, - int numTopk, int maxTokensPerRank, int numExperts, + int numTopk, int maxTokensPerRank, int numExperts, mscclpp::ep::DispatchLayout dispatchLayout, mscclpp::ep::low_latency::DispatchDataType dispatchDataType, mscclpp::ep::low_latency::CombineMode mode, int numBlocks, uintptr_t streamPtr) { self.combine(ptr(outputPtr), ptr(expertOutputPtr), reinterpret_cast(ptr(topkIdxPtr)), reinterpret_cast(ptr(topkWeightsPtr)), reinterpret_cast(ptr(srcInfoPtr)), reinterpret_cast(ptr(layoutRangePtr)), numTokens, hidden, numTopk, maxTokensPerRank, - numExperts, dispatchDataType, mode, numBlocks, stream(streamPtr)); + numExperts, dispatchLayout, dispatchDataType, mode, numBlocks, stream(streamPtr)); }, nb::arg("expert_output_ptr"), nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("src_info_ptr"), nb::arg("layout_range_ptr"), nb::arg("output_ptr"), nb::arg("num_tokens"), nb::arg("hidden"), - nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("dispatch_data_type"), - nb::arg("mode"), nb::arg("num_blocks"), nb::arg("stream_ptr")); + nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("dispatch_layout"), + nb::arg("dispatch_data_type"), nb::arg("mode"), nb::arg("num_blocks"), nb::arg("stream_ptr")); nb::class_(m, "Config") .def(nb::init(), nb::arg("num_sms") = 20, nb::arg("num_max_nvl_chunked_send_tokens") = 6, diff --git a/src/ext/ep/include/api.cuh b/src/ext/ep/include/api.cuh index 05390245..e6ae4e44 100644 --- a/src/ext/ep/include/api.cuh +++ b/src/ext/ep/include/api.cuh @@ -29,8 +29,10 @@ enum class MoEMode { enum class DispatchLayout { /// [num_local_experts, num_ranks * max_tokens_per_rank, hidden]. EXPERT_MAJOR, - /// [num_local_experts * num_ranks * max_tokens_per_rank, hidden]. - FLAT + /// 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 }; // =========================================================================== @@ -140,6 +142,8 @@ struct Workload { int numExperts_; /// Maximum tokens per rank in the packed layout. int maxTokensPerRank_; + /// User-visible dispatch output layout. + DispatchLayout outputLayout_; /// Dispatch payload data format. DispatchDataType dispatchDataType_; }; @@ -171,14 +175,16 @@ struct CommContext { size_t workspaceSize(int numRanks, int numExperts); /// Low-latency dispatch that distributes tokens to experts across ranks. -/// @param[out] output Expert-major packed output -/// [num_local_experts, num_ranks * max_tokens_per_rank, hidden]. -/// @param[out] outputScales FP8 block scales in -/// [num_local_experts, hidden / 128, num_ranks * max_tokens_per_rank], -/// or nullptr for BF16 dispatch. -/// @param[out] outputSrcInfo Original source-token index for every packed expert row. -/// @param[out] outputLayout Per-[local expert, source rank] packed count and offset. -/// @param[out] outputCount Total packed token count for every local expert. +/// @param[out] output Expert-major or token-major packed output selected by +/// Workload::outputLayout_. +/// @param[out] outputScales Layout-matched FP8 block scales, or nullptr for BF16 dispatch. +/// @param[out] outputSrcInfo Original source-token index for every output row. +/// @param[out] outputTopkIdx Token-major local expert indices [num_ranks * max_tokens_per_rank, num_topk], or nullptr. +/// @param[out] outputTopkWeights Token-major routing weights +/// [num_ranks * max_tokens_per_rank, num_topk], or nullptr. +/// @param[out] outputLayout Per-[local expert, source rank] packed count and offset for expert-major output, or +/// nullptr. +/// @param[out] outputCount Per-local-expert counts for expert-major output or per-source-rank counts for token-major. /// @param[in] input Local input tokens [num_tokens, hidden]. /// @param[in] topkIdx Global expert indices [num_tokens, num_topk]. /// @param[in] topkWeights Routing weights [num_tokens, num_topk], or nullptr for unit weights. @@ -188,13 +194,15 @@ size_t workspaceSize(int numRanks, int numExperts); /// @param[in,out] workspace Persistent counters, task storage, semaphores, and device barriers. /// @param[in] numBlocks Total dispatch grid size, including one scheduler and one metadata-notify block. /// @param[in] stream CUDA stream. -void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, - const void* input, const int64_t* topkIdx, const float* topkWeights, const Workload& workload, - void* recvBuffer, const CommContext& comm, void* workspace, int numBlocks, cudaStream_t stream); +void dispatch(void* output, float* outputScales, int* outputSrcInfo, int* outputTopkIdx, float* outputTopkWeights, + int64_t* outputLayout, int* outputCount, const void* input, const int64_t* topkIdx, + const float* topkWeights, const Workload& workload, void* recvBuffer, const CommContext& comm, + void* workspace, int numBlocks, cudaStream_t stream); /// Low-latency combine that aggregates expert outputs back to tokens. /// @param[out] output Combined local tokens [num_tokens, hidden]. -/// @param[in] input Expert outputs [num_local_experts, num_ranks * max_tokens_per_rank, hidden]. +/// @param[in] input Expert-major expert outputs or token-major pre-weighted +/// rank-local partials, matching Workload::outputLayout_. /// @param[in] topkIdx Global expert indices [num_tokens, num_topk]. /// @param[in] topkWeights Routing weights [num_tokens, num_topk], or nullptr for unit weights. /// @param[in] srcInfo Original source-token index for every packed expert row. diff --git a/src/ext/ep/low_latency/combine.cu b/src/ext/ep/low_latency/combine.cu index 450f545c..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 +template MSCCLPP_HOST_DEVICE_INLINE size_t combineSharedBytes(int nLocalExperts) { constexpr size_t TileBytes = static_cast(Hidden) * sizeof(Bf16); if constexpr (Mode == low_latency::CombineMode::DIRECT_SEND) { return directSendControlBytes(nLocalExperts) + static_cast(directSendWorkerCount(nLocalExperts)) * directSendWorkerBytes(); } + if constexpr (Layout == DispatchLayout::TOKEN_MAJOR) { + return CombineNStages * TileBytes; + } return CombineNStages * TileBytes; } @@ -187,6 +190,62 @@ MSCCLPP_DEVICE_INLINE void sendRankReducedPartials(const void* expertOutput, int if (tokenIteration > 0 && threadId == 0) waitBulkGroup(); } +template