[EPLB] Optmize eplb mapping and record in router for prefill (#36261)

Signed-off-by: ilmarkov <markovilya197@gmail.com>
2026-03-30 21:48:33 +02:00
parent 494636b29d
commit 12701e8af2
6 changed files with 338 additions and 66 deletions
--- a/tests/model_executor/test_routed_experts_capture.py
+++ b/tests/model_executor/test_routed_experts_capture.py
@@ -62,6 +62,7 @@ def test_base_router_capture_with_eplb_enabled():
    router.eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
    router.eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
    router.eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
+    router.eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)

    captured = []