[EPLB] Simplify EPLB rearrange by only returning one map

[SageMoore]

Purpose

On main EPLB rearrange returns 3 maps: physical expert -> logical expert, logical expert -> physical expert, and logical expert -> replica count. The latter two are derivable from the physical expert -> logical expert mapping.

With this change, we only have to keep track of the physical expert -> logical expert map in when running EPLB. Once weight transfer has finished and EPLB is ready to make the changes visible to the model, it will generate the other two maps and write them directly into the model state.

This change deletes a lot of code and reduces the amount of state that EPLB has to maintain when running async. Functionally the only change is "delaying" the creation of the second two maps until after weight transfer has completed when running async. The rest is purely cosmetic.

It also opens the door for post-processing the result of EPLB rearrange without having to re-update all of the maps.

Test Plan

Sync EPLB Server Command

vllm serve deepseek-ai/DeepSeek-V2-Lite -dp 2 --enable-expert-parallel --all2all_backend deepep_low_latency --enable-eplb --eplb-config '{"window_size":20,"step_interval":20}'

Async EPLB Server Command

vllm serve deepseek-ai/DeepSeek-V2-Lite -dp 2 --enable-expert-parallel --all2all_backend deepep_low_latency --enable-eplb --eplb-config '{"window_size":20,"step_interval":20,"use_async":"true"}'

Test Result

Sync EPLB

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.3867|±  |0.0282|
|     |       |strict-match    |     5|exact_match|↑  |0.3833|±  |0.0281|

Async EPLB

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.3733|±  | 0.028|
|     |       |strict-match    |     5|exact_match|↑  |0.3733|±  | 0.028|

[gemini-code-assist]

Code Review

This pull request refactors the EPLB rearrange logic to simplify state management by deriving two expert mapping tensors from a single one, instead of computing and storing all three. This is a good simplification that reduces complexity. The overall changes look correct and align with the stated purpose. However, I've found a couple of critical issues in the new compute_logical_maps helper function where it doesn't correctly handle negative indices for unused experts, which could lead to incorrect behavior or crashes. I've provided detailed comments and suggestions to address these bugs.

_{Note: Security Review is unavailable for this PR.}

[ilmarkov]

Makes sense. Thank you!