[Feature] Integrate Elastic NIXL-EP into SGLang (#19248)

Signed-off-by: Barak Biber <bbiber@nvidia.com>
Signed-off-by: Yoray Zack <yorayz@nvidia.com>
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Co-authored-by: Barak Biber <bbiber@nvidia.com>

python/sglang/srt/eplb/eplb_algorithms/__init__.py

@@ -13,6 +13,7 @@ class EplbAlgorithm(Enum):
     deepseek_vec = auto()
     deepseek_vec_hierarchical = auto()
     elasticity_aware = auto()
+    elasticity_aware_hierarchical = auto()
     # TODO may have more algorithm later
 
 
@@ -47,14 +48,19 @@ def rebalance_experts(
             enable_hierarchical=algorithm == EplbAlgorithm.deepseek_vec_hierarchical,
         )
 
-    if algorithm == EplbAlgorithm.elasticity_aware:
+    if algorithm in [
+        EplbAlgorithm.elasticity_aware,
+        EplbAlgorithm.elasticity_aware_hierarchical,
+    ]:
         return elasticity_aware.rebalance_experts(
             weight=tokens_per_expert.sum(dim=0),
             num_replicas=num_physical_experts,
             num_groups=num_groups,
             num_nodes=num_nodes,
             num_gpus=num_physical_experts // num_local_physical_experts,
-            enable_hierarchical=False,
+            enable_hierarchical=(
+                algorithm == EplbAlgorithm.elasticity_aware_hierarchical
+            ),
             active_ranks=(
                 ElasticEPStateManager.instance().active_ranks
                 if ElasticEPStateManager.instance() is not None