mirror of
https://github.com/deepseek-ai/DeepEP.git
synced 2026-11-05 00:47:32 +00:00
[PR #685] Fix TypeError in cached ElasticBuffer dispatch with default num_sms #705
Labels
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
No due date set.
Dependencies
No dependencies set.
Reference
deepseek-ai/DeepEP#705
Loading…
Add table
Add a link
Reference in a new issue
No description provided.
Delete branch "%!s()"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
📋 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)
3caa7edFix TypeError in cached ElasticBuffer dispatch with default num_sms3ef4840Run yapf and ruff262bcb2Update 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.dispatchreadsnum_expertsatdeep_ep/buffers/elastic.py:928butonly infers it from the handle 14 lines later, at line 942:
Both defaults make this reachable:
num_experts: Optional[int] = None(L860) andnum_sms: int = 0(L863). The docstring documents exactly this call —"
num_experts: ... Inferred fromhandleif provided" (L893) and "num_sms: ...0 for automatic" (L897) — so the documented cached call
enters L928 with
num_experts=None, giving:Callers that pass
num_smsexplicitly are unaffected, which is why the test suitenever hits it:
tests/elastic/test_ep.py:65resolvesnum_smsitself up front andpasses it into every dispatch call (L146, L165, L212, L223), so
num_smsis never 0on 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(onlydeep_ep._Cstubbed) and calling the realget_theoretical_num_smswith thenum_expertsthat L928 passes it in cached mode: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
combinealready does for the same value at
elastic.py:1086:EPHandlestoresnum_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_expertsis populated before L928. I did not take it: the handle'schannel metadata (
channel_linked_list,dst_buffer_slot_idx, the psum tensors) waslaid 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_smskeeps cached dispatch consistent with both the handle and withcombine, which must agree on channel geometry for the same handle.Scope
One line in
dispatch, plus a comment. Behavior changes only on the branch thatcurrently raises:
num_sms=0-> now useshandle.num_sms(was:TypeError)handle is None) -> unchanged, stillget_theoretical_num_sms(num_experts, ...)num_smsexplicitly -> unchangedTest plan
num_smsno longer raises; resolves tohandle.num_smsnum_smsstill resolves viaget_theoretical_num_smsnum_smsstill respected in both modestests/elastic/test_ep.pypassesnum_smsexplicitly, so it does not cover this branch either wayNote on
format-checkCI is red here on pre-existing lint in this file, not on this change.
format.shlints only files that differ from
origin/main, somainnever lintselastic.pyand these never surface there. Running the pinned tooling againstunmodified
origin/main(60d4403):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.