[Kernel] Expand MoE weight loading + Add Fused Marlin MoE Kernel (vll…

…m-project#7766) Co-authored-by: ElizaWszola <[email protected]>
triple-Mu · Sep 4, 2024 · 89ec134 · 89ec134
1 parent d58e692
commit 89ec134
Show file tree

Hide file tree

Showing 16 changed files with 2,382 additions and 85 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -302,6 +302,11 @@ set(VLLM_MOE_EXT_SRC
   "csrc/moe/torch_bindings.cpp"
   "csrc/moe/topk_softmax_kernels.cu")
 
+if(VLLM_GPU_LANG STREQUAL "CUDA")
+  list(APPEND VLLM_MOE_EXT_SRC
+      "csrc/moe/marlin_moe_ops.cu")
+endif()
+
 define_gpu_extension_target(
   _moe_C
   DESTINATION vllm