[PR #787] perf: reuse compact dispatch sort permutations for cached handles #785

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/787
Author: @yuweih205
Created: 10/17/2026
Status: 🔄 Open

Base: main ← Head: perf/cache-compact-sort-1017

📝 Commits (1)

b363ade Sort once per handle in deterministic dispatch

📊 Changes

1 file changed (+13 additions, -10 deletions)

View changed files

📝 deep_ep/buffers/ep.py (+13 -10)

📄 Description

Compact deterministic dispatch sorts the received rows by source index on every call. For a cached dispatch the handle already fixes that order: the same slots are refilled, the receive counts are the same, and topk_idx is asserted immutable. This stores the permutation on the handle next to cached_recv_src_metadata_before_sort at the first dispatch and reuses it, so cached compact dispatch no longer builds the sort keys, masks the unused capacity, or sorts.

The stored permutation follows the same stream and lifetime rules as the handle's existing cached metadata, so no per-call checks are needed and the reuse also holds under CUDA Graph capture and for deferred epilogues waited on another stream. Each handle keeps one extra int64 tensor of capacity rows (128 KiB at 16k rows). No API, default, native or test-file changes; tests/ep/test_ep.py already compares repeated cached dispatch against the reference, and the EP8 suite passes on this branch with CPU counts (direct and deferred epilogues) and in the sync-free mode.

Measured on 8×H200 (PyTorch 2.13.0+cu132), one cached dispatch per rank, 2048 tokens/rank, hidden 4096, top-2, BF16, deterministic, do_cpu_sync=False, under both EP_AVOID_RECORD_STREAM policies; outputs (payload, topk_idx, weights, source index) are bitwise identical to main across 608 fresh/cached/deferred/graph cases per policy:

Per cached dispatch	main	this PR
Device kernels and memsets	24	6
Host enqueue time	328 µs	129 µs
Synchronized wall time, 8-rank max	439 µs	280 µs
Kernels in a captured cached dispatch	15 (incl. sort)	6

The removed kernels (radix sort, arange, compare, clone, masked_fill, memsets) account for about 46 µs of device time per call; the wall-time reduction is mostly host overhead and is a component measurement, not a training-throughput claim. Independent of #786, which removes the unused sorts of the expanded layout; both touch deterministic_sort and the later one needs a trivial rebase.

_{🔄 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/787 **Author:** [@yuweih205](https://github.com/yuweih205) **Created:** 10/17/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `perf/cache-compact-sort-1017` --- ### 📝 Commits (1) - [`b363ade`](https://github.com/deepseek-ai/DeepEP/commit/b363aded73a2d2eb632bd895802830d4a747dcab) Sort once per handle in deterministic dispatch ### 📊 Changes **1 file changed** (+13 additions, -10 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/buffers/ep.py` (+13 -10) </details> ### 📄 Description Compact deterministic dispatch sorts the received rows by source index on every call. For a cached dispatch the handle already fixes that order: the same slots are refilled, the receive counts are the same, and `topk_idx` is asserted immutable. This stores the permutation on the handle next to `cached_recv_src_metadata_before_sort` at the first dispatch and reuses it, so cached compact dispatch no longer builds the sort keys, masks the unused capacity, or sorts. The stored permutation follows the same stream and lifetime rules as the handle's existing cached metadata, so no per-call checks are needed and the reuse also holds under CUDA Graph capture and for deferred epilogues waited on another stream. Each handle keeps one extra int64 tensor of capacity rows (128 KiB at 16k rows). No API, default, native or test-file changes; `tests/ep/test_ep.py` already compares repeated cached dispatch against the reference, and the EP8 suite passes on this branch with CPU counts (direct and deferred epilogues) and in the sync-free mode. Measured on 8×H200 (PyTorch 2.13.0+cu132), one cached dispatch per rank, 2048 tokens/rank, hidden 4096, top-2, BF16, deterministic, `do_cpu_sync=False`, under both `EP_AVOID_RECORD_STREAM` policies; outputs (payload, `topk_idx`, weights, source index) are bitwise identical to main across 608 fresh/cached/deferred/graph cases per policy: | Per cached dispatch | main | this PR | |---|---:|---:| | Device kernels and memsets | 24 | 6 | | Host enqueue time | 328 µs | 129 µs | | Synchronized wall time, 8-rank max | 439 µs | 280 µs | | Kernels in a captured cached dispatch | 15 (incl. sort) | 6 | The removed kernels (radix sort, `arange`, compare, `clone`, `masked_fill`, memsets) account for about 46 µs of device time per call; the wall-time reduction is mostly host overhead and is a component measurement, not a training-throughput claim. Independent of #786, which removes the unused sorts of the expanded layout; both touch `deterministic_sort` and the later one needs a trivial rebase. --- <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:19 +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#785

No description provided.

Rows
Columns