feat: update other MoE models deps (#2156)

python/sglang/srt/layers/fused_moe/layer.py

@@ -153,12 +153,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         num_expert_group: Optional[int],
         topk_group: Optional[int],
     ) -> torch.Tensor:
-        from vllm.model_executor.layers.fused_moe.moe_pallas import fused_moe
-
-        assert not use_grouped_topk
-        assert num_expert_group is None
-        assert topk_group is None
-        return fused_moe(x, w1, w2, router_logits, top_k, renormalize)
+        raise NotImplementedError("The TPU backend currently does not support MoE.")
 
 
 class FusedMoE(torch.nn.Module):

python/sglang/srt/layers/triton_fused_moe/fused_moe.py

@@ -1,3 +1,5 @@
+# Adapted from https://github.com/vllm-project/vllm/blob/a6221a144af772fd1a68fe7e627935dc53e81738/vllm/model_executor/layers/fused_moe/fused_moe.py
+
 """Fused MoE kernel."""
 
 import functools

python/sglang/srt/layers/triton_fused_moe/layer.py

@@ -1,3 +1,5 @@
+# Adapted from https://github.com/vllm-project/vllm/blob/a6221a144af772fd1a68fe7e627935dc53e81738/vllm/model_executor/layers/fused_moe/layer.py
+
 from abc import abstractmethod
 from enum import Enum
 from typing import Callable, List, Optional, Tuple
@@ -18,7 +20,7 @@ from sglang.srt.layers.quantization.base_config import (
 from sglang.srt.utils import set_weight_attrs
 
 if torch.cuda.is_available() or torch.hip.is_available():
-    from .fused_moe import fused_experts
+    from sglang.srt.layers.triton_fused_moe.fused_moe import fused_experts
 else:
     fused_experts = None  # type: ignore
 
@@ -512,7 +514,7 @@ class FusedMoE(torch.nn.Module):
         num_expert_group: Optional[int] = None,
         custom_routing_function: Optional[Callable] = None,
     ):
-        from vllm.model_executor.layers.fused_moe.fused_moe import (
+        from sglang.srt.layers.triton_fused_moe.fused_moe import (
             fused_topk,
             grouped_topk,
         )