[Attention][Kernel]moe support for llama4 and mllama4

Co-authored-by: evian <eviantai@u.nus.edu>
Signed-off-by: chenxu <chenxu68@huawei.com>

Author: chenxu

vllm_ascend/attention/attention.py

@@ -668,6 +668,7 @@ def __init__(
         blocksparse_params: Optional[Dict[str, Any]] = None,
         logits_soft_cap: Optional[float] = None,
         attn_type: str = AttentionType.DECODER,
+        use_irope: bool = False,
     ) -> None:
         self.num_heads = num_heads
         self.head_size = head_size

vllm_ascend/attention/attention_v1.py

@@ -172,6 +172,7 @@ def __init__(
         blocksparse_params: Optional[Dict[str, Any]] = None,
         logits_soft_cap: Optional[float] = None,
         attn_type: str = AttentionType.DECODER,
+        use_irope: bool = False,
     ) -> None:
         self.num_heads = num_heads
         self.head_size = head_size

vllm_ascend/ops/common_fused_moe.py

@@ -61,7 +61,8 @@ def forward_oot(
                          topk_weights=topk_weights,
                          topk_ids=topk_ids,
                          top_k=top_k,
-                         expert_map=expert_map)
+                         expert_map=expert_map,
+                         apply_router_weight_on_input=apply_router_weight_on_input)
 
 
 UnquantizedFusedMoEMethod.forward_oot = forward_oot

vllm_ascend/ops/fused_moe.py

@@ -150,6 +150,7 @@ def fused_experts(
     topk_ids: torch.Tensor,
     top_k: int,
     expert_map: torch.Tensor = None,
+    apply_router_weight_on_input: bool = False,
 ) -> torch.Tensor:
     """
     Fused experts with top-k routing.
@@ -188,6 +189,15 @@ def fused_experts(
     # assert dtype in [torch.float32, torch.float16, torch.bfloat16
     #                  ], "Only float32, float16, and bfloat16 are supported"
 
+    if apply_router_weight_on_input:
+        assert (topk_weights.dim() == 2
+                ), "`topk_weights` should be in shape (num_tokens, topk)"
+        _, topk = topk_weights.shape
+        assert (
+            topk == 1
+        ), "Only support topk=1 when `apply_router_weight_on_input` is True"
+        hidden_states = hidden_states * topk_weights.to(hidden_states.dtype)
+
     if expert_map is not None:
         # Generate token indices and flatten
         token_indices = (torch.arange(num_tokens,
@@ -289,14 +299,15 @@ def fused_experts(
             torch.zeros_like(weighted_down_out)).to(dtype)
         final_hidden_states.index_add_(0, sorted_token_indices, valid_output)
     else:
+        scales = torch.ones_like(topk_weights) if apply_router_weight_on_input else topk_weights
         # TODO: Reorder device memory 2 times here, replace the current
         # implementation here when suitable operators become available.
         final_hidden_states = torch_npu.npu_moe_finalize_routing(
             down_out_list,
             skip1=None,
             skip2=None,
             bias=None,
-            scales=topk_weights,
+            scales=scales,
             expanded_src_to_dst_row=expanded_row_idx,
             export_for_source_row=topk_ids,
         )
@@ -363,9 +374,6 @@ def select_experts(
     Raises:
         ValueError: If an unsupported scoring function is provided.
     """
-    if custom_routing_function is not None:
-        raise NotImplementedError(
-            "Custom routing function is not supported now")
 
     if scoring_func == "softmax":
         # NOTE: vLLM use dtype=torch.float here
@@ -402,9 +410,18 @@ def select_experts(
                                                 k=top_k,
                                                 dim=-1,
                                                 sorted=False)
-    else:
+    elif custom_routing_function is None:
         topk_weights, topk_ids = topk_weights.topk(top_k, dim=-1)
         topk_weights = topk_weights.to(hidden_states.dtype)
+    else:
+        topk_weights, topk_ids = custom_routing_function(
+                hidden_states=hidden_states,
+                gating_output=router_logits,
+                topk=top_k,
+                renormalize=renormalize)
+        # Required by npu_moe_init_routing
+        topk_ids = topk_ids.to(torch.int32)
+        return topk_weights, topk_ids
 
     # Required by npu_moe_init_routing
     topk_ids = topk_ids.to(torch.int32)