microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-15 03:40:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf test

Browse Source
This commit is contained in:
Binyang Li
2026-07-08 23:16:00 +00:00
parent b1d0893da9
commit 5612702fd5
10 changed files with 391 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/mscclpp/ep/_cpp.py
View File

@@ -18,6 +18,8 @@ Config = getattr(_cpp, "Config", None)
def get_low_latency_rdma_size_hint(
num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int
num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int, num_topk: int
) -> int:
return _cpp.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)
return _cpp.get_low_latency_rdma_size_hint(
num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts, num_topk
)

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

@@ -117,7 +117,7 @@ class LowLatencyBackend:
self.dispatch_requires_quantization = self.quant_dtype is not None
num_rdma_bytes = get_low_latency_rdma_size_hint(
self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts
self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts, self.topk
)
self._dispatch_scales: Optional[torch.Tensor] = None
self._dispatch_src_info: Optional[torch.Tensor] = None
@@ -131,9 +131,9 @@ class LowLatencyBackend:
mode=self.mode,
num_qps_per_rank=config.num_rdma_qps_per_rank,
)
# LL always uses the RDMA transport, but a single-node LL job is not
# internode topology-wise. num_rdma_ranks > 1 iff world_size spans more
# than one local NVLink domain.
# LL uses the registered symmetric buffer for both IPC/NVLink and RDMA-backed
# modes. A single-node LL job is not internode topology-wise.
# num_rdma_ranks > 1 iff world_size spans more than one local NVLink domain.
self._is_internode = self._runtime.get_num_rdma_ranks() > 1
def is_available(self) -> bool:
@@ -165,6 +165,7 @@ class LowLatencyBackend:
self._runtime.cpp_runtime.dispatch(
input.data_ptr(),
topk_ids.data_ptr(),
weights.data_ptr(),
out_buf.data_ptr(),
0 if packed_scales is None else packed_scales.data_ptr(),
src_info.data_ptr(),