deepseek-ai/DeepEP

Fork 0

mirror of https://github.com/deepseek-ai/DeepEP.git synced 2026-11-05 00:47:32 +00:00

[PR #716] perf: reuse HybridEP routing allgather buffer #729

New issue

Open

opened 2026-10-27 21:05:40 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:05:40 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/716
Author: @akoumpa
Created: 8/5/2026
Status: 🔄 Open

Base: hybrid-ep ← Head: akoumpa/perf/reuse-hybridep-routing-gather-buffer

📝 Commits (3)

fd14215 perf: reuse HybridEP routing gather buffer
4a12969 test: cover HybridEP allgather buffer reuse
1f2c371 test: assert HybridEP validation topology

📊 Changes

2 files changed (+272 additions, -1 deletions)

View changed files

📝 csrc/hybrid_ep/executor/executor.cu (+5 -1)
➕ tests/test_allgather_grad_validation.py (+267 -0)

📄 Description

Summary

Reuse HybridEP's preallocated allgather output buffer for the NCCL/inter-node routing-map allgather instead of allocating a temporary CUDA tensor during every dispatch.

The routing map is consumed by metadata preprocessing on the same stream before the next dispatch can reuse the buffer, so its lifetime does not overlap the following collective.

Memory impact

The removed temporary allocation contains the globally gathered routing payload on every rank.

For sparse Boolean routing, the saved allocation per rank is:

tokens_per_rank x world_size x num_experts bytes

For dense int16 routing, it is:

tokens_per_rank x world_size x topk x 2 bytes

For the stress configuration used during validation (8192 tokens/rank, 8 ranks, 256 experts, and top-k 36), this avoids:

Sparse routing: 16 MiB per rank, or 128 MiB across the process group.
Dense routing: 4.5 MiB per rank, or 36 MiB across the process group.

Implementation

Wrap CustomAllgather::get_output_buffer() as the NCCL allgather output tensor.
Preserve the existing NCCL behavior for inter-node execution, disabled custom allgather, and unaligned payloads.
Preserve the upstream dense-routing behavior, including the dynamic bool/int16 dtype, routing shape, and byte views used for NCCL.
Add a distributed output and gradient regression test.
Assert the requested node topology so multi-node validation cannot silently execute as a single-node test.

Testing

Validation was performed on H100 GPUs using a multi-node DOCA/RDMA build in both of these topologies:

One node with eight local ranks.
Two nodes with four local ranks per node, for eight global ranks.

Output and gradient regression

Added tests/test_allgather_grad_validation.py, which compares the reused-buffer path with the existing allgather behavior.

The test covers the Cartesian product of:

Dense expert indices and sparse Boolean routing maps.
Routing with and without probabilities.
Three independently executed dispatch/combine cycles.

This produces 12 cases per topology. Each case validates:

Dispatched hidden states and routing probabilities.
Per-expert token counts and dispatch-handle metadata.
Forward combine hidden-state and probability outputs.
Hidden-state and probability gradients for the backward-of-combine communication.
Hidden-state and probability gradients for the backward-of-dispatch communication.
Tensor shapes, dtypes, and optional None behavior.

All 12 cases passed on both the single-node and two-node topologies.

Full HybridEP reference validation

Ran tests/test_hybrid_ep.py with 512 tokens/rank, eight local experts, top-k 8, hidden dimension 1024, and padding multiple 32.

The following passed for both BF16 and FP8 on both topologies:

Standard dispatch and combine.
Sparse and dense routing.
Routing with and without probabilities.
Dispatch with permutation and combine with unpermutation.
Non-fused and fused permute/dispatch execution.
Torch reference comparisons.
Intra-node NVLink and inter-node RDMA execution.

The API-level and kernel-level benchmark loops also completed successfully.

CUDA graph and nonblocking validation

Ran tests/test_graphed_hybrid_ep.py across:

BF16 and FP8.
With and without routing probabilities.
Fused and non-fused execution.
CUDA graph capture/replay and nonblocking dispatch/combine behavior.

All combinations passed on every rank in both topologies.

Harry Zhou's standalone allgather stress test

Thanks to Harry Zhou for providing test_allgather.py, which was used as an additional correctness and stress test.

It was run in both single-node and two-node configurations with its original stress dimensions:

tokens per rank  = 8192
hidden dimension = 7168
global experts   = 256
local experts    = 32
top-k            = 36
global ranks     = 8

For both dense-index and sparse-Boolean routing:

Dispatched hidden tokens matched bitwise.
Dispatched probabilities matched bitwise.
The custom, NCCL, and legacy benchmark loops completed successfully.
Repeated execution completed without exposing buffer-lifetime or reuse issues.

Thanks again to Harry for the focused test coverage and stress configuration.

Static checks

The new regression test also passed Ruff linting, Ruff formatting, Python bytecode compilation, and repository pre-commit checks.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/716 **Author:** [@akoumpa](https://github.com/akoumpa) **Created:** 8/5/2026 **Status:** 🔄 Open **Base:** `hybrid-ep` ← **Head:** `akoumpa/perf/reuse-hybridep-routing-gather-buffer` --- ### 📝 Commits (3) - [`fd14215`](https://github.com/deepseek-ai/DeepEP/commit/fd14215339805eb5768ecd7d57e0be65b3f7e805) perf: reuse HybridEP routing gather buffer - [`4a12969`](https://github.com/deepseek-ai/DeepEP/commit/4a1296958cea04a098acd734983e9d2fff4e25f1) test: cover HybridEP allgather buffer reuse - [`1f2c371`](https://github.com/deepseek-ai/DeepEP/commit/1f2c37154ec7b0416217a0d87649691638279d1a) test: assert HybridEP validation topology ### 📊 Changes **2 files changed** (+272 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `csrc/hybrid_ep/executor/executor.cu` (+5 -1) ➕ `tests/test_allgather_grad_validation.py` (+267 -0) </details> ### 📄 Description ## Summary Reuse HybridEP's preallocated allgather output buffer for the NCCL/inter-node routing-map allgather instead of allocating a temporary CUDA tensor during every dispatch. The routing map is consumed by metadata preprocessing on the same stream before the next dispatch can reuse the buffer, so its lifetime does not overlap the following collective. ## Memory impact The removed temporary allocation contains the globally gathered routing payload on every rank. For sparse Boolean routing, the saved allocation per rank is: ```text tokens_per_rank x world_size x num_experts bytes ``` For dense `int16` routing, it is: ```text tokens_per_rank x world_size x topk x 2 bytes ``` For the stress configuration used during validation (`8192` tokens/rank, `8` ranks, `256` experts, and top-k `36`), this avoids: - Sparse routing: 16 MiB per rank, or 128 MiB across the process group. - Dense routing: 4.5 MiB per rank, or 36 MiB across the process group. ## Implementation - Wrap `CustomAllgather::get_output_buffer()` as the NCCL allgather output tensor. - Preserve the existing NCCL behavior for inter-node execution, disabled custom allgather, and unaligned payloads. - Preserve the upstream dense-routing behavior, including the dynamic `bool`/`int16` dtype, routing shape, and byte views used for NCCL. - Add a distributed output and gradient regression test. - Assert the requested node topology so multi-node validation cannot silently execute as a single-node test. ## Testing Validation was performed on H100 GPUs using a multi-node DOCA/RDMA build in both of these topologies: - One node with eight local ranks. - Two nodes with four local ranks per node, for eight global ranks. ### Output and gradient regression Added `tests/test_allgather_grad_validation.py`, which compares the reused-buffer path with the existing allgather behavior. The test covers the Cartesian product of: - Dense expert indices and sparse Boolean routing maps. - Routing with and without probabilities. - Three independently executed dispatch/combine cycles. This produces 12 cases per topology. Each case validates: - Dispatched hidden states and routing probabilities. - Per-expert token counts and dispatch-handle metadata. - Forward combine hidden-state and probability outputs. - Hidden-state and probability gradients for the backward-of-combine communication. - Hidden-state and probability gradients for the backward-of-dispatch communication. - Tensor shapes, dtypes, and optional `None` behavior. All 12 cases passed on both the single-node and two-node topologies. ### Full HybridEP reference validation Ran `tests/test_hybrid_ep.py` with 512 tokens/rank, eight local experts, top-k 8, hidden dimension 1024, and padding multiple 32. The following passed for both BF16 and FP8 on both topologies: - Standard dispatch and combine. - Sparse and dense routing. - Routing with and without probabilities. - Dispatch with permutation and combine with unpermutation. - Non-fused and fused permute/dispatch execution. - Torch reference comparisons. - Intra-node NVLink and inter-node RDMA execution. The API-level and kernel-level benchmark loops also completed successfully. ### CUDA graph and nonblocking validation Ran `tests/test_graphed_hybrid_ep.py` across: - BF16 and FP8. - With and without routing probabilities. - Fused and non-fused execution. - CUDA graph capture/replay and nonblocking dispatch/combine behavior. All combinations passed on every rank in both topologies. ### Harry Zhou's standalone allgather stress test Thanks to Harry Zhou for providing [`test_allgather.py`](https://github.com/harryzhou2000/DeepEP/blob/hhanyu/hybrid-ep-sparse-opt-2/tests/test_allgather.py), which was used as an additional correctness and stress test. It was run in both single-node and two-node configurations with its original stress dimensions: ```text tokens per rank = 8192 hidden dimension = 7168 global experts = 256 local experts = 32 top-k = 36 global ranks = 8 ``` For both dense-index and sparse-Boolean routing: - Dispatched hidden tokens matched bitwise. - Dispatched probabilities matched bitwise. - The custom, NCCL, and legacy benchmark loops completed successfully. - Repeated execution completed without exposing buffer-lifetime or reuse issues. Thanks again to Harry for the focused test coverage and stress configuration. ### Static checks The new regression test also passed Ruff linting, Ruff formatting, Python bytecode compilation, and repository pre-commit checks. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-27 21:05:40 +00:00

bsakaguchi referenced this issue

2026-10-27 21:05:43 +00:00

[PR #729] docs: clarify cached dispatch weights #735