[Kernel] Expand MoE weight loading + Add Fused Marlin MoE Kernel (#7766)

Co-authored-by: ElizaWszola <[email protected]>
vllm-project · Aug 27, 2024 · fc91188 · fc91188
1 parent ed6f002
commit fc91188
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
@@ -296,6 +296,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