Use torch top-k routing fallback

zonos2/layers/moe/fused_moe/topk.py

@@ -1,14 +1,7 @@
-from typing import Any, Optional
+from typing import Optional
 
 import torch
 
-try:
-    from sgl_kernel import topk_softmax
-except ImportError:
-
-    def topk_softmax(*args: Any, **kwargs: Any) -> Any:
-        raise NotImplementedError("sgl_kernel not found.")
-
 
 def fused_topk(
     hidden_states: torch.Tensor,
@@ -21,15 +14,9 @@ def fused_topk(
 
     M, _ = hidden_states.shape
 
-    topk_weights = torch.empty(M, topk, dtype=torch.float32, device=hidden_states.device)
-    topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device)
-
-    topk_softmax(
-        topk_weights,
-        topk_ids,
-        gating_output.float(),
-        renormalize,
-    )
+    scores = torch.softmax(gating_output.float(), dim=-1)
+    topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
+    topk_ids = topk_ids.to(dtype=torch.int32)
 
     return _fused_topk_postprocess(
         topk_weights=topk_weights,