rebase wip

CMakeLists.txt

@@ -288,6 +288,46 @@ define_gpu_extension_target(
   USE_SABI 3
   WITH_SOABI)
 
+#
+# _vmm extension (VMM)
+#
+
+set(VLLM_VMM_EXT_SRC
+  "csrc/vmm/torch_bindings.cpp"
+  "csrc/vmm/vmm.cu"
+  )
+
+define_gpu_extension_target(
+  _vmm_C
+  DESTINATION vllm
+  LANGUAGE ${VLLM_GPU_LANG}
+  SOURCES ${VLLM_VMM_EXT_SRC}
+  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+  ARCHITECTURES ${VLLM_GPU_ARCHES}
+  USE_SABI 3
+  WITH_SOABI)
+
+#
+# dattn extension (DATTN)
+#
+
+set(VLLM_DATTN_EXT_SRC
+  "csrc/dattn/torch_bindings.cpp"
+  "csrc/dattn/dattn.cu"
+  )
+
+define_gpu_extension_target(
+  _dattn_C
+  DESTINATION vllm
+  LANGUAGE ${VLLM_GPU_LANG}
+  SOURCES ${VLLM_DATTN_EXT_SRC}
+  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+  ARCHITECTURES ${VLLM_GPU_ARCHES}
+  USE_SABI 3
+  WITH_SOABI)
+
+
+
 #
 # _moe_C extension
 #
@@ -315,4 +355,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   message(STATUS "Enabling moe extension.")
   add_dependencies(default _moe_C)
 
+  message(STATUS "Enabling VMM extension.")
+  add_dependencies(default _vmm_C)
+
+  message(STATUS "Enabling dAttention extension.")
+  add_dependencies(default _dattn_C)
+
 endif()