[bvolpato]

Summary

The topk_weights parameter documentation says weights must be None when a cached handle is supplied, but the handle documentation explicitly allows them. The implementation and existing EP test also support optional weights with a cached handle.

Changes

Correct the parameter docstring to match the existing behavior. This changes documentation only.

Testing

git diff --check passes. Checked Python syntax, confirmed executable code is unchanged, and validated cached weight handling with a runtime stub. GPU tests were not run.

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

该变更正确转发了缓存分发所需的 top-k 权重，并返回接收端实际生成的索引与权重，符合 ElasticBuffer.dispatch 的接口和现有测试约定。

v5

该 MR（提交 ab9b5da）为纯文档修复，仅改动 README.md 中推理解码示例 decode_dispatch 的缓存分支：修复前示例丢弃了 ElasticBuffer.dispatch 的接收端路由输出，错误地返回 cached_handle.topk_idx（发送端全局 expert 选择）且权重恒为 None；修复后转发当前 topk_weights 并返回真实的 recv_topk_idx / recv_topk_weights。已核对：(1) deep_ep/buffers/elastic.py 中 dispatch 的签名与文档表明缓存 handle 路径下允许可选传入 topk_weights，返回值即接收端的 recv_topk_idx / recv_topk_weights，修复与运行时契约一致；(2) tests/elastic/test_ep.py 的缓存 dispatch 测试同样以接收端输出（cached_recv_topk_idx）和缓存权重转发为契约；(3) 多 rank 场景下发送端 token 数/全局 expert ID 与接收端 token 数/本地 expert ID 可能不同，旧示例会误导用户用错误的路由元数据做后续专家计算，且返回 None 权重与示例函数声明的返回类型 torch.Tensor 不符，此修复均予以纠正。无运行时代码改动，风险极低，建议合并。

v4p

本 MR 修复 README 中缓存解码 dispatch 示例：在 cached handle 分支传入当前 topk_weights，并返回 dispatch 实际输出的 recv_topk_idx 和 recv_topk_weights，而不是缓存的 sender 端 topk_idx 和 None。变更与 ElasticBuffer.dispatch 的接口契约一致，属于最小且正确的文档修复，无运行时代码影响。

Files reviewed: 1
Issues found: 🔵 1 suggestion
Inline comments posted: 0
General comments (无法定位到 diff): 1

---

📍 未定位到 diff 的评论

🔵 suggestion deep_ep/buffers/elastic.py:L890: 非本 MR 范围的小建议：dispatch 文档字符串中 topk_weights 参数处写着 'Must be None if handle is provided'，而 handle 参数处又说 'topk_weights can be optionally provided (e.g. for backward pass with cached expand)'，两处表述互相矛盾（实际行为是后者，README 新示例也依赖该行为）。建议在后续 docs 提交中把 topk_weights 的描述统一为“提供 handle 时可选传入”，避免读者困惑。 _{🤖 v5}

[bvolpato]

@sphish, would you mind taking a look at this? It fixes the cached ElasticBuffer.dispatch README example to match the receiver-side outputs and top-k weight handling in the elastic path you worked on. If someone else owns this path now, happy to redirect.

[bvolpato]

V2.5 removed the original README example, so I updated this PR to fix the cached-weight docstring contradiction that still exists in EPBuffer.dispatch, as suggested in the earlier review.

I executed the current Python wrapper with a recording runtime stub: cached dispatch accepts optional weights and preserves them through ordinary/expanded layouts and immediate/deferred epilogues. The C++ implementation and existing cached-expanded test agree. Pinned formatting and lint checks pass; the required file cleanup is in a separate commit. GPU kernels were not run.

Would someone familiar with the EP dispatch path take a look?

[bvolpato]

@LyricZhao, would you mind taking a look? This clarifies the optional-weight behavior of cached EPBuffer.dispatch calls in V2.5. If someone else owns this API now, happy to redirect.