[PR #693] Pipeline direct NVLink dispatch with triple-buffered prefetch #711

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/693
Author: @usernamehaha2022
Created: 7/22/2026
Status: 🔄 Open

Base: main ← Head: main

📝 Commits (3)

b9ef9b3 perf: double-buffer direct NVLink dispatch
d2be99a perf: prefetch direct dispatch tokens across NVLink stores
45c96fa fix: keep dispatch TMA buffer count in sync

📊 Changes

2 files changed (+81 additions, -18 deletions)

View changed files

📝 csrc/kernels/elastic/dispatch.hpp (+13 -2)
📝 deep_ep/include/deep_ep/impls/dispatch.cuh (+68 -16)

📄 Description

Summary

This PR pipelines direct NVLink dispatch using three shared-memory TMA buffers per dispatch warp.

Previously, a dispatch warp prepared tokens sequentially around remote S2G stores. With this change, the next token's hidden data, scale factors, metadata, and routing information are prefetched while previous remote stores remain in flight.

At steady state, the three buffers are used for:

A previous token whose remote S2G store may still be in flight.
The current token being sent to remote ranks.
The next token being prefetched from global memory.

Before recycling the oldest buffer, tma_store_wait<1>() reduces the number of outstanding store groups to one. This keeps buffer reuse safe without fully draining all remote stores between tokens.

The shared-memory calculation now accounts for all three staging buffers, and direct NVLink dispatch is limited to at most 14 dispatch warps. The non-NVLink path remains single-buffered.

Validation

Tested on 8 x NVIDIA H100 GPUs with direct NVLink dispatch.

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
EP_DISABLE_GIN=1 \
python -u tests/elastic/test_ep.py \
  --num-processes 8 \
  --allow-hybrid-mode 0 \
  --test-first-only \
  --ignore-local-traffic

Average across 8 ranks:

Metric	`main`	This PR	Change
Dispatch	250GB/s	280 GB/s	30
Combine	300 GB/s	300 GB/s	-

With rank unbalanced ratio set to 1.5:

Metric	`main`	This PR	Change
Dispatch	203GB/s	220 GB/s	17
Combine	209 GB/s	209 GB/s	-

_{🔄 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/693 **Author:** [@usernamehaha2022](https://github.com/usernamehaha2022) **Created:** 7/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `main` --- ### 📝 Commits (3) - [`b9ef9b3`](https://github.com/deepseek-ai/DeepEP/commit/b9ef9b3f2e8dc33e5acf8e4826aab65c41fd5b77) perf: double-buffer direct NVLink dispatch - [`d2be99a`](https://github.com/deepseek-ai/DeepEP/commit/d2be99a22990606775dd21623ddfa1a702e42e91) perf: prefetch direct dispatch tokens across NVLink stores - [`45c96fa`](https://github.com/deepseek-ai/DeepEP/commit/45c96fa945081bd14d598ef294daf64cbded4b5f) fix: keep dispatch TMA buffer count in sync ### 📊 Changes **2 files changed** (+81 additions, -18 deletions) <details> <summary>View changed files</summary> 📝 `csrc/kernels/elastic/dispatch.hpp` (+13 -2) 📝 `deep_ep/include/deep_ep/impls/dispatch.cuh` (+68 -16) </details> ### 📄 Description ## Summary This PR pipelines direct NVLink dispatch using three shared-memory TMA buffers per dispatch warp. Previously, a dispatch warp prepared tokens sequentially around remote S2G stores. With this change, the next token's hidden data, scale factors, metadata, and routing information are prefetched while previous remote stores remain in flight. At steady state, the three buffers are used for: 1. A previous token whose remote S2G store may still be in flight. 2. The current token being sent to remote ranks. 3. The next token being prefetched from global memory. Before recycling the oldest buffer, `tma_store_wait<1>()` reduces the number of outstanding store groups to one. This keeps buffer reuse safe without fully draining all remote stores between tokens. The shared-memory calculation now accounts for all three staging buffers, and direct NVLink dispatch is limited to at most 14 dispatch warps. The non-NVLink path remains single-buffered. ## Validation Tested on 8 x NVIDIA H100 GPUs with direct NVLink dispatch. ```bash CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ EP_DISABLE_GIN=1 \ python -u tests/elastic/test_ep.py \ --num-processes 8 \ --allow-hybrid-mode 0 \ --test-first-only \ --ignore-local-traffic ``` Average across 8 ranks: | Metric | `main` | This PR | Change | |---|---:|---:|---:| | Dispatch | 250GB/s | 280 GB/s | 30 | | Combine | 300 GB/s | 300 GB/s |- | With rank unbalanced ratio set to 1.5: | Metric | `main` | This PR | Change | |---|---:|---:|---:| | Dispatch | 203GB/s | 220 GB/s | 17 | | Combine | 209 GB/s | 209 GB/s |- | --- <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:29 +00:00

bsakaguchi referenced this issue

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

[PR #711] [CLOSED] Adapt nvshmem37 #724