deepseek-ai/DeepEP

Fork 0

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

[PR #800] Support optional `topk_weights` in the reference dispatch #794

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/800
Author: @lsyx6399
Created: 10/27/2026
Status: 🔄 Open

Base: main ← Head: fix/refs-dispatch-optional-topk-weights

📝 Commits (1)

767af27 Support optional topk_weights in the reference dispatch

📊 Changes

3 files changed (+219 additions, -6 deletions)

View changed files

📝 deep_ep/utils/refs.py (+7 -6)
➕ tests/utils/_run_refs_test.py (+121 -0)
➕ tests/utils/test_refs.py (+91 -0)

📄 Description

Problem

The real EPBuffer.dispatch (deep_ep/buffers/ep.py) already supports optional top-k weights:

def dispatch(self, ..., topk_weights: Optional[torch.Tensor] = None, ...) -> ...:
    """...
    topk_weights can be optionally provided (e.g. for backward pass with cached expand).
    ...
    recv_topk_weights: received expert weights (None if topk_weights was not provided).
    """

The pure-torch reference implementation in deep_ep/utils/refs.py did not match this contract: dispatch assumed topk_weights was always present (it called topk_weights[...], torch.cat(..., topk_weights.dtype), and dist.all_to_all_single(recv_topk_weights, ...) unconditionally), so passing topk_weights=None — a documented, supported input of the real kernel — raised AttributeError: 'NoneType' object has no attribute ....

This left the optional-weights path:

uncovered by the reference (it couldn't be used as a correctness oracle for the weights=None case), and
uncovered by tests (tests/ep/test_ep.py always generates topk_weights via torch.topk, never None).

The stale # TODO: make top-k weight fully optional in refs.py confirms this was a known gap.

Fix

Align the reference with the real kernel:

When topk_weights is None, skip the weights send/recv and return recv_topk_weights=None.
Keep the with-weights path identical (same tensors exchanged, same return shape).
Update the docstring to state recv_topk_weights is None when topk_weights is not provided, matching ep.py's wording.
Remove the now-resolved # TODO: make top-k weight fully optional.

combine / generate_pre_combine_data / ordered_accumulate are unchanged — they don't carry weights.

Tests

Added tests/utils/test_refs.py (+ worker tests/utils/_run_refs_test.py): a 2-rank gloo group on CPU, with no compiled deep_ep._C, NCCL, GPU, or NPU required.

The worker loads refs via importlib with the package-level deep_ep._C / NCCL / JIT init stubbed, so it runs anywhere.
A file-based gloo init is used (avoids the libuv dependency of the c10d TCP store, which some CPU-only torch builds lack).
All ranks are launched concurrently (subprocess.Popen) since gloo collectives block until every rank joins.

Cases verified:

recv_topk_weights is None when topk_weights=None.
Routing is identical with vs. without weights — recv_x, recv_topk_idx, recv_src_token_idx, num_recv_tokens_per_rank all match.
With weights provided, recv_topk_weights has the expected shape.
Tuple input (the use_fp8 = isinstance(x, tuple) path) + weights=None still exchanges the SF tensor correctly. (Real FP8 dtypes need a GPU; here both tuple members are float32 so gloo can exchange them, exercising the same SF branch on CPU.)

Verified: the test passes on this fix and fails on the pre-fix refs.py (the worker raises AttributeError and emits no RESULT line), confirming the test exercises the gap.

Behavior impact

Only the topk_weights=None path changes (previously crashed, now returns None weights as the real kernel does).
The with-weights path is byte-for-byte identical.
This is a reference-implementation + test change; no kernel or runtime behaviour changes.

_{🔄 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/800 **Author:** [@lsyx6399](https://github.com/lsyx6399) **Created:** 10/27/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/refs-dispatch-optional-topk-weights` --- ### 📝 Commits (1) - [`767af27`](https://github.com/deepseek-ai/DeepEP/commit/767af27b5d8b78be5b923f0556620cfce27b9b6a) Support optional `topk_weights` in the reference dispatch ### 📊 Changes **3 files changed** (+219 additions, -6 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/utils/refs.py` (+7 -6) ➕ `tests/utils/_run_refs_test.py` (+121 -0) ➕ `tests/utils/test_refs.py` (+91 -0) </details> ### 📄 Description ## Problem The real `EPBuffer.dispatch` (`deep_ep/buffers/ep.py`) already supports optional top-k weights: ```python def dispatch(self, ..., topk_weights: Optional[torch.Tensor] = None, ...) -> ...: """... topk_weights can be optionally provided (e.g. for backward pass with cached expand). ... recv_topk_weights: received expert weights (None if topk_weights was not provided). """ ``` The pure-torch reference implementation in `deep_ep/utils/refs.py` did **not** match this contract: `dispatch` assumed `topk_weights` was always present (it called `topk_weights[...]`, `torch.cat(..., topk_weights.dtype)`, and `dist.all_to_all_single(recv_topk_weights, ...)` unconditionally), so passing `topk_weights=None` — a documented, supported input of the real kernel — raised `AttributeError: 'NoneType' object has no attribute ...`. This left the optional-weights path: - **uncovered by the reference** (it couldn't be used as a correctness oracle for the `weights=None` case), and - **uncovered by tests** (`tests/ep/test_ep.py` always generates `topk_weights` via `torch.topk`, never `None`). The stale `# TODO: make top-k weight fully optional` in `refs.py` confirms this was a known gap. ## Fix Align the reference with the real kernel: - When `topk_weights is None`, skip the weights send/recv and return `recv_topk_weights=None`. - Keep the with-weights path identical (same tensors exchanged, same return shape). - Update the docstring to state `recv_topk_weights is None when topk_weights is not provided`, matching `ep.py`'s wording. - Remove the now-resolved `# TODO: make top-k weight fully optional`. `combine` / `generate_pre_combine_data` / `ordered_accumulate` are unchanged — they don't carry weights. ## Tests Added `tests/utils/test_refs.py` (+ worker `tests/utils/_run_refs_test.py`): a 2-rank `gloo` group on **CPU**, with **no compiled `deep_ep._C`, NCCL, GPU, or NPU** required. - The worker loads `refs` via `importlib` with the package-level `deep_ep._C` / NCCL / JIT init stubbed, so it runs anywhere. - A file-based gloo init is used (avoids the libuv dependency of the c10d TCP store, which some CPU-only torch builds lack). - All ranks are launched concurrently (`subprocess.Popen`) since gloo collectives block until every rank joins. Cases verified: 1. `recv_topk_weights is None` when `topk_weights=None`. 2. Routing is identical with vs. without weights — `recv_x`, `recv_topk_idx`, `recv_src_token_idx`, `num_recv_tokens_per_rank` all match. 3. With weights provided, `recv_topk_weights` has the expected shape. 4. Tuple input (the `use_fp8 = isinstance(x, tuple)` path) + `weights=None` still exchanges the SF tensor correctly. (Real FP8 dtypes need a GPU; here both tuple members are `float32` so gloo can exchange them, exercising the same SF branch on CPU.) Verified: the test passes on this fix and **fails on the pre-fix `refs.py`** (the worker raises `AttributeError` and emits no `RESULT` line), confirming the test exercises the gap. ## Behavior impact - **Only the `topk_weights=None` path changes** (previously crashed, now returns `None` weights as the real kernel does). - The with-weights path is byte-for-byte identical. - This is a reference-implementation + test change; no kernel or runtime behaviour changes. --- <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:26 +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#794

No description provided.

Rows
Columns

[PR #800] Support optional topk_weights in the reference dispatch #794