[Kernel][ROCm][AMD] enable fused topk_softmax kernel for moe layer (#4927)

This PR enables the fused topk_softmax kernel used in moe layer for HIP
2024-06-02 16:13:26 -05:00
parent f790ad3c50
commit a66cf40b20
6 changed files with 45 additions and 43 deletions
--- a/setup.py
+++ b/setup.py
@ -382,7 +382,7 @@ def get_requirements() -> List[str]:

 ext_modules = []

-if _is_cuda():
+if _is_cuda() or _is_hip():
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

 if not _is_neuron():