[almogtavor]

Purpose

Fixes #60685.

EPLB records load per physical expert. When step_interval < window_size, part of the window still holds load recorded under the old map, and the next rearrange counts it for the expert that now sits at that physical index.

This PR records load by logical expert instead. eplb_map_and_record already has the logical id, so the atomic_add now goes to the logical id rather than the physical index. expert_load_pass and the window have logical width, so old steps stay with the right expert after the map changes, and rearrange only sums the window. No extra kernel or scatter is added.

The PR also removes expert_load_pass_buffer, last_expert_load and use_last_expert_load. The number of logical experts does not change when Elastic EP scales, so the window stays valid across a scale and the scale rearrange uses it directly.

log_balancedness needs per rank load, so the same kernel also records physical counts into a new physical_load_pass when log_balancedness is on. The window never reads these counts, and they are cleared in the same step as the logical pass, so each logged value comes from one forward under one map. With log_balancedness off the buffer is not allocated and Triton compiles the second atomic_add out. The MegaMoE kernels share this helper, so they get the same recording.

No open PR fixes this. #49700 only clears the window on an elastic EP topology change, and #64485 changes when rearrangements run, not how the window is recorded.

Test Plan

All runs are on H100, on main at 24a49203ee with this PR on top. The native build is the nightly wheel ef70c6a23f, 4 commits before that, with the Python files replaced. The only native difference in those commits is vllm_flash_attn.cmake.

• pytest tests/kernels/moe/test_routing.py tests/distributed/test_eplb_utils.py
• pytest tests/kernels/moe/test_deep_gemm_mega_moe.py -k test_prepare_megamoe_eplb_maps_and_records_under_graph_replay, with @requires_sm100 removed locally because this test only runs the Triton staging kernels.
• Atomic contention microbench of the record kernel, main vs this PR with logging off and on, CUDA graph of 100 calls, median of 20, top_k 8, extra replicas on the 8 hottest experts, 2 runs per side. "hot" sends half of all picks to one expert.
• gsm8k on 2xH100: DeepSeek-V2-Lite-Chat, TP 2, EP + EPLB async, window_size=1000, step_interval=100, num_redundant_experts=2, log_balancedness=true, 5-shot, 2 runs per side.
• pytest tests/distributed/test_elastic_ep.py -k "enforce_eager_heavy or cuda_graphs_heavy" on 4xH100 with NIXL 1.5.2, main vs this PR with logging off and on. For logging on I added log_balancedness to the server args locally. This scales dp 2 to 4 and back under heavy traffic and checks gsm8k after each scale.
• Temporary check, not part of this PR: on every log_balancedness sync, fold the physical counts through the current map and compare them with the logical pass.

Test Result

test_routing.py + test_eplb_utils.py: 518 passed. Running the same tests on main gives 10 failed, which are the 8 changed routing cases, test_eplb_map_records_exact_physical_load and test_log_balancedness_from_physical_load. The MegaMoE test: 8 passed, with logging off and on.

The physical counts folded through the map matched the logical pass in every logged check: 364 in gsm8k and 61 in Elastic EP, including at physical widths 64 (dp 2) and 128 (dp 4). No check reported a mismatch.

Record kernel time per call (us):

model	dist	tokens	main	logging off	logging on
256+32 experts	uniform	1024	4.37	4.38	6.00
256+32 experts	uniform	16384	35.23	36.27	50.83
256+32 experts	zipf	16384	37.80	38.50	58.59
256+32 experts	hot	1024	4.30	4.84	6.36
256+32 experts	hot	16384	34.11	43.66	58.89
128+32 experts	uniform	16384	40.84	27.21	59.09
128+32 experts	hot	16384	40.63	34.06	64.14

With logging off, the only clear slowdown is one very hot expert out of 256, up to about 10 us per layer at 16k tokens and under 1 us at 1k tokens. The 128+32 case got faster. Logging on adds 1.4 to 2.5 us per layer up to 1k tokens and 15 to 32 us at 16k tokens, and only when log_balancedness is enabled. The two runs of main differed by at most 1.1%.

gsm8k accuracy with log_balancedness on:

	run 1	run 2
main	0.6315	0.6293
this PR	0.6391	0.6308

Every run did 3 to 6 async rearranges with no errors. Run 2 of this PR also had the temporary check and log_balancedness_interval=10, while the other runs used 50.

Elastic EP scaling: 3 passed on main, 3 passed on this PR with logging off, 3 passed on this PR with logging on. gsm8k on 256 questions:

case		initial	after 2 to 4
enforce_eager_heavy	main	0.668	0.684
enforce_eager_heavy	logging off	0.684	0.672
enforce_eager_heavy	logging on	0.652	0.668
cuda_graphs_heavy	main	0.656	0.656
cuda_graphs_heavy	logging off	0.684	0.641
cuda_graphs_heavy	logging on	0.668	0.645
cuda_graphs_heavy_nixl_ep	main	0.625	0.641
cuda_graphs_heavy_nixl_ep	logging off	0.656	0.660
cuda_graphs_heavy_nixl_ep	logging on	0.637	0.652

The test also checks accuracy after scaling back to 2, and that check passed in every case.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[itayalroy]

This is a real issue with our current EPLB implementation (I mentioned it in #58473’s description as well), but I see two problems with this fix:

1. This is dangerous for async EPLB. Async EPLB transfers one layer at a time, so _commit_eplb_maps_for_layer() for the last layers might happen very close to the next EPLB reshuffle, depending on transfer time and window/step configs. That reshuffle could use very little data for those layers, or even no data if the last layer is committed in the same step that triggers a reshuffle. This can result in really bad mappings which severely hurt e2e performance.

2. For sync EPLB, this change effectively silently sets window_size to step_interval whenever window_size > step_interval, which is pretty misleading

I spent some time thinking about this while fixing #58473. I think the safest solution is to record load by logical expert, so changing the physical mapping does not invalidate the history. But this is not an easy change as we need to make sure we do it without slowing down statistics collection.

[almogtavor]

Thanks @itayalroy it makes sense. What do you think about doing the logical conversion once per step instead of in the forward pass?

I mean leaving the forward pass as is so it keeps counting per physical slot in expert_load_pass. Then, in step() when the expert_load_pass is copied into expert_load_window we can applyscatter_add_ by the logical ID using the physical_to_logical_map so the window will hold load per logical expert. The scatter index is cached and rebuilt only when the map changes.

I think this is safer than counting by logical ID in the forward pass, since recording code stays untouched, including the code that lives in _eplb_map_and_record_i32_kernel in base_router.py and also in model specific paths like deepseek_v4 and kimi_k3 that would otherwise need to change. log_balancedness also needs the per-physical counts to compute load per rank, and this keeps them. The extra work runs once per step for all layers together and not once per layer per forward. It solves both of your points because the copy happens before any commit. Nothing has to be cleared when the map changes.

I measured it on 4xH100 with DeepSeek-V2-Lite-Chat and EP with rearranges every 200 steps. The per-step copy went from about 9.2us to 10.5us, rearrange() got about 2x faster because the window is already logical, and E2E throughput was 80.3 req/s on main vs 81.3 with this change over 6 runs each. Already updated the PR with this, but happy to change it if you prefer a different approach

[itayalroy]

Thanks @itayalroy it makes sense. What do you think about doing the logical conversion once per step instead of in the forward pass?

I mean leaving the forward pass as is so it keeps counting per physical slot in expert_load_pass. Then, in step() when the expert_load_pass is copied into expert_load_window we can applyscatter_add_ by the logical ID using the physical_to_logical_map so the window will hold load per logical expert. The scatter index is cached and rebuilt only when the map changes.

I think this is safer than counting by logical ID in the forward pass, since recording code stays untouched, including the code that lives in _eplb_map_and_record_i32_kernel in base_router.py and also in model specific paths like deepseek_v4 and kimi_k3 that would otherwise need to change. log_balancedness also needs the per-physical counts to compute load per rank, and this keeps them. The extra work runs once per step for all layers together and not once per layer per forward. It solves both of your points because the copy happens before any commit. Nothing has to be cleared when the map changes.

I measured it on 4xH100 with DeepSeek-V2-Lite-Chat and EP with rearranges every 200 steps. The per-step copy went from about 9.2us to 10.5us, rearrange() got about 2x faster because the window is already logical, and E2E throughput was 80.3 req/s on main vs 81.3 with this change over 6 runs each. Already updated the PR with this, but happy to change it if you prefer a different approach

I benchmarked this PR with DeepSeek V4 Pro on 20 GB200 GPUs. Compared with upstream, the median gap between forward passes (measured with CUDA events immediately after one forward call and before the next) increased by about 11 µs, and the EPLB step itself increased by the same amount. No measurable effect on e2e performance though.

I suspect the extra kernel launch accounts for much of that 11 µs. A fused Triton kernel would probably help here (it could even improve performance over upstream as we can fuse three kernels into one, compared with two kernel launches in upstream), but it would add code complexity, so not sure if it is worth it.

@tlrmchlsmth @ilmarkov for vis

[ilmarkov]

Thank you for the fix, it does make sense!

Regarding adding extra torch calls on the hot path.
Could we record logical-expert load in the existing _eplb_map_and_record_i32_kernel? It already has the logical expert_id before selecting physical_id, so this could let step() copy the counts into the window and avoid the new per-step zero_().scatter_add_ and cached index.
Physical counts are still needed for optional log_balancedness; we could keep physical recording for that mode. The tradeoff to benchmark is atomic contention when several replicas of a popular logical expert share one counter.

[ilmarkov]

Also, could you please add accuracy validation results (e.g. gsm8k)?

[itayalroy]

Also note that if we record expert load by logical expert, we don't need this EEP w/a anymore, we can remove it and clean up some code

[almogtavor]

@ilmarkov yes I'll do that even though I don't see any signal for accuracy being affected. I'll soon look into your suggestion as well thanks

[almogtavor]

I moved the recording into _eplb_map_and_record_i32_kernel so it counts by logical id, and I dropped the step() scatter. This also removes the extra launch from the forward gap.

Results on H100:

• gsm8k on DeepSeek-V2-Lite, TP 2, EP + async EPLB, 2 runs per side: main 0.6368 / 0.6323, PR 0.6399 / 0.6384. Each run did 19 rearranges with no errors.
• Atomic contention on the record kernel: uniform and zipf routing are within noise. With one very hot expert out of 256 it costs up to about 10 us per layer at 16k tokens and under 1 us at 1k tokens.
• test_elastic_ep.py (eager, CUDA graphs and nixl_ep, heavy traffic) passes on 4xH100. The scale rearrange now uses the logical window instead of last_expert_load.

The full tables are in the description.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @almogtavor.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ilmarkov]

We should use actual physical-expert counts for log_balancedness. We should record both logical and physical expert load in the existing routing kernel when log_balancedness is enabled.

Logical counts keep the rearrangement window valid across map changes, physical counts let us report actual per-rank balance. Splitting logical load evenly among replicas can hide a skewed replica and make balancedness look better than it is.

[almogtavor]

I changed this. The routing kernel now records both logical and physical counts when log_balancedness is on, and the logged balance comes from the physical counts. The window still reads only the logical counts. With logging off, the physical buffer isn't allocated and the second add compiles out. The test results are in the PR description.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @almogtavor.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @almogtavor.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ilmarkov]

Thank you for the update. Looks good. Could you add e2e throughput and latency comparison to the main? 10us overhead per layer in case of hot experts sounds not good. We need to check that e2e performance in normal conditions is not severely affected.

Also, is there a test verifying that the initial issue is solved? Test that checks that a partially refreshed window retains the right logical load after a map change.

[ilmarkov]

Nit: I think it's 2D now

[ilmarkov]

Nit: For log_balancedness it can be num_physical_experts too

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @almogtavor.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork