[BugFix]Revise the EPLB feature guide content and optimize the EPLB algorithm.

docs/source/user_guide/configuration/additional_config.md

@@ -36,6 +36,12 @@ The following table lists the additional configuration options available in vLLM
 | `lmhead_tensor_parallel_size` | int | `None` | The custom tensor parallel size of lmhead. |
 | `oproj_tensor_parallel_size` | int | `None` | The custom tensor parallel size of oproj. |
 | `multistream_overlap_shared_expert`| bool | `False` | Whether to enable multistream shared expert. This option only takes effects on moe models with shared experts. |
+| `dynamic_eplb`              | bool | `False` | Whether to enable dynamic eplb                                                                                                                |
+|`num_iterations_eplb_update`| int  | `400` | Forward iterations when eplb would begin                                                                                                      |
+|`gate_eplb`| bool | `False` | Whether to enale eplb only once.                                                                                                              |
+|`num_wait_worker_iterations`| int  | `30` | The  forward iterations when eplb worker will finish cpu task. In our test default value 30 would cover most cases.                           |
+|`expert_map_record_path`| str  | `None` |              When dynamic eplb is completed, save the current expert load heatmap to the specified path.                                                                                                                                 |
+|`init_redundancy_expert`| int  | `0`  |Specify redundant experts during initialization.|
 
 The details of each config option are as follows:
 

docs/source/user_guide/feature_guide/eplb_swift_balancer.md

@@ -1,7 +1,7 @@
 # Swift Balancer
 
 ## Overview
-Experts rebalancing of MoE models for LLM serving is a mandatory option.Changing experts dynamically would have a negative impact on TTFT and TPOT while stop-the-world.
+Experts balancing of MoE models for LLM serving is a mandatory option.Changing experts dynamically would have a negative impact on TTFT and TPOT while stop-the-world.
 Asynchronously expert load balancing would be a better choice.
 We have launched SwiftBalancer to support dynamic experts load balancing with Zero-overhead experts movement.
 
@@ -15,31 +15,34 @@ The overall workflow involves:
 2. Do all-gather for experts distribution. Using all-gather instead of all-reduce as less traffic volume.
 3. Wake up eplb worker process with experts distribution when num_iterations comes. Run eplb algorithm in eplb worker.
 4. Generate p2p send/recv ops and other operator such as log2phy would cost long cpu time.
-5. Lanch ibatch_send_recv in async_stream before forward.
-6. After forward, wait for the ibatch_send_recv finish, then do uapte expert map and expert weights.
+5. Launch ibatch_send_recv in async_stream before forward.
+6. After forward, wait for the ibatch_send_recv finish, then do update expert map and expert weights.
 
 In our profiling shows experts transforming is hidden in the bubble between forward iterations. Cpu time cost of eplb algo. and other python operator such as log2phy
 would be hidden by eplb worker process too.
 
-## Config Params
-
-Currently swift balancer optimize 5ms TPOT with ep size 64 while cost less than 2ms for every layer expert movement.
-
-We add new parameters for eplb:
-"dynamic_eplb":true ---  enable dynamic eplb
-"num_iterations_eplb_update": 400 -- forward iterations when eplb would begin
-"gate_eplb":true -- eplb would update only once, false by default.
-"num_wait_worker_iterations":30 -- forward iterations when eplb worker will finish cpu task. In our test default value 30 would cover most cases.
-"expert_map_record_path" -- When dynamic eplb is completed, save the current expert load heatmap to the specified path.
-"init_redundancy_expert" -- Specify redundant experts during initialization.
 
 ## Examples
 ### Dynamic eplb
 Enable dynamic eplb and specify the trigger rounds.
---additional-config '{ "dynamic_eplb":true,"num_iterations_eplb_update":400, "gate_eplb":true, "num_wait_worker_iterations":30}'
-### Record expert map for static eplb
-Specify the path for the static eplb initialization file.
---additional-config '{ "expert_map_record_path": "/xx/xx.json", "init_redundancy_expert": 16, dynamic_eplb":true,"num_iterations_eplb_update":400, "gate_eplb":true, "num_wait_worker_iterations":30}'
+```shell
+vllm serve Qwen/Qwen3-235B-A22 \
+  --tensor-parallel-size 16 \
+  --enable-expert-parallel \
+  --additional-config '{ "dynamic_eplb":true,"num_iterations_eplb_update":400, "gate_eplb":true, "num_wait_worker_iterations":30}'
+```
 ### Static eplb
-If expert map has been recorded, enable static eplb with expert map path.
---additional-config '{ "expert_map_path": "/xx/xx.json"}'
+1. Specify the path for the static eplb initialization file.
+```shell
+vllm serve Qwen/Qwen3-235B-A22 \
+  --tensor-parallel-size 16 \
+  --enable-expert-parallel \
+  --additional-config '{ "expert_map_record_path": "/path/to/eplb.json", "init_redundancy_expert": 16, dynamic_eplb":true,"num_iterations_eplb_update":400, "gate_eplb":true, "num_wait_worker_iterations":30}'
+```
+2. If expert map has been recorded, enable static eplb with expert map path.
+```shell
+vllm serve Qwen/Qwen3-235B-A22 \
+  --tensor-parallel-size 16 \
+  --enable-expert-parallel \
+  --additional-config '{ "expert_map_path": "/path/to/eplb.json"}'
+```

vllm_ascend/eplb/core/eplb_utils.py

@@ -15,10 +15,7 @@
 # This file is a part of the vllm-ascend project.
 #
 # Todo: Once https://github.com/vllm-project/vllm/issues/22246 is merged in vllm. Remove eplb utils.
-import random
-
 import torch
-from vllm.logger import logger
 
 
 def determine_default_expert_map(global_expert_num, world_size, rank_id,
@@ -56,40 +53,59 @@ def determine_default_expert_map(global_expert_num, world_size, rank_id,
 
 
 def generate_log2phy_map(expert_map):
+    """
+    Generate a log-to-physical map for experts in a fully vectorized manner.
+
+    Args:
+        expert_map: Tensor of shape [num_ranks, num_global_expert], with -1 indicating
+                    rank does not hold the expert.
+
+    Returns:
+        log2phy_map: Tensor of same shape, mapping logical experts to physical IDs.
+    """
+    num_ranks, _ = expert_map.shape
     num_local_experts = expert_map.max() + 1
+    device = expert_map.device
+
+    # Step 1: linear mapping based on rank
     log2phy_map = expert_map.clone()
-    num_ranks, num_global_expert = log2phy_map.shape
-
-    row_indices = torch.arange(num_ranks).view(-1, 1).expand(num_ranks, \
-                                                             num_global_expert) * num_local_experts
-    log2phy_map[log2phy_map != -1] += row_indices[log2phy_map != -1]
-
-    for idx in range(num_global_expert):
-        positive_rank_idx = torch.where(log2phy_map[:, idx] != -1)[0]
-        negative_rank_idx = torch.where(log2phy_map[:, idx] == -1)[0]
-        num_rank_holding_expert = positive_rank_idx.size(0)
-
-        if num_rank_holding_expert == 0:
-            log2phy_map[:, idx] = torch.full((num_ranks, ),
-                                             0,
-                                             dtype=log2phy_map.dtype)
-
-        if num_rank_holding_expert == 1:
-            log2phy_map[negative_rank_idx, idx] = torch.full(
-                (num_ranks - 1, ),
-                log2phy_map[positive_rank_idx, idx].item(),
-                dtype=log2phy_map.dtype)
-        else:
-            try:
-                random_list = [
-                    random.choice(log2phy_map[positive_rank_idx, idx])
-                    for _ in range(num_ranks - num_rank_holding_expert)
-                ]
-                log2phy_map[negative_rank_idx,
-                            idx] = torch.tensor(random_list,
-                                                dtype=log2phy_map.dtype)
-            except Exception as e:
-                logger.error(f"Fail to get log2phy_map: {str(e)}")
+    row_indices = torch.arange(num_ranks, device=device).view(
+        -1, 1) * num_local_experts
+    mask = log2phy_map != -1
+    # broadcast addition
+    log2phy_map = log2phy_map + row_indices * mask.long()
+
+    # Step 2: find positive/negative positions
+    positive_mask = log2phy_map != -1
+    negative_mask = ~positive_mask
+
+    # Count number of ranks holding each global expert
+    num_positive_per_col = positive_mask.sum(dim=0)  # [num_global_expert]
+
+    # Step 3: handle columns with only one rank holding the expert
+    single_pos_mask = num_positive_per_col == 1
+    if single_pos_mask.any():
+        # get row indices for the positive element in these columns
+        # pos_idx = torch.nonzero(positive_mask[:, single_pos_mask], as_tuple=True)
+        # broadcast to fill negative positions
+        for col_idx in torch.nonzero(single_pos_mask, as_tuple=True)[0]:
+            pos_row = torch.nonzero(positive_mask[:, col_idx])[0]
+            neg_rows = torch.nonzero(negative_mask[:, col_idx])[:, 0]
+            log2phy_map[neg_rows, col_idx] = log2phy_map[pos_row, col_idx]
+
+    # Step 4: handle columns with multiple ranks holding the expert
+    multi_pos_mask = num_positive_per_col > 1
+    if multi_pos_mask.any():
+        for col_idx in torch.nonzero(multi_pos_mask, as_tuple=True)[0]:
+            pos_rows = torch.nonzero(positive_mask[:, col_idx])[:, 0]
+            neg_rows = torch.nonzero(negative_mask[:, col_idx])[:, 0]
+            if len(neg_rows) > 0:
+                # random assignment from available positive ranks
+                rand_idx = torch.randint(0,
+                                         len(pos_rows), (len(neg_rows), ),
+                                         device=device)
+                log2phy_map[neg_rows,
+                            col_idx] = log2phy_map[pos_rows[rand_idx], col_idx]
 
     return log2phy_map
 

vllm_ascend/eplb/core/eplb_worker.py

@@ -243,6 +243,7 @@ def compose_expert_update_info_greedy(self, updated_expert_maps,
                 yield (expert_send_info_this_layer,
                        expert_recv_info_this_layer,
                        updated_expert_maps_this_layer, layer_id)
+                continue
 
             # Parse expert_ids each rank needs to receive from other ranks
             dst_rank_indices, experts_to_recv = torch.where((current_expert_maps_this_layer == -1) \

vllm_ascend/ops/common_fused_moe.py

@@ -346,7 +346,8 @@ def forward_impl(self, hidden_states: torch.Tensor,
             expert_load_view=self.expert_load_view,
             logical_to_physical_map=self.logical_to_physical_map,
             logical_replica_count=self.logical_replica_count,
-        )
+            log2phy=self.log2phy,
+            global_redundant_expert_num=self.global_redundant_expert_num)
         if isinstance(final_hidden_states, tuple):
             final_hidden_states, group_list_type, expert_tokens = final_hidden_states
 

vllm_ascend/quantization/w8a8_dynamic.py

@@ -230,7 +230,9 @@ def apply(
                 w1_scale=layer.w13_weight_scale,
                 w2_scale=layer.w2_weight_scale,
                 expert_map=expert_map,
-                dynamic_eplb=self.dynamic_eplb)
+                dynamic_eplb=self.dynamic_eplb,
+                log2phy=log2phy,
+                global_redundant_expert_num=global_redundant_expert_num)
 
         # this is a naive implementation for experts load balance so as
         # to avoid accumulating too much tokens on a single rank.