add more comments for force load balance

Signed-off-by: angazenn <zengyanjia@huawei.com>

Author: angazenn

vllm_ascend/models/deepseek_v2.py

@@ -25,7 +25,6 @@
 # # vllm-project/vllm/vllm/model_executor/models/deepseek_v2.py
 # """Inference-only DeepseekV2/DeepseekV3 model."""
 
-import os
 from typing import Any, Dict, List, Optional, Union
 
 import torch
@@ -213,8 +212,8 @@ def __init__(
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         attn_metadata = get_forward_context().attn_metadata
-        # when profile runs, force experts load balance to avoid high memory
-        # consumption from 1 rank.
+        # when profile runs, force experts to load balanced tokens
+        # to avoid high memory consumption on a single rank.
         # TODO: need a better flag to indicate whether in profile run or not.
         if attn_metadata is None or attn_metadata.slot_mapping[-1] < 0:
             # for profile run

vllm_ascend/ops/fused_moe.py

@@ -15,7 +15,6 @@
 # This file is a part of the vllm-ascend project.
 # Adapted from vllm/tests/kernels/test_moe.py
 
-import os
 from typing import Callable, Optional
 
 import torch

vllm_ascend/quantization/w8a8_dynamic.py

@@ -15,16 +15,15 @@
 # limitations under the License.
 #
 
-import os
 from typing import Any, Callable, Dict, List, Optional
 
 import torch
 import torch.distributed as dist
 import torch_npu
 from vllm.distributed import GroupCoordinator
 
-from vllm_ascend.distributed.parallel_state import get_ep_group
 import vllm_ascend.envs as envs_ascend
+from vllm_ascend.distributed.parallel_state import get_ep_group
 from vllm_ascend.ops.fused_moe import select_experts
 
 VLLM_ENABLE_MC2: bool = envs_ascend.VLLM_ENABLE_MC2
@@ -200,6 +199,8 @@ def fused_experts_with_mc2(
     return hidden_states
 
 
+# currently expert parallelism implemented with all2all
+# is under-optimized.
 def fused_experts_with_all2all(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
@@ -616,6 +617,9 @@ def apply(
                 e_score_correction_bias=e_score_correction_bias,
             )
 
+        # this is a naive implementation for experts load balance so as
+        # to avoid accumulating too much tokens on a single rank.
+        # currently it is only activated when doing profile runs.
         if enable_force_load_balance:
             topk_ids = torch.randint_like(topk_ids, 0, global_num_experts)