Skip to content

fix: derive ep_rank from inner_rank / mlp_tp_size instead of inner_rank % ep_size#4756

Closed
irexyc with Copilot wants to merge 2 commits into
mainfrom
copilot/fix-code-review-comment
Closed

fix: derive ep_rank from inner_rank / mlp_tp_size instead of inner_rank % ep_size#4756
irexyc with Copilot wants to merge 2 commits into
mainfrom
copilot/fix-code-review-comment

Conversation

Copilot AI commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

Motivation

When both mlp_tp_size > 1 and ep_size > 1, computing ep_rank = inner_rank % ep_size collides with mlp_tp_rank = inner_rank % mlp_tp_size — both consume the same modulo space, producing duplicated/incorrect EP ranks. The invariant inner_rank == ep_rank * mlp_tp_size + mlp_tp_rank must hold.

Modification

  • src/turbomind/turbomind.cc: Fix ep_rank derivation to inner_rank / mlp_tp_size; fix mlp_tp_rank to inner_rank % mlp_tp_size; restore assertion to mlp_tp_size * ep_size == comm_size_; add GetEpRank() implementation
  • src/turbomind/turbomind.h: Declare GetEpRank(int index)
  • src/turbomind/models/llama/llama_params.h: Add int ep_rank = 0 to EngineParam
  • src/turbomind/engine/engine_config.h: Add ep_size (default 1) to EngineConfig
  • src/turbomind/python/bind.cpp: Expose ep_rank via Python binding

The correct decomposition:

inner_rank = ep_rank * mlp_tp_size + mlp_tp_rank
           ⟹  mlp_tp_rank = inner_rank % mlp_tp_size
               ep_rank     = inner_rank / mlp_tp_size

BC-breaking (Optional)

No BC-breaking changes. ep_size defaults to 1, preserving existing single-EP behavior.

Use cases (Optional)

N/A

Checklist

  1. Pre-commit or other linting tools are used to fix the potential lint issues.
  2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
  3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
  4. The documentation has been modified accordingly, like docstring or example tutorials.

Copilot AI changed the title [WIP] Fix code based on review comment fix: derive ep_rank from inner_rank / mlp_tp_size instead of inner_rank % ep_size Jul 16, 2026
Copilot AI requested a review from irexyc July 16, 2026 09:45
@irexyc irexyc closed this Jul 16, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants