deepep prefill support a2

build.sh

@@ -2,6 +2,7 @@
 set -e
 
 BUILD_DEEPEP_MODULE="ON"
+BUILD_DEEPEP_OPS="ON"
 BUILD_KERNELS_MODULE="ON"
 BUILD_MEMORY_SAVER_MODULE="ON"
 
@@ -20,6 +21,11 @@ while getopts ":a:hd" opt; do
             case "$OPTARG" in
                 deepep )
                     BUILD_DEEPEP_MODULE="ON"
+                    BUILD_DEEPEP_OPS="ON"
+                    ;;
+                deepep2 )
+                    BUILD_DEEPEP_MODULE="ON"
+                    BUILD_DEEPEP_OPS="OFF"
                     ;;
                 kernels )
                     BUILD_KERNELS_MODULE="ON"
@@ -120,7 +126,11 @@ function build_deepep_kernels()
     if [[ "$ONLY_BUILD_DEEPEP_ADAPTER_MODULE" == "ON" ]]; then return 0; fi
     if [[ "$BUILD_DEEPEP_MODULE" != "ON" ]]; then return 0; fi
 
-    KERNEL_DIR="csrc/deepep/ops"
+    if [[ "$BUILD_DEEPEP_OPS" == "ON" ]]; then
+        KERNEL_DIR="csrc/deepep/ops"
+    else
+        KERNEL_DIR="csrc/deepep/ops2"
+    fi
     CUSTOM_OPP_DIR="${CURRENT_DIR}/python/deep_ep/deep_ep"
 
     cd "$KERNEL_DIR" || exit
@@ -137,6 +147,7 @@ function build_deepep_kernels()
         echo "find run package: $custom_opp_file"
         chmod +x "$custom_opp_file"
     fi
+    rm -rf "$CUSTOM_OPP_DIR"/vendors
     ./build_out/custom_opp_*.run --install-path=$CUSTOM_OPP_DIR
     cd -
 }

csrc/deepep/CMakeLists.txt

@@ -1,6 +1,14 @@
 # this is the cmakelist file for deepep build
 # deepep will be built as separated wheel package
 
+if(${CMAKE_SYSTEM_PROCESSOR} STREQUAL "x86_64")
+    set(DEEPEP_ARCH "x86_64")
+elseif(${CMAKE_SYSTEM_PROCESSOR} STREQUAL "aarch64")
+    set(DEEPEP_ARCH "aarch64")
+else()
+    message(FATAL_ERROR "Unsupported host processor: ${CMAKE_SYSTEM_PROCESSOR}")
+endif()
+
 set(PROJECT_BUILD_PATH ${PROJECT_BINARY_DIR})
 set(TARGET_INSTALL_DIR ${CMAKE_INSTALL_PREFIX})
 set(ASCEND_HOME_PATH ${ASCEND_HOME_PATH})
@@ -27,6 +35,7 @@ target_include_directories( deep_ep_cpp PRIVATE
     ${TORCH_NPU_DIR}/include/third_party/acl/inc/acl
     ${TORCH_NPU_DIR}/include/third_party/acl/inc
     ${ASCEND_HOME_PATH}/include
+    ${ASCEND_HOME_PATH}/${DEEPEP_ARCH}-linux/include/experiment/platform
 )
 target_link_directories(deep_ep_cpp PRIVATE
         ${TORCH_DIR}/lib
@@ -38,6 +47,7 @@ target_link_libraries(deep_ep_cpp PRIVATE
     ascendcl
     hccl
     torch_npu
+    opapi
 )
 
 message(STATUS "TARGET_INSTALL_DIR = ${TARGET_INSTALL_DIR}")

csrc/deepep/deep_ep.hpp

@@ -7,15 +7,17 @@
 #include <optional>
 #include "hccl/hccl.h"
 #include "hccl/hccl_types.h"
+#include "aclnn/opdev/platform.h"
 
 #include "config.hpp"
 #include "event.hpp"
 
 namespace deep_ep {
 
 struct Buffer {
-    int64_t rank, rdma_rank;
-    int64_t num_ranks;
+    int64_t rank, rdma_rank, nvl_rank;
+    int64_t num_ranks, num_rdma_ranks, num_nvl_ranks;
+    op::SocVersion soc_version;
 
     int64_t num_nvl_bytes;
     int64_t num_rdma_bytes;
@@ -26,6 +28,7 @@ struct Buffer {
     at::Tensor ori_x;
     at::Tensor new_topk_idx;
     at::Tensor new_scales;
+    at::Tensor send_data;
 
     int64_t shared_expert_rank_num;
     int64_t shared_expert_num = 1;
@@ -47,6 +50,10 @@ struct Buffer {
 
     bool is_available() const;
 
+    int get_num_rdma_ranks() const;
+
+    int get_rdma_rank() const;
+
     std::tuple<torch::Tensor, std::optional<torch::Tensor>, torch::Tensor, torch::Tensor, std::optional<EventHandle>>
     get_dispatch_layout(const torch::Tensor &topk_idx, int num_experts, std::optional<EventHandle> &previous_event,
                         bool async, bool allocate_on_comm_stream);
@@ -70,15 +77,29 @@ struct Buffer {
                       const std::optional<torch::Tensor> &topk_weights, const torch::Tensor &src_idx,
                       const torch::Tensor &send_head, const std::optional<at::Tensor> &combine_send_cost_stats);
 
+    std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<torch::Tensor>, std::optional<torch::Tensor>,
+               std::vector<int>, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+               torch::Tensor, std::optional<EventHandle>>
+    internode_dispatch(const torch::Tensor &x, const std::optional<torch::Tensor> &x_scales,
+                       const std::optional<torch::Tensor> &topk_idx, const std::optional<torch::Tensor> &topk_weights,
+                       const std::optional<torch::Tensor> &num_tokens_per_rank,
+                       const std::optional<torch::Tensor> &num_tokens_per_rdma_rank,
+                       const torch::Tensor &is_token_in_rank, const std::optional<torch::Tensor> &num_tokens_per_expert,
+                       const Config &config, std::optional<EventHandle> &previous_event, bool async,
+                       bool allocate_on_comm_stream, bool use_quant);
+
+    std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandle>> internode_combine(
+        const torch::Tensor &x, const torch::Tensor &topk_idx, const std::optional<torch::Tensor> &topk_weights,
+        const torch::Tensor &src_idx, const torch::Tensor &send_head, const torch::Tensor &offsetInner,
+        const torch::Tensor &offsetOuter, const torch::Tensor &countOuter, const torch::Tensor &expand_scales);
+
     std::tuple<at::Tensor, std::optional<at::Tensor>, at::Tensor, at::Tensor, at::Tensor, std::optional<EventHandle>,
                std::optional<std::function<void()>>>
     low_latency_dispatch(const at::Tensor &x, const at::Tensor &topk_idx,
                          const std::optional<at::Tensor> &cumulative_local_expert_recv_stats,
                          int64_t num_max_dispatch_tokens_per_rank, int64_t num_experts, bool use_fp8, bool round_scale,
                          bool use_ue8m0, bool async, bool return_recv_hook);
 
-    int get_rdma_rank() const;
-
     std::tuple<at::Tensor, std::optional<EventHandle>, std::optional<std::function<void()>>> low_latency_combine(
         const at::Tensor &x, const at::Tensor &topk_idx, const at::Tensor &topk_weights, const at::Tensor &src_info,
         const at::Tensor &layout_range, int64_t num_max_dispatch_tokens_per_rank, int64_t num_experts,

csrc/deepep/ops2/CMakeLists.txt

@@ -0,0 +1,67 @@
+cmake_minimum_required(VERSION 3.16.0)
+project(opp)
+
+if(${CMAKE_SYSTEM_PROCESSOR} STREQUAL "x86_64")
+    set(CANN_HOST_ARCH "x86_64")
+elseif(${CMAKE_SYSTEM_PROCESSOR} STREQUAL "aarch64")
+    set(CANN_HOST_ARCH "aarch64")
+else()
+    message(FATAL_ERROR "Unsupported host processor: ${CMAKE_SYSTEM_PROCESSOR}")
+endif()
+
+include(cmake/config.cmake)
+include(cmake/func.cmake)
+include(cmake/intf.cmake)
+
+set(CMAKE_COMPILE ${CMAKE_CXX_COMPILER})
+
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/framework)
+    add_subdirectory(framework)
+endif()
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/op_host)
+    add_subdirectory(op_host)
+endif()
+if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/op_kernel)
+    add_subdirectory(op_kernel)
+endif()
+if(ENABLE_TEST AND EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/testcases)
+    add_subdirectory(testcases)
+endif()
+
+add_custom_command(OUTPUT ${CMAKE_BINARY_DIR}/scripts/install.sh ${CMAKE_BINARY_DIR}/scripts/upgrade.sh
+    COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_BINARY_DIR}/scripts
+    COMMAND ${CMAKE_COMMAND} -E copy_directory ${CMAKE_SOURCE_DIR}/scripts/ ${CMAKE_BINARY_DIR}/scripts/
+    COMMAND sed -i "s/vendor_name=customize/vendor_name=${vendor_name}/g" ${CMAKE_BINARY_DIR}/scripts/install.sh ${CMAKE_BINARY_DIR}/scripts/upgrade.sh
+    VERBATIM
+)
+add_custom_target(modify_vendor ALL DEPENDS ${CMAKE_BINARY_DIR}/scripts/install.sh ${CMAKE_BINARY_DIR}/scripts/upgrade.sh)
+
+get_system_info(SYSTEM_INFO)
+
+# gen version.info
+add_custom_target(gen_version_info ALL
+        COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/cmake/util/gen_version_info.sh ${ASCEND_CANN_PACKAGE_PATH} ${CMAKE_CURRENT_BINARY_DIR}
+)
+
+if(NOT ASCEND_PACK_SHARED_LIBRARY)
+    install(DIRECTORY ${CMAKE_BINARY_DIR}/scripts/ DESTINATION . FILE_PERMISSIONS OWNER_EXECUTE OWNER_READ GROUP_READ)
+
+    install(FILES ${CMAKE_SOURCE_DIR}/custom.proto DESTINATION packages OPTIONAL)
+    install(FILES ${CMAKE_CURRENT_BINARY_DIR}/version.info
+            DESTINATION packages/vendors/${vendor_name}/)
+
+    # CPack config
+    set(CPACK_PACKAGE_NAME ${CMAKE_PROJECT_NAME})
+    set(CPACK_PACKAGE_VERSION ${CMAKE_PROJECT_VERSION})
+    set(CPACK_PACKAGE_DESCRIPTION "CPack opp project")
+    set(CPACK_PACKAGE_DESCRIPTION_SUMMARY "CPack opp project")
+    set(CPACK_PACKAGE_DIRECTORY ${CMAKE_INSTALL_PREFIX})
+
+    set(CPACK_PACKAGE_FILE_NAME "custom_opp_${SYSTEM_INFO}_${CMAKE_SYSTEM_PROCESSOR}.run")
+    set(CPACK_GENERATOR External)
+    set(CPACK_CMAKE_GENERATOR "Unix Makefiles")
+    set(CPACK_EXTERNAL_ENABLE_STAGING TRUE)
+    set(CPACK_EXTERNAL_PACKAGE_SCRIPT ${CMAKE_SOURCE_DIR}/cmake/makeself.cmake)
+    set(CPACK_EXTERNAL_BUILT_PACKAGES ${CPACK_PACKAGE_DIRECTORY}/_CPack_Packages/Linux/External/${CPACK_PACKAGE_FILE_NAME}/${CPACK_PACKAGE_FILE_NAME})
+    include(CPack)
+endif()

csrc/deepep/ops2/CMakePresets.json

@@ -0,0 +1,59 @@
+{
+    "version": 1,
+    "cmakeMinimumRequired": {
+        "major": 3,
+        "minor": 19,
+        "patch": 0
+    },
+    "configurePresets": [
+        {
+            "name": "default",
+            "displayName": "Default Config",
+            "description": "Default build using Unix Makefiles generator for native compilation",
+            "generator": "Unix Makefiles",
+            "binaryDir": "${sourceDir}/build_out",
+            "cacheVariables": {
+                "CMAKE_BUILD_TYPE": {
+                    "type": "STRING",
+                    "value": "Release"
+                },
+                "ENABLE_SOURCE_PACKAGE": {
+                    "type": "BOOL",
+                    "value": "True"
+                },
+                "ENABLE_BINARY_PACKAGE": {
+                    "type": "BOOL",
+                    "value": "True"
+                },
+                "ASCEND_COMPUTE_UNIT": {
+                    "type": "STRING",
+                    "value": "ascend910b"
+                },
+                "ENABLE_TEST": {
+                    "type": "BOOL",
+                    "value": "True"
+                },
+                "vendor_name": {
+                    "type": "STRING",
+                    "value": "hwcomputing"
+                },
+                "ASCEND_CANN_PACKAGE_PATH": {
+                    "type": "PATH",
+                    "value": "/usr/local/Ascend/ascend-toolkit/latest"
+                },
+                "ASCEND_PYTHON_EXECUTABLE": {
+                    "type": "STRING",
+                    "value": "python3"
+                },
+                "CMAKE_INSTALL_PREFIX": {
+                    "type": "PATH",
+                    "value": "${sourceDir}/build_out"
+                },
+                "ASCEND_PACK_SHARED_LIBRARY": {
+                    "type": "BOOL",
+                    "value": "False"
+                }
+            }
+        }
+    ]
+}

csrc/deepep/ops2/README.md

@@ -0,0 +1,10 @@
+# moe_dispatch_combine
+
+# compile
+bash build.sh
+
+# install
+./build_out/custom_opp_ubuntu_aarch64.run --install-path=/usr/local/Ascend/ascend-toolkit/latest/opp/
+
+# require import env parameters before running
+source /usr/local/Ascend/ascend-toolkit/latest/opp/vendors/hwcomputing/bin/set_env.bash

csrc/deepep/ops2/build.sh

@@ -0,0 +1,69 @@
+#!/bin/bash
+
+export OPS_PROJECT_NAME=aclnnInner
+
+SCRIPTS_DIR=$(cd "$(dirname "$0")" && pwd)
+
+if [ -n "$BASE_LIBS_PATH" ]; then
+  export ASCEND_HOME_PATH="$BASE_LIBS_PATH"
+elif [ -z "$ASCEND_HOME_PATH" ]; then
+  if [ -n "$ASCEND_AICPU_PATH" ]; then
+    export ASCEND_HOME_PATH="$ASCEND_AICPU_PATH"
+  else
+    echo "please set env." >&2
+    exit 1
+  fi
+fi
+echo "using ASCEND_HOME_PATH: $ASCEND_HOME_PATH"
+script_path=$(realpath $(dirname $0))
+
+BUILD_DIR="build_out"
+HOST_NATIVE_DIR="host_native_tiling"
+
+chmod +x cmake/util/gen_ops_filter.sh
+mkdir -p build_out
+rm -rf build_out/*
+
+opts=$(python3 $script_path/cmake/util/preset_parse.py $script_path/CMakePresets.json)
+ENABLE_CROSS="-DENABLE_CROSS_COMPILE=True"
+ENABLE_BINARY="-DENABLE_BINARY_PACKAGE=True"
+ENABLE_LIBRARY="-DASCEND_PACK_SHARED_LIBRARY=True"
+cmake_version=$(cmake --version | grep "cmake version" | awk '{print $3}')
+
+target=package
+if [ -n "$1" ]; then target="$1"; fi
+if [[ $opts =~ $ENABLE_LIBRARY ]]; then target=install; fi
+
+if [[ $opts =~ $ENABLE_CROSS ]] && [[ $opts =~ $ENABLE_BINARY ]]
+then
+  if [ "$cmake_version" \< "3.19.0" ] ; then
+    cmake -S . -B "$BUILD_DIR" $opts -DENABLE_CROSS_COMPILE=0
+  else
+    cmake -S . -B "$BUILD_DIR" --preset=default -DENABLE_CROSS_COMPILE=0
+  fi
+  cmake --build "$BUILD_DIR" --target cust_optiling
+  mkdir $BUILD_DIR/$HOST_NATIVE_DIR
+  lib_path=$(find "$BUILD_DIR" -name "libcust_opmaster_rt2.0.so")
+  if [ -z "$lib_path" ] || [ $(echo "$lib_path" | wc -l) -ne 1 ]; then
+    echo "Error: Expected to find exactly one libcust_opmaster_rt2.0.so, but found none or multiple." >&2
+    exit 1
+  fi
+  mv "$lib_path" "$BUILD_DIR/$HOST_NATIVE_DIR/"
+  find "$BUILD_DIR" -mindepth 1 -maxdepth 1 ! -name "$HOST_NATIVE_DIR" -exec rm -rf {} +
+  host_native_tiling_lib=$(realpath $(find $BUILD_DIR -type f -name "libcust_opmaster_rt2.0.so"))
+  if [ "$cmake_version" \< "3.19.0" ] ; then
+    cmake -S . -B "$BUILD_DIR" $opts -DHOST_NATIVE_TILING_LIB=$host_native_tiling_lib
+  else
+    cmake -S . -B "$BUILD_DIR" --preset=default -DHOST_NATIVE_TILING_LIB=$host_native_tiling_lib
+  fi
+  cmake --build "$BUILD_DIR" --target binary -j$(nproc)
+  cmake --build "$BUILD_DIR" --target $target -j$(nproc)
+else
+  if [ "$cmake_version" \< "3.19.0" ] ; then
+    cmake -S . -B "$BUILD_DIR" $opts
+  else
+      cmake -S . -B "$BUILD_DIR" --preset=default
+  fi
+  cmake --build "$BUILD_DIR" --target binary -j$(nproc)
+  cmake --build "$BUILD_DIR" --target $target -j$(nproc)
+fi

csrc/deepep/ops2/cmake/config.cmake

@@ -0,0 +1,42 @@
+
+set(CMAKE_CXX_FLAGS_DEBUG "")
+set(CMAKE_CXX_FLAGS_RELEASE "")
+
+if (NOT DEFINED CMAKE_BUILD_TYPE)
+    set(CMAKE_BUILD_TYPE Release CACHE STRING "")
+endif()
+if (CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
+    set(CMAKE_INSTALL_PREFIX "${CMAKE_SOURCE_DIR}/build_out" CACHE PATH "" FORCE)
+endif()
+if (NOT DEFINED ASCEND_CANN_PACKAGE_PATH)
+    set(ASCEND_CANN_PACKAGE_PATH /usr/local/Ascend/latest CACHE PATH "")
+endif()
+if (NOT DEFINED ASCEND_PYTHON_EXECUTABLE)
+    set(ASCEND_PYTHON_EXECUTABLE python3 CACHE STRING "")
+endif()
+if (NOT DEFINED ASCEND_COMPUTE_UNIT)
+    set(ASCEND_COMPUTE_UNIT ascend910_93 CACHE STRING "")
+endif()
+if (NOT DEFINED ENABLE_TEST)
+    set(ENABLE_TEST FALSE CACHE BOOL "")
+endif()
+if (NOT DEFINED ENABLE_CROSS_COMPILE)
+    set(ENABLE_CROSS_COMPILE  FALSE CACHE BOOL "")
+endif()
+if (NOT DEFINED CMAKE_CROSS_PLATFORM_COMPILER)
+    set(CMAKE_CROSS_PLATFORM_COMPILER "/your/cross/compiler/path" CACHE PATH "")
+endif()
+if (NOT DEFINED CMAKE_CROSS_LIBRARY_PATH)
+    set(CMAKE_CROSS_LIBRARY_PATH "" CACHE PATH "")
+endif()
+if (NOT DEFINED ASCEND_PACK_SHARED_LIBRARY)
+    set(ASCEND_PACK_SHARED_LIBRARY False CACHE BOOL "")
+endif()
+set(ASCEND_TENSOR_COMPILER_PATH ${ASCEND_CANN_PACKAGE_PATH}/compiler)
+set(ASCEND_CCEC_COMPILER_PATH ${ASCEND_TENSOR_COMPILER_PATH}/ccec_compiler/bin)
+set(ASCEND_AUTOGEN_PATH ${CMAKE_BINARY_DIR}/autogen)
+file(MAKE_DIRECTORY ${ASCEND_AUTOGEN_PATH})
+set(CUSTOM_COMPILE_OPTIONS "custom_compile_options.ini")
+set(CUSTOM_OPC_OPTIONS "custom_opc_options.ini")
+file(WRITE ${ASCEND_AUTOGEN_PATH}/${CUSTOM_COMPILE_OPTIONS} "")
+file(WRITE ${ASCEND_AUTOGEN_PATH}/${CUSTOM_OPC_OPTIONS} "")