support eplb for qwen3 (#6533)

2025-05-24 09:31:30 +08:00
parent 7b02c32679
commit e6f113569e
3 changed files with 46 additions and 25 deletions
--- a/python/sglang/srt/managers/expert_distribution.py
+++ b/python/sglang/srt/managers/expert_distribution.py
@@ -690,7 +690,9 @@ def _convert_global_physical_count_to_logical_count(
    )
    logical_count.scatter_add_(
        dim=2,
-        index=physical_to_logical_map.unsqueeze(0).expand(dim_extra, -1, -1),
+        index=physical_to_logical_map.unsqueeze(0)
+        .expand(dim_extra, -1, -1)
+        .to(torch.int64),
        src=global_physical_count,
    )
    return logical_count