[Kernel] Split Marlin MoE kernels into multiple files (vllm-project#8661

) Co-authored-by: mgoin <[email protected]>
axonn-ai · Sep 30, 2024 · 8e4e318 · 8e4e318
1 parent c98cb62
commit 8e4e318
Show file tree

Hide file tree

Showing 7 changed files with 1,552 additions and 1,427 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -316,6 +316,11 @@ set(VLLM_MOE_EXT_SRC
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_MOE_EXT_SRC
+      "csrc/moe/marlin_kernels/marlin_moe_kernel.h"
+      "csrc/moe/marlin_kernels/marlin_moe_kernel_ku4b8.h"
+      "csrc/moe/marlin_kernels/marlin_moe_kernel_ku4b8.cu"
+      "csrc/moe/marlin_kernels/marlin_moe_kernel_ku8b128.h"
+      "csrc/moe/marlin_kernels/marlin_moe_kernel_ku8b128.cu"
       "csrc/moe/marlin_moe_ops.cu")
 endif()