Update setup.py

Keeping in sync with upstream vllm - Preventing breakage for intel.
ROCm · May 21, 2024 · e7ebd71 · e7ebd71
1 parent 54d6fda
commit e7ebd71
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -365,10 +365,12 @@ def _read_requirements(filename: str) -> List[str]:
 if _is_cuda() and _install_punica():
     ext_modules.append(CMakeExtension(name="vllm._punica_C"))
 
+if _is_cuda() or _is_hip():
+    ext_modules.append(CMakeExtension(name="vllm._moe_C"))
+
 if not _is_neuron():
     ext_modules.append(CMakeExtension(name="vllm._C"))
     ext_modules.append(CMakeExtension(name="vllm._custom_C"))
-    ext_modules.append(CMakeExtension(name="vllm._moe_C"))
 
 package_data = {
     "vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]