[Public release 26/04] Introducing EPv2: faster EP, and Engram/PP/CP supports

.editorconfig

@@ -0,0 +1,58 @@
+# https://editorconfig.org/
+
+root = true
+
+[*]
+charset = utf-8
+end_of_line = lf
+indent_style = space
+indent_size = 4
+trim_trailing_whitespace = true
+insert_final_newline = true
+
+[*.{py,pyi}]
+indent_size = 4
+
+[*.{cpp,hpp,cxx,cc,c,h,cu,cuh}]
+indent_size = 4
+
+[*.rs]
+indent_size = 4
+
+[*.go]
+indent_style = tab
+
+[*.{yaml,yml}]
+indent_size = 2
+
+[.clang-{format,tidy}]
+indent_size = 2
+
+[Makefile]
+indent_style = tab
+
+[*.sh]
+indent_size = 4
+
+[*.bat]
+indent_size = 4
+end_of_line = crlf
+
+[*.md]
+indent_size = 2
+x-soft-wrap-text = true
+
+[*.rst]
+indent_size = 4
+x-soft-wrap-text = true
+
+[*.{html,xml,css,scss,js,jsx,ts,tsx,vue}]
+indent_size = 2
+
+[**/test{,s,ing}/**/*.txt]
+trim_trailing_whitespace = false
+insert_final_newline = false
+
+[**/example{,s}/**/*.txt]
+trim_trailing_whitespace = false
+insert_final_newline = false

.gitignore

@@ -1,8 +1,28 @@
+# VS Code
 compile_commands.json
-.idea
-.DS_Store
-*.pyc
-build/
+.clangd
 .cache/
 .vscode/
+
+# macOS
+.DS_Store
+
+# JetBrains
+.idea
 */cmake-build-*/
+
+# Python
+*.pyc
+
+# Build
+build/
+*.so
+deep_ep.egg-info/
+dist/
+
+# Coredump
+coredump/
+
+# OpenCode
+.sisyphus/
+AGENTS.md

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third-party/fmt"]
+	path = third-party/fmt
+	url = https://github.com/fmtlib/fmt.git

CMakeLists.txt

@@ -0,0 +1,52 @@
+# NOTES: this CMake is only for debugging; for setup, please use Torch extension
+cmake_minimum_required(VERSION 3.10)
+project(deep_ep LANGUAGES CUDA CXX)
+set(CMAKE_VERBOSE_MAKEFILE ON)
+
+set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -O3 -fPIC")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC")
+set(CUDA_SEPARABLE_COMPILATION ON)
+list(APPEND CUDA_NVCC_FLAGS "-DENABLE_FAST_DEBUG")
+list(APPEND CUDA_NVCC_FLAGS "-O3")
+list(APPEND CUDA_NVCC_FLAGS "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage")
+list(APPEND CUDA_NVCC_FLAGS "-Xcompiler=-rdynamic")
+list(APPEND CUDA_NVCC_FLAGS "-lineinfo")
+# Suppress warnings for the `fmt` library
+list(APPEND CUDA_NVCC_FLAGS "--diag-suppress=128,2417")
+
+set(USE_SYSTEM_NVTX on)
+set(CUDA_ARCH_LIST "9.0" CACHE STRING "List of CUDA architectures to compile")
+set(TORCH_CUDA_ARCH_LIST "${CUDA_ARCH_LIST}")
+
+find_package(CUDAToolkit REQUIRED)
+find_package(pybind11 REQUIRED)
+find_package(Torch REQUIRED)
+
+# NVSHMEM
+find_package(NVSHMEM REQUIRED HINTS ${NVSHMEM_ROOT_DIR}/lib/cmake/nvshmem)
+add_library(nvshmem ALIAS nvshmem::nvshmem)
+add_library(nvshmem_host ALIAS nvshmem::nvshmem_host)
+add_library(nvshmem_device ALIAS nvshmem::nvshmem_device)
+
+# NCCL
+# TODO: use `find_package` instead of manual checks
+if (NOT NCCL_ROOT_DIR)
+    message(FATAL_ERROR "NCCL_ROOT_DIR is not set.")
+endif()
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CUDA_STANDARD 20)
+
+include_directories(deep_ep/include third-party/fmt/include)
+include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include ${CUDA_TOOLKIT_ROOT_DIR}/include/cccl ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS} ${NVSHMEM_INCLUDE_DIR} ${NCCL_ROOT_DIR}/include .)
+link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib ${NVSHMEM_LIB_DIR} ${NCCL_ROOT_DIR}/lib)
+
+# Add kernels
+add_subdirectory(csrc)
+
+# Link CPP and CUDA together
+pybind11_add_module(_C csrc/python_api.cpp)
+target_link_libraries(_C PRIVATE nccl ${RUNTIME_CUDA_LIBRARIES} ${LEGACY_CUDA_LIBRARIES} ${TORCH_LIBRARIES} torch_python)
+
+# Enable kernel code indexing with CMake-based IDEs
+cuda_add_library(deep_ep_indexing_cuda STATIC csrc/indexing/main.cu)