docs: fix cached dispatch outputs

[bvolpato]

Summary

Cached decoding example discarded receiver-side routing outputs from ElasticBuffer.dispatch. It returned cached_handle.topk_idx, which stores sender-side global expert selections, and always returned None for weights.

Changes

• Forward current topk_weights during cached dispatch.
• Return actual recv_topk_idx and recv_topk_weights outputs.

Repro

Executed exact README code block against a contract-faithful fake buffer that returns distinct sender and receiver metadata.

Before this change:

AssertionError: returned source top-k indices instead of received top-k indices; did not forward and return top-k weights

After this change:

README cached decode example preserves dispatch outputs

This distinction matters in multi-rank use because sender token count and global expert IDs can differ from receiver token count and receiver-local expert IDs.

Testing

• uv run --isolated --with yapf==0.40.2 --with ruff==0.6.5 --with clang-format==15.0.7 -- bash ./format.sh
• Executed exact README cached path with distinct source and received metadata before and after fix.
• Verified existing cached-dispatch test contract compares cached_recv_topk_idx with receiver output and exercises cached weight forwarding.

GPU suite not run. Change only corrects README example; no runtime code changed.

Did this cause any problems?

Docs-only change. Revert this commit to restore previous example.

[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]

Thanks for catching this. Agreed, the topk_weights docstring is contradictory. I will keep this PR focused on the README fix and address the docstring in a follow-up.