[Question]: nccl-ep: Low-Latency Dispatch communication buffer optimization

[mumupika]

Question

Hello! Thanks for your great contributions in providing the nccl-ep extension. I have a few questions regarding the low-latency dispatch buffer optimization described in your paper, which I couldn't quite figure out from the current codebase. I am looking forward to your insights!

In your paper NCCL EP: Towards a Unified Expert Parallel Communication API for NCCL Section IV.D, your team provided a scheme to optimize the layout to reduce low-latency dispatch buffer usage. I quote:

To address the issue in the dispatch, we extend the message header with routing information $R(r, t)$, avoiding layout-based routing. In the Send Tokens step (Section IV-C0a), a token is sent once per destination rank, as opposed to once per expert. The sub-region indexing is unified for dispatch $idx_D(e, r) = r$, not depending on the calling side. The dispatch buffer is reduced by a factor of number of local experts $L$, scaling as $O(N\cdot B\cdot P)$.

Based on my understanding of the original DeepEP implementation, the dispatch receive buffer is organized in a 3D layout: [expert_id][src_rank_id][slot_id]. Thus, the dispatch buffer size scales as O(Experts * numRanks * maxTokensPerRank * MsgSize) (which is also multiplied by 2 for double buffering).

However, when I looked into your open-source implementation, it seems that the code is still applying the previous DeepEP-like expert-major layout. I quote:

nccl/contrib/nccl_ep/device/low_latency.cu

Lines 483 to 489 in 49839df

	const auto recvPtr = reinterpret_cast<uint64_t>(recvBuf) +
	dstExpertLocalIdx * numRanks * maxTokensPerRank * numBytesPerMsg +
	currRank * maxTokensPerRank * numBytesPerMsg +
	slotIdx * numBytesPerMsg;
	size_t expectedDstOffset = recvOff +
	dstExpertLocalIdx * numRanks * maxTokensPerRank * numBytesPerMsg +
	currRank * maxTokensPerRank * numBytesPerMsg + slotIdx * numBytesPerMsg;

Here, the destination pointer is still calculated as dstExpertLocalIdx * numRanks * maxTokensPerRank * ..., which perfectly reflects the $O(E \cdot B \cdot P)$ scaling rather than the optimized $O(N \cdot B \cdot P)$ mentioned in the paper.

Given this context, I have two questions:

1. Is the $O(N \cdot B \cdot P)$ optimization with the $R(r,t)$ header routing planned for a future PR, or is it currently maintained in an internal branch?
2. From a technical perspective, how exactly will the scheme described in the paper be implemented using the $R(r,t)$ header? Furthermore, how is the memory layout on the destination rank organized so that the experts can fetch the data efficiently without introducing extra overhead for shuffling or re-layout? (The previous DeepEP combine send has used complex TMA instructions to efficiently fetch the data)

Thanks again for your time and excellent work!

[xiaofanl-nvidia]

@mumupika maybe you haven't seen this PR: #2040
++ @artpol84 is actively working on this, in case you have anything to add.

[artpol84]

Thank you for the question!

#2040 is the first iteration of this that shows the concept with minimal changes.

Performance-wise it works well for RDMA-enabled configurations but is less efficient for NVLink only. The inefficiency for NVLink comes from findNthToken during permutation.
I’m working on a new version that will close that gap. Will publish it once it’s ready.

The only thing that will remain is to get rid of the staging buffer that 2040 and its successor are still using in the combine stage:
https://github.com/artpol84/nccl/blob/topic/ncclEP/ll_mem_opt_v2/contrib/nccl_ep/device/low_latency.cu#L1373

To address this - partial reduction on the expert side are required. We plan to support this as well in one of the next releases.

Finally, I believe you assumed that a non-permitted output can be provided directly to the user. This is another direction we are exploring currently. But it will only benefit frameworks that can work with unpermuted data.

[mumupika]

Thank you very much for your detailed and patient reply. I will take a closer look at #2040 to fully understand the great optimization your team has implemented. I look forward to your upcoming release!

[FireKingY]

Thank you for the question!

#2040 is the first iteration of this that shows the concept with minimal changes.

Performance-wise it works well for RDMA-enabled configurations but is less efficient for NVLink only. The inefficiency for NVLink comes from findNthToken during permutation. I’m working on a new version that will close that gap. Will publish it once it’s ready.

The only thing that will remain is to get rid of the staging buffer that 2040 and its successor are still using in the combine stage: https://github.com/artpol84/nccl/blob/topic/ncclEP/ll_mem_opt_v2/contrib/nccl_ep/device/low_latency.cu#L1373

To address this - partial reduction on the expert side are required. We plan to support this as well in one of the next releases.

Finally, I believe you assumed that a non-permitted output can be provided directly to the user. This is another direction we are exploring currently. But it will only benefit frameworks that can work with unpermuted data.

Thanks for the clarification, this is very helpful.

I have a follow-up question regarding the performance characteristics you mentioned for RDMA-enabled configurations.

From my understanding, one potential source of the gain is that with the routing header ( R(r, t) ), a token is sent once per destination rank instead of once per expert, which effectively eliminates redundant transmissions when multiple experts are colocated on the same rank.

Is it correct to interpret that the RDMA benefit comes from this per-rank deduplication of token dispatch？

Would appreciate your insights on how these factors break down in practice.

[mumupika]

Is it correct to interpret that the RDMA benefit comes from this per-rank deduplication of token dispatch?

In my understanding, what artpol84 wanted to express is that NVLink provides vastly higher bandwidth and lower latency, meaning that the local re-layout overhead (i.e., findNthToken) becomes much more expensive than simply transmitting duplicate tokens directly to their expert slots.

Therefore, the deduplication scheme provides a net positive performance gain under RDMA configurations because the reduced network transmission load outweighs the permutation overhead. In fact, in my opinion, it's an inevitable trade-off between network deduplication savings and the local findNthToken compute operation.