Fix TypeError in cached ElasticBuffer dispatch with default num_sms

[yurekami]

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.

[ds-review-bot]

🟡 warning: The public dispatch docstring still says num_sms: ... 0 for automatic via get_theoretical_num_sms, which is now inaccurate/inconsistent with combine (0 to reuse the SM count from the dispatch handle). It should state that a cached handle reuses handle.num_sms. Pure documentation; non-blocking.

_{🤖 v4}

[ds-review-bot]

🔵 suggestion: The nested ternary (handle.num_sms if handle is not None else get_theoretical_num_sms(num_experts, num_topk)) if num_sms == 0 else num_sms is hard to read and duplicates logic. Consider splitting into a clear few lines. Behaviour already correct; non-blocking.

_{🤖 v4}

[ds-review-bot]

🤖 ds-review-bot Code Review

v4

The candidate PR fixes the cached-dispatch TypeError (None % int) at elastic.py:928, where get_theoretical_num_sms(num_experts, num_topk) runs before num_experts is inferred from the handle. The applied one-liner reuses handle.num_sms for a cached dispatch (mirroring combine, which stores/reads the same value) and is guarded against a None handle so a fresh dispatch keeps its get_theoretical_num_sms(num_experts, num_topk) fallback. The patch is correct, minimal, scope-preserving (three behaviour cases in the issue hold, plus the pre-existing handle=None+num_experts=None user error is unchanged), and strictly safer than the issue's literal single-line proposal which would AttributeError on a fresh dispatch.

v3

The change fixes a TypeError in ElasticBuffer.dispatch where num_experts was read at line 928 (get_theoretical_num_sms) before being inferred from the handle 14 lines later. In cached-dispatch mode with the default num_sms=0 and num_experts=None, this raised 'TypeError: unsupported operand type(s) for %: NoneType and int'. The fix (line 930) reuses handle.num_sms when a handle is provided (as combine already does at line 1088 for the same value), instead of recomputing a theoretical SM count that could disagree with the channel geometry the handle was built under. All three behavioral branches are preserved correctly via the nested ternary: cached+num_sms=0 -> handle.num_sms; fresh dispatch (handle is None) -> get_theoretical_num_sms(num_experts, num_topk); explicit num_sms -> respected in both modes via the outer 'if num_sms == 0' guard. The fix is minimal (one line + comment), well-reasoned, and consistent with the existing combine() path. The analysis in the description is accurate and the design tradeoff (reuse handle.num_sms vs. hoisting the unpack block) is sound. Minor, non-blocking: the single-line nested ternary is fairly dense and could be split into an if/else block for readability. No correctness issues found.

Files reviewed: 1
Issues found: 🟡 1 warning | 🔵 1 suggestion
Inline comments posted: 2