Merge branch 'v0.9.2-dev-ep' into 'v0.9.2-dev'

feat: moe_align_block_size 更新lightop 接口,加入对ep的支持

See merge request dcutoolkit/deeplearing/vllm!216

vllm/model_executor/layers/fused_moe/moe_align_block_size.py

@@ -151,7 +151,8 @@ def moe_align_block_size(
     num_experts: int,
     expert_map: Optional[torch.Tensor] = None,
     pad_sorted_ids: bool = False,
-    num_token: Optional[int] = None
+    num_token: Optional[int] = None,
+    expert_mask: Optional[torch.Tensor] = None
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     Aligns the token distribution across experts to be compatible with block
@@ -231,14 +232,17 @@ def moe_align_block_size(
                                       dtype=torch.int32,
                                       device=topk_ids.device)
 
-    if envs.VLLM_USE_LIGHTOP:
+    if envs.VLLM_USE_LIGHTOP or expert_mask is not None:
         from lightop import op as op
         op.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
-                                expert_ids, num_tokens_post_pad, None, None, None)
+                                expert_ids, num_tokens_post_pad,
+                                expert_map = expert_map,
+                                expert_mask = expert_mask,
+                                num_local_tokens = None)
     else:
         ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
-                                expert_ids, num_tokens_post_pad)
-    if expert_map is not None:
-        expert_ids = expert_map[expert_ids]
+                                 expert_ids, num_tokens_post_pad)
+        if expert_map is not None:
+            expert_ids = expert_map[expert_ids]
 
     return sorted_ids, expert_ids, num_tokens_post_pad
\ No newline at end of file