microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-18 01:37:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adjust token major layout (#844)

Browse Source 
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Binyang Li
2026-07-16 18:56:10 -07:00
committed by GitHub
parent 39813bfff1
commit dd9fa728ad
15 changed files with 190 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -60,11 +60,11 @@ class MoECommunicatorConfig:
hidden_size: int = 0
topk: int = 0
max_tokens_per_rank: int = 0
max_recv_tokens_per_rank: Optional[int] = None
# Runtime mode and output layout
mode: MoEMode = MoEMode.LOW_LATENCY
output_layout: Optional[DispatchLayout] = None # default is derived from mode
token_major_init_padding: bool = False
# Quantization defaults
quant: Optional[QuantConfig] = None
@@ -137,8 +137,8 @@ a later version can add an explicit `expert_map` for arbitrary placement.
|---|---|
| `mode` | Backend selection (`MoEMode.LOW_LATENCY` or `MoEMode.HIGH_THROUGHPUT`) |
| `output_layout` | MLP input layout returned by dispatch |
| `token_major_init_padding` | Initialize token-major padding metadata for fixed-capacity Triton kernels |
| `max_tokens_per_rank` | dispatch capacity |
| `max_recv_tokens_per_rank` | recv buffer capacity |
| scratch buffers | internally sized from mode, capacity, topology, and shape |
| `num_sms` | backend launch/resource tuning |
| `dispatch_config`, `combine_config` | backend-specific tuning configs |
@@ -310,6 +310,7 @@ class TokenMajorCombineContext:
hidden_size: int
source_token_ids: torch.Tensor
num_tokens_per_rank: torch.Tensor
rank_offsets: torch.Tensor
num_max_dispatch_tokens_per_rank: int
@@ -514,9 +515,24 @@ For token-major LL layout:
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 caller chooses `max_tokens_per_rank`, allocates that fixed capacity, and may
pass the resulting row capacity directly to a Triton kernel. This avoids a CPU
synchronization while keeping a static CUDA Graph shape. All valid rows are
compacted into one contiguous prefix. For source rank `r`, its rows are:
```python
begin = dispatch_out.layout.offsets[r]
end = dispatch_out.layout.offsets[r + 1]
```
`offsets[-1]` is the total number of valid rows.
Rows after `offsets[-1]` are padding. Set
`token_major_init_padding=True` when a fixed-capacity Triton kernel should
process the entire allocation: padding `topk_ids` are then initialized to `-1`
and weights to `0`, so the kernel can skip a row when all expert IDs are `-1`.
The option defaults to `False` to avoid initialization overhead when the MLP
uses the compact valid length.
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
@@ -560,8 +576,10 @@ 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` and weight `0`. The valid row count in each
source-rank region is returned in `dispatch_out.layout.num_tokens_per_rank`.
Only the prefix ending at `dispatch_out.layout.offsets[-1]` contains valid
tokens. Non-local entries always use expert ID `-1` and weight `0`; padding uses
the same sentinel values when `token_major_init_padding=True`. Per-source-rank
counts are 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:

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

@@ -54,11 +54,19 @@ class LowLatencyRuntime:
hidden: int,
num_experts: int,
num_topk: int,
initialize_token_major_padding: bool,
) -> None:
self.rank: int = comm.my_rank
self.group_size: int = comm.nranks
self.comm = comm
self.cpp_runtime = MoERuntime(comm.communicator, max_tokens_per_rank, hidden, num_experts, num_topk)
self.cpp_runtime = MoERuntime(
comm.communicator,
max_tokens_per_rank,
hidden,
num_experts,
num_topk,
initialize_token_major_padding,
)
def is_available(self) -> bool:
return self.cpp_runtime.is_available()
@@ -90,6 +98,7 @@ class LowLatencyBackend:
self.num_blocks = config.low_latency_num_blocks
self.num_sms = self.num_blocks - 2
self.combine_mode = config.low_latency_combine_mode
self.token_major_init_padding = config.token_major_init_padding
self.enable_overlap = config.enable_overlap
if self.output_layout not in (DispatchLayout.EXPERT_MAJOR, DispatchLayout.TOKEN_MAJOR):
@@ -100,6 +109,10 @@ class LowLatencyBackend:
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 not isinstance(self.token_major_init_padding, bool):
raise TypeError("token_major_init_padding must be a bool")
if self.token_major_init_padding and self.output_layout != DispatchLayout.TOKEN_MAJOR:
raise ValueError("token_major_init_padding requires TOKEN_MAJOR output")
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")
@@ -111,8 +124,6 @@ class LowLatencyBackend:
local_expert_start=config.local_expert_start,
)
if config.max_recv_tokens_per_rank not in (None, self.max_tokens_per_rank):
raise NotImplementedError("low-latency mode currently uses max_tokens_per_rank as recv capacity")
self.dispatch_data_type = _resolve_dispatch_data_type(config.quant)
self._dispatch_scales: Optional[torch.Tensor] = None
@@ -128,6 +139,7 @@ class LowLatencyBackend:
hidden=self.hidden_size,
num_experts=self.num_experts,
num_topk=self.topk,
initialize_token_major_padding=self.token_major_init_padding,
)
self._is_internode = self._runtime.is_internode_available()
@@ -188,8 +200,15 @@ class LowLatencyBackend:
)
if self.output_layout == DispatchLayout.EXPERT_MAJOR:
layout_info = DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count)
elif self.output_layout == DispatchLayout.TOKEN_MAJOR:
assert layout_range is not None
layout_info = DispatchLayoutInfo(
kind=self.output_layout,
num_tokens_per_rank=count,
offsets=layout_range,
)
else:
layout_info = DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_rank=count)
raise ValueError(f"unsupported low-latency output layout: {self.output_layout}")
output_info = DispatchOutputInfo(layout=layout_info, quant=output_quant)
dispatch_out = DispatchOutput(
tokens=out_buf,
@@ -213,7 +232,8 @@ class LowLatencyBackend:
num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank,
),
)
else:
elif self.output_layout == DispatchLayout.TOKEN_MAJOR:
assert layout_range is not None
handle = TokenMajorDispatchHandle(
output_info=output_info,
combine_context=TokenMajorCombineContext(
@@ -223,9 +243,12 @@ class LowLatencyBackend:
hidden_size=self.hidden_size,
source_token_ids=src_info,
num_tokens_per_rank=count,
rank_offsets=layout_range,
num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank,
),
)
else:
raise ValueError(f"unsupported low-latency output layout: {self.output_layout}")
return dispatch_out, handle
def combine(
@@ -246,7 +269,7 @@ class LowLatencyBackend:
context = handle.combine_context
topk_weights = None
src_info = context.source_token_ids
layout_range = None
layout_range = context.rank_offsets
else:
raise ValueError("DispatchHandle does not contain low-latency combine context")
if out is None:
@@ -292,17 +315,19 @@ class LowLatencyBackend:
(self.num_local_experts, num_scales, slots_per_expert), dtype=torch.float32, device=device
)
self._dispatch_scales = scale_storage.transpose(1, 2)
else:
elif self.output_layout == DispatchLayout.TOKEN_MAJOR:
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_layout_range = torch.empty((self.world_size + 1,), dtype=torch.int64, device=device)
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
)
else:
raise ValueError(f"unsupported low-latency output layout: {self.output_layout}")
assert self._dispatch_src_info is not None
assert self._dispatch_count is not None
return (
@@ -346,8 +371,10 @@ class LowLatencyBackend:
slots_per_expert = self.world_size * self.max_tokens_per_rank
if self.output_layout == DispatchLayout.EXPERT_MAJOR:
expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size)
else:
elif self.output_layout == DispatchLayout.TOKEN_MAJOR:
expected_shape = (self.world_size * self.max_tokens_per_rank, self.hidden_size)
else:
raise ValueError(f"unsupported low-latency output layout: {self.output_layout}")
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
@@ -371,8 +398,12 @@ class LowLatencyBackend:
slots_per_expert = self.world_size * self.max_tokens_per_rank
if handle.output_info.layout.kind == DispatchLayout.EXPERT_MAJOR:
expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size)
else:
elif handle.output_info.layout.kind == DispatchLayout.TOKEN_MAJOR:
expected_shape = (self.world_size * self.max_tokens_per_rank, self.hidden_size)
if handle.output_info.layout.offsets is None:
raise ValueError("TOKEN_MAJOR DispatchHandle is missing compact rank offsets")
else:
raise ValueError(f"unsupported low-latency output layout: {handle.output_info.layout.kind}")
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:

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

@@ -47,11 +47,11 @@ class MoECommunicatorConfig:
hidden_size: int = 0
topk: int = 0
max_tokens_per_rank: int = 0
max_recv_tokens_per_rank: Optional[int] = None
# Runtime mode and output layout
mode: MoEMode = MoEMode.LOW_LATENCY
output_layout: Optional[DispatchLayout] = None
token_major_init_padding: bool = False
# Quantization defaults
quant: Optional[QuantConfig] = None
@@ -127,6 +127,7 @@ class TokenMajorCombineContext:
hidden_size: int
source_token_ids: torch.Tensor
num_tokens_per_rank: torch.Tensor
rank_offsets: torch.Tensor
num_max_dispatch_tokens_per_rank: int