microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-17 17:27:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add token major layout support (#841)

Browse Source
This commit is contained in:
Binyang Li
2026-07-15 23:06:02 -07:00
committed by GitHub
parent 5ebbd340bb
commit 39813bfff1
17 changed files with 1090 additions and 423 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
python/mscclpp/ep/README.md
View File

@@ -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

12
python/mscclpp/ep/__init__.py
View File

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

18
python/mscclpp/ep/communicator.py
View File

@@ -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

60
python/mscclpp/ep/high_throughput.py
View File

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

145
python/mscclpp/ep/low_latency.py
View File

@@ -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:

31
python/mscclpp/ep/types.py
View File

@@ -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.