Repository navigation
[Perf] [NVIDIA] eplb: Vectorize the static per-rank dispatch map - #48784
Open
weireweire wants to merge 2 commits into
Open
weireweire wants to merge 2 commits into
weireweire wants to merge 2 commits into
Conversation
Root cause: with ep_dispatch_algorithm=static, every rebalance rebuilds the per-rank dispatch map in compute_logical_to_rank_dispatch_physical_map, which calls _find_nearest_expert in Python for every (layer, logical expert, rank): 40 x 384 x 32 calls for DeepSeek-V4.1-Flash at DEP32, 250-290 ms on every rank per rebalance while the whole DP group waits. Fix: compact each row's valid replicas to the front (cumsum scatter, original order kept), then pick the same-GPU replica, else the narrowing same-node one, with tensor ops over the replica slots; rows with one replica broadcast. Only the (layer, expert, rank) entries left unresolved still call _fair_choices on the seeded Random, in the same row-major order, so the result is bitwise identical to the previous implementation. Validation (GB300): - Bitwise check against the previous function: 96 cases (12 layouts x 2 shapes x 4 ranks, scattered -1 padding, ep_dispatch prefer-same-node on and off), all equal. - Time per call: 40 x 384 logical / 416 physical, ep 32: 255 -> 39 ms; 3 x 128 / 160: 5.6 -> 9.1 ms. - DeepSeek-V4.1-Flash PD decode DEP32 (8 nodes), EPLB every 15000 iterations, AgentX 3600 s at concurrency 3018, same-day runs: rebalance pause p50 0.734 -> 0.467 s; total tput/GPU 458.7k -> 466.6k (+1.7%), interactivity p90 326.7 -> 339.8 tok/s (+4.0%), TTFT p95 9.33 -> 7.81 s.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Root cause: the vectorized dispatch-map helpers created their scratch tensors
with torch.zeros / torch.arange and no device, so they followed the ambient
default device. Called under a torch.device("cuda") context, they would mix
CUDA and CPU tensors in scatter_ / where. _logical_to_all_physical_raw lost its
last caller in the vectorized path.
Fix: create the scratch tensors on the device of the input map, and delete the
unused _logical_to_all_physical_raw.
Validation: no behaviour change on the current call paths, which run on CPU
with no device context. pre-commit passes. An independent equivalence check
of the vectorized map against the previous loop (17,430 random layouts x
ranks: ep 1-16, 1-8 nodes, scattered -1 padding, duplicate ids, int32/int64)
matched bitwise.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
With
ep_dispatch_algorithm=static, every EPLB rebalance rebuilds the per-rank dispatch map incompute_logical_to_rank_dispatch_physical_map. It calls_find_nearest_expertin Python once per (layer, logical expert, rank). For DeepSeek-V4.1-Flash at DEP32 that is 40 x 384 x 32 calls, which take 250-290 ms on every rank on each rebalance while the whole DP group waits.Modifications
_compact_candidates: move each row's valid physical ids to the front, keeping their original order._nearest_expert_dispatch: pick the same-GPU replica, otherwise the same-node replica when it narrows the candidate set, using tensor ops over the replica slots. Rows with a single replica broadcast it. This is the same rule as_find_nearest_expert._fair_choiceswith the seededrandom.Random, in the same row-major order. The output is therefore bitwise identical to the previous implementation.Accuracy
Placement is unchanged: a bitwise comparison against the previous function over 96 cases (12 layouts x 2 shapes x 4 ranks, scattered -1 padding, with and without the node topology) matched in every case.
Performance
Time per call:
The small shape gets a few ms slower from tensor-op overhead; the large one, where the time matters, is 6.5x faster.
End to end, DeepSeek-V4.1-Flash disaggregated serving on GB300 (DEP32 decode, 8 nodes), EPLB every 15000 iterations, AgentX 3600 s at concurrency 3018, same-day runs:
Related
#45282 (non-blocking EPLB rebalance) moves the rebalance off the forward thread but keeps calling this function, so the two are complementary; it adds a
deviceargument to the same function, which is a small conflict.CI States
Latest PR Test (Base): ❌ Run #42154060052
Latest PR Test (Extra): ❌ Run #42154059980
Latest PR Test (AMD ROCm 10): ❌ Run #42154060006
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: