[PR #769] Add opt-in borrowed receive storage for compact BF16 dispatch #769

New issue

Open

opened 2026-10-27 21:06:11 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:06:11 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/769
Author: @yuweih205
Created: 9/30/2026
Status: 🔄 Open

Base: main ← Head: perf/borrowed-recv-view

📝 Commits (1)

100350f Add opt-in borrowed receive storage for compact BF16 dispatch

📊 Changes

8 files changed (+345 additions, -51 deletions)

View changed files

📝 README.md (+6 -0)
📝 csrc/buffers/ep.hpp (+27 -5)
📝 csrc/kernels/ep/dispatch.hpp (+6 -6)
📝 deep_ep/__init__.py (+1 -0)
📝 deep_ep/buffers/ep.py (+52 -22)
➕ deep_ep/buffers/recv_view.py (+62 -0)
📝 deep_ep/include/deep_ep/impls/ep/dispatch_copy_epilogue.cuh (+36 -18)
➕ tests/ep/test_recv_view.py (+155 -0)

📄 Description

dispatch(..., borrow_recv=True) skips the activation copy of the dispatch epilogue: the epilogue writes only metadata plus a row map and returns a DispatchRecvView (slab, row_indices) over the communication buffer, which an indexed consumer such as a gathered GEMM reads directly. The view holds an exclusive lease: release(*streams) orders the next communication after the consumers, and dispatch/combine/LB calls/destroy() reject an active view. Supported: fresh compact BF16 dispatch within one NVLink domain with CPU counts and defer_epilogue=False; everything else is rejected up front. Default behavior is unchanged.

8×H200, deterministic, hidden 3072, top-8, 16 local experts, dispatch plus the same indexed consumer: −13.1% at 8192 tokens/rank and −13.5% at 16384, but +15.9% at 1024, hence opt-in. Fixed-route serial forward only; no overlap or training-throughput claim. Combined with DeepGEMM #463 in a live EP8 + gate/up GEMM + SwiGLU step: −17%. Raw cycles and the reproducer: gist.

Test: python tests/ep/test_recv_view.py --num-processes 8 (bitwise payload/metadata against the gathered inputs, replay/combine after release, delayed cross-stream readers, invalid-mode rejection).

Re-validated 2026-10-16 at 71157be (100350f differs only in documentation; native code byte-identical to the measured build) on 8×H200, PyTorch 2.13.0+cu132: test_recv_view.py 16/16 cases under both EP_AVOID_RECORD_STREAM policies; the EP8 functional suite passes with CPU counts for direct and deferred epilogues. With --do-cpu-sync 0 the unmodified suite stops at its full-capacity recv_x_twice comparison on this branch exactly as on upstream main; with the valid-prefix comparison from #786 it passes on both.

_{🔄 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/769 **Author:** [@yuweih205](https://github.com/yuweih205) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `perf/borrowed-recv-view` --- ### 📝 Commits (1) - [`100350f`](https://github.com/deepseek-ai/DeepEP/commit/100350fd82da8485d4fd904c88883d91dd3b8604) Add opt-in borrowed receive storage for compact BF16 dispatch ### 📊 Changes **8 files changed** (+345 additions, -51 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+6 -0) 📝 `csrc/buffers/ep.hpp` (+27 -5) 📝 `csrc/kernels/ep/dispatch.hpp` (+6 -6) 📝 `deep_ep/__init__.py` (+1 -0) 📝 `deep_ep/buffers/ep.py` (+52 -22) ➕ `deep_ep/buffers/recv_view.py` (+62 -0) 📝 `deep_ep/include/deep_ep/impls/ep/dispatch_copy_epilogue.cuh` (+36 -18) ➕ `tests/ep/test_recv_view.py` (+155 -0) </details> ### 📄 Description `dispatch(..., borrow_recv=True)` skips the activation copy of the dispatch epilogue: the epilogue writes only metadata plus a row map and returns a `DispatchRecvView` (`slab`, `row_indices`) over the communication buffer, which an indexed consumer such as a gathered GEMM reads directly. The view holds an exclusive lease: `release(*streams)` orders the next communication after the consumers, and dispatch/combine/LB calls/`destroy()` reject an active view. Supported: fresh compact BF16 dispatch within one NVLink domain with CPU counts and `defer_epilogue=False`; everything else is rejected up front. Default behavior is unchanged. 8×H200, deterministic, hidden 3072, top-8, 16 local experts, dispatch plus the same indexed consumer: −13.1% at 8192 tokens/rank and −13.5% at 16384, but +15.9% at 1024, hence opt-in. Fixed-route serial forward only; no overlap or training-throughput claim. Combined with DeepGEMM #463 in a live EP8 + gate/up GEMM + SwiGLU step: −17%. Raw cycles and the reproducer: [gist](https://gist.github.com/yuweih205/d3fd1e4c3fcb97364f36ec8fda996cfe). Test: `python tests/ep/test_recv_view.py --num-processes 8` (bitwise payload/metadata against the gathered inputs, replay/combine after release, delayed cross-stream readers, invalid-mode rejection). Re-validated 2026-10-16 at 71157be (100350f differs only in documentation; native code byte-identical to the measured build) on 8×H200, PyTorch 2.13.0+cu132: `test_recv_view.py` 16/16 cases under both `EP_AVOID_RECORD_STREAM` policies; the EP8 functional suite passes with CPU counts for direct and deferred epilogues. With `--do-cpu-sync 0` the unmodified suite stops at its full-capacity `recv_x_twice` comparison on this branch exactly as on upstream main; with the valid-prefix comparison from #786 it passes on both. --- <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:06:11 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepEP#769

No description provided.

Rows
Columns