diff --git a/docs/superpowers/specs/2026-07-14-a5-dispatch-reachability-pruning-design.md b/docs/superpowers/specs/2026-07-14-a5-dispatch-reachability-pruning-design.md new file mode 100644 index 00000000..bd331346 --- /dev/null +++ b/docs/superpowers/specs/2026-07-14-a5-dispatch-reachability-pruning-design.md @@ -0,0 +1,152 @@ +# A5 Dispatch 可达性裁剪设计 + +## 目标 + +精简 `reference/moe_distribute_dispatch_v2_full_mesh_3510_simplified.h`，仅保留 +A5 `Init()` 和 `Process()` 流程可能使用的代码，同时保证所有受支持路径的行为不变。 + +## 范围 + +- 以 `Init()` 和 `Process()` 作为调用图根节点。 +- 除下文明确固定的特性输入外，所有模板实例和运行时分支均视为可能生效。 +- 删除不可达成员函数的类内声明和类外定义。 +- 删除函数裁剪后没有引用的类成员变量。 +- 删除函数和成员裁剪后没有引用的文件级常量。 +- 保持现有 `#include` 指令不变。 +- 保留公开类型别名、构造函数以及 `Init()`、`Process()` 的受支持行为。 + +## 固定特性输入 + +简化后的 A5 参考代码采用以下固定输入： + +- `zeroComputeExpertNum` 永远为 `0`。 +- `isPerformance` 永远为 `false`。 +- `hasElasticInfo` 永远为 `false`。 + +代码不再读取这些 tiling 字段，也不再保存对应类成员。依赖这些输入的条件将折叠到 +固定分支，由此变为不可达的特性代码将被删除。 + +为了保持调用接口兼容，`Init()` 声明和定义中的 `elasticInfo`、`performanceInfo` +形参继续保留，但不再绑定 GlobalTensor，也不会在函数体中使用。 + +## 函数候选 + +应用固定特性输入前，当前词法调用图包含 49 个类外成员函数定义。从 `Init()` 和 +`Process()` 出发的传递可达闭包包含 42 个定义。以下 7 个函数已经不可达： + +- `AllToAllDispatchA3` +- `CalcBSTokenRange` +- `CalExpertSendNum` +- `SendBSExpertLoop` +- `SendToMoeExpertByBS` +- `SetExpertTokenNums` +- `SplitExpertNumToCore` + +这些函数的类内声明和类外定义都将删除。 + +折叠三个固定特性输入后，以下 6 个函数也会变为不可达： + +- `InitElasticInfo` +- `CalAndSendCntByExp` +- `RecordRankCommDuration` +- `GenerateGatherMaskTensor` +- `MaskZeroComputeExpert` +- `ZeroComputeExpertMaskCal` + +分支折叠完成后会重新计算调用图。只有确认不属于 `Init()`、`Process()` 传递可达 +闭包的函数才会继续删除。 + +## 分支折叠 + +实现时进行以下等价替换： + +- Mask Buffer 的分配条件简化为 `isTokenMaskFlag_ || isExpertMaskFlag_`，删除 + zero 专家 Mask 初始化和计算。 +- `CalCumSum()` 直接调用 `CalAndSendCntByRank()`。 +- Rank 地址计算直接使用非扩缩容场景的 Rank ID，删除所有 + `isScalingDownFlag_` 重映射分支。 +- LocalWindow 源数据索引条件简化为 `if (!isShareExpertRankFlag_)`。 +- 删除性能 Buffer、计时调用和性能输出拷贝。 +- `Init()` 保留兼容形参，但不创建弹性信息或性能信息 GlobalTensor。 + +## 成员变量候选 + +删除原有不可达函数后，以下 18 个私有成员不再有引用： + +- `axisHExpandXAlignSize_` +- `cleanStatusTensor_` +- `cumSumTime1Tensor_` +- `cumSumTime2Tensor_` +- `cumSumTimes_` +- `cumSumUB_` +- `dealRankPerCore_` +- `delLastExpertId_` +- `flagPadOffset_` +- `gatherTmpTensor_` +- `maskSizePerExpert_` +- `remainderExpertNum_` +- `sharedTmpBufTensor_` +- `statusSumOutTensor_` +- `syncOnCoreTensor_` +- `tempTime1Tensor_` +- `tempTime2Tensor_` +- `tokenNumToExpertTensor_` + +编辑后会重新进行引用分析。只有声明成为唯一剩余引用时，成员才会被删除。 + +固定特性输入还会使以下 12 个成员失去用途： + +- `zeroComputeExpertNum_` +- `hasElasticInfoFlag_` +- `isScalingDownFlag_` +- `isPerformanceFlag_` +- `elasticInfoGMTensor_` +- `elasticInfoTensor_` +- `elasticInfoBuf_` +- `performanceInfoGMTensor_` +- `performanceInfoTensor_` +- `performanceFlagTensor_` +- `performanceInfoBuf_` +- `performanceFlagBuf_` + +如果共享 Mask 或临时 Buffer 仍被 Token Mask、Expert Mask 或正常 Dispatch 路径 +使用，则继续保留。 + +## 文件级常量候选 + +以下 4 个常量在可达代码中没有引用： + +- `AIV_STATE_SIZE` +- `MIN_ACTIVE_BS_FOR_BS_MODE` +- `SFFVALUE_SIZE` +- `SYNC_OFFSET` + +删除性能打点后，`DURATION_OFFSET` 也会失去引用。 + +只有在编辑后的引用扫描确认常量仅剩声明时，才会将其删除。 + +## 编辑方法 + +删除函数定义时使用大括号配平后的精确范围，不通过下一个函数签名推断删除边界， +避免误删相邻的模板声明和注释。函数声明、成员变量和常量使用局部补丁删除。 + +可达函数体只允许进行上文明确列出的固定输入分支折叠。其他可达代码不进行格式化 +或重写。 + +## 验证标准 + +处理结果必须满足以下检查： + +1. `Init()` 和 `Process()` 仍然存在，其可达调用闭包不存在缺失的类成员函数定义。 +2. 每个保留的类外函数定义都有对应的类内声明。 +3. 原有 7 个不可达函数和新增 6 个特性函数均不存在声明、定义或调用点。 +4. `zeroComputeExpertNum_`、`hasElasticInfoFlag_`、`isScalingDownFlag_`、 + `isPerformanceFlag_` 不再出现。 +5. `elasticInfo`、`performanceInfo` 仅作为 `Init()` 兼容形参存在，不在函数体中使用。 +6. 已删除的成员变量和常量不存在残留引用。 +7. 可达函数体中使用的成员标识符仍有对应类成员声明。 +8. 预处理指令和大括号保持配平。 +9. 文件保持无 BOM 的 UTF-8 编码和 LF 换行。 + +该参考头文件不参与当前 Windows 构建，并依赖外部 Ascend C 头文件，因此本阶段采用 +结构化验证，不执行本地编译。 diff --git a/docs/superpowers/specs/2026-07-15-ep-dispatch-memory-reference-port-design.md b/docs/superpowers/specs/2026-07-15-ep-dispatch-memory-reference-port-design.md new file mode 100644 index 00000000..7304a0cc --- /dev/null +++ b/docs/superpowers/specs/2026-07-15-ep-dispatch-memory-reference-port-design.md @@ -0,0 +1,243 @@ +# EP Dispatch Memory Reference Port Design + +## Goal + +Port the complete reachable non-quantized A5 full-mesh dispatch behavior from +`reference/moe_distribute_dispatch_v2_full_mesh_3510_simplified.h` into +`src/ep/kernels/tilexr_ep_dispatch_memory_kernel.cpp`. + +The TileXR kernel is launched directly through +`launch_tilexr_ep_dispatch_memory_kernel`. It does not expose the reference +implementation's separate `Init()` and `Process()` calling convention. + +## Supported Scope + +The first version supports: + +- FP16 and BF16 input/output with identical input and output types; +- normal MoE experts; +- shared experts, including multiple ranks per shared expert; +- no active mask, token mask `[bs]`, or expert mask `[bs, topK]`; +- count and prefix-sum forms of `expertTokenNumsOut`; +- the reference AIV split between dispatch and count/cumsum work; +- the reference 512-byte DataAsFlag token slots and state-window protocol; +- the reference local output compaction order. + +The first version rejects: + +- all quantization modes; +- smooth-scale input and scale outputs; +- differing input and output types; +- TP execution; +- elastic rank remapping, zero-compute experts, and performance recording, + which are already removed from the simplified reference source. + +## Kernel Organization + +`tilexr_ep_dispatch_memory_kernel.cpp` contains the ported implementation. +The obsolete `tilexr_ep_dispatch_memory_helpers.h` chunk/source-slot protocol +is deleted and is not used as an implementation reference. + +The kernel constructs a `TPipe` and invokes one internal `Run()` entry. `Run()` +contains the reference initialization work followed by the reference process +work. The following reference functions retain their algorithms, AIV ownership, +and call relationships: + +- `SetTilingDataAndCal` +- `SetDataStatus` +- `TokenToExpert` +- `SplitToCore` +- `SendToSharedExpert` +- `SendToMoeExpert` +- `CalcSendTokenBufNum` +- `AllToAllDispatch` +- `AllToAllDispatchA5` +- `CalTokenSendExpertCnt` +- `CalAndSendCntByRank` +- `BufferInit` +- `WaitDispatchClearStatus` +- `GatherSumRecvCnt` +- `GetCumSum` +- `WaitDispatch` +- `CalRecvAndSetFlag` +- `CalCumSum` +- `WaitCumSumFlag` +- `SetValidExpertInfo` +- `CheckDataArriveWithFlag` +- `CopyInAndOut` +- `WaitAndFormatOutput` +- `RunPosRecord` +- `LocalWindowCopy` +- token-mask and expert-mask calculation helpers + +Quantization-only functions and branches are removed rather than stubbed. + +## Communication Context Mapping + +The reference A5 context places a one-MiB state area before the data window. +TileXR reproduces that relationship inside each peer's IPC data region: + +```text +peerWindowBase(rank) = commArgs.peerMems[rank] + IPC_DATA_OFFSET +stateWindow(rank) = peerWindowBase(rank) +dataWindow(rank) = peerWindowBase(rank) + 1 MiB +statusDataSpace = stateWindow(selfRank) +``` + +The first one MiB therefore keeps the reference layout: + +- ping state 0: `[0, 384 KiB)`; +- ping state 1: `[384 KiB, 768 KiB)`; +- per-AIV run state from `768 KiB`; +- cumsum exchange from `868 KiB`; +- cumsum completion flags from `876 KiB`. + +The mapping does not use or overwrite TileXR's flag region before +`IPC_DATA_OFFSET`. + +The data window starts after the one-MiB state area. Its two halves use the +reference `dataState * (totalWinSize / 2)` selection and retain the combine +reserve prefix and the dispatch DataAsFlag layout. + +## Local Workspace + +The reference cumsum path requires `workspaceGM`, replicated once per AIV. +TileXR reserves this workspace at the tail of the local IPC data region: + +```text +workspaceStatusNum = epWorldSize * moeExpertNumPerRank +workspaceBytes = align32(aivNum * workspaceStatusNum * sizeof(int32_t)) +totalWinSize = IPC_BUFF_MAX_SIZE - 1 MiB - workspaceBytes +workspaceGM = local dataWindow base + totalWinSize +``` + +Only the local rank accesses this workspace. The reservation nevertheless uses +the communicator-wide maximum receive-status count rather than the local +`rscvStatusNum`. This keeps `totalWinSize` and both dispatch-half offsets +identical on shared-expert and MoE-expert ranks, so remote payload addresses +match the receiver's polling addresses. Host validation checks that each data +state half can contain the combine reserve and every reference expert segment +before launch. + +## Launch Interface + +The internal launch interface becomes: + +```cpp +void launch_tilexr_ep_dispatch_memory_kernel( + uint32_t blockDim, + void *stream, + GM_ADDR commArgs, + GM_ADDR x, + GM_ADDR expertIds, + GM_ADDR xActiveMask, + GM_ADDR expandXOut, + GM_ADDR expertTokenNumsOut, + GM_ADDR epRecvCountsOut, + GM_ADDR assistInfoForCombineOut, + int64_t bs, + int64_t h, + int64_t topK, + int64_t moeExpertNum, + int64_t sharedExpertNum, + int64_t sharedExpertRankNum, + int64_t globalBs, + int64_t expertTokenNumsType, + int64_t activeMaskType, + int64_t dtype, + int64_t magic); +``` + +`CommArgs` supplies rank, world size, and peer addresses. `blockDim` supplies +the reference `aivNum`. Window sizes, UB sizes, core-group sizes, and offsets +are derived by the exact reference formulas and validated on the host. + +`magic` is the communicator-wide invocation sequence. Its low bit selects the +reference ping-pong state half consistently across ranks; TileXR IPC buffers do +not provide the pre-synchronized state bit assumed by the HCCL/MC2 context. + +The old route-count, chunk-count, offset, source-slot, payload, and total-size +arguments are deleted. `magic` is retained solely for communicator-wide state +selection. + +Each receive-count status occupies one 32-byte block. The last two `int32` +fields store the count followed by the `float(1.0)` arrival flag. This adapts +the reference status block to TileXR's existing payload-before-tail-flag +ordering. The sender copies the complete block, or strided complete blocks, +with one `DataCopy`; the receiver polls the tail flags before consuming the +adjacent counts. The +memory dispatch path must not add a separate `SyncCollectives` ready message, +because that would split one reference transaction into a count write followed +by another remote flag write. + +## Active Mask Contract + +The public API adds an active-mask type because a pointer alone cannot identify +the reference tiling flags: + +```text +NONE = xActiveMask must be null +TOKEN = xActiveMask points to [bs] +EXPERT = xActiveMask points to [bs, topK] +``` + +Token and expert mask modes are mutually exclusive, matching the supported +reference tiling cases. + +## Shared Expert Contract + +Shared expert ranks are the leading EP ranks, as in the reference code. +Validation requires: + +- `sharedExpertNum == 0` iff `sharedExpertRankNum == 0`; +- `sharedExpertRankNum % sharedExpertNum == 0` when shared experts exist; +- `sharedExpertRankNum < rankSize`; +- `moeExpertNum % (rankSize - sharedExpertRankNum) == 0`. + +The port retains `rankNumPerSharedExpert`, `idInSharedGroup`, the shared-expert +AIV allocation formula, and the reference destination-rank formula. + +## Output Compatibility + +`expandXOut`, `expertTokenNumsOut`, and `epRecvCountsOut` keep the reference +ordering and values. Reference `sendCountsOut` maps to TileXR +`epRecvCountsOut`. + +TileXR keeps its four-int `assistInfoForCombineOut` record: + +```text +[sourceRank, sourceTokenIndex, topKIndex, expertId] +``` + +The first three fields are the reference `expandIdxOut` triple. The fourth +field is the TileXR extension required by the current combine path. Shared +expert records use the same `topK + sharedExpertIndex` convention as the +reference code. + +## Host Validation + +The memory-dispatch host path: + +- rejects quantization and scale-related inputs; +- rejects TP; +- validates EP rank/world values against `CommArgs` when explicitly supplied; +- derives `globalBs` as `bs * rankSize` when the public API passes zero; +- validates the active-mask pointer/type pair; +- validates all reference UB and IPC-window size formulas; +- obtains the A5 vector-core count and launches that exact block dimension. + +## Testing + +Tests are added or updated before implementation to cover: + +- the reduced launch signature and removal of chunk/source-slot parameters; +- deletion of `tilexr_ep_dispatch_memory_helpers.h`; +- presence of the reference dispatch, count/cumsum, DataAsFlag, compaction, + shared-expert, token-mask, and expert-mask paths in the target kernel; +- active-mask pointer/type validation; +- general shared-expert rank grouping; +- rejection of quantization and TP; +- reference state/data/workspace layout calculations and overflow rejection; +- compatibility of the four-field assist tuple; +- host and source-guard unit suites; +- A5 kernel compilation when the CANN environment is available. diff --git a/src/ep/CMakeLists.txt b/src/ep/CMakeLists.txt index c2fc93d5..6dd240ad 100644 --- a/src/ep/CMakeLists.txt +++ b/src/ep/CMakeLists.txt @@ -1,5 +1,8 @@ include(GNUInstallDirs) +list(APPEND CMAKE_MODULE_PATH "${CMAKE_SOURCE_DIR}/src/collectives/kernels/cmake") +enable_language(CCE) + find_program(BISHENG_EXECUTABLE bisheng) if(NOT BISHENG_EXECUTABLE) message(FATAL_ERROR "bisheng not found; source scripts/common_env.sh before configuring with -DTILEXR_BUILD_EP=ON") @@ -39,9 +42,15 @@ set(TILEXR_EP_KERNEL_COMPILE_OPTIONS set(TILEXR_EP_KERNEL_LINK_OPTIONS ${TILEXR_EP_AICORE_ARCH}) set(TILEXR_EP_DISPATCH_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_dispatch_kernel.cpp") +set(TILEXR_EP_DISPATCH_MEMORY_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_dispatch_memory_kernel.cpp") set(TILEXR_EP_COMBINE_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_combine_kernel.cpp") +set(TILEXR_EP_COMBINE_MEMORY_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_combine_memory_kernel.cpp") set(TILEXR_EP_DISPATCH_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_dispatch_kernel.so") +set(TILEXR_EP_DISPATCH_MEMORY_OP "${CMAKE_CURRENT_BINARY_DIR}/tilexr_ep_dispatch_memory_kernel.o") +set(TILEXR_EP_DISPATCH_MEMORY_EMBED_CPP + "${CMAKE_CURRENT_BINARY_DIR}/tilexr_ep_dispatch_memory_kernel_embed.cpp") set(TILEXR_EP_COMBINE_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_combine_kernel.so") +set(TILEXR_EP_COMBINE_MEMORY_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_combine_memory_kernel.so") set(TILEXR_EP_KERNEL_INCLUDES -I${ASCEND_HOME_PATH}/compiler/tikcpp -I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw @@ -51,6 +60,7 @@ set(TILEXR_EP_KERNEL_INCLUDES -I${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp/tikcfw -I${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp/tikcfw/impl -I${ASCEND_HOME_PATH}/${ARCH}-linux/tikcpp/tikcfw/interface + -I${ASCEND_HOME_PATH}/${ARCH}-linux/asc/include -I${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc -I${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/runtime -I${ASCEND_HOME_PATH}/${ARCH}-linux/runtime/include @@ -59,6 +69,7 @@ set(TILEXR_EP_KERNEL_INCLUDES -I${CMAKE_SOURCE_DIR}/3rdparty -I${CMAKE_SOURCE_DIR}/src/include -I${CMAKE_CURRENT_SOURCE_DIR}/common + -I${CMAKE_CURRENT_SOURCE_DIR}/host ) add_custom_command( @@ -98,6 +109,50 @@ add_custom_command( ) add_custom_target(tilexr_ep_dispatch_kernel ALL DEPENDS "${TILEXR_EP_DISPATCH_KERNEL_SO}") +set_source_files_properties("${TILEXR_EP_DISPATCH_MEMORY_KERNEL_SOURCE}" PROPERTIES LANGUAGE CCE) +set_source_files_properties("${TILEXR_EP_DISPATCH_MEMORY_KERNEL_SOURCE}" PROPERTIES + OBJECT_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_mxfp8_quant.h;${CMAKE_SOURCE_DIR}/src/include/comm_args.h;${CMAKE_SOURCE_DIR}/src/include/tilexr_types.h") +add_library(tilexr_ep_dispatch_memory_kernel_tmp OBJECT + "${TILEXR_EP_DISPATCH_MEMORY_KERNEL_SOURCE}") +target_compile_options(tilexr_ep_dispatch_memory_kernel_tmp PRIVATE + -O2 + -std=gnu++17 + --cce-aicore-only + -Wno-deprecated-declarations + "SHELL:-mllvm -cce-aicore-stack-size=0x8000" + "SHELL:-mllvm -cce-aicore-function-stack-size=0x8000" + "SHELL:-mllvm -cce-aicore-record-overflow=true" + "SHELL:-mllvm -cce-aicore-addr-transform" + "SHELL:-mllvm -cce-aicore-dcci-insert-for-scalar=false" + ${TILEXR_EP_AICORE_ARCH} + ${TILEXR_EP_KERNEL_INCLUDES}) +target_compile_definitions(tilexr_ep_dispatch_memory_kernel_tmp PRIVATE + CATLASS_ARCH=${TILEXR_EP_CATLASS_ARCH}) + +set(TILEXR_EP_DISPATCH_MEMORY_TMP_OBJECT + "${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/tilexr_ep_dispatch_memory_kernel_tmp.dir/kernels/tilexr_ep_dispatch_memory_kernel.cpp.o") +add_custom_command( + OUTPUT "${TILEXR_EP_DISPATCH_MEMORY_OP}" + COMMAND ${CMAKE_CCE_LINKER} -m aicorelinux -Ttext=0 + "${TILEXR_EP_DISPATCH_MEMORY_TMP_OBJECT}" + --static -o "${TILEXR_EP_DISPATCH_MEMORY_OP}" --allow-multiple-definition + COMMAND truncate -c -s 10485760 "${TILEXR_EP_DISPATCH_MEMORY_OP}" + DEPENDS tilexr_ep_dispatch_memory_kernel_tmp "${TILEXR_EP_DISPATCH_MEMORY_TMP_OBJECT}" + VERBATIM + COMMENT "Linking TileXR EP memory dispatch AICore binary") + +add_custom_command( + OUTPUT "${TILEXR_EP_DISPATCH_MEMORY_EMBED_CPP}" + COMMAND ${CMAKE_COMMAND} + -DTILEXR_EP_DISPATCH_MEMORY_OP=${TILEXR_EP_DISPATCH_MEMORY_OP} + -DTILEXR_EP_DISPATCH_MEMORY_EMBED_CPP=${TILEXR_EP_DISPATCH_MEMORY_EMBED_CPP} + -P "${CMAKE_CURRENT_SOURCE_DIR}/cmake/embed_dispatch_memory_kernel.cmake" + DEPENDS "${TILEXR_EP_DISPATCH_MEMORY_OP}" + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/embed_dispatch_memory_kernel.cmake" + VERBATIM) +add_custom_target(tilexr_ep_dispatch_memory_kernel ALL + DEPENDS "${TILEXR_EP_DISPATCH_MEMORY_OP}" "${TILEXR_EP_DISPATCH_MEMORY_EMBED_CPP}") + add_custom_command( OUTPUT "${TILEXR_EP_COMBINE_KERNEL_SO}" COMMAND ${BISHENG_EXECUTABLE} @@ -135,15 +190,58 @@ add_custom_command( ) add_custom_target(tilexr_ep_combine_kernel ALL DEPENDS "${TILEXR_EP_COMBINE_KERNEL_SO}") +add_custom_command( + OUTPUT "${TILEXR_EP_COMBINE_MEMORY_KERNEL_SO}" + COMMAND ${BISHENG_EXECUTABLE} + ${TILEXR_EP_KERNEL_COMPILE_OPTIONS} + -std=gnu++17 + -fPIC + -shared + ${TILEXR_EP_KERNEL_LINK_OPTIONS} + -DCATLASS_ARCH=${TILEXR_EP_CATLASS_ARCH} + ${TILEXR_EP_KERNEL_INCLUDES} + "${TILEXR_EP_COMBINE_MEMORY_KERNEL_SOURCE}" + -L${ASCEND_DRIVER_PATH}/lib64/driver + -L${ASCEND_HOME_PATH}/${ARCH}-linux/lib64 + -L${ASCEND_HOME_PATH}/${ARCH}-linux/devlib + -lruntime + -lascendcl + -lstdc++ + -lm + -ltiling_api + -lplatform + -lc_sec + -ldl + -lnnopbase + -lpthread + -o "${TILEXR_EP_COMBINE_MEMORY_KERNEL_SO}" + DEPENDS + "${TILEXR_EP_COMBINE_MEMORY_KERNEL_SOURCE}" + "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_mxfp8_quant.h" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_window.h" + "${CMAKE_SOURCE_DIR}/src/include/comm_args.h" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_data_as_flag.h" + VERBATIM + COMMENT "Building TileXR EP memory combine kernel with bisheng" +) +add_custom_target(tilexr_ep_combine_memory_kernel ALL DEPENDS "${TILEXR_EP_COMBINE_MEMORY_KERNEL_SO}") + add_library(tilexr-ep SHARED host/ep_layout.cpp + host/ep_memory_layout.cpp host/ep_dispatch_host.cpp + host/ep_dispatch_memory_host.cpp + host/ep_combine_memory_host.cpp host/ep_launch_context.cpp host/ep_kernel_launch.cpp host/tilexr_ep_dispatch.cpp + host/tilexr_ep_memory_dispatch.cpp + host/tilexr_ep_memory_combine.cpp + "${TILEXR_EP_DISPATCH_MEMORY_EMBED_CPP}" ) -add_dependencies(tilexr-ep tilexr_ep_dispatch_kernel tilexr_ep_combine_kernel) +add_dependencies(tilexr-ep tilexr_ep_dispatch_kernel tilexr_ep_dispatch_memory_kernel tilexr_ep_combine_kernel + tilexr_ep_combine_memory_kernel) target_include_directories(tilexr-ep PUBLIC @@ -165,6 +263,7 @@ target_link_libraries(tilexr-ep PRIVATE -l:libtilexr_ep_dispatch_kernel.so -l:libtilexr_ep_combine_kernel.so + -l:libtilexr_ep_combine_memory_kernel.so tile-comm ascendcl runtime @@ -180,6 +279,7 @@ install(TARGETS tilexr-ep LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}) install(FILES "${TILEXR_EP_DISPATCH_KERNEL_SO}" "${TILEXR_EP_COMBINE_KERNEL_SO}" + "${TILEXR_EP_COMBINE_MEMORY_KERNEL_SO}" DESTINATION ${CMAKE_INSTALL_LIBDIR} ) install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_ep.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) diff --git a/src/ep/cmake/embed_dispatch_memory_kernel.cmake b/src/ep/cmake/embed_dispatch_memory_kernel.cmake new file mode 100644 index 00000000..35bb475f --- /dev/null +++ b/src/ep/cmake/embed_dispatch_memory_kernel.cmake @@ -0,0 +1,34 @@ +if(NOT DEFINED TILEXR_EP_DISPATCH_MEMORY_OP) + message(FATAL_ERROR "TILEXR_EP_DISPATCH_MEMORY_OP is required") +endif() +if(NOT DEFINED TILEXR_EP_DISPATCH_MEMORY_EMBED_CPP) + message(FATAL_ERROR "TILEXR_EP_DISPATCH_MEMORY_EMBED_CPP is required") +endif() +if(NOT EXISTS "${TILEXR_EP_DISPATCH_MEMORY_OP}") + message(FATAL_ERROR "Dispatch memory CCE object not found: ${TILEXR_EP_DISPATCH_MEMORY_OP}") +endif() + +file(SIZE "${TILEXR_EP_DISPATCH_MEMORY_OP}" _tilexr_ep_dispatch_memory_size) + +file(WRITE "${TILEXR_EP_DISPATCH_MEMORY_EMBED_CPP}" "/* + * Generated by src/ep/cmake/embed_dispatch_memory_kernel.cmake. + */ +#include +#include + +extern \"C\" { +extern const std::size_t TileXREpDispatchMemoryKernelBinarySize = ${_tilexr_ep_dispatch_memory_size}; +} + +asm(R\"( +.section .rodata, \"a\", @progbits +.balign 16 +.global TileXREpDispatchMemoryKernelBinaryData +.type TileXREpDispatchMemoryKernelBinaryData, @object +TileXREpDispatchMemoryKernelBinaryData: +.incbin \"${TILEXR_EP_DISPATCH_MEMORY_OP}\" +TileXREpDispatchMemoryKernelBinaryDataEnd: +.size TileXREpDispatchMemoryKernelBinaryData, TileXREpDispatchMemoryKernelBinaryDataEnd - TileXREpDispatchMemoryKernelBinaryData +.previous +)\"); +") diff --git a/src/ep/host/ep_combine_memory_host.cpp b/src/ep/host/ep_combine_memory_host.cpp new file mode 100644 index 00000000..95b22af3 --- /dev/null +++ b/src/ep/host/ep_combine_memory_host.cpp @@ -0,0 +1,103 @@ +#include "ep_dispatch_host.h" + +#include +#include + +#include "comm_args.h" +#include "ep_memory_layout.h" +#include "tilexr_types.h" + +namespace TileXREp { +namespace { + +bool ResolveGlobalBs(const EpCombineParams ¶ms, const TileXR::CommArgs &commArgs, int64_t *globalBs) +{ + if (globalBs == nullptr || params.bs <= 0 || commArgs.rankSize <= 0 || + params.bs > std::numeric_limits::max() / commArgs.rankSize) { + return false; + } + const int64_t expected = params.bs * commArgs.rankSize; + if (params.globalBs != 0 && params.globalBs != expected) { + return false; + } + *globalBs = expected; + return true; +} + +int ValidateMask(const EpCombineParams ¶ms) +{ + if (params.activeMaskType == TILEXR_EP_ACTIVE_MASK_NONE) { + return params.xActiveMask == nullptr ? TileXR::TILEXR_SUCCESS : TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (params.activeMaskType == TILEXR_EP_ACTIVE_MASK_TOKEN) { + return params.xActiveMask != nullptr ? TileXR::TILEXR_SUCCESS : TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (params.activeMaskType == TILEXR_EP_ACTIVE_MASK_EXPERT) { + return TileXR::TILEXR_ERROR_NOT_SUPPORT; + } + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; +} + +} // namespace + +int TileXREpValidateCombineMemoryConfig(const EpCombineParams ¶ms, const TileXR::CommArgs &commArgs, + uint32_t blockDim, EpMemoryCombineReferenceConfig *config) +{ + if (config == nullptr || commArgs.rankSize <= 0 || commArgs.rankSize > TileXR::TILEXR_MAX_RANK_SIZE || + commArgs.rank < 0 || commArgs.rank >= commArgs.rankSize) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + for (int rank = 0; rank < commArgs.rankSize; ++rank) { + if (commArgs.peerMems[rank] == nullptr) { + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + } + + const int64_t effectiveTpWorldSize = params.tpWorldSize == 0 ? 1 : params.tpWorldSize; + const bool useMxfp8 = params.quantMode == 3 || params.quantMode == 4; + if (effectiveTpWorldSize != 1 || params.tpRankId != 0 || + (params.quantMode != 0 && !useMxfp8)) { + return TileXR::TILEXR_ERROR_NOT_SUPPORT; + } + if (useMxfp8 && params.expertScales == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (params.expertShardType != 0 || + (params.epWorldSize != 0 && params.epWorldSize != commArgs.rankSize) || + (params.epRankId != 0 && params.epRankId != commArgs.rank)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + const int maskRet = ValidateMask(params); + if (maskRet != TileXR::TILEXR_SUCCESS) { + return maskRet; + } + + int64_t globalBs = 0; + if (!ResolveGlobalBs(params, commArgs, &globalBs)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + return TileXREpBuildMemoryCombineReferenceConfig(commArgs.rankSize, commArgs.rank, params.bs, params.h, + params.topK, params.moeExpertNum, params.sharedExpertNum, params.sharedExpertRankNum, globalBs, + params.dtype, params.quantMode, blockDim, config); +} + +int TileXREpPrepareMemoryCombineLaunchContext(const EpCombineParams ¶ms, EpHostLaunchContext *context) +{ + if (context == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + *context = EpHostLaunchContext {}; + int ret = TileXRGetCommArgsHost(params.comm, context->hostArgs); + if (ret != TileXR::TILEXR_SUCCESS || context->hostArgs == nullptr) { + *context = EpHostLaunchContext {}; + return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + } + ret = TileXRGetCommArgsDev(params.comm, context->devArgs); + if (ret != TileXR::TILEXR_SUCCESS || context->devArgs == nullptr) { + *context = EpHostLaunchContext {}; + return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + } + return TileXR::TILEXR_SUCCESS; +} + +} // namespace TileXREp diff --git a/src/ep/host/ep_dispatch_host.h b/src/ep/host/ep_dispatch_host.h index 903e1a79..9a693104 100644 --- a/src/ep/host/ep_dispatch_host.h +++ b/src/ep/host/ep_dispatch_host.h @@ -5,7 +5,9 @@ #include "acl/acl_base.h" #include "ep_layout.h" +#include "ep_memory_layout.h" #include "tilexr_api.h" +#include "tilexr_ep.h" namespace TileXREp { @@ -17,6 +19,7 @@ struct EpDispatchParams { int32_t *expertIds = nullptr; void *scales = nullptr; bool *xActiveMask = nullptr; + int64_t activeMaskType = TILEXR_EP_ACTIVE_MASK_NONE; void *expertScales = nullptr; TileXRCommPtr comm = nullptr; int64_t bs = 0; @@ -42,6 +45,7 @@ struct EpDispatchParams { void *expandScalesOut = nullptr; void *workspace = nullptr; TileXR::TileXRDataType dtype = TileXR::TILEXR_DATA_TYPE_RESERVED; + TileXR::TileXRDataType expandXOutDtype = TileXR::TILEXR_DATA_TYPE_RESERVED; aclrtStream stream = nullptr; }; @@ -49,11 +53,24 @@ struct EpCombineParams { void *expertOut = nullptr; int32_t *assistInfoForCombine = nullptr; int32_t *epRecvCounts = nullptr; + float *expertScales = nullptr; + bool *xActiveMask = nullptr; + int64_t activeMaskType = TILEXR_EP_ACTIVE_MASK_NONE; + void *sharedExpertX = nullptr; TileXRCommPtr comm = nullptr; int64_t bs = 0; int64_t h = 0; int64_t topK = 0; int64_t moeExpertNum = 0; + int64_t epWorldSize = 0; + int64_t epRankId = 0; + int64_t tpWorldSize = 0; + int64_t tpRankId = 0; + int64_t expertShardType = 0; + int64_t sharedExpertNum = 0; + int64_t sharedExpertRankNum = 0; + int64_t quantMode = 0; + int64_t globalBs = 0; void *yOut = nullptr; void *workspace = nullptr; TileXR::TileXRDataType dtype = TileXR::TILEXR_DATA_TYPE_RESERVED; @@ -69,13 +86,21 @@ struct EpHostLaunchContext { int TileXREpValidateBasicDispatchParams(const EpDispatchParams ¶ms); int TileXREpValidateDispatchConfig(const EpDispatchParams ¶ms, const TileXR::CommArgs &commArgs, EpWindowConfig *window); +int TileXREpValidateDispatchMemoryConfig(const EpDispatchParams ¶ms, const TileXR::CommArgs &commArgs, + EpWindowConfig *window); +int TileXREpValidateDispatchMemoryConfig(const EpDispatchParams ¶ms, const TileXR::CommArgs &commArgs, + uint32_t blockDim, EpMemoryDispatchReferenceConfig *config); int TileXREpValidateDispatchV2Config(const EpDispatchParams ¶ms, const TileXR::CommArgs &commArgs); int TileXREpPrepareLaunchContext(const EpDispatchParams ¶ms, EpHostLaunchContext *context); +int TileXREpPrepareMemoryLaunchContext(const EpDispatchParams ¶ms, EpHostLaunchContext *context); int TileXREpValidateBasicCombineParams(const EpCombineParams ¶ms); int TileXREpValidateCombineConfig(const EpCombineParams ¶ms, const TileXR::CommArgs &commArgs, EpWindowConfig *window); +int TileXREpValidateCombineMemoryConfig(const EpCombineParams ¶ms, const TileXR::CommArgs &commArgs, + uint32_t blockDim, EpMemoryCombineReferenceConfig *config); int TileXREpPrepareCombineLaunchContext(const EpCombineParams ¶ms, EpHostLaunchContext *context); +int TileXREpPrepareMemoryCombineLaunchContext(const EpCombineParams ¶ms, EpHostLaunchContext *context); } // namespace TileXREp diff --git a/src/ep/host/ep_dispatch_memory_host.cpp b/src/ep/host/ep_dispatch_memory_host.cpp new file mode 100644 index 00000000..830c7385 --- /dev/null +++ b/src/ep/host/ep_dispatch_memory_host.cpp @@ -0,0 +1,166 @@ +#include "ep_dispatch_host.h" + +#include +#include + +#include "comm_args.h" +#include "ep_memory_layout.h" +#include "tilexr_types.h" + +namespace TileXREp { +namespace { + +int64_t TileXREpEffectiveTpWorldSize(int64_t tpWorldSize) +{ + return tpWorldSize == 0 ? 1 : tpWorldSize; +} + +TileXR::TileXRDataType TileXREpMemoryExpandXOutDtype(const EpDispatchParams ¶ms) +{ + return params.expandXOutDtype == TileXR::TILEXR_DATA_TYPE_RESERVED ? params.dtype : params.expandXOutDtype; +} + +bool TileXREpMemoryGlobalBs(const EpDispatchParams ¶ms, const TileXR::CommArgs &commArgs, int64_t *globalBs) +{ + if (globalBs == nullptr || params.bs <= 0 || commArgs.rankSize <= 0 || + params.bs > std::numeric_limits::max() / commArgs.rankSize) { + return false; + } + const int64_t expected = params.bs * commArgs.rankSize; + if (params.globalBs != 0 && params.globalBs != expected) { + return false; + } + *globalBs = expected; + return true; +} + +int TileXREpValidateMemoryMask(const EpDispatchParams ¶ms) +{ + if (params.activeMaskType == TILEXR_EP_ACTIVE_MASK_NONE) { + return params.xActiveMask == nullptr ? TileXR::TILEXR_SUCCESS : TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (params.activeMaskType == TILEXR_EP_ACTIVE_MASK_TOKEN || + params.activeMaskType == TILEXR_EP_ACTIVE_MASK_EXPERT) { + return params.xActiveMask != nullptr ? TileXR::TILEXR_SUCCESS : TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; +} + +} // namespace + +int TileXREpValidateDispatchMemoryConfig(const EpDispatchParams ¶ms, const TileXR::CommArgs &commArgs, + uint32_t blockDim, EpMemoryDispatchReferenceConfig *config) +{ + if (config == nullptr || commArgs.rankSize <= 0 || commArgs.rankSize > TileXR::TILEXR_MAX_RANK_SIZE || + commArgs.rank < 0 || commArgs.rank >= commArgs.rankSize) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + for (int rank = 0; rank < commArgs.rankSize; ++rank) { + if (commArgs.peerMems[rank] == nullptr) { + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + } + + const int64_t effectiveTpWorldSize = TileXREpEffectiveTpWorldSize(params.tpWorldSize); + if (effectiveTpWorldSize != 1 || params.tpRankId != 0 || params.tpRecvCountsOut != nullptr || + params.scales != nullptr || params.expertScales != nullptr || params.expandScalesOut != nullptr) { + return TileXR::TILEXR_ERROR_NOT_SUPPORT; + } + const TileXR::TileXRDataType expandXOutDtype = TileXREpMemoryExpandXOutDtype(params); + if (params.quantMode == 0) { + if (params.dynamicScalesOut != nullptr || expandXOutDtype != params.dtype) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } else if (params.quantMode == 4) { + if (params.dynamicScalesOut == nullptr || + (expandXOutDtype != TileXR::TILEXR_DATA_TYPE_FP8E4M3 && + expandXOutDtype != TileXR::TILEXR_DATA_TYPE_FP8E5M2)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } else { + return TileXR::TILEXR_ERROR_NOT_SUPPORT; + } + if (params.expertShardType != 0 || params.expertTokenNumsType < kEpExpertTokenNumsTypePrefixSum || + params.expertTokenNumsType > kEpExpertTokenNumsTypeCount || + (params.epWorldSize != 0 && params.epWorldSize != commArgs.rankSize) || + (params.epRankId != 0 && params.epRankId != commArgs.rank)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + const int maskRet = TileXREpValidateMemoryMask(params); + if (maskRet != TileXR::TILEXR_SUCCESS) { + return maskRet; + } + int64_t globalBs = 0; + if (!TileXREpMemoryGlobalBs(params, commArgs, &globalBs)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + const int ret = TileXREpBuildMemoryDispatchReferenceConfig(commArgs.rankSize, commArgs.rank, params.bs, + params.h, params.topK, params.moeExpertNum, params.sharedExpertNum, params.sharedExpertRankNum, + globalBs, params.dtype, expandXOutDtype, params.quantMode, blockDim, config); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpValidateDispatchMemoryConfig(const EpDispatchParams ¶ms, const TileXR::CommArgs &commArgs, + EpWindowConfig *window) +{ + if (window == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + EpMemoryDispatchReferenceConfig memoryConfig {}; + const int ret = TileXREpValidateDispatchMemoryConfig( + params, commArgs, kEpMemoryDefaultVectorCoreNum, &memoryConfig); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + + EpWindowConfig compatible {}; + compatible.rankSize = commArgs.rankSize; + compatible.bs = params.bs; + compatible.h = params.h; + compatible.topK = params.topK; + compatible.moeExpertNum = params.moeExpertNum; + compatible.localExpertNum = memoryConfig.localExpertNum; + compatible.dtypeBytes = params.quantMode == 4 ? 1 : TileXREpDataTypeSize(params.dtype); + compatible.maxRoutesPerSrc = params.bs * (params.topK + params.sharedExpertNum); + compatible.rowBytes = params.h * compatible.dtypeBytes; + compatible.payloadRowBytes = compatible.rowBytes + memoryConfig.scaleOutBytes; + compatible.totalBytes = memoryConfig.totalWinSize + kEpMemoryStateWindowBytes + memoryConfig.workspaceBytes; + *window = compatible; + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpPrepareMemoryLaunchContext(const EpDispatchParams ¶ms, EpHostLaunchContext *context) +{ + if (context == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + *context = EpHostLaunchContext {}; + + int ret = TileXRGetCommArgsHost(params.comm, context->hostArgs); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + if (context->hostArgs == nullptr) { + *context = EpHostLaunchContext {}; + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + + ret = TileXRGetCommArgsDev(params.comm, context->devArgs); + if (ret != TileXR::TILEXR_SUCCESS) { + *context = EpHostLaunchContext {}; + return ret; + } + if (context->devArgs == nullptr) { + *context = EpHostLaunchContext {}; + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + + return TileXR::TILEXR_SUCCESS; +} + +} // namespace TileXREp diff --git a/src/ep/host/ep_kernel_launch.cpp b/src/ep/host/ep_kernel_launch.cpp index f7095626..c3b9e7cc 100644 --- a/src/ep/host/ep_kernel_launch.cpp +++ b/src/ep/host/ep_kernel_launch.cpp @@ -1,11 +1,21 @@ #include "ep_kernel_launch.h" +#include #include +#include +#include #include "acl/acl_rt.h" +#include "ep_memory_layout.h" #include "ep_window.h" #include "tilexr_api.h" #include "tilexr_types.h" +#include "runtime/kernel.h" + +extern "C" { +extern const unsigned char TileXREpDispatchMemoryKernelBinaryData[]; +extern const std::size_t TileXREpDispatchMemoryKernelBinarySize; +} extern void launch_tilexr_ep_dispatch_kernel(uint32_t blockDim, void *stream, GM_ADDR commArgs, GM_ADDR x, GM_ADDR expertIds, GM_ADDR scales, GM_ADDR xActiveMask, GM_ADDR expandXOut, GM_ADDR dynamicScalesOut, @@ -27,6 +37,12 @@ extern void launch_tilexr_ep_dispatch_cross_node_kernel(uint32_t blockDim, void int64_t expertTokenNumsType, int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t quantMode, int64_t tpWorldSize, int64_t tpRankId, int64_t magic); +extern void launch_tilexr_ep_combine_memory_kernel(uint32_t blockDim, void *stream, GM_ADDR commArgs, + GM_ADDR expertOut, GM_ADDR assistInfoForCombine, GM_ADDR sendCounts, GM_ADDR expertScales, + GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR yOut, int64_t bs, int64_t h, int64_t topK, + int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, + int64_t activeMaskType, int64_t quantMode, int64_t dtype); + extern void launch_tilexr_ep_combine_kernel(uint32_t blockDim, void *stream, GM_ADDR commArgs, GM_ADDR expertOut, GM_ADDR assistInfoForCombine, GM_ADDR epRecvCounts, GM_ADDR yOut, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, int64_t dtype, int64_t dtypeBytes, int64_t maxRoutesPerSrc, int64_t rowBytes, @@ -47,6 +63,95 @@ namespace TileXREp { namespace { +constexpr uint32_t kAivBinaryMagic = 0x41415246U; +constexpr char kDispatchMemoryKernelName[] = "tilexr_ep_dispatch_memory_kernel"; + +std::mutex gDispatchMemoryRegistrationMutex; +bool gDispatchMemoryRegistered = false; +int gDispatchMemoryRegistrationStatus = TileXR::TILEXR_ERROR_NOT_INITIALIZED; +void *gDispatchMemoryBinaryHandle = nullptr; +uint8_t gDispatchMemoryKernelStub = 0; + +struct DispatchMemoryKernelArgs { + GM_ADDR commArgs; + GM_ADDR x; + GM_ADDR expertIds; + GM_ADDR xActiveMask; + GM_ADDR expandXOut; + GM_ADDR dynamicScalesOut; + GM_ADDR expertTokenNumsOut; + GM_ADDR sendCountsOut; + GM_ADDR assistInfoForCombineOut; + int64_t bs; + int64_t h; + int64_t topK; + int64_t moeExpertNum; + int64_t sharedExpertNum; + int64_t sharedExpertRankNum; + int64_t globalBs; + int64_t expertTokenNumsType; + int64_t activeMaskType; + int64_t quantMode; + int64_t dtype; + int64_t expandXOutDtype; + int64_t magic; +}; + +static_assert(sizeof(DispatchMemoryKernelArgs) == 22U * sizeof(uint64_t), + "dispatch memory kernel argument ABI changed"); + +int EnsureDispatchMemoryKernelRegistered() +{ + std::lock_guard guard(gDispatchMemoryRegistrationMutex); + if (gDispatchMemoryRegistered) { + return gDispatchMemoryRegistrationStatus; + } + if (TileXREpDispatchMemoryKernelBinarySize == 0) { + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + + rtDevBinary_t binary {}; + binary.data = TileXREpDispatchMemoryKernelBinaryData; + binary.length = static_cast(TileXREpDispatchMemoryKernelBinarySize); + binary.magic = kAivBinaryMagic; + binary.version = 0; + + rtError_t rtRet = rtDevBinaryRegister(&binary, &gDispatchMemoryBinaryHandle); + if (rtRet == RT_ERROR_NONE) { + rtRet = rtFunctionRegister(gDispatchMemoryBinaryHandle, &gDispatchMemoryKernelStub, + kDispatchMemoryKernelName, kDispatchMemoryKernelName, 0); + } + if (rtRet != RT_ERROR_NONE) { + std::cerr << "TileXR EP dispatch memory kernel registration failed, ret=" << rtRet << std::endl; + gDispatchMemoryRegistrationStatus = TileXR::TILEXR_ERROR_MKIRT; + return gDispatchMemoryRegistrationStatus; + } + + gDispatchMemoryRegistered = true; + gDispatchMemoryRegistrationStatus = TileXR::TILEXR_SUCCESS; + return gDispatchMemoryRegistrationStatus; +} + +int LaunchDispatchMemoryKernel(uint32_t blockDim, aclrtStream stream, DispatchMemoryKernelArgs *args) +{ + const int registerRet = EnsureDispatchMemoryKernelRegistered(); + if (registerRet != TileXR::TILEXR_SUCCESS) { + return registerRet; + } + + rtArgsEx_t argsInfo {}; + argsInfo.args = args; + argsInfo.argsSize = sizeof(*args); + rtTaskCfgInfo_t cfgInfo {}; + const rtError_t rtRet = rtKernelLaunchWithFlagV2(&gDispatchMemoryKernelStub, blockDim, &argsInfo, nullptr, + static_cast(stream), 0, &cfgInfo); + if (rtRet != RT_ERROR_NONE) { + std::cerr << "TileXR EP dispatch memory kernel launch failed, ret=" << rtRet << std::endl; + return TileXR::TILEXR_ERROR_MKIRT; + } + return TileXR::TILEXR_SUCCESS; +} + bool TileXREpUsesCrossNodeKernel(const EpHostLaunchContext &context) { return context.hostArgs != nullptr && context.hostArgs->localRankSize > 0 && @@ -103,6 +208,76 @@ int TileXREpLaunchDispatchKernel(const EpDispatchParams ¶ms, const EpHostLau return TileXR::TILEXR_SUCCESS; } +int TileXREpLaunchDispatchMemoryKernel(const EpDispatchParams ¶ms, const EpHostLaunchContext &context) +{ + int32_t deviceId = 0; + if (aclrtGetDevice(&deviceId) != ACL_SUCCESS) { + return TileXR::TILEXR_ERROR_INTERNAL; + } + int64_t vectorCoreNum = 0; + if (aclrtGetDeviceInfo(static_cast(deviceId), ACL_DEV_ATTR_VECTOR_CORE_NUM, &vectorCoreNum) != + ACL_SUCCESS || + vectorCoreNum < 2 || vectorCoreNum > static_cast(UINT32_MAX)) { + return TileXR::TILEXR_ERROR_INTERNAL; + } + + EpMemoryDispatchReferenceConfig memoryConfig {}; + const uint32_t blockDim = static_cast(vectorCoreNum); + const int configRet = TileXREpValidateDispatchMemoryConfig( + params, *context.hostArgs, blockDim, &memoryConfig); + if (configRet != TileXR::TILEXR_SUCCESS) { + return configRet; + } + + const int64_t globalBs = params.globalBs == 0 ? params.bs * context.hostArgs->rankSize : params.globalBs; + int64_t magic = 0; + const int magicRet = TileXRCommNextMagic(params.comm, &magic); + if (magicRet != TileXR::TILEXR_SUCCESS) { + return magicRet; + } + + DispatchMemoryKernelArgs args { context.devArgs, static_cast(params.x), + reinterpret_cast(params.expertIds), reinterpret_cast(params.xActiveMask), + static_cast(params.expandXOut), static_cast(params.dynamicScalesOut), + reinterpret_cast(params.expertTokenNumsOut), reinterpret_cast(params.epRecvCountsOut), + reinterpret_cast(params.assistInfoForCombineOut), params.bs, params.h, params.topK, + params.moeExpertNum, params.sharedExpertNum, params.sharedExpertRankNum, globalBs, + params.expertTokenNumsType, params.activeMaskType, params.quantMode, static_cast(params.dtype), + static_cast(params.expandXOutDtype == TileXR::TILEXR_DATA_TYPE_RESERVED ? + params.dtype : params.expandXOutDtype), magic }; + return LaunchDispatchMemoryKernel(blockDim, params.stream, &args); +} + +int TileXREpLaunchCombineMemoryKernel(const EpCombineParams ¶ms, const EpHostLaunchContext &context) +{ + int32_t deviceId = 0; + if (aclrtGetDevice(&deviceId) != ACL_SUCCESS) { + return TileXR::TILEXR_ERROR_INTERNAL; + } + int64_t vectorCoreNum = 0; + if (aclrtGetDeviceInfo(static_cast(deviceId), ACL_DEV_ATTR_VECTOR_CORE_NUM, &vectorCoreNum) != + ACL_SUCCESS || + vectorCoreNum <= 0 || vectorCoreNum > static_cast(UINT32_MAX)) { + return TileXR::TILEXR_ERROR_INTERNAL; + } + const uint32_t blockDim = static_cast(vectorCoreNum); + EpMemoryCombineReferenceConfig memoryConfig {}; + const int configRet = TileXREpValidateCombineMemoryConfig( + params, *context.hostArgs, blockDim, &memoryConfig); + if (configRet != TileXR::TILEXR_SUCCESS) { + return configRet; + } + const int64_t globalBs = params.globalBs == 0 ? params.bs * context.hostArgs->rankSize : params.globalBs; + launch_tilexr_ep_combine_memory_kernel(blockDim, params.stream, context.devArgs, + static_cast(params.expertOut), reinterpret_cast(params.assistInfoForCombine), + reinterpret_cast(params.epRecvCounts), reinterpret_cast(params.expertScales), + reinterpret_cast(params.xActiveMask), static_cast(params.sharedExpertX), + static_cast(params.yOut), params.bs, params.h, params.topK, params.moeExpertNum, + params.sharedExpertNum, params.sharedExpertRankNum, globalBs, params.activeMaskType, + params.quantMode, static_cast(params.dtype)); + return TileXR::TILEXR_SUCCESS; +} + int TileXREpLaunchCombineKernel(const EpCombineParams ¶ms, const EpHostLaunchContext &context) { int64_t magic = 0; diff --git a/src/ep/host/ep_kernel_launch.h b/src/ep/host/ep_kernel_launch.h index 93c3b826..315b16ec 100644 --- a/src/ep/host/ep_kernel_launch.h +++ b/src/ep/host/ep_kernel_launch.h @@ -6,7 +6,9 @@ namespace TileXREp { int TileXREpLaunchDispatchKernel(const EpDispatchParams ¶ms, const EpHostLaunchContext &context); +int TileXREpLaunchDispatchMemoryKernel(const EpDispatchParams ¶ms, const EpHostLaunchContext &context); int TileXREpLaunchCombineKernel(const EpCombineParams ¶ms, const EpHostLaunchContext &context); +int TileXREpLaunchCombineMemoryKernel(const EpCombineParams ¶ms, const EpHostLaunchContext &context); } // namespace TileXREp diff --git a/src/ep/host/ep_memory_layout.cpp b/src/ep/host/ep_memory_layout.cpp new file mode 100644 index 00000000..e884fb91 --- /dev/null +++ b/src/ep/host/ep_memory_layout.cpp @@ -0,0 +1,463 @@ +#include "ep_memory_layout.h" + +#include +#include + +#include "comm_args.h" +#include "ep_layout.h" + +namespace TileXREp { +namespace { + +bool MulInt64(int64_t lhs, int64_t rhs, int64_t *out) +{ + if (out == nullptr || lhs < 0 || rhs < 0 || (lhs != 0 && rhs > std::numeric_limits::max() / lhs)) { + return false; + } + *out = lhs * rhs; + return true; +} + +bool AddInt64(int64_t lhs, int64_t rhs, int64_t *out) +{ + if (out == nullptr || lhs < 0 || rhs < 0 || rhs > std::numeric_limits::max() - lhs) { + return false; + } + *out = lhs + rhs; + return true; +} + +bool AlignUpInt64(int64_t value, int64_t alignment, int64_t *out) +{ + if (out == nullptr) { + return false; + } + const int64_t aligned = TileXREpAlignUp(value, alignment); + if (aligned == TileXR::TILEXR_INVALID_VALUE) { + return false; + } + *out = aligned; + return true; +} + +bool CeilDivInt64(int64_t value, int64_t divisor, int64_t *out) +{ + if (out == nullptr || value < 0 || divisor <= 0) { + return false; + } + *out = value / divisor + (value % divisor == 0 ? 0 : 1); + return true; +} + +} // namespace + +uint32_t TileXREpMemoryCountCoreNum(int64_t totalExpertNum, int64_t rscvStatusNum, uint32_t blockDim) +{ + if (totalExpertNum <= 0 || rscvStatusNum <= 0 || blockDim < 2) { + return 0; + } + uint64_t count = static_cast(totalExpertNum / 16); + count = std::max(count, 1); + count = std::min(count, blockDim / 2); + count = std::min(count, kEpMemoryMaxCountCoreNum); + count = std::min(count, static_cast(rscvStatusNum)); + return static_cast(count); +} + +int TileXREpBuildMemoryDispatchReferenceConfig(int64_t rankSize, int64_t rank, int64_t bs, int64_t h, + int64_t topK, int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, TileXR::TileXRDataType dtype, TileXR::TileXRDataType expandXOutDtype, + int64_t quantMode, uint32_t blockDim, + EpMemoryDispatchReferenceConfig *out) +{ + const int64_t moeRankNum = rankSize - sharedExpertRankNum; + const bool useMxfp8 = quantMode == 4; + const bool validOutputType = useMxfp8 ? + (expandXOutDtype == TileXR::TILEXR_DATA_TYPE_FP8E4M3 || + expandXOutDtype == TileXR::TILEXR_DATA_TYPE_FP8E5M2) : expandXOutDtype == dtype; + if (out == nullptr || rankSize <= 0 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || rank < 0 || + rank >= rankSize || bs <= 0 || h <= 0 || topK <= 0 || moeExpertNum <= 0 || sharedExpertNum < 0 || + sharedExpertRankNum < 0 || globalBs <= 0 || globalBs % rankSize != 0 || blockDim < 2 || + blockDim > kEpMemoryMaxVectorCoreNum || + (dtype != TileXR::TILEXR_DATA_TYPE_FP16 && dtype != TileXR::TILEXR_DATA_TYPE_BFP16) || + (quantMode != 0 && quantMode != 4) || !validOutputType || + ((sharedExpertNum == 0) != (sharedExpertRankNum == 0)) || sharedExpertRankNum >= rankSize || + (sharedExpertNum > 0 && sharedExpertRankNum % sharedExpertNum != 0) || moeRankNum <= 0 || + moeExpertNum % moeRankNum != 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + EpMemoryDispatchReferenceConfig next {}; + next.rankSize = rankSize; + next.rank = rank; + next.blockDim = blockDim; + const int64_t moeExpertNumPerRank = moeExpertNum / moeRankNum; + const bool isSharedExpertRank = rank < sharedExpertRankNum; + next.localExpertNum = isSharedExpertRank ? 1 : moeExpertNumPerRank; + next.rscvStatusNum = isSharedExpertRank ? rankSize : rankSize * moeExpertNumPerRank; + + int64_t totalExpertNum = 0; + if (!AddInt64(sharedExpertRankNum, moeExpertNum, &totalExpertNum)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.aivUsedCumSum = TileXREpMemoryCountCoreNum(totalExpertNum, next.rscvStatusNum, blockDim); + if (next.aivUsedCumSum == 0 || next.aivUsedCumSum >= blockDim) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.aivUsedAllToAll = blockDim - next.aivUsedCumSum; + if (sharedExpertRankNum > 0) { + const int64_t routeKinds = topK + sharedExpertNum; + next.sharedUsedAivNum = static_cast( + (static_cast(next.aivUsedAllToAll) * static_cast(sharedExpertNum)) / + static_cast(routeKinds)); + if (next.sharedUsedAivNum == 0) { + next.sharedUsedAivNum = 1; + } + } + if (next.sharedUsedAivNum >= next.aivUsedAllToAll) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.moeUsedAivNum = next.aivUsedAllToAll - next.sharedUsedAivNum; + + const int64_t dtypeBytes = TileXREpDataTypeSize(dtype); + const int64_t outputDtypeBytes = useMxfp8 ? 1 : dtypeBytes; + int64_t inputBytes = 0; + int64_t hOutBytes = 0; + int64_t tokenPayloadBytes = 0; + int64_t payloadAndScaleBytes = 0; + int64_t hOutSizeAlign = 0; + int64_t blockCntPerToken = 0; + int64_t scaleBlockCount = 0; + if (!MulInt64(h, dtypeBytes, &inputBytes) || + !MulInt64(h, outputDtypeBytes, &hOutBytes) || + !AlignUpInt64(hOutBytes, useMxfp8 ? 256 : kEpMemoryWindowAlignmentBytes, &tokenPayloadBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (useMxfp8 && (!CeilDivInt64(h, 32, &scaleBlockCount) || + !AlignUpInt64(scaleBlockCount, 2, &next.scaleOutBytes))) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (!AddInt64(tokenPayloadBytes, next.scaleOutBytes, &payloadAndScaleBytes) || + !AlignUpInt64(payloadAndScaleBytes, kEpMemoryWindowAlignmentBytes, &next.tokenQuantAlignBytes) || + !AddInt64(next.tokenQuantAlignBytes, kEpMemoryWindowAlignmentBytes, &hOutSizeAlign) || + !CeilDivInt64(hOutSizeAlign, kEpMemorySplitPayloadBytes, &blockCntPerToken) || + !MulInt64(blockCntPerToken, kEpMemorySplitBlockBytes, &next.hCommuSize)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.hOutSize = hOutBytes; + + const int64_t axisMaxBs = globalBs / rankSize; + if (!MulInt64(axisMaxBs, next.hCommuSize, &next.expertPerSizeOnWin)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t combinePayloadBytes = 0; + int64_t combineBlockCount = 0; + int64_t combineRowBytes = 0; + int64_t routesPerToken = 0; + int64_t combineRows = 0; + if (!MulInt64(h, dtypeBytes, &combinePayloadBytes) || + !CeilDivInt64(combinePayloadBytes, kEpMemorySplitPayloadBytes, &combineBlockCount) || + !MulInt64(combineBlockCount, kEpMemorySplitBlockBytes, &combineRowBytes) || + !AddInt64(topK, sharedExpertNum, &routesPerToken) || + !MulInt64(axisMaxBs, routesPerToken, &combineRows) || + !MulInt64(combineRows, combineRowBytes, &next.combineReserveBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t workspaceStatusNum = 0; + int64_t workspaceRawBytes = 0; + if (!MulInt64(rankSize, moeExpertNumPerRank, &workspaceStatusNum) || + !MulInt64(static_cast(blockDim), workspaceStatusNum, &workspaceRawBytes) || + !MulInt64(workspaceRawBytes, static_cast(sizeof(int32_t)), &workspaceRawBytes) || + !AlignUpInt64(workspaceRawBytes, kEpMemoryWindowAlignmentBytes, &next.workspaceBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.totalWinSize = TileXR::IPC_BUFF_MAX_SIZE - kEpMemoryStateWindowBytes - next.workspaceBytes; + if (next.totalWinSize <= 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.dispatchHalfBytes = next.totalWinSize / 2; + + int64_t dispatchDataBytes = 0; + int64_t requiredHalfBytes = 0; + if (!MulInt64(next.rscvStatusNum, next.expertPerSizeOnWin, &dispatchDataBytes) || + !AddInt64(next.combineReserveBytes, dispatchDataBytes, &requiredHalfBytes) || + requiredHalfBytes > next.dispatchHalfBytes) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t expertIdsBytes = 0; + int64_t expertIdsAligned = 0; + int64_t expertMaskRowBytes = 0; + int64_t expertMaskBytes = 0; + int64_t expertMaskHalfBytes = 0; + if (!MulInt64(bs, topK, &expertIdsBytes) || + !MulInt64(expertIdsBytes, static_cast(sizeof(int32_t)), &expertIdsBytes) || + !AlignUpInt64(expertIdsBytes, kEpMemoryWindowAlignmentBytes, &expertIdsAligned) || + !AlignUpInt64(topK * static_cast(sizeof(bool)), kEpMemoryWindowAlignmentBytes, + &expertMaskRowBytes) || + !MulInt64(bs, expertMaskRowBytes, &expertMaskBytes) || + !MulInt64(expertMaskBytes, static_cast(sizeof(uint16_t)), &expertMaskHalfBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.maxSizeForUbBuffer = std::max(expertIdsAligned, expertMaskHalfBytes); + if (useMxfp8) { + int64_t alignedScaleCount = 0; + int64_t quantWorkBytes = 0; + if (!AlignUpInt64(next.scaleOutBytes, 32, &alignedScaleCount) || + !MulInt64(alignedScaleCount, static_cast(sizeof(float)), &quantWorkBytes) || + !AddInt64(quantWorkBytes, next.scaleOutBytes * static_cast(sizeof(uint16_t)), + &quantWorkBytes) || + !AlignUpInt64(quantWorkBytes, kEpMemoryWindowAlignmentBytes, &quantWorkBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.maxSizeForUbBuffer = std::max(next.maxSizeForUbBuffer, quantWorkBytes); + } + next.totalUbSize = kEpMemoryFullUbBytes; + if (next.maxSizeForUbBuffer <= 0 || next.maxSizeForUbBuffer >= next.totalUbSize) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t hAlignBytes = 0; + int64_t expertIdsVecBytes = 0; + int64_t allToAllUbBytes = 2 * kEpMemoryWindowAlignmentBytes; + int64_t term = 0; + if (!AlignUpInt64(inputBytes, useMxfp8 ? 128 : kEpMemoryWindowAlignmentBytes, &hAlignBytes) || + !AlignUpInt64(expertIdsBytes, 256, &expertIdsVecBytes) || + !MulInt64(4, hAlignBytes, &term) || !AddInt64(allToAllUbBytes, term, &allToAllUbBytes) || + !MulInt64(2, expertIdsVecBytes, &term) || !AddInt64(allToAllUbBytes, term, &allToAllUbBytes) || + !MulInt64(2, next.maxSizeForUbBuffer, &term) || !AddInt64(allToAllUbBytes, term, &allToAllUbBytes) || + !AddInt64(allToAllUbBytes, expertIdsAligned, &allToAllUbBytes) || + (useMxfp8 && !AddInt64(allToAllUbBytes, hOutSizeAlign, &allToAllUbBytes)) || + !AddInt64(allToAllUbBytes, next.hCommuSize, &allToAllUbBytes) || + allToAllUbBytes > next.totalUbSize) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + const int64_t recvStatusPerCore = next.rscvStatusNum / next.aivUsedCumSum + + (next.rscvStatusNum % next.aivUsedCumSum == 0 ? 0 : 1); + int64_t waitStatusBytes = 0; + int64_t recvStatusScalarBytes = 0; + int64_t sumContinueBytes = 0; + int64_t tokenNumBytes = 0; + int64_t workLocalBytes = 0; + int64_t statusCountAligned = 0; + int64_t countUbBytes = kEpMemoryWindowAlignmentBytes; + if (!MulInt64(recvStatusPerCore, kEpMemoryWindowAlignmentBytes, &term) || + !AlignUpInt64(term, 256, &waitStatusBytes) || + !MulInt64(recvStatusPerCore, static_cast(sizeof(float)), &term) || + !AlignUpInt64(term, kEpMemoryWindowAlignmentBytes, &recvStatusScalarBytes) || + !MulInt64(blockDim, static_cast(sizeof(float)), &term) || + !AlignUpInt64(term, kEpMemoryWindowAlignmentBytes, &sumContinueBytes) || + !MulInt64(moeExpertNumPerRank, static_cast(sizeof(int64_t)), &term) || + !AlignUpInt64(term, kEpMemoryWindowAlignmentBytes, &tokenNumBytes) || + !MulInt64(rankSize, static_cast(sizeof(float)), &term) || + !AlignUpInt64(term, kEpMemoryWindowAlignmentBytes, &workLocalBytes) || + !AlignUpInt64(totalExpertNum, 8, &statusCountAligned) || + !MulInt64(2, next.maxSizeForUbBuffer, &term) || !AddInt64(countUbBytes, term, &countUbBytes) || + !MulInt64(2, expertIdsVecBytes, &term) || !AddInt64(countUbBytes, term, &countUbBytes) || + !MulInt64(statusCountAligned, kEpMemoryWindowAlignmentBytes, &term) || + !AddInt64(countUbBytes, term, &countUbBytes) || + !AddInt64(countUbBytes, waitStatusBytes, &countUbBytes) || + !AddInt64(countUbBytes, recvStatusScalarBytes, &countUbBytes) || + !MulInt64(recvStatusPerCore, kEpMemoryWindowAlignmentBytes, &term) || + !AddInt64(countUbBytes, term, &countUbBytes) || + !MulInt64(2 * static_cast(blockDim), kEpMemoryWindowAlignmentBytes, &term) || + !AddInt64(countUbBytes, term, &countUbBytes) || + !AddInt64(countUbBytes, sumContinueBytes, &countUbBytes) || + !AddInt64(countUbBytes, 3 * kEpMemoryWindowAlignmentBytes, &countUbBytes) || + !MulInt64(next.rscvStatusNum, kEpMemoryWindowAlignmentBytes, &term) || + !AddInt64(countUbBytes, term, &countUbBytes) || + !AddInt64(countUbBytes, tokenNumBytes, &countUbBytes) || + !AddInt64(countUbBytes, workLocalBytes, &countUbBytes) || countUbBytes > next.totalUbSize) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t receiveCountBytes = 0; + int64_t sourceInfoBytes = 0; + int64_t localCopyFixedBytes = kEpMemoryWindowAlignmentBytes; + const int64_t statesPerAiv = next.rscvStatusNum / blockDim + + (next.rscvStatusNum % blockDim == 0 ? 0 : 1); + const int64_t blockCountPerToken = next.hCommuSize / kEpMemorySplitBlockBytes; + if (!MulInt64(next.rscvStatusNum, static_cast(sizeof(int32_t)), &term) || + !AlignUpInt64(term, kEpMemoryWindowAlignmentBytes, &receiveCountBytes) || + !MulInt64(statesPerAiv, static_cast(sizeof(uint32_t)), &term) || + !AlignUpInt64(term, kEpMemoryWindowAlignmentBytes, &sourceInfoBytes) || + !AddInt64(localCopyFixedBytes, receiveCountBytes, &localCopyFixedBytes) || + !MulInt64(2 * static_cast(next.aivUsedCumSum), kEpMemoryWindowAlignmentBytes, &term) || + !AddInt64(localCopyFixedBytes, term, &localCopyFixedBytes) || + !MulInt64(3, sourceInfoBytes, &term) || !AddInt64(localCopyFixedBytes, term, &localCopyFixedBytes) || + !AddInt64(localCopyFixedBytes, 2 * kEpMemoryWindowAlignmentBytes, &localCopyFixedBytes) || + !MulInt64(blockCountPerToken, kEpMemoryWindowAlignmentBytes, &term) || + !AddInt64(localCopyFixedBytes, term, &localCopyFixedBytes) || + localCopyFixedBytes >= next.totalUbSize || + next.totalUbSize - localCopyFixedBytes < next.hCommuSize) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + *out = next; + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpBuildMemoryDispatchReferenceConfig(int64_t rankSize, int64_t rank, int64_t bs, int64_t h, + int64_t topK, int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, TileXR::TileXRDataType dtype, uint32_t blockDim, + EpMemoryDispatchReferenceConfig *out) +{ + return TileXREpBuildMemoryDispatchReferenceConfig(rankSize, rank, bs, h, topK, moeExpertNum, + sharedExpertNum, sharedExpertRankNum, globalBs, dtype, dtype, 0, blockDim, out); +} + +int TileXREpBuildMemoryCombineReferenceConfig(int64_t rankSize, int64_t rank, int64_t bs, int64_t h, + int64_t topK, int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, TileXR::TileXRDataType dtype, int64_t quantMode, uint32_t blockDim, + EpMemoryCombineReferenceConfig *out) +{ + const int64_t moeRankNum = rankSize - sharedExpertRankNum; + const bool useMxfp8 = quantMode == 3 || quantMode == 4; + if (out == nullptr || rankSize <= 0 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || rank < 0 || + rank >= rankSize || bs <= 0 || h <= 0 || topK <= 0 || moeExpertNum <= 0 || sharedExpertNum < 0 || + sharedExpertRankNum < 0 || globalBs <= 0 || globalBs % rankSize != 0 || blockDim == 0 || + blockDim > kEpMemoryMaxVectorCoreNum || + (dtype != TileXR::TILEXR_DATA_TYPE_FP16 && dtype != TileXR::TILEXR_DATA_TYPE_BFP16) || + (quantMode != 0 && !useMxfp8) || + ((sharedExpertNum == 0) != (sharedExpertRankNum == 0)) || sharedExpertRankNum >= rankSize || + (sharedExpertNum > 0 && sharedExpertRankNum % sharedExpertNum != 0) || moeRankNum <= 0 || + moeExpertNum % moeRankNum != 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + EpMemoryCombineReferenceConfig next {}; + next.rankSize = rankSize; + next.rank = rank; + next.blockDim = blockDim; + next.moeExpertNumPerRank = moeExpertNum / moeRankNum; + next.sendCountNum = rank < sharedExpertRankNum ? rankSize : rankSize * next.moeExpertNumPerRank; + + const int64_t dtypeBytes = TileXREpDataTypeSize(dtype); + int64_t rowBytes = 0; + int64_t commDataBytes = 0; + int64_t inputAlignBytes = 0; + int64_t scaleCount = 0; + if (!MulInt64(h, dtypeBytes, &rowBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (useMxfp8) { + int64_t fp8PayloadBytes = 0; + if (!AlignUpInt64(h, 256, &fp8PayloadBytes) || + !CeilDivInt64(h, 32, &scaleCount) || + !AlignUpInt64(scaleCount, 2, &scaleCount) || + !AddInt64(fp8PayloadBytes, scaleCount, &commDataBytes) || + !AlignUpInt64(h, 128, &inputAlignBytes) || + !MulInt64(inputAlignBytes, dtypeBytes, &inputAlignBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } else if (!AlignUpInt64(rowBytes, kEpMemoryWindowAlignmentBytes, &inputAlignBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } else { + commDataBytes = rowBytes; + } + if (!CeilDivInt64(commDataBytes, kEpMemorySplitPayloadBytes, &next.blockCntPerToken) || + !MulInt64(next.blockCntPerToken, kEpMemorySplitBlockBytes, &next.packedRowBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t routesPerToken = 0; + int64_t combineRows = 0; + if (!AddInt64(topK, sharedExpertNum, &routesPerToken) || + !MulInt64(globalBs / rankSize, routesPerToken, &combineRows) || + !MulInt64(combineRows, next.packedRowBytes, &next.combineReserveBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t workspaceStatusNum = 0; + int64_t workspaceBytes = 0; + if (!MulInt64(rankSize, next.moeExpertNumPerRank, &workspaceStatusNum) || + !MulInt64(static_cast(blockDim), workspaceStatusNum, &workspaceBytes) || + !MulInt64(workspaceBytes, static_cast(sizeof(int32_t)), &workspaceBytes) || + !AlignUpInt64(workspaceBytes, kEpMemoryWindowAlignmentBytes, &next.workspaceBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.totalWinSize = TileXR::IPC_BUFF_MAX_SIZE - kEpMemoryStateWindowBytes - next.workspaceBytes; + next.combineHalfBytes = next.totalWinSize / 2; + if (next.totalWinSize <= 0 || next.combineReserveBytes > next.combineHalfBytes) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t totalPackedBlocks = 0; + int64_t checkFlagsBytes = 0; + int64_t compactPayloadBytes = 0; + int64_t floatRowBytes = 0; + int64_t alignedFloatRowBytes = 0; + int64_t alignedOutputBytes = 0; + int64_t scaleBytes = 0; + int64_t maskBytes = 0; + int64_t clearFlagBytes = 0; + int64_t receiveBytes = 4 * kEpMemoryWindowAlignmentBytes; + int64_t term = 0; + if (!MulInt64(routesPerToken, next.blockCntPerToken, &totalPackedBlocks) || + !MulInt64(totalPackedBlocks, 2 * kEpMemoryWindowAlignmentBytes, &checkFlagsBytes) || + !AddInt64(checkFlagsBytes, kEpMemoryWindowAlignmentBytes, &checkFlagsBytes) || + !MulInt64(next.blockCntPerToken, kEpMemorySplitPayloadBytes, &compactPayloadBytes) || + !MulInt64(h, static_cast(sizeof(float)), &floatRowBytes) || + !AlignUpInt64(floatRowBytes, useMxfp8 ? 512 : kEpMemoryWindowAlignmentBytes, + &alignedFloatRowBytes) || + !AlignUpInt64(rowBytes, kEpMemoryWindowAlignmentBytes, &alignedOutputBytes) || + !MulInt64(topK, static_cast(sizeof(float)), &term) || + !AlignUpInt64(term, kEpMemoryWindowAlignmentBytes, &scaleBytes) || + !AlignUpInt64(bs * static_cast(sizeof(bool)), kEpMemoryWindowAlignmentBytes, &maskBytes) || + !MulInt64(next.blockCntPerToken, kEpMemoryWindowAlignmentBytes, &clearFlagBytes) || + !AddInt64(receiveBytes, checkFlagsBytes, &receiveBytes) || + !AddInt64(receiveBytes, compactPayloadBytes, &receiveBytes) || + !MulInt64(3, alignedFloatRowBytes, &term) || !AddInt64(receiveBytes, term, &receiveBytes) || + !AddInt64(receiveBytes, alignedOutputBytes, &receiveBytes) || + !AddInt64(receiveBytes, scaleBytes, &receiveBytes) || + !AddInt64(receiveBytes, maskBytes, &receiveBytes) || + !AddInt64(receiveBytes, clearFlagBytes, &receiveBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (useMxfp8) { + int64_t dequantScaleBytes = 0; + if (!AlignUpInt64(scaleCount, 128, &dequantScaleBytes) || + !MulInt64(dequantScaleBytes, dtypeBytes * 2, &dequantScaleBytes) || + !AddInt64(receiveBytes, dequantScaleBytes, &receiveBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + next.receiveUbBytes = receiveBytes; + int64_t sendUbBytes = kEpMemoryWindowAlignmentBytes; + if (!AddInt64(sendUbBytes, inputAlignBytes, &sendUbBytes) || + !AddInt64(sendUbBytes, next.packedRowBytes, &sendUbBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (useMxfp8) { + int64_t alignedScaleCount = 0; + int64_t quantWorkBytes = 0; + if (!AlignUpInt64(scaleCount, 32, &alignedScaleCount) || + !MulInt64(alignedScaleCount, static_cast(sizeof(float)), &quantWorkBytes) || + !AddInt64(quantWorkBytes, scaleCount * static_cast(sizeof(uint16_t)), &quantWorkBytes) || + !AlignUpInt64(quantWorkBytes, kEpMemoryWindowAlignmentBytes, &quantWorkBytes) || + !AddInt64(sendUbBytes, compactPayloadBytes, &sendUbBytes) || + !AddInt64(sendUbBytes, quantWorkBytes, &sendUbBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + if (sendUbBytes > kEpMemoryFullUbBytes || next.receiveUbBytes > kEpMemoryFullUbBytes) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + *out = next; + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpBuildMemoryCombineReferenceConfig(int64_t rankSize, int64_t rank, int64_t bs, int64_t h, + int64_t topK, int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, TileXR::TileXRDataType dtype, uint32_t blockDim, + EpMemoryCombineReferenceConfig *out) +{ + return TileXREpBuildMemoryCombineReferenceConfig(rankSize, rank, bs, h, topK, moeExpertNum, + sharedExpertNum, sharedExpertRankNum, globalBs, dtype, 0, blockDim, out); +} + +} // namespace TileXREp diff --git a/src/ep/host/ep_memory_layout.h b/src/ep/host/ep_memory_layout.h new file mode 100644 index 00000000..59b836eb --- /dev/null +++ b/src/ep/host/ep_memory_layout.h @@ -0,0 +1,82 @@ +#ifndef TILEXR_EP_HOST_EP_MEMORY_LAYOUT_H +#define TILEXR_EP_HOST_EP_MEMORY_LAYOUT_H + +#include + +#include "tilexr_types.h" + +namespace TileXREp { + +constexpr int64_t kEpMemoryStateWindowBytes = 1024 * 1024; +constexpr int64_t kEpMemoryWindowAlignmentBytes = 32; +constexpr int64_t kEpMemorySplitBlockBytes = 512; +constexpr int64_t kEpMemorySplitPayloadBytes = 480; +constexpr int64_t kEpMemoryFullUbBytes = 190 * 1024; +constexpr uint32_t kEpMemoryDefaultVectorCoreNum = 48; +constexpr uint32_t kEpMemoryMaxCountCoreNum = 8; +constexpr uint32_t kEpMemoryMaxVectorCoreNum = 200; + +struct EpMemoryDispatchReferenceConfig { + int64_t rankSize = 0; + int64_t rank = 0; + int64_t localExpertNum = 0; + int64_t rscvStatusNum = 0; + uint32_t blockDim = 0; + uint32_t aivUsedCumSum = 0; + uint32_t aivUsedAllToAll = 0; + uint32_t sharedUsedAivNum = 0; + uint32_t moeUsedAivNum = 0; + int64_t hCommuSize = 0; + int64_t hOutSize = 0; + int64_t scaleOutBytes = 0; + int64_t tokenQuantAlignBytes = 0; + int64_t expertPerSizeOnWin = 0; + int64_t combineReserveBytes = 0; + int64_t workspaceBytes = 0; + int64_t totalWinSize = 0; + int64_t dispatchHalfBytes = 0; + int64_t maxSizeForUbBuffer = 0; + int64_t totalUbSize = 0; +}; + +struct EpMemoryCombineReferenceConfig { + int64_t rankSize = 0; + int64_t rank = 0; + int64_t moeExpertNumPerRank = 0; + int64_t sendCountNum = 0; + uint32_t blockDim = 0; + int64_t blockCntPerToken = 0; + int64_t packedRowBytes = 0; + int64_t combineReserveBytes = 0; + int64_t workspaceBytes = 0; + int64_t totalWinSize = 0; + int64_t combineHalfBytes = 0; + int64_t receiveUbBytes = 0; +}; + +uint32_t TileXREpMemoryCountCoreNum(int64_t totalExpertNum, int64_t rscvStatusNum, uint32_t blockDim); + +int TileXREpBuildMemoryDispatchReferenceConfig(int64_t rankSize, int64_t rank, int64_t bs, int64_t h, + int64_t topK, int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, TileXR::TileXRDataType dtype, TileXR::TileXRDataType expandXOutDtype, + int64_t quantMode, uint32_t blockDim, + EpMemoryDispatchReferenceConfig *out); + +int TileXREpBuildMemoryDispatchReferenceConfig(int64_t rankSize, int64_t rank, int64_t bs, int64_t h, + int64_t topK, int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, TileXR::TileXRDataType dtype, uint32_t blockDim, + EpMemoryDispatchReferenceConfig *out); + +int TileXREpBuildMemoryCombineReferenceConfig(int64_t rankSize, int64_t rank, int64_t bs, int64_t h, + int64_t topK, int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, TileXR::TileXRDataType dtype, int64_t quantMode, uint32_t blockDim, + EpMemoryCombineReferenceConfig *out); + +int TileXREpBuildMemoryCombineReferenceConfig(int64_t rankSize, int64_t rank, int64_t bs, int64_t h, + int64_t topK, int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, TileXR::TileXRDataType dtype, uint32_t blockDim, + EpMemoryCombineReferenceConfig *out); + +} // namespace TileXREp + +#endif // TILEXR_EP_HOST_EP_MEMORY_LAYOUT_H diff --git a/src/ep/host/tilexr_ep_memory_combine.cpp b/src/ep/host/tilexr_ep_memory_combine.cpp new file mode 100644 index 00000000..5a2e0a39 --- /dev/null +++ b/src/ep/host/tilexr_ep_memory_combine.cpp @@ -0,0 +1,68 @@ +#include "tilexr_ep.h" + +#include "ep_dispatch_host.h" +#include "ep_kernel_launch.h" +#include "tilexr_types.h" + +namespace { + +int LaunchEpCombineMemory(const TileXREp::EpCombineParams ¶ms) +{ + int ret = TileXREp::TileXREpValidateBasicCombineParams(params); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + TileXREp::EpHostLaunchContext context {}; + ret = TileXREp::TileXREpPrepareMemoryCombineLaunchContext(params, &context); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + return TileXREp::TileXREpLaunchCombineMemoryKernel(params, context); +} + +} // namespace + +int TileXRMoeEpCombineMemoryV2(void *expertOut, int32_t *assistInfoForCombine, int32_t *sendCounts, + float *expertScales, bool *xActiveMask, int64_t activeMaskType, void *sharedExpertX, + TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + int64_t epWorldSize, int64_t epRankId, int64_t tpWorldSize, int64_t tpRankId, + int64_t expertShardType, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t quantMode, int64_t globalBs, void *yOut, TileXR::TileXRDataType dtype, aclrtStream stream) +{ + TileXREp::EpCombineParams params {}; + params.expertOut = expertOut; + params.assistInfoForCombine = assistInfoForCombine; + params.epRecvCounts = sendCounts; + params.expertScales = expertScales; + params.xActiveMask = xActiveMask; + params.activeMaskType = activeMaskType; + params.sharedExpertX = sharedExpertX; + params.comm = comm; + params.bs = bs; + params.h = h; + params.topK = topK; + params.moeExpertNum = moeExpertNum; + params.epWorldSize = epWorldSize; + params.epRankId = epRankId; + params.tpWorldSize = tpWorldSize; + params.tpRankId = tpRankId; + params.expertShardType = expertShardType; + params.sharedExpertNum = sharedExpertNum; + params.sharedExpertRankNum = sharedExpertRankNum; + params.quantMode = quantMode; + params.globalBs = globalBs; + params.yOut = yOut; + params.workspace = nullptr; + params.dtype = dtype; + params.stream = stream; + return LaunchEpCombineMemory(params); +} + +int TileXRMoeEpCombineMemory(void *expertOut, int32_t *assistInfoForCombine, int32_t *sendCounts, + TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + void *yOut, TileXR::TileXRDataType dtype, aclrtStream stream) +{ + return TileXRMoeEpCombineMemoryV2(expertOut, assistInfoForCombine, sendCounts, nullptr, nullptr, + TileXREp::TILEXR_EP_ACTIVE_MASK_NONE, nullptr, comm, bs, h, topK, moeExpertNum, + 0, 0, 0, 0, 0, 0, 0, 0, 0, yOut, dtype, stream); +} diff --git a/src/ep/host/tilexr_ep_memory_dispatch.cpp b/src/ep/host/tilexr_ep_memory_dispatch.cpp new file mode 100644 index 00000000..cdc97264 --- /dev/null +++ b/src/ep/host/tilexr_ep_memory_dispatch.cpp @@ -0,0 +1,82 @@ +#include "tilexr_ep.h" + +#include "ep_dispatch_host.h" +#include "ep_kernel_launch.h" +#include "tilexr_types.h" + +namespace { + +int LaunchEpDispatchMemory(const TileXREp::EpDispatchParams ¶ms) +{ + int ret = TileXREp::TileXREpValidateBasicDispatchParams(params); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + + TileXREp::EpHostLaunchContext context {}; + ret = TileXREp::TileXREpPrepareMemoryLaunchContext(params, &context); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + + return TileXREp::TileXREpLaunchDispatchMemoryKernel(params, context); +} + +} // namespace + +int TileXRMoeEpDispatchMemoryV2(void *x, int32_t *expertIds, void *scales, bool *xActiveMask, + int64_t activeMaskType, void *expertScales, TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, + int64_t moeExpertNum, int64_t epWorldSize, + int64_t epRankId, int64_t tpWorldSize, int64_t tpRankId, int64_t expertShardType, int64_t sharedExpertNum, + int64_t sharedExpertRankNum, int64_t quantMode, int64_t globalBs, int64_t expertTokenNumsType, void *expandXOut, + void *dynamicScalesOut, int32_t *assistInfoForCombineOut, int64_t *expertTokenNumsOut, int32_t *sendCountsOut, + int32_t *tpRecvCountsOut, void *expandScalesOut, TileXR::TileXRDataType dtype, + TileXR::TileXRDataType expandXOutDtype, aclrtStream stream) +{ + TileXREp::EpDispatchParams params {}; + params.x = x; + params.expertIds = expertIds; + params.scales = scales; + params.xActiveMask = xActiveMask; + params.activeMaskType = activeMaskType; + params.expertScales = expertScales; + params.comm = comm; + params.bs = bs; + params.h = h; + params.topK = topK; + params.moeExpertNum = moeExpertNum; + params.epWorldSize = epWorldSize; + params.epRankId = epRankId; + params.tpWorldSize = tpWorldSize; + params.tpRankId = tpRankId; + params.expertShardType = expertShardType; + params.sharedExpertNum = sharedExpertNum; + params.sharedExpertRankNum = sharedExpertRankNum; + params.quantMode = quantMode; + params.globalBs = globalBs; + params.expertTokenNumsType = expertTokenNumsType; + params.expandXOut = expandXOut; + params.dynamicScalesOut = dynamicScalesOut; + params.assistInfoForCombineOut = assistInfoForCombineOut; + params.expertTokenNumsOut = expertTokenNumsOut; + params.epRecvCountsOut = sendCountsOut; + params.tpRecvCountsOut = tpRecvCountsOut; + params.expandScalesOut = expandScalesOut; + params.workspace = nullptr; + params.dtype = dtype; + params.expandXOutDtype = expandXOutDtype; + params.stream = stream; + + return LaunchEpDispatchMemory(params); +} + +int TileXRMoeEpDispatchMemory(void *x, int32_t *expertIds, TileXRCommPtr comm, + int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + void *expandXOut, int64_t *expertTokenNumsOut, int32_t *sendCountsOut, + int32_t *assistInfoForCombineOut, TileXR::TileXRDataType dtype, aclrtStream stream) +{ + return TileXRMoeEpDispatchMemoryV2(x, expertIds, nullptr, nullptr, TileXREp::TILEXR_EP_ACTIVE_MASK_NONE, + nullptr, comm, bs, h, topK, moeExpertNum, 0, 0, 0, 0, 0, 0, 0, 0, 0, + TileXREp::kEpExpertTokenNumsTypeCount, expandXOut, nullptr, + assistInfoForCombineOut, expertTokenNumsOut, sendCountsOut, nullptr, nullptr, dtype, dtype, stream); +} diff --git a/src/ep/kernels/tilexr_ep_combine_memory_kernel.cpp b/src/ep/kernels/tilexr_ep_combine_memory_kernel.cpp new file mode 100644 index 00000000..468ba4ed --- /dev/null +++ b/src/ep/kernels/tilexr_ep_combine_memory_kernel.cpp @@ -0,0 +1,710 @@ +#include + +#include "comm_args.h" +#include "kernel_operator.h" +#include "tilexr_data_as_flag.h" +#include "tilexr_ep_mxfp8_quant.h" +#include "tilexr_types.h" + +#define FLOAT_OVERFLOW_MODE_CTRL 60 + +namespace Mc2Kernel { + +using namespace AscendC; + +constexpr uint32_t UB_ALIGN = 32U; +constexpr uint32_t STATE_WINDOW_BYTES = 1024U * 1024U; +constexpr uint32_t DISPATCH_STATE_OFFSET = 768U * 1024U; +constexpr uint32_t STATE_SLOT_STRIDE_BYTES = 512U; +constexpr uint32_t ASSIST_INFO_WIDTH = 4U; +constexpr int64_t ACTIVE_MASK_NONE = 0; +constexpr int64_t ACTIVE_MASK_TOKEN = 1; +constexpr int64_t MXFP8_E5M2_COMM_QUANT = 3; +constexpr int64_t MXFP8_E4M3_COMM_QUANT = 4; + +template +__aicore__ inline void SyncFunc() +{ + AscendC::TEventID eventId = GetTPipePtr()->FetchEventID(event); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +__aicore__ inline uint64_t CeilDiv(uint64_t value, uint64_t divisor) +{ + return value / divisor + (value % divisor == 0U ? 0U : 1U); +} + +__aicore__ inline uint64_t AlignUp(uint64_t value, uint64_t alignment) +{ + return CeilDiv(value, alignment) * alignment; +} + +class TileXRCombineMemoryContext { +public: + __aicore__ inline void Init(GM_ADDR commArgsGM) + { + args_ = reinterpret_cast<__gm__ TileXR::CommArgs *>(commArgsGM); + } + + __aicore__ inline uint32_t Rank() const + { + return static_cast(args_->rank); + } + + __aicore__ inline uint32_t RankSize() const + { + return static_cast(args_->rankSize); + } + + __aicore__ inline GM_ADDR PeerDataBase(uint32_t rank) const + { + return args_->peerMems[rank] + TileXR::IPC_DATA_OFFSET; + } + +private: + __gm__ TileXR::CommArgs *args_{nullptr}; +}; + +template +class MoeDistributeCombineV2A5Mte { +public: + __aicore__ inline void Init(GM_ADDR commArgsGM, GM_ADDR expertOutGM, GM_ADDR assistInfoGM, + GM_ADDR sendCountsGM, GM_ADDR expertScalesGM, GM_ADDR xActiveMaskGM, GM_ADDR sharedExpertXGM, + GM_ADDR yOutGM, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, + int64_t activeMaskType, int64_t quantMode, TPipe *pipe); + __aicore__ inline void Process(); + +private: + __aicore__ inline void ReadDispatchStateAndSendCount(); + __aicore__ inline void SplitSendRows(); + __aicore__ inline void InitSendBuffers(); + __aicore__ inline void ExpertAlltoAllDispatchCopyAdd(); + __aicore__ inline void SendOneRow(uint32_t row, LocalTensor assistLocal); + template + __aicore__ inline void QuantMxfp8(LocalTensor quantLocal, LocalTensor inputLocal); + __aicore__ inline void InitReceiveBuffers(); + __aicore__ inline bool TokenActive(uint32_t tokenIndex); + __aicore__ inline bool CheckPackedTokenArrive(uint32_t tokenIndex); + __aicore__ inline void CopyAndAccumulateSlot(uint32_t tokenIndex, uint32_t slotIndex, float scale, + LocalTensor sumLocal, LocalTensor rowFloat, LocalTensor mulLocal); + __aicore__ inline void AddSharedExpertX(uint32_t tokenIndex, LocalTensor sumLocal, + LocalTensor rowFloat); + __aicore__ inline void ClearTokenFlags(uint32_t tokenIndex, LocalTensor clearLocal); + __aicore__ inline void WriteOutput(uint32_t tokenIndex, LocalTensor sumLocal); + __aicore__ inline void LocalWindowCopy(); + __aicore__ inline GM_ADDR CombineWindowBase(uint32_t rank) const; + + TPipe *pipe_{nullptr}; + TileXRCombineMemoryContext context_; + uint32_t coreIdx_{0}; + uint32_t aivNum_{0}; + uint32_t epRankId_{0}; + uint32_t epWorldSize_{0}; + uint32_t axisBS_{0}; + uint32_t axisH_{0}; + uint32_t axisK_{0}; + uint32_t moeExpertNum_{0}; + uint32_t sharedExpertNum_{0}; + uint32_t sharedExpertRankNum_{0}; + uint32_t moeExpertNumPerRank_{0}; + uint32_t sendCountNum_{0}; + uint32_t selfSendCnt_{0}; + uint32_t startSendRow_{0}; + uint32_t sendRowCount_{0}; + uint32_t dataState_{0}; + uint32_t blockCntPerToken_{0}; + uint32_t packedRowBytes_{0}; + uint32_t compactPayloadBytes_{0}; + uint32_t commDataBytes_{0}; + uint32_t rowBytes_{0}; + uint32_t rowAlignBytes_{0}; + uint32_t inputAlignBytes_{0}; + uint32_t floatRowAlignBytes_{0}; + uint32_t quantPayloadBytes_{0}; + uint32_t quantScaleCount_{0}; + uint32_t quantWorkBytes_{0}; + uint32_t slotCount_{0}; + uint64_t totalWinSize_{0}; + uint64_t halfWinSize_{0}; + int64_t activeMaskType_{ACTIVE_MASK_NONE}; + int64_t quantMode_{0}; + bool useMxfp8_{false}; + bool hasExpertScales_{false}; + bool hasSharedExpertX_{false}; + GlobalTensor expertOutGM_; + GlobalTensor assistInfoGM_; + GlobalTensor sendCountsGM_; + GlobalTensor expertScalesGM_; + GlobalTensor xActiveMaskGM_; + GlobalTensor sharedExpertXGM_; + GlobalTensor yOutGM_; + + TBuf<> metaBuf_; + TBuf<> assistBuf_; + TQue sendInputQueue_; + TQue sendOutputQueue_; + TBuf<> quantResultBuf_; + TBuf<> quantWorkBuf_; + TBuf<> packedCheckFlagBuf_; + TBuf<> packedCheckCompareBuf_; + TQue packedInputQueue_; + TBuf<> rowFloatBuf_; + TBuf<> mulFloatBuf_; + TBuf<> sumFloatBuf_; + TBuf<> dequantScaleBuf_; + TQue outputQueue_; + TBuf<> expertScaleBuf_; + TBuf<> activeMaskBuf_; + TBuf<> clearFlagBuf_; + TBuf<> tokenStatusBuf_; +}; + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::Init(GM_ADDR commArgsGM, GM_ADDR expertOutGM, + GM_ADDR assistInfoGM, GM_ADDR sendCountsGM, GM_ADDR expertScalesGM, GM_ADDR xActiveMaskGM, + GM_ADDR sharedExpertXGM, GM_ADDR yOutGM, int64_t bs, int64_t h, int64_t topK, + int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t globalBs, int64_t activeMaskType, int64_t quantMode, TPipe *pipe) +{ + AscendC::SetCtrlSpr(0); + pipe_ = pipe; + context_.Init(commArgsGM); + coreIdx_ = GetBlockIdx(); + aivNum_ = GetBlockNum(); + epRankId_ = context_.Rank(); + epWorldSize_ = context_.RankSize(); + axisBS_ = static_cast(bs); + axisH_ = static_cast(h); + axisK_ = static_cast(topK); + moeExpertNum_ = static_cast(moeExpertNum); + sharedExpertNum_ = static_cast(sharedExpertNum); + sharedExpertRankNum_ = static_cast(sharedExpertRankNum); + activeMaskType_ = activeMaskType; + quantMode_ = quantMode; + useMxfp8_ = quantMode_ == MXFP8_E5M2_COMM_QUANT || quantMode_ == MXFP8_E4M3_COMM_QUANT; + slotCount_ = axisK_ + sharedExpertNum_; + const uint32_t moeRankNum = epWorldSize_ - sharedExpertRankNum_; + moeExpertNumPerRank_ = moeExpertNum_ / moeRankNum; + sendCountNum_ = epRankId_ < sharedExpertRankNum_ ? epWorldSize_ : epWorldSize_ * moeExpertNumPerRank_; + + rowBytes_ = axisH_ * sizeof(XType); + rowAlignBytes_ = static_cast(AlignUp(rowBytes_, UB_ALIGN)); + inputAlignBytes_ = useMxfp8_ ? static_cast(AlignUp(axisH_, 128U) * sizeof(XType)) : rowAlignBytes_; + floatRowAlignBytes_ = static_cast( + AlignUp(axisH_ * sizeof(float), useMxfp8_ ? 512U : UB_ALIGN)); + if (useMxfp8_) { + quantPayloadBytes_ = static_cast(AlignUp(axisH_, 256U)); + quantScaleCount_ = static_cast(AlignUp(CeilDiv(axisH_, 32U), 2U)); + commDataBytes_ = quantPayloadBytes_ + quantScaleCount_; + quantWorkBytes_ = static_cast(AlignUp( + AlignUp(quantScaleCount_, 32U) * sizeof(float) + + quantScaleCount_ * sizeof(uint16_t), UB_ALIGN)); + } else { + commDataBytes_ = rowBytes_; + } + blockCntPerToken_ = static_cast( + CeilDiv(commDataBytes_, TileXR::DATA_AS_FLAG_PAYLOAD_BYTES)); + packedRowBytes_ = blockCntPerToken_ * TileXR::DATA_AS_FLAG_BLOCK_BYTES; + compactPayloadBytes_ = blockCntPerToken_ * TileXR::DATA_AS_FLAG_PAYLOAD_BYTES; + + const uint64_t workspaceStatusNum = static_cast(epWorldSize_) * moeExpertNumPerRank_; + const uint64_t workspaceBytes = AlignUp( + static_cast(aivNum_) * workspaceStatusNum * sizeof(int32_t), UB_ALIGN); + totalWinSize_ = TileXR::IPC_BUFF_MAX_SIZE - STATE_WINDOW_BYTES - workspaceBytes; + halfWinSize_ = totalWinSize_ / 2U; + + expertOutGM_.SetGlobalBuffer(reinterpret_cast<__gm__ XType *>(expertOutGM)); + assistInfoGM_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(assistInfoGM)); + sendCountsGM_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(sendCountsGM)); + expertScalesGM_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(expertScalesGM)); + xActiveMaskGM_.SetGlobalBuffer(reinterpret_cast<__gm__ bool *>(xActiveMaskGM)); + sharedExpertXGM_.SetGlobalBuffer(reinterpret_cast<__gm__ XType *>(sharedExpertXGM)); + yOutGM_.SetGlobalBuffer(reinterpret_cast<__gm__ XType *>(yOutGM)); + hasExpertScales_ = expertScalesGM != nullptr; + hasSharedExpertX_ = sharedExpertXGM != nullptr; + + pipe_->InitBuffer(metaBuf_, 2U * UB_ALIGN); + ReadDispatchStateAndSendCount(); + SplitSendRows(); +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::ReadDispatchStateAndSendCount() +{ + LocalTensor meta = metaBuf_.Get(); + GlobalTensor stateGlobal; + stateGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t *>(context_.PeerDataBase(epRankId_) + + DISPATCH_STATE_OFFSET + static_cast(coreIdx_) * STATE_SLOT_STRIDE_BYTES)); + DataCopy(meta, stateGlobal, UB_ALIGN / sizeof(uint32_t)); + SyncFunc(); + dataState_ = meta.GetValue(0) == 0U ? 1U : 0U; + + LocalTensor countLocal = metaBuf_.GetWithOffset(UB_ALIGN / sizeof(int32_t), UB_ALIGN); + const DataCopyExtParams countParams {1U, sizeof(int32_t), 0U, 0U, 0U}; + const DataCopyPadExtParams padParams {false, 0U, 0U, 0U}; + DataCopyPad(countLocal, sendCountsGM_[sendCountNum_ - 1U], countParams, padParams); + SyncFunc(); + const int32_t count = countLocal.GetValue(0); + selfSendCnt_ = count > 0 ? static_cast(count) : 0U; +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::SplitSendRows() +{ + sendRowCount_ = selfSendCnt_ / aivNum_; + const uint32_t remainder = selfSendCnt_ % aivNum_; + startSendRow_ = sendRowCount_ * coreIdx_; + if (coreIdx_ < remainder) { + ++sendRowCount_; + startSendRow_ += coreIdx_; + } else { + startSendRow_ += remainder; + } +} + +template +__aicore__ inline GM_ADDR MoeDistributeCombineV2A5Mte::CombineWindowBase(uint32_t rank) const +{ + return context_.PeerDataBase(rank) + STATE_WINDOW_BYTES + dataState_ * halfWinSize_; +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::InitSendBuffers() +{ + pipe_->Reset(); + AscendC::SetCtrlSpr(0); + pipe_->InitBuffer(assistBuf_, UB_ALIGN); + pipe_->InitBuffer(sendInputQueue_, 1, inputAlignBytes_); + pipe_->InitBuffer(sendOutputQueue_, 1, packedRowBytes_); + if (useMxfp8_) { + pipe_->InitBuffer(quantResultBuf_, compactPayloadBytes_); + pipe_->InitBuffer(quantWorkBuf_, quantWorkBytes_); + } +} + +template +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::QuantMxfp8( + LocalTensor quantLocal, LocalTensor inputLocal) +{ + __ubuf__ XType *srcAddr = reinterpret_cast<__ubuf__ XType *>(inputLocal.GetPhyAddr()); + __ubuf__ uint8_t *workAddr = reinterpret_cast<__ubuf__ uint8_t *>( + quantWorkBuf_.Get().GetPhyAddr()); + __ubuf__ uint16_t *maxExpAddr = reinterpret_cast<__ubuf__ uint16_t *>(workAddr); + __ubuf__ uint16_t *halfScaleAddr = reinterpret_cast<__ubuf__ uint16_t *>( + workAddr + AlignUp(quantScaleCount_, 32U) * sizeof(float)); + __ubuf__ int8_t *outAddr = reinterpret_cast<__ubuf__ int8_t *>(quantLocal.GetPhyAddr()); + __ubuf__ uint16_t *scaleAddr = reinterpret_cast<__ubuf__ uint16_t *>( + quantLocal[quantPayloadBytes_].GetPhyAddr()); + + TileXRMxfp8Quant::ComputeMaxExp(srcAddr, maxExpAddr, axisH_); + TileXRMxfp8Quant::ComputeScale( + maxExpAddr, scaleAddr, halfScaleAddr, quantScaleCount_); + TileXRMxfp8Quant::ComputeFp8Data( + srcAddr, halfScaleAddr, outAddr, axisH_); +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::SendOneRow(uint32_t row, + LocalTensor assistLocal) +{ + const DataCopyExtParams assistParams {1U, ASSIST_INFO_WIDTH * sizeof(int32_t), 0U, 0U, 0U}; + const DataCopyPadExtParams assistPad {false, 0U, 0U, 0U}; + DataCopyPad(assistLocal, assistInfoGM_[row * ASSIST_INFO_WIDTH], assistParams, assistPad); + SyncFunc(); + const uint32_t toRankId = static_cast(assistLocal.GetValue(0)); + const uint32_t tokenId = static_cast(assistLocal.GetValue(1)); + const uint32_t topkId = static_cast(assistLocal.GetValue(2)); + if (toRankId >= epWorldSize_ || tokenId >= axisBS_ || topkId >= slotCount_) { + return; + } + const uint64_t slot = static_cast(tokenId) * slotCount_ + topkId; + GM_ADDR dst = CombineWindowBase(toRankId) + slot * packedRowBytes_; + + LocalTensor inputLocal = sendInputQueue_.AllocTensor(); + const DataCopyExtParams inputParams {1U, rowBytes_, 0U, 0U, 0U}; + const DataCopyPadExtParams inputPad {useMxfp8_, 0U, 0U, 0U}; + if (useMxfp8_) { + Duplicate(inputLocal.template ReinterpretCast(), 0, inputAlignBytes_); + SyncFunc(); + } + DataCopyPad(inputLocal, expertOutGM_[static_cast(row) * axisH_], inputParams, inputPad); + sendInputQueue_.EnQue(inputLocal); + inputLocal = sendInputQueue_.DeQue(); + + LocalTensor outputLocal = sendOutputQueue_.AllocTensor(); + LocalTensor sourceFloat = inputLocal.template ReinterpretCast(); + if (useMxfp8_) { + LocalTensor quantLocal = quantResultBuf_.Get(); + Duplicate(quantLocal, 0, compactPayloadBytes_); + PipeBarrier(); + if (quantMode_ == MXFP8_E5M2_COMM_QUANT) { + QuantMxfp8(quantLocal, inputLocal); + } else { + QuantMxfp8(quantLocal, inputLocal); + } + PipeBarrier(); + sourceFloat = quantLocal.template ReinterpretCast(); + } + LocalTensor packedFloat = outputLocal.template ReinterpretCast(); + Duplicate(packedFloat, TileXR::DATA_AS_FLAG_READY_VALUE, packedRowBytes_ / sizeof(float)); + PipeBarrier(); + Copy(packedFloat, sourceFloat, 64U, static_cast(blockCntPerToken_), {1, 1, 16, 15}); + Copy(packedFloat[64], sourceFloat[64], 56U, static_cast(blockCntPerToken_), {1, 1, 16, 15}); + sendOutputQueue_.EnQue(outputLocal); + outputLocal = sendOutputQueue_.DeQue(); + + GlobalTensor dstPackedGlobal; + dstPackedGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(dst)); + DataCopy(dstPackedGlobal, packedFloat, packedRowBytes_ / sizeof(float)); + sendOutputQueue_.FreeTensor(outputLocal); + sendInputQueue_.FreeTensor(inputLocal); +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::ExpertAlltoAllDispatchCopyAdd() +{ + if (sendRowCount_ == 0U) { + return; + } + LocalTensor assistLocal = assistBuf_.Get(); + uint32_t permStride = 1U; + if (sendRowCount_ > 2U) { + permStride = sendRowCount_ / 2U + 1U; + if (((sendRowCount_ & 1U) == 0U) && ((permStride & 1U) == 0U)) { + ++permStride; + } + } + const uint32_t rankOffset = (epRankId_ * sendRowCount_) / epWorldSize_; + uint32_t permIdx = rankOffset % sendRowCount_; + for (uint32_t loop = 0U; loop < sendRowCount_; ++loop) { + SendOneRow(startSendRow_ + permIdx, assistLocal); + permIdx += permStride; + if (permIdx >= sendRowCount_) { + permIdx -= sendRowCount_; + } + } + SyncFunc(); +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::InitReceiveBuffers() +{ + pipe_->Reset(); + AscendC::SetCtrlSpr(0); + const uint32_t totalBlocks = slotCount_ * blockCntPerToken_; + const uint32_t flagFloatCount = totalBlocks * TileXR::DATA_AS_FLAG_FLAG_FLOATS; + const uint32_t compareCount = static_cast(AlignUp(flagFloatCount, 64U)); + pipe_->InitBuffer(packedCheckFlagBuf_, compareCount * sizeof(float)); + pipe_->InitBuffer(packedCheckCompareBuf_, AlignUp(compareCount * sizeof(uint8_t), 256U)); + pipe_->InitBuffer(packedInputQueue_, 1, compactPayloadBytes_); + pipe_->InitBuffer(rowFloatBuf_, floatRowAlignBytes_); + pipe_->InitBuffer(mulFloatBuf_, floatRowAlignBytes_); + pipe_->InitBuffer(sumFloatBuf_, floatRowAlignBytes_); + if (useMxfp8_) { + pipe_->InitBuffer(dequantScaleBuf_, + AlignUp(quantScaleCount_, 128U) * sizeof(XType) * 2U); + } + pipe_->InitBuffer(outputQueue_, 1, rowAlignBytes_); + pipe_->InitBuffer(expertScaleBuf_, AlignUp(axisK_ * sizeof(float), UB_ALIGN)); + pipe_->InitBuffer(activeMaskBuf_, AlignUp(axisBS_ * sizeof(bool), UB_ALIGN)); + pipe_->InitBuffer(clearFlagBuf_, blockCntPerToken_ * TileXR::DATA_AS_FLAG_FLAG_BYTES); + const uint32_t maxTokenCountPerCore = static_cast(CeilDiv(axisBS_, aivNum_)); + pipe_->InitBuffer(tokenStatusBuf_, AlignUp(maxTokenCountPerCore * sizeof(int32_t), UB_ALIGN)); + if (activeMaskType_ == ACTIVE_MASK_TOKEN) { + LocalTensor maskLocal = activeMaskBuf_.Get(); + const DataCopyExtParams maskParams { + 1U, static_cast(axisBS_ * sizeof(bool)), 0U, 0U, 0U}; + const DataCopyPadExtParams maskPad {false, 0U, 0U, 0U}; + DataCopyPad(maskLocal, xActiveMaskGM_, maskParams, maskPad); + SyncFunc(); + } + LocalTensor clearLocal = clearFlagBuf_.Get(); + Duplicate(clearLocal, 0.0f, + blockCntPerToken_ * TileXR::DATA_AS_FLAG_FLAG_BYTES / sizeof(float)); + SyncFunc(); +} + +template +__aicore__ inline bool MoeDistributeCombineV2A5Mte::TokenActive(uint32_t tokenIndex) +{ + if (activeMaskType_ != ACTIVE_MASK_TOKEN) { + return true; + } + return activeMaskBuf_.Get().GetValue(tokenIndex); +} + +template +__aicore__ inline bool MoeDistributeCombineV2A5Mte::CheckPackedTokenArrive(uint32_t tokenIndex) +{ + const uint32_t totalBlocks = slotCount_ * blockCntPerToken_; + const uint32_t flagFloatCount = totalBlocks * TileXR::DATA_AS_FLAG_FLAG_FLOATS; + const uint32_t compareCount = static_cast(AlignUp(flagFloatCount, 64U)); + const uint32_t compareU64Count = static_cast(CeilDiv(flagFloatCount, 64U)); + GM_ADDR tokenBase = CombineWindowBase(epRankId_) + + static_cast(tokenIndex) * slotCount_ * packedRowBytes_; + + LocalTensor flagLocal = packedCheckFlagBuf_.Get(); + LocalTensor compareLocal = packedCheckCompareBuf_.Get(); + LocalTensor compareU64 = packedCheckCompareBuf_.Get(); + Duplicate(flagLocal, 0.0f, compareCount); + PipeBarrier(); + + GlobalTensor flagGlobal; + flagGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( + tokenBase + TileXR::DATA_AS_FLAG_FLAG_OFFSET_BYTES)); + const DataCopyExtParams flagParams {static_cast(totalBlocks), + TileXR::DATA_AS_FLAG_FLAG_BYTES, TileXR::DATA_AS_FLAG_PAYLOAD_BYTES, 0U, 0U}; + const DataCopyPadExtParams flagPad {false, 0U, 0U, 0U}; + DataCopyPad(flagLocal, flagGlobal, flagParams, flagPad); + SyncFunc(); + CompareScalar(compareLocal, flagLocal, TileXR::DATA_AS_FLAG_READY_VALUE, + AscendC::CMPMODE::EQ, compareCount); + SyncFunc(); + + uint32_t arrived = 0U; + for (uint32_t index = 0U; index < compareU64Count; ++index) { + const uint64_t mask = compareU64.GetValue(index); + const int64_t firstInvalid = ScalarGetSFFValue<0>(mask); + if (firstInvalid == -1) { + arrived += 64U; + } else { + arrived += static_cast(firstInvalid); + break; + } + } + return (arrived > flagFloatCount ? flagFloatCount : arrived) == flagFloatCount; +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::CopyAndAccumulateSlot(uint32_t tokenIndex, + uint32_t slotIndex, float scale, LocalTensor sumLocal, LocalTensor rowFloat, + LocalTensor mulLocal) +{ + GM_ADDR slotAddr = CombineWindowBase(epRankId_) + + (static_cast(tokenIndex) * slotCount_ + slotIndex) * packedRowBytes_; + GlobalTensor packedGlobal; + packedGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ XType *>(slotAddr)); + LocalTensor packedLocal = packedInputQueue_.AllocTensor(); + const DataCopyExtParams copyParams {static_cast(blockCntPerToken_), + TileXR::DATA_AS_FLAG_PAYLOAD_BYTES, TileXR::DATA_AS_FLAG_FLAG_BYTES, 0U, 0U}; + const DataCopyPadExtParams copyPad {false, 0U, 0U, 0U}; + DataCopyPad(packedLocal, packedGlobal, copyParams, copyPad); + packedInputQueue_.EnQue(packedLocal); + packedLocal = packedInputQueue_.DeQue(); + if (useMxfp8_) { + __ubuf__ uint8_t *tokenAddr = reinterpret_cast<__ubuf__ uint8_t *>(packedLocal.GetPhyAddr()); + __ubuf__ fp8_e8m0_t *scaleAddr = reinterpret_cast<__ubuf__ fp8_e8m0_t *>( + tokenAddr + quantPayloadBytes_); + __ubuf__ float *scaleWorkAddr = reinterpret_cast<__ubuf__ float *>( + dequantScaleBuf_.Get().GetPhyAddr()); + __ubuf__ float *sumAddr = reinterpret_cast<__ubuf__ float *>(sumLocal.GetPhyAddr()); + if (quantMode_ == MXFP8_E5M2_COMM_QUANT) { + TileXRMxfp8Quant::DequantizeAndAccumulate( + tokenAddr, scaleAddr, scaleWorkAddr, sumAddr, + axisH_, quantScaleCount_, scale); + } else { + TileXRMxfp8Quant::DequantizeAndAccumulate( + tokenAddr, scaleAddr, scaleWorkAddr, sumAddr, + axisH_, quantScaleCount_, scale); + } + packedInputQueue_.FreeTensor(packedLocal); + return; + } + Cast(rowFloat, packedLocal, RoundMode::CAST_NONE, axisH_); + PipeBarrier(); + if (scale == 1.0f) { + Add(sumLocal, sumLocal, rowFloat, axisH_); + } else { + Muls(mulLocal, rowFloat, scale, axisH_); + PipeBarrier(); + Add(sumLocal, sumLocal, mulLocal, axisH_); + } + PipeBarrier(); + packedInputQueue_.FreeTensor(packedLocal); +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::AddSharedExpertX(uint32_t tokenIndex, + LocalTensor sumLocal, LocalTensor rowFloat) +{ + GlobalTensor sharedGlobal = sharedExpertXGM_; + LocalTensor packedLocal = packedInputQueue_.AllocTensor(); + const DataCopyExtParams copyParams {1U, rowBytes_, 0U, 0U, 0U}; + const DataCopyPadExtParams copyPad {false, 0U, 0U, 0U}; + DataCopyPad(packedLocal, sharedGlobal[tokenIndex * axisH_], copyParams, copyPad); + packedInputQueue_.EnQue(packedLocal); + packedLocal = packedInputQueue_.DeQue(); + Cast(rowFloat, packedLocal, RoundMode::CAST_NONE, axisH_); + PipeBarrier(); + Add(sumLocal, sumLocal, rowFloat, axisH_); + PipeBarrier(); + packedInputQueue_.FreeTensor(packedLocal); +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::ClearTokenFlags(uint32_t tokenIndex, + LocalTensor clearLocal) +{ + const DataCopyExtParams clearParams {static_cast(blockCntPerToken_), + TileXR::DATA_AS_FLAG_FLAG_BYTES, 0U, TileXR::DATA_AS_FLAG_PAYLOAD_BYTES, 0U}; + for (uint32_t slot = 0U; slot < slotCount_; ++slot) { + GM_ADDR slotAddr = CombineWindowBase(epRankId_) + + (static_cast(tokenIndex) * slotCount_ + slot) * packedRowBytes_; + GlobalTensor flagGlobal; + flagGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( + slotAddr + TileXR::DATA_AS_FLAG_FLAG_OFFSET_BYTES)); + DataCopyPad(flagGlobal, clearLocal, clearParams); + } +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::WriteOutput(uint32_t tokenIndex, + LocalTensor sumLocal) +{ + PipeBarrier(); + LocalTensor outputLocal = outputQueue_.AllocTensor(); + Cast(outputLocal, sumLocal, RoundMode::CAST_RINT, axisH_); + outputQueue_.EnQue(outputLocal); + outputLocal = outputQueue_.DeQue(); + const DataCopyExtParams outputParams {1U, rowBytes_, 0U, 0U, 0U}; + DataCopyPad(yOutGM_[tokenIndex * axisH_], outputLocal, outputParams); + outputQueue_.FreeTensor(outputLocal); +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::LocalWindowCopy() +{ + LocalTensor rowFloat = rowFloatBuf_.Get(); + LocalTensor mulLocal = mulFloatBuf_.Get(); + LocalTensor sumLocal = sumFloatBuf_.Get(); + LocalTensor scalesLocal = expertScaleBuf_.Get(); + LocalTensor clearLocal = clearFlagBuf_.Get(); + const DataCopyExtParams scaleParams { + 1U, static_cast(axisK_ * sizeof(float)), 0U, 0U, 0U}; + const DataCopyPadExtParams scalePad {false, 0U, 0U, 0U}; + + uint32_t tokenCount = axisBS_ / aivNum_; + const uint32_t remainder = axisBS_ % aivNum_; + uint32_t beginIndex = tokenCount * coreIdx_; + if (coreIdx_ < remainder) { + ++tokenCount; + beginIndex += coreIdx_; + } else { + beginIndex += remainder; + } + if (tokenCount == 0U) { + return; + } + const uint32_t endIndex = beginIndex + tokenCount; + LocalTensor tokenStatus = tokenStatusBuf_.Get(); + Duplicate(tokenStatus, 0, tokenCount); + SyncFunc(); + + uint32_t completed = 0U; + while (completed != tokenCount) { + for (uint32_t tokenIndex = beginIndex; tokenIndex < endIndex; ++tokenIndex) { + const uint32_t localIndex = tokenIndex - beginIndex; + if (tokenStatus.GetValue(localIndex) == 1) { + continue; + } + if (!TokenActive(tokenIndex)) { + Duplicate(sumLocal, 0.0f, axisH_); + WriteOutput(tokenIndex, sumLocal); + } else { + if (!CheckPackedTokenArrive(tokenIndex)) { + continue; + } + Duplicate(sumLocal, 0.0f, axisH_); + if (hasExpertScales_) { + DataCopyPad(scalesLocal, expertScalesGM_[tokenIndex * axisK_], scaleParams, scalePad); + SyncFunc(); + } + for (uint32_t topk = 0U; topk < axisK_; ++topk) { + const float scale = hasExpertScales_ ? scalesLocal.GetValue(topk) : 1.0f; + CopyAndAccumulateSlot(tokenIndex, topk, scale, sumLocal, rowFloat, mulLocal); + } + for (uint32_t shared = 0U; shared < sharedExpertNum_; ++shared) { + CopyAndAccumulateSlot(tokenIndex, axisK_ + shared, 1.0f, sumLocal, rowFloat, mulLocal); + } + if (hasSharedExpertX_) { + AddSharedExpertX(tokenIndex, sumLocal, rowFloat); + } + ClearTokenFlags(tokenIndex, clearLocal); + WriteOutput(tokenIndex, sumLocal); + } + tokenStatus.SetValue(localIndex, 1); + ++completed; + } + } + SyncFunc(); +} + +template +__aicore__ inline void MoeDistributeCombineV2A5Mte::Process() +{ + if ASCEND_IS_AIV { + InitSendBuffers(); + ExpertAlltoAllDispatchCopyAdd(); + PipeBarrier(); + InitReceiveBuffers(); + LocalWindowCopy(); + } +} + +template +__aicore__ inline void RunCombine(GM_ADDR commArgsGM, GM_ADDR expertOutGM, GM_ADDR assistInfoGM, + GM_ADDR sendCountsGM, GM_ADDR expertScalesGM, GM_ADDR xActiveMaskGM, GM_ADDR sharedExpertXGM, + GM_ADDR yOutGM, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, + int64_t activeMaskType, int64_t quantMode) +{ + TPipe pipe; + MoeDistributeCombineV2A5Mte op; + op.Init(commArgsGM, expertOutGM, assistInfoGM, sendCountsGM, expertScalesGM, xActiveMaskGM, + sharedExpertXGM, yOutGM, bs, h, topK, moeExpertNum, sharedExpertNum, sharedExpertRankNum, + globalBs, activeMaskType, quantMode, &pipe); + op.Process(); +} + +} // namespace Mc2Kernel + +extern "C" __global__ __aicore__ void tilexr_ep_combine_memory_kernel(GM_ADDR commArgsGM, GM_ADDR expertOutGM, + GM_ADDR assistInfoForCombineGM, GM_ADDR sendCountsGM, GM_ADDR expertScalesGM, GM_ADDR xActiveMaskGM, + GM_ADDR sharedExpertXGM, GM_ADDR yOutGM, int64_t bs, int64_t h, int64_t topK, + int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, + int64_t activeMaskType, int64_t quantMode, int64_t dtype) +{ + if (commArgsGM == nullptr || expertOutGM == nullptr || assistInfoForCombineGM == nullptr || + sendCountsGM == nullptr || yOutGM == nullptr || bs <= 0 || h <= 0 || topK <= 0 || + moeExpertNum <= 0 || sharedExpertNum < 0 || sharedExpertRankNum < 0 || globalBs <= 0 || + (activeMaskType != Mc2Kernel::ACTIVE_MASK_NONE && activeMaskType != Mc2Kernel::ACTIVE_MASK_TOKEN) || + (quantMode != 0 && quantMode != Mc2Kernel::MXFP8_E5M2_COMM_QUANT && + quantMode != Mc2Kernel::MXFP8_E4M3_COMM_QUANT) || + (quantMode != 0 && expertScalesGM == nullptr)) { + return; + } + if (dtype == static_cast(TileXR::TILEXR_DATA_TYPE_FP16)) { + Mc2Kernel::RunCombine(commArgsGM, expertOutGM, assistInfoForCombineGM, sendCountsGM, + expertScalesGM, xActiveMaskGM, sharedExpertXGM, yOutGM, bs, h, topK, moeExpertNum, + sharedExpertNum, sharedExpertRankNum, globalBs, activeMaskType, quantMode); + } else if (dtype == static_cast(TileXR::TILEXR_DATA_TYPE_BFP16)) { + Mc2Kernel::RunCombine(commArgsGM, expertOutGM, assistInfoForCombineGM, sendCountsGM, + expertScalesGM, xActiveMaskGM, sharedExpertXGM, yOutGM, bs, h, topK, moeExpertNum, + sharedExpertNum, sharedExpertRankNum, globalBs, activeMaskType, quantMode); + } +} + +void launch_tilexr_ep_combine_memory_kernel(uint32_t blockDim, void *stream, GM_ADDR commArgs, + GM_ADDR expertOut, GM_ADDR assistInfoForCombine, GM_ADDR sendCounts, GM_ADDR expertScales, + GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR yOut, int64_t bs, int64_t h, int64_t topK, + int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, + int64_t activeMaskType, int64_t quantMode, int64_t dtype) +{ + tilexr_ep_combine_memory_kernel<<>>(commArgs, expertOut, assistInfoForCombine, + sendCounts, expertScales, xActiveMask, sharedExpertX, yOut, bs, h, topK, moeExpertNum, + sharedExpertNum, sharedExpertRankNum, globalBs, activeMaskType, quantMode, dtype); +} diff --git a/src/ep/kernels/tilexr_ep_dispatch_memory_kernel.cpp b/src/ep/kernels/tilexr_ep_dispatch_memory_kernel.cpp new file mode 100644 index 00000000..bd460ddd --- /dev/null +++ b/src/ep/kernels/tilexr_ep_dispatch_memory_kernel.cpp @@ -0,0 +1,1589 @@ +#include + +#include "adv_api/reduce/sum.h" +#include "comm_args.h" +#include "kernel_operator.h" +#include "tilexr_ep_mxfp8_quant.h" +#include "tilexr_types.h" + +#define FLOAT_OVERFLOW_MODE_CTRL 60 + +namespace Mc2Kernel { +constexpr uint64_t FLAG_FIELD_OFFSET = 768UL * 1024UL; +constexpr uint64_t CUMSUM_CAL_OFFSET = 868UL * 1024UL; +constexpr uint64_t CUMSUM_FLAG_OFFSET = 876UL * 1024UL; +constexpr uint64_t SPLIT_BLOCK_SIZE = 512UL; +constexpr uint64_t SPLIT_BLOCK_COUNT = 128UL; +constexpr int32_t FULL_MESH_MAX_UB_SIZE = 190 * 1024; +constexpr uint32_t COMPARE_COUNT_PER_BLOCK = 256 / sizeof(int32_t); +constexpr uint32_t SPLIT_BLOCK_DATA_SIZE = 480U; +constexpr uint32_t SPLIT_BLOCK_DATA_COUNT = 120U; +constexpr uint32_t SIZE_ALIGN_256 = 256U; +constexpr uint32_t CUMSUM_MAX_CORE_NUM = 8U; +constexpr uint32_t RUNPOS_CALCUMSUM = 2U; +constexpr uint32_t RUNPOS_CUMSUMFLAG = 3U; +constexpr uint32_t RUNPOS_ARRIVECNT = 4U; +constexpr uint8_t VALID_EVENT_FLAG_NUM = 8U; +constexpr uint8_t UB_ALIGN_DATA_COUNT = 8U; +constexpr uint8_t STATUS_COUNT_INDEX = UB_ALIGN_DATA_COUNT - 2U; +constexpr uint8_t STATUS_FLAG_INDEX = UB_ALIGN_DATA_COUNT - 1U; +constexpr uint32_t STATUS_COUNT_PATTERN = 1U << STATUS_COUNT_INDEX; +constexpr uint32_t STATUS_FLAG_PATTERN = 1U << STATUS_FLAG_INDEX; +constexpr uint64_t STATUS_FLAG_DUPLICATE_MASK = 0x8080808080808080ULL; +constexpr uint32_t UB_ALIGN = 32U; +constexpr uint8_t BUFFER_NUM = 2U; +constexpr uint32_t STATE_OFFSET = 32U; +constexpr uint8_t COMBINE_IN_DATA_SIZE = 2U; +constexpr uint64_t WIN_STATE_OFFSET = 384UL * 1024UL; +constexpr uint64_t WIN_ADDR_ALIGN = 512UL; +constexpr uint32_t EXPAND_IDX_INFO = 3U; +constexpr int32_t BITS_PER_BYTE = 8; +constexpr uint64_t A5_MTE_STATE_WIN_SIZE = 1024UL * 1024UL; +constexpr uint64_t OP_CNT_POSUL = 3UL; +constexpr uint32_t ZERONE_STATE_POS = 0U; +constexpr uint32_t OPOSITION_POS = 1U; +constexpr uint32_t TILING_EPRANKID_POS = 2U; +constexpr uint32_t MOE_NUM_POS = 3U; +constexpr uint32_t TILING_WORLDSIZE_POS = 4U; +constexpr uint32_t GLOBALBS_POS = 5U; +constexpr int64_t ACTIVE_MASK_NONE = 0; +constexpr int64_t ACTIVE_MASK_TOKEN = 1; +constexpr int64_t ACTIVE_MASK_EXPERT = 2; +constexpr int64_t MX_QUANT = 4; + +using namespace AscendC; + +template +__aicore__ inline void SyncFunc() +{ + AscendC::TEventID eventId = GetTPipePtr()->FetchEventID(event); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); +} + +class TileXRMemoryContext { +public: + __aicore__ inline void Init(GM_ADDR commArgsGM) + { + args_ = reinterpret_cast<__gm__ TileXR::CommArgs *>(commArgsGM); + rank_ = args_->rank; + rankSize_ = args_->rankSize; + } + + __aicore__ inline uint32_t GetEpRankId() + { + return static_cast(rank_); + } + + __aicore__ inline uint32_t GetEpWorldSize() + { + return static_cast(rankSize_); + } + + __aicore__ inline GM_ADDR GetStatusDataSpaceGm() + { + return GetPeerStateBase(rank_); + } + + __aicore__ inline GM_ADDR GetWindAddrByRankId(int32_t rankId, int32_t) + { + return GetPeerStateBase(rankId) + A5_MTE_STATE_WIN_SIZE; + } + + __aicore__ inline GM_ADDR GetWindStateAddrByRankId(int32_t rankId, int32_t) + { + return GetPeerStateBase(rankId); + } + + __aicore__ inline GM_ADDR GetPeerMemBase(int32_t rankId) + { + return args_->peerMems[rankId]; + } + +private: + __aicore__ inline GM_ADDR GetPeerStateBase(int32_t rankId) + { + return args_->peerMems[rankId] + TileXR::IPC_DATA_OFFSET; + } + + __gm__ TileXR::CommArgs *args_{nullptr}; + int32_t rank_{0}; + int32_t rankSize_{0}; +}; + + + +template +class MoeDistributeDispatchV2FullMesh { +public: + __aicore__ inline MoeDistributeDispatchV2FullMesh() {} + __aicore__ inline void Init(GM_ADDR commArgs, GM_ADDR x, GM_ADDR expertIds, GM_ADDR xActiveMask, + GM_ADDR expandXOut, GM_ADDR dynamicScalesOut, GM_ADDR expertTokenNumsOut, GM_ADDR sendCountsOut, + GM_ADDR assistInfoForCombineOut, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, + int64_t expertTokenNumsType, int64_t activeMaskType, int64_t quantMode, + int64_t expandXOutDtype, int64_t magic, TPipe *pipe); + __aicore__ inline void Run(); + +private: + __aicore__ inline void ExpIdsCopyAndMaskCal(); + __aicore__ inline void TokenActiveMaskCal(); + __aicore__ inline void SetDataStatus(); + __aicore__ inline void CalValidBSCnt(LocalTensor maskStrideTensor); + __aicore__ inline void CalValidExpIdx(LocalTensor maskInputTensor); + __aicore__ inline void SetTilingDataAndCal(); + __aicore__ inline uint32_t InitWinState(GlobalTensor selfDataStatusGMTensor, + uint32_t epRankId, uint32_t epWorldSize, uint32_t moeExpertNum, uint32_t globalBs, uint32_t dataStateSeed, + TBuf<> dataStateBuf); + __aicore__ inline void SendToSharedExpert(TQue inQueue, TBuf<> outBuf); + __aicore__ inline void SendToMoeExpert(TQue inQueue, TBuf<> expertMaskBuf, TBuf<> outBuf); + __aicore__ inline void ExpertActiveMaskInit(); + __aicore__ inline void ExpertActiveMaskCal(); + __aicore__ inline void CalcSendTokenBufNum(TBuf<>& outBuf); + __aicore__ inline void AllToAllDispatchA5(TQue inQueue, TBuf<> expertMaskBuf, TBuf<> outBuf); + __aicore__ inline void AllToAllDispatch(); + __aicore__ inline void CalCumSum(); + __aicore__ inline void WaitCumSumFlag(); + __aicore__ inline void CalAndSendCntByRank(); + __aicore__ inline void BufferInit(); + __aicore__ inline void WaitDispatchClearStatus(); + __aicore__ inline void GatherSumRecvCnt(LocalTensor &gatherMaskOutTensor, + LocalTensor &gatherTmpTensor, LocalTensor &statusSumOutTensor); + __aicore__ inline void CalRecvAndSetFlag(); + __aicore__ inline void WaitDispatch(); + __aicore__ inline void GetCumSum(LocalTensor &outLocal, uint32_t newAivId); + + __aicore__ inline void RunPosRecord(const uint32_t runPos); + __aicore__ inline void LocalWindowCopy(); + __aicore__ inline void SetValidExpertInfo(uint32_t expInfoSize, uint32_t &validNum); + __aicore__ inline uint32_t CheckDataArriveWithFlag(uint32_t srcExpDataIdx, int32_t beginIdx, int32_t copyCnt); + __aicore__ inline void CopyInAndOut(LocalTensor xOutInt32Tensor, + GM_ADDR wAddr, uint32_t index, uint32_t srcExpertId, uint32_t dstPosition, uint32_t arriveCount); + __aicore__ inline void WaitAndFormatOutput(TBuf<> tBuf, uint32_t validNum); + __aicore__ inline void SplitToCore(uint32_t curSendCnt, uint32_t curUseAivNum, uint32_t &startTokenId, + uint32_t &endTokenId, uint32_t &sendTokenNum, bool isFront = true); + __aicore__ inline void FillTriple(LocalTensor &xOutTensor, uint32_t tokenIndex, uint32_t k); + __aicore__ inline void CalTokenSendExpertCnt(uint32_t dstExpertId, int32_t calCnt, int32_t &curExpertCnt); + __aicore__ inline void TokenToExpert(GlobalTensor dstWinGMTensor, TQue inQueue, + uint32_t srcTokenIndex, uint32_t toExpertIndex); + __aicore__ inline void TokenToExpertInQuant(GlobalTensor dstWinGMTensor, + TQue inQueue, uint32_t srcTokenIndex, uint32_t toExpertIndex); + template + __aicore__ inline void QuantMxfp8(LocalTensor &outLocal, LocalTensor &inLocal); + __aicore__ inline GM_ADDR GetWindAddrByRankId(const int32_t rankId) + { + return ctx_.GetWindAddrByRankId(rankId, epRankIdOriginal_) + winDataSizeOffset_; + } + + + __aicore__ inline GM_ADDR GetWindStateAddrByRankId(const int32_t rankId) + { + return ctx_.GetWindStateAddrByRankId(rankId, epRankIdOriginal_) + dataState_ * WIN_STATE_OFFSET; + } + + TPipe *tpipe_{nullptr}; + GlobalTensor xGMTensor_; + GlobalTensor expertIdsGMTensor_; + GlobalTensor expertTokenNumsOutGMTensor_; + GlobalTensor dynamicScalesOutGMTensor_; + GlobalTensor windowInstatusFp32Tensor_; + GlobalTensor xActiveMaskGMTensor_; + GlobalTensor selfRankWinInGMTensor_; + GlobalTensor selfDataStatusGMTensor_; + + LocalTensor statusTensor_; + LocalTensor waitStatusTensor_; + LocalTensor workLocalTensor_; + LocalTensor validExpertIdsTensor_; + LocalTensor validBsIndexTensor_; + LocalTensor statusFp32Tensor_; + LocalTensor gatherMaskTensor_; + LocalTensor statusCleanFp32Tensor_; + LocalTensor sendCntTensor_; + LocalTensor outTensor_; + LocalTensor expertMapTensor_; + LocalTensor expertFinishNumTensor_; + LocalTensor expertLeftNumTensor_; + LocalTensor flagCompResultU8_; + LocalTensor flagCompResultLtU64_; + LocalTensor flagRecvGatherMask_; + LocalTensor cleanUpTensor_; + LocalTensor dataStateLocalTensor_; + LocalTensor xTmpTensor_; + LocalTensor quantTensor_; + LocalTensor quantWorkTensor_; + + LocalTensor flagGatherOutTensor_; + LocalTensor flagRecvTensor_; + + TBuf<> statusBuf_; + TBuf<> recvStatusBuf_; + TBuf<> tokenNumBuf_; + TBuf<> workLocalBuf_; + TBuf<> dstExpBuf_; + TBuf<> subExpBuf_; + TBuf<> gatherMaskTBuf_; + TBuf<> expertIdsBuf_; + TBuf<> waitStatusBuf_; + TBuf<> gatherMaskOutBuf_; + TBuf<> sumCoreBuf_; + TBuf<> sumLocalBuf_; + TBuf<> sumContinueBuf_; + TBuf<> scalarBuf_; + TBuf<> validExpertIndexBuf_; + TBuf<> validBsIndexTBuf_; + TBuf<> calBeginBuf_; + TBuf<> calEndBuf_; + GM_ADDR expandXOutGM_; + GM_ADDR assistInfoForCombineOutGM_; + GM_ADDR sendCountsOutGM_; + GM_ADDR statusSpaceGM_; + GM_ADDR windowGM_; + GM_ADDR recvCntWorkspaceGM_; + GM_ADDR statusDataSpaceGM_; + + + uint32_t syncFlagId_{0}; + uint8_t sendTokenBufNum_{0}; + uint32_t axisBS_{0}; + uint32_t axisMaxBS_{0}; + uint32_t axisH_{0}; + uint32_t axisK_{0}; + uint32_t aivNum_{0}; + uint32_t sharedUsedAivNum_{0}; + uint32_t moeUsedAivNum_{0}; + uint32_t epWorldSize_{0}; + + uint32_t epWorldSizeOriginal_{0}; + int32_t epRankId_{0}; + int32_t epRankIdOriginal_{0}; + uint32_t aivId_{0}; + uint64_t usedTime_[10]; + uint32_t sharedExpertNum_{0}; + uint32_t sharedExpertRankNum_{0}; + uint32_t rankNumPerSharedExpert_{0}; + uint32_t moeExpertNum_{0}; + uint32_t moeExpertRankNum_{0}; + uint32_t moeExpertNumPerRank_{0}; + uint32_t totalExpertNum_{0}; + uint32_t hOutSize_{0}; + uint32_t hOutSizeAlign_{0}; + uint32_t hAlignSize_{0}; + uint32_t startId_; + uint32_t endId_; + uint32_t sendNum_; + uint32_t statusCntAlign_; + uint32_t dataState_{0}; + uint32_t dataStateSeed_{0}; + uint32_t tBufRealSize_{0}; + uint64_t winDataSizeOffset_{0}; + uint64_t expertPerSizeOnWin_{0}; + uint64_t activeMaskBsCnt_{0}; + uint64_t sendToMoeExpTokenCnt_{0}; + bool isTokenMaskFlag_ = false; + bool isExpertMaskFlag_ = false; + bool isShareExpertRankFlag_ = false; + uint64_t totalWinSize_{0}; + uint32_t expertTokenNumsType_{1}; + int32_t expertIdsCnt_{0}; + int32_t tokenQuantAlign_{0}; + uint32_t blockCntPerToken_{0}; + uint32_t axisHCommu_{0}; + uint32_t hCommuSize_{0}; + uint32_t scaleOutBytes_{0}; + uint32_t quantWorkBytes_{0}; + uint32_t quantTensorBytes_{0}; + bool useMxfp8_{false}; + bool useMxfp8E4M3_{false}; + + uint32_t expertIdsBufSize_{0}; + uint32_t rscvStatusNum_{0}; + uint32_t startStatusIndex_{0}; + uint32_t endStatusIndex_{0}; + uint32_t recStatusNumPerCore_{0}; + uint32_t aivUsedCumSum_{0}; + uint32_t aivUsedAllToAll_{0}; + uint32_t maxSize_{0}; + uint32_t expertIdsSize_{0}; + uint32_t globalBS_{0}; + uint32_t copyInAxisH_{0}; + uint32_t copyOutAxisH_{0}; + uint64_t totalUbSize_{0}; + TileXRMemoryContext ctx_; + + DataCopyParams hCopyParams_; + DataCopyParams dataStateParams_{1U, sizeof(uint32_t), 0U, 0U}; +}; + +template +__aicore__ inline uint32_t MoeDistributeDispatchV2FullMesh::InitWinState(GlobalTensor selfDataStatusGMTensor, + uint32_t epRankId, uint32_t epWorldSize, uint32_t moeExpertNum, uint32_t globalBs, uint32_t dataStateSeed, + TBuf<> dataStateBuf) +{ + // usedTime_[0] = AscendC::GetSystemCycle(); + LocalTensor dataStateLocalTensor64 = dataStateBuf.Get(); + LocalTensor dataStateLocalTensor = dataStateBuf.Get(); + DataCopy(dataStateLocalTensor, selfDataStatusGMTensor, UB_ALIGN / sizeof(uint32_t)); + SyncFunc(); + // usedTime_[1] = AscendC::GetSystemCycle(); + uint32_t dataState = dataStateSeed; + dataStateLocalTensor.SetValue(ZERONE_STATE_POS, dataState == 0 ? 1 : 0); + dataStateLocalTensor.SetValue(OPOSITION_POS, 1); + dataStateLocalTensor.SetValue(TILING_EPRANKID_POS, epRankId); + dataStateLocalTensor.SetValue(MOE_NUM_POS, moeExpertNum); + dataStateLocalTensor.SetValue(TILING_WORLDSIZE_POS, epWorldSize); + dataStateLocalTensor.SetValue(GLOBALBS_POS, globalBs); + uint32_t opCnt = dataStateLocalTensor64.GetValue(OP_CNT_POSUL); + dataStateLocalTensor64.SetValue(OP_CNT_POSUL, opCnt + 1); + // usedTime_[2] = AscendC::GetSystemCycle(); + SyncFunc(); + DataCopy(selfDataStatusGMTensor, dataStateLocalTensor, UB_ALIGN / sizeof(uint32_t)); + return dataState; +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::SetTilingDataAndCal() +{ + copyInAxisH_ = axisH_; + copyOutAxisH_ = axisH_; + isShareExpertRankFlag_ = epRankId_ < static_cast(sharedExpertRankNum_); + if (sharedExpertNum_ > 0) { + rankNumPerSharedExpert_ = sharedExpertRankNum_ / sharedExpertNum_; + } + moeExpertRankNum_ = epWorldSize_ - sharedExpertRankNum_; + moeExpertNumPerRank_ = moeExpertNum_ / moeExpertRankNum_; + expertIdsCnt_ = axisBS_ * axisK_; + hOutSize_ = copyOutAxisH_ * (useMxfp8_ ? sizeof(uint8_t) : sizeof(XType)); + if (useMxfp8_) { + scaleOutBytes_ = ((Ceil(axisH_, 32U) + 1U) / 2U) * 2U; + hAlignSize_ = Ceil(axisH_, 128U) * 128U * sizeof(XType); + const uint32_t quantPayloadBytes = Ceil(hOutSize_, 256U) * 256U + scaleOutBytes_; + tokenQuantAlign_ = Ceil(quantPayloadBytes, UB_ALIGN) * UB_ALIGN / sizeof(int32_t); + quantWorkBytes_ = Ceil( + Ceil(scaleOutBytes_, 32U) * 32U * sizeof(float) + scaleOutBytes_ * sizeof(uint16_t), UB_ALIGN) * UB_ALIGN; + } else { + hAlignSize_ = Ceil(axisH_ * sizeof(XType), UB_ALIGN) * UB_ALIGN; + tokenQuantAlign_ = hAlignSize_ / sizeof(int32_t); + } + hOutSizeAlign_ = tokenQuantAlign_ * sizeof(int32_t) + UB_ALIGN; + quantTensorBytes_ = Ceil(hOutSizeAlign_, UB_ALIGN) * UB_ALIGN; + blockCntPerToken_ = Ceil(hOutSizeAlign_, SPLIT_BLOCK_DATA_SIZE); + hCommuSize_ = blockCntPerToken_ * SPLIT_BLOCK_SIZE; + axisHCommu_ = hCommuSize_; + expertPerSizeOnWin_ = axisMaxBS_ * hCommuSize_; + rscvStatusNum_ = isShareExpertRankFlag_ ? epWorldSize_ : epWorldSize_ * moeExpertNumPerRank_; + totalExpertNum_ = sharedExpertRankNum_ + moeExpertNum_; + statusCntAlign_ = Ceil(totalExpertNum_, UB_ALIGN_DATA_COUNT) * UB_ALIGN_DATA_COUNT; + aivUsedCumSum_ = totalExpertNum_ / 16; + aivUsedCumSum_ = aivUsedCumSum_ == 0 ? 1 : aivUsedCumSum_; + aivUsedCumSum_ = aivUsedCumSum_ >= aivNum_ / 2 ? aivNum_ / 2 : aivUsedCumSum_; + aivUsedCumSum_ = aivUsedCumSum_ >= CUMSUM_MAX_CORE_NUM ? CUMSUM_MAX_CORE_NUM : aivUsedCumSum_; + aivUsedCumSum_ = aivUsedCumSum_ >= rscvStatusNum_ ? rscvStatusNum_ : aivUsedCumSum_; + aivUsedAllToAll_ = aivNum_ - aivUsedCumSum_; + if (sharedExpertRankNum_ != 0U) { + sharedUsedAivNum_ = aivUsedAllToAll_ * sharedExpertNum_ / (axisK_ + sharedExpertNum_); + if (sharedUsedAivNum_ == 0) { + sharedUsedAivNum_ = 1; + } + } + moeUsedAivNum_ = aivUsedAllToAll_ - sharedUsedAivNum_; + + expertIdsSize_ = Ceil(expertIdsCnt_ * sizeof(int32_t), UB_ALIGN) * UB_ALIGN; + uint32_t expertMaskRowBytes = Ceil(axisK_ * sizeof(bool), UB_ALIGN) * UB_ALIGN; + uint32_t expertMaskBytes = axisBS_ * expertMaskRowBytes * sizeof(half); + maxSize_ = expertIdsSize_ > expertMaskBytes ? expertIdsSize_ : expertMaskBytes; + if (useMxfp8_ && maxSize_ < quantWorkBytes_) { + maxSize_ = quantWorkBytes_; + } + totalUbSize_ = FULL_MESH_MAX_UB_SIZE; +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::SetDataStatus() +{ + statusDataSpaceGM_ = ctx_.GetStatusDataSpaceGm(); + selfDataStatusGMTensor_.SetGlobalBuffer( + reinterpret_cast<__gm__ uint32_t *>(statusDataSpaceGM_ + FLAG_FIELD_OFFSET + aivId_ * WIN_ADDR_ALIGN)); + TBuf<> dataStateBuf; + tpipe_->InitBuffer(dataStateBuf, UB_ALIGN); + + dataState_ = InitWinState(selfDataStatusGMTensor_, epRankId_, epWorldSize_, moeExpertNum_, globalBS_, + dataStateSeed_, dataStateBuf); + uint64_t hSizeAlignCombine = + Ceil(axisH_ * COMBINE_IN_DATA_SIZE, SPLIT_BLOCK_DATA_SIZE) * SPLIT_BLOCK_SIZE; + winDataSizeOffset_ = dataState_ * (totalWinSize_ / BUFFER_NUM) + + axisMaxBS_ * (axisK_ + sharedExpertNum_) * hSizeAlignCombine; +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::Init(GM_ADDR commArgs, GM_ADDR x, + GM_ADDR expertIds, GM_ADDR xActiveMask, GM_ADDR expandXOut, GM_ADDR dynamicScalesOut, + GM_ADDR expertTokenNumsOut, + GM_ADDR sendCountsOut, GM_ADDR assistInfoForCombineOut, int64_t bs, int64_t h, int64_t topK, + int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, + int64_t expertTokenNumsType, int64_t activeMaskType, int64_t quantMode, + int64_t expandXOutDtype, int64_t magic, TPipe *pipe) +{ + AscendC::SetCtrlSpr(0); + tpipe_ = pipe; + tpipe_->InitBuffer(calBeginBuf_, UB_ALIGN); + aivId_ = GetBlockIdx(); + ctx_.Init(commArgs); + + axisBS_ = static_cast(bs); + axisH_ = static_cast(h); + axisK_ = static_cast(topK); + moeExpertNum_ = static_cast(moeExpertNum); + sharedExpertNum_ = static_cast(sharedExpertNum); + sharedExpertRankNum_ = static_cast(sharedExpertRankNum); + globalBS_ = static_cast(globalBs); + expertTokenNumsType_ = static_cast(expertTokenNumsType); + useMxfp8_ = quantMode == MX_QUANT; + useMxfp8E4M3_ = expandXOutDtype == TileXR::TILEXR_DATA_TYPE_FP8E4M3; + dataStateSeed_ = static_cast(magic) & 1U; + epRankId_ = static_cast(ctx_.GetEpRankId()); + epRankIdOriginal_ = epRankId_; + epWorldSize_ = ctx_.GetEpWorldSize(); + epWorldSizeOriginal_ = epWorldSize_; + aivNum_ = GetBlockNum(); + axisMaxBS_ = globalBS_ / epWorldSize_; + isTokenMaskFlag_ = activeMaskType == ACTIVE_MASK_TOKEN; + isExpertMaskFlag_ = activeMaskType == ACTIVE_MASK_EXPERT; + + xGMTensor_.SetGlobalBuffer(reinterpret_cast<__gm__ XType *>(x)); + xActiveMaskGMTensor_.SetGlobalBuffer(reinterpret_cast<__gm__ bool *>(xActiveMask)); + expertIdsGMTensor_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(expertIds)); + expertTokenNumsOutGMTensor_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(expertTokenNumsOut)); + dynamicScalesOutGMTensor_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(dynamicScalesOut)); + + SetTilingDataAndCal(); + uint64_t workspaceStatusNum = static_cast(epWorldSize_) * moeExpertNumPerRank_; + uint64_t workspaceBytes = Ceil(aivNum_ * workspaceStatusNum * sizeof(int32_t), UB_ALIGN) * UB_ALIGN; + totalWinSize_ = TileXR::IPC_BUFF_MAX_SIZE - A5_MTE_STATE_WIN_SIZE - workspaceBytes; + SetDataStatus(); + + expandXOutGM_ = expandXOut; + assistInfoForCombineOutGM_ = assistInfoForCombineOut; + sendCountsOutGM_ = sendCountsOut; + recvCntWorkspaceGM_ = ctx_.GetStatusDataSpaceGm() + TileXR::IPC_BUFF_MAX_SIZE - workspaceBytes; + statusSpaceGM_ = GetWindStateAddrByRankId(epRankIdOriginal_); + windowInstatusFp32Tensor_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(statusSpaceGM_)); + selfRankWinInGMTensor_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(statusDataSpaceGM_)); + windowGM_ = GetWindAddrByRankId(epRankIdOriginal_); + hCopyParams_ = {1U, static_cast(copyInAxisH_ * sizeof(XType)), 0U, 0U}; + dataStateParams_ = {1U, sizeof(uint32_t), 0U, 0U}; + // usedTime_[3] = AscendC::GetSystemCycle(); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::FillTriple( + LocalTensor &xOutTensor, uint32_t tokenIndex, uint32_t k) +{ + LocalTensor xOutTint32 = xOutTensor.template ReinterpretCast(); + xOutTint32(tokenQuantAlign_) = epRankId_; + xOutTint32(tokenQuantAlign_ + 1) = tokenIndex; + xOutTint32(tokenQuantAlign_ + 2) = k; +} + +template +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::QuantMxfp8( + LocalTensor &outLocal, LocalTensor &inLocal) +{ + __ubuf__ XType *srcAddr = reinterpret_cast<__ubuf__ XType *>(inLocal.GetPhyAddr()); + __ubuf__ uint16_t *maxExpAddr = reinterpret_cast<__ubuf__ uint16_t *>(quantWorkTensor_.GetPhyAddr()); + __ubuf__ uint16_t *halfScaleAddr = reinterpret_cast<__ubuf__ uint16_t *>( + quantWorkTensor_[Ceil(scaleOutBytes_, 32U) * 32U].GetPhyAddr()); + __ubuf__ int8_t *outAddr = reinterpret_cast<__ubuf__ int8_t *>(outLocal.GetPhyAddr()); + __ubuf__ uint16_t *scaleAddr = reinterpret_cast<__ubuf__ uint16_t *>( + outLocal[Ceil(axisH_, 256U) * 256U].GetPhyAddr()); + + TileXRMxfp8Quant::ComputeMaxExp(srcAddr, maxExpAddr, axisH_); + TileXRMxfp8Quant::ComputeScale(maxExpAddr, scaleAddr, halfScaleAddr, scaleOutBytes_); + TileXRMxfp8Quant::ComputeFp8Data( + srcAddr, halfScaleAddr, outAddr, axisH_); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::TokenToExpertInQuant( + GlobalTensor dstWinGMTensor, TQue inQueue, + uint32_t srcTokenIndex, uint32_t toExpertIndex) +{ + DataCopyPadParams copyPadParams{true, 0U, 0U, 0U}; + LocalTensor xInTensor = inQueue.AllocTensor(); + DataCopyPad(xInTensor, xGMTensor_[srcTokenIndex * axisH_], hCopyParams_, copyPadParams); + inQueue.EnQue(xInTensor); + xInTensor = inQueue.DeQue(); + if (useMxfp8E4M3_) { + QuantMxfp8(quantTensor_, xInTensor); + } else { + QuantMxfp8(quantTensor_, xInTensor); + } + inQueue.FreeTensor(xInTensor); + SyncFunc(); + FillTriple(quantTensor_, srcTokenIndex, toExpertIndex); + SyncFunc(); + AscendC::WaitFlag(syncFlagId_ % sendTokenBufNum_); + LocalTensor quantTensorInt32 = quantTensor_.template ReinterpretCast(); + LocalTensor outTensorInt32 = + outTensor_[(syncFlagId_ % sendTokenBufNum_) * axisHCommu_].template ReinterpretCast(); + Copy(outTensorInt32, quantTensorInt32, uint64_t(64), uint8_t(blockCntPerToken_), {1, 1, 16, 15}); + Copy(outTensorInt32[64], quantTensorInt32[64], uint64_t(56), + uint8_t(blockCntPerToken_), {1, 1, 16, 15}); + AscendC::SetFlag(syncFlagId_ % sendTokenBufNum_); + AscendC::WaitFlag(syncFlagId_ % sendTokenBufNum_); + DataCopy(dstWinGMTensor, outTensor_[(syncFlagId_ % sendTokenBufNum_) * axisHCommu_], axisHCommu_); + AscendC::SetFlag(syncFlagId_ % sendTokenBufNum_); + ++syncFlagId_; +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::TokenToExpert( + GlobalTensor dstWinGMTensor, TQue inQueue, + uint32_t srcTokenIndex, uint32_t toExpertIndex) +{ + DataCopyPadParams copyPadParams{false, 0U, 0U, 0U}; + LocalTensor xInTensor = inQueue.AllocTensor(); + DataCopyPad(xInTensor, xGMTensor_[srcTokenIndex * axisH_], hCopyParams_, copyPadParams); + inQueue.EnQue(xInTensor); + xInTensor = inQueue.DeQue(); + SyncFunc(); + LocalTensor xInTensorBytes = xInTensor.template ReinterpretCast(); + FillTriple(xInTensorBytes, srcTokenIndex, toExpertIndex); + SyncFunc(); + AscendC::WaitFlag(syncFlagId_ % sendTokenBufNum_); + LocalTensor xInTensorInt32 = xInTensorBytes.template ReinterpretCast(); + LocalTensor outTensorInt32 = + (outTensor_[(syncFlagId_ % sendTokenBufNum_) * axisHCommu_]).template ReinterpretCast(); + + Copy(outTensorInt32, xInTensorInt32, uint64_t(64), uint8_t(blockCntPerToken_), {1, 1, 16, 15}); + + Copy(outTensorInt32[64], xInTensorInt32[64], uint64_t(56), uint8_t(blockCntPerToken_), {1, 1, 16, 15}); + inQueue.FreeTensor(xInTensor); + AscendC::SetFlag(syncFlagId_ % sendTokenBufNum_); + AscendC::WaitFlag(syncFlagId_ % sendTokenBufNum_); + DataCopy(dstWinGMTensor, outTensor_[(syncFlagId_ % sendTokenBufNum_) * axisHCommu_], axisHCommu_); + AscendC::SetFlag(syncFlagId_ % sendTokenBufNum_); + syncFlagId_ ++; +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::SplitToCore( + uint32_t curSendCnt, uint32_t curUseAivNum, uint32_t &startTokenId, + uint32_t &endTokenId, uint32_t &sendTokenNum, bool isFront) + +{ + sendTokenNum = curSendCnt / curUseAivNum; + uint32_t remainderTokenNum = curSendCnt % curUseAivNum; + uint32_t newAivId; + if (isFront) { + newAivId = aivId_; + } else if (aivId_ >= aivUsedAllToAll_) { + newAivId = aivId_ - aivUsedAllToAll_; + } else { + newAivId = aivId_ - moeUsedAivNum_; + } + startTokenId = sendTokenNum * newAivId; + if (newAivId < remainderTokenNum) { + sendTokenNum += 1; + startTokenId += newAivId; + } else { + startTokenId += remainderTokenNum; + } + endTokenId = startTokenId + sendTokenNum; +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::SendToSharedExpert(TQue inQueue, TBuf<> outBuf) +{ + + LocalTensor outTensorFp32 = outBuf.Get(); + Duplicate(outTensorFp32, float(1), hCommuSize_ * sendTokenBufNum_ / sizeof(float)); + PipeBarrier(); + + uint32_t startTokenId, endTokenId, sendTokenNum; + uint32_t curSendCnt = activeMaskBsCnt_ * sharedExpertNum_; + SplitToCore(curSendCnt, sharedUsedAivNum_, startTokenId, endTokenId, sendTokenNum, false); + if (startTokenId >= curSendCnt) {return;} + + GlobalTensor dstWinGMTensor; + uint32_t idInSharedGroup = epRankId_ % rankNumPerSharedExpert_; + syncFlagId_ = 0; + for (int i = 0; i < sendTokenBufNum_; i ++) { + AscendC::SetFlag(i % sendTokenBufNum_); + } + for (uint32_t virtualTokenIndex = startTokenId; virtualTokenIndex < endTokenId; ++virtualTokenIndex) { + uint32_t sendTokenIndex = virtualTokenIndex % activeMaskBsCnt_; + uint32_t toSharedExpertIndex = virtualTokenIndex / activeMaskBsCnt_; + int32_t toRankId = idInSharedGroup + toSharedExpertIndex * rankNumPerSharedExpert_; + dstWinGMTensor.SetGlobalBuffer((__gm__ uint8_t *)(uint64_t(GetWindAddrByRankId(toRankId)) + + expertPerSizeOnWin_ * epRankId_ + sendTokenIndex * hCommuSize_)); + uint32_t srcTokenIndex = sendTokenIndex; + if (isExpertMaskFlag_) { + srcTokenIndex = validBsIndexTensor_.GetValue(sendTokenIndex); + } + if (useMxfp8_) { + TokenToExpertInQuant(dstWinGMTensor, inQueue, srcTokenIndex, axisK_ + toSharedExpertIndex); + } else { + TokenToExpert(dstWinGMTensor, inQueue, srcTokenIndex, axisK_ + toSharedExpertIndex); + } + } + for (int i = 0; i < sendTokenBufNum_; i ++) { + AscendC::WaitFlag(i % sendTokenBufNum_); + } +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::SendToMoeExpert(TQue inQueue, + TBuf<> expertMaskBuf, TBuf<> outBuf) +{ + + LocalTensor outTensorFp32 = outBuf.Get(); + Duplicate(outTensorFp32, float(1), hCommuSize_ * sendTokenBufNum_ / sizeof(float)); + uint32_t validTokenNum = isTokenMaskFlag_ ? (activeMaskBsCnt_ * axisK_) : expertIdsCnt_; + GlobalTensor dstWinGMTensor; + + int32_t dstTokenIdx = 0; + syncFlagId_ = 0; + + for (int i = 0; i < sendTokenBufNum_; i ++) { + AscendC::SetFlag(i % sendTokenBufNum_); + } + + for (int32_t index = aivId_; index < validTokenNum; index += moeUsedAivNum_) { + int32_t tokenId = index / axisK_; + int32_t topKId = index % axisK_; + int32_t expertId = validExpertIdsTensor_(index); + if (expertId >= moeExpertNum_ || expertId < 0) + continue; + int32_t toRankId = expertId / moeExpertNumPerRank_ + sharedExpertRankNum_; + CalTokenSendExpertCnt(expertId, index, dstTokenIdx); + + dstWinGMTensor.SetGlobalBuffer((__gm__ uint8_t *)(uint64_t(GetWindAddrByRankId(toRankId)) + + expertPerSizeOnWin_ * ((epRankId_ + toRankId) % epWorldSize_ * moeExpertNumPerRank_ + + expertId % moeExpertNumPerRank_) + dstTokenIdx * hCommuSize_)); + if (useMxfp8_) { + TokenToExpertInQuant(dstWinGMTensor, inQueue, tokenId, topKId); + } else { + TokenToExpert(dstWinGMTensor, inQueue, tokenId, topKId); + } + } + + for (int i = 0; i < sendTokenBufNum_; i ++) { + AscendC::WaitFlag(i % sendTokenBufNum_); + } +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::CalcSendTokenBufNum(TBuf<>& outBuf) +{ + tpipe_->InitBuffer(calEndBuf_, UB_ALIGN); + uint64_t beiginUbAddr = (calBeginBuf_.Get()).GetPhyAddr(); + uint64_t endUbAddr = (calEndBuf_.Get()).GetPhyAddr(); + uint64_t remainUbSize = totalUbSize_ - (endUbAddr - beiginUbAddr + UB_ALIGN); + + sendTokenBufNum_ = remainUbSize / hCommuSize_; + if (sendTokenBufNum_ > VALID_EVENT_FLAG_NUM) + sendTokenBufNum_ = VALID_EVENT_FLAG_NUM; + if (sendTokenBufNum_ == 0) { return; } + + + tpipe_->InitBuffer(outBuf, hCommuSize_ * sendTokenBufNum_); + outTensor_ = outBuf.Get(); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::AllToAllDispatch() +{ + TQue inQueue; + TBuf<> outBuf, expertMaskBuf, inQueueCleanBuf, calTempBuf; + TBuf<> quantBuf, quantWorkBuf, quantAuxBuf; + expertIdsBufSize_ = Ceil(expertIdsCnt_ * sizeof(int32_t), SIZE_ALIGN_256) * SIZE_ALIGN_256; + AscendC::TBufPool tbufPool0, tbufPool1; + tpipe_->InitBufPool(tbufPool0, BUFFER_NUM * hAlignSize_); + tpipe_->InitBufPool(tbufPool1, BUFFER_NUM * hAlignSize_, tbufPool0); + tbufPool0.InitBuffer(inQueue, BUFFER_NUM, hAlignSize_); + tbufPool1.InitBuffer(inQueueCleanBuf, BUFFER_NUM * hAlignSize_); + if (useMxfp8_) { + LocalTensor cleanTensor = inQueueCleanBuf.Get(); + Duplicate(cleanTensor, 0, BUFFER_NUM * hAlignSize_); + } + + uint32_t calTokenIdxBuffSize = Ceil(axisBS_ * axisK_ * sizeof(int32_t), UB_ALIGN) * UB_ALIGN; + tpipe_->InitBuffer(expertIdsBuf_, expertIdsBufSize_); + bool needMaskCalFlag = isTokenMaskFlag_ || isExpertMaskFlag_; + if (needMaskCalFlag) { + tpipe_->InitBuffer(gatherMaskTBuf_, expertIdsBufSize_); + } + if (useMxfp8_) { + tpipe_->InitBuffer(quantBuf, quantTensorBytes_); + tpipe_->InitBuffer(quantWorkBuf, maxSize_); + tpipe_->InitBuffer(quantAuxBuf, maxSize_); + quantTensor_ = quantBuf.Get(); + quantWorkTensor_ = quantWorkBuf.Get(); + dstExpBuf_ = quantWorkBuf; + subExpBuf_ = quantAuxBuf; + } else if (needMaskCalFlag) { + tpipe_->InitBuffer(dstExpBuf_, maxSize_); + tpipe_->InitBuffer(subExpBuf_, maxSize_); + } else { + tpipe_->InitBuffer(dstExpBuf_, calTokenIdxBuffSize); + tpipe_->InitBuffer(subExpBuf_, calTokenIdxBuffSize); + } + tpipe_->InitBuffer(calTempBuf, calTokenIdxBuffSize); + workLocalTensor_ = calTempBuf.Get(); + ExpIdsCopyAndMaskCal(); + if (activeMaskBsCnt_ == 0) { + return; + } + AllToAllDispatchA5(inQueue, expertMaskBuf, outBuf); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::AllToAllDispatchA5( + TQue inQueue, TBuf<> expertMaskBuf, TBuf<> outBuf) +{ + CalcSendTokenBufNum(outBuf); + if ((aivId_ >= moeUsedAivNum_) && (sharedExpertRankNum_ != 0)) { + SendToSharedExpert(inQueue, outBuf); + } else { + SendToMoeExpert(inQueue, expertMaskBuf, outBuf); + } +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::CalTokenSendExpertCnt(uint32_t dstExpertId, int32_t calCnt, int32_t &curExpertCnt) +{ + if (calCnt < axisK_) { + curExpertCnt = 0; + return; + } + LocalTensor dstExpIdTensor = dstExpBuf_.Get(); + LocalTensor subExpIdTensor = subExpBuf_.Get(); + Duplicate(dstExpIdTensor, dstExpertId, calCnt); + PipeBarrier(); + Sub(subExpIdTensor, validExpertIdsTensor_, dstExpIdTensor, calCnt); + PipeBarrier(); + LocalTensor tmpFp32 = subExpIdTensor.ReinterpretCast(); + LocalTensor tmpoutFp32 = dstExpIdTensor.ReinterpretCast(); + Abs(tmpoutFp32, tmpFp32, calCnt); + PipeBarrier(); + Mins(subExpIdTensor, dstExpIdTensor, 1, calCnt); + PipeBarrier(); + ReduceSum(tmpoutFp32, tmpFp32, workLocalTensor_, calCnt); + SyncFunc(); + int32_t curOtherExpertCnt = dstExpIdTensor(0); + if (calCnt >= curOtherExpertCnt) { + curExpertCnt = calCnt - curOtherExpertCnt; + } else { + curExpertCnt = 0; + } +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::CalAndSendCntByRank() +{ + + uint32_t startRankId, endRankId, sendRankNum; + uint32_t startExpertId, endExpertId, sendExpertNum; + uint32_t maskCnt = isTokenMaskFlag_ ? (activeMaskBsCnt_ * axisK_) : expertIdsCnt_; + uint64_t mask[2] = {STATUS_FLAG_DUPLICATE_MASK, 0}; + Duplicate(statusTensor_, 0, statusCntAlign_ * UB_ALIGN_DATA_COUNT); + PipeBarrier(); + Duplicate(statusTensor_, 0x3F800000, mask, statusCntAlign_ / 8, 1, 8); + SyncFunc(); + + GlobalTensor rankGMTensor; + uint32_t newAivId = aivId_ - aivUsedAllToAll_; + + for (uint32_t dstRankId = newAivId; dstRankId < epWorldSize_; dstRankId += aivUsedCumSum_) { + + if (dstRankId >= sharedExpertRankNum_) { + startExpertId = (dstRankId - sharedExpertRankNum_) * moeExpertNumPerRank_; + endExpertId = startExpertId + moeExpertNumPerRank_; + for (uint32_t curMoeExpertId = startExpertId; curMoeExpertId < endExpertId; ++curMoeExpertId) { + int32_t curExpertCnt = 0; + int32_t cntPosIndex = + (curMoeExpertId + sharedExpertRankNum_) * UB_ALIGN_DATA_COUNT + STATUS_COUNT_INDEX; + + if (sendToMoeExpTokenCnt_ > 0) { + CalTokenSendExpertCnt(curMoeExpertId, maskCnt, curExpertCnt); + } + statusTensor_.SetValue(cntPosIndex, curExpertCnt); + } + } else { + int32_t curExpertCnt = 0; + int32_t cntPosIndex = dstRankId * UB_ALIGN_DATA_COUNT + STATUS_COUNT_INDEX; + + if (activeMaskBsCnt_ > 0) { + if (dstRankId % rankNumPerSharedExpert_ == epRankId_ % rankNumPerSharedExpert_) { + curExpertCnt = activeMaskBsCnt_; + } + } + statusTensor_.SetValue(cntPosIndex, curExpertCnt); + } + } + if (newAivId < epWorldSize_) + SyncFunc(); + for (uint32_t dstRankId = newAivId; dstRankId < epWorldSize_; dstRankId += aivUsedCumSum_) { + uint32_t offset = STATE_OFFSET * epRankId_; + GM_ADDR rankGM = (__gm__ uint8_t*)(GetWindStateAddrByRankId(dstRankId) + offset); + rankGMTensor.SetGlobalBuffer((__gm__ int32_t*)rankGM); + if (dstRankId >= sharedExpertRankNum_) { + uint32_t startStatusIdx = + (dstRankId - sharedExpertRankNum_) * moeExpertNumPerRank_ + sharedExpertRankNum_; + DataCopyParams cntCopyParams = {uint16_t(moeExpertNumPerRank_), 1U, 0U, uint16_t(epWorldSize_ - 1)}; + DataCopy(rankGMTensor, statusTensor_[startStatusIdx * UB_ALIGN_DATA_COUNT], cntCopyParams); + } else { + DataCopy(rankGMTensor, statusTensor_[dstRankId * UB_ALIGN_DATA_COUNT], UB_ALIGN_DATA_COUNT); + } + } +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::BufferInit() +{ + uint32_t waitStatusBufSize = Ceil((recStatusNumPerCore_ * UB_ALIGN), SIZE_ALIGN_256) * SIZE_ALIGN_256; + tpipe_->InitBuffer(waitStatusBuf_, waitStatusBufSize); + uint64_t recStatusNumPerCoreSpace = Ceil(recStatusNumPerCore_ * sizeof(float), UB_ALIGN) * UB_ALIGN; + uint64_t recvWinBlockNumSpace = epWorldSize_ * moeExpertNumPerRank_ * sizeof(float); + uint64_t gatherMaskOutSize = (recStatusNumPerCoreSpace > recvWinBlockNumSpace) ? recStatusNumPerCoreSpace : recvWinBlockNumSpace; + uint64_t sumContinueAlignSize = Ceil((aivNum_ * sizeof(float)), UB_ALIGN) * UB_ALIGN; + tpipe_->InitBuffer(gatherMaskOutBuf_, gatherMaskOutSize); + tpipe_->InitBuffer(sumCoreBuf_, aivNum_ * UB_ALIGN); + tpipe_->InitBuffer(sumLocalBuf_, aivNum_ * UB_ALIGN); + tpipe_->InitBuffer(sumContinueBuf_, sumContinueAlignSize); + tpipe_->InitBuffer(scalarBuf_, UB_ALIGN * 3); + uint32_t statusBufSize = rscvStatusNum_ * UB_ALIGN; + uint32_t tokenNumBufSize = Ceil(moeExpertNumPerRank_ * sizeof(int64_t), UB_ALIGN) * UB_ALIGN; + uint32_t workLocalBufSize = Ceil(epWorldSize_ * sizeof(float), UB_ALIGN) * UB_ALIGN; + tpipe_->InitBuffer(recvStatusBuf_, statusBufSize); + tpipe_->InitBuffer(tokenNumBuf_, tokenNumBufSize); + tpipe_->InitBuffer(workLocalBuf_, workLocalBufSize); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::WaitDispatchClearStatus() +{ + SyncFunc(); + DataCopyParams intriOutParams{static_cast(recStatusNumPerCore_), 1, 0, 0}; + uint64_t duplicateMask[2] = {STATUS_FLAG_DUPLICATE_MASK, 0}; + LocalTensor cleanStateTensor = waitStatusBuf_.Get(); + SyncFunc(); + Duplicate(cleanStateTensor, 0, duplicateMask, Ceil(recStatusNumPerCore_, 8), 1, 8); + SyncFunc(); + DataCopy(windowInstatusFp32Tensor_[startStatusIndex_ * STATE_OFFSET / sizeof(float)], + cleanStateTensor.ReinterpretCast(), intriOutParams); + SyncFunc(); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::GatherSumRecvCnt( + LocalTensor &gatherMaskOutTensor, LocalTensor &gatherTmpTensor, + LocalTensor &statusSumOutTensor) +{ + gatherTmpTensor.SetValue(0, STATUS_COUNT_PATTERN); + uint32_t mask = UB_ALIGN_DATA_COUNT; + SyncFunc(); + + + + uint64_t recvCnt = 0; + GatherMask(gatherMaskOutTensor, statusFp32Tensor_, gatherTmpTensor, true, mask, + {1, (uint16_t)recStatusNumPerCore_, 1, 0}, recvCnt); + PipeBarrier(); + + + uint32_t recStatusNumPerCoreInner = Ceil(recStatusNumPerCore_ * sizeof(float), UB_ALIGN) + * UB_ALIGN / sizeof(float); + SumParams sumParams{1, recStatusNumPerCoreInner, recStatusNumPerCore_}; + Sum(statusSumOutTensor, gatherMaskOutTensor, sumParams); + SyncFunc(); + float sumOfRecvCnt = statusSumOutTensor.ReinterpretCast().GetValue(0); + + + uint32_t newAivId = aivId_ - aivUsedAllToAll_; + + LocalTensor sumCoreFP32Tensor = sumCoreBuf_.Get(); + uint64_t maskArrayCount[2] = {0x0101010101010101, 0}; + uint8_t repeatTimes = Ceil(aivUsedCumSum_, 8); + + Duplicate(sumCoreFP32Tensor, sumOfRecvCnt, maskArrayCount, repeatTimes, 1, 8); + uint64_t maskArrayFlag[2] = {0x0202020202020202, 0}; + Duplicate(sumCoreFP32Tensor, static_cast(1.0), maskArrayFlag, repeatTimes, 1, 8); + DataCopyParams sumIntriParams{static_cast(aivUsedCumSum_), 1, 0, 0}; + SyncFunc(); + DataCopy(selfRankWinInGMTensor_[(CUMSUM_CAL_OFFSET + newAivId * aivUsedCumSum_ * UB_ALIGN) / sizeof(float)], sumCoreFP32Tensor, sumIntriParams); + SyncFunc(); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::GetCumSum(LocalTensor &outLocal, uint32_t newAivId) + +{ + outLocal = gatherMaskOutBuf_.Get(); + DataCopyParams sumIntriParams{static_cast(aivUsedCumSum_), 1, static_cast(aivUsedCumSum_ - 1), 0}; + LocalTensor sumLocalTensor = sumLocalBuf_.Get(); + LocalTensor gatherSumPattern = scalarBuf_.GetWithOffset(UB_ALIGN / sizeof(uint32_t), 0); + LocalTensor sumContinueTensor = sumContinueBuf_.Get(); + LocalTensor recvCntSumOutTensor = scalarBuf_.GetWithOffset(UB_ALIGN / sizeof(float), UB_ALIGN); + + uint32_t mask = 2; + uint64_t recvCnt = 0; + uint32_t innerSumParams = Ceil(aivUsedCumSum_ * sizeof(float), UB_ALIGN) * UB_ALIGN / sizeof(float); + SumParams sumParams{1, innerSumParams, aivUsedCumSum_}; + int32_t cumSumFlag = 0; + gatherSumPattern.SetValue(0, 2); + SyncFunc(); + + + while (true) { + DataCopy(sumLocalTensor, selfRankWinInGMTensor_[(CUMSUM_CAL_OFFSET + newAivId * UB_ALIGN) / sizeof(float)], sumIntriParams); + SyncFunc(); + GatherMask(sumContinueTensor, sumLocalTensor, gatherSumPattern, true, mask, {1, static_cast(aivUsedCumSum_), 1, 0}, recvCnt); + PipeBarrier(); + Sum(recvCntSumOutTensor, sumContinueTensor, sumParams); + SyncFunc(); + cumSumFlag = static_cast(recvCntSumOutTensor.GetValue(0)); + if (cumSumFlag == aivUsedCumSum_) { + break; + } + } + + + if (newAivId == 0) { + outLocal.SetValue(0, 0); + } else { + mask = 1; + recvCnt = 0; + gatherSumPattern.SetValue(0, 1); + SyncFunc(); + GatherMask(sumContinueTensor, sumLocalTensor, gatherSumPattern, true, mask, {1, static_cast(newAivId), 1, 0}, recvCnt); + PipeBarrier(); + uint32_t innerCumSumParams = Ceil(newAivId * sizeof(float), UB_ALIGN) * UB_ALIGN / sizeof(float); + SumParams cumSumParams{1, innerCumSumParams, newAivId}; + Sum(recvCntSumOutTensor, sumContinueTensor, cumSumParams); + SyncFunc(); + outLocal.SetValue(0, recvCntSumOutTensor.ReinterpretCast().GetValue(0)); + } + + LocalTensor sumCoreFp32Tensor = sumLocalBuf_.Get(); + + uint8_t repeatTimes = Ceil(aivUsedCumSum_, 8); + + Duplicate(sumCoreFp32Tensor, static_cast(0), 64, repeatTimes, 1, 8); + DataCopyParams cleanParams{static_cast(aivUsedCumSum_), 1, 0, static_cast(aivUsedCumSum_ - 1)}; + SyncFunc(); + DataCopy(selfRankWinInGMTensor_[(CUMSUM_CAL_OFFSET + newAivId * UB_ALIGN) / sizeof(float)], sumCoreFp32Tensor, cleanParams); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::WaitDispatch() +{ + LocalTensor gatherMaskOutTensor = gatherMaskOutBuf_.Get(); + LocalTensor gatherTmpTensor = scalarBuf_.GetWithOffset(UB_ALIGN / sizeof(uint32_t), 0); + LocalTensor statusSumOutTensor = scalarBuf_.GetWithOffset(UB_ALIGN / sizeof(float), UB_ALIGN); + statusFp32Tensor_ = waitStatusBuf_.Get(); + uint32_t mask = UB_ALIGN_DATA_COUNT; + gatherTmpTensor.SetValue(0, STATUS_FLAG_PATTERN); + float compareTarget = static_cast(1.0) * recStatusNumPerCore_; + float sumOfFlag = static_cast(-1.0); + uint64_t gatheredFlagCount = 0; + uint32_t flagSumInner = Ceil(recStatusNumPerCore_ * sizeof(float), UB_ALIGN) * UB_ALIGN / sizeof(float); + SumParams flagSumParams{1, flagSumInner, recStatusNumPerCore_}; + + DataCopyParams intriParams{static_cast(recStatusNumPerCore_), 1, 0, 0}; + SyncFunc(); + while (sumOfFlag != compareTarget) { + DataCopy(statusFp32Tensor_, windowInstatusFp32Tensor_[startStatusIndex_ * STATE_OFFSET / sizeof(float)], intriParams); + SyncFunc(); + gatheredFlagCount = 0; + GatherMask(gatherMaskOutTensor, statusFp32Tensor_, gatherTmpTensor, true, mask, + {1, static_cast(recStatusNumPerCore_), 1, 0}, gatheredFlagCount); + PipeBarrier(); + Sum(statusSumOutTensor, gatherMaskOutTensor, flagSumParams); + SyncFunc(); + sumOfFlag = statusSumOutTensor.GetValue(0); + } + RunPosRecord(RUNPOS_CALCUMSUM); + + WaitDispatchClearStatus(); + GatherSumRecvCnt(gatherMaskOutTensor, gatherTmpTensor, statusSumOutTensor); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::CalRecvAndSetFlag() +{ + + LocalTensor outCountLocal; + uint32_t newAivId = aivId_ - aivUsedAllToAll_; + GetCumSum(outCountLocal, newAivId); + + uint32_t preSum = outCountLocal.GetValue(0); + uint32_t curCnt = preSum; + waitStatusTensor_ = waitStatusBuf_.Get(); + for (uint32_t index = startStatusIndex_; index < endStatusIndex_; index++) { + uint32_t i = index - startStatusIndex_; + uint32_t count = waitStatusTensor_.GetValue(i * UB_ALIGN_DATA_COUNT + STATUS_COUNT_INDEX); + curCnt += count; + outCountLocal.SetValue(i, curCnt); + } + SyncFunc(); + GM_ADDR wAddr = (__gm__ uint8_t*)(recvCntWorkspaceGM_); + GlobalTensor sendCountsGlobal, workspaceGlobal; + sendCountsGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(sendCountsOutGM_)); + workspaceGlobal.SetGlobalBuffer((__gm__ int32_t*)wAddr); + DataCopyExtParams dataCopyOutParams{1U, static_cast(recStatusNumPerCore_ * sizeof(int32_t)), 0U, 0U, 0U}; + DataCopyPad(sendCountsGlobal[startStatusIndex_], outCountLocal, dataCopyOutParams); + + for (uint32_t index = 0; index < aivNum_; index++) { + DataCopyPad(workspaceGlobal[index * rscvStatusNum_ + startStatusIndex_], outCountLocal, dataCopyOutParams); + } + uint8_t repeatTimes = Ceil(aivNum_, 8); + DataCopyParams sumIntriParams{static_cast(aivNum_), 1, 0, static_cast(aivUsedCumSum_ - 1)}; + LocalTensor syncOnCoreTensor = sumCoreBuf_.Get(); + LocalTensor syncOnCoreFP32Tensor = sumCoreBuf_.Get(); + + Duplicate(syncOnCoreTensor, static_cast(1), SIZE_ALIGN_256 / sizeof(int32_t), repeatTimes, 1, 8); + SyncFunc(); + PipeBarrier(); + + DataCopy(selfRankWinInGMTensor_[(CUMSUM_FLAG_OFFSET + newAivId * UB_ALIGN) / sizeof(float)], syncOnCoreFP32Tensor, sumIntriParams); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::CalCumSum() +{ + + expertIdsBufSize_ = Ceil(expertIdsCnt_ * sizeof(int32_t), SIZE_ALIGN_256) * SIZE_ALIGN_256; + tpipe_->InitBuffer(dstExpBuf_, maxSize_); + tpipe_->InitBuffer(subExpBuf_, maxSize_); + tpipe_->InitBuffer(gatherMaskTBuf_, expertIdsBufSize_); + tpipe_->InitBuffer(expertIdsBuf_, expertIdsBufSize_); + tpipe_->InitBuffer(statusBuf_, statusCntAlign_ * UB_ALIGN); + workLocalTensor_ = gatherMaskTBuf_.Get(); + statusTensor_ = statusBuf_.Get(); + ExpIdsCopyAndMaskCal(); + CalAndSendCntByRank(); + SplitToCore(rscvStatusNum_, aivUsedCumSum_, startStatusIndex_, endStatusIndex_, recStatusNumPerCore_, false); + BufferInit(); + WaitDispatch(); + CalRecvAndSetFlag(); + +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::WaitCumSumFlag() +{ + + int32_t cumSumFlag = 0; + int32_t targetFlag = aivUsedCumSum_ * UB_ALIGN_DATA_COUNT; + uint32_t cumSumFlagOffset = (CUMSUM_FLAG_OFFSET + aivId_ * aivUsedCumSum_ * UB_ALIGN) / sizeof(float); + uint32_t innerSumParams = aivUsedCumSum_ * UB_ALIGN / sizeof(float); + SumParams sumFlagParams{1, innerSumParams, aivUsedCumSum_ * UB_ALIGN_DATA_COUNT}; + LocalTensor statusSumOutTensor = scalarBuf_.Get(); + + while (true) { + DataCopy(statusFp32Tensor_, selfRankWinInGMTensor_[cumSumFlagOffset], aivUsedCumSum_ * UB_ALIGN_DATA_COUNT); + SyncFunc(); + Sum(statusSumOutTensor, statusFp32Tensor_, sumFlagParams); + SyncFunc(); + cumSumFlag = statusSumOutTensor.ReinterpretCast().GetValue(0); + if (cumSumFlag == targetFlag) { + break; + } + } + RunPosRecord(RUNPOS_CUMSUMFLAG); + + Duplicate(statusCleanFp32Tensor_, static_cast(0), aivUsedCumSum_ * UB_ALIGN_DATA_COUNT); + SyncFunc(); + + SyncFunc(); + DataCopy(selfRankWinInGMTensor_[cumSumFlagOffset], statusCleanFp32Tensor_, aivUsedCumSum_ * UB_ALIGN_DATA_COUNT); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::SetValidExpertInfo(uint32_t expInfoSize, uint32_t &validNum) +{ + + GlobalTensor workspaceGlobal; + workspaceGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(recvCntWorkspaceGM_)); + DataCopyExtParams scalesCopyInParams{1U, static_cast(rscvStatusNum_ * sizeof(int32_t)), 0U, 0U, 0U}; + DataCopyPadExtParams copyPadExtParams{false, 0U, 0U, 0U}; + DataCopyPad(sendCntTensor_, workspaceGlobal[aivId_ * rscvStatusNum_], scalesCopyInParams, copyPadExtParams); + PipeBarrier(); + + if (aivId_ == 0) { + uint32_t localExpertNum = isShareExpertRankFlag_ ? 1 : moeExpertNumPerRank_; + int64_t lastVal = 0; + uint32_t tokenNumBufSize = Ceil(moeExpertNumPerRank_ * sizeof(int64_t), UB_ALIGN) * UB_ALIGN; + tpipe_->InitBuffer(tokenNumBuf_, tokenNumBufSize); + LocalTensor expertTokenNumsLocalTensor = tokenNumBuf_.Get(); + for (uint32_t localExpertIdx = 0; localExpertIdx < localExpertNum; ++localExpertIdx) { + if (expertTokenNumsType_ == 0) { + expertTokenNumsLocalTensor(localExpertIdx) = int64_t(sendCntTensor_(localExpertIdx * epWorldSize_ + + epWorldSize_ - 1)); + } else { + expertTokenNumsLocalTensor(localExpertIdx) = int64_t(sendCntTensor_(localExpertIdx * epWorldSize_ + + epWorldSize_ - 1)) - lastVal; + lastVal = int64_t(sendCntTensor_(localExpertIdx * epWorldSize_ + epWorldSize_ -1)); + } + } + SyncFunc(); + DataCopyExtParams expertTokenNumsCopyParams{1U, static_cast(localExpertNum * sizeof(int64_t)), + 0U, 0U, 0U}; + DataCopyPad(expertTokenNumsOutGMTensor_, expertTokenNumsLocalTensor, expertTokenNumsCopyParams); + } + + Duplicate(expertFinishNumTensor_, 0, expInfoSize / sizeof(uint32_t)); + for (uint32_t index = startId_; index < endId_; index++) { + expertMapTensor_(validNum) = index; + if (index == 0) { + expertLeftNumTensor_(validNum) = sendCntTensor_(index); + } else { + expertLeftNumTensor_(validNum) = sendCntTensor_(index) - sendCntTensor_(index - 1); + } + if (expertLeftNumTensor_(validNum) != 0) { + validNum += 1; + } + } +} + +template +__aicore__ inline uint32_t MoeDistributeDispatchV2FullMesh::CheckDataArriveWithFlag(uint32_t srcExpDataIdx, + int32_t beginIdx, int32_t copyCnt) +{ + uint64_t rsvdCnt = 0; + uint32_t arriveFlagNum = 0; + uint32_t flagNum = blockCntPerToken_ * uint32_t(copyCnt); + uint32_t compareCount = Ceil(flagNum, COMPARE_COUNT_PER_BLOCK) * COMPARE_COUNT_PER_BLOCK; + uint32_t compResultU64Num = Ceil(flagNum, 64); + DataCopyExtParams expFlagCopyParams{static_cast(flagNum), static_cast(sizeof(float)), + static_cast(SPLIT_BLOCK_SIZE - sizeof(float)), 0, 0}; + DataCopyPadExtParams expFlagPadParams{false, 0U, 0U, 0U}; + GlobalTensor dataFlagGlobal; + GM_ADDR wAddr = (__gm__ uint8_t*)(windowGM_) + srcExpDataIdx * expertPerSizeOnWin_ + + beginIdx * hCommuSize_ + SPLIT_BLOCK_DATA_SIZE; + dataFlagGlobal.SetGlobalBuffer((__gm__ float *)(wAddr)); + DataCopyPad(flagRecvTensor_, dataFlagGlobal, expFlagCopyParams, expFlagPadParams); + SyncFunc(); + GatherMask(flagGatherOutTensor_, flagRecvTensor_, flagRecvGatherMask_, true, uint32_t(1), + + {1, (uint16_t)(flagNum), 1, 0}, rsvdCnt); + PipeBarrier(); + CompareScalar(flagCompResultU8_, flagGatherOutTensor_, float(1), AscendC::CMPMODE::EQ, compareCount); + SyncFunc(); + + for (uint32_t i = 0; i < compResultU64Num; i++) { + uint64_t flagCompMask = flagCompResultLtU64_(i); + int64_t firstValidIdx = ScalarGetSFFValue<0>(flagCompMask); + if (firstValidIdx == -1) { + arriveFlagNum += 64U; + } else { + arriveFlagNum += uint32_t(firstValidIdx); + break; + } + } + if (arriveFlagNum > flagNum) { + arriveFlagNum = flagNum; + } + return uint32_t(arriveFlagNum / blockCntPerToken_); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::CopyInAndOut( + LocalTensor xOutInt32Tensor, GM_ADDR wAddr, uint32_t index, uint32_t srcExpertId, + uint32_t dstPosition, uint32_t arriveCount) +{ + GlobalTensor dataFlagGlobal, expandXOutGlobal; + GlobalTensor assistGlobal; + dataFlagGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(wAddr)); + expandXOutGlobal.SetGlobalBuffer( + reinterpret_cast<__gm__ uint8_t *>(expandXOutGM_) + dstPosition * hOutSize_); + assistGlobal.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t *>(assistInfoForCombineOutGM_) + dstPosition * 4); + + DataCopyParams srcTokenCopyParams{static_cast(blockCntPerToken_ * arriveCount), + static_cast(SPLIT_BLOCK_DATA_SIZE), static_cast(UB_ALIGN), 0}; + DataCopyExtParams tokenCopyParams{static_cast(arriveCount), hOutSize_, + static_cast((blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE - hOutSize_) / UB_ALIGN), 0U, 0U}; + DataCopyExtParams scalesCopyParams{static_cast(arriveCount), scaleOutBytes_, + static_cast((blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE - scaleOutBytes_) / UB_ALIGN), 0U, 0U}; + DataCopyPadParams srcTokenPadParams{false, 0U, 0U, 0U}; + + DataCopyPad(xTmpTensor_, + dataFlagGlobal[expertFinishNumTensor_(index) * hCommuSize_], + srcTokenCopyParams, srcTokenPadParams); + SyncFunc(); + if (useMxfp8_) { + LocalTensor scalesLocal = xTmpTensor_[Ceil(axisH_, 256U) * 256U]; + DataCopyPad(dynamicScalesOutGMTensor_[dstPosition * scaleOutBytes_], scalesLocal, scalesCopyParams); + } + DataCopyPad(expandXOutGlobal, xTmpTensor_, tokenCopyParams); + SyncFunc(); + + uint32_t tokenStride = blockCntPerToken_ * SPLIT_BLOCK_DATA_SIZE / sizeof(int32_t); + int32_t moeExpertId = static_cast(sharedExpertNum_) + + (epRankId_ - static_cast(sharedExpertRankNum_)) * + static_cast(moeExpertNumPerRank_) + + static_cast(srcExpertId / epWorldSize_); + for (uint32_t token = 0; token < arriveCount; ++token) { + uint32_t srcOffset = token * tokenStride + tokenQuantAlign_; + uint32_t dstOffset = token * 4; + int32_t topKId = xOutInt32Tensor.GetValue(srcOffset + 2); + xOutInt32Tensor.SetValue(dstOffset, xOutInt32Tensor.GetValue(srcOffset)); + xOutInt32Tensor.SetValue(dstOffset + 1, xOutInt32Tensor.GetValue(srcOffset + 1)); + xOutInt32Tensor.SetValue(dstOffset + 2, topKId); + xOutInt32Tensor.SetValue(dstOffset + 3, + isShareExpertRankFlag_ ? topKId - static_cast(axisK_) : moeExpertId); + } + SyncFunc(); + DataCopyExtParams assistCopyParams{1U, + static_cast(arriveCount * 4 * sizeof(int32_t)), 0U, 0U, 0U}; + DataCopyPad(assistGlobal, xOutInt32Tensor, assistCopyParams); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::WaitAndFormatOutput(TBuf<> tBuf, uint32_t validNum) +{ + uint32_t index = 0; + uint32_t finishNum = 0; + uint32_t maxCopyTokenCnt = tBufRealSize_ / hCommuSize_; + uint32_t localExpertNum = isShareExpertRankFlag_ ? 1 : moeExpertNumPerRank_; + uint32_t srcExpertId, dstPosition, arriveCount, copyCnt, srcDataBlockIdx; + uint32_t flagMaxRecvNum = (blockCntPerToken_ * maxCopyTokenCnt * UB_ALIGN) / sizeof(uint32_t); + uint32_t gatherOutSize = Ceil(blockCntPerToken_ * maxCopyTokenCnt * sizeof(uint32_t), SIZE_ALIGN_256) * SIZE_ALIGN_256; + GlobalTensor cleanGlobal; + flagGatherOutTensor_ = tBuf.GetWithOffset(gatherOutSize / sizeof(float), 0); + flagRecvTensor_ = tBuf.GetWithOffset(flagMaxRecvNum, gatherOutSize); + LocalTensor xOutInt32Tensor = xTmpTensor_.template ReinterpretCast(); + DataCopyParams cleanUpParams = {uint16_t(blockCntPerToken_), 1U, 0U, SPLIT_BLOCK_DATA_SIZE / UB_ALIGN}; + while (true) { + if (expertLeftNumTensor_(index) == 0) { + index = (index + 1) % validNum; + continue; + } + srcExpertId = expertMapTensor_(index); + copyCnt = expertLeftNumTensor_(index) > maxCopyTokenCnt ? maxCopyTokenCnt : expertLeftNumTensor_(index); + srcDataBlockIdx = srcExpertId % epWorldSize_ * localExpertNum + srcExpertId / epWorldSize_; + if (!isShareExpertRankFlag_) { + srcDataBlockIdx = (srcExpertId + epRankId_) % epWorldSize_ * localExpertNum + srcExpertId / epWorldSize_; + } + arriveCount = CheckDataArriveWithFlag(srcDataBlockIdx, expertFinishNumTensor_(index), copyCnt); + if (arriveCount == copyCnt) { + dstPosition = srcExpertId != 0 ? sendCntTensor_(srcExpertId - 1) : 0; + dstPosition += expertFinishNumTensor_(index); + GM_ADDR wAddr = (__gm__ uint8_t*)(windowGM_) + srcDataBlockIdx * expertPerSizeOnWin_; + CopyInAndOut(xOutInt32Tensor, wAddr, index, srcExpertId, dstPosition, arriveCount); + + expertFinishNumTensor_(index) += arriveCount; + expertLeftNumTensor_(index) -= arriveCount; + PipeBarrier(); + if (expertLeftNumTensor_(index) == 0) { + cleanGlobal.SetGlobalBuffer((__gm__ float *)(wAddr)); + for (uint32_t i = 0; i < expertFinishNumTensor_(index); i++){ + uint32_t flagIndex = i * SPLIT_BLOCK_COUNT * blockCntPerToken_ + SPLIT_BLOCK_DATA_COUNT; + DataCopy(cleanGlobal[flagIndex], cleanUpTensor_, cleanUpParams); + } + finishNum++; + } + } else { + index = (index + 1) % validNum; + } + if (validNum == finishNum) { + break; + } + } +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::RunPosRecord(const uint32_t runPos) +{ + TBuf<> runPosBuf; + tpipe_->InitBuffer(runPosBuf, UB_ALIGN); + dataStateLocalTensor_ = runPosBuf.Get(); + dataStateLocalTensor_.SetValue(0, runPos); + SyncFunc(); + DataCopyPad(selfDataStatusGMTensor_[1], dataStateLocalTensor_, dataStateParams_); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::LocalWindowCopy() +{ + + tpipe_->Reset(); + TBuf<> cumSumBuf, statusWaitBuf, statusCleanBuf; + uint32_t rscvNumAlign = Ceil(rscvStatusNum_ * sizeof(int32_t), UB_ALIGN) * UB_ALIGN; + + tpipe_->InitBuffer(scalarBuf_, UB_ALIGN); + tpipe_->InitBuffer(statusWaitBuf, aivUsedCumSum_ * UB_ALIGN); + tpipe_->InitBuffer(cumSumBuf, rscvNumAlign); + tpipe_->InitBuffer(statusCleanBuf, aivUsedCumSum_ * UB_ALIGN); + statusFp32Tensor_ = statusWaitBuf.Get(); + statusCleanFp32Tensor_ = statusCleanBuf.Get(); + sendCntTensor_ = cumSumBuf.Get(); + SplitToCore(rscvStatusNum_, aivNum_, startId_, endId_, sendNum_, true); + + WaitCumSumFlag(); + if (sendNum_ == 0) { + + return; + } + + TBuf<> expertMapBuf, expertFinishBuf, expertLeftBuf, flagMaskBuf, cleanUpBuf, tBuf; + uint32_t validNum = 0; + uint32_t expInfoSize = Ceil(sendNum_ * sizeof(uint32_t), UB_ALIGN) * UB_ALIGN; + tpipe_->InitBuffer(expertMapBuf, expInfoSize); + tpipe_->InitBuffer(expertFinishBuf, expInfoSize); + tpipe_->InitBuffer(expertLeftBuf, expInfoSize); + tpipe_->InitBuffer(flagMaskBuf, BUFFER_NUM * UB_ALIGN); + tpipe_->InitBuffer(cleanUpBuf, blockCntPerToken_ * UB_ALIGN); + tBufRealSize_ = FULL_MESH_MAX_UB_SIZE - (UB_ALIGN + rscvNumAlign + 2 * aivUsedCumSum_ * UB_ALIGN) - + (expInfoSize * 3) - BUFFER_NUM * UB_ALIGN - blockCntPerToken_ * UB_ALIGN; + tpipe_->InitBuffer(tBuf, tBufRealSize_); + expertMapTensor_ = expertMapBuf.Get(); + expertFinishNumTensor_ = expertFinishBuf.Get(); + expertLeftNumTensor_ = expertLeftBuf.Get(); + SetValidExpertInfo(expInfoSize, validNum); + if (validNum == 0) { + return; + } + flagCompResultU8_ = flagMaskBuf.Get(); + flagCompResultLtU64_ = flagMaskBuf.Get(); + flagRecvGatherMask_ = statusCleanBuf.GetWithOffset(UB_ALIGN / sizeof(uint32_t), 0); + cleanUpTensor_ = cleanUpBuf.Get(); + xTmpTensor_ = tBuf.Get(); + LocalTensor flagCompResultLtU32 = flagMaskBuf.Get(); + Duplicate(flagCompResultLtU32, 0, BUFFER_NUM * UB_ALIGN / sizeof(uint32_t)); + Duplicate(flagRecvGatherMask_, 0, UB_ALIGN / sizeof(uint32_t)); + Duplicate(cleanUpTensor_, float(0), blockCntPerToken_ * UB_ALIGN_DATA_COUNT); + SyncFunc(); + flagRecvGatherMask_.SetValue(0, 1); + SyncFunc(); + WaitAndFormatOutput(tBuf, validNum); + RunPosRecord(RUNPOS_ARRIVECNT); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::TokenActiveMaskCal() +{ + + LocalTensor maskTmpTensor; + LocalTensor sumOutTensor; + LocalTensor maskInputTensor; + uint32_t axisBsAlignSize = Ceil(axisBS_ * sizeof(bool), UB_ALIGN) * UB_ALIGN; + maskInputTensor = dstExpBuf_.Get(); + maskTmpTensor = subExpBuf_.Get(); + sumOutTensor = gatherMaskTBuf_.Get(); + DataCopyExtParams maskParams = {1U, static_cast(axisBS_ * sizeof(bool)), 0U, 0U, 0U}; + DataCopyPadExtParams maskCopyPadParams{false, 0U, 0U, 0U}; + DataCopyPad(maskInputTensor, xActiveMaskGMTensor_, maskParams, maskCopyPadParams); + SyncFunc(); + LocalTensor maskInputInt8Tensor = maskInputTensor.ReinterpretCast(); + Cast(maskTmpTensor, maskInputInt8Tensor, RoundMode::CAST_NONE, axisBS_); + PipeBarrier(); + SumParams params{1, axisBsAlignSize, axisBS_}; + Sum(sumOutTensor, maskTmpTensor, params); + SyncFunc(); + activeMaskBsCnt_ = static_cast(sumOutTensor.GetValue(0)); + sendToMoeExpTokenCnt_ = activeMaskBsCnt_ * axisK_; +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::CalValidBSCnt(LocalTensor maskStrideTensor) +{ + uint64_t rsvdCnt = 0; + uint32_t mask = axisBS_; + uint32_t activeMaskAlignSize = axisBS_ * (Ceil(axisK_ * sizeof(bool), UB_ALIGN) * UB_ALIGN); + uint32_t calCnt = Ceil(axisBS_ * sizeof(half), SIZE_ALIGN_256) * SIZE_ALIGN_256 / sizeof(half); + uint32_t innerAlign = Ceil(axisK_ * sizeof(half), UB_ALIGN) * UB_ALIGN / sizeof(half) * BUFFER_NUM; + LocalTensor tempTensor = validExpertIndexBuf_.Get(); + LocalTensor maskTempTensor = expertIdsBuf_.Get(); + LocalTensor tokenTargetTensor = validBsIndexTBuf_.Get(); + LocalTensor maskTensor = gatherMaskTBuf_.Get(); + LocalTensor bsIndexTensor = subExpBuf_.Get(); + LocalTensor maskTensorInt32 = gatherMaskTBuf_.Get(); + SumParams axisKSumParams{axisBS_, innerAlign, axisK_}; + SumParams axisBsSumParams{1, static_cast(Ceil(axisBS_ * sizeof(half), UB_ALIGN) * UB_ALIGN / sizeof(half)), axisBS_}; + + Duplicate(maskTempTensor, (half)0, calCnt); + SyncFunc(); + LocalTensor maskStrideInt8Tensor = maskStrideTensor.ReinterpretCast(); + Cast(tempTensor, maskStrideInt8Tensor, RoundMode::CAST_NONE, activeMaskAlignSize); + PipeBarrier(); + Sum(tokenTargetTensor, tempTensor, axisKSumParams); + PipeBarrier(); + Mins(maskTempTensor, tokenTargetTensor, static_cast(1), axisBS_); + PipeBarrier(); + CompareScalar(maskTensor, maskTempTensor, static_cast(1), AscendC::CMPMODE::EQ, calCnt); + CreateVecIndex(bsIndexTensor, 0, axisBS_); + PipeBarrier(); + GatherMask(validBsIndexTensor_, bsIndexTensor, maskTensorInt32, true, mask, {1, 1, 0, 0}, activeMaskBsCnt_); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::CalValidExpIdx(LocalTensor maskInputTensor) +{ + uint32_t mask = expertIdsCnt_; + uint32_t curMaskCnt = axisBS_ * axisK_; + uint32_t calCnt = Ceil(curMaskCnt * sizeof(half), SIZE_ALIGN_256) * SIZE_ALIGN_256 / sizeof(half); + + LocalTensor validExpertIndexTensor = validExpertIndexBuf_.Get(); + LocalTensor tempTensor = subExpBuf_.Get(); + LocalTensor gatherMaskTensorInt8 = gatherMaskTBuf_.Get(); + LocalTensor expertsIndexTensor = expertIdsBuf_.Get(); + + Duplicate(tempTensor, (half)0, calCnt); + PipeBarrier(); + + SyncFunc(); + LocalTensor maskInputInt8Tensor = maskInputTensor.ReinterpretCast(); + Cast(tempTensor, maskInputInt8Tensor, RoundMode::CAST_NONE, curMaskCnt); + PipeBarrier(); + Duplicate(gatherMaskTensor_, 0, Ceil(expertIdsCnt_, SIZE_ALIGN_256) * SIZE_ALIGN_256 / BITS_PER_BYTE / sizeof(uint32_t)); + PipeBarrier(); + CompareScalar(gatherMaskTensorInt8, tempTensor, static_cast(1), AscendC::CMPMODE::EQ, calCnt); + CreateVecIndex(expertsIndexTensor, 0, curMaskCnt); + PipeBarrier(); + GatherMask(validExpertIndexTensor, expertsIndexTensor, gatherMaskTensor_, true, mask, {1, 1, 0, 0}, sendToMoeExpTokenCnt_); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::ExpertActiveMaskInit() +{ + uint32_t axisBSAlign = Ceil(axisBS_ * sizeof(int32_t), UB_ALIGN) * UB_ALIGN; + uint32_t xActivateMaskSize = axisBS_ * (Ceil(axisK_ * sizeof(bool), UB_ALIGN) * UB_ALIGN) * sizeof(half); + tpipe_->InitBuffer(validBsIndexTBuf_, axisBSAlign); + uint32_t validBufferSize = expertIdsSize_ > xActivateMaskSize ? expertIdsSize_ : xActivateMaskSize; + tpipe_->InitBuffer(validExpertIndexBuf_, validBufferSize); + validBsIndexTensor_ = validBsIndexTBuf_.Get(); + gatherMaskTensor_ = gatherMaskTBuf_.Get(); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::ExpertActiveMaskCal() +{ + + LocalTensor maskStrideTensor = dstExpBuf_.Get(); + DataCopyPadExtParams maskStrideCopyPadParams{false, 0U, 0U, 0U}; + DataCopyExtParams maskStrideParams{ + static_cast(axisBS_), static_cast(axisK_ * sizeof(bool)), 0U, 0U, 0U}; + DataCopyPad(maskStrideTensor, xActiveMaskGMTensor_, maskStrideParams, maskStrideCopyPadParams); + CalValidBSCnt(maskStrideTensor); + + LocalTensor maskInputTensor = dstExpBuf_.Get(); + DataCopyPadExtParams maskCopyPadParams{false, 0U, 0U, 0U}; + DataCopyExtParams maskParams{1U, static_cast(expertIdsCnt_ * sizeof(bool)), 0U, 0U, 0U}; + DataCopyPad(maskInputTensor, xActiveMaskGMTensor_, maskParams, maskCopyPadParams); + CalValidExpIdx(maskInputTensor); + SyncFunc(); +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::ExpIdsCopyAndMaskCal() +{ + activeMaskBsCnt_ = axisBS_; + sendToMoeExpTokenCnt_ = axisBS_ * axisK_; + validExpertIdsTensor_ = expertIdsBuf_.Get(); + + if (isExpertMaskFlag_) { + ExpertActiveMaskInit(); + } + + if (isTokenMaskFlag_) { + TokenActiveMaskCal(); + } + + if (isExpertMaskFlag_) { + ExpertActiveMaskCal(); + } + if (activeMaskBsCnt_ == 0) { + return; + } + Duplicate(validExpertIdsTensor_, -1, int32_t(expertIdsBufSize_ / sizeof(int32_t))); + + if (isExpertMaskFlag_) { + LocalTensor tmpExpertIdsTensor = subExpBuf_.Get(); + LocalTensor tmpExpertIdsTensorFloat = subExpBuf_.Get(); + LocalTensor gatherMaskTensorInt8 = gatherMaskTensor_.ReinterpretCast(); + DataCopyExtParams expertIdsMaskParams{1U, static_cast(expertIdsCnt_ * sizeof(uint32_t)), 0U, 0U, 0U}; + DataCopyPadExtParams expertIdsMaskCopyPadParams{false, 0U, 0U, 0U}; + DataCopyPad(tmpExpertIdsTensor, expertIdsGMTensor_, expertIdsMaskParams, expertIdsMaskCopyPadParams); + SyncFunc(); + PipeBarrier(); + LocalTensor validExpertIdsFloat = validExpertIdsTensor_.ReinterpretCast(); + Select(validExpertIdsFloat, gatherMaskTensorInt8, tmpExpertIdsTensorFloat, static_cast(-1), SELMODE::VSEL_TENSOR_SCALAR_MODE, expertIdsCnt_); + SyncFunc(); + } else { + uint32_t expertIdsMask = activeMaskBsCnt_ * axisK_; + uint32_t expertIdsAlignCnt = Ceil(expertIdsMask, BITS_PER_BYTE) * BITS_PER_BYTE; + uint32_t rightPadding = expertIdsAlignCnt - expertIdsMask; + DataCopyPadExtParams expertIdsCntCopyPadParams{true, 0U, uint8_t(rightPadding), -1}; + DataCopyExtParams expertIdsCntParams{1U, static_cast(expertIdsMask * sizeof(uint32_t)), 0U, 0U, 0U}; + SyncFunc(); + DataCopyPad(validExpertIdsTensor_, expertIdsGMTensor_, expertIdsCntParams, expertIdsCntCopyPadParams); + SyncFunc(); + } +} + +template +__aicore__ inline void MoeDistributeDispatchV2FullMesh::Run() +{ + if ASCEND_IS_AIV { + // printf("aivId_ %d, time %d, %d, %d\n", aivId_, + // (usedTime_[1] - usedTime_[0])/ 1000 , (usedTime_[2] - usedTime_[1])/ 1000, (usedTime_[3] - usedTime_[2])/ 1000); + // printf("aivId_ %d, time %d\n", aivId_, (usedTime_[3] - usedTime_[0])/ 1000); + return; + if (aivId_ < aivUsedAllToAll_) { + AllToAllDispatch(); + } else { + CalCumSum(); + } + + PipeBarrier(); + LocalWindowCopy(); + } +} + +} + +extern "C" __global__ __aicore__ void tilexr_ep_dispatch_memory_kernel(GM_ADDR commArgsGM, GM_ADDR xGM, + GM_ADDR expertIdsGM, GM_ADDR xActiveMaskGM, GM_ADDR expandXOutGM, GM_ADDR dynamicScalesOutGM, + GM_ADDR expertTokenNumsOutGM, + GM_ADDR sendCountsOutGM, GM_ADDR assistInfoForCombineOutGM, int64_t bs, int64_t h, int64_t topK, + int64_t moeExpertNum, int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t globalBs, + int64_t expertTokenNumsType, int64_t activeMaskType, int64_t quantMode, + int64_t dtype, int64_t expandXOutDtype, int64_t magic) +{ + if (commArgsGM == nullptr || xGM == nullptr || expertIdsGM == nullptr || expandXOutGM == nullptr || + expertTokenNumsOutGM == nullptr || sendCountsOutGM == nullptr || assistInfoForCombineOutGM == nullptr || + bs <= 0 || h <= 0 || topK <= 0 || moeExpertNum <= 0 || sharedExpertNum < 0 || + sharedExpertRankNum < 0 || globalBs <= 0 || expertTokenNumsType < 0 || expertTokenNumsType > 1 || magic <= 0 || + activeMaskType < Mc2Kernel::ACTIVE_MASK_NONE || activeMaskType > Mc2Kernel::ACTIVE_MASK_EXPERT || + (activeMaskType == Mc2Kernel::ACTIVE_MASK_NONE && xActiveMaskGM != nullptr) || + (activeMaskType != Mc2Kernel::ACTIVE_MASK_NONE && xActiveMaskGM == nullptr) || + (quantMode != 0 && quantMode != Mc2Kernel::MX_QUANT) || + (quantMode == Mc2Kernel::MX_QUANT && (dynamicScalesOutGM == nullptr || + (expandXOutDtype != TileXR::TILEXR_DATA_TYPE_FP8E4M3 && + expandXOutDtype != TileXR::TILEXR_DATA_TYPE_FP8E5M2)))) { + return; + } + + AscendC::TPipe pipe; + if (dtype == TileXR::TILEXR_DATA_TYPE_FP16) { + + Mc2Kernel::MoeDistributeDispatchV2FullMesh op; + op.Init(commArgsGM, xGM, expertIdsGM, xActiveMaskGM, expandXOutGM, dynamicScalesOutGM, expertTokenNumsOutGM, + sendCountsOutGM, assistInfoForCombineOutGM, bs, h, topK, moeExpertNum, sharedExpertNum, + sharedExpertRankNum, globalBs, expertTokenNumsType, activeMaskType, quantMode, + expandXOutDtype, magic, &pipe); + op.Run(); + } else if (dtype == TileXR::TILEXR_DATA_TYPE_BFP16) { + Mc2Kernel::MoeDistributeDispatchV2FullMesh op; + op.Init(commArgsGM, xGM, expertIdsGM, xActiveMaskGM, expandXOutGM, dynamicScalesOutGM, expertTokenNumsOutGM, + sendCountsOutGM, assistInfoForCombineOutGM, bs, h, topK, moeExpertNum, sharedExpertNum, + sharedExpertRankNum, globalBs, expertTokenNumsType, activeMaskType, quantMode, + expandXOutDtype, magic, &pipe); + op.Run(); + } +} diff --git a/src/ep/kernels/tilexr_ep_mxfp8_quant.h b/src/ep/kernels/tilexr_ep_mxfp8_quant.h new file mode 100644 index 00000000..4d3fbb17 --- /dev/null +++ b/src/ep/kernels/tilexr_ep_mxfp8_quant.h @@ -0,0 +1,356 @@ +#ifndef TILEXR_EP_KERNELS_TILEXR_EP_MXFP8_QUANT_H +#define TILEXR_EP_KERNELS_TILEXR_EP_MXFP8_QUANT_H + +#include "kernel_operator.h" + +namespace TileXRMxfp8Quant { + +using namespace AscendC; + +constexpr int kPairCount = 2; +constexpr uint16_t kBf16ExponentMask = 0x7f80; +constexpr uint16_t kBf16ExponentBias = 0x7f00; +constexpr uint16_t kFp8ExponentMask = 0x00ff; +constexpr uint16_t kCustomizedNan = 0x7f81; +constexpr uint16_t kSpecialExponentThreshold = 0x0040; +constexpr int16_t kBf16ExponentShift = 7; +constexpr uint16_t kFp8E4M3MaxExponent = 0x0400; +constexpr uint16_t kFp8E5M2MaxExponent = 0x0780; +constexpr uint16_t kInvalidFp16 = 0x7c00; +constexpr int64_t kOutputElementsPerBlock = 64; + +__aicore__ inline constexpr uint32_t GetUbBlockSize() +{ + return 32U; +} + +__aicore__ inline constexpr uint32_t GetVectorRegisterSize() +{ +#if __CCE_AICORE__ == 310 + return AscendC::VECTOR_REG_WIDTH; +#else + return 256U; +#endif +} + +template +__aicore__ inline void ComputeMaxExp( + __ubuf__ T *srcAddr, __ubuf__ uint16_t *maxExpAddr, uint32_t totalCountInUb) +{ + const uint32_t elementsPerRegister = GetVectorRegisterSize() / sizeof(T); + const uint16_t reducedElements = GetVectorRegisterSize() / GetUbBlockSize(); + const uint16_t loopCount = static_cast( + (totalCountInUb + 2 * elementsPerRegister - 1) / (2 * elementsPerRegister)); + + __VEC_SCOPE__ + { + MicroAPI::RegTensor input0; + MicroAPI::RegTensor input1; + MicroAPI::RegTensor inputBf160; + MicroAPI::RegTensor inputBf161; + MicroAPI::RegTensor selected0; + MicroAPI::RegTensor selected1; + MicroAPI::RegTensor exponent0; + MicroAPI::RegTensor exponent1; + MicroAPI::RegTensor exponentMask; + MicroAPI::Duplicate(exponentMask, kBf16ExponentMask); + MicroAPI::RegTensor invalidFp16; + MicroAPI::Duplicate(invalidFp16, kInvalidFp16); + MicroAPI::RegTensor maxExponent; + MicroAPI::MaskReg mask0; + MicroAPI::MaskReg mask1; + MicroAPI::MaskReg valid0; + MicroAPI::MaskReg valid1; + MicroAPI::UnalignReg unalign; + static constexpr MicroAPI::CastTrait kHalfToBf16 = { + MicroAPI::RegLayout::UNKNOWN, + MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::CAST_TRUNC}; + + for (uint16_t loop = 0; loop < loopCount; ++loop) { + mask0 = MicroAPI::UpdateMask(totalCountInUb); + mask1 = MicroAPI::UpdateMask(totalCountInUb); + MicroAPI::DataCopy( + input0, input1, srcAddr, elementsPerRegister * kPairCount); + if constexpr (Std::IsSame::value) { + MicroAPI::And(selected0, reinterpret_cast &>(input0), + invalidFp16, mask0); + MicroAPI::And(selected1, reinterpret_cast &>(input1), + invalidFp16, mask0); + MicroAPI::Compare(valid0, selected0, invalidFp16, mask0); + MicroAPI::Compare(valid1, selected1, invalidFp16, mask0); + MicroAPI::Cast(inputBf160, input0, mask0); + MicroAPI::Cast(inputBf161, input1, mask0); + MicroAPI::And(exponent0, reinterpret_cast &>(inputBf160), + exponentMask, mask0); + MicroAPI::And(exponent1, reinterpret_cast &>(inputBf161), + exponentMask, mask0); + MicroAPI::Select(exponent0, exponent0, exponentMask, valid0); + MicroAPI::Select(exponent1, exponent1, exponentMask, valid1); + } else { + MicroAPI::And(exponent0, reinterpret_cast &>(input0), + exponentMask, mask0); + MicroAPI::And(exponent1, reinterpret_cast &>(input1), + exponentMask, mask0); + } + MicroAPI::Max(maxExponent, exponent0, exponent1, mask0); + MicroAPI::ReduceMaxWithDataBlock(maxExponent, maxExponent, mask0); + MicroAPI::DataCopyUnAlign( + maxExpAddr, maxExponent, unalign, reducedElements); + } + MicroAPI::DataCopyUnAlignPost(maxExpAddr, unalign, 0); + } +} + +template +__aicore__ inline void ComputeScale(__ubuf__ uint16_t *maxExpAddr, + __ubuf__ uint16_t *mxScaleLocalAddr, __ubuf__ uint16_t *halfScaleLocalAddr, + uint32_t totalScaleInUb) +{ + const uint32_t elementsPerRegister = GetVectorRegisterSize() / sizeof(uint16_t); + const uint16_t loopCount = static_cast( + (totalScaleInUb + elementsPerRegister - 1) / elementsPerRegister); + const uint16_t maxFp8Exponent = Std::IsSame::value ? + kFp8E4M3MaxExponent : kFp8E5M2MaxExponent; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor exponentMask; + MicroAPI::RegTensor maxExponent; + MicroAPI::Duplicate(exponentMask, kBf16ExponentMask); + MicroAPI::MaskReg validMask; + MicroAPI::MaskReg nonzeroMask; + MicroAPI::MaskReg boundedMask; + MicroAPI::MaskReg specialMask; + MicroAPI::RegTensor maxFp8ExponentTensor; + MicroAPI::Duplicate(maxFp8ExponentTensor, maxFp8Exponent); + MicroAPI::RegTensor sharedExponent; + MicroAPI::RegTensor scaleValue; + MicroAPI::RegTensor exponentBias; + MicroAPI::Duplicate(exponentBias, kBf16ExponentBias); + MicroAPI::RegTensor reciprocalScale; + MicroAPI::RegTensor fp8Nan; + MicroAPI::Duplicate(fp8Nan, kFp8ExponentMask); + MicroAPI::RegTensor zero; + MicroAPI::Duplicate(zero, 0); + MicroAPI::RegTensor nan; + MicroAPI::Duplicate(nan, kCustomizedNan); + MicroAPI::RegTensor specialExponent; + MicroAPI::Duplicate(specialExponent, kSpecialExponentThreshold); + + for (uint16_t loop = 0; loop < loopCount; ++loop) { + MicroAPI::MaskReg mask = MicroAPI::UpdateMask(totalScaleInUb); + MicroAPI::DataCopy( + maxExponent, maxExpAddr, elementsPerRegister); + MicroAPI::Compare(validMask, maxExponent, exponentMask, mask); + MicroAPI::Compare(nonzeroMask, maxExponent, zero, mask); + MicroAPI::Compare(boundedMask, maxExponent, maxFp8ExponentTensor, mask); + MicroAPI::Select(maxExponent, maxFp8ExponentTensor, maxExponent, boundedMask); + MicroAPI::Sub(sharedExponent, maxExponent, maxFp8ExponentTensor, mask); + MicroAPI::ShiftRights(scaleValue, sharedExponent, kBf16ExponentShift, mask); + MicroAPI::Select(scaleValue, scaleValue, fp8Nan, validMask); + MicroAPI::Select(scaleValue, scaleValue, zero, nonzeroMask); + MicroAPI::DataCopy( + mxScaleLocalAddr, scaleValue, elementsPerRegister / kPairCount, mask); + + MicroAPI::Compare(specialMask, sharedExponent, exponentBias, mask); + MicroAPI::Sub(reciprocalScale, exponentBias, sharedExponent, mask); + MicroAPI::Select(reciprocalScale, reciprocalScale, nan, validMask); + MicroAPI::Select(reciprocalScale, reciprocalScale, zero, nonzeroMask); + MicroAPI::Select(reciprocalScale, specialExponent, reciprocalScale, specialMask); + MicroAPI::DataCopy( + halfScaleLocalAddr, reciprocalScale, elementsPerRegister, mask); + } + } +} + +template +__aicore__ inline void ComputeFp8Data(__ubuf__ InputType *srcAddr, + __ubuf__ uint16_t *halfScaleLocalAddr, __ubuf__ int8_t *outLocalAddr, + uint32_t totalCountInUb) +{ + const uint32_t elementsPerRegister = GetVectorRegisterSize() / sizeof(InputType); + const uint16_t scaleElementsPerBlock = GetVectorRegisterSize() / GetUbBlockSize(); + uint32_t doubledCount = totalCountInUb * kPairCount; + const uint16_t loopCount = static_cast( + (totalCountInUb + 2 * elementsPerRegister - 1) / (2 * elementsPerRegister)); + + __VEC_SCOPE__ + { + MicroAPI::MaskReg inputMask0; + MicroAPI::MaskReg inputMask1; + MicroAPI::MaskReg fp32Mask0; + MicroAPI::MaskReg fp32Mask1; + MicroAPI::MaskReg allMask = + MicroAPI::CreateMask(); + MicroAPI::RegTensor packedScale; + MicroAPI::RegTensor fp32Scale; + MicroAPI::RegTensor input0; + MicroAPI::RegTensor input1; + MicroAPI::RegTensor fp32Input00; + MicroAPI::RegTensor fp32Input01; + MicroAPI::RegTensor fp32Input10; + MicroAPI::RegTensor fp32Input11; + MicroAPI::RegTensor fp8Input00; + MicroAPI::RegTensor fp8Input01; + MicroAPI::RegTensor fp8Input10; + MicroAPI::RegTensor fp8Input11; + static constexpr MicroAPI::CastTrait kEvenElements = { + MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + static constexpr MicroAPI::CastTrait kOddElements = { + MicroAPI::RegLayout::ONE, + MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, + RoundMode::UNKNOWN}; + static constexpr MicroAPI::CastTrait kFp32ToFp8 = { + MicroAPI::RegLayout::ZERO, + MicroAPI::SatMode::SAT, + MicroAPI::MaskMergeMode::ZEROING, + Fp8RoundMode}; + + for (uint16_t loop = 0; loop < loopCount; ++loop) { + inputMask0 = MicroAPI::UpdateMask(totalCountInUb); + inputMask1 = MicroAPI::UpdateMask(totalCountInUb); + fp32Mask0 = MicroAPI::UpdateMask(doubledCount); + fp32Mask1 = MicroAPI::UpdateMask(doubledCount); + MicroAPI::DataCopy( + input0, input1, srcAddr, elementsPerRegister * kPairCount); + MicroAPI::DataCopy( + packedScale, halfScaleLocalAddr, scaleElementsPerBlock); + + if constexpr (Std::IsSame::value) { + MicroAPI::Cast(fp32Input00, input0, inputMask0); + MicroAPI::Cast(fp32Input01, input0, inputMask0); + MicroAPI::Cast( + fp32Scale, reinterpret_cast &>(packedScale), allMask); + MicroAPI::Mul(fp32Input00, fp32Input00, fp32Scale, fp32Mask0); + MicroAPI::Mul(fp32Input01, fp32Input01, fp32Scale, fp32Mask1); + MicroAPI::Interleave(fp32Input00, fp32Input01, fp32Input00, fp32Input01); + MicroAPI::Cast(fp32Input10, input1, inputMask0); + MicroAPI::Cast(fp32Input11, input1, inputMask0); + MicroAPI::Mul(fp32Input10, fp32Input10, fp32Scale, fp32Mask0); + MicroAPI::Mul(fp32Input11, fp32Input11, fp32Scale, fp32Mask1); + MicroAPI::Interleave(fp32Input10, fp32Input11, fp32Input10, fp32Input11); + MicroAPI::Interleave(fp32Input00, fp32Input10, fp32Input00, fp32Input10); + MicroAPI::Interleave(fp32Input01, fp32Input11, fp32Input01, fp32Input11); + MicroAPI::Cast(fp8Input00, fp32Input00, fp32Mask0); + MicroAPI::Cast(fp8Input01, fp32Input10, fp32Mask0); + MicroAPI::Cast(fp8Input10, fp32Input01, fp32Mask1); + MicroAPI::Cast(fp8Input11, fp32Input11, fp32Mask1); + } else { + MicroAPI::Mul(input0, input0, reinterpret_cast &>(packedScale), + inputMask0); + MicroAPI::Mul(input1, input1, reinterpret_cast &>(packedScale), + inputMask0); + MicroAPI::Interleave(input0, input1, input0, input1); + MicroAPI::Cast(fp32Input00, input0, inputMask0); + MicroAPI::Cast(fp32Input01, input0, inputMask0); + MicroAPI::Interleave(fp32Input00, fp32Input01, fp32Input00, fp32Input01); + MicroAPI::Cast(fp8Input00, fp32Input00, fp32Mask0); + MicroAPI::Cast(fp8Input01, fp32Input01, fp32Mask0); + MicroAPI::Cast(fp32Input10, input1, inputMask1); + MicroAPI::Cast(fp32Input11, input1, inputMask1); + MicroAPI::Interleave(fp32Input10, fp32Input11, fp32Input10, fp32Input11); + MicroAPI::Cast(fp8Input10, fp32Input10, fp32Mask1); + MicroAPI::Cast(fp8Input11, fp32Input11, fp32Mask1); + } + + MicroAPI::DataCopy(outLocalAddr, + reinterpret_cast &>(fp8Input00), + kOutputElementsPerBlock, fp32Mask0); + MicroAPI::DataCopy(outLocalAddr, + reinterpret_cast &>(fp8Input01), + kOutputElementsPerBlock, fp32Mask0); + MicroAPI::DataCopy(outLocalAddr, + reinterpret_cast &>(fp8Input10), + kOutputElementsPerBlock, fp32Mask1); + MicroAPI::DataCopy(outLocalAddr, + reinterpret_cast &>(fp8Input11), + kOutputElementsPerBlock, fp32Mask1); + } + } +} + +template +__aicore__ inline void DequantizeAndAccumulate(__ubuf__ uint8_t *tokenAddr, + __ubuf__ fp8_e8m0_t *scaleAddr, __ubuf__ float *scaleWorkAddr, + __ubuf__ float *sumAddr, uint32_t axisH, uint32_t scaleCount, float expertScale) +{ + const uint32_t fp32RepeatSize = GetVectorRegisterSize() / sizeof(float); + const uint16_t scaleRepeatTimes = static_cast( + (scaleCount + fp32RepeatSize - 1U) / fp32RepeatSize); + const uint16_t dataRepeatTimes = static_cast( + (axisH + fp32RepeatSize * 2U - 1U) / (fp32RepeatSize * 2U)); + uint32_t scaleMaskCount = scaleCount; + uint32_t dataMaskCount = axisH; + uint32_t fp8MaskCount = axisH * 4U; + constexpr int16_t kFp32ExponentShift = 23; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor scaleReg; + MicroAPI::RegTensor tokenReg; + MicroAPI::RegTensor tokenFp320; + MicroAPI::RegTensor tokenFp321; + MicroAPI::RegTensor scaleFp32; + MicroAPI::RegTensor weighted0; + MicroAPI::RegTensor weighted1; + MicroAPI::RegTensor sum0; + MicroAPI::RegTensor sum1; + static constexpr MicroAPI::CastTrait kCastLane0 = { + MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + static constexpr MicroAPI::CastTrait kCastLane2 = { + MicroAPI::RegLayout::TWO, MicroAPI::SatMode::UNKNOWN, + MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; + + for (uint16_t loop = 0; loop < scaleRepeatTimes; ++loop) { + MicroAPI::MaskReg mask = MicroAPI::UpdateMask(scaleMaskCount); + MicroAPI::DataCopy( + scaleReg, scaleAddr + loop * fp32RepeatSize); + MicroAPI::ShiftLefts(reinterpret_cast &>(scaleFp32), + reinterpret_cast &>(scaleReg), kFp32ExponentShift, mask); + MicroAPI::DataCopy( + scaleWorkAddr + loop * fp32RepeatSize * 2U, scaleFp32, scaleFp32, mask); + } + + MicroAPI::LocalMemBar(); + for (uint16_t loop = 0; loop < dataRepeatTimes; ++loop) { + MicroAPI::MaskReg dataMask = MicroAPI::UpdateMask(dataMaskCount); + MicroAPI::MaskReg fp8Mask = MicroAPI::UpdateMask(fp8MaskCount); + MicroAPI::DataCopy( + scaleFp32, scaleWorkAddr + loop * 8U); + MicroAPI::DataCopy( + tokenReg, reinterpret_cast<__ubuf__ Fp8Type *>(tokenAddr) + + 2U * loop * fp32RepeatSize); + MicroAPI::DataCopy( + sum0, sum1, sumAddr + 2U * loop * fp32RepeatSize); + MicroAPI::Cast(tokenFp320, tokenReg, fp8Mask); + MicroAPI::Cast(tokenFp321, tokenReg, fp8Mask); + MicroAPI::Mul(weighted0, scaleFp32, tokenFp320, dataMask); + MicroAPI::Mul(weighted1, scaleFp32, tokenFp321, dataMask); + MicroAPI::Muls(weighted0, weighted0, expertScale, dataMask); + MicroAPI::Muls(weighted1, weighted1, expertScale, dataMask); + MicroAPI::Add(sum0, sum0, weighted0, dataMask); + MicroAPI::Add(sum1, sum1, weighted1, dataMask); + MicroAPI::DataCopy( + sumAddr + loop * fp32RepeatSize * 2U, sum0, sum1, dataMask); + } + } +} + +} // namespace TileXRMxfp8Quant + +#endif // TILEXR_EP_KERNELS_TILEXR_EP_MXFP8_QUANT_H diff --git a/src/include/tilexr_data_as_flag.h b/src/include/tilexr_data_as_flag.h index 93c64604..228dcf40 100644 --- a/src/include/tilexr_data_as_flag.h +++ b/src/include/tilexr_data_as_flag.h @@ -193,16 +193,12 @@ __aicore__ inline void DataAsFlagCopyScratchToDataAsFlagGM( AscendC::LocalTensor& sendScratch, uint32_t batchBlocks) { - AscendC::GlobalTensor dstGlobal; - dstGlobal.SetGlobalBuffer( - dstDataAsFlagGM + static_cast(dstBlockOffset) * DATA_AS_FLAG_BLOCK_BYTES); - AscendC::DataCopyExtParams outParams { - 1U, - batchBlocks * DATA_AS_FLAG_BLOCK_BYTES, - 0U, - 0U, - 0U}; - AscendC::DataCopyPad(dstGlobal, sendScratch, outParams); + AscendC::GlobalTensor dstGlobal; + dstGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ float*>( + dstDataAsFlagGM + static_cast(dstBlockOffset) * DATA_AS_FLAG_BLOCK_BYTES)); + AscendC::LocalTensor sendFloat = sendScratch.template ReinterpretCast(); + AscendC::DataCopy(dstGlobal, sendFloat, + batchBlocks * DATA_AS_FLAG_BLOCK_BYTES / sizeof(float)); } __aicore__ inline uint32_t DataAsFlagSend( diff --git a/src/include/tilexr_ep.h b/src/include/tilexr_ep.h index 82c79cde..7e7c4d7b 100644 --- a/src/include/tilexr_ep.h +++ b/src/include/tilexr_ep.h @@ -9,6 +9,16 @@ #include "tilexr_api.h" #include "tilexr_types.h" +namespace TileXREp { + +enum TileXREpActiveMaskType : int64_t { + TILEXR_EP_ACTIVE_MASK_NONE = 0, + TILEXR_EP_ACTIVE_MASK_TOKEN = 1, + TILEXR_EP_ACTIVE_MASK_EXPERT = 2, +}; + +} // namespace TileXREp + // This public API is C++ header-compatible because it reuses TileXR namespace datatypes. extern "C" { @@ -17,14 +27,30 @@ int TileXRMoeEpDispatch(void *x, int32_t *expertIds, TileXRCommPtr comm, void *expandXOut, int64_t *expertTokenNumsOut, int32_t *epRecvCountsOut, int32_t *assistInfoForCombineOut, TileXR::TileXRDataType dtype, aclrtStream stream); +int TileXRMoeEpDispatchMemory(void *x, int32_t *expertIds, TileXRCommPtr comm, + int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + void *expandXOut, int64_t *expertTokenNumsOut, int32_t *sendCountsOut, + int32_t *assistInfoForCombineOut, TileXR::TileXRDataType dtype, aclrtStream stream); + int TileXRMoeEpCombine(void *expertOut, int32_t *assistInfoForCombine, int32_t *epRecvCounts, TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, void *yOut, TileXR::TileXRDataType dtype, aclrtStream stream); +int TileXRMoeEpCombineMemory(void *expertOut, int32_t *assistInfoForCombine, int32_t *epRecvCounts, + TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + void *yOut, TileXR::TileXRDataType dtype, aclrtStream stream); + int TileXRMoeEpCombineV2(void *expertOut, int32_t *assistInfoForCombine, int32_t *epRecvCounts, TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, void *yOut, void *workspace, TileXR::TileXRDataType dtype, aclrtStream stream); +int TileXRMoeEpCombineMemoryV2(void *expertOut, int32_t *assistInfoForCombine, int32_t *sendCounts, + float *expertScales, bool *xActiveMask, int64_t activeMaskType, void *sharedExpertX, + TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, + int64_t epWorldSize, int64_t epRankId, int64_t tpWorldSize, int64_t tpRankId, + int64_t expertShardType, int64_t sharedExpertNum, int64_t sharedExpertRankNum, + int64_t quantMode, int64_t globalBs, void *yOut, TileXR::TileXRDataType dtype, aclrtStream stream); + int TileXRMoeEpDispatchV2(void *x, int32_t *expertIds, void *scales, bool *xActiveMask, void *expertScales, TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, int64_t epWorldSize, int64_t epRankId, int64_t tpWorldSize, int64_t tpRankId, int64_t expertShardType, int64_t sharedExpertNum, @@ -33,6 +59,15 @@ int TileXRMoeEpDispatchV2(void *x, int32_t *expertIds, void *scales, bool *xActi int32_t *tpRecvCountsOut, void *expandScalesOut, void *workspace, TileXR::TileXRDataType dtype, aclrtStream stream); +int TileXRMoeEpDispatchMemoryV2(void *x, int32_t *expertIds, void *scales, bool *xActiveMask, + int64_t activeMaskType, void *expertScales, TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, + int64_t moeExpertNum, + int64_t epWorldSize, int64_t epRankId, int64_t tpWorldSize, int64_t tpRankId, int64_t expertShardType, + int64_t sharedExpertNum, int64_t sharedExpertRankNum, int64_t quantMode, int64_t globalBs, + int64_t expertTokenNumsType, void *expandXOut, void *dynamicScalesOut, int32_t *assistInfoForCombineOut, + int64_t *expertTokenNumsOut, int32_t *sendCountsOut, int32_t *tpRecvCountsOut, void *expandScalesOut, + TileXR::TileXRDataType dtype, TileXR::TileXRDataType expandXOutDtype, aclrtStream stream); + } #endif diff --git a/tests/ep/CMakeLists.txt b/tests/ep/CMakeLists.txt index e505f1d8..7d624ed6 100644 --- a/tests/ep/CMakeLists.txt +++ b/tests/ep/CMakeLists.txt @@ -50,6 +50,7 @@ set(TILEXR_EP_TEST_INCLUDE_DIRS ${TILEXR_ROOT}/src/include ${TILEXR_ROOT}/src/ep/common ${TILEXR_ROOT}/src/ep/host + ${CMAKE_CURRENT_SOURCE_DIR}/demo ) foreach(dir "${ASCEND_DRIVER_PATH}/kernel/inc" @@ -77,6 +78,7 @@ endif() add_executable(test_tilexr_ep_layout unit/test_tilexr_ep_layout.cpp ${TILEXR_ROOT}/src/ep/host/ep_layout.cpp + ${TILEXR_ROOT}/src/ep/host/ep_memory_layout.cpp ) add_executable(test_tilexr_ep_api_sources @@ -88,7 +90,10 @@ add_executable(test_tilexr_ep_kernel_sources unit/test_tilexr_ep_kernel_sources. add_executable(test_tilexr_ep_host_validation unit/test_tilexr_ep_host_validation.cpp ${TILEXR_ROOT}/src/ep/host/ep_layout.cpp + ${TILEXR_ROOT}/src/ep/host/ep_memory_layout.cpp ${TILEXR_ROOT}/src/ep/host/ep_dispatch_host.cpp + ${TILEXR_ROOT}/src/ep/host/ep_dispatch_memory_host.cpp + ${TILEXR_ROOT}/src/ep/host/ep_combine_memory_host.cpp ) target_include_directories(test_tilexr_ep_layout PRIVATE diff --git a/tests/ep/README.md b/tests/ep/README.md index cace06ef..acb0334f 100644 --- a/tests/ep/README.md +++ b/tests/ep/README.md @@ -1,6 +1,6 @@ # TileXR EP Dispatch/Combine Tests -This tree tests the standalone TileXR EP module under `src/ep`. It is independent from hcomm, HCCL window helpers, and `ops-transformer`; the same-node route uses TileXR IPC peer-memory windows and `SyncCollectives`, while cross-node dispatch/combine use TileXR-registered UDMA workspaces. +This tree tests the standalone TileXR EP module under `src/ep`. It is independent from hcomm, HCCL window helpers, `ops-transformer`, and the removed `examples/mc2` tree; the same-node route uses TileXR IPC peer-memory windows, while cross-node dispatch/combine use TileXR-registered UDMA workspaces. ## Source-Only Tests @@ -27,7 +27,124 @@ bash build.sh full bash demo/run_tilexr_ep_dispatch_demo.sh 2 ``` -`full` mode builds and installs `tile-comm`, `tilexr-ep`, `libtilexr_ep_dispatch_kernel.so`, and `libtilexr_ep_combine_kernel.so` under the repository `install` directory, then builds the EP demo. +`full` mode builds and installs `tile-comm`, `tilexr-ep`, `libtilexr_ep_dispatch_kernel.so`, +`libtilexr_ep_dispatch_memory_kernel.so`, `libtilexr_ep_combine_kernel.so`, and +`libtilexr_ep_combine_memory_kernel.so` under the repository `install` directory, then builds the EP demo. + +The demo has exactly two execution backends, selected with `TILEXR_EP_DEMO_IMPL=udma|memory`. API version and +same-node/cross-node topology are handled inside the TileXR library and are not demo branches. Select the operator +sequence with `TILEXR_EP_DEMO_RUN_MODE=dispatch|combine|dispatch_combine`. + +The runner accepts: + +```text +run_tilexr_ep_dispatch_demo.sh rank_size npu_count first_npu loop_count impl bs h topk expert_ids \ + run_mode expert_mode expert_seed quant_mode mxfp8_format comm_quant_mode +``` + +`expert_ids` is a comma-, semicolon-, or space-separated list with exactly `bs * topk` entries. The same settings +can be supplied with `TILEXR_EP_DEMO_IMPL`, `TILEXR_EP_DEMO_BS`, `TILEXR_EP_DEMO_H`, +`TILEXR_EP_DEMO_TOPK`, `TILEXR_EP_DEMO_EXPERT_IDS`, `TILEXR_EP_DEMO_RUN_MODE`, +`TILEXR_EP_DEMO_EXPERT_MODE`, `TILEXR_EP_DEMO_EXPERT_SEED`, `TILEXR_EP_DEMO_QUANT_MODE`, +`TILEXR_EP_DEMO_MXFP8_FORMAT`, and `TILEXR_EP_DEMO_COMM_QUANT_MODE`. All values are runtime inputs; changing them +does not require rebuilding the operator. + +The run modes behave as follows: + +- `dispatch`: host constructs dispatch inputs and expected dispatch outputs, runs dispatch `loop_count` times on + one communicator, and validates only the final device outputs. +- `combine`: host constructs `expertOut`, assist tuples, send/receive counts, and expected `yOut`, runs combine + `loop_count` times, and validates only the final `yOut`. +- `dispatch_combine`: host constructs dispatch inputs and expected final output, runs dispatch `loop_count` times, + then passes the final `expandXOut`, assist tuples, and counts directly to one combine call. + +The expert modes are `uniform`, `random`, and `explicit`. `uniform` assigns the flattened token/topK routes to +experts in deterministic round-robin order. `random` uses a reproducible per-token random permutation and selects +the first `topk` expert IDs; set `expert_seed` to change it. `explicit` uses `expert_ids` and requires exactly +`bs * topk` IDs. Expert IDs are in `[0, moeExpertNum - 1]`, and `topk` must not exceed `moeExpertNum`. + +### MXFP8 Golden Tensors + +The test-side MXFP8 golden path uses `quant_mode=4` and supports `mxfp8_format=e4m3|e5m2`. It mirrors the +`pta-moe-test-main/quantize.py::mx_quantize` contract: + +- each 32 hidden elements share one E8M0 scale; +- the scale exponent is `floor(log2(maxAbs)) - emax`, with `emax=8` for E4M3 and `emax=15` for E5M2; +- MXFP8 element mantissas use round-to-nearest, ties-to-even; +- the per-row scale count is `align_up(ceil(h / 32), 2)`, with zero-block/padding scale byte `0x00`. + +The host constructs expected routed `expandX` as FP8 bytes and expected `dynamicScalesOut` as E8M0 bytes in the +same row order as the selected UDMA or memory dispatch backend. Dispatch MXFP8 remains restricted to dispatch-only +mode until the expert-compute/dequant stage is available. The byte-level golden implementation is covered by +`test_tilexr_ep_layout` for both E4M3 and E5M2. + +Combine communication quantization uses the independent `comm_quant_mode` setting: `0` disables it, `3` selects +MXFP8 E5M2, and `4` selects MXFP8 E4M3. For combine-only tests the host fills every valid `expertOut` row from its +assist tuple's source rank/token, rather than using constant ones. The golden path pads hidden rows to 32 elements, +performs the MXFP8 quantize/dequantize round trip, multiplies MoE routes by deterministic FP32 `expertScales`, adds +shared-expert routes with scale 1, and accumulates in token/topK order before comparing `yOut`. The same golden is +used when non-quantized dispatch output is passed directly to combine. Nonzero `comm_quant_mode` is limited to the +memory backend and supports MXFP8 E5M2 (`3`) and E4M3 (`4`). + +For example, run memory dispatch-only with deterministic uniform routes: + +```bash +bash demo/run_tilexr_ep_dispatch_demo.sh 8 8 0 100 memory 4 256 2 '' dispatch uniform 1 +``` + +Run UDMA combine-only with random routes generated from seed 2026: + +```bash +bash demo/run_tilexr_ep_dispatch_demo.sh 8 8 0 100 udma 4 256 2 '' combine random 2026 +``` + +Run dispatch+combine with an explicit expert list: + +```bash +bash demo/run_tilexr_ep_dispatch_demo.sh 8 8 0 100 memory 4 256 2 \ + '0,1,2,3,4,5,6,7' dispatch_combine explicit 1 +``` + +The memory dispatch path ports the A5 full-mesh dispatch algorithm and launches it in one kernel call. It supports +same-type FP16/BF16 dispatch and MXFP8 quantized dispatch (`quant_mode=4`) to E4M3 or E5M2. MXFP8 uses one E8M0 +scale per 32 hidden elements, pads the scale count to an even number, and returns the scales through +`dynamicScalesOut`. The path currently requires TP size one. It supports ordinary MoE experts, shared experts +deployed on one or more ranks per shared expert, token masks shaped `[bs]`, and expert masks shaped `[bs, topK]`. +The reference state/data windows and receive-count workspace are carved from TileXR IPC +peer memory internally; callers do not provide a workspace. + +For the memory API, `sendCountsOut` keeps the reference cumulative expert-major layout. Its element count is +`epWorldSize` on a shared-expert rank and `epWorldSize * localMoeExpertNum` on a MoE-expert rank. + +The memory combine path ports the A5 MTE combine flow into one kernel launch. Its current scope is non-quantized +FP16/BF16, TP size one, token masks, and shared experts. Expert masks, quantization, TP, and AddRmsNorm remain +unsupported. Run dispatch and combine together with: + +```bash +TILEXR_EP_DEMO_IMPL=memory \ + bash demo/run_tilexr_ep_dispatch_demo.sh 2 +``` + +The fourth argument is the loop count and defaults to `100`. In dispatch and dispatch+combine modes it controls the +number of dispatch calls; in combine-only mode it controls the number of combine calls. Each repeated call reaches +a stream/rank completion point before its communication window is reused, and only the final output is copied back. +For example, this runs 100 memory dispatches followed by one memory combine: + +```bash +bash demo/run_tilexr_ep_dispatch_demo.sh 8 8 0 100 memory +``` + +Set `TILEXR_EP_DEMO_LOOP=` when using the environment instead of the fourth argument. Use a fourth argument of +`1` for a single dispatch+combine smoke test. + +Use `TILEXR_EP_DEMO_ACTIVE_MASK_TYPE=none|token|expert` to select the active-mask contract. The legacy +`TILEXR_EP_DEMO_ACTIVE_MASK=1` setting remains an alias for token masking. Use +`TILEXR_EP_DEMO_DTYPE=fp16|bf16` to select the input/output data type. For example: + +```bash +TILEXR_EP_DEMO_IMPL=memory TILEXR_EP_DEMO_ACTIVE_MASK_TYPE=token TILEXR_EP_DEMO_DTYPE=bf16 \ + bash demo/run_tilexr_ep_dispatch_demo.sh 2 +``` ## Remote Verification @@ -39,6 +156,7 @@ bash demo/deploy_and_run_remote.sh The remote verification script syncs the complete repository into `${TILEXR_EP_REMOTE_BASE}/TileXR` on `${TILEXR_EP_REMOTE}`, initializes submodules, sources `scripts/common_env.sh`, builds the full EP artifacts, and runs the two-rank dispatch demo. -## Cross-Node UDMA +## UDMA Workspace -Cross-node dispatch/combine require a workspace allocated by the caller and registered with `TileXRUDMARegister`. The demo allocates a cache-line-aligned workspace, registers it before dispatch/combine, and validates both dispatch and combine outputs after all ranks synchronize. +The UDMA backend always allocates and registers its aligned workspace. The same demo path is used for same-node and +cross-node runs; topology-specific behavior remains inside TileXR. diff --git a/tests/ep/demo/mxfp8_golden.h b/tests/ep/demo/mxfp8_golden.h new file mode 100644 index 00000000..9162a8a2 --- /dev/null +++ b/tests/ep/demo/mxfp8_golden.h @@ -0,0 +1,201 @@ +#ifndef TILEXR_TESTS_EP_DEMO_MXFP8_GOLDEN_H +#define TILEXR_TESTS_EP_DEMO_MXFP8_GOLDEN_H + +#include +#include +#include +#include +#include +#include + +namespace TileXREpDemo { + +constexpr std::size_t kMxfp8BlockSize = 32; + +enum class Mxfp8Format { + E4M3, + E5M2, +}; + +struct Mxfp8Tensor { + std::vector elements; + std::vector scales; + std::size_t scaleCountPerRow = 0; +}; + +inline std::size_t Mxfp8ScaleCountPerRow(std::size_t h) +{ + const std::size_t blockCount = (h + kMxfp8BlockSize - 1) / kMxfp8BlockSize; + return blockCount + (blockCount & 1U); +} + +inline float RoundToNearestEven(float value) +{ + const float lower = std::floor(value); + const float fraction = value - lower; + if (fraction < 0.5f) { + return lower; + } + if (fraction > 0.5f) { + return lower + 1.0f; + } + return std::fmod(lower, 2.0f) == 0.0f ? lower : lower + 1.0f; +} + +inline uint8_t EncodeFp8(float value, Mxfp8Format format) +{ + const int exponentBits = format == Mxfp8Format::E4M3 ? 4 : 5; + const int mantissaBits = format == Mxfp8Format::E4M3 ? 3 : 2; + const int exponentBias = format == Mxfp8Format::E4M3 ? 7 : 15; + const float maxFinite = format == Mxfp8Format::E4M3 ? 448.0f : 57344.0f; + const uint8_t sign = std::signbit(value) ? 0x80U : 0U; + float magnitude = std::min(std::fabs(value), maxFinite); + if (magnitude == 0.0f) { + return sign; + } + + const int minNormalExponent = 1 - exponentBias; + const float minNormal = std::ldexp(1.0f, minNormalExponent); + int encodedExponent = 0; + int encodedMantissa = 0; + if (magnitude < minNormal) { + const float subnormalStep = std::ldexp(1.0f, minNormalExponent - mantissaBits); + encodedMantissa = static_cast(RoundToNearestEven(magnitude / subnormalStep)); + if (encodedMantissa >= (1 << mantissaBits)) { + encodedExponent = 1; + encodedMantissa = 0; + } + } else { + int exponent = static_cast(std::floor(std::log2(magnitude))); + const float normalized = std::ldexp(magnitude, -exponent); + encodedMantissa = static_cast(RoundToNearestEven( + (normalized - 1.0f) * static_cast(1 << mantissaBits))); + if (encodedMantissa == (1 << mantissaBits)) { + ++exponent; + encodedMantissa = 0; + } + encodedExponent = exponent + exponentBias; + } + const int exponentMask = (1 << exponentBits) - 1; + return static_cast(sign | + (static_cast(encodedExponent & exponentMask) << mantissaBits) | + static_cast(encodedMantissa)); +} + +inline float DecodeFp8(uint8_t value, Mxfp8Format format) +{ + const int mantissaBits = format == Mxfp8Format::E4M3 ? 3 : 2; + const int exponentBits = format == Mxfp8Format::E4M3 ? 4 : 5; + const int exponentBias = format == Mxfp8Format::E4M3 ? 7 : 15; + const int exponentMask = (1 << exponentBits) - 1; + const int mantissaMask = (1 << mantissaBits) - 1; + const int exponent = (value >> mantissaBits) & exponentMask; + const int mantissa = value & mantissaMask; + const float sign = (value & 0x80U) == 0 ? 1.0f : -1.0f; + + if (exponent == 0) { + if (mantissa == 0) { + return std::copysign(0.0f, sign); + } + return sign * std::ldexp(static_cast(mantissa), + 1 - exponentBias - mantissaBits); + } + if (format == Mxfp8Format::E5M2 && exponent == exponentMask) { + return mantissa == 0 ? sign * std::numeric_limits::infinity() : + std::numeric_limits::quiet_NaN(); + } + if (format == Mxfp8Format::E4M3 && exponent == exponentMask && mantissa == mantissaMask) { + return std::numeric_limits::quiet_NaN(); + } + return sign * std::ldexp(1.0f + static_cast(mantissa) / + static_cast(1 << mantissaBits), exponent - exponentBias); +} + +inline void QuantizeMxfp8Row(const float *input, std::size_t h, Mxfp8Format format, + uint8_t *elements, uint8_t *scales) +{ + const std::size_t scaleCount = Mxfp8ScaleCountPerRow(h); + std::fill(scales, scales + scaleCount, 0U); + const int elementMaxExponent = format == Mxfp8Format::E4M3 ? 8 : 15; + const int exponentBits = format == Mxfp8Format::E4M3 ? 4 : 5; + const int mantissaBits = format == Mxfp8Format::E4M3 ? 3 : 2; + const int minPrivateExponent = -(1 << (exponentBits - 1)) + 2; + const float maxFinite = format == Mxfp8Format::E4M3 ? 448.0f : 57344.0f; + + const std::size_t blockCount = (h + kMxfp8BlockSize - 1) / kMxfp8BlockSize; + for (std::size_t block = 0; block < blockCount; ++block) { + const std::size_t begin = block * kMxfp8BlockSize; + const std::size_t end = std::min(begin + kMxfp8BlockSize, h); + float maxAbs = 0.0f; + for (std::size_t index = begin; index < end; ++index) { + maxAbs = std::max(maxAbs, std::fabs(input[index])); + } + int sharedExponent = maxAbs == 0.0f ? -127 : + static_cast(std::floor(std::log2(maxAbs))) - elementMaxExponent; + sharedExponent = std::max(-127, std::min(127, sharedExponent)); + scales[block] = static_cast(sharedExponent + 127); + const float sharedScale = std::ldexp(1.0f, sharedExponent); + + for (std::size_t index = begin; index < end; ++index) { + const float scaled = input[index] / sharedScale; + int privateExponent = scaled == 0.0f ? 0 : + static_cast(std::floor(std::log2(std::fabs(scaled)))); + privateExponent = std::max(privateExponent, minPrivateExponent); + const float privateScale = std::ldexp(1.0f, privateExponent); + float quantized = RoundToNearestEven( + scaled / privateScale * static_cast(1 << mantissaBits)); + quantized = quantized / static_cast(1 << mantissaBits) * privateScale; + quantized = std::max(-maxFinite, std::min(maxFinite, quantized)); + elements[index] = EncodeFp8(quantized, format); + } + } +} + +inline Mxfp8Tensor QuantizeMxfp8(const std::vector &input, std::size_t rows, + std::size_t h, Mxfp8Format format) +{ + Mxfp8Tensor result; + if (rows == 0 || h == 0 || input.size() != rows * h) { + return result; + } + result.scaleCountPerRow = Mxfp8ScaleCountPerRow(h); + result.elements.resize(rows * h); + result.scales.resize(rows * result.scaleCountPerRow); + for (std::size_t row = 0; row < rows; ++row) { + QuantizeMxfp8Row(input.data() + row * h, h, format, + result.elements.data() + row * h, + result.scales.data() + row * result.scaleCountPerRow); + } + return result; +} + +inline std::vector DequantizeMxfp8(const Mxfp8Tensor &input, std::size_t rows, + std::size_t h, Mxfp8Format format) +{ + const std::size_t scaleCount = Mxfp8ScaleCountPerRow(h); + if (rows == 0 || h == 0 || input.scaleCountPerRow != scaleCount || + input.elements.size() != rows * h || input.scales.size() != rows * scaleCount) { + return {}; + } + + std::vector result(rows * h, 0.0f); + for (std::size_t row = 0; row < rows; ++row) { + for (std::size_t index = 0; index < h; ++index) { + const std::size_t block = index / kMxfp8BlockSize; + const uint8_t scale = input.scales[row * scaleCount + block]; + result[row * h + index] = DecodeFp8(input.elements[row * h + index], format) * + std::ldexp(1.0f, static_cast(scale) - 127); + } + } + return result; +} + +inline std::vector RoundTripMxfp8(const std::vector &input, std::size_t rows, + std::size_t h, Mxfp8Format format) +{ + return DequantizeMxfp8(QuantizeMxfp8(input, rows, h, format), rows, h, format); +} + +} // namespace TileXREpDemo + +#endif // TILEXR_TESTS_EP_DEMO_MXFP8_GOLDEN_H diff --git a/tests/ep/demo/run_tilexr_ep_dispatch_demo.sh b/tests/ep/demo/run_tilexr_ep_dispatch_demo.sh index f7baff33..02b980d2 100755 --- a/tests/ep/demo/run_tilexr_ep_dispatch_demo.sh +++ b/tests/ep/demo/run_tilexr_ep_dispatch_demo.sh @@ -9,6 +9,87 @@ INSTALL_DIR="${EP_DIR}/install" rank_size="${1:-2}" npu_count="${2:-${rank_size}}" first_npu="${3:-0}" +loop_count="${4:-${TILEXR_EP_DEMO_LOOP:-100}}" +impl="${5:-${TILEXR_EP_DEMO_IMPL:-udma}}" +bs="${6:-${TILEXR_EP_DEMO_BS:-4}}" +h="${7:-${TILEXR_EP_DEMO_H:-8}}" +topk="${8:-${TILEXR_EP_DEMO_TOPK:-2}}" +expert_ids="${9:-${TILEXR_EP_DEMO_EXPERT_IDS:-}}" +run_mode="${10:-${TILEXR_EP_DEMO_RUN_MODE:-dispatch_combine}}" +expert_mode="${11:-${TILEXR_EP_DEMO_EXPERT_MODE:-}}" +expert_seed="${12:-${TILEXR_EP_DEMO_EXPERT_SEED:-1}}" +quant_mode="${13:-${TILEXR_EP_DEMO_QUANT_MODE:-0}}" +mxfp8_format="${14:-${TILEXR_EP_DEMO_MXFP8_FORMAT:-e4m3}}" +comm_quant_mode="${15:-${TILEXR_EP_DEMO_COMM_QUANT_MODE:-0}}" + +if [[ -z "${expert_mode}" ]]; then + if [[ -n "${expert_ids}" ]]; then + expert_mode="explicit" + else + expert_mode="uniform" + fi +fi + +for value_name in loop_count bs h topk; do + value="${!value_name}" + if [[ ! "${value}" =~ ^[1-9][0-9]*$ ]]; then + echo "${value_name} must be a positive integer, got: ${value}" >&2 + exit 2 + fi +done +if [[ "${impl}" != "udma" && "${impl}" != "memory" ]]; then + echo "impl must be udma or memory, got: ${impl}" >&2 + exit 2 +fi +if [[ "${run_mode}" != "dispatch" && "${run_mode}" != "combine" && "${run_mode}" != "dispatch_combine" ]]; then + echo "run_mode must be dispatch, combine, or dispatch_combine, got: ${run_mode}" >&2 + exit 2 +fi +if [[ "${expert_mode}" != "uniform" && "${expert_mode}" != "random" && "${expert_mode}" != "explicit" ]]; then + echo "expert_mode must be uniform, random, or explicit, got: ${expert_mode}" >&2 + exit 2 +fi +if [[ ! "${expert_seed}" =~ ^[0-9]+$ ]]; then + echo "expert_seed must be a non-negative integer, got: ${expert_seed}" >&2 + exit 2 +fi +if [[ "${quant_mode}" != "0" && "${quant_mode}" != "4" ]]; then + echo "quant_mode must be 0 or 4 (MXFP8), got: ${quant_mode}" >&2 + exit 2 +fi +if [[ "${mxfp8_format}" != "e4m3" && "${mxfp8_format}" != "e5m2" && + "${mxfp8_format}" != "fp8_e4m3fn" && "${mxfp8_format}" != "fp8_e5m2" ]]; then + echo "mxfp8_format must be e4m3 or e5m2, got: ${mxfp8_format}" >&2 + exit 2 +fi +if [[ "${quant_mode}" == "4" && "${run_mode}" != "dispatch" ]]; then + echo "MXFP8 golden generation currently supports dispatch-only mode" >&2 + exit 2 +fi +if [[ "${quant_mode}" == "4" && "${impl}" != "memory" ]]; then + echo "MXFP8 dispatch currently requires the memory backend" >&2 + exit 2 +fi +if [[ "${comm_quant_mode}" != "0" && "${comm_quant_mode}" != "3" && "${comm_quant_mode}" != "4" ]]; then + echo "comm_quant_mode must be 0, 3 (E5M2), or 4 (E4M3), got: ${comm_quant_mode}" >&2 + exit 2 +fi +if [[ "${comm_quant_mode}" != "0" && "${run_mode}" == "dispatch" ]]; then + echo "comm_quant_mode requires combine or dispatch_combine mode" >&2 + exit 2 +fi +if [[ "${comm_quant_mode}" != "0" && "${impl}" != "memory" ]]; then + echo "nonzero comm_quant_mode currently requires the memory backend" >&2 + exit 2 +fi +if [[ "${expert_mode}" == "explicit" && -z "${expert_ids}" ]]; then + echo "expert_ids is required when expert_mode=explicit" >&2 + exit 2 +fi +if [[ "${expert_mode}" != "explicit" && -n "${expert_ids}" ]]; then + echo "expert_ids must be empty unless expert_mode=explicit" >&2 + exit 2 +fi : "${ASCEND_HOME_PATH:=}" : "${LD_LIBRARY_PATH:=}" @@ -17,6 +98,18 @@ source "${TILEXR_ROOT}/scripts/common_env.sh" export TILEXR_COMM_ID="${TILEXR_COMM_ID:-127.0.0.1:10077}" export TILEXR_DEMO_NPUS="${npu_count}" export TILEXR_DEMO_FIRST_NPU="${first_npu}" +export TILEXR_EP_DEMO_LOOP="${loop_count}" +export TILEXR_EP_DEMO_IMPL="${impl}" +export TILEXR_EP_DEMO_BS="${bs}" +export TILEXR_EP_DEMO_H="${h}" +export TILEXR_EP_DEMO_TOPK="${topk}" +export TILEXR_EP_DEMO_EXPERT_IDS="${expert_ids}" +export TILEXR_EP_DEMO_RUN_MODE="${run_mode}" +export TILEXR_EP_DEMO_EXPERT_MODE="${expert_mode}" +export TILEXR_EP_DEMO_EXPERT_SEED="${expert_seed}" +export TILEXR_EP_DEMO_QUANT_MODE="${quant_mode}" +export TILEXR_EP_DEMO_MXFP8_FORMAT="${mxfp8_format}" +export TILEXR_EP_DEMO_COMM_QUANT_MODE="${comm_quant_mode}" export LD_LIBRARY_PATH="${TILEXR_ROOT}/install/lib64:${TILEXR_ROOT}/install/lib:${INSTALL_DIR}/lib64:${INSTALL_DIR}/lib:${LD_LIBRARY_PATH:-}" bin="${INSTALL_DIR}/bin/tilexr_ep_dispatch_demo" diff --git a/tests/ep/demo/tilexr_ep_dispatch_demo.cpp b/tests/ep/demo/tilexr_ep_dispatch_demo.cpp index e3f34b77..2335f49b 100644 --- a/tests/ep/demo/tilexr_ep_dispatch_demo.cpp +++ b/tests/ep/demo/tilexr_ep_dispatch_demo.cpp @@ -6,6 +6,8 @@ #include #include #include +#include +#include #include #include #include @@ -16,20 +18,19 @@ #include #include "acl/acl.h" +#include "mxfp8_golden.h" #include "tilexr_api.h" #include "tilexr_ep.h" #include "tilexr_types.h" namespace { -constexpr int64_t kBs = 4; -constexpr int64_t kH = 8; -constexpr int64_t kTopK = 2; -constexpr int64_t kRoutes = kBs * kTopK; -constexpr int64_t kXElements = kBs * kH; +constexpr int64_t kDefaultBs = 4; +constexpr int64_t kDefaultH = 8; +constexpr int64_t kDefaultTopK = 2; constexpr int64_t kAssistInts = 4; constexpr uint16_t kFp16One = 0x3c00; -constexpr uint16_t kFp16Two = 0x4000; +constexpr uint16_t kBf16One = 0x3f80; constexpr std::size_t kUdmaCacheLineBytes = 64; constexpr std::size_t kUdmaRegistrationAlignment = 2 * 1024 * 1024; @@ -37,6 +38,9 @@ TileXRUDMAMemHandle g_workspaceHandle = 0; bool g_workspaceRegistered = false; struct DemoConfig { + int64_t bs = kDefaultBs; + int64_t h = kDefaultH; + int64_t topK = kDefaultTopK; int64_t moeExpertNum = 8; int64_t sharedExpertNum = 0; int64_t sharedExpertRankNum = 0; @@ -45,7 +49,7 @@ struct DemoConfig { int64_t maxRoutesPerRank() const { - return kBs * (kTopK + sharedExpertNum); + return bs * (topK + sharedExpertNum); } int64_t effectiveTpWorldSize() const @@ -55,8 +59,27 @@ struct DemoConfig { int64_t expandedElements() const { - return maxRoutesPerRank() * effectiveTpWorldSize() * kH; + return maxRoutesPerRank() * effectiveTpWorldSize() * h; } + + std::vector expertIds; +}; + +enum class DemoBackend { + UDMA, + MEMORY, +}; + +enum class DemoRunMode { + DISPATCH, + COMBINE, + DISPATCH_COMBINE, +}; + +enum class ExpertListMode { + UNIFORM, + RANDOM, + EXPLICIT, }; struct HostPort { @@ -79,6 +102,57 @@ int GetEnvInt(const char *name, int fallback) return std::atoi(value); } +bool GetEnvUint32(const char *name, uint32_t fallback, uint32_t *value) +{ + if (value == nullptr) { + return false; + } + const char *text = std::getenv(name); + if (text == nullptr || text[0] == '\0') { + *value = fallback; + return true; + } + errno = 0; + char *end = nullptr; + const unsigned long parsed = std::strtoul(text, &end, 10); + if (end == text || *end != '\0' || errno != 0 || parsed > std::numeric_limits::max()) { + return false; + } + *value = static_cast(parsed); + return true; +} + +bool ParseInt32List(const char *text, std::vector *values) +{ + if (text == nullptr || values == nullptr) { + return false; + } + values->clear(); + const char *cursor = text; + while (*cursor != '\0') { + while (*cursor == ',' || *cursor == ';' || *cursor == ' ' || *cursor == '\t' || *cursor == '\n') { + ++cursor; + } + if (*cursor == '\0') { + break; + } + errno = 0; + char *end = nullptr; + const long parsed = std::strtol(cursor, &end, 10); + if (end == cursor || errno != 0 || parsed < std::numeric_limits::min() || + parsed > std::numeric_limits::max()) { + return false; + } + values->push_back(static_cast(parsed)); + cursor = end; + if (*cursor != '\0' && *cursor != ',' && *cursor != ';' && *cursor != ' ' && *cursor != '\t' && + *cursor != '\n') { + return false; + } + } + return !values->empty(); +} + bool ParseHostPort(const std::string &text, HostPort *out) { const std::size_t pos = text.rfind(':'); @@ -310,10 +384,11 @@ std::size_t AlignSize(std::size_t value, std::size_t alignment) return remainder == 0 ? value : value + alignment - remainder; } -std::size_t EpSlotBytes(const DemoConfig &config, std::size_t payloadRowBytes, bool usePerTokenDynamicQuant) +std::size_t EpSlotBytes( + const DemoConfig &config, std::size_t payloadRowBytes, std::size_t payloadScaleBytesPerRow) { - const std::size_t payloadScaleBytes = usePerTokenDynamicQuant ? - static_cast(config.maxRoutesPerRank()) * sizeof(float) : 0; + const std::size_t payloadScaleBytes = + static_cast(config.maxRoutesPerRank()) * payloadScaleBytesPerRow; const std::size_t payloadBytes = AlignSize( static_cast(config.maxRoutesPerRank()) * payloadRowBytes + payloadScaleBytes, 32); const std::size_t assistWindowBytes = AlignSize( @@ -322,35 +397,45 @@ std::size_t EpSlotBytes(const DemoConfig &config, std::size_t payloadRowBytes, b } std::size_t EpWindowBytes(int rankSize, const DemoConfig &config, std::size_t payloadRowBytes, - bool usePerTokenDynamicQuant) + std::size_t payloadScaleBytesPerRow) { return AlignSize(64 + static_cast(rankSize) * - EpSlotBytes(config, payloadRowBytes, usePerTokenDynamicQuant), 32); + EpSlotBytes(config, payloadRowBytes, payloadScaleBytesPerRow), 32); } std::size_t EpOperationBytes(int rankSize, const DemoConfig &config, std::size_t payloadRowBytes, - bool usePerTokenDynamicQuant) + std::size_t payloadScaleBytesPerRow) { - const std::size_t windowBytes = EpWindowBytes(rankSize, config, payloadRowBytes, usePerTokenDynamicQuant); + const std::size_t windowBytes = EpWindowBytes(rankSize, config, payloadRowBytes, payloadScaleBytesPerRow); const std::size_t readyOffset = windowBytes * 2 + static_cast(rankSize) * sizeof(uint64_t); const std::size_t relayOffset = AlignSize(readyOffset, kUdmaCacheLineBytes); const std::size_t relayBytes = static_cast(rankSize) * static_cast(rankSize) * - EpSlotBytes(config, payloadRowBytes, usePerTokenDynamicQuant); + EpSlotBytes(config, payloadRowBytes, payloadScaleBytesPerRow); const std::size_t relayReadyOffset = AlignSize(relayOffset + relayBytes, kUdmaCacheLineBytes); const std::size_t relayReadyBytes = static_cast(rankSize) * sizeof(uint64_t); return AlignSize(relayReadyOffset + relayReadyBytes, kUdmaCacheLineBytes); } std::size_t EpRequiredWorkspaceBytes(int rankSize, const DemoConfig &config, std::size_t payloadRowBytes, - bool usePerTokenDynamicQuant) + std::size_t payloadScaleBytesPerRow) { - const std::size_t operationBytes = EpOperationBytes(rankSize, config, payloadRowBytes, usePerTokenDynamicQuant); + const std::size_t operationBytes = EpOperationBytes(rankSize, config, payloadRowBytes, payloadScaleBytesPerRow); return AlignSize(operationBytes * 2 + sizeof(uint64_t), kUdmaCacheLineBytes); } uint16_t XValue(int rank, int64_t token, int64_t h) { - return static_cast(0x3c00 + rank * 0x0400 + token * 0x0100 + h * 0x0010); + return static_cast( + 0x3c00 + rank * 0x0400 + token * 0x0100 + (h % 32) * 0x0010); +} + +uint16_t InputValue(TileXR::TileXRDataType dtype, int rank, int64_t token, int64_t h) +{ + if (dtype == TileXR::TILEXR_DATA_TYPE_BFP16) { + return static_cast( + 0x3f80 + rank * 0x0080 + token * 0x0010 + (h % 32)); + } + return XValue(rank, token, h); } float HalfBitsToFloat(uint16_t bits) @@ -382,6 +467,17 @@ float HalfBitsToFloat(uint16_t bits) return static_cast(sign) * value; } +float DataBitsToFloat(uint16_t bits, TileXR::TileXRDataType dtype) +{ + if (dtype == TileXR::TILEXR_DATA_TYPE_BFP16) { + const uint32_t fp32Bits = static_cast(bits) << 16U; + float value = 0.0f; + std::memcpy(&value, &fp32Bits, sizeof(value)); + return value; + } + return HalfBitsToFloat(bits); +} + int8_t QuantizedXValue(int rank, int64_t token, int64_t h, float scale) { const float value = HalfBitsToFloat(XValue(rank, token, h)) * scale; @@ -394,40 +490,118 @@ int8_t QuantizedXValue(int rank, int64_t token, int64_t h, float scale) return static_cast(rounded); } -float DynamicScaleForXValue(int rank, int64_t token) +float DynamicScaleForXValue(int rank, int64_t token, int64_t hSize) { float maxAbs = 0.0f; - for (int64_t h = 0; h < kH; ++h) { + for (int64_t h = 0; h < hSize; ++h) { const float value = std::fabs(HalfBitsToFloat(XValue(rank, token, h))); maxAbs = std::max(maxAbs, value); } return maxAbs > 0.0f ? maxAbs / 127.0f : 1.0f; } -int8_t DynamicQuantizedXValue(int rank, int64_t token, int64_t h) +int8_t DynamicQuantizedXValue(int rank, int64_t token, int64_t h, int64_t hSize) { - const float scale = DynamicScaleForXValue(rank, token); + const float scale = DynamicScaleForXValue(rank, token, hSize); return QuantizedXValue(rank, token, h, scale > 0.0f ? 1.0f / scale : 1.0f); } +float ExpertScaleValue(int64_t token, int64_t topKId) +{ + return 0.5f + static_cast(token) * 0.0625f + static_cast(topKId) * 0.125f; +} + +std::vector BuildExpertScales(const DemoConfig &config) +{ + std::vector scales(static_cast(config.bs * config.topK)); + for (int64_t token = 0; token < config.bs; ++token) { + for (int64_t topKId = 0; topKId < config.topK; ++topKId) { + scales[static_cast(token * config.topK + topKId)] = + ExpertScaleValue(token, topKId); + } + } + return scales; +} + std::vector ExpertIds(const DemoConfig &config) { - std::vector expertIds(kRoutes); - for (int64_t route = 0; route < kRoutes; ++route) { - expertIds[route] = static_cast(route % config.moeExpertNum); + return config.expertIds; +} + +bool BuildExpertIds(DemoConfig *config, ExpertListMode mode, uint32_t seed) +{ + if (config == nullptr || config->bs <= 0 || config->topK <= 0 || config->moeExpertNum <= 0 || + config->topK > config->moeExpertNum) { + return false; } - return expertIds; + const std::size_t routeCount = static_cast(config->bs * config->topK); + if (mode == ExpertListMode::EXPLICIT) { + if (config->expertIds.size() != routeCount) { + return false; + } + return std::all_of(config->expertIds.begin(), config->expertIds.end(), + [config](int32_t expertId) { return expertId >= 0 && expertId < config->moeExpertNum; }); + } + + config->expertIds.assign(routeCount, 0); + if (mode == ExpertListMode::UNIFORM) { + for (std::size_t route = 0; route < routeCount; ++route) { + config->expertIds[route] = static_cast(route % config->moeExpertNum); + } + return true; + } + + std::mt19937 generator(seed); + std::vector candidates(static_cast(config->moeExpertNum)); + for (int64_t expertId = 0; expertId < config->moeExpertNum; ++expertId) { + candidates[static_cast(expertId)] = static_cast(expertId); + } + for (int64_t token = 0; token < config->bs; ++token) { + std::shuffle(candidates.begin(), candidates.end(), generator); + for (int64_t topKId = 0; topKId < config->topK; ++topKId) { + config->expertIds[static_cast(token * config->topK + topKId)] = + candidates[static_cast(topKId)]; + } + } + return true; } -std::vector ActiveMask(bool enabled) +std::vector ActiveMask(int64_t activeMaskType, const DemoConfig &config) { - std::vector mask(kBs, 1); - if (enabled && kBs > 0) { - mask[kBs - 1] = 0; + if (activeMaskType == TileXREp::TILEXR_EP_ACTIVE_MASK_NONE) { + return {}; + } + const int64_t elementCount = activeMaskType == TileXREp::TILEXR_EP_ACTIVE_MASK_EXPERT ? + config.bs * config.topK : config.bs; + std::vector mask(static_cast(elementCount), 1); + if (!mask.empty()) { + mask.back() = 0; } return mask; } +bool IsRouteActive(const std::vector &activeMask, const DemoConfig &config, + int64_t token, int64_t topKId) +{ + if (activeMask.empty()) { + return true; + } + if (activeMask.size() == static_cast(config.bs)) { + return activeMask[static_cast(token)] != 0; + } + return activeMask[static_cast(token * config.topK + topKId)] != 0; +} + +bool IsTokenActive(const std::vector &activeMask, const DemoConfig &config, int64_t token) +{ + for (int64_t topKId = 0; topKId < config.topK; ++topKId) { + if (IsRouteActive(activeMask, config, token, topKId)) { + return true; + } + } + return false; +} + int64_t LocalExpertNum(int rankSize, const DemoConfig &config) { const int64_t expertRankSize = static_cast(rankSize) / config.effectiveTpWorldSize(); @@ -438,12 +612,22 @@ int64_t LocalExpertNum(int rankSize, const DemoConfig &config) return config.moeExpertNum / moeRankNum; } +int64_t OutputLocalExpertNum(int rank, int rankSize, const DemoConfig &config) +{ + return rank < config.sharedExpertRankNum ? 1 : LocalExpertNum(rankSize, config); +} + +std::size_t MemorySendCountsCount(int rank, int rankSize, const DemoConfig &config) +{ + return static_cast(OutputLocalExpertNum(rank, rankSize, config) * rankSize); +} + int ExpertRankForRank(int rank, const DemoConfig &config) { return static_cast(rank / config.effectiveTpWorldSize()); } -int64_t DstRankForExpert(int32_t globalExpertId, int rankSize, const DemoConfig &config) +int64_t DstRankForExpert(int32_t globalExpertId, int srcRank, int rankSize, const DemoConfig &config) { const int64_t localExpertNum = LocalExpertNum(rankSize, config); if (globalExpertId < 0 || localExpertNum <= 0) { @@ -451,7 +635,8 @@ int64_t DstRankForExpert(int32_t globalExpertId, int rankSize, const DemoConfig } const int64_t expertRankSize = static_cast(rankSize) / config.effectiveTpWorldSize(); if (globalExpertId < config.sharedExpertNum) { - return globalExpertId < config.sharedExpertRankNum ? globalExpertId : -1; + const int64_t rankNumPerSharedExpert = config.sharedExpertRankNum / config.sharedExpertNum; + return srcRank % rankNumPerSharedExpert + globalExpertId * rankNumPerSharedExpert; } const int64_t moeExpertId = static_cast(globalExpertId) - config.sharedExpertNum; const int64_t dstRank = config.sharedExpertRankNum + moeExpertId / localExpertNum; @@ -465,14 +650,14 @@ int64_t LocalExpertForExpert(int32_t globalExpertId, int rankSize, const DemoCon return -1; } if (globalExpertId < config.sharedExpertNum) { - return globalExpertId < config.sharedExpertRankNum ? globalExpertId : -1; + return 0; } return (static_cast(globalExpertId) - config.sharedExpertNum) % localExpertNum; } -bool RouteBelongsToRank(int32_t globalExpertId, int rank, int rankSize, const DemoConfig &config) +bool RouteBelongsToRank(int32_t globalExpertId, int srcRank, int rank, int rankSize, const DemoConfig &config) { - return DstRankForExpert(globalExpertId, rankSize, config) == ExpertRankForRank(rank, config); + return DstRankForExpert(globalExpertId, srcRank, rankSize, config) == ExpertRankForRank(rank, config); } struct ExpectedRoute { @@ -493,20 +678,24 @@ std::vector BuildExpectedRoutes( if (effectiveTpWorldSize > 1 && srcRank % effectiveTpWorldSize != targetTpRankId) { continue; } - for (int64_t token = 0; token < kBs; ++token) { - if (!activeMask.empty() && activeMask[token] == 0) { + for (int64_t token = 0; token < config.bs; ++token) { + if (!IsTokenActive(activeMask, config, token)) { continue; } for (int64_t sharedExpertId = 0; sharedExpertId < config.sharedExpertNum; ++sharedExpertId) { - if (RouteBelongsToRank(static_cast(sharedExpertId), rank, rankSize, config)) { + if (RouteBelongsToRank( + static_cast(sharedExpertId), srcRank, rank, rankSize, config)) { expected.push_back(ExpectedRoute {srcRank, static_cast(token), - static_cast(kTopK + sharedExpertId), static_cast(sharedExpertId)}); + static_cast(config.topK + sharedExpertId), static_cast(sharedExpertId)}); } } - for (int64_t topKId = 0; topKId < kTopK; ++topKId) { - const int64_t route = token * kTopK + topKId; + for (int64_t topKId = 0; topKId < config.topK; ++topKId) { + if (!IsRouteActive(activeMask, config, token, topKId)) { + continue; + } + const int64_t route = token * config.topK + topKId; const int32_t expertId = static_cast(config.sharedExpertNum) + expertIds[route]; - if (RouteBelongsToRank(expertId, rank, rankSize, config)) { + if (RouteBelongsToRank(expertId, srcRank, rank, rankSize, config)) { expected.push_back(ExpectedRoute {srcRank, static_cast(token), static_cast(topKId), expertId}); } @@ -536,20 +725,125 @@ std::vector BuildExpectedTpRoutes( return expected; } +std::vector BuildExpectedMemoryRoutes( + int rank, int rankSize, const DemoConfig &config, const std::vector &activeMask) +{ + const std::vector sourceMajor = BuildExpectedRoutes(rank, rankSize, config, activeMask); + const int64_t localExpertNum = OutputLocalExpertNum(rank, rankSize, config); + std::vector expected; + expected.reserve(sourceMajor.size()); + for (int64_t localExpert = 0; localExpert < localExpertNum; ++localExpert) { + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + for (const ExpectedRoute &route : sourceMajor) { + if (route.srcRank == srcRank && + LocalExpertForExpert(route.expertId, rankSize, config) == localExpert) { + expected.push_back(route); + } + } + } + } + return expected; +} + +std::vector BuildExpectedRecvCounts( + int rank, int rankSize, const DemoConfig &config, const std::vector &activeMask, bool useMemory) +{ + const std::vector localExpected = BuildExpectedRoutes(rank, rankSize, config, activeMask); + std::vector expected(useMemory ? MemorySendCountsCount(rank, rankSize, config) : + static_cast(rankSize), 0); + if (!useMemory) { + for (const ExpectedRoute &route : localExpected) { + ++expected[static_cast(route.srcRank)]; + } + return expected; + } + + int32_t running = 0; + const int64_t localExpertNum = OutputLocalExpertNum(rank, rankSize, config); + for (int64_t localExpert = 0; localExpert < localExpertNum; ++localExpert) { + for (int srcRank = 0; srcRank < rankSize; ++srcRank) { + for (const ExpectedRoute &route : localExpected) { + if (route.srcRank == srcRank && + LocalExpertForExpert(route.expertId, rankSize, config) == localExpert) { + ++running; + } + } + expected[static_cast(localExpert * rankSize + srcRank)] = running; + } + } + return expected; +} + +struct StandaloneCombineInputs { + std::vector expertOut; + std::vector assist; + std::vector recvCounts; +}; + +bool BuildStandaloneCombineInputs(int rank, int rankSize, const DemoConfig &config, + const std::vector &activeMask, bool useMemory, TileXR::TileXRDataType dtype, + bool useCombineMxfp8, std::size_t expandedRows, StandaloneCombineInputs *inputs) +{ + if (inputs == nullptr) { + return false; + } + const std::vector routes = useMemory ? + BuildExpectedMemoryRoutes(rank, rankSize, config, activeMask) : + BuildExpectedTpRoutes(rank, rankSize, config, activeMask); + if (routes.size() > expandedRows) { + return false; + } + + const uint16_t one = dtype == TileXR::TILEXR_DATA_TYPE_BFP16 ? kBf16One : kFp16One; + inputs->expertOut.assign(expandedRows * static_cast(config.h), one); + inputs->assist.assign(expandedRows * kAssistInts, 0); + inputs->recvCounts = BuildExpectedRecvCounts(rank, rankSize, config, activeMask, useMemory); + for (std::size_t row = 0; row < routes.size(); ++row) { + const ExpectedRoute &route = routes[row]; + if (useCombineMxfp8) { + for (int64_t h = 0; h < config.h; ++h) { + inputs->expertOut[row * static_cast(config.h) + static_cast(h)] = + InputValue(dtype, route.srcRank, route.tokenId, h); + } + } + const std::size_t offset = row * kAssistInts; + inputs->assist[offset] = route.srcRank; + inputs->assist[offset + 1] = route.tokenId; + inputs->assist[offset + 2] = route.topKId; + inputs->assist[offset + 3] = route.expertId; + } + return true; +} + +TileXREpDemo::Mxfp8Tensor BuildExpectedMxfp8Dispatch(const std::vector &routes, + const DemoConfig &config, TileXR::TileXRDataType dtype, std::size_t expandedRows, + TileXREpDemo::Mxfp8Format format) +{ + std::vector routedInput(expandedRows * static_cast(config.h), 0.0f); + for (std::size_t row = 0; row < routes.size(); ++row) { + for (int64_t h = 0; h < config.h; ++h) { + routedInput[row * static_cast(config.h) + static_cast(h)] = + DataBitsToFloat(InputValue(dtype, routes[row].srcRank, routes[row].tokenId, h), dtype); + } + } + return TileXREpDemo::QuantizeMxfp8( + routedInput, expandedRows, static_cast(config.h), format); +} + bool ValidateOutputs(int rank, int rankSize, const DemoConfig &config, const std::vector &expandX, const std::vector &expertTokenNums, const std::vector &recvCounts, - const std::vector &assist, const std::vector &dynamicScalesOut, + const std::vector &assist, const std::vector &dynamicScalesOut, const std::vector &activeMask, int expertTokenNumsType, bool useStaticQuant, - bool usePerTokenDynamicQuant, float staticQuantScale) + bool usePerTokenDynamicQuant, bool useMxfp8, float staticQuantScale, bool useMemoryDispatch, + TileXR::TileXRDataType dtype, const TileXREpDemo::Mxfp8Tensor &expectedMxfp8) { - const int64_t localExpertNum = LocalExpertNum(rankSize, config); - const std::vector expected = BuildExpectedTpRoutes(rank, rankSize, config, activeMask); - const std::vector localExpected = BuildExpectedRoutes(rank, rankSize, config, activeMask); - std::vector expectedRecv(rankSize, 0); + const int64_t localExpertNum = OutputLocalExpertNum(rank, rankSize, config); + const std::vector expected = useMemoryDispatch ? + BuildExpectedMemoryRoutes(rank, rankSize, config, activeMask) : + BuildExpectedTpRoutes(rank, rankSize, config, activeMask); + const std::vector expectedRecv = + BuildExpectedRecvCounts(rank, rankSize, config, activeMask, useMemoryDispatch); std::vector expectedExpertCounts(localExpertNum, 0); - for (const ExpectedRoute &route : localExpected) { - ++expectedRecv[route.srcRank]; - } for (const ExpectedRoute &route : expected) { const int64_t localExpert = LocalExpertForExpert(route.expertId, rankSize, config); if (localExpert >= 0 && localExpert < localExpertNum) { @@ -564,10 +858,18 @@ bool ValidateOutputs(int rank, int rankSize, const DemoConfig &config, const std } } - for (int srcRank = 0; srcRank < rankSize; ++srcRank) { - if (recvCounts[srcRank] != expectedRecv[srcRank]) { - std::cerr << "rank " << rank << " recvCounts[" << srcRank << "] expected " - << expectedRecv[srcRank] << " got " << recvCounts[srcRank] << std::endl; + for (std::size_t index = 0; index < expectedRecv.size(); ++index) { + if (recvCounts[index] != expectedRecv[index]) { + std::cerr << "rank " << rank << (useMemoryDispatch ? " sendCounts[" : " recvCounts[") + << index << "] expected " << expectedRecv[index] << " got " << recvCounts[index] + << std::endl; + if (useMemoryDispatch) { + std::cerr << "rank " << rank << " sendCounts expected/got:"; + for (std::size_t countIndex = 0; countIndex < expectedRecv.size(); ++countIndex) { + std::cerr << " " << expectedRecv[countIndex] << "/" << recvCounts[countIndex]; + } + std::cerr << std::endl; + } return false; } } @@ -592,8 +894,9 @@ bool ValidateOutputs(int rank, int rankSize, const DemoConfig &config, const std } if (usePerTokenDynamicQuant) { - const float expectedScale = DynamicScaleForXValue(route.srcRank, route.tokenId); - const float actualScale = dynamicScalesOut[row]; + const float expectedScale = DynamicScaleForXValue(route.srcRank, route.tokenId, config.h); + float actualScale = 0.0f; + std::memcpy(&actualScale, dynamicScalesOut.data() + row * sizeof(float), sizeof(float)); if (std::fabs(actualScale - expectedScale) > 1.0e-5f) { std::cerr << "rank " << rank << " dynamicScalesOut[" << row << "] expected " << expectedScale << " got " << actualScale << std::endl; @@ -601,15 +904,26 @@ bool ValidateOutputs(int rank, int rankSize, const DemoConfig &config, const std } } - for (int64_t h = 0; h < kH; ++h) { + for (int64_t h = 0; h < config.h; ++h) { + if (useMxfp8) { + const std::size_t offset = row * static_cast(config.h) + + static_cast(h); + if (expandX[offset] != expectedMxfp8.elements[offset]) { + std::cerr << "rank " << rank << " MXFP8 expandX[" << row << "][" << h + << "] expected 0x" << std::hex << static_cast(expectedMxfp8.elements[offset]) + << " got 0x" << static_cast(expandX[offset]) << std::dec << std::endl; + return false; + } + continue; + } const bool useInt8Output = useStaticQuant || usePerTokenDynamicQuant; - const std::size_t byteOffset = row * kH * (useInt8Output ? sizeof(int8_t) : sizeof(uint16_t)) + + const std::size_t byteOffset = row * config.h * (useInt8Output ? sizeof(int8_t) : sizeof(uint16_t)) + h * (useInt8Output ? sizeof(int8_t) : sizeof(uint16_t)); const int expectedValue = useStaticQuant ? static_cast(QuantizedXValue(route.srcRank, route.tokenId, h, staticQuantScale)) : (usePerTokenDynamicQuant ? - static_cast(DynamicQuantizedXValue(route.srcRank, route.tokenId, h)) : - static_cast(XValue(route.srcRank, route.tokenId, h))); + static_cast(DynamicQuantizedXValue(route.srcRank, route.tokenId, h, config.h)) : + static_cast(InputValue(dtype, route.srcRank, route.tokenId, h))); const int actualValue = useInt8Output ? static_cast(*reinterpret_cast(&expandX[byteOffset])) : static_cast(*reinterpret_cast(&expandX[byteOffset])); @@ -619,19 +933,91 @@ bool ValidateOutputs(int rank, int rankSize, const DemoConfig &config, const std return false; } } + if (useMxfp8) { + for (std::size_t scale = 0; scale < expectedMxfp8.scaleCountPerRow; ++scale) { + const std::size_t offset = row * expectedMxfp8.scaleCountPerRow + scale; + if (dynamicScalesOut[offset] != expectedMxfp8.scales[offset]) { + std::cerr << "rank " << rank << " MXFP8 dynamicScalesOut[" << row << "][" << scale + << "] expected 0x" << std::hex << static_cast(expectedMxfp8.scales[offset]) + << " got 0x" << static_cast(dynamicScalesOut[offset]) << std::dec << std::endl; + return false; + } + } + } } return true; } -bool ValidateCombineOutputs(int rank, const std::vector &yOut) +int64_t ExpectedCombineRouteCount( + const DemoConfig &config, const std::vector &activeMask, int64_t token) { - for (int64_t token = 0; token < kBs; ++token) { - for (int64_t h = 0; h < kH; ++h) { - const uint16_t actualValue = yOut[token * kH + h]; - if (actualValue != kFp16Two) { - std::cerr << "rank " << rank << " yOut[" << token << "][" << h << "] expected 0x" - << std::hex << kFp16Two << " got 0x" << actualValue << std::dec << std::endl; + if (!IsTokenActive(activeMask, config, token)) { + return 0; + } + int64_t count = config.sharedExpertNum; + const std::vector expertIds = ExpertIds(config); + for (int64_t topKId = 0; topKId < config.topK; ++topKId) { + const int32_t expertId = expertIds[static_cast(token * config.topK + topKId)]; + if (IsRouteActive(activeMask, config, token, topKId) && expertId >= 0 && expertId < config.moeExpertNum) { + ++count; + } + } + return count; +} + +bool ValidateCombineOutputs(int rank, const DemoConfig &config, const std::vector &yOut, + const std::vector &activeMask, TileXR::TileXRDataType dtype, bool usesDispatchOutput, + int commQuantMode, const std::vector &expertScales) +{ + const bool useCombineMxfp8 = commQuantMode == 3 || commQuantMode == 4; + if (useCombineMxfp8 && expertScales.size() != static_cast(config.bs * config.topK)) { + std::cerr << "rank " << rank << " invalid combine expertScales size" << std::endl; + return false; + } + const TileXREpDemo::Mxfp8Format format = commQuantMode == 3 ? + TileXREpDemo::Mxfp8Format::E5M2 : TileXREpDemo::Mxfp8Format::E4M3; + const std::vector expertIds = ExpertIds(config); + for (int64_t token = 0; token < config.bs; ++token) { + const int64_t routeCount = ExpectedCombineRouteCount(config, activeMask, token); + std::vector roundTrip; + if (useCombineMxfp8) { + std::vector row(static_cast(config.h)); + for (int64_t h = 0; h < config.h; ++h) { + row[static_cast(h)] = + DataBitsToFloat(InputValue(dtype, rank, token, h), dtype); + } + roundTrip = TileXREpDemo::RoundTripMxfp8( + row, 1, static_cast(config.h), format); + } + for (int64_t h = 0; h < config.h; ++h) { + float expectedValue = 0.0f; + if (useCombineMxfp8) { + if (IsTokenActive(activeMask, config, token)) { + for (int64_t topKId = 0; topKId < config.topK; ++topKId) { + const std::size_t route = static_cast(token * config.topK + topKId); + if (IsRouteActive(activeMask, config, token, topKId) && expertIds[route] >= 0 && + expertIds[route] < config.moeExpertNum) { + expectedValue += roundTrip[static_cast(h)] * expertScales[route]; + } + } + for (int64_t sharedExpertId = 0; sharedExpertId < config.sharedExpertNum; ++sharedExpertId) { + expectedValue += roundTrip[static_cast(h)]; + } + } + } else { + const float inputValue = usesDispatchOutput ? + DataBitsToFloat(InputValue(dtype, rank, token, h), dtype) : 1.0f; + expectedValue = static_cast(routeCount) * inputValue; + } + const uint16_t actualBits = yOut[token * config.h + h]; + const float actualValue = DataBitsToFloat(actualBits, dtype); + const float tolerance = (usesDispatchOutput || useCombineMxfp8) ? + std::max(1.0e-3f, std::fabs(expectedValue) * 0.01f) : 0.0f; + if (std::fabs(actualValue - expectedValue) > tolerance) { + std::cerr << "rank " << rank << " yOut[" << token << "][" << h << "] expected " + << expectedValue << " got " << actualValue << " (bits 0x" << std::hex + << actualBits << std::dec << ")" << std::endl; return false; } } @@ -710,47 +1096,160 @@ int main(int argc, char **argv) const int rank = argc > 2 ? std::atoi(argv[2]) : GetEnvInt("RANK", 0); const int npuCount = argc > 3 ? std::atoi(argv[3]) : GetEnvInt("TILEXR_DEMO_NPUS", rankSize); const int firstNpu = argc > 4 ? std::atoi(argv[4]) : GetEnvInt("TILEXR_DEMO_FIRST_NPU", 0); - const bool dispatchOnly = argc > 5 ? std::atoi(argv[5]) != 0 : GetEnvInt("TILEXR_DEMO_DISPATCH_ONLY", 0) != 0; - const bool useActiveMask = EnvEnabled("TILEXR_EP_DEMO_ACTIVE_MASK"); + const int loopCount = GetEnvInt("TILEXR_EP_DEMO_LOOP", 100); + const char *backendText = std::getenv("TILEXR_EP_DEMO_IMPL"); + bool validBackend = true; + DemoBackend backend = DemoBackend::UDMA; + if (backendText != nullptr && backendText[0] != '\0') { + if (std::strcmp(backendText, "memory") == 0) { + backend = DemoBackend::MEMORY; + } else if (std::strcmp(backendText, "udma") != 0) { + validBackend = false; + } + } + const bool useMemory = backend == DemoBackend::MEMORY; + const char *runModeText = std::getenv("TILEXR_EP_DEMO_RUN_MODE"); + bool validRunMode = true; + DemoRunMode runMode = DemoRunMode::DISPATCH_COMBINE; + if (runModeText != nullptr && runModeText[0] != '\0') { + if (std::strcmp(runModeText, "dispatch") == 0) { + runMode = DemoRunMode::DISPATCH; + } else if (std::strcmp(runModeText, "combine") == 0) { + runMode = DemoRunMode::COMBINE; + } else if (std::strcmp(runModeText, "dispatch_combine") != 0) { + validRunMode = false; + } + } + const bool runDispatch = runMode != DemoRunMode::COMBINE; + const bool runCombine = runMode != DemoRunMode::DISPATCH; + const char *activeMaskTypeText = std::getenv("TILEXR_EP_DEMO_ACTIVE_MASK_TYPE"); + bool validActiveMaskType = true; + int64_t activeMaskType = TileXREp::TILEXR_EP_ACTIVE_MASK_NONE; + if (activeMaskTypeText == nullptr || activeMaskTypeText[0] == '\0') { + activeMaskType = EnvEnabled("TILEXR_EP_DEMO_ACTIVE_MASK") ? + TileXREp::TILEXR_EP_ACTIVE_MASK_TOKEN : TileXREp::TILEXR_EP_ACTIVE_MASK_NONE; + } else if (std::strcmp(activeMaskTypeText, "token") == 0) { + activeMaskType = TileXREp::TILEXR_EP_ACTIVE_MASK_TOKEN; + } else if (std::strcmp(activeMaskTypeText, "expert") == 0) { + activeMaskType = TileXREp::TILEXR_EP_ACTIVE_MASK_EXPERT; + } else if (std::strcmp(activeMaskTypeText, "none") != 0) { + validActiveMaskType = false; + } + const bool useActiveMask = activeMaskType != TileXREp::TILEXR_EP_ACTIVE_MASK_NONE; + const char *dtypeText = std::getenv("TILEXR_EP_DEMO_DTYPE"); + bool validDtype = true; + TileXR::TileXRDataType dtype = TileXR::TILEXR_DATA_TYPE_FP16; + if (dtypeText != nullptr && dtypeText[0] != '\0') { + if (std::strcmp(dtypeText, "bf16") == 0) { + dtype = TileXR::TILEXR_DATA_TYPE_BFP16; + } else if (std::strcmp(dtypeText, "fp16") != 0) { + validDtype = false; + } + } const bool requestedTpRecvCounts = EnvEnabled("TILEXR_EP_DEMO_TP_RECV_COUNTS"); const int expertTokenNumsType = GetEnvInt("TILEXR_EP_DEMO_EXPERT_TOKEN_NUMS_TYPE", 1); const int quantMode = GetEnvInt("TILEXR_EP_DEMO_QUANT_MODE", 0); const bool useStaticQuant = quantMode == 1; const bool usePerTokenDynamicQuant = quantMode == 2; + const bool useMxfp8 = quantMode == 4; + const int commQuantMode = GetEnvInt("TILEXR_EP_DEMO_COMM_QUANT_MODE", 0); + const bool useCombineMxfp8 = commQuantMode == 3 || commQuantMode == 4; + const char *mxfp8FormatText = std::getenv("TILEXR_EP_DEMO_MXFP8_FORMAT"); + bool validMxfp8Format = true; + TileXREpDemo::Mxfp8Format mxfp8Format = TileXREpDemo::Mxfp8Format::E4M3; + if (mxfp8FormatText != nullptr && mxfp8FormatText[0] != '\0') { + if (std::strcmp(mxfp8FormatText, "e5m2") == 0 || + std::strcmp(mxfp8FormatText, "fp8_e5m2") == 0) { + mxfp8Format = TileXREpDemo::Mxfp8Format::E5M2; + } else if (std::strcmp(mxfp8FormatText, "e4m3") != 0 && + std::strcmp(mxfp8FormatText, "fp8_e4m3fn") != 0) { + validMxfp8Format = false; + } + } + const TileXR::TileXRDataType expandXOutDtype = useMxfp8 ? + (mxfp8Format == TileXREpDemo::Mxfp8Format::E4M3 ? TileXR::TILEXR_DATA_TYPE_FP8E4M3 : + TileXR::TILEXR_DATA_TYPE_FP8E5M2) : dtype; const float staticQuantScale = static_cast(GetEnvInt("TILEXR_EP_DEMO_STATIC_QUANT_SCALE", 1)); DemoConfig config {}; + config.bs = GetEnvInt("TILEXR_EP_DEMO_BS", static_cast(config.bs)); + config.h = GetEnvInt("TILEXR_EP_DEMO_H", static_cast(config.h)); + config.topK = GetEnvInt("TILEXR_EP_DEMO_TOPK", static_cast(config.topK)); config.moeExpertNum = GetEnvInt("TILEXR_EP_DEMO_MOE_EXPERT_NUM", static_cast(config.moeExpertNum)); config.sharedExpertNum = GetEnvInt("TILEXR_EP_DEMO_SHARED_EXPERT_NUM", 0); config.sharedExpertRankNum = GetEnvInt("TILEXR_EP_DEMO_SHARED_EXPERT_RANK_NUM", 0); config.tpWorldSize = GetEnvInt("TILEXR_EP_DEMO_TP_WORLD_SIZE", 0); config.tpRankId = GetEnvInt("TILEXR_EP_DEMO_TP_RANK_ID", config.effectiveTpWorldSize() > 1 ? rank % config.effectiveTpWorldSize() : 0); + const char *expertIdsText = std::getenv("TILEXR_EP_DEMO_EXPERT_IDS"); + const bool hasConfiguredExpertIds = expertIdsText != nullptr && expertIdsText[0] != '\0'; + const bool validExpertIds = !hasConfiguredExpertIds || ParseInt32List(expertIdsText, &config.expertIds); + const char *expertModeText = std::getenv("TILEXR_EP_DEMO_EXPERT_MODE"); + bool validExpertMode = true; + ExpertListMode expertMode = hasConfiguredExpertIds ? ExpertListMode::EXPLICIT : ExpertListMode::UNIFORM; + if (expertModeText != nullptr && expertModeText[0] != '\0') { + if (std::strcmp(expertModeText, "random") == 0) { + expertMode = ExpertListMode::RANDOM; + } else if (std::strcmp(expertModeText, "explicit") == 0) { + expertMode = ExpertListMode::EXPLICIT; + } else if (std::strcmp(expertModeText, "uniform") != 0) { + validExpertMode = false; + } + } + uint32_t expertSeed = 1; + const bool validExpertSeed = GetEnvUint32("TILEXR_EP_DEMO_EXPERT_SEED", 1, &expertSeed); + const bool unambiguousExpertConfig = + (expertMode == ExpertListMode::EXPLICIT) == hasConfiguredExpertIds; + const bool validExpertConfiguration = validExpertIds && validExpertMode && validExpertSeed && + unambiguousExpertConfig && BuildExpertIds(&config, expertMode, expertSeed); const bool useTpRecvCounts = requestedTpRecvCounts || config.effectiveTpWorldSize() != 1; const int64_t expertRankSize = static_cast(rankSize) / config.effectiveTpWorldSize(); const int64_t moeRankNum = expertRankSize - config.sharedExpertRankNum; - if (rankSize <= 0 || rank < 0 || rank >= rankSize || config.effectiveTpWorldSize() <= 0 || - rankSize % config.effectiveTpWorldSize() != 0 || moeRankNum <= 0 || + if (rankSize <= 0 || rank < 0 || rank >= rankSize || loopCount <= 0 || + config.effectiveTpWorldSize() <= 0 || + rankSize % config.effectiveTpWorldSize() != 0 || config.bs <= 0 || config.h <= 0 || config.topK <= 0 || + moeRankNum <= 0 || config.moeExpertNum <= 0 || config.moeExpertNum % moeRankNum != 0 || config.sharedExpertNum < 0 || config.sharedExpertRankNum < 0 || - config.sharedExpertNum != config.sharedExpertRankNum || + ((config.sharedExpertNum == 0) != (config.sharedExpertRankNum == 0)) || + (config.sharedExpertNum > 0 && config.sharedExpertRankNum % config.sharedExpertNum != 0) || (config.effectiveTpWorldSize() > 1 && config.tpRankId != rank % config.effectiveTpWorldSize()) || + !validBackend || !validRunMode || !validActiveMaskType || !validDtype || !validMxfp8Format || + !validExpertConfiguration || (expertTokenNumsType != 0 && expertTokenNumsType != 1) || - (quantMode != 0 && quantMode != 1 && quantMode != 2) || - ((useStaticQuant || usePerTokenDynamicQuant) && !dispatchOnly)) { + (quantMode != 0 && quantMode != 1 && quantMode != 2 && quantMode != 4) || + (commQuantMode != 0 && commQuantMode != 3 && commQuantMode != 4) || + activeMaskType == TileXREp::TILEXR_EP_ACTIVE_MASK_EXPERT || + useStaticQuant || usePerTokenDynamicQuant || (useMxfp8 && (runCombine || !useMemory)) || + (useCombineMxfp8 && (!runCombine || !useMemory)) || + (runCombine && config.effectiveTpWorldSize() != 1)) { std::cerr << "This demo expects a valid rank and moeExpertNum divisible by MoE rank num, got moeExpertNum=" << config.moeExpertNum << " rankSize=" << rankSize + << " bs=" << config.bs + << " h=" << config.h + << " topK=" << config.topK << " sharedExpertNum=" << config.sharedExpertNum << " sharedExpertRankNum=" << config.sharedExpertRankNum << " tpWorldSize=" << config.tpWorldSize << " tpRankId=" << config.tpRankId << ", and expertTokenNumsType 0 or 1, got rankSize=" << rankSize << " rank=" << rank - << " expertTokenNumsType=" << expertTokenNumsType - << " quantMode=" << quantMode << std::endl; + << " expertTokenNumsType=" << expertTokenNumsType + << " quantMode=" << quantMode << " activeMaskType=" << activeMaskType + << " commQuantMode=" << commQuantMode + << " backend=" << (useMemory ? "memory" : "udma") + << " runMode=" << (runMode == DemoRunMode::DISPATCH ? "dispatch" : + (runMode == DemoRunMode::COMBINE ? "combine" : "dispatch_combine")) + << " expertMode=" << (expertMode == ExpertListMode::UNIFORM ? "uniform" : + (expertMode == ExpertListMode::RANDOM ? "random" : "explicit")) + << " expertSeed=" << expertSeed + << " mxfp8Format=" << (mxfp8Format == TileXREpDemo::Mxfp8Format::E4M3 ? "e4m3" : "e5m2") + << " expertIds=" << config.expertIds.size() + << " dtype=" << static_cast(dtype) + << " loopCount=" << loopCount << std::endl; return 2; } - const int64_t localExpertNum = LocalExpertNum(rankSize, config); + const int64_t localExpertNum = OutputLocalExpertNum(rank, rankSize, config); const int deviceId = GetDeviceIdFromEnv(rank, npuCount, firstNpu); bool aclReady = false; bool deviceSet = false; @@ -777,16 +1276,18 @@ int main(int argc, char **argv) return 1; } - std::vector hostX(kXElements); - for (int64_t token = 0; token < kBs; ++token) { - for (int64_t h = 0; h < kH; ++h) { - hostX[token * kH + h] = XValue(rank, token, h); + std::vector hostX(static_cast(config.bs * config.h)); + for (int64_t token = 0; token < config.bs; ++token) { + for (int64_t h = 0; h < config.h; ++h) { + hostX[token * config.h + h] = InputValue(dtype, rank, token, h); } } const std::vector hostExpertIds = ExpertIds(config); - const std::vector hostActiveMask = ActiveMask(useActiveMask); - const std::size_t expectedRouteCount = - BuildExpectedTpRoutes(rank, rankSize, config, hostActiveMask).size(); + const std::vector hostActiveMask = ActiveMask(activeMaskType, config); + const std::vector expectedRoutes = useMemory ? + BuildExpectedMemoryRoutes(rank, rankSize, config, hostActiveMask) : + BuildExpectedTpRoutes(rank, rankSize, config, hostActiveMask); + const std::size_t expectedRouteCount = expectedRoutes.size(); void *xDev = nullptr; void *expertIdsDev = nullptr; @@ -799,6 +1300,7 @@ int main(int argc, char **argv) void *tpRecvCountsDev = nullptr; void *assistDev = nullptr; void *expertOutDev = nullptr; + void *expertScalesDev = nullptr; void *yOutDev = nullptr; void *workspaceDev = nullptr; void *rawWorkspaceDev = nullptr; @@ -809,26 +1311,47 @@ int main(int argc, char **argv) const std::size_t expertIdsBytes = hostExpertIds.size() * sizeof(int32_t); const std::size_t xActiveMaskBytes = hostActiveMask.size() * sizeof(uint8_t); const std::size_t expandedElements = std::max(static_cast(config.expandedElements()), - expectedRouteCount * static_cast(kH)); + expectedRouteCount * static_cast(config.h)); const std::size_t maxRoutesPerRank = static_cast(config.maxRoutesPerRank()); const std::size_t expandElementBytes = - (useStaticQuant || usePerTokenDynamicQuant) ? sizeof(int8_t) : sizeof(uint16_t); + (useStaticQuant || usePerTokenDynamicQuant || useMxfp8) ? sizeof(int8_t) : sizeof(uint16_t); const std::size_t expandXBytes = expandedElements * expandElementBytes; - const std::size_t expandedRows = expandedElements / kH; - const std::size_t dynamicScalesBytes = expandedRows * sizeof(float); + const std::size_t expertOutBytes = expandedElements * sizeof(uint16_t); + const std::size_t expandedRows = expandedElements / config.h; + const std::size_t payloadScaleBytesPerRow = usePerTokenDynamicQuant ? sizeof(float) : + (useMxfp8 ? TileXREpDemo::Mxfp8ScaleCountPerRow(static_cast(config.h)) : 0U); + const std::size_t dynamicScalesBytes = expandedRows * payloadScaleBytesPerRow; const std::size_t expertTokenNumsBytes = localExpertNum * sizeof(int64_t); - const std::size_t recvCountsBytes = rankSize * sizeof(int32_t); - const std::size_t tpRecvCountsBytes = recvCountsBytes; - const std::size_t assistBytes = (expandedElements / kH) * kAssistInts * sizeof(int32_t); - const std::size_t yOutBytes = kXElements * sizeof(uint16_t); - const std::size_t payloadRowBytes = kH * expandElementBytes; - const std::size_t dispatchWindowBytes = EpWindowBytes(rankSize, config, payloadRowBytes, usePerTokenDynamicQuant); + const std::size_t recvCountsElements = useMemory ? + MemorySendCountsCount(rank, rankSize, config) : static_cast(rankSize); + const std::size_t recvCountsBytes = recvCountsElements * sizeof(int32_t); + const std::size_t tpRecvCountsBytes = rankSize * sizeof(int32_t); + const std::size_t assistBytes = (expandedElements / config.h) * kAssistInts * sizeof(int32_t); + const std::size_t yOutBytes = static_cast(config.bs * config.h) * sizeof(uint16_t); + const std::vector hostExpertScales = useCombineMxfp8 ? + BuildExpertScales(config) : std::vector {}; + const std::size_t expertScalesBytes = hostExpertScales.size() * sizeof(float); + const std::size_t payloadRowBytes = config.h * expandElementBytes; + const std::size_t dispatchWindowBytes = + EpWindowBytes(rankSize, config, payloadRowBytes, payloadScaleBytesPerRow); const std::size_t dispatchPayloadBytes = AlignSize(dispatchWindowBytes, 32) * static_cast(config.effectiveTpWorldSize() + 2); const std::size_t workspacePayloadBytes = std::max(dispatchPayloadBytes, - EpRequiredWorkspaceBytes(rankSize, config, payloadRowBytes, usePerTokenDynamicQuant)); + EpRequiredWorkspaceBytes(rankSize, config, payloadRowBytes, payloadScaleBytesPerRow)); const std::size_t workspaceBytes = ((workspacePayloadBytes + kUdmaRegistrationAlignment - 1) / kUdmaRegistrationAlignment) * kUdmaRegistrationAlignment; + const TileXREpDemo::Mxfp8Tensor expectedMxfp8 = useMxfp8 ? + BuildExpectedMxfp8Dispatch(expectedRoutes, config, dtype, expandedRows, mxfp8Format) : + TileXREpDemo::Mxfp8Tensor {}; + + StandaloneCombineInputs standaloneCombineInputs; + if (runMode == DemoRunMode::COMBINE && + !BuildStandaloneCombineInputs(rank, rankSize, config, hostActiveMask, useMemory, dtype, + useCombineMxfp8, expandedRows, &standaloneCombineInputs)) { + std::cerr << "rank " << rank << " failed to construct standalone combine inputs" << std::endl; + Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); + return 1; + } if (!CheckAcl(aclrtMalloc(&xDev, xBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc x") || !CheckAcl(aclrtMalloc(&expertIdsDev, expertIdsBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc expertIds") || @@ -837,7 +1360,7 @@ int main(int argc, char **argv) (useActiveMask && !CheckAcl(aclrtMalloc(&xActiveMaskDev, xActiveMaskBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc xActiveMask")) || !CheckAcl(aclrtMalloc(&expandXDev, expandXBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc expandX") || - (usePerTokenDynamicQuant && !CheckAcl(aclrtMalloc(&dynamicScalesDev, dynamicScalesBytes, + ((usePerTokenDynamicQuant || useMxfp8) && !CheckAcl(aclrtMalloc(&dynamicScalesDev, dynamicScalesBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc dynamicScales")) || !CheckAcl(aclrtMalloc(&expertTokenNumsDev, expertTokenNumsBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc expertTokenNums") || @@ -845,20 +1368,23 @@ int main(int argc, char **argv) (useTpRecvCounts && !CheckAcl(aclrtMalloc(&tpRecvCountsDev, tpRecvCountsBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc tpRecvCounts")) || !CheckAcl(aclrtMalloc(&assistDev, assistBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc assist") || - !CheckAcl(aclrtMalloc(&expertOutDev, expandXBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc expertOut") || + !CheckAcl(aclrtMalloc(&expertOutDev, expertOutBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc expertOut") || + (useCombineMxfp8 && !CheckAcl(aclrtMalloc(&expertScalesDev, expertScalesBytes, + ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc expertScales")) || !CheckAcl(aclrtMalloc(&yOutDev, yOutBytes, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc yOut") || !CheckAcl(aclrtMalloc(&rawWorkspaceDev, workspaceBytes + kUdmaRegistrationAlignment - 1, ACL_MEM_MALLOC_HUGE_FIRST), "aclrtMalloc workspace")) { workspaceDev = rawWorkspaceDev; buffers = {xDev, expertIdsDev, scalesDev, xActiveMaskDev, expandXDev, dynamicScalesDev, - expertTokenNumsDev, recvCountsDev, tpRecvCountsDev, assistDev, expertOutDev, yOutDev, workspaceDev}; + expertTokenNumsDev, recvCountsDev, tpRecvCountsDev, assistDev, expertOutDev, expertScalesDev, + yOutDev, workspaceDev}; Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } workspaceDev = reinterpret_cast(AlignAddress(reinterpret_cast(rawWorkspaceDev), kUdmaRegistrationAlignment)); buffers = {xDev, expertIdsDev, scalesDev, xActiveMaskDev, expandXDev, dynamicScalesDev, expertTokenNumsDev, - recvCountsDev, tpRecvCountsDev, assistDev, expertOutDev, yOutDev, rawWorkspaceDev}; + recvCountsDev, tpRecvCountsDev, assistDev, expertOutDev, expertScalesDev, yOutDev, rawWorkspaceDev}; if (!CheckAcl(aclrtMemcpy(xDev, xBytes, hostX.data(), xBytes, ACL_MEMCPY_HOST_TO_DEVICE), "copy x") || !CheckAcl(aclrtMemcpy(expertIdsDev, expertIdsBytes, hostExpertIds.data(), expertIdsBytes, @@ -868,181 +1394,167 @@ int main(int argc, char **argv) (useActiveMask && !CheckAcl(aclrtMemcpy(xActiveMaskDev, xActiveMaskBytes, hostActiveMask.data(), xActiveMaskBytes, ACL_MEMCPY_HOST_TO_DEVICE), "copy xActiveMask")) || !CheckAcl(aclrtMemset(expandXDev, expandXBytes, 0, expandXBytes), "memset expandX") || - (usePerTokenDynamicQuant && !CheckAcl(aclrtMemset(dynamicScalesDev, dynamicScalesBytes, 0, + ((usePerTokenDynamicQuant || useMxfp8) && !CheckAcl(aclrtMemset(dynamicScalesDev, dynamicScalesBytes, 0, dynamicScalesBytes), "memset dynamicScales")) || !CheckAcl(aclrtMemset(expertTokenNumsDev, expertTokenNumsBytes, 0, expertTokenNumsBytes), "memset expertTokenNums") || - !CheckAcl(aclrtMemset(recvCountsDev, recvCountsBytes, 0, recvCountsBytes), "memset recvCounts") || + (runMode == DemoRunMode::COMBINE ? + !CheckAcl(aclrtMemcpy(recvCountsDev, recvCountsBytes, standaloneCombineInputs.recvCounts.data(), + recvCountsBytes, ACL_MEMCPY_HOST_TO_DEVICE), "copy standalone recvCounts") : + !CheckAcl(aclrtMemset(recvCountsDev, recvCountsBytes, 0, recvCountsBytes), "memset recvCounts")) || (useTpRecvCounts && !CheckAcl(aclrtMemset(tpRecvCountsDev, tpRecvCountsBytes, 0, tpRecvCountsBytes), "memset tpRecvCounts")) || - !CheckAcl(aclrtMemset(assistDev, assistBytes, 0, assistBytes), "memset assist") || + (runMode == DemoRunMode::COMBINE ? + !CheckAcl(aclrtMemcpy(assistDev, assistBytes, standaloneCombineInputs.assist.data(), assistBytes, + ACL_MEMCPY_HOST_TO_DEVICE), "copy standalone assist") : + !CheckAcl(aclrtMemset(assistDev, assistBytes, 0, assistBytes), "memset assist")) || + (useCombineMxfp8 && !CheckAcl(aclrtMemcpy(expertScalesDev, expertScalesBytes, + hostExpertScales.data(), expertScalesBytes, ACL_MEMCPY_HOST_TO_DEVICE), "copy expertScales")) || !CheckAcl(aclrtMemset(yOutDev, yOutBytes, 0, yOutBytes), "memset yOut") || !CheckAcl(aclrtMemset(workspaceDev, workspaceBytes, 0, workspaceBytes), "memset workspace")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - TileXR::CommArgs *commArgsHost = nullptr; - if (!CheckTileXR(TileXRGetCommArgsHost(comm, commArgsHost), "TileXRGetCommArgsHost")) { - Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); - return 1; - } - const bool crossNode = commArgsHost != nullptr && commArgsHost->localRankSize > 0 && - commArgsHost->localRankSize < commArgsHost->rankSize; - if (crossNode && !CheckTileXR(TileXRUDMARegister(comm, static_cast(workspaceDev), workspaceBytes, + if (!useMemory && !CheckTileXR(TileXRUDMARegister(comm, static_cast(workspaceDev), workspaceBytes, &workspaceHandle), "TileXRUDMARegister workspace")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - if (crossNode) { + if (!useMemory) { g_workspaceHandle = workspaceHandle; g_workspaceRegistered = true; } - const std::vector hostExpertOut(expandedElements, kFp16One); - if (!CheckAcl(aclrtMemcpy(expertOutDev, expandXBytes, hostExpertOut.data(), expandXBytes, + const uint16_t expertOutOne = dtype == TileXR::TILEXR_DATA_TYPE_BFP16 ? kBf16One : kFp16One; + const std::vector hostExpertOut = runMode == DemoRunMode::COMBINE ? + standaloneCombineInputs.expertOut : std::vector(expandedElements, expertOutOne); + if (!CheckAcl(aclrtMemcpy(expertOutDev, expertOutBytes, hostExpertOut.data(), expertOutBytes, ACL_MEMCPY_HOST_TO_DEVICE), "copy expertOut")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - const bool useSharedExperts = config.sharedExpertNum != 0 || config.sharedExpertRankNum != 0; - const bool useTp = config.effectiveTpWorldSize() != 1; - const bool useDispatchV2 = crossNode || useActiveMask || useTpRecvCounts || expertTokenNumsType != 1 || - useSharedExperts || useTp || useStaticQuant || usePerTokenDynamicQuant; - const int dispatchRet = useDispatchV2 ? - TileXRMoeEpDispatchV2(xDev, static_cast(expertIdsDev), scalesDev, - static_cast(xActiveMaskDev), nullptr, comm, kBs, kH, kTopK, config.moeExpertNum, + const auto DispatchOnce = [&]() -> int { + if (useMemory) { + return TileXRMoeEpDispatchMemoryV2(xDev, static_cast(expertIdsDev), scalesDev, + static_cast(xActiveMaskDev), activeMaskType, nullptr, comm, config.bs, config.h, config.topK, + config.moeExpertNum, + expertRankSize, ExpertRankForRank(rank, config), config.tpWorldSize, config.tpRankId, 0, + config.sharedExpertNum, config.sharedExpertRankNum, quantMode, config.bs * rankSize, + expertTokenNumsType, expandXDev, dynamicScalesDev, static_cast(assistDev), + static_cast(expertTokenNumsDev), static_cast(recvCountsDev), + static_cast(tpRecvCountsDev), nullptr, dtype, expandXOutDtype, stream); + } + return TileXRMoeEpDispatchV2(xDev, static_cast(expertIdsDev), scalesDev, + static_cast(xActiveMaskDev), nullptr, comm, config.bs, config.h, config.topK, config.moeExpertNum, expertRankSize, ExpertRankForRank(rank, config), config.tpWorldSize, config.tpRankId, 0, - config.sharedExpertNum, config.sharedExpertRankNum, quantMode, kBs * rankSize, expertTokenNumsType, - expandXDev, dynamicScalesDev, - static_cast(assistDev), static_cast(expertTokenNumsDev), - static_cast(recvCountsDev), static_cast(tpRecvCountsDev), nullptr, workspaceDev, - (useStaticQuant || usePerTokenDynamicQuant) ? TileXR::TILEXR_DATA_TYPE_INT8 : - TileXR::TILEXR_DATA_TYPE_FP16, - stream) : - TileXRMoeEpDispatch(xDev, static_cast(expertIdsDev), comm, kBs, kH, kTopK, config.moeExpertNum, - expandXDev, static_cast(expertTokenNumsDev), static_cast(recvCountsDev), - static_cast(assistDev), TileXR::TILEXR_DATA_TYPE_FP16, stream); - if (!CheckTileXR(dispatchRet, "TileXRMoeEpDispatch") || - !CheckAcl(aclrtSynchronizeStream(stream), "aclrtSynchronizeStream")) { - Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); - return 1; - } - if (!DemoBarrierAll(rank, rankSize, "dispatch synchronized")) { + config.sharedExpertNum, config.sharedExpertRankNum, quantMode, config.bs * rankSize, + expertTokenNumsType, expandXDev, dynamicScalesDev, static_cast(assistDev), + static_cast(expertTokenNumsDev), static_cast(recvCountsDev), + static_cast(tpRecvCountsDev), nullptr, workspaceDev, dtype, stream); + }; + + const auto CombineOnce = [&]() -> int { + void *combineExpertOutDev = runMode == DemoRunMode::DISPATCH_COMBINE ? expandXDev : expertOutDev; + if (useMemory) { + return TileXRMoeEpCombineMemoryV2(combineExpertOutDev, static_cast(assistDev), + static_cast(recvCountsDev), static_cast(expertScalesDev), + static_cast(xActiveMaskDev), activeMaskType, nullptr, comm, config.bs, config.h, + config.topK, config.moeExpertNum, rankSize, rank, + config.tpWorldSize, config.tpRankId, 0, config.sharedExpertNum, config.sharedExpertRankNum, + commQuantMode, config.bs * rankSize, yOutDev, dtype, stream); + } + return TileXRMoeEpCombineV2(combineExpertOutDev, static_cast(assistDev), + static_cast(recvCountsDev), comm, config.bs, config.h, config.topK, config.moeExpertNum, + yOutDev, workspaceDev, dtype, stream); + }; + + if (runMode == DemoRunMode::COMBINE && + !DemoBarrierAll(rank, rankSize, "standalone combine inputs ready")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - if (EnvEnabled("TILEXR_EP_DEMO_DUMP_WINDOW") && commArgsHost != nullptr) { - const std::size_t rowBytes = kH * expandElementBytes; - const std::size_t payloadBytes = AlignSize(maxRoutesPerRank * rowBytes + - (usePerTokenDynamicQuant ? maxRoutesPerRank * sizeof(float) : 0), 32); - const std::size_t assistWindowBytes = AlignSize(maxRoutesPerRank * kAssistInts * sizeof(int32_t), 32); - const std::size_t slotBytes = AlignSize(64 + payloadBytes + assistWindowBytes, 32); - const std::size_t windowBytes = AlignSize(64 + static_cast(rankSize) * slotBytes, 32); - if (crossNode) { - for (int slotRank = 0; slotRank < rankSize; ++slotRank) { - uint64_t slotHeader[8] = {}; - const GM_ADDR slotAddr = static_cast(workspaceDev) + windowBytes + 64 + - static_cast(slotRank) * slotBytes; - if (CheckAcl(aclrtMemcpy(slotHeader, sizeof(slotHeader), slotAddr, sizeof(slotHeader), - ACL_MEMCPY_DEVICE_TO_HOST), "dump slot header")) { - const uint32_t count = static_cast(slotHeader[0] & 0xffffffffULL); - const uint32_t slotSrc = static_cast((slotHeader[0] >> 32) & 0xffffffffULL); - std::cerr << "rank " << rank << " dump workspace slotRank " << slotRank - << " count " << count << " slotSrc " << slotSrc - << " payloadBytes " << slotHeader[1] << " assistBytes " << slotHeader[2] - << " magic " << slotHeader[3] - << std::endl; - } + if (runDispatch) { + for (int loop = 0; loop < loopCount; ++loop) { + const int dispatchRet = DispatchOnce(); + if (!CheckTileXR(dispatchRet, useMemory ? "memory dispatch" : "udma dispatch")) { + std::cerr << "rank " << rank << " dispatch loop " << loop << " failed" << std::endl; + Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); + return 1; } - } else { - for (int srcRank = 0; srcRank < rankSize; ++srcRank) { - for (int slotRank = 0; slotRank < rankSize; ++slotRank) { - uint64_t slotHeader[8] = {}; - const GM_ADDR slotAddr = commArgsHost->peerMems[srcRank] + TileXR::IPC_DATA_OFFSET + 64 + - static_cast(slotRank) * slotBytes; - if (CheckAcl(aclrtMemcpy(slotHeader, sizeof(slotHeader), slotAddr, sizeof(slotHeader), - ACL_MEMCPY_DEVICE_TO_HOST), "dump slot header")) { - const uint32_t count = static_cast(slotHeader[0] & 0xffffffffULL); - const uint32_t slotSrc = static_cast((slotHeader[0] >> 32) & 0xffffffffULL); - std::cerr << "rank " << rank << " dump sourceWindow " << srcRank << " slotRank " << slotRank - << " count " << count << " slotSrc " << slotSrc - << " payloadBytes " << slotHeader[1] << " assistBytes " << slotHeader[2] - << std::endl; - } - } + if (!CheckAcl(aclrtSynchronizeStream(stream), "aclrtSynchronizeStream dispatch") || + !DemoBarrierAll(rank, rankSize, "dispatch synchronized")) { + std::cerr << "rank " << rank << " dispatch completion loop " << loop << " failed" << std::endl; + Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); + return 1; + } + } + } + const int combineLoopCount = runMode == DemoRunMode::COMBINE ? loopCount : 1; + if (runCombine) { + for (int loop = 0; loop < combineLoopCount; ++loop) { + const int combineRet = CombineOnce(); + if (!CheckTileXR(combineRet, useMemory ? "memory combine" : "udma combine") || + !CheckAcl(aclrtSynchronizeStream(stream), "aclrtSynchronizeStream combine") || + !DemoBarrierAll(rank, rankSize, "combine synchronized")) { + std::cerr << "rank " << rank << " combine loop " << loop << " failed" << std::endl; + Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); + return 1; } } } std::vector hostExpandX(expandXBytes); std::vector hostExpertTokenNums(localExpertNum); - std::vector hostRecvCounts(rankSize); + std::vector hostRecvCounts(recvCountsElements); std::vector hostTpRecvCounts(rankSize); - std::vector hostAssist((expandedElements / kH) * kAssistInts); - std::vector hostDynamicScales(expandedRows); + std::vector hostAssist((expandedElements / config.h) * kAssistInts); + std::vector hostDynamicScales(dynamicScalesBytes); + std::vector hostYOut(static_cast(config.bs * config.h)); - if (!CheckAcl(aclrtMemcpy(hostExpandX.data(), expandXBytes, expandXDev, expandXBytes, + if (runDispatch && (!CheckAcl(aclrtMemcpy(hostExpandX.data(), expandXBytes, expandXDev, expandXBytes, ACL_MEMCPY_DEVICE_TO_HOST), "copy expandX") || !CheckAcl(aclrtMemcpy(hostExpertTokenNums.data(), expertTokenNumsBytes, expertTokenNumsDev, expertTokenNumsBytes, ACL_MEMCPY_DEVICE_TO_HOST), "copy expertTokenNums") || !CheckAcl(aclrtMemcpy(hostRecvCounts.data(), recvCountsBytes, recvCountsDev, recvCountsBytes, ACL_MEMCPY_DEVICE_TO_HOST), "copy recvCounts") || !CheckAcl(aclrtMemcpy(hostAssist.data(), assistBytes, assistDev, assistBytes, - ACL_MEMCPY_DEVICE_TO_HOST), "copy assist")) { + ACL_MEMCPY_DEVICE_TO_HOST), "copy assist"))) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - if (useTpRecvCounts && !CheckAcl(aclrtMemcpy(hostTpRecvCounts.data(), tpRecvCountsBytes, tpRecvCountsDev, - tpRecvCountsBytes, ACL_MEMCPY_DEVICE_TO_HOST), "copy tpRecvCounts")) { + if (runCombine && !CheckAcl(aclrtMemcpy(hostYOut.data(), yOutBytes, yOutDev, yOutBytes, + ACL_MEMCPY_DEVICE_TO_HOST), "copy yOut")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - if (usePerTokenDynamicQuant && !CheckAcl(aclrtMemcpy(hostDynamicScales.data(), dynamicScalesBytes, - dynamicScalesDev, dynamicScalesBytes, ACL_MEMCPY_DEVICE_TO_HOST), "copy dynamicScales")) { + if (runDispatch && useTpRecvCounts && !CheckAcl(aclrtMemcpy(hostTpRecvCounts.data(), tpRecvCountsBytes, tpRecvCountsDev, + tpRecvCountsBytes, ACL_MEMCPY_DEVICE_TO_HOST), "copy tpRecvCounts")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - - const bool dispatchOk = ValidateOutputs(rank, rankSize, config, hostExpandX, hostExpertTokenNums, hostRecvCounts, - hostAssist, hostDynamicScales, hostActiveMask, expertTokenNumsType, useStaticQuant, - usePerTokenDynamicQuant, staticQuantScale) && - (!useTpRecvCounts || ValidateTpRecvCounts(rank, rankSize, config, hostActiveMask, hostRecvCounts, - hostTpRecvCounts)); - std::cout << "rank " << rank << " validation " << (dispatchOk ? "PASS" : "FAIL") << std::endl; - if (!dispatchOk || dispatchOnly) { - Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); - return dispatchOk ? 0 : 1; - } - - const int combineRet = crossNode ? - TileXRMoeEpCombineV2(expertOutDev, static_cast(assistDev), - static_cast(recvCountsDev), comm, kBs, kH, kTopK, config.moeExpertNum, yOutDev, workspaceDev, - TileXR::TILEXR_DATA_TYPE_FP16, stream) : - TileXRMoeEpCombine(expertOutDev, static_cast(assistDev), - static_cast(recvCountsDev), comm, kBs, kH, kTopK, config.moeExpertNum, yOutDev, - TileXR::TILEXR_DATA_TYPE_FP16, stream); - if (!CheckTileXR(combineRet, "TileXRMoeEpCombine") || - !CheckAcl(aclrtSynchronizeStream(stream), "aclrtSynchronizeStream combine")) { + if (runDispatch && (usePerTokenDynamicQuant || useMxfp8) && + !CheckAcl(aclrtMemcpy(hostDynamicScales.data(), dynamicScalesBytes, + dynamicScalesDev, dynamicScalesBytes, ACL_MEMCPY_DEVICE_TO_HOST), "copy dynamicScales")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - if (!DemoBarrierAll(rank, rankSize, "combine synchronized")) { - Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); - return 1; + const bool dispatchOk = !runDispatch || + (ValidateOutputs(rank, rankSize, config, hostExpandX, hostExpertTokenNums, hostRecvCounts, + hostAssist, hostDynamicScales, hostActiveMask, expertTokenNumsType, useStaticQuant, + usePerTokenDynamicQuant, useMxfp8, staticQuantScale, useMemory, dtype, expectedMxfp8) && + (!useTpRecvCounts || ValidateTpRecvCounts(rank, rankSize, config, hostActiveMask, hostRecvCounts, + hostTpRecvCounts))); + if (runDispatch) { + std::cout << "rank " << rank << " dispatch validation " << (dispatchOk ? "PASS" : "FAIL") << std::endl; } - - std::vector hostYOut(kXElements); - if (!CheckAcl(aclrtMemcpy(hostYOut.data(), yOutBytes, yOutDev, yOutBytes, - ACL_MEMCPY_DEVICE_TO_HOST), "copy yOut")) { - Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); - return 1; + const bool combineOk = !runCombine || ValidateCombineOutputs(rank, config, hostYOut, hostActiveMask, dtype, + runMode == DemoRunMode::DISPATCH_COMBINE, commQuantMode, hostExpertScales); + if (runCombine) { + std::cout << "rank " << rank << " combine validation " << (combineOk ? "PASS" : "FAIL") << std::endl; } - - const bool combineOk = ValidateCombineOutputs(rank, hostYOut); - std::cout << "rank " << rank << " combine validation " << (combineOk ? "PASS" : "FAIL") << std::endl; Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); - return combineOk ? 0 : 1; + return dispatchOk && combineOk ? 0 : 1; } diff --git a/tests/ep/unit/test_tilexr_ep_host_validation.cpp b/tests/ep/unit/test_tilexr_ep_host_validation.cpp index d999bbcc..dc859a2c 100644 --- a/tests/ep/unit/test_tilexr_ep_host_validation.cpp +++ b/tests/ep/unit/test_tilexr_ep_host_validation.cpp @@ -5,6 +5,18 @@ #include "ep_dispatch_host.h" #include "tilexr_types.h" +int TileXRGetCommArgsDev(TileXRCommPtr, GM_ADDR &commArgsPtr) +{ + commArgsPtr = nullptr; + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; +} + +int TileXRGetCommArgsHost(TileXRCommPtr, TileXR::CommArgs *&commArgsPtr) +{ + commArgsPtr = nullptr; + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; +} + namespace { int g_failures = 0; @@ -48,6 +60,7 @@ TileXREp::EpDispatchParams ValidV2Params() static bool xActiveMask[4] = {true, false, true, true}; TileXREp::EpDispatchParams params = ValidParams(); params.xActiveMask = xActiveMask; + params.activeMaskType = TileXREp::TILEXR_EP_ACTIVE_MASK_TOKEN; params.epWorldSize = 2; params.epRankId = 0; params.tpWorldSize = 1; @@ -85,6 +98,18 @@ TileXREp::EpDispatchParams ValidPerTokenDynamicQuantParams() return params; } +TileXREp::EpDispatchParams ValidMemoryMxfp8Params() +{ + TileXREp::EpDispatchParams params = ValidV2Params(); + static uint8_t expandXOutFp8[64] = {}; + static uint8_t dynamicScalesOut[16] = {}; + params.expandXOut = expandXOutFp8; + params.dynamicScalesOut = dynamicScalesOut; + params.quantMode = 4; + params.expandXOutDtype = TileXR::TILEXR_DATA_TYPE_FP8E4M3; + return params; +} + TileXREp::EpCombineParams ValidCombineParams() { static uint16_t expertOut[64] = {}; @@ -107,6 +132,16 @@ TileXREp::EpCombineParams ValidCombineParams() return params; } +TileXREp::EpCombineParams ValidMemoryCombineMxfp8Params(int64_t quantMode) +{ + TileXREp::EpCombineParams params = ValidCombineParams(); + static float expertScales[8] = {1.0f, 0.5f, 1.0f, 0.5f, 1.0f, 0.5f, 1.0f, 0.5f}; + params.expertScales = expertScales; + params.quantMode = quantMode; + params.globalBs = 8; + return params; +} + TileXR::CommArgs ValidCommArgs() { TileXR::CommArgs args {}; @@ -189,6 +224,104 @@ void TestCommValidation() CheckInt("cross-node dispatch needs udma registry", TileXREp::TileXREpValidateDispatchConfig(params, commArgs, &window), TileXR::TILEXR_ERROR_NOT_INITIALIZED); + commArgs = CrossNodeCommArgs(); + CheckInt("cross-node memory dispatch config", TileXREp::TileXREpValidateDispatchMemoryConfig(params, commArgs, + &window), TileXR::TILEXR_SUCCESS); + + commArgs = CrossNodeCommArgs(); + commArgs.peerMems[1] = nullptr; + CheckInt("cross-node memory dispatch missing peer mem", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_ERROR_NOT_INITIALIZED); + + params = ValidV2Params(); + commArgs = ValidCommArgs(); + params.tpWorldSize = 2; + params.epWorldSize = 1; + static int32_t memoryTpRecvCountsOut[2] = {}; + params.tpRecvCountsOut = memoryTpRecvCountsOut; + CheckInt("memory dispatch rejects tp", TileXREp::TileXREpValidateDispatchMemoryConfig(params, commArgs, + &window), TileXR::TILEXR_ERROR_NOT_SUPPORT); + + params = ValidStaticQuantParams(); + commArgs = ValidCommArgs(); + CheckInt("memory dispatch rejects static quant", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_ERROR_NOT_SUPPORT); + + params = ValidMemoryMxfp8Params(); + commArgs = ValidCommArgs(); + CheckInt("memory dispatch supports mxfp8 e4m3", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_SUCCESS); + + params = ValidMemoryMxfp8Params(); + params.expandXOutDtype = TileXR::TILEXR_DATA_TYPE_FP8E5M2; + CheckInt("memory dispatch supports mxfp8 e5m2", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_SUCCESS); + + params = ValidMemoryMxfp8Params(); + params.dynamicScalesOut = nullptr; + CheckInt("memory mxfp8 requires dynamic scales output", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + + params = ValidMemoryMxfp8Params(); + params.expandXOutDtype = TileXR::TILEXR_DATA_TYPE_INT8; + CheckInt("memory mxfp8 requires fp8 output dtype", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + + params = ValidV2Params(); + params.epWorldSize = 4; + params.globalBs = 16; + params.sharedExpertNum = 1; + params.sharedExpertRankNum = 1; + params.moeExpertNum = 6; + commArgs = ValidCommArgs(); + commArgs.rankSize = 4; + commArgs.localRankSize = 4; + commArgs.peerMems[2] = reinterpret_cast(0x30000000); + commArgs.peerMems[3] = reinterpret_cast(0x40000000); + CheckInt("memory dispatch supports shared experts", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_SUCCESS); + + params = ValidV2Params(); + params.xActiveMask = nullptr; + params.activeMaskType = TileXREp::TILEXR_EP_ACTIVE_MASK_TOKEN; + commArgs = ValidCommArgs(); + CheckInt("memory token mask requires pointer", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + + params = ValidV2Params(); + params.activeMaskType = TileXREp::TILEXR_EP_ACTIVE_MASK_NONE; + commArgs = ValidCommArgs(); + CheckInt("memory none mask rejects pointer", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + + params = ValidV2Params(); + static bool expertMask[8] = {true, true, false, true, true, false, true, true}; + params.xActiveMask = expertMask; + params.activeMaskType = TileXREp::TILEXR_EP_ACTIVE_MASK_EXPERT; + commArgs = ValidCommArgs(); + CheckInt("memory expert mask supported", TileXREp::TileXREpValidateDispatchMemoryConfig(params, + commArgs, &window), TileXR::TILEXR_SUCCESS); + + params = ValidV2Params(); + params.epWorldSize = 8; + params.globalBs = 32; + params.sharedExpertNum = 2; + params.sharedExpertRankNum = 4; + params.moeExpertNum = 8; + commArgs = ValidCommArgs(); + commArgs.rankSize = 8; + commArgs.localRankSize = 8; + for (int rank = 0; rank < commArgs.rankSize; ++rank) { + commArgs.peerMems[rank] = reinterpret_cast(0x10000000 + rank * 0x01000000); + } + CheckInt("memory supports multiple ranks per shared expert", + TileXREp::TileXREpValidateDispatchMemoryConfig(params, commArgs, &window), TileXR::TILEXR_SUCCESS); + + params.sharedExpertRankNum = 3; + CheckInt("memory rejects uneven shared expert groups", + TileXREp::TileXREpValidateDispatchMemoryConfig(params, commArgs, &window), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + commArgs = CrossNodeUdmaCommArgs(); params = ValidParams(); static uint8_t workspace[1024] = {}; @@ -251,6 +384,28 @@ void TestCombineValidation() params.workspace = reinterpret_cast(0x50000000); CheckInt("cross-node combine config", TileXREp::TileXREpValidateCombineConfig(params, commArgs, &window), TileXR::TILEXR_SUCCESS); + + TileXREp::EpMemoryCombineReferenceConfig memoryConfig {}; + params = ValidMemoryCombineMxfp8Params(3); + commArgs = ValidCommArgs(); + CheckInt("memory combine supports mxfp8 e5m2", + TileXREp::TileXREpValidateCombineMemoryConfig(params, commArgs, 48, &memoryConfig), + TileXR::TILEXR_SUCCESS); + + params = ValidMemoryCombineMxfp8Params(4); + CheckInt("memory combine supports mxfp8 e4m3", + TileXREp::TileXREpValidateCombineMemoryConfig(params, commArgs, 48, &memoryConfig), + TileXR::TILEXR_SUCCESS); + + params.expertScales = nullptr; + CheckInt("memory combine mxfp8 requires expert scales", + TileXREp::TileXREpValidateCombineMemoryConfig(params, commArgs, 48, &memoryConfig), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + + params = ValidMemoryCombineMxfp8Params(2); + CheckInt("memory combine rejects int8 communication quantization", + TileXREp::TileXREpValidateCombineMemoryConfig(params, commArgs, 48, &memoryConfig), + TileXR::TILEXR_ERROR_NOT_SUPPORT); } void TestV2CapabilityValidation() diff --git a/tests/ep/unit/test_tilexr_ep_layout.cpp b/tests/ep/unit/test_tilexr_ep_layout.cpp index 69c0e7b9..26a5bfc2 100644 --- a/tests/ep/unit/test_tilexr_ep_layout.cpp +++ b/tests/ep/unit/test_tilexr_ep_layout.cpp @@ -1,8 +1,12 @@ +#include #include #include +#include #include "comm_args.h" #include "ep_layout.h" +#include "ep_memory_layout.h" +#include "mxfp8_golden.h" #include "tilexr_types.h" namespace { @@ -100,13 +104,202 @@ void TestRejectsInvalidConfig() TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); } +void CheckUint8(const char *label, uint8_t actual, uint8_t expected) +{ + if (actual != expected) { + std::cerr << label << " actual=0x" << std::hex << static_cast(actual) + << " expected=0x" << static_cast(expected) << std::dec << std::endl; + ++g_failures; + } +} + +void CheckFloat(const char *label, float actual, float expected) +{ + if (std::fabs(actual - expected) > 1.0e-6f) { + std::cerr << label << " actual=" << actual << " expected=" << expected << std::endl; + ++g_failures; + } +} + +void TestMxfp8Golden() +{ + std::vector e4Input(32, 0.0f); + e4Input[0] = 1.0f; + e4Input[1] = -1.0f; + e4Input[2] = 448.0f; + e4Input[3] = -448.0f; + e4Input[4] = std::ldexp(1.0f, -9); + e4Input[5] = std::ldexp(1.0f, -10); + e4Input[6] = 3.0f * std::ldexp(1.0f, -10); + const TileXREpDemo::Mxfp8Tensor e4 = TileXREpDemo::QuantizeMxfp8( + e4Input, 1, e4Input.size(), TileXREpDemo::Mxfp8Format::E4M3); + CheckInt64("mxfp8 e4 scale count", static_cast(e4.scaleCountPerRow), 2); + CheckUint8("mxfp8 e4 one", e4.elements[0], 0x38); + CheckUint8("mxfp8 e4 minus one", e4.elements[1], 0xb8); + CheckUint8("mxfp8 e4 max", e4.elements[2], 0x7e); + CheckUint8("mxfp8 e4 minus max", e4.elements[3], 0xfe); + CheckUint8("mxfp8 e4 min subnormal", e4.elements[4], 0x01); + CheckUint8("mxfp8 e4 tie to even zero", e4.elements[5], 0x00); + CheckUint8("mxfp8 e4 tie to even two", e4.elements[6], 0x02); + CheckUint8("mxfp8 e4 scale", e4.scales[0], 0x7f); + CheckUint8("mxfp8 e4 padded scale", e4.scales[1], 0x00); + + std::vector e5Input(32, 0.0f); + e5Input[0] = 1.0f; + e5Input[1] = -1.0f; + e5Input[2] = 57344.0f; + e5Input[3] = -57344.0f; + e5Input[4] = std::ldexp(1.0f, -16); + e5Input[5] = std::ldexp(1.0f, -17); + e5Input[6] = 3.0f * std::ldexp(1.0f, -17); + const TileXREpDemo::Mxfp8Tensor e5 = TileXREpDemo::QuantizeMxfp8( + e5Input, 1, e5Input.size(), TileXREpDemo::Mxfp8Format::E5M2); + CheckUint8("mxfp8 e5 one", e5.elements[0], 0x3c); + CheckUint8("mxfp8 e5 minus one", e5.elements[1], 0xbc); + CheckUint8("mxfp8 e5 max", e5.elements[2], 0x7b); + CheckUint8("mxfp8 e5 minus max", e5.elements[3], 0xfb); + CheckUint8("mxfp8 e5 min subnormal", e5.elements[4], 0x01); + CheckUint8("mxfp8 e5 tie to even zero", e5.elements[5], 0x00); + CheckUint8("mxfp8 e5 tie to even two", e5.elements[6], 0x02); + CheckUint8("mxfp8 e5 scale", e5.scales[0], 0x7f); + + std::vector blockInput(33, 1.0f); + blockInput[32] = 2.0f; + const TileXREpDemo::Mxfp8Tensor blocks = TileXREpDemo::QuantizeMxfp8( + blockInput, 1, blockInput.size(), TileXREpDemo::Mxfp8Format::E4M3); + CheckInt64("mxfp8 two block scale count", static_cast(blocks.scaleCountPerRow), 2); + CheckUint8("mxfp8 first block scale", blocks.scales[0], 0x77); + CheckUint8("mxfp8 second block scale", blocks.scales[1], 0x78); + + CheckFloat("mxfp8 e4 decode one", + TileXREpDemo::DecodeFp8(0x38, TileXREpDemo::Mxfp8Format::E4M3), 1.0f); + CheckFloat("mxfp8 e4 decode max", + TileXREpDemo::DecodeFp8(0x7e, TileXREpDemo::Mxfp8Format::E4M3), 448.0f); + CheckFloat("mxfp8 e5 decode min subnormal", + TileXREpDemo::DecodeFp8(0x01, TileXREpDemo::Mxfp8Format::E5M2), std::ldexp(1.0f, -16)); + + const std::vector roundTrip = TileXREpDemo::DequantizeMxfp8( + blocks, 1, blockInput.size(), TileXREpDemo::Mxfp8Format::E4M3); + CheckInt64("mxfp8 round trip size", static_cast(roundTrip.size()), 33); + CheckFloat("mxfp8 first block round trip", roundTrip[0], 1.0f); + CheckFloat("mxfp8 second block round trip", roundTrip[32], 2.0f); +} + +void TestMemoryCoreAllocation() +{ + CheckInt("8 experts count cores", TileXREp::TileXREpMemoryCountCoreNum(8, 8, 48), 1); + CheckInt("32 experts count cores", TileXREp::TileXREpMemoryCountCoreNum(32, 32, 48), 2); + CheckInt("count cores capped at eight", TileXREp::TileXREpMemoryCountCoreNum(256, 128, 48), 8); + CheckInt("count cores capped by receive states", TileXREp::TileXREpMemoryCountCoreNum(128, 2, 48), 2); + CheckInt("one block rejected", TileXREp::TileXREpMemoryCountCoreNum(8, 8, 1), 0); +} + +void TestMemoryWindowConfig() +{ + TileXREp::EpMemoryDispatchReferenceConfig config {}; + const int ret = TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 2, 0, 4, 8, 2, 8, 0, 0, 8, TileXR::TILEXR_DATA_TYPE_FP16, 48, &config); + CheckInt("memory config ret", ret, TileXR::TILEXR_SUCCESS); + CheckInt64("memory local experts", config.localExpertNum, 4); + CheckInt64("memory receive states", config.rscvStatusNum, 8); + CheckInt("memory block dim", config.blockDim, 48); + CheckInt("memory all-to-all cores", config.aivUsedAllToAll, 47); + CheckInt("memory count cores", config.aivUsedCumSum, 1); + CheckInt("memory moe cores", config.moeUsedAivNum, 47); + CheckInt("memory shared cores", config.sharedUsedAivNum, 0); + CheckInt64("memory token slot", config.hCommuSize, 512); + CheckInt64("memory expert segment", config.expertPerSizeOnWin, 2048); + CheckInt64("memory combine reserve", config.combineReserveBytes, 4096); + CheckInt64("memory workspace", config.workspaceBytes, 1536); + CheckInt64("memory total window", config.totalWinSize, 103807488); + CheckInt64("memory half window", config.dispatchHalfBytes, 51903744); + + const int sharedRet = TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 8, 0, 4, 8, 2, 8, 2, 4, 32, TileXR::TILEXR_DATA_TYPE_FP16, 48, &config); + CheckInt("shared memory config ret", sharedRet, TileXR::TILEXR_SUCCESS); + CheckInt64("shared rank local experts", config.localExpertNum, 1); + CheckInt64("shared rank receive states", config.rscvStatusNum, 8); + CheckInt("shared dispatch cores", config.sharedUsedAivNum, 23); + CheckInt("shared moe cores", config.moeUsedAivNum, 24); + CheckInt64("shared combine reserve", config.combineReserveBytes, 8192); + + TileXREp::EpMemoryDispatchReferenceConfig sharedRankConfig {}; + TileXREp::EpMemoryDispatchReferenceConfig moeRankConfig {}; + CheckInt("shared role config ret", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 8, 0, 4, 8, 2, 8, 2, 4, 32, TileXR::TILEXR_DATA_TYPE_FP16, 48, &sharedRankConfig), + TileXR::TILEXR_SUCCESS); + CheckInt("moe role config ret", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 8, 4, 4, 8, 2, 8, 2, 4, 32, TileXR::TILEXR_DATA_TYPE_FP16, 48, &moeRankConfig), + TileXR::TILEXR_SUCCESS); + CheckInt64("communicator workspace reservation", sharedRankConfig.workspaceBytes, 3072); + CheckInt64("shared and moe workspace match", sharedRankConfig.workspaceBytes, moeRankConfig.workspaceBytes); + CheckInt64("shared and moe total window match", sharedRankConfig.totalWinSize, moeRankConfig.totalWinSize); + CheckInt64("shared and moe half window match", sharedRankConfig.dispatchHalfBytes, + moeRankConfig.dispatchHalfBytes); + + const int mxfp8Ret = TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 2, 0, 4, 33, 2, 8, 0, 0, 8, TileXR::TILEXR_DATA_TYPE_FP16, + TileXR::TILEXR_DATA_TYPE_FP8E4M3, 4, 48, &config); + CheckInt("mxfp8 memory config ret", mxfp8Ret, TileXR::TILEXR_SUCCESS); + CheckInt64("mxfp8 payload bytes", config.hOutSize, 33); + CheckInt64("mxfp8 scale bytes", config.scaleOutBytes, 2); + CheckInt64("mxfp8 triple offset", config.tokenQuantAlignBytes, 288); + CheckInt64("mxfp8 token slot", config.hCommuSize, 512); + + TileXREp::EpMemoryCombineReferenceConfig combineConfig {}; + const int combineMxfp8Ret = TileXREp::TileXREpBuildMemoryCombineReferenceConfig( + 8, 0, 4, 1024, 2, 8, 0, 0, 32, TileXR::TILEXR_DATA_TYPE_FP16, 4, 48, + &combineConfig); + CheckInt("combine mxfp8 memory config ret", combineMxfp8Ret, TileXR::TILEXR_SUCCESS); + CheckInt64("combine mxfp8 blocks", combineConfig.blockCntPerToken, 3); + CheckInt64("combine mxfp8 packed row", combineConfig.packedRowBytes, 1536); + + const int combineInvalidQuantRet = TileXREp::TileXREpBuildMemoryCombineReferenceConfig( + 8, 0, 4, 1024, 2, 8, 0, 0, 32, TileXR::TILEXR_DATA_TYPE_FP16, 2, 48, + &combineConfig); + CheckInt("combine memory rejects int8 quant", combineInvalidQuantRet, + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); +} + +void TestMemoryWindowRejectsInvalidConfig() +{ + TileXREp::EpMemoryDispatchReferenceConfig config {}; + CheckInt("memory null out", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 2, 0, 4, 8, 2, 8, 0, 0, 8, TileXR::TILEXR_DATA_TYPE_FP16, 48, nullptr), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("memory one block", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 2, 0, 4, 8, 2, 8, 0, 0, 8, TileXR::TILEXR_DATA_TYPE_FP16, 1, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("memory too many blocks", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 2, 0, 4, 8, 2, 8, 0, 0, 8, TileXR::TILEXR_DATA_TYPE_FP16, 201, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("memory int8 rejected", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 2, 0, 4, 8, 2, 8, 0, 0, 8, TileXR::TILEXR_DATA_TYPE_INT8, 48, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("memory uneven shared groups", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 8, 0, 4, 8, 2, 8, 2, 3, 32, TileXR::TILEXR_DATA_TYPE_FP16, 48, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("memory oversized", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 8, 4, 4096, 7168, 8, 64, 0, 0, 32768, + TileXR::TILEXR_DATA_TYPE_FP16, 48, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("memory oversized ub", TileXREp::TileXREpBuildMemoryDispatchReferenceConfig( + 2, 0, 1, 60000, 1, 2, 0, 0, 2, + TileXR::TILEXR_DATA_TYPE_FP16, 48, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); +} + } // namespace int main() { + TestMxfp8Golden(); TestExpertMapping(); TestDataTypes(); TestWindowConfig(); TestRejectsInvalidConfig(); + TestMemoryCoreAllocation(); + TestMemoryWindowConfig(); + TestMemoryWindowRejectsInvalidConfig(); return g_failures == 0 ? 0 : 1; }