[PR #685] Fix TypeError in cached ElasticBuffer dispatch with default num_sms #705

Open
opened 2026-10-27 21:05:26 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/685
Author: @yurekami
Created: 7/12/2026
Status: 🔄 Open

Base: main ← Head: ghcontrib/elastic-cached-num-sms


📝 Commits (3)

  • 3caa7ed Fix TypeError in cached ElasticBuffer dispatch with default num_sms
  • 3ef4840 Run yapf and ruff
  • 262bcb2 Update dispatch num_sms docstring for cached-handle reuse

📊 Changes

1 file changed (+103 additions, -168 deletions)

View changed files

📝 deep_ep/buffers/elastic.py (+103 -168)

📄 Description

Summary

ElasticBuffer.dispatch reads num_experts at deep_ep/buffers/elastic.py:928 but
only infers it from the handle 14 lines later, at line 942:

926  # Automatic decide SM and QP count
927  num_topk = (handle.topk_idx if topk_idx is None else topk_idx).shape[1]
928  num_sms = self.get_theoretical_num_sms(num_experts, num_topk) if num_sms == 0 else num_sms
...
937  if handle is not None:
942      num_experts = value_or(num_experts, handle.num_experts)   # too late

Both defaults make this reachable: num_experts: Optional[int] = None (L860) and
num_sms: int = 0 (L863). The docstring documents exactly this call —
"num_experts: ... Inferred from handle if provided" (L893) and "num_sms: ...
0 for automatic" (L897) — so the documented cached call

recv_x, _, _, handle, _ = buffer.dispatch(x, topk_idx=topk_idx, num_experts=256)
recv_x, _, _, _, _      = buffer.dispatch(x, handle=handle)   # <-- raises

enters L928 with num_experts=None, giving:

get_theoretical_num_sms(None, num_topk)     elastic.py:928
  -> get_expected_topk(self.num_ranks)      elastic.py:778
     -> assert num_experts % num_groups == 0  elastic.py:771
TypeError: unsupported operand type(s) for %: 'NoneType' and 'int'

Callers that pass num_sms explicitly are unaffected, which is why the test suite
never hits it: tests/elastic/test_ep.py:65 resolves num_sms itself up front and
passes it into every dispatch call (L146, L165, L212, L223), so num_sms is never 0
on entry.

Reproduction

I do not have an NVIDIA device, so this is reproduced host-side — the failure occurs
in pure Python at L928, before any native/NVSHMEM call. Loading the unmodified
elastic.py (only deep_ep._C stubbed) and calling the real
get_theoretical_num_sms with the num_experts that L928 passes it in cached mode:

[control] num_experts=256   -> OK, returns num_sms = 4
[defect]  num_experts=None  -> TypeError: unsupported operand type(s) for %: 'NoneType' and 'int'

  File "deep_ep/buffers/elastic.py", line 778, in get_theoretical_num_sms
    num_expected_topk = get_expected_topk(self.num_ranks)
  File "deep_ep/buffers/elastic.py", line 771, in get_expected_topk
    assert num_experts % num_groups == 0
TypeError: unsupported operand type(s) for %: 'NoneType' and 'int'

On a GPU box the two-line snippet above reproduces it directly.

Fix

Take the SM count from the handle when one is provided, which is what combine
already does for the same value at elastic.py:1086:

1086  num_sms = handle.num_sms if num_sms == 0 else num_sms

EPHandle stores num_sms (L63, L82) precisely so the cached path can reuse it.

Design context

The alternative fix is to hoist the handle-unpacking block (L937-948) above the SM/QP
computation so num_experts is populated before L928. I did not take it: the handle's
channel metadata (channel_linked_list, dst_buffer_slot_idx, the psum tensors) was
laid out under the SM count the original dispatch actually ran with, and recomputing a
theoretical value here could return a different count than the handle was built under.
Reusing handle.num_sms keeps cached dispatch consistent with both the handle and with
combine, which must agree on channel geometry for the same handle.

Scope

One line in dispatch, plus a comment. Behavior changes only on the branch that
currently raises:

  • cached + num_sms=0 -> now uses handle.num_sms (was: TypeError)
  • fresh dispatch (handle is None) -> unchanged, still get_theoretical_num_sms(num_experts, ...)
  • any caller passing num_sms explicitly -> unchanged

Test plan

  • cached dispatch with default num_sms no longer raises; resolves to handle.num_sms
  • fresh dispatch with default num_sms still resolves via get_theoretical_num_sms
  • explicit num_sms still respected in both modes
  • not run on device (no NVIDIA hardware available); tests/elastic/test_ep.py passes
    num_sms explicitly, so it does not cover this branch either way

Note on format-check

CI is red here on pre-existing lint in this file, not on this change. format.sh
lints only files that differ from origin/main, so main never lints
elastic.py and these never surface there. Running the pinned tooling against
unmodified origin/main (60d4403):

$ ruff 0.6.5 check deep_ep/buffers/elastic.py     # repo config
elastic.py:6:54:    F401  `typing.List` imported but unused
elastic.py:930:51:  F541  f-string without any placeholders
elastic.py:1088:51: F541  f-string without any placeholders
Found 3 errors.

$ yapf 0.40.2 --diff deep_ep/buffers/elastic.py   # column_limit=140
439 diff lines

This PR adds zero new violations — ruff reports the same 3 errors and yapf the
same 439 diff lines with and without it, and my added lines are all <=115 chars.
I deliberately did not fix the pre-existing lint here: it would either be a
1-of-3 fix that still leaves CI red, or a ~439-line reformat unrelated to a
4-line bugfix. Happy to send that as a separate PR if you want the file made
lint-clean.


🔄 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/685 **Author:** [@yurekami](https://github.com/yurekami) **Created:** 7/12/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `ghcontrib/elastic-cached-num-sms` --- ### 📝 Commits (3) - [`3caa7ed`](https://github.com/deepseek-ai/DeepEP/commit/3caa7ed3108645cbe87fec37886ceb11c4fdfcd0) Fix TypeError in cached ElasticBuffer dispatch with default num_sms - [`3ef4840`](https://github.com/deepseek-ai/DeepEP/commit/3ef48401085ae2d291a75ed9bcda351838006280) Run yapf and ruff - [`262bcb2`](https://github.com/deepseek-ai/DeepEP/commit/262bcb2ba4fcf51bd47fedcd202db1df9f09142f) Update dispatch num_sms docstring for cached-handle reuse ### 📊 Changes **1 file changed** (+103 additions, -168 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/buffers/elastic.py` (+103 -168) </details> ### 📄 Description ## Summary `ElasticBuffer.dispatch` reads `num_experts` at `deep_ep/buffers/elastic.py:928` but only infers it from the handle 14 lines later, at line 942: ```python 926 # Automatic decide SM and QP count 927 num_topk = (handle.topk_idx if topk_idx is None else topk_idx).shape[1] 928 num_sms = self.get_theoretical_num_sms(num_experts, num_topk) if num_sms == 0 else num_sms ... 937 if handle is not None: 942 num_experts = value_or(num_experts, handle.num_experts) # too late ``` Both defaults make this reachable: `num_experts: Optional[int] = None` (L860) and `num_sms: int = 0` (L863). The docstring documents exactly this call — "`num_experts`: ... Inferred from `handle` if provided" (L893) and "`num_sms`: ... 0 for automatic" (L897) — so the documented cached call ```python recv_x, _, _, handle, _ = buffer.dispatch(x, topk_idx=topk_idx, num_experts=256) recv_x, _, _, _, _ = buffer.dispatch(x, handle=handle) # <-- raises ``` enters L928 with `num_experts=None`, giving: ``` get_theoretical_num_sms(None, num_topk) elastic.py:928 -> get_expected_topk(self.num_ranks) elastic.py:778 -> assert num_experts % num_groups == 0 elastic.py:771 TypeError: unsupported operand type(s) for %: 'NoneType' and 'int' ``` Callers that pass `num_sms` explicitly are unaffected, which is why the test suite never hits it: `tests/elastic/test_ep.py:65` resolves `num_sms` itself up front and passes it into every dispatch call (L146, L165, L212, L223), so `num_sms` is never 0 on entry. ## Reproduction I do not have an NVIDIA device, so this is reproduced host-side — the failure occurs in pure Python at L928, before any native/NVSHMEM call. Loading the unmodified `elastic.py` (only `deep_ep._C` stubbed) and calling the real `get_theoretical_num_sms` with the `num_experts` that L928 passes it in cached mode: ``` [control] num_experts=256 -> OK, returns num_sms = 4 [defect] num_experts=None -> TypeError: unsupported operand type(s) for %: 'NoneType' and 'int' File "deep_ep/buffers/elastic.py", line 778, in get_theoretical_num_sms num_expected_topk = get_expected_topk(self.num_ranks) File "deep_ep/buffers/elastic.py", line 771, in get_expected_topk assert num_experts % num_groups == 0 TypeError: unsupported operand type(s) for %: 'NoneType' and 'int' ``` On a GPU box the two-line snippet above reproduces it directly. ## Fix Take the SM count from the handle when one is provided, which is what `combine` already does for the same value at `elastic.py:1086`: ```python 1086 num_sms = handle.num_sms if num_sms == 0 else num_sms ``` `EPHandle` stores `num_sms` (L63, L82) precisely so the cached path can reuse it. ## Design context The alternative fix is to hoist the handle-unpacking block (L937-948) above the SM/QP computation so `num_experts` is populated before L928. I did not take it: the handle's channel metadata (`channel_linked_list`, `dst_buffer_slot_idx`, the psum tensors) was laid out under the SM count the original dispatch actually ran with, and recomputing a theoretical value here could return a different count than the handle was built under. Reusing `handle.num_sms` keeps cached dispatch consistent with both the handle and with `combine`, which must agree on channel geometry for the same handle. ## Scope One line in `dispatch`, plus a comment. Behavior changes only on the branch that currently raises: - cached + `num_sms=0` -> now uses `handle.num_sms` (was: `TypeError`) - fresh dispatch (`handle is None`) -> unchanged, still `get_theoretical_num_sms(num_experts, ...)` - any caller passing `num_sms` explicitly -> unchanged ## Test plan - [x] cached dispatch with default `num_sms` no longer raises; resolves to `handle.num_sms` - [x] fresh dispatch with default `num_sms` still resolves via `get_theoretical_num_sms` - [x] explicit `num_sms` still respected in both modes - [ ] not run on device (no NVIDIA hardware available); `tests/elastic/test_ep.py` passes `num_sms` explicitly, so it does not cover this branch either way ## Note on `format-check` CI is red here on pre-existing lint in this file, not on this change. `format.sh` lints only files that differ from `origin/main`, so `main` never lints `elastic.py` and these never surface there. Running the pinned tooling against **unmodified `origin/main`** (`60d4403`): ``` $ ruff 0.6.5 check deep_ep/buffers/elastic.py # repo config elastic.py:6:54: F401 `typing.List` imported but unused elastic.py:930:51: F541 f-string without any placeholders elastic.py:1088:51: F541 f-string without any placeholders Found 3 errors. $ yapf 0.40.2 --diff deep_ep/buffers/elastic.py # column_limit=140 439 diff lines ``` This PR adds **zero** new violations — ruff reports the same 3 errors and yapf the same 439 diff lines with and without it, and my added lines are all <=115 chars. I deliberately did not fix the pre-existing lint here: it would either be a 1-of-3 fix that still leaves CI red, or a ~439-line reformat unrelated to a 4-line bugfix. Happy to send that as a separate PR if you want the file made lint-clean. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
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#705
No description provided.