From 5b9994d4dcce210a7ceb241322329bb5a8f56640 Mon Sep 17 00:00:00 2001 From: "zhujunan (A)" Date: Thu, 30 Jul 2026 09:54:50 +0800 Subject: [PATCH 1/9] feat(ep): add Ascend 950 S22 URMA combine --- docs/EP_URMA_COMBINE_S22.md | 59 + .../tilexr_ep_urma_combine_s22_bs128.html | 70 + src/comm/CMakeLists.txt | 18 + src/comm/udma/tilexr_udma_layout.cpp | 25 +- src/comm/udma/tilexr_udma_layout.h | 11 + src/comm/udma/tilexr_udma_transport.cpp | 264 +- src/ep/CMakeLists.txt | 91 +- src/ep/common/ep_urma_combine.h | 72 + src/ep/common/ep_urma_combine_config.h | 177 ++ src/ep/common/ep_urma_perf_trace.h | 176 ++ src/ep/common/ep_urma_udma.h | 160 ++ src/ep/host/ep_urma_combine_host.cpp | 175 ++ src/ep/host/ep_urma_combine_host.h | 51 + src/ep/host/ep_urma_combine_layout.cpp | 125 + src/ep/host/ep_urma_combine_layout.h | 39 + .../host/ep_urma_combine_start_gate_window.h | 26 + src/ep/host/tilexr_ep_urma_combine.cpp | 110 + .../kernels/tilexr_ep_urma_combine_kernel.cpp | 2513 +++++++++++++++++ src/include/tilexr_ep.h | 18 + tests/ep/CMakeLists.txt | 11 + tests/ep/unit/test_tilexr_ep_api_sources.cpp | 30 + .../ep/unit/test_tilexr_ep_kernel_sources.cpp | 32 + tests/ep/unit/test_tilexr_ep_layout.cpp | 150 + .../unit/test_tilexr_ep_start_gate_window.cpp | 45 + 24 files changed, 4327 insertions(+), 121 deletions(-) create mode 100644 docs/EP_URMA_COMBINE_S22.md create mode 100644 docs/performance/tilexr_ep_urma_combine_s22_bs128.html create mode 100644 src/ep/common/ep_urma_combine.h create mode 100644 src/ep/common/ep_urma_combine_config.h create mode 100644 src/ep/common/ep_urma_perf_trace.h create mode 100644 src/ep/common/ep_urma_udma.h create mode 100644 src/ep/host/ep_urma_combine_host.cpp create mode 100644 src/ep/host/ep_urma_combine_host.h create mode 100644 src/ep/host/ep_urma_combine_layout.cpp create mode 100644 src/ep/host/ep_urma_combine_layout.h create mode 100644 src/ep/host/ep_urma_combine_start_gate_window.h create mode 100644 src/ep/host/tilexr_ep_urma_combine.cpp create mode 100644 src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp create mode 100644 tests/ep/unit/test_tilexr_ep_start_gate_window.cpp diff --git a/docs/EP_URMA_COMBINE_S22.md b/docs/EP_URMA_COMBINE_S22.md new file mode 100644 index 00000000..c4d7fc60 --- /dev/null +++ b/docs/EP_URMA_COMBINE_S22.md @@ -0,0 +1,59 @@ +# EP URMA Combine for Ascend 950 + +This change adds a TileXR-native Mixture-of-Experts combine operator backed by +URMA/UDMA. The production implementation is fixed to the validated S22 layout: + +- 42 Pack/Receive AIVs and 22 Send AIVs; +- 22 independent UDMA queue pairs; +- one WQE per doorbell; +- QDC-v3 quant/dequant; +- round-robin receive scheduling with sticky ready state; +- parallel release publication and deferred same-parity credit checks; +- first-launch start gate after a stream synchronization; +- BiSheng `-O2`, with profiling disabled by default. + +The fixed configuration avoids carrying experiment-only variant switches into +the production build. + +## Build + +Configure and build through CMake. No variant or wrapper script is required. + +```bash +cmake -S . -B build-s22 \ + -DTILEXR_BUILD_EP=ON \ + -DTILEXR_EP_SOC_TYPE=ascend950 \ + -DBUILD_TESTING=OFF \ + -DCMAKE_INSTALL_PREFIX="$PWD/install-s22" +cmake --build build-s22 -j --target tilexr-ep +cmake --install build-s22 +``` + +`ASCEND_HOME_PATH` and `ASCEND_DRIVER_PATH` must point to the active CANN +toolkit and driver before configuration. A profiling build is available only +when explicitly requested with `-DTILEXR_EP_ENABLE_PROFILING=ON`; it is not the +production default. + +## API + +The public header `src/include/tilexr_ep.h` provides: + +- `TileXRMoeEpCombineUrmaGetWorkspaceSize` for the registered workspace size; +- `TileXRMoeEpCombineUrma` for the production launch; +- `TileXRMoeEpCombineUrmaProfile` for an explicitly profiled launch. + +The workspace must be aligned and registered with TileXR UDMA before launch. +The operator supports an eight-rank Ascend 950 deployment and uses the existing +TileXR communicator for rank information and registered-memory exchange. + +## Performance Evidence + +The production comparison metric is profiling-free `strictKernelCycles`, +aggregated as max core per rank, max rank per launch, then median across 100 +launches. For BS128, H=5120, top-k=6, rank-size=8 and enqueue-window=1, S22 +measured 94,889 cycles (94.889 us using 1000 cycles/us). + +The single retained detailed report is +[BS128 S22 profile](performance/tilexr_ep_urma_combine_s22_bs128.html). The +default proportional timeline excludes the start-gate bar and rebases the axis +to steady-state work; start-gate measurements remain in report metadata. diff --git a/docs/performance/tilexr_ep_urma_combine_s22_bs128.html b/docs/performance/tilexr_ep_urma_combine_s22_bs128.html new file mode 100644 index 00000000..771f2a66 --- /dev/null +++ b/docs/performance/tilexr_ep_urma_combine_s22_bs128.html @@ -0,0 +1,70 @@ + + + + + +TileXR EP URMA Combine - S22 P42/S22/QP22 O2 - 42 Pack/Receive + 22 Send - 22 QPs, 1-WQE doorbell, TX-ready batch 1, shared ready flag off, TX ready in data off, TX metadata full prefetch off, TX early ready off, RX sticky on, RX batched MTE2 off, RX Vector ready off, balanced send routes on, RX token round robin sticky, round publish parallel, start gate enabled: first after stream synchronize (executed in representative launch), Q/D v3 - capture + + + +
capture

TileXR EP URMA Combine · S22 P42/S22/QP22 O2 · 42 Pack/Receive + 22 Send · 22 QPs, 1-WQE doorbell, TX-ready batch 1, shared ready flag off, TX ready in data off, TX metadata full prefetch off, TX early ready off, RX sticky on, RX batched MTE2 off, RX Vector ready off, balanced send routes on, RX token round robin sticky, round publish parallel, start gate enabled: first after stream synchronize (executed in representative launch), Q/D v3

+
8 ranks · 64 AIV · URMA · 10 paired profile samples · representative launch8 · median cluster kernel 89.43 us (coarse attribution profile) · coarse/fine attribution medians 89.43/97.51 us · Gate executed 10/10, executed-only median 259.99 us / 56 WQEs / 56 DBs
+
bs=128, topK=6, h=5120, selfSendCnt=768, routeStride=5632B, enqueueWindow=1 · INT8 per-route quantization · deterministic randomized routes, seed=20260728 · start gate collapsed from steady-state views
+
+
+
+

Coarse Stage Maxima Across AIVs

Each value is the maximum coarse sum_us for that stage across the selected rank's AIVs. Maxima can come from different cores, overlap in time, and must not be added as a critical path.

+

Fine Explicit DCCI Breakdown

Fine diagnostic capture; not the coarse timing launch.

+

64-Core Coarse Timeline

Outlined bars are the exact per-stage maxima listed above. Receive bar width and position are coarse; its internal colors are fine-derived composition, not chronology. The steady-state axis is rebased after the collapsed start gate.

+

Core x Fine Exposed-Time Heatmap

Fine diagnostic capture; absolute durations include instrumentation overhead.

+

Fine Route Transport AIV Diagnostics

Fine diagnostic capture; route counters and exposed stage time are not from the coarse launch.

+

Rank Critical Paths

+
The total-only capture provides the timing headline when supplied; it disables stage attribution inside the measured kernel. Coarse bars are per-core activity envelopes for attribution. Each receive_total envelope is divided into RX flag wait, unpack exposed wait, dequant+clear, and output using the same core's fine exposed wall-time. Fine profiling preserves the quant/dequant double-buffer pipeline: hidden asynchronous transfer time is not charged to an exposed-wait stage. Dequant+clear includes next-route unpack submission bookkeeping, but not that transfer's asynchronous in-flight time. Explicit DCCI spans cover combine-kernel UDMACleanCacheLines calls and exclude their surrounding barriers; SQ/CQ cache maintenance inside the UDMA helper remains part of UDMA post/quiet. DCCI category values are accumulated core-time, not additional coarse wall-time. The colored parts show composition, not chronological ordering, and must not be added to infer hardware-engine occupancy. Legacy fine captures without rx_unpack_wait predate this contract. Repeated fine stages use accumulated exposed wall-time in the heatmap, and their durations include material instrumentation overhead. The detailed timeline uses launch8, selected as the coarse sample nearest the median after maximum core per rank and maximum rank per launch. The sample CSV contains all 10 paired launches. Each rank is independently normalized; the report does not assume synchronized device clocks across NPUs. Open perfetto_trace.json for the standard TileXR drill-down.
+
+ + \ No newline at end of file diff --git a/src/comm/CMakeLists.txt b/src/comm/CMakeLists.txt index e0fd5253..a4468383 100644 --- a/src/comm/CMakeLists.txt +++ b/src/comm/CMakeLists.txt @@ -122,6 +122,24 @@ set(TILEXR_SOURCE_FILE tilexr_comm.cpp add_library(tile-comm SHARED ${TILEXR_SOURCE_FILE}) +# S22 maps each of its 22 Send AIVs to an independent UDMA QP. Keep the +# transport default unchanged for builds that do not target the A5 EP path. +set(_tilexr_udma_qp_count 1) +set(_tilexr_ep_soc_type "${TILEXR_EP_SOC_TYPE}") +if(NOT _tilexr_ep_soc_type) + set(_tilexr_ep_soc_type "$ENV{TILEXR_SOC_NAME}") +endif() +string(TOLOWER "${_tilexr_ep_soc_type}" _tilexr_ep_soc_type_lower) +if(TILEXR_BUILD_EP AND + _tilexr_ep_soc_type_lower MATCHES "ascend950|ascend910_9|ascend910-9|a5") + set(_tilexr_udma_qp_count 22) +endif() +target_compile_definitions(tile-comm PRIVATE + TILEXR_UDMA_QP_COUNT_VALUE=${_tilexr_udma_qp_count}) +unset(_tilexr_udma_qp_count) +unset(_tilexr_ep_soc_type) +unset(_tilexr_ep_soc_type_lower) + if(TILEXR_HAVE_PTO_SDMA) target_compile_definitions(tile-comm PRIVATE TILEXR_HAVE_PTO_SDMA=1) else() diff --git a/src/comm/udma/tilexr_udma_layout.cpp b/src/comm/udma/tilexr_udma_layout.cpp index 401b5f6e..2c4c8cc1 100644 --- a/src/comm/udma/tilexr_udma_layout.cpp +++ b/src/comm/udma/tilexr_udma_layout.cpp @@ -10,8 +10,6 @@ namespace TileXR { namespace { -constexpr uint32_t TILEXR_UDMA_QP_NUM = 1; - template void CopyVector(std::vector& dst, size_t offset, const std::vector& src) { @@ -32,9 +30,24 @@ int BuildUDMAInfoImage( UDMAInfo& info, std::vector& bytes) { - const size_t rankCount = sq.size(); - if (rankCount == 0 || rq.size() != rankCount || scq.size() != rankCount || - rcq.size() != rankCount || mem.size() != rankCount) { + return BuildUDMAInfoImage(deviceBase, 1, sq, rq, scq, rcq, mem, info, bytes); +} + +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpCount, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + UDMAInfo& info, + std::vector& bytes) +{ + const size_t queueEntryCount = sq.size(); + if (qpCount == 0 || queueEntryCount == 0 || queueEntryCount % qpCount != 0 || + rq.size() != queueEntryCount || scq.size() != queueEntryCount || + rcq.size() != queueEntryCount || mem.size() != queueEntryCount) { return TILEXR_UDMA_LAYOUT_INVALID; } @@ -46,7 +59,7 @@ int BuildUDMAInfoImage( const size_t totalBytes = memOffset + mem.size() * sizeof(UDMAMemInfo); info = {}; - info.qpNum = TILEXR_UDMA_QP_NUM; + info.qpNum = qpCount; info.sqPtr = deviceBase + sqOffset; info.rqPtr = deviceBase + rqOffset; info.scqPtr = deviceBase + scqOffset; diff --git a/src/comm/udma/tilexr_udma_layout.h b/src/comm/udma/tilexr_udma_layout.h index fefd49e4..39c94e43 100644 --- a/src/comm/udma/tilexr_udma_layout.h +++ b/src/comm/udma/tilexr_udma_layout.h @@ -26,6 +26,17 @@ int BuildUDMAInfoImage( UDMAInfo& info, std::vector& bytes); +int BuildUDMAInfoImage( + uintptr_t deviceBase, + uint32_t qpCount, + const std::vector& sq, + const std::vector& rq, + const std::vector& scq, + const std::vector& rcq, + const std::vector& mem, + UDMAInfo& info, + std::vector& bytes); + } // namespace TileXR #endif // TILEXR_UDMA_LAYOUT_H diff --git a/src/comm/udma/tilexr_udma_transport.cpp b/src/comm/udma/tilexr_udma_transport.cpp index 7bafa5f1..5f8dcfbb 100644 --- a/src/comm/udma/tilexr_udma_transport.cpp +++ b/src/comm/udma/tilexr_udma_transport.cpp @@ -23,6 +23,11 @@ namespace TileXR { namespace { +#ifndef TILEXR_UDMA_QP_COUNT_VALUE +#define TILEXR_UDMA_QP_COUNT_VALUE 1 +#endif +constexpr uint32_t kTileXRUdmaQpCount = TILEXR_UDMA_QP_COUNT_VALUE; + uint32_t Log2Uint64(uint64_t value) { uint32_t result = 0; @@ -288,11 +293,7 @@ bool ResolveLocalEidRoute( } // namespace -struct TileXRUDMATransport::PerEidState { - uint32_t eidIndex = 0; - void* ctxHandle = nullptr; - void* tokenHandle = nullptr; - void* chanHandle = nullptr; +struct PerQpState { void* cqHandle = nullptr; void* qpHandle = nullptr; CqInfoT cqInfo {}; @@ -309,6 +310,14 @@ struct TileXRUDMATransport::PerEidState { UDMACQCtx localCq {}; }; +struct TileXRUDMATransport::PerEidState { + uint32_t eidIndex = 0; + void* ctxHandle = nullptr; + void* tokenHandle = nullptr; + void* chanHandle = nullptr; + std::vector qps; +}; + TileXRUDMATransport::TileXRUDMATransport() = default; TileXRUDMATransport::~TileXRUDMATransport() @@ -582,12 +591,11 @@ void TileXRUDMATransport::FreeDeviceScalar(void*& ptr) const int TileXRUDMATransport::CreateQueues() { for (const auto& ctxEntry : ctxHandleByEid_) { - PerEidState state {}; + auto& state = states_[ctxEntry.first]; state.eidIndex = ctxEntry.first; state.ctxHandle = ctxEntry.second; state.tokenHandle = tokenHandleByEid_[ctxEntry.first]; - state.remoteQpHandles.assign(options_.rankSize, nullptr); - state.tpnList.assign(options_.rankSize, 0); + state.qps.resize(kTileXRUdmaQpCount); ChanInfoT chanInfo {}; chanInfo.in.dataPlaneFlag.bs.poolCqCstm = 1; @@ -596,91 +604,103 @@ int TileXRUDMATransport::CreateQueues() return TILEXR_ERROR_INTERNAL; } - state.cqInfo.in.chanHandle = state.chanHandle; - state.cqInfo.in.depth = TILEXR_UDMA_CQ_DEPTH; - state.cqInfo.in.ub.mode = JFC_MODE_USER_CTL_NORMAL; - ret = loader_.RaCtxCqCreate(state.ctxHandle, &state.cqInfo, &state.cqHandle); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; - } - state.localCq.cqn = 0; - state.localCq.bufAddr = state.cqInfo.out.bufAddr; - state.localCq.baseBkShift = Log2Uint64(state.cqInfo.out.cqeSize); - state.localCq.depth = state.cqInfo.in.depth; - if (AllocDeviceScalar(&state.cqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.cqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS) { - return TILEXR_ERROR_INTERNAL; - } - state.localCq.headAddr = reinterpret_cast(state.cqPiAddr); - state.localCq.tailAddr = reinterpret_cast(state.cqCiAddr); - state.localCq.dbMode = UDMADBMode::SW_DB; - state.localCq.dbAddr = state.cqInfo.out.swdbAddr; - - QpCreateAttr qpAttr {}; - qpAttr.scqHandle = state.cqHandle; - qpAttr.rcqHandle = state.cqHandle; - qpAttr.srqHandle = state.cqHandle; - qpAttr.sqDepth = TILEXR_UDMA_SQ_DEPTH; - qpAttr.rqDepth = TILEXR_UDMA_RQ_DEPTH_DEFAULT; - qpAttr.transportMode = CONN_RM; - qpAttr.ub.mode = JETTY_MODE_USER_CTL_NORMAL; - qpAttr.ub.flag.value = 1; - qpAttr.ub.jfsFlag.value = 2; - qpAttr.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - qpAttr.ub.rnrRetry = 7; - qpAttr.ub.extMode.piType = 0; - qpAttr.ub.extMode.cstmFlag.bs.sqCstm = 0; - qpAttr.ub.extMode.sqebbNum = TILEXR_UDMA_SQ_DEPTH; - qpAttr.ub.tokenIdHandle = state.tokenHandle; - ret = loader_.RaCtxQpCreate(state.ctxHandle, &qpAttr, &state.qpInfo, &state.qpHandle); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; - } - state.localWq.wqn = 0; - state.localWq.bufAddr = state.qpInfo.ub.sqBuffVa; - state.localWq.baseBkShift = Log2Uint64(state.qpInfo.ub.wqebbSize); - state.localWq.depth = TILEXR_UDMA_SQ_BB_COUNT; - if (AllocDeviceScalar(&state.sqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.sqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.wqeCntAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || - AllocDeviceScalar(&state.amoAddr, sizeof(uint64_t)) != TILEXR_SUCCESS) { - return TILEXR_ERROR_INTERNAL; + for (uint32_t qpIdx = 0; qpIdx < kTileXRUdmaQpCount; ++qpIdx) { + auto& qp = state.qps[qpIdx]; + qp.remoteQpHandles.assign(options_.rankSize, nullptr); + qp.tpnList.assign(options_.rankSize, 0); + + qp.cqInfo.in.chanHandle = state.chanHandle; + qp.cqInfo.in.depth = TILEXR_UDMA_CQ_DEPTH; + qp.cqInfo.in.ub.mode = JFC_MODE_USER_CTL_NORMAL; + ret = loader_.RaCtxCqCreate(state.ctxHandle, &qp.cqInfo, &qp.cqHandle); + if (ret != 0) { + return TILEXR_ERROR_INTERNAL; + } + qp.localCq.cqn = 0; + qp.localCq.bufAddr = qp.cqInfo.out.bufAddr; + qp.localCq.baseBkShift = Log2Uint64(qp.cqInfo.out.cqeSize); + qp.localCq.depth = qp.cqInfo.in.depth; + if (AllocDeviceScalar(&qp.cqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&qp.cqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + qp.localCq.headAddr = reinterpret_cast(qp.cqPiAddr); + qp.localCq.tailAddr = reinterpret_cast(qp.cqCiAddr); + qp.localCq.dbMode = UDMADBMode::SW_DB; + qp.localCq.dbAddr = qp.cqInfo.out.swdbAddr; + + QpCreateAttr qpAttr {}; + qpAttr.scqHandle = qp.cqHandle; + qpAttr.rcqHandle = qp.cqHandle; + qpAttr.srqHandle = qp.cqHandle; + qpAttr.sqDepth = TILEXR_UDMA_SQ_DEPTH; + qpAttr.rqDepth = TILEXR_UDMA_RQ_DEPTH_DEFAULT; + qpAttr.transportMode = CONN_RM; + qpAttr.ub.mode = JETTY_MODE_USER_CTL_NORMAL; + qpAttr.ub.flag.value = 1; + qpAttr.ub.jfsFlag.value = 2; + qpAttr.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + qpAttr.ub.rnrRetry = 7; + qpAttr.ub.extMode.piType = 0; + qpAttr.ub.extMode.cstmFlag.bs.sqCstm = 0; + qpAttr.ub.extMode.sqebbNum = TILEXR_UDMA_SQ_DEPTH; + qpAttr.ub.tokenIdHandle = state.tokenHandle; + ret = loader_.RaCtxQpCreate(state.ctxHandle, &qpAttr, &qp.qpInfo, &qp.qpHandle); + if (ret != 0) { + return TILEXR_ERROR_INTERNAL; + } + qp.localWq.wqn = 0; + qp.localWq.bufAddr = qp.qpInfo.ub.sqBuffVa; + qp.localWq.baseBkShift = Log2Uint64(qp.qpInfo.ub.wqebbSize); + qp.localWq.depth = TILEXR_UDMA_SQ_BB_COUNT; + if (AllocDeviceScalar(&qp.sqPiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&qp.sqCiAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&qp.wqeCntAddr, sizeof(uint32_t)) != TILEXR_SUCCESS || + AllocDeviceScalar(&qp.amoAddr, sizeof(uint64_t)) != TILEXR_SUCCESS) { + return TILEXR_ERROR_INTERNAL; + } + qp.localWq.headAddr = reinterpret_cast(qp.sqPiAddr); + qp.localWq.tailAddr = reinterpret_cast(qp.sqCiAddr); + qp.localWq.dbMode = UDMADBMode::SW_DB; + qp.localWq.dbAddr = qp.qpInfo.ub.dbAddr; + qp.localWq.wqeCntAddr = reinterpret_cast(qp.wqeCntAddr); + qp.localWq.amoAddr = reinterpret_cast(qp.amoAddr); } - state.localWq.headAddr = reinterpret_cast(state.sqPiAddr); - state.localWq.tailAddr = reinterpret_cast(state.sqCiAddr); - state.localWq.dbMode = UDMADBMode::SW_DB; - state.localWq.dbAddr = state.qpInfo.ub.dbAddr; - state.localWq.wqeCntAddr = reinterpret_cast(state.wqeCntAddr); - state.localWq.amoAddr = reinterpret_cast(state.amoAddr); - states_[state.eidIndex] = state; } return states_.empty() ? TILEXR_ERROR_INTERNAL : TILEXR_SUCCESS; } int TileXRUDMATransport::ImportQueues() { - std::vector localImports(eidCount_); - std::vector localKeys(eidCount_); + const size_t entriesPerRank = static_cast(eidCount_) * kTileXRUdmaQpCount; + std::vector localImports(entriesPerRank); + std::vector localKeys(entriesPerRank); for (const auto& stateEntry : states_) { const auto& state = stateEntry.second; if (state.eidIndex >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - localImports[state.eidIndex].in.ub.mode = JETTY_IMPORT_MODE_NORMAL; - localImports[state.eidIndex].in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; - localImports[state.eidIndex].in.ub.policy = JETTY_GRP_POLICY_RR; - localImports[state.eidIndex].in.ub.type = TARGET_TYPE_JETTY; - localImports[state.eidIndex].in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; - localImports[state.eidIndex].in.ub.tpType = 1; - localKeys[state.eidIndex] = state.qpInfo.key; + if (state.qps.size() != kTileXRUdmaQpCount) { + return TILEXR_ERROR_INTERNAL; + } + for (uint32_t qpIdx = 0; qpIdx < kTileXRUdmaQpCount; ++qpIdx) { + const size_t index = static_cast(state.eidIndex) * kTileXRUdmaQpCount + qpIdx; + localImports[index].in.ub.mode = JETTY_IMPORT_MODE_NORMAL; + localImports[index].in.ub.tokenValue = TILEXR_UDMA_TOKEN_VALUE; + localImports[index].in.ub.policy = JETTY_GRP_POLICY_RR; + localImports[index].in.ub.type = TARGET_TYPE_JETTY; + localImports[index].in.ub.flag.bs.tokenPolicy = TOKEN_POLICY_PLAIN_TEXT; + localImports[index].in.ub.tpType = 1; + localKeys[index] = state.qps[qpIdx].qpInfo.key; + } } - std::vector allImports(options_.rankSize * eidCount_); + std::vector allImports(options_.rankSize * entriesPerRank); int ret = options_.exchange->AllGather(localImports.data(), localImports.size(), allImports.data()); if (ret != TILEXR_SUCCESS) { return ret; } - std::vector allKeys(options_.rankSize * eidCount_); + std::vector allKeys(options_.rankSize * entriesPerRank); ret = options_.exchange->AllGather(localKeys.data(), localKeys.size(), allKeys.data()); if (ret != TILEXR_SUCCESS) { return ret; @@ -700,13 +720,18 @@ int TileXRUDMATransport::ImportQueues() if (remoteEid >= eidCount_) { return TILEXR_ERROR_INTERNAL; } - QpImportInfoT importInfo = allImports[peer * eidCount_ + remoteEid]; - importInfo.in.key = allKeys[peer * eidCount_ + remoteEid]; - ret = loader_.RaCtxQpImport(state.ctxHandle, &importInfo, &state.remoteQpHandles[peer]); - if (ret != 0) { - return TILEXR_ERROR_INTERNAL; + for (uint32_t qpIdx = 0; qpIdx < kTileXRUdmaQpCount; ++qpIdx) { + const size_t index = static_cast(peer) * entriesPerRank + + static_cast(remoteEid) * kTileXRUdmaQpCount + qpIdx; + QpImportInfoT importInfo = allImports[index]; + importInfo.in.key = allKeys[index]; + auto& qp = state.qps[qpIdx]; + ret = loader_.RaCtxQpImport(state.ctxHandle, &importInfo, &qp.remoteQpHandles[peer]); + if (ret != 0 || qp.remoteQpHandles[peer] == nullptr) { + return TILEXR_ERROR_INTERNAL; + } + qp.tpnList[peer] = importInfo.out.ub.tpn; } - state.tpnList[peer] = importInfo.out.ub.tpn; } } return TILEXR_SUCCESS; @@ -768,11 +793,12 @@ int TileXRUDMATransport::RefreshUDMAInfo() return TILEXR_ERROR_INTERNAL; } - std::vector sq(options_.rankSize); - std::vector rq(options_.rankSize); - std::vector scq(options_.rankSize); - std::vector rcq(options_.rankSize); - std::vector mem(options_.rankSize); + const size_t queueEntryCount = static_cast(options_.rankSize) * kTileXRUdmaQpCount; + std::vector sq(queueEntryCount); + std::vector rq(queueEntryCount); + std::vector scq(queueEntryCount); + std::vector rcq(queueEntryCount); + std::vector mem(queueEntryCount); for (int rank = 0; rank < options_.rankSize; ++rank) { uint32_t localEid = fallbackEid; @@ -786,26 +812,37 @@ int TileXRUDMATransport::RefreshUDMAInfo() stateIt = fallbackIt; } const auto& state = stateIt->second; - sq[rank] = state.localWq; - rq[rank] = state.localWq; - scq[rank] = state.localCq; - rcq[rank] = state.localCq; + if (state.qps.size() != kTileXRUdmaQpCount) { + return TILEXR_ERROR_INTERNAL; + } + UDMAMemInfo rankMem {}; if (rank == options_.rank) { const auto localMemIt = localMemInfoByEid_.find(localEid); if (localMemIt != localMemInfoByEid_.end()) { - mem[rank] = localMemIt->second; + rankMem = localMemIt->second; } } else { - mem[rank] = allMem[rank * eidCount_ + remoteEid]; - mem[rank].tpn = state.tpnList[rank]; + rankMem = allMem[rank * eidCount_ + remoteEid]; } - mem[rank].eidAddr = reinterpret_cast( + rankMem.eidAddr = reinterpret_cast( eidTableDev_ + (rank * eidCount_ + remoteEid) * sizeof(HccpEid)); + for (uint32_t qpIdx = 0; qpIdx < kTileXRUdmaQpCount; ++qpIdx) { + const size_t index = static_cast(rank) * kTileXRUdmaQpCount + qpIdx; + const auto& qp = state.qps[qpIdx]; + sq[index] = qp.localWq; + rq[index] = qp.localWq; + scq[index] = qp.localCq; + rcq[index] = qp.localCq; + mem[index] = rankMem; + if (rank != options_.rank) { + mem[index].tpn = qp.tpnList[rank]; + } + } } if (udmaInfoDev_ == nullptr) { - const size_t oneRankSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo); - udmaInfoSize_ = static_cast(sizeof(UDMAInfo) + oneRankSize * options_.rankSize); + const size_t oneQueueSize = 2 * sizeof(UDMAWQCtx) + 2 * sizeof(UDMACQCtx) + sizeof(UDMAMemInfo); + udmaInfoSize_ = static_cast(sizeof(UDMAInfo) + oneQueueSize * queueEntryCount); ret = aclrtMalloc(reinterpret_cast(&udmaInfoDev_), udmaInfoSize_, ACL_MEM_MALLOC_HUGE_FIRST); if (ret != ACL_SUCCESS) { return TILEXR_ERROR_INTERNAL; @@ -814,7 +851,8 @@ int TileXRUDMATransport::RefreshUDMAInfo() UDMAInfo info {}; std::vector image; - ret = BuildUDMAInfoImage(reinterpret_cast(udmaInfoDev_), sq, rq, scq, rcq, mem, info, image); + ret = BuildUDMAInfoImage(reinterpret_cast(udmaInfoDev_), kTileXRUdmaQpCount, + sq, rq, scq, rcq, mem, info, image); if (ret != TILEXR_UDMA_LAYOUT_SUCCESS) { return TILEXR_ERROR_PARA_CHECK_FAIL; } @@ -1002,26 +1040,28 @@ void TileXRUDMATransport::CleanupQueues() { for (auto& stateEntry : states_) { auto& state = stateEntry.second; - for (void* remoteQp : state.remoteQpHandles) { - if (remoteQp != nullptr && state.ctxHandle != nullptr) { - loader_.RaCtxQpUnimport(state.ctxHandle, remoteQp); + for (auto& qp : state.qps) { + for (void* remoteQp : qp.remoteQpHandles) { + if (remoteQp != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxQpUnimport(state.ctxHandle, remoteQp); + } } - } - if (state.qpHandle != nullptr) { - loader_.RaCtxQpDestroy(state.qpHandle); - } - if (state.cqHandle != nullptr && state.ctxHandle != nullptr) { - loader_.RaCtxCqDestroy(state.ctxHandle, state.cqHandle); + if (qp.qpHandle != nullptr) { + loader_.RaCtxQpDestroy(qp.qpHandle); + } + if (qp.cqHandle != nullptr && state.ctxHandle != nullptr) { + loader_.RaCtxCqDestroy(state.ctxHandle, qp.cqHandle); + } + FreeDeviceScalar(qp.cqPiAddr); + FreeDeviceScalar(qp.cqCiAddr); + FreeDeviceScalar(qp.sqPiAddr); + FreeDeviceScalar(qp.sqCiAddr); + FreeDeviceScalar(qp.wqeCntAddr); + FreeDeviceScalar(qp.amoAddr); } if (state.chanHandle != nullptr && state.ctxHandle != nullptr) { loader_.RaCtxChanDestroy(state.ctxHandle, state.chanHandle); } - FreeDeviceScalar(state.cqPiAddr); - FreeDeviceScalar(state.cqCiAddr); - FreeDeviceScalar(state.sqPiAddr); - FreeDeviceScalar(state.sqCiAddr); - FreeDeviceScalar(state.wqeCntAddr); - FreeDeviceScalar(state.amoAddr); } states_.clear(); } diff --git a/src/ep/CMakeLists.txt b/src/ep/CMakeLists.txt index c2fc93d5..dfc02ede 100644 --- a/src/ep/CMakeLists.txt +++ b/src/ep/CMakeLists.txt @@ -1,5 +1,14 @@ include(GNUInstallDirs) +option(TILEXR_EP_ENABLE_PROFILING "Enable EP URMA combine kernel profiling" OFF) + +# The protocol constants live in ep_urma_combine_config.h and are fixed +# to the validated Ascend 950 S22 production configuration. +set(TILEXR_EP_PROFILE_COMPILE_DEFINITION) +if(TILEXR_EP_ENABLE_PROFILING) + set(TILEXR_EP_PROFILE_COMPILE_DEFINITION -DTILEXR_EP_ENABLE_PROFILING=1) +endif() + 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") @@ -12,7 +21,9 @@ endif() set(TILEXR_EP_SOC_TYPE "${TILEXR_EP_DEFAULT_SOC_TYPE}" CACHE STRING "TileXR EP kernel SOC type") string(TOLOWER "${TILEXR_EP_SOC_TYPE}" TILEXR_EP_SOC_TYPE_LOWER) +set(TILEXR_EP_BUILD_URMA_COMBINE OFF) if(TILEXR_EP_SOC_TYPE_LOWER MATCHES "ascend950|ascend910_9|ascend910-9|a5") + set(TILEXR_EP_BUILD_URMA_COMBINE ON) set(TILEXR_EP_NPU_ARCH "dav-3510") set(TILEXR_EP_AICORE_ARCH "--cce-aicore-arch=dav-c310-vec") set(TILEXR_EP_CATLASS_ARCH "3510") @@ -40,8 +51,10 @@ 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_COMBINE_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_combine_kernel.cpp") +set(TILEXR_EP_URMA_COMBINE_KERNEL_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/tilexr_ep_urma_combine_kernel.cpp") set(TILEXR_EP_DISPATCH_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_dispatch_kernel.so") set(TILEXR_EP_COMBINE_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_combine_kernel.so") +set(TILEXR_EP_URMA_COMBINE_KERNEL_SO "${CMAKE_CURRENT_BINARY_DIR}/libtilexr_ep_urma_combine_kernel.so") set(TILEXR_EP_KERNEL_INCLUDES -I${ASCEND_HOME_PATH}/compiler/tikcpp -I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw @@ -55,6 +68,8 @@ set(TILEXR_EP_KERNEL_INCLUDES -I${ASCEND_HOME_PATH}/${ARCH}-linux/pkg_inc/runtime -I${ASCEND_HOME_PATH}/${ARCH}-linux/runtime/include -I${ASCEND_HOME_PATH}/${ARCH}-linux/include + -I${ASCEND_HOME_PATH}/${ARCH}-linux/asc/include + -I${ASCEND_HOME_PATH}/${ARCH}-linux/asc/impl -I${ASCEND_DRIVER_PATH}/kernel/inc -I${CMAKE_SOURCE_DIR}/3rdparty -I${CMAKE_SOURCE_DIR}/src/include @@ -135,15 +150,77 @@ add_custom_command( ) add_custom_target(tilexr_ep_combine_kernel ALL DEPENDS "${TILEXR_EP_COMBINE_KERNEL_SO}") -add_library(tilexr-ep SHARED +if(TILEXR_EP_BUILD_URMA_COMBINE) + add_custom_command( + OUTPUT "${TILEXR_EP_URMA_COMBINE_KERNEL_SO}" + COMMAND ${BISHENG_EXECUTABLE} + ${TILEXR_EP_KERNEL_COMPILE_OPTIONS} + -O2 + -std=gnu++17 + -fPIC + -shared + ${TILEXR_EP_KERNEL_LINK_OPTIONS} + -DCATLASS_ARCH=${TILEXR_EP_CATLASS_ARCH} + ${TILEXR_EP_PROFILE_COMPILE_DEFINITION} + ${TILEXR_EP_KERNEL_INCLUDES} + "${TILEXR_EP_URMA_COMBINE_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_URMA_COMBINE_KERNEL_SO}" + DEPENDS + "${TILEXR_EP_URMA_COMBINE_KERNEL_SOURCE}" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_urma_combine_config.h" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_urma_combine.h" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_urma_perf_trace.h" + "${CMAKE_CURRENT_SOURCE_DIR}/common/ep_urma_udma.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" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_perf_trace.h" + "${CMAKE_SOURCE_DIR}/src/include/tilexr_udma.h" + VERBATIM + COMMENT "Building TileXR EP URMA combine kernel (Ascend 950 S22, O2)" + ) + add_custom_target(tilexr_ep_urma_combine_kernel ALL + DEPENDS "${TILEXR_EP_URMA_COMBINE_KERNEL_SO}") +endif() + +set(TILEXR_EP_HOST_SOURCES host/ep_layout.cpp host/ep_dispatch_host.cpp host/ep_launch_context.cpp host/ep_kernel_launch.cpp host/tilexr_ep_dispatch.cpp ) +if(TILEXR_EP_BUILD_URMA_COMBINE) + list(APPEND TILEXR_EP_HOST_SOURCES + host/ep_urma_combine_layout.cpp + host/ep_urma_combine_host.cpp + host/tilexr_ep_urma_combine.cpp) +endif() +add_library(tilexr-ep SHARED ${TILEXR_EP_HOST_SOURCES}) -add_dependencies(tilexr-ep tilexr_ep_dispatch_kernel tilexr_ep_combine_kernel) +add_dependencies(tilexr-ep + tilexr_ep_dispatch_kernel + tilexr_ep_combine_kernel) +if(TILEXR_EP_BUILD_URMA_COMBINE) + add_dependencies(tilexr-ep tilexr_ep_urma_combine_kernel) +endif() + +if(TILEXR_EP_BUILD_URMA_COMBINE AND TILEXR_EP_ENABLE_PROFILING) + target_compile_definitions(tilexr-ep PRIVATE TILEXR_EP_ENABLE_PROFILING=1) +endif() target_include_directories(tilexr-ep PUBLIC @@ -170,6 +247,9 @@ target_link_libraries(tilexr-ep runtime ascend_hal ) +if(TILEXR_EP_BUILD_URMA_COMBINE) + target_link_libraries(tilexr-ep PRIVATE -l:libtilexr_ep_urma_combine_kernel.so) +endif() set_target_properties(tilexr-ep PROPERTIES BUILD_RPATH "${CMAKE_CURRENT_BINARY_DIR}" @@ -182,4 +262,9 @@ install(FILES "${TILEXR_EP_COMBINE_KERNEL_SO}" DESTINATION ${CMAKE_INSTALL_LIBDIR} ) -install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_ep.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +if(TILEXR_EP_BUILD_URMA_COMBINE) + install(FILES "${TILEXR_EP_URMA_COMBINE_KERNEL_SO}" + DESTINATION ${CMAKE_INSTALL_LIBDIR}) +endif() +install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../include/tilexr_ep.h + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) diff --git a/src/ep/common/ep_urma_combine.h b/src/ep/common/ep_urma_combine.h new file mode 100644 index 00000000..382a7a2d --- /dev/null +++ b/src/ep/common/ep_urma_combine.h @@ -0,0 +1,72 @@ +#ifndef TILEXR_EP_COMMON_EP_URMA_COMBINE_H +#define TILEXR_EP_COMMON_EP_URMA_COMBINE_H + +#include +#include + +#include "ep_urma_combine_config.h" + +namespace TileXREp { + +constexpr int64_t kEpUrmaCombineWorkspaceAlignment = 2 * 1024 * 1024; +constexpr int64_t kEpUrmaCombineHeaderBytes = 512; +constexpr int64_t kEpUrmaCombineCacheLineBytes = 64; +constexpr int64_t kEpUrmaCombineDataBlockBytes = 512; +constexpr int64_t kEpUrmaCombinePayloadBytes = 480; +constexpr int64_t kEpUrmaCombineFlagBytes = 32; +constexpr int64_t kEpUrmaCombineQuantHeaderBytes = 32; +constexpr int64_t kEpUrmaCombineAivCount = kEpUrmaCombineProfileCoreCount; +constexpr int64_t kEpUrmaCombinePackLaneCount = kEpUrmaCombineProfilePackReceiveCoreCount; +constexpr int64_t kEpUrmaCombineSendLaneCount = kEpUrmaCombineProfileSendCoreCount; +constexpr int64_t kEpUrmaCombineRequiredQpCount = kEpUrmaCombineSendLaneCount; +constexpr int64_t kEpUrmaCombineMaxHidden = 8192; +constexpr int64_t kEpUrmaCombineMaxTopK = 16; +constexpr int64_t kEpUrmaCombineMaxBlocksPerRoute = 255; +constexpr int64_t kEpUrmaCombineQuantModeInt8PerRoute = 1; + +constexpr uint32_t kEpUrmaCombineTxRouteReady = 1; +constexpr uint32_t kEpUrmaCombineRxLaneDone = 2; +constexpr uint32_t kEpUrmaCombineSenderDone = 3; +constexpr uint32_t kEpUrmaCombineRxBufferReleased = 4; +constexpr uint32_t kEpUrmaCombinePublishDone = 5; +constexpr uint32_t kEpUrmaCombineStartLocalReady = 6; +constexpr uint32_t kEpUrmaCombineStartRankReady = 7; +constexpr uint32_t kEpUrmaCombineStartPublishDone = 8; +constexpr uint32_t kEpUrmaCombineStartRun = 9; +constexpr uint32_t kEpUrmaCombineCreditExpectedReady = 10; +constexpr uint32_t kEpUrmaCombineCreditShardDone = 11; +constexpr uint32_t kEpUrmaCombineCreditRun = 12; +constexpr uint32_t kEpUrmaCombineRxReleaseShardDone = 13; + +constexpr uint64_t kEpUrmaCombineStatusOk = 0; +constexpr uint64_t kEpUrmaCombineStatusInvalidRoute = 1; +constexpr uint64_t kEpUrmaCombineStatusInvalidQuantHeader = 2; + +struct alignas(kEpUrmaCombineQuantHeaderBytes) EpUrmaCombineQuantHeader { + float scale; + int32_t quantMode; + int64_t reserved0; + int64_t reserved1; + int64_t reserved2; +}; + +constexpr int64_t kEpUrmaCombineTxReadyHeaderOffset = + offsetof(EpUrmaCombineQuantHeader, reserved2); + +struct alignas(kEpUrmaCombineCacheLineBytes) EpUrmaCombineControlLine { + uint64_t value; + uint64_t reserved[7]; +}; + +static_assert(sizeof(EpUrmaCombineQuantHeader) == kEpUrmaCombineQuantHeaderBytes, + "URMA combine quant header must be 32 bytes"); +static_assert(kEpUrmaCombineTxReadyHeaderOffset == 24, + "URMA combine in-data TX-ready word must remain at byte 24 of the quant header"); +static_assert(sizeof(EpUrmaCombineControlLine) == kEpUrmaCombineCacheLineBytes, + "URMA combine control line must be one cache line"); +static_assert(kEpUrmaCombinePackLaneCount + kEpUrmaCombineSendLaneCount == kEpUrmaCombineAivCount, + "URMA combine must occupy all configured AIV cores"); + +} // namespace TileXREp + +#endif // TILEXR_EP_COMMON_EP_URMA_COMBINE_H diff --git a/src/ep/common/ep_urma_combine_config.h b/src/ep/common/ep_urma_combine_config.h new file mode 100644 index 00000000..492cc675 --- /dev/null +++ b/src/ep/common/ep_urma_combine_config.h @@ -0,0 +1,177 @@ +#ifndef TILEXR_EP_COMMON_EP_URMA_COMBINE_CONFIG_H +#define TILEXR_EP_COMMON_EP_URMA_COMBINE_CONFIG_H + +#include + +namespace TileXREp { + +// Fixed Ascend 950 S22 production configuration. +#define TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT 22 +#define TILEXR_EP_URMA_DOORBELL_BATCH_SIZE 1 +#define TILEXR_EP_URMA_TX_READY_BATCH_SIZE 1 +#define TILEXR_EP_URMA_TX_READY_SHARED_FLAG 0 +#define TILEXR_EP_URMA_TX_READY_IN_DATA 0 +#define TILEXR_EP_URMA_RX_SCHEDULER 1 +#define TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH 1 +#define TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT 1 +#define TILEXR_EP_URMA_START_GATE 1 +#define TILEXR_EP_URMA_QDC_VERSION 3 +#define TILEXR_EP_URMA_TX_META_PREFETCH_FULL 0 +#define TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH 0 +#define TILEXR_EP_URMA_RX_READY_STICKY_MASK 1 +#define TILEXR_EP_URMA_RX_READY_BATCH_MTE2 0 +#define TILEXR_EP_URMA_RX_READY_BATCH_VECTOR 0 + +constexpr uint32_t kEpUrmaCombineProfileCoreCount = 64; +constexpr uint32_t kEpUrmaCombineProfileSendCoreCount = TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT; +constexpr uint32_t kEpUrmaCombineProfilePackReceiveCoreCount = + kEpUrmaCombineProfileCoreCount - kEpUrmaCombineProfileSendCoreCount; +constexpr uint32_t kEpUrmaCombineDoorbellBatchSize = TILEXR_EP_URMA_DOORBELL_BATCH_SIZE; +constexpr uint32_t kEpUrmaCombineTxReadyBatchSize = TILEXR_EP_URMA_TX_READY_BATCH_SIZE; +constexpr bool kEpUrmaCombineTxReadySharedFlag = TILEXR_EP_URMA_TX_READY_SHARED_FLAG != 0; +constexpr bool kEpUrmaCombineTxReadyInData = TILEXR_EP_URMA_TX_READY_IN_DATA != 0; +constexpr uint32_t kEpUrmaCombineRxScheduler = TILEXR_EP_URMA_RX_SCHEDULER; +constexpr bool kEpUrmaCombineRxRoundRobin = kEpUrmaCombineRxScheduler >= 1; +constexpr bool kEpUrmaCombineRxStickyReady = + kEpUrmaCombineRxScheduler >= 2 || TILEXR_EP_URMA_RX_READY_STICKY_MASK != 0; +constexpr bool kEpUrmaCombineParallelRoundPublish = TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH != 0; +constexpr bool kEpUrmaCombineDeferredRoundCredit = TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT != 0; +constexpr bool kEpUrmaCombineStartGate = TILEXR_EP_URMA_START_GATE != 0; +constexpr uint32_t kEpUrmaCombineQdcVersion = TILEXR_EP_URMA_QDC_VERSION; +constexpr bool kEpUrmaCombineTxMetaPrefetchFull = TILEXR_EP_URMA_TX_META_PREFETCH_FULL != 0; +constexpr bool kEpUrmaCombineTxReadyEarlyPublish = TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH != 0; +constexpr bool kEpUrmaCombineRxReadyStickyMask = TILEXR_EP_URMA_RX_READY_STICKY_MASK != 0; +constexpr bool kEpUrmaCombineRxReadyBatchMte2 = TILEXR_EP_URMA_RX_READY_BATCH_MTE2 != 0; +constexpr bool kEpUrmaCombineRxReadyBatchVector = TILEXR_EP_URMA_RX_READY_BATCH_VECTOR != 0; +constexpr bool kEpUrmaCombineBalancedSendRoutes = true; + +enum class EpUrmaCombinePerfStage : uint32_t { + KERNEL_TOTAL = 0, + PACK_TOTAL = 1, + PACK_INPUT_WAIT = 2, + PACK_QUANTIZE = 3, + PACK_TX_PUBLISH = 4, + RECEIVE_TOTAL = 5, + RX_FLAG_POLL_WAIT = 6, + RX_UNPACK_WAIT = 7, + RX_UNPACK_DEQUANT_CLEAR = 8, + RX_OUTPUT = 9, + SEND_TOTAL = 10, + TX_META_SCAN = 11, + TX_READY_POLL = 12, + SELF_COPY = 13, + UDMA_POST = 14, + UDMA_QUIET = 15, + LOCAL_SENDER_WAIT = 16, + LOCAL_RX_WAIT = 17, + ROUND_PUBLISH = 18, + GLOBAL_ROUND_WAIT = 19, + DCCI_TOTAL = 20, + START_GATE = 21, + PACK_TX_DATA_SUBMIT = 22, + PACK_FIRST_TX_READY = 23, + PACK_MTE3_EXPOSED_WAIT = 24, + RX_READY_MTE2_WAIT = 25, + RX_READY_VECTOR = 26, +}; + +constexpr uint32_t kEpUrmaCombinePerfStageCount = 27; + +#if !defined(__CCE__) || !defined(__CCE_IS_AICORE__) +constexpr const char *const kEpUrmaCombinePerfStageNames[kEpUrmaCombinePerfStageCount] = { + "kernel_total", + "pack_total", + "pack_input_wait", + "pack_quantize", + "pack_tx_publish", + "receive_total", + "rx_flag_poll_wait", + "rx_unpack_wait", + "rx_unpack_dequant_clear", + "rx_output", + "send_total", + "tx_meta_scan", + "tx_ready_poll", + "self_copy", + "udma_post", + "udma_quiet", + "local_sender_wait", + "local_rx_wait", + "round_publish", + "global_round_wait", + "dcci_total", + "start_gate", + "pack_tx_data_submit", + "pack_first_tx_ready", + "pack_mte3_exposed_wait", + "rx_ready_mte2_wait", + "rx_ready_vector", +}; +#endif + +static_assert(kEpUrmaCombineProfilePackReceiveCoreCount + kEpUrmaCombineProfileSendCoreCount == + kEpUrmaCombineProfileCoreCount, + "URMA combine profile core roles must cover all AIV cores"); +static_assert(kEpUrmaCombineProfileSendCoreCount > 0 && + kEpUrmaCombineProfileSendCoreCount < kEpUrmaCombineProfileCoreCount, + "URMA combine Send core count must leave at least one Pack/Receive core"); +static_assert(kEpUrmaCombineDoorbellBatchSize == 1 || kEpUrmaCombineDoorbellBatchSize == 2 || + kEpUrmaCombineDoorbellBatchSize == 4 || kEpUrmaCombineDoorbellBatchSize == 8, + "URMA combine doorbell batch size must be 1, 2, 4, or 8"); +static_assert(kEpUrmaCombineTxReadyBatchSize == 1 || kEpUrmaCombineTxReadyBatchSize == 2 || + kEpUrmaCombineTxReadyBatchSize == 4, + "URMA combine TX-ready batch size must be 1, 2, or 4"); +static_assert(TILEXR_EP_URMA_TX_READY_SHARED_FLAG == 0 || + TILEXR_EP_URMA_TX_READY_SHARED_FLAG == 1, + "URMA combine shared TX-ready flag must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_TX_READY_IN_DATA == 0 || TILEXR_EP_URMA_TX_READY_IN_DATA == 1, + "URMA combine in-data TX-ready flag must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_TX_READY_IN_DATA == 0 || + (TILEXR_EP_URMA_TX_READY_BATCH_SIZE == 1 && TILEXR_EP_URMA_TX_READY_SHARED_FLAG == 0), + "URMA combine in-data TX-ready requires TX batch 1 and shared-ready disabled"); +static_assert(TILEXR_EP_URMA_RX_SCHEDULER >= 0 && TILEXR_EP_URMA_RX_SCHEDULER <= 2, + "URMA combine Receive scheduler must be sequential, round-robin, or round-robin+sticky"); +static_assert(TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH == 0 || + TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH == 1, + "URMA combine parallel round publish must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT == 0 || + TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT == 1, + "URMA combine deferred round credit must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT == 0 || + TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH == 1, + "URMA combine deferred round credit requires parallel round publish"); +static_assert(TILEXR_EP_URMA_START_GATE == 0 || TILEXR_EP_URMA_START_GATE == 1, + "URMA combine start gate must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_QDC_VERSION == 0 || TILEXR_EP_URMA_QDC_VERSION == 1 || + TILEXR_EP_URMA_QDC_VERSION == 2 || TILEXR_EP_URMA_QDC_VERSION == 3, + "URMA combine Quant/Dequant implementation version must be in [0, 3]"); +static_assert(TILEXR_EP_URMA_TX_META_PREFETCH_FULL == 0 || + TILEXR_EP_URMA_TX_META_PREFETCH_FULL == 1, + "URMA combine full metadata prefetch must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH == 0 || + TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH == 1, + "URMA combine early TX-ready publish must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH == 0 || + (TILEXR_EP_URMA_TX_READY_IN_DATA == 0 && TILEXR_EP_URMA_TX_READY_BATCH_SIZE == 1 && + TILEXR_EP_URMA_TX_READY_SHARED_FLAG == 0), + "URMA combine early TX-ready publish requires separate per-route ready lines"); +static_assert(TILEXR_EP_URMA_RX_READY_STICKY_MASK == 0 || + TILEXR_EP_URMA_RX_READY_STICKY_MASK == 1, + "URMA combine RX sticky ready mask must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_RX_READY_BATCH_MTE2 == 0 || + TILEXR_EP_URMA_RX_READY_BATCH_MTE2 == 1, + "URMA combine RX ready MTE2 batching must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_RX_READY_BATCH_VECTOR == 0 || + TILEXR_EP_URMA_RX_READY_BATCH_VECTOR == 1, + "URMA combine RX ready vector reduction must be disabled or enabled"); +static_assert(TILEXR_EP_URMA_RX_READY_STICKY_MASK == 0 || TILEXR_EP_URMA_RX_SCHEDULER > 0, + "URMA combine RX sticky ready mask requires a round-robin scheduler"); +static_assert(TILEXR_EP_URMA_RX_READY_BATCH_MTE2 == 0 || TILEXR_EP_URMA_RX_SCHEDULER > 0, + "URMA combine RX ready MTE2 batching requires a round-robin scheduler"); +static_assert(TILEXR_EP_URMA_RX_READY_BATCH_VECTOR == 0 || + TILEXR_EP_URMA_RX_READY_BATCH_MTE2 == 1, + "URMA combine RX ready vector reduction requires batched MTE2 reads"); + +} // namespace TileXREp + +#endif // TILEXR_EP_COMMON_EP_URMA_COMBINE_CONFIG_H diff --git a/src/ep/common/ep_urma_perf_trace.h b/src/ep/common/ep_urma_perf_trace.h new file mode 100644 index 00000000..7433984c --- /dev/null +++ b/src/ep/common/ep_urma_perf_trace.h @@ -0,0 +1,176 @@ +#ifndef TILEXR_EP_COMMON_EP_URMA_PERF_TRACE_H +#define TILEXR_EP_COMMON_EP_URMA_PERF_TRACE_H + +#include "kernel_operator.h" +#include "tilexr_perf_trace.h" + +namespace TileXR { + +constexpr uint32_t TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT = 27; +constexpr uint32_t TILEXR_PERF_TRACE_LOCAL_STATS_UB_END = 195712; +constexpr uint32_t TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET = + TILEXR_PERF_TRACE_LOCAL_STATS_UB_END - + TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT * sizeof(TileXRPerfCoreStageStats); +constexpr uint32_t TILEXR_PERF_TRACE_MIN_UB_BYTES = 192 * 1024; + +static_assert(TILEXR_PERF_TRACE_LOCAL_STATS_UB_END <= TILEXR_PERF_TRACE_MIN_UB_BYTES, + "local perf stats must fit the minimum supported AIV UB"); +static_assert(TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET % 32 == 0, + "local perf stats must remain 32-byte aligned"); + +#if defined(TILEXR_EP_ENABLE_PROFILING) + +__attribute__((always_inline)) inline __aicore__ bool TileXRPerfTraceEnabled(GM_ADDR trace) +{ + return trace != nullptr; +} + +__attribute__((always_inline)) inline __aicore__ uint64_t TileXRPerfCycleNow(GM_ADDR trace) +{ + return trace == nullptr ? 0 : static_cast(AscendC::GetSystemCycle()); +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalStatsInit( + GM_ADDR trace, __ubuf__ TileXRPerfCoreStageStats *statsUB, uint32_t rank, uint32_t core, + uint32_t stageCount) +{ + if (trace == nullptr || stageCount > TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT) { + return; + } + for (uint32_t stage = 0; stage < stageCount; ++stage) { + statsUB[stage].rank = rank; + statsUB[stage].core = core; + statsUB[stage].stageId = stage; + statsUB[stage].reserved = 0; + statsUB[stage].count = 0; + statsUB[stage].sumCycles = 0; + statsUB[stage].minCycles = 0; + statsUB[stage].maxCycles = 0; + statsUB[stage].firstStartCycle = 0; + statsUB[stage].lastEndCycle = 0; + statsUB[stage].aux0 = 0; + statsUB[stage].aux1 = 0; + statsUB[stage].aux2 = 0; + statsUB[stage].aux3 = 0; + } +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalRecord( + GM_ADDR trace, __ubuf__ TileXRPerfCoreStageStats *statsUB, uint32_t stageCount, + uint32_t stage, uint64_t startCycle, uint64_t endCycle) +{ + if (trace == nullptr || stageCount > TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT || + stage >= stageCount || endCycle < startCycle) { + return; + } + __ubuf__ TileXRPerfCoreStageStats *stat = statsUB + stage; + const uint64_t duration = endCycle - startCycle; + if (stat->count == 0) { + stat->minCycles = duration; + stat->maxCycles = duration; + stat->firstStartCycle = startCycle; + } else { + if (duration < stat->minCycles) { + stat->minCycles = duration; + } + if (duration > stat->maxCycles) { + stat->maxCycles = duration; + } + if (startCycle < stat->firstStartCycle) { + stat->firstStartCycle = startCycle; + } + } + stat->count += 1; + stat->sumCycles += duration; + if (endCycle > stat->lastEndCycle) { + stat->lastEndCycle = endCycle; + } +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalAddAux( + GM_ADDR trace, __ubuf__ TileXRPerfCoreStageStats *statsUB, uint32_t stageCount, + uint32_t stage, uint32_t auxIndex, uint64_t value) +{ + if (trace == nullptr || stageCount > TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT || stage >= stageCount) { + return; + } + __ubuf__ TileXRPerfCoreStageStats *stat = statsUB + stage; + if (auxIndex == 0) { + stat->aux0 += value; + } else if (auxIndex == 1) { + stat->aux1 += value; + } else if (auxIndex == 2) { + stat->aux2 += value; + } else if (auxIndex == 3) { + stat->aux3 += value; + } +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalStatsFlush( + GM_ADDR trace, uint32_t rank, uint32_t core, uint32_t maxCoreCount, uint32_t stageCount, + __ubuf__ TileXRPerfCoreStageStats *statsUB) +{ + if (trace == nullptr || stageCount == 0 || + stageCount > TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT) { + return; + } + const size_t firstSlot = PerfTraceStatsOffset(rank, core, 0, maxCoreCount, stageCount); + __gm__ uint8_t *dstAddr = reinterpret_cast<__gm__ uint8_t *>( + reinterpret_cast<__gm__ TileXRPerfCoreStageStats *>( + trace + TILEXR_PERF_TRACE_STATS_OFFSET) + firstSlot); + AscendC::LocalTensor local; + AscendC::TBuffAddr localAddr; + localAddr.logicPos = static_cast(AscendC::TPosition::VECIN); + localAddr.bufferAddr = reinterpret_cast(statsUB); + local.SetAddr(localAddr); + AscendC::GlobalTensor dst; + dst.SetGlobalBuffer(dstAddr); + AscendC::DataCopyExtParams params { + 1, stageCount * static_cast(sizeof(TileXRPerfCoreStageStats)), 0, 0, 0}; + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyPad(dst, local, params); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + (void)statsUB; +} + +#else + +__attribute__((always_inline)) inline __aicore__ bool TileXRPerfTraceEnabled(GM_ADDR trace) +{ + (void)trace; + return false; +} + +__attribute__((always_inline)) inline __aicore__ uint64_t TileXRPerfCycleNow(GM_ADDR trace) +{ + (void)trace; + return 0; +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalStatsInit( + GM_ADDR, __ubuf__ TileXRPerfCoreStageStats *, uint32_t, uint32_t, uint32_t) +{ +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalRecord( + GM_ADDR, __ubuf__ TileXRPerfCoreStageStats *, uint32_t, uint32_t, uint64_t, uint64_t) +{ +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalAddAux( + GM_ADDR, __ubuf__ TileXRPerfCoreStageStats *, uint32_t, uint32_t, uint32_t, uint64_t) +{ +} + +__attribute__((always_inline)) inline __aicore__ void TileXRPerfLocalStatsFlush( + GM_ADDR, uint32_t, uint32_t, uint32_t, uint32_t, __ubuf__ TileXRPerfCoreStageStats *) +{ +} + +#endif + +} // namespace TileXR + +#endif // TILEXR_EP_COMMON_EP_URMA_PERF_TRACE_H diff --git a/src/ep/common/ep_urma_udma.h b/src/ep/common/ep_urma_udma.h new file mode 100644 index 00000000..5d575925 --- /dev/null +++ b/src/ep/common/ep_urma_udma.h @@ -0,0 +1,160 @@ +#ifndef TILEXR_EP_COMMON_EP_URMA_UDMA_H +#define TILEXR_EP_COMMON_EP_URMA_UDMA_H + +#include "tilexr_udma.h" + +namespace TileXREp { + +// S22 gives each Send AIV an independent QP. Keep the optimized WQE/CQE path +// private to the operator instead of changing TileXR's shared UDMA helpers. +__aicore__ inline __gm__ TileXR::UDMAMemInfo *EpUrmaGetRemoteMemInfo( + __gm__ TileXR::UDMAInfo *info, uint32_t rank, uint32_t qpIdx) +{ + return reinterpret_cast<__gm__ TileXR::UDMAMemInfo *>( + info->memPtr + sizeof(TileXR::UDMAMemInfo) * (rank * info->qpNum + qpIdx)); +} + +__aicore__ inline uint32_t EpUrmaPollCQ( + __gm__ TileXR::UDMAInfo *info, uint32_t rank, uint32_t qpIdx, uint32_t target) +{ + if (target == 0) { + return 0; + } + __gm__ TileXR::UDMACQCtx *cq = TileXR::UDMAGetSCQCtx(info, rank, qpIdx); + __gm__ TileXR::UDMAWQCtx *sq = TileXR::UDMAGetWQCtx(info, rank, qpIdx); + const uint64_t cqBase = cq->bufAddr; + const uint32_t cqeSize = 1U << cq->baseBkShift; + uint32_t tail = ld_dev(reinterpret_cast<__gm__ uint32_t *>(cq->tailAddr), 0); + while (tail != target) { + __gm__ uint32_t *cqeWords = reinterpret_cast<__gm__ uint32_t *>( + cqBase + cqeSize * (tail & (TileXR::TILEXR_UDMA_CQ_DEPTH - 1))); + const bool validOwner = ((tail / TileXR::TILEXR_UDMA_CQ_DEPTH) & 1U) != 0; + uint32_t word0 = 0; + uint32_t retries = 0; + while (retries < TileXR::TILEXR_UDMA_MAX_RETRY_TIMES) { + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t *>(cqeWords), sizeof(TileXR::UDMACqeCtx)); + __asm__ __volatile__("" ::: "memory"); + word0 = cqeWords[0]; + __asm__ __volatile__("" ::: "memory"); + if ((validOwner ^ ((word0 & (1U << 2U)) != 0)) != 0) { + break; + } + ++retries; + } + if (retries >= TileXR::TILEXR_UDMA_MAX_RETRY_TIMES) { + return 0xFFU; + } + const uint8_t status = static_cast((word0 >> 24U) & 0xFFU); + const uint8_t subStatus = static_cast((word0 >> 16U) & 0xFFU); + if (status != 0 || subStatus != 0) { + return (static_cast(status) << 8U) | subStatus; + } + ++tail; + } + TileXR::UDMAPollCQUpdateInfo(tail, cq, sq); + return 0; +} + +__aicore__ inline void EpUrmaFillWriteWqe( + __gm__ uint8_t *wqe, __gm__ uint8_t *remoteAddr, + __gm__ TileXR::UDMAMemInfo *remote, uint32_t head, uint32_t depth, + uint64_t bytes, __gm__ uint8_t *localAddr) +{ + __gm__ uint32_t *sqe = reinterpret_cast<__gm__ uint32_t *>(wqe); + const uint32_t owner = (head & depth) == 0U ? 1U : 0U; + sqe[0] = (head % depth) | (0b00100010U << 16U) | + ((remote->tokenValueValid ? 1U : 0U) << 28U) | + ((remote->rmtJettyType & 0x3U) << 29U) | (owner << 31U); + sqe[1] = (static_cast(remote->targetHint) & 0xFFU) | + (static_cast(TileXR::UDMAOpcode::WRITE) << 8U); + sqe[2] = (remote->tpn & 0xFFFFFFU) | (1U << 24U); + sqe[3] = remote->tid & 0xFFFFFU; + __gm__ uint32_t *eid = reinterpret_cast<__gm__ uint32_t *>(remote->eidAddr); + sqe[4] = eid[0]; + sqe[5] = eid[1]; + sqe[6] = eid[2]; + sqe[7] = eid[3]; + sqe[8] = remote->rmtTokenValue; + sqe[9] = 0; + const uint64_t remoteValue = reinterpret_cast(remoteAddr); + sqe[10] = static_cast(remoteValue); + sqe[11] = static_cast(remoteValue >> 32U); + + __gm__ uint32_t *sge = sqe + sizeof(TileXR::UDMASqeCtx) / sizeof(uint32_t); + const uint64_t localValue = reinterpret_cast(localAddr); + sge[0] = static_cast(bytes); + sge[1] = 0; + sge[2] = static_cast(localValue); + sge[3] = static_cast(localValue >> 32U); +} + +__aicore__ inline void EpUrmaPostWrite( + __gm__ TileXR::UDMAInfo *info, __gm__ uint8_t *remoteAddr, + __gm__ uint8_t *localAddr, uint32_t rank, uint32_t qpIdx, uint64_t bytes) +{ + __gm__ TileXR::UDMAWQCtx *sq = TileXR::UDMAGetWQCtx(info, rank, qpIdx); + const uint32_t wqeSize = 1U << sq->baseBkShift; + uint32_t head = ld_dev(reinterpret_cast<__gm__ uint32_t *>(sq->headAddr), 0); + uint32_t count = ld_dev(reinterpret_cast<__gm__ uint32_t *>(sq->wqeCntAddr), 0); + const uint32_t tail = ld_dev(reinterpret_cast<__gm__ uint32_t *>(sq->tailAddr), 0); + if ((count + 10U) % TileXR::TILEXR_UDMA_SQ_BB_COUNT == + tail % TileXR::TILEXR_UDMA_SQ_BB_COUNT) { + const uint32_t target = tail + TileXR::TILEXR_UDMA_NUM_CQE_PER_POLL > count ? + count : tail + TileXR::TILEXR_UDMA_NUM_CQE_PER_POLL; + (void)EpUrmaPollCQ(info, rank, qpIdx, target); + } + + __gm__ uint8_t *wqe = reinterpret_cast<__gm__ uint8_t *>( + sq->bufAddr + wqeSize * (head % TileXR::TILEXR_UDMA_SQ_BB_COUNT)); + EpUrmaFillWriteWqe(wqe, remoteAddr, EpUrmaGetRemoteMemInfo(info, rank, qpIdx), + head, sq->depth, bytes, localAddr); + AscendC::PipeBarrier(); + TileXR::UDMACleanCacheLines(wqe, wqeSize); + AscendC::PipeBarrier(); + ++head; + st_dev(head, reinterpret_cast<__gm__ uint32_t *>(sq->dbAddr), 0); + st_dev(head, reinterpret_cast<__gm__ uint32_t *>(sq->headAddr), 0); + st_dev(++count, reinterpret_cast<__gm__ uint32_t *>(sq->wqeCntAddr), 0); +} + +template +__aicore__ inline void EpUrmaUDMAPutNbi( + const __gm__ TileXR::CommArgs *args, int targetRank, const __gm__ T *localSrc, + uint64_t byteOffset, uint32_t byteCount, uint32_t qpIdx = 0) +{ + if (!TileXR::UDMARegistryEnabled(args)) { + return; + } + __gm__ TileXR::UDMAInfo *info = TileXR::GetUDMAInfo(args); + if (qpIdx >= info->qpNum) { + return; + } + __gm__ TileXR::TileXRUDMARegistry *registry = TileXR::GetUDMARegistry(args); + if (!TileXR::UDMARegisteredRangeValid(registry, targetRank, byteOffset, byteCount)) { + return; + } + EpUrmaPostWrite(info, TileXR::UDMARegisteredRemoteAddr(registry, targetRank, byteOffset), + reinterpret_cast<__gm__ uint8_t *>(const_cast<__gm__ T *>(localSrc)), + static_cast(targetRank), qpIdx, byteCount); +} + +__aicore__ inline uint32_t EpUrmaUDMAQuiet( + const __gm__ TileXR::CommArgs *args, int targetRank, uint32_t qpIdx = 0) +{ + if (!TileXR::UDMAEnabled(args)) { + return 0; + } + __gm__ TileXR::UDMAInfo *info = TileXR::GetUDMAInfo(args); + if (qpIdx >= info->qpNum) { + return 0; + } + __gm__ TileXR::UDMAWQCtx *sq = + TileXR::UDMAGetWQCtx(info, static_cast(targetRank), qpIdx); + const uint32_t count = ld_dev(reinterpret_cast<__gm__ uint32_t *>(sq->wqeCntAddr), 0); + return EpUrmaPollCQ(info, static_cast(targetRank), qpIdx, count); +} + +} // namespace TileXREp + +#endif // TILEXR_EP_COMMON_EP_URMA_UDMA_H diff --git a/src/ep/host/ep_urma_combine_host.cpp b/src/ep/host/ep_urma_combine_host.cpp new file mode 100644 index 00000000..ec67f6bc --- /dev/null +++ b/src/ep/host/ep_urma_combine_host.cpp @@ -0,0 +1,175 @@ +#include "ep_urma_combine_host.h" + +#include +#include + + +#include "ep_urma_combine.h" +#include "tilexr_perf_trace.h" +#include "tilexr_udma_reg.h" + +extern void launch_tilexr_ep_urma_combine_kernel(uint32_t blockDim, void *stream, GM_ADDR commArgs, + GM_ADDR expertOut, GM_ADDR assistInfoForCombine, GM_ADDR topKWeights, GM_ADDR yOut, GM_ADDR workspace, + int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, int64_t workspaceBytes, int64_t magic, + int64_t commBytes, int64_t blockCount, int64_t routeStride, int64_t rxWindowBytes, int64_t rxWindowOffset0, + int64_t rxWindowOffset1, int64_t roundDoneOffset0, int64_t roundDoneOffset1, int64_t rxLaneDoneOffset, + int64_t senderDoneOffset, int64_t roundPublishOffset, int64_t roundCreditOffset, + int64_t startGateOffset, int64_t runStartGate, + int64_t errorStatusOffset, int64_t txReadyOffset, int64_t txDataOffset, GM_ADDR perfTrace, int64_t perfTraceBytes, + GM_ADDR strictKernelCycles); + +namespace TileXREp { + + +int TileXREpGetUrmaCombineProfileSize(int64_t rankSize, int64_t *profileBytes) +{ + if (profileBytes == nullptr || rankSize <= 0 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + constexpr int64_t kStatsPerRank = kEpUrmaCombineAivCount * kEpUrmaCombinePerfStageCount; + constexpr int64_t kStatsBytesPerRank = + kStatsPerRank * static_cast(sizeof(TileXR::TileXRPerfCoreStageStats)); + if (rankSize > (std::numeric_limits::max() - + static_cast(TileXR::TILEXR_PERF_TRACE_STATS_OFFSET)) / kStatsBytesPerRank) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + *profileBytes = static_cast(TileXR::TILEXR_PERF_TRACE_STATS_OFFSET) + + rankSize * kStatsBytesPerRank; + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpValidateBasicUrmaCombineParams(const EpUrmaCombineParams ¶ms) +{ + if ((params.selfSendCnt > 0 && (params.expertOut == nullptr || params.assistInfoForCombine == nullptr)) || + params.topKWeights == nullptr || params.comm == nullptr || params.yOut == nullptr || + params.workspace == nullptr || params.stream == nullptr || params.selfSendCnt < 0 || params.bs <= 0 || + params.h <= 0 || params.h > kEpUrmaCombineMaxHidden || params.topK <= 0 || + params.topK > kEpUrmaCombineMaxTopK || params.workspaceBytes <= 0 || params.perfTraceBytes < 0 || + ((params.perfTrace == nullptr) != (params.perfTraceBytes == 0)) || + params.strictKernelCyclesBytes < 0 || + ((params.strictKernelCycles == nullptr) != (params.strictKernelCyclesBytes == 0)) || + params.dtype != TileXR::TILEXR_DATA_TYPE_FP16) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if ((reinterpret_cast(params.workspace) % kEpUrmaCombineWorkspaceAlignment) != 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (params.perfTrace != nullptr && (reinterpret_cast(params.perfTrace) % 32) != 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + constexpr int64_t kStrictKernelCyclesBytes = + kEpUrmaCombineAivCount * static_cast(sizeof(uint64_t)); + if (params.strictKernelCycles != nullptr && (params.perfTrace != nullptr || + params.strictKernelCyclesBytes < kStrictKernelCyclesBytes || + (reinterpret_cast(params.strictKernelCycles) % 32) != 0)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpPrepareUrmaCombineLaunchContext( + const EpUrmaCombineParams ¶ms, EpUrmaCombineLaunchContext *context) +{ + if (context == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + *context = EpUrmaCombineLaunchContext {}; + + int ret = TileXREpValidateBasicUrmaCombineParams(params); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + ret = TileXRGetCommArgsHost(params.comm, context->hostArgs); + if (ret != TileXR::TILEXR_SUCCESS || context->hostArgs == nullptr) { + *context = EpUrmaCombineLaunchContext {}; + 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 = EpUrmaCombineLaunchContext {}; + return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + } + + const TileXR::CommArgs &args = *context->hostArgs; + if (args.rankSize <= 0 || args.rankSize > TileXR::TILEXR_MAX_RANK_SIZE || args.rank < 0 || + args.rank >= args.rankSize || (args.rankSize > 1 && + ((args.extraFlag & TileXR::ExtraFlag::UDMA) == 0 || args.udmaInfoPtr == nullptr || + args.udmaRegistryPtr == nullptr))) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_NOT_INITIALIZED; + } + + if (params.perfTrace != nullptr) { + int64_t requiredProfileBytes = 0; + ret = TileXREpGetUrmaCombineProfileSize(args.rankSize, &requiredProfileBytes); + if (ret != TileXR::TILEXR_SUCCESS || params.perfTraceBytes < requiredProfileBytes) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + + ret = TileXREpBuildUrmaCombineWorkspaceConfig( + args.rankSize, params.bs, params.h, params.topK, params.selfSendCnt, &context->workspace); + if (ret != TileXR::TILEXR_SUCCESS || context->workspace.requiredBytes > params.workspaceBytes) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + if (args.rankSize > 1) { + const TileXR::TileXRUDMARegistry *registry = nullptr; + ret = TileXRGetUDMARegistryHost(params.comm, ®istry); + if (ret != TileXR::TILEXR_SUCCESS || !TileXR::UDMARegistryValid(registry, args.rankSize)) { + *context = EpUrmaCombineLaunchContext {}; + return ret == TileXR::TILEXR_SUCCESS ? TileXR::TILEXR_ERROR_NOT_INITIALIZED : ret; + } + if (registry->regions[args.rank].base != static_cast(params.workspace) || + !TileXR::UDMARegionContains(registry, args.rank, 0, + static_cast(context->workspace.requiredBytes))) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + for (int32_t peer = 0; peer < args.rankSize; ++peer) { + if (!TileXR::UDMARegionContains( + registry, peer, 0, static_cast(context->workspace.fixedBytes))) { + *context = EpUrmaCombineLaunchContext {}; + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + } + } + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpLaunchPreparedUrmaCombineKernel( + const EpUrmaCombineParams ¶ms, const EpUrmaCombineLaunchContext &context, int64_t magic, + bool runStartGate) +{ + launch_tilexr_ep_urma_combine_kernel(static_cast(kEpUrmaCombineAivCount), + params.stream, context.devArgs, + static_cast(params.expertOut), reinterpret_cast(params.assistInfoForCombine), + reinterpret_cast(params.topKWeights), static_cast(params.yOut), + static_cast(params.workspace), params.selfSendCnt, params.bs, params.h, params.topK, + params.workspaceBytes, magic, context.workspace.commBytes, context.workspace.blockCount, + context.workspace.routeStride, context.workspace.rxWindowBytes, context.workspace.rxWindowOffsets[0], + context.workspace.rxWindowOffsets[1], context.workspace.roundDoneOffsets[0], + context.workspace.roundDoneOffsets[1], context.workspace.rxLaneDoneOffset, + context.workspace.senderDoneOffset, context.workspace.roundPublishOffset, + context.workspace.roundCreditOffset, + context.workspace.startGateOffset, runStartGate ? 1 : 0, context.workspace.errorStatusOffset, + context.workspace.txReadyOffset, context.workspace.txDataOffset, + static_cast(params.perfTrace), params.perfTraceBytes, + static_cast(params.strictKernelCycles)); + return TileXR::TILEXR_SUCCESS; +} + +int TileXREpLaunchUrmaCombineKernel( + const EpUrmaCombineParams ¶ms, const EpUrmaCombineLaunchContext &context) +{ + int64_t magic = 0; + const int ret = TileXRCommNextMagic(params.comm, &magic); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + return TileXREpLaunchPreparedUrmaCombineKernel(params, context, magic, true); +} + +} // namespace TileXREp diff --git a/src/ep/host/ep_urma_combine_host.h b/src/ep/host/ep_urma_combine_host.h new file mode 100644 index 00000000..6e7b531b --- /dev/null +++ b/src/ep/host/ep_urma_combine_host.h @@ -0,0 +1,51 @@ +#ifndef TILEXR_EP_HOST_EP_URMA_COMBINE_HOST_H +#define TILEXR_EP_HOST_EP_URMA_COMBINE_HOST_H + +#include + +#include "acl/acl_base.h" +#include "ep_urma_combine_layout.h" +#include "tilexr_api.h" +#include "tilexr_types.h" + +namespace TileXREp { + +struct EpUrmaCombineParams { + void *expertOut = nullptr; + int32_t *assistInfoForCombine = nullptr; + float *topKWeights = nullptr; + TileXRCommPtr comm = nullptr; + int64_t selfSendCnt = 0; + int64_t bs = 0; + int64_t h = 0; + int64_t topK = 0; + void *yOut = nullptr; + void *workspace = nullptr; + int64_t workspaceBytes = 0; + void *perfTrace = nullptr; + int64_t perfTraceBytes = 0; + void *strictKernelCycles = nullptr; + int64_t strictKernelCyclesBytes = 0; + TileXR::TileXRDataType dtype = TileXR::TILEXR_DATA_TYPE_RESERVED; + aclrtStream stream = nullptr; +}; + +struct EpUrmaCombineLaunchContext { + TileXR::CommArgs *hostArgs = nullptr; + GM_ADDR devArgs = nullptr; + EpUrmaCombineWorkspaceConfig workspace {}; +}; + +int TileXREpGetUrmaCombineProfileSize(int64_t rankSize, int64_t *profileBytes); +int TileXREpValidateBasicUrmaCombineParams(const EpUrmaCombineParams ¶ms); +int TileXREpPrepareUrmaCombineLaunchContext( + const EpUrmaCombineParams ¶ms, EpUrmaCombineLaunchContext *context); +int TileXREpLaunchPreparedUrmaCombineKernel( + const EpUrmaCombineParams ¶ms, const EpUrmaCombineLaunchContext &context, int64_t magic, + bool runStartGate = true); +int TileXREpLaunchUrmaCombineKernel( + const EpUrmaCombineParams ¶ms, const EpUrmaCombineLaunchContext &context); + +} // namespace TileXREp + +#endif // TILEXR_EP_HOST_EP_URMA_COMBINE_HOST_H diff --git a/src/ep/host/ep_urma_combine_layout.cpp b/src/ep/host/ep_urma_combine_layout.cpp new file mode 100644 index 00000000..1958c60c --- /dev/null +++ b/src/ep/host/ep_urma_combine_layout.cpp @@ -0,0 +1,125 @@ +#include "ep_urma_combine_layout.h" + +#include + +#include "comm_args.h" +#include "ep_urma_combine.h" +#include "tilexr_types.h" + +namespace TileXREp { +namespace { + +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 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 AlignUpInt64(int64_t value, int64_t alignment, int64_t *out) +{ + if (out == nullptr || value < 0 || alignment <= 0) { + return false; + } + const int64_t remainder = value % alignment; + return remainder == 0 ? (*out = value, true) : AddInt64(value, alignment - remainder, out); +} + +bool AddAlignedRegion(int64_t *offset, int64_t bytes, int64_t alignment, int64_t *regionOffset) +{ + int64_t aligned = 0; + if (offset == nullptr || regionOffset == nullptr || bytes < 0 || + !AlignUpInt64(*offset, alignment, &aligned) || !AddInt64(aligned, bytes, offset)) { + return false; + } + *regionOffset = aligned; + return true; +} + +} // namespace + +int TileXREpBuildUrmaCombineWorkspaceConfig(int64_t rankSize, int64_t bs, int64_t h, int64_t topK, + int64_t selfSendCnt, EpUrmaCombineWorkspaceConfig *out) +{ + if (out == nullptr || rankSize <= 0 || rankSize > TileXR::TILEXR_MAX_RANK_SIZE || bs <= 0 || h <= 0 || + h > kEpUrmaCombineMaxHidden || topK <= 0 || topK > kEpUrmaCombineMaxTopK || selfSendCnt < 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + EpUrmaCombineWorkspaceConfig next {}; + next.rankSize = rankSize; + next.bs = bs; + next.h = h; + next.topK = topK; + next.selfSendCnt = selfSendCnt; + + if (!AlignUpInt64(h, kEpUrmaCombineQuantHeaderBytes, &next.quantDataBytes) || + !AddInt64(kEpUrmaCombineQuantHeaderBytes, next.quantDataBytes, &next.commBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.blockCount = (next.commBytes + kEpUrmaCombinePayloadBytes - 1) / kEpUrmaCombinePayloadBytes; + if (next.blockCount <= 0 || next.blockCount > kEpUrmaCombineMaxBlocksPerRoute || + !MulInt64(next.blockCount, kEpUrmaCombineDataBlockBytes, &next.routeStride) || + !MulInt64(bs, topK, &next.routeCount) || + !MulInt64(next.routeCount, next.routeStride, &next.rxWindowBytes)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + int64_t offset = kEpUrmaCombineHeaderBytes; + int64_t roundDoneBytes = 0; + int64_t rxLaneDoneBytes = 0; + int64_t senderDoneBytes = 0; + int64_t txReadyBytes = 0; + int64_t txDataBytes = 0; + if (!MulInt64(rankSize, kEpUrmaCombineCacheLineBytes, &roundDoneBytes) || + !MulInt64(kEpUrmaCombinePackLaneCount, kEpUrmaCombineCacheLineBytes, &rxLaneDoneBytes) || + !MulInt64(kEpUrmaCombineSendLaneCount, kEpUrmaCombineCacheLineBytes, &senderDoneBytes) || + !MulInt64(selfSendCnt, kEpUrmaCombineCacheLineBytes, &txReadyBytes) || + !MulInt64(selfSendCnt, next.routeStride, &txDataBytes) || + !AddAlignedRegion(&offset, next.rxWindowBytes, kEpUrmaCombineDataBlockBytes, &next.rxWindowOffsets[0]) || + !AddAlignedRegion(&offset, next.rxWindowBytes, kEpUrmaCombineDataBlockBytes, &next.rxWindowOffsets[1]) || + !AddAlignedRegion(&offset, roundDoneBytes, kEpUrmaCombineCacheLineBytes, &next.roundDoneOffsets[0]) || + !AddAlignedRegion(&offset, roundDoneBytes, kEpUrmaCombineCacheLineBytes, &next.roundDoneOffsets[1]) || + !AddAlignedRegion(&offset, rxLaneDoneBytes, kEpUrmaCombineCacheLineBytes, &next.rxLaneDoneOffset) || + !AddAlignedRegion(&offset, senderDoneBytes, kEpUrmaCombineCacheLineBytes, &next.senderDoneOffset) || + !AddAlignedRegion(&offset, kEpUrmaCombineCacheLineBytes, kEpUrmaCombineCacheLineBytes, + &next.roundPublishOffset)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (kEpUrmaCombineDeferredRoundCredit && + !AddAlignedRegion(&offset, kEpUrmaCombineCacheLineBytes, kEpUrmaCombineCacheLineBytes, + &next.roundCreditOffset)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (kEpUrmaCombineStartGate && + !AddAlignedRegion(&offset, roundDoneBytes, kEpUrmaCombineCacheLineBytes, &next.startGateOffset)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + if (!AddAlignedRegion(&offset, kEpUrmaCombineCacheLineBytes, kEpUrmaCombineCacheLineBytes, + &next.errorStatusOffset)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + + next.fixedBytes = offset; + if (!AddAlignedRegion(&offset, txReadyBytes, kEpUrmaCombineCacheLineBytes, &next.txReadyOffset) || + !AddAlignedRegion(&offset, txDataBytes, kEpUrmaCombineDataBlockBytes, &next.txDataOffset)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + next.requiredBytes = offset; + *out = next; + return TileXR::TILEXR_SUCCESS; +} + +} // namespace TileXREp diff --git a/src/ep/host/ep_urma_combine_layout.h b/src/ep/host/ep_urma_combine_layout.h new file mode 100644 index 00000000..fc71bd85 --- /dev/null +++ b/src/ep/host/ep_urma_combine_layout.h @@ -0,0 +1,39 @@ +#ifndef TILEXR_EP_HOST_EP_URMA_COMBINE_LAYOUT_H +#define TILEXR_EP_HOST_EP_URMA_COMBINE_LAYOUT_H + +#include + +namespace TileXREp { + +struct EpUrmaCombineWorkspaceConfig { + int64_t rankSize = 0; + int64_t bs = 0; + int64_t h = 0; + int64_t topK = 0; + int64_t selfSendCnt = 0; + int64_t quantDataBytes = 0; + int64_t commBytes = 0; + int64_t blockCount = 0; + int64_t routeStride = 0; + int64_t routeCount = 0; + int64_t rxWindowBytes = 0; + int64_t rxWindowOffsets[2] = {0, 0}; + int64_t roundDoneOffsets[2] = {0, 0}; + int64_t rxLaneDoneOffset = 0; + int64_t senderDoneOffset = 0; + int64_t roundPublishOffset = 0; + int64_t roundCreditOffset = 0; + int64_t startGateOffset = 0; + int64_t errorStatusOffset = 0; + int64_t fixedBytes = 0; + int64_t txReadyOffset = 0; + int64_t txDataOffset = 0; + int64_t requiredBytes = 0; +}; + +int TileXREpBuildUrmaCombineWorkspaceConfig(int64_t rankSize, int64_t bs, int64_t h, int64_t topK, + int64_t selfSendCnt, EpUrmaCombineWorkspaceConfig *out); + +} // namespace TileXREp + +#endif // TILEXR_EP_HOST_EP_URMA_COMBINE_LAYOUT_H diff --git a/src/ep/host/ep_urma_combine_start_gate_window.h b/src/ep/host/ep_urma_combine_start_gate_window.h new file mode 100644 index 00000000..78c9d907 --- /dev/null +++ b/src/ep/host/ep_urma_combine_start_gate_window.h @@ -0,0 +1,26 @@ +#ifndef TILEXR_EP_HOST_EP_URMA_COMBINE_START_GATE_WINDOW_H +#define TILEXR_EP_HOST_EP_URMA_COMBINE_START_GATE_WINDOW_H + +namespace TileXREp { + +class EpUrmaCombineStartGateWindow { +public: + bool BeginLaunch() + { + const bool runStartGate = firstLaunchAfterSynchronization_; + firstLaunchAfterSynchronization_ = false; + return runStartGate; + } + + void StreamSynchronized() + { + firstLaunchAfterSynchronization_ = true; + } + +private: + bool firstLaunchAfterSynchronization_ = true; +}; + +} // namespace TileXREp + +#endif // TILEXR_EP_HOST_EP_URMA_COMBINE_START_GATE_WINDOW_H diff --git a/src/ep/host/tilexr_ep_urma_combine.cpp b/src/ep/host/tilexr_ep_urma_combine.cpp new file mode 100644 index 00000000..f7df518f --- /dev/null +++ b/src/ep/host/tilexr_ep_urma_combine.cpp @@ -0,0 +1,110 @@ +#include "tilexr_ep.h" + +#include + +#include "ep_urma_combine.h" +#include "ep_urma_combine_host.h" + +int TileXRMoeEpCombineUrmaGetWorkspaceSize(int64_t rankSize, int64_t bs, int64_t h, int64_t topK, + int64_t selfSendCapacity, int64_t *workspaceBytes) +{ + if (workspaceBytes == nullptr) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + TileXREp::EpUrmaCombineWorkspaceConfig layout {}; + const int ret = TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig( + rankSize, bs, h, topK, selfSendCapacity, &layout); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + const int64_t alignment = TileXREp::kEpUrmaCombineWorkspaceAlignment; + const int64_t remainder = layout.requiredBytes % alignment; + if (remainder != 0 && layout.requiredBytes > + std::numeric_limits::max() - (alignment - remainder)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + *workspaceBytes = remainder == 0 ? layout.requiredBytes : layout.requiredBytes + alignment - remainder; + return TileXR::TILEXR_SUCCESS; +} + +int TileXRMoeEpCombineUrmaGetProfileSize(int64_t rankSize, int64_t *profileBytes) +{ + return TileXREp::TileXREpGetUrmaCombineProfileSize(rankSize, profileBytes); +} + +namespace { + +int CombineUrmaImpl(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, void *perfTrace, int64_t perfTraceBytes, + TileXR::TileXRDataType dtype, aclrtStream stream) +{ + TileXREp::EpUrmaCombineParams params {}; + params.expertOut = expertOut; + params.assistInfoForCombine = assistInfoForCombine; + params.topKWeights = topKWeights; + params.comm = comm; + params.selfSendCnt = selfSendCnt; + params.bs = bs; + params.h = h; + params.topK = topK; + params.yOut = yOut; + params.workspace = workspace; + params.workspaceBytes = workspaceBytes; + params.perfTrace = perfTrace; + params.perfTraceBytes = perfTraceBytes; + params.dtype = dtype; + params.stream = stream; + + int ret = TileXREp::TileXREpValidateBasicUrmaCombineParams(params); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + TileXREp::EpUrmaCombineLaunchContext context {}; + ret = TileXREp::TileXREpPrepareUrmaCombineLaunchContext(params, &context); + if (ret != TileXR::TILEXR_SUCCESS) { + return ret; + } + return TileXREp::TileXREpLaunchUrmaCombineKernel(params, context); +} + +} // namespace + +int TileXRMoeEpCombineUrma(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, TileXR::TileXRDataType dtype, aclrtStream stream) +{ + return CombineUrmaImpl(expertOut, assistInfoForCombine, topKWeights, comm, selfSendCnt, bs, h, topK, + yOut, workspace, workspaceBytes, nullptr, 0, dtype, stream); +} + +int TileXRMoeEpCombineUrmaProfile(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, void *perfTrace, int64_t perfTraceBytes, + TileXR::TileXRDataType dtype, aclrtStream stream) +{ +#if !defined(TILEXR_EP_ENABLE_PROFILING) + (void)expertOut; + (void)assistInfoForCombine; + (void)topKWeights; + (void)comm; + (void)selfSendCnt; + (void)bs; + (void)h; + (void)topK; + (void)yOut; + (void)workspace; + (void)workspaceBytes; + (void)perfTrace; + (void)perfTraceBytes; + (void)dtype; + (void)stream; + return TileXR::TILEXR_ERROR_NOT_SUPPORT; +#else + if (perfTrace == nullptr || perfTraceBytes <= 0) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } + return CombineUrmaImpl(expertOut, assistInfoForCombine, topKWeights, comm, selfSendCnt, bs, h, topK, + yOut, workspace, workspaceBytes, perfTrace, perfTraceBytes, dtype, stream); +#endif +} diff --git a/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp new file mode 100644 index 00000000..4948d0d7 --- /dev/null +++ b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp @@ -0,0 +1,2513 @@ +#include "comm_args.h" +#include "ep_urma_combine.h" +#include "ep_urma_udma.h" +#include "ep_window.h" +#include "kernel_operator.h" +#include "tilexr_data_as_flag.h" +#include "ep_urma_perf_trace.h" +#include "tilexr_udma.h" + +namespace { + +#if !defined(TILEXR_EP_URMA_CACHELESS) +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) +#define TILEXR_EP_URMA_CACHELESS 1 +#else +#define TILEXR_EP_URMA_CACHELESS 0 +#endif +#endif + +constexpr uint32_t kCursorAlignment = 32; +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL +constexpr uint32_t kRouteMetaPrefetchCapacityBytes = 16 * 1024; +static_assert(sizeof(TileXREp::EpAssistTuple) == 16, + "URMA combine metadata prefetch requires the existing 16-byte tuple ABI"); +#endif +#if !TILEXR_EP_URMA_CACHELESS +constexpr uint32_t kRouteMetaBufferBytes = 32; +#endif +constexpr int64_t kPipelineBufferCount = 2; +using PackInputQueue = AscendC::TQue; +constexpr int64_t kRxTokenScheduleWindow = 3; + + +#define TILEXR_EP_SEND_FUNCTION inline + +#define TILEXR_EP_LOCAL_FUNCTION inline + +using PerfStats = TileXR::TileXRPerfCoreStageStats; +using PerfStage = TileXREp::EpUrmaCombinePerfStage; + +enum class DcciProfileCategory : uint32_t { + TX_DATA = 0, + RX_FLAG_POLL = 1, + RX_DATA = 2, + CONTROL_OTHER = 3, +}; + +static_assert(TileXREp::kEpUrmaCombinePerfStageCount <= + TileXR::TILEXR_PERF_TRACE_LOCAL_MAX_STAGE_COUNT, + "URMA combine profile stages must fit the reserved UB trace region"); +static_assert(TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET + + TileXREp::kEpUrmaCombinePerfStageCount * sizeof(PerfStats) <= + TileXR::TILEXR_PERF_TRACE_MIN_UB_BYTES, + "URMA combine profile stats exceed the minimum supported AIV UB"); + +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint32_t PerfStageId(PerfStage stage) +{ + return static_cast(stage); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t ProfileBegin(GM_ADDR perfTrace) +{ + return TileXR::TileXRPerfCycleNow(perfTrace); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void ProfileEnd( + GM_ADDR perfTrace, __ubuf__ PerfStats *stats, PerfStage stage, uint64_t startCycle) +{ + TileXR::TileXRPerfLocalRecord(perfTrace, stats, TileXREp::kEpUrmaCombinePerfStageCount, + PerfStageId(stage), startCycle, TileXR::TileXRPerfCycleNow(perfTrace)); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void ProfileAux( + GM_ADDR perfTrace, __ubuf__ PerfStats *stats, PerfStage stage, uint32_t auxIndex, uint64_t value) +{ + TileXR::TileXRPerfLocalAddAux(perfTrace, stats, TileXREp::kEpUrmaCombinePerfStageCount, + PerfStageId(stage), auxIndex, value); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void ProfileDcciEnd( + GM_ADDR perfTrace, __ubuf__ PerfStats *stats, + DcciProfileCategory category, uint64_t startCycle) +{ + const uint64_t endCycle = TileXR::TileXRPerfCycleNow(perfTrace); + TileXR::TileXRPerfLocalRecord(perfTrace, stats, TileXREp::kEpUrmaCombinePerfStageCount, + PerfStageId(PerfStage::DCCI_TOTAL), startCycle, endCycle); + if (endCycle >= startCycle) { + TileXR::TileXRPerfLocalAddAux(perfTrace, stats, TileXREp::kEpUrmaCombinePerfStageCount, + PerfStageId(PerfStage::DCCI_TOTAL), static_cast(category), endCycle - startCycle); + } +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t ProfileKernelTimingBegin(GM_ADDR perfTrace) +{ + if (TileXR::TileXRPerfTraceEnabled(perfTrace)) { + AscendC::PipeBarrier(); + return static_cast(AscendC::GetSystemCycle()); + } + return 0; +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void ProfileFinish( + GM_ADDR perfTrace, uint32_t rank, uint32_t core, + __ubuf__ PerfStats *stats, uint64_t kernelStart) +{ + if (TileXR::TileXRPerfTraceEnabled(perfTrace)) { + AscendC::PipeBarrier(); + const uint64_t kernelEnd = static_cast(AscendC::GetSystemCycle()); + TileXR::TileXRPerfLocalRecord(perfTrace, stats, + TileXREp::kEpUrmaCombinePerfStageCount, + PerfStageId(PerfStage::KERNEL_TOTAL), kernelStart, kernelEnd); + } + TileXR::TileXRPerfLocalStatsFlush(perfTrace, rank, core, + static_cast(TileXREp::kEpUrmaCombineAivCount), + TileXREp::kEpUrmaCombinePerfStageCount, stats); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t StrictKernelTimingBegin( + GM_ADDR strictKernelCycles) +{ + if (strictKernelCycles == nullptr) { + return 0; + } + AscendC::PipeBarrier(); + return static_cast(AscendC::GetSystemCycle()); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void StrictKernelTimingFinish( + GM_ADDR strictKernelCycles, uint32_t core, uint64_t startCycle) +{ + if (strictKernelCycles == nullptr) { + return; + } + AscendC::PipeBarrier(); + const uint64_t endCycle = static_cast(AscendC::GetSystemCycle()); + const uint64_t duration = endCycle >= startCycle ? endCycle - startCycle : 0; +#if TILEXR_EP_URMA_CACHELESS + AscendC::WriteGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>(strictKernelCycles) + core, duration); + AscendC::DataSyncBarrier(); +#else + reinterpret_cast<__gm__ uint64_t *>(strictKernelCycles)[core] = duration; +#endif +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool ProfileBufferValid( + GM_ADDR perfTrace, int64_t perfTraceBytes, int32_t rank, int32_t rankSize) +{ + if (perfTrace == nullptr || perfTraceBytes < + static_cast(sizeof(TileXR::TileXRPerfTraceHeader)) || rank < 0 || rankSize <= 0) { + return false; + } + const uint64_t requiredStatsBytes = static_cast(rankSize) * + TileXREp::kEpUrmaCombineAivCount * TileXREp::kEpUrmaCombinePerfStageCount * sizeof(PerfStats); + const uint64_t requiredBytes = TileXR::TILEXR_PERF_TRACE_STATS_OFFSET + requiredStatsBytes; + if (requiredBytes > static_cast(perfTraceBytes)) { + return false; + } + const auto header = reinterpret_cast<__gm__ TileXR::TileXRPerfTraceHeader *>(perfTrace); + return header->magic == TileXR::TILEXR_PERF_TRACE_MAGIC && + header->version == TileXR::TILEXR_PERF_TRACE_VERSION && + header->headerSize == sizeof(TileXR::TileXRPerfTraceHeader) && + header->coreStageStatsSize == sizeof(PerfStats) && + header->rank == static_cast(rank) && + header->rankSize == static_cast(rankSize) && + header->blockDim == TileXREp::kEpUrmaCombineAivCount && + header->maxCoreCount == TileXREp::kEpUrmaCombineAivCount && + header->stageCount == TileXREp::kEpUrmaCombinePerfStageCount && + header->flags <= 2 && header->cycleToUsDivisor != 0 && + header->statsOffset == TileXR::TILEXR_PERF_TRACE_STATS_OFFSET && + header->statsBytes == requiredStatsBytes; +} + +static_assert(TileXREp::kEpUrmaCombineDataBlockBytes == TileXR::DATA_AS_FLAG_BLOCK_BYTES, + "URMA combine and DataAsFlag block sizes must match"); +static_assert(TileXREp::kEpUrmaCombinePayloadBytes == TileXR::DATA_AS_FLAG_PAYLOAD_BYTES, + "URMA combine and DataAsFlag payload sizes must match"); +static_assert(TileXREp::kEpUrmaCombineFlagBytes == TileXR::DATA_AS_FLAG_FLAG_BYTES, + "URMA combine and DataAsFlag flag sizes must match"); +constexpr int64_t kMaxInputBytes = TileXREp::kEpUrmaCombineMaxHidden * sizeof(half); +constexpr int64_t kMaxFloatBytes = TileXREp::kEpUrmaCombineMaxHidden * sizeof(float); +constexpr float kMaxFiniteHalf = 65504.0f; +constexpr int64_t kMaxLogicalBytes = + ((TileXREp::kEpUrmaCombineQuantHeaderBytes + TileXREp::kEpUrmaCombineMaxHidden + + TileXREp::kEpUrmaCombinePayloadBytes - 1) / + TileXREp::kEpUrmaCombinePayloadBytes) * TileXREp::kEpUrmaCombinePayloadBytes; +constexpr int64_t kMaxBlockCount = kMaxLogicalBytes / TileXREp::kEpUrmaCombinePayloadBytes; +constexpr int64_t kMaxRouteStride = kMaxBlockCount * TileXREp::kEpUrmaCombineDataBlockBytes; +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 +constexpr int64_t kMaxPackUbBytes = kPipelineBufferCount * kMaxInputBytes + + kPipelineBufferCount * kMaxRouteStride + + kMaxFloatBytes + 2 * kMaxInputBytes + kCursorAlignment + kMaxLogicalBytes; +#else +constexpr int64_t kMaxPackUbBytes = kPipelineBufferCount * kMaxInputBytes + + kPipelineBufferCount * kMaxRouteStride + + 3 * kMaxFloatBytes + kCursorAlignment + kMaxLogicalBytes; +#endif +constexpr int64_t kMaxRouteFlagBytes = + kMaxBlockCount * TileXREp::kEpUrmaCombineFlagBytes; +#if TILEXR_EP_URMA_RX_READY_BATCH_MTE2 +constexpr int64_t kMaxReadyFlagBufferBytes = + TileXREp::kEpUrmaCombineMaxTopK * kMaxRouteFlagBytes; +#else +constexpr int64_t kMaxReadyFlagBufferBytes = kMaxRouteFlagBytes; +#endif +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR +constexpr int64_t kMaxReadyVectorBufferBytes = + TileXREp::kEpUrmaCombineMaxTopK * kMaxRouteFlagBytes + kMaxRouteFlagBytes + + TileXREp::kEpUrmaCombineMaxTopK * kCursorAlignment; +#else +constexpr int64_t kMaxReadyVectorBufferBytes = 0; +#endif +constexpr int64_t kMaxReceiveUbBytes = kPipelineBufferCount * kMaxLogicalBytes + + kMaxReadyFlagBufferBytes + kMaxRouteFlagBytes + kMaxReadyVectorBufferBytes + + 2 * kMaxFloatBytes + kMaxInputBytes + + TileXREp::kEpUrmaCombineMaxTopK * static_cast(sizeof(float)); +static_assert(kMaxPackUbBytes <= TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET, + "URMA combine Pack ping-pong buffers overlap the profiling UB region"); +static_assert(kMaxReceiveUbBytes <= TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET, + "URMA combine Receive ping-pong buffers overlap the profiling UB region"); +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL +constexpr int64_t kMaxSendCursorBytes = + ((TileXREp::kEpUrmaCombinePackLaneCount * static_cast(sizeof(uint32_t)) + + kCursorAlignment - 1) / kCursorAlignment) * kCursorAlignment; +constexpr int64_t kMaxSendUsedPeerBytes = + ((TileXR::TILEXR_MAX_RANK_SIZE + kCursorAlignment - 1) / kCursorAlignment) * kCursorAlignment; +constexpr int64_t kMaxSendSelfCopyBytes = kMaxBlockCount * + (TileXREp::kEpUrmaCombinePayloadBytes + TileXREp::kEpUrmaCombineFlagBytes); +constexpr int64_t kMaxSendUbBytes = kMaxSendCursorBytes + kMaxSendUsedPeerBytes + + kMaxSendSelfCopyBytes + kRouteMetaPrefetchCapacityBytes; +static_assert(kMaxSendUbBytes <= TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET, + "URMA combine Send metadata prefetch overlaps the profiling UB region"); +#endif + +__aicore__ TILEXR_EP_LOCAL_FUNCTION int64_t AlignUpInt64(int64_t value, int64_t alignment) +{ + const int64_t remainder = value % alignment; + return remainder == 0 ? value : value + alignment - remainder; +} + +template +__aicore__ TILEXR_EP_LOCAL_FUNCTION void SetNoCacheRead(AscendC::GlobalTensor &tensor) +{ +#if TILEXR_EP_URMA_CACHELESS + tensor.template SetL2CacheHint( + AscendC::CacheMode::CACHE_MODE_DISABLE); +#else + (void)tensor; +#endif +} + +template +__aicore__ TILEXR_EP_LOCAL_FUNCTION void SetNoCacheWrite(AscendC::GlobalTensor &tensor) +{ +#if TILEXR_EP_URMA_CACHELESS + tensor.template SetL2CacheHint( + AscendC::CacheMode::CACHE_MODE_DISABLE); +#else + (void)tensor; +#endif +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void CachelessAcquireBarrier() +{ +#if TILEXR_EP_URMA_CACHELESS + AscendC::DataSyncBarrier(); + AscendC::PipeBarrier(); +#endif +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t EncodeControlValue(int64_t magic, uint32_t step) +{ + return (static_cast(static_cast(magic)) << 32) | static_cast(step); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void StoreControlValue( + GM_ADDR lineAddr, uint64_t value, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ +#if TILEXR_EP_URMA_CACHELESS + AscendC::PipeBarrier(); + AscendC::WriteGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>(lineAddr), value); + AscendC::DataSyncBarrier(); +#else + auto line = reinterpret_cast<__gm__ uint64_t *>(lineAddr); + line[0] = value; + AscendC::PipeBarrier(); + const uint64_t dcciStart = ProfileBegin(perfTrace); + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t *>(lineAddr), TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::CONTROL_OTHER, dcciStart); + AscendC::PipeBarrier(); +#endif +} + + +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t LoadControlValue( + GM_ADDR lineAddr, GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ +#if TILEXR_EP_URMA_CACHELESS + const uint64_t value = AscendC::ReadGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>(lineAddr)); + return value; +#else + const uint64_t dcciStart = ProfileBegin(perfTrace); + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t *>(lineAddr), TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::CONTROL_OTHER, dcciStart); + AscendC::PipeBarrier(); + return reinterpret_cast<__gm__ uint64_t *>(lineAddr)[0]; +#endif +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void StoreError( + GM_ADDR workspaceGM, int64_t errorStatusOffset, int64_t magic, uint64_t status, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + if (status != TileXREp::kEpUrmaCombineStatusOk) { + StoreControlValue(workspaceGM + errorStatusOffset, + EncodeControlValue(magic, static_cast(status)), perfTrace, perfStats); + } +} + +#if TILEXR_EP_URMA_CACHELESS || TILEXR_EP_URMA_TX_META_PREFETCH_FULL +__aicore__ TILEXR_EP_LOCAL_FUNCTION TileXREp::EpAssistTuple LoadRouteMetaBypass( + GM_ADDR assistInfoGM, int64_t index) +{ + __gm__ int32_t *src = reinterpret_cast<__gm__ int32_t *>(assistInfoGM) + + index * TileXREp::kEpAssistTupleInts; + TileXREp::EpAssistTuple tuple; + tuple.srcRank = AscendC::ReadGmByPassDCache(src); + tuple.tokenId = AscendC::ReadGmByPassDCache(src + 1); + tuple.topKId = AscendC::ReadGmByPassDCache(src + 2); + tuple.expertId = 0; + return tuple; +} +#endif + +#if TILEXR_EP_URMA_CACHELESS +__aicore__ TILEXR_EP_LOCAL_FUNCTION TileXREp::EpAssistTuple LoadRouteMeta( + GM_ADDR assistInfoGM, int64_t index) +{ + return LoadRouteMetaBypass(assistInfoGM, index); +} +#else +__aicore__ TILEXR_EP_LOCAL_FUNCTION TileXREp::EpAssistTuple LoadRouteMeta( + GM_ADDR assistInfoGM, int64_t index, AscendC::LocalTensor &local) +{ + AscendC::GlobalTensor src; + src.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(assistInfoGM) + + index * TileXREp::kEpAssistTupleInts, TileXREp::kEpAssistTupleInts); + AscendC::DataCopyExtParams copyParams { + 1, static_cast(sizeof(TileXREp::EpAssistTuple)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams {false, 0, 0, 0}; + AscendC::DataCopyPad(local, src, copyParams, padParams); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + TileXREp::EpAssistTuple tuple; + tuple.srcRank = local.GetValue(0); + tuple.tokenId = local.GetValue(1); + tuple.topKId = local.GetValue(2); + tuple.expertId = local.GetValue(3); + return tuple; +} +#endif + +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL +__aicore__ TILEXR_EP_LOCAL_FUNCTION void StartRouteMetaPrefetch( + GM_ADDR assistInfoGM, int64_t tupleCount, + AscendC::LocalTensor &local, event_t eventId) +{ + AscendC::GlobalTensor src; + src.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(assistInfoGM), + tupleCount * TileXREp::kEpAssistTupleInts); + src.SetL2CacheHint( + AscendC::CacheMode::CACHE_MODE_DISABLE); + AscendC::DataCopyExtParams copyParams { + 1, static_cast(tupleCount * static_cast(sizeof(TileXREp::EpAssistTuple))), + 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams {false, 0, 0, 0}; + AscendC::DataCopyPad(local, src, copyParams, padParams); + AscendC::SetFlag(eventId); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION TileXREp::EpAssistTuple LoadPrefetchedRouteMeta( + AscendC::LocalTensor &local, int64_t index) +{ + const uint32_t begin = static_cast(index * TileXREp::kEpAssistTupleInts); + TileXREp::EpAssistTuple tuple; + tuple.srcRank = local.GetValue(begin); + tuple.tokenId = local.GetValue(begin + 1); + tuple.topKId = local.GetValue(begin + 2); + tuple.expertId = 0; + return tuple; +} +#endif + +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool RouteMetaValid( + const TileXREp::EpAssistTuple &tuple, int32_t rankSize, int64_t bs, int64_t topK) +{ + return tuple.srcRank >= 0 && tuple.srcRank < rankSize && tuple.tokenId >= 0 && tuple.tokenId < bs && + tuple.topKId >= 0 && tuple.topKId < topK; +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void LoadTopKWeights( + GM_ADDR weightsGM, int64_t begin, int64_t topK, AscendC::LocalTensor &local) +{ + AscendC::GlobalTensor src; + src.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(weightsGM) + begin, topK); + const uint32_t alignedElements = static_cast( + AlignUpInt64(topK * static_cast(sizeof(float)), kCursorAlignment) / sizeof(float)); + AscendC::DataCopyExtParams copyParams { + 1, static_cast(topK * static_cast(sizeof(float))), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams { + true, 0, static_cast(alignedElements - static_cast(topK)), 0.0f}; + AscendC::DataCopyPad(local, src, copyParams, padParams); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION event_t PipelineEvent(int64_t slot) +{ + return slot == 0 ? EVENT_ID0 : EVENT_ID1; +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void EnqueuePackInput( + const AscendC::GlobalTensor &expertOut, int64_t route, + int64_t h, uint8_t inputPadding, PackInputQueue &inputQueue) +{ + AscendC::LocalTensor input = inputQueue.AllocTensor(); + AscendC::DataCopyExtParams inputParams { + 1, static_cast(h * static_cast(sizeof(half))), 0, 0, 0}; + AscendC::DataCopyPadExtParams inputPad {true, 0, inputPadding, static_cast(0.0f)}; + AscendC::DataCopyPad(input, expertOut[route * h], inputParams, inputPad); + inputQueue.EnQue(input); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void StartPackPublish( + GM_ADDR txRouteAddr, int64_t routeStride, + AscendC::LocalTensor &packed, event_t eventId) +{ + AscendC::GlobalTensor txRoute; + txRoute.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(txRouteAddr), routeStride / sizeof(float)); + SetNoCacheWrite(txRoute); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(txRoute, packed, static_cast(routeStride / sizeof(float))); + AscendC::SetFlag(eventId); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void FinishPackCopy( + GM_ADDR txRouteAddr, int64_t routeStride, event_t eventId, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t waitStart = ProfileBegin(perfTrace); + AscendC::WaitFlag(eventId); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_MTE3_EXPOSED_WAIT, waitStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_MTE3_EXPOSED_WAIT, 0, + static_cast(routeStride)); +#if !TILEXR_EP_URMA_CACHELESS + const uint64_t dcciStart = ProfileBegin(perfTrace); + TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t *>(txRouteAddr), routeStride); + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::TX_DATA, dcciStart); +#endif +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void PublishTxReadyInData( + GM_ADDR txRouteAddr, uint64_t readyValue, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ +#if TILEXR_EP_URMA_CACHELESS + AscendC::PipeBarrier(); + AscendC::WriteGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>( + txRouteAddr + TileXREp::kEpUrmaCombineTxReadyHeaderOffset), + readyValue); + AscendC::DataSyncBarrier(); +#else + reinterpret_cast<__gm__ uint64_t *>( + txRouteAddr + TileXREp::kEpUrmaCombineTxReadyHeaderOffset)[0] = readyValue; + AscendC::PipeBarrier(); + const uint64_t dcciStart = ProfileBegin(perfTrace); + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t *>(txRouteAddr), + TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::CONTROL_OTHER, dcciStart); + AscendC::PipeBarrier(); +#endif +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t LoadTxReadyInData( + GM_ADDR txRouteAddr, GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ +#if TILEXR_EP_URMA_CACHELESS + const uint64_t value = AscendC::ReadGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>( + txRouteAddr + TileXREp::kEpUrmaCombineTxReadyHeaderOffset)); + return value; +#else + const uint64_t dcciStart = ProfileBegin(perfTrace); + TileXR::UDMACleanCacheLines( + reinterpret_cast<__gm__ uint8_t *>(txRouteAddr), + TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::CONTROL_OTHER, dcciStart); + AscendC::PipeBarrier(); + return reinterpret_cast<__gm__ uint64_t *>( + txRouteAddr + TileXREp::kEpUrmaCombineTxReadyHeaderOffset)[0]; +#endif +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void PublishTxReadyBatch( + GM_ADDR workspaceGM, int64_t txReadyOffset, + int64_t txDataOffset, int64_t routeStride, int64_t firstRoute, int64_t routeCount, + uint64_t readyValue, GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + if (routeCount <= 0) { + return; + } +#if TILEXR_EP_URMA_TX_READY_IN_DATA + (void)txReadyOffset; + for (int64_t index = 0; index < routeCount; ++index) { + PublishTxReadyInData(workspaceGM + txDataOffset + + (firstRoute + index) * routeStride, readyValue, perfTrace, perfStats); + } +#else + (void)txDataOffset; + (void)routeStride; + // In shared mode the one published line releases the whole lane-local batch. PackRoutes + // reaches this call only after every route represented by that line has completed MTE3. +#if TILEXR_EP_URMA_TX_READY_SHARED_FLAG && TILEXR_EP_URMA_TX_READY_BATCH_SIZE > 1 + const int64_t publishedRouteCount = 1; +#else + const int64_t publishedRouteCount = routeCount; +#endif + GM_ADDR firstLineAddr = workspaceGM + txReadyOffset + + firstRoute * TileXREp::kEpUrmaCombineCacheLineBytes; +#if TILEXR_EP_URMA_CACHELESS + AscendC::PipeBarrier(); + for (int64_t index = 0; index < publishedRouteCount; ++index) { + AscendC::WriteGmByPassDCache( + reinterpret_cast<__gm__ uint64_t *>(firstLineAddr + + index * TileXREp::kEpUrmaCombineCacheLineBytes), + readyValue); + } + AscendC::DataSyncBarrier(); +#else + for (int64_t index = 0; index < publishedRouteCount; ++index) { + reinterpret_cast<__gm__ uint64_t *>(firstLineAddr + + index * TileXREp::kEpUrmaCombineCacheLineBytes)[0] = readyValue; + } + AscendC::PipeBarrier(); + const uint64_t dcciStart = ProfileBegin(perfTrace); + TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t *>(firstLineAddr), + publishedRouteCount * TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::CONTROL_OTHER, dcciStart); + AscendC::PipeBarrier(); +#endif +#endif +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void FinishAndPublishPackBatch( + GM_ADDR workspaceGM, int64_t txReadyOffset, + int64_t txDataOffset, int64_t routeStride, int64_t firstRoute, int64_t routeCount, + event_t eventId, uint64_t readyValue, uint64_t firstReadyStart, bool &firstReadyPublished, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + GM_ADDR txRouteAddr = workspaceGM + txDataOffset + firstRoute * routeStride; + FinishPackCopy(txRouteAddr, routeStride, eventId, perfTrace, perfStats); + + const uint64_t publishStart = ProfileBegin(perfTrace); + PublishTxReadyBatch(workspaceGM, txReadyOffset, txDataOffset, routeStride, + firstRoute, routeCount, readyValue, perfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, publishStart); +#if TILEXR_EP_URMA_TX_READY_SHARED_FLAG && TILEXR_EP_URMA_TX_READY_BATCH_SIZE > 1 + constexpr uint64_t publishedLineCount = 1; +#else + const uint64_t publishedLineCount = static_cast(routeCount); +#endif + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 1, publishedLineCount); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 2, + static_cast(routeCount)); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 3, + publishedLineCount * TileXREp::kEpUrmaCombineCacheLineBytes); + if (!firstReadyPublished) { + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, firstReadyStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, 0, + static_cast(firstRoute)); + firstReadyPublished = true; + } +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION int64_t TxReadyPollRoute( + int64_t route, int64_t laneBegin) +{ +#if TILEXR_EP_URMA_TX_READY_SHARED_FLAG && TILEXR_EP_URMA_TX_READY_BATCH_SIZE > 1 + // Batches restart at each Pack lane boundary; global route/TXB rounding is incorrect when + // selfSendCnt is not divisible by the number of Pack lanes. + return laneBegin + ((route - laneBegin) / + static_cast(TileXREp::kEpUrmaCombineTxReadyBatchSize)) * + static_cast(TileXREp::kEpUrmaCombineTxReadyBatchSize); +#else + (void)laneBegin; + return route; +#endif +} + +__aicore__ __attribute__((always_inline)) inline void StartPackQuantization( + AscendC::LocalTensor &routeInput, + AscendC::LocalTensor &routePacked, + AscendC::LocalTensor &quantFloat, +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + AscendC::LocalTensor &absHalf, + AscendC::LocalTensor &reduceOut, + AscendC::LocalTensor &reduceTmp, +#else + AscendC::LocalTensor &absFloat, + AscendC::LocalTensor &reduceOut, + AscendC::LocalTensor &reduceTmp, +#endif + AscendC::LocalTensor &logical, + int64_t h, int64_t logicalBytes, int64_t routeStride, AscendC::TEventID quantizeEvent) +{ +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + (void)routePacked; + (void)logical; + (void)logicalBytes; + (void)routeStride; + AscendC::Abs(absHalf, routeInput, static_cast(h)); + AscendC::PipeBarrier(); + AscendC::ReduceMax(reduceOut, absHalf, reduceTmp, static_cast(h), false); +#if TILEXR_EP_URMA_QDC_VERSION != 3 + AscendC::Cast(quantFloat, routeInput, AscendC::RoundMode::CAST_NONE, + static_cast(h)); +#else + (void)quantFloat; +#endif +#else + AscendC::Duplicate(logical, static_cast(0), static_cast(logicalBytes)); + AscendC::Duplicate(routePacked, TileXR::DATA_AS_FLAG_READY_VALUE, + static_cast(routeStride / sizeof(float))); + AscendC::Cast(quantFloat, routeInput, AscendC::RoundMode::CAST_NONE, + static_cast(h)); + AscendC::PipeBarrier(); + AscendC::Abs(absFloat, quantFloat, static_cast(h)); + AscendC::PipeBarrier(); + AscendC::ReduceMax(reduceOut, absFloat, reduceTmp, static_cast(h), false); +#endif + AscendC::SetFlag(quantizeEvent); +} + +__aicore__ __attribute__((always_inline)) inline void FinishPackQuantization( + AscendC::LocalTensor &routeInput, + AscendC::LocalTensor &routePacked, + AscendC::LocalTensor &quantFloat, +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + AscendC::LocalTensor &reduceOut, +#else + AscendC::LocalTensor &reduceOut, +#endif + AscendC::LocalTensor &logical, + AscendC::LocalTensor &logicalFloat, + AscendC::LocalTensor &logicalInt, + int64_t h, int64_t blockCount, uint8_t repeats, + AscendC::TEventID quantizeEvent, AscendC::TEventID scalarToVectorEvent) +{ + AscendC::WaitFlag(quantizeEvent); +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + const float maxAbs = static_cast(reduceOut(0)); +#else + const float maxAbs = reduceOut(0); +#endif + const float scale = maxAbs > 0.0f ? maxAbs / 127.0f : 1.0f; + const float inverseScale = maxAbs > 0.0f ? 127.0f / maxAbs : 1.0f; + logicalFloat.SetValue(0, scale); + logicalInt.SetValue(1, static_cast(TileXREp::kEpUrmaCombineQuantModeInt8PerRoute)); +#if TILEXR_EP_URMA_TX_READY_IN_DATA + constexpr uint32_t readyWordIndex = static_cast( + TileXREp::kEpUrmaCombineTxReadyHeaderOffset / sizeof(int32_t)); + logicalInt.SetValue(readyWordIndex, 0); + logicalInt.SetValue(readyWordIndex + 1, 0); +#endif +#if TILEXR_EP_URMA_QDC_VERSION == 3 + const bool halfScaleSafe = maxAbs >= 0.0f && maxAbs <= kMaxFiniteHalf && + inverseScale > 0.0f && inverseScale <= kMaxFiniteHalf; + if (halfScaleSafe) { + AscendC::Muls(routeInput, routeInput, static_cast(inverseScale), + static_cast(h)); + AscendC::PipeBarrier(); + } else { + AscendC::Cast(quantFloat, routeInput, AscendC::RoundMode::CAST_NONE, + static_cast(h)); + AscendC::PipeBarrier(); + AscendC::Muls(quantFloat, quantFloat, inverseScale, static_cast(h)); + AscendC::PipeBarrier(); + AscendC::Cast(routeInput, quantFloat, AscendC::RoundMode::CAST_RINT, + static_cast(h)); + AscendC::PipeBarrier(); + } +#else + AscendC::Muls(quantFloat, quantFloat, inverseScale, static_cast(h)); + AscendC::PipeBarrier(); + AscendC::Cast(routeInput, quantFloat, AscendC::RoundMode::CAST_RINT, + static_cast(h)); + AscendC::PipeBarrier(); +#endif + AscendC::Cast(logical[TileXREp::kEpUrmaCombineQuantHeaderBytes], routeInput, + AscendC::RoundMode::CAST_RINT, static_cast(h)); + AscendC::PipeBarrier(); + + AscendC::SetFlag(scalarToVectorEvent); + AscendC::WaitFlag(scalarToVectorEvent); + AscendC::Copy(routePacked, logicalFloat, static_cast(64), repeats, {1, 1, 16, 15}); + AscendC::Copy(routePacked[64], logicalFloat[64], static_cast(56), repeats, {1, 1, 16, 15}); + AscendC::PipeBarrier(); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void PackRoutes( + GM_ADDR expertOutGM, GM_ADDR workspaceGM, int64_t selfSendCnt, int64_t h, + int64_t magic, int64_t blockCount, int64_t routeStride, int64_t txReadyOffset, int64_t txDataOffset, + int64_t laneId, AscendC::TPipe &pipe, GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t firstReadyStart = ProfileBegin(perfTrace); + const int64_t inputBytes = AlignUpInt64(h * static_cast(sizeof(half)), kCursorAlignment); + const int64_t floatBytes = AlignUpInt64(h * static_cast(sizeof(float)), kCursorAlignment); + const int64_t logicalBytes = blockCount * TileXREp::kEpUrmaCombinePayloadBytes; + + PackInputQueue inputQueue; + AscendC::TBuf quantFloatBuf; +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + AscendC::TBuf absHalfBuf; +#else + AscendC::TBuf absFloatBuf; +#endif + AscendC::TBuf reduceOutBuf; + AscendC::TBuf reduceTmpBuf; + AscendC::TBuf logicalBuf; + AscendC::TBuf packedBuf; + pipe.InitBuffer(inputQueue, static_cast(kPipelineBufferCount), static_cast(inputBytes)); + pipe.InitBuffer(quantFloatBuf, static_cast(floatBytes)); +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + pipe.InitBuffer(absHalfBuf, static_cast(inputBytes)); + pipe.InitBuffer(reduceOutBuf, kCursorAlignment); + pipe.InitBuffer(reduceTmpBuf, static_cast(inputBytes)); +#else + pipe.InitBuffer(absFloatBuf, static_cast(floatBytes)); + pipe.InitBuffer(reduceOutBuf, kCursorAlignment); + pipe.InitBuffer(reduceTmpBuf, static_cast(floatBytes)); +#endif + pipe.InitBuffer(logicalBuf, static_cast(logicalBytes)); + pipe.InitBuffer(packedBuf, static_cast(kPipelineBufferCount * routeStride)); + AscendC::LocalTensor packed0 = packedBuf.GetWithOffset( + static_cast(routeStride / sizeof(float)), 0); + AscendC::LocalTensor packed1 = packedBuf.GetWithOffset( + static_cast(routeStride / sizeof(float)), static_cast(routeStride)); + AscendC::LocalTensor quantFloat = quantFloatBuf.Get(); +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + AscendC::LocalTensor absHalf = absHalfBuf.Get(); + AscendC::LocalTensor reduceOutHalf = reduceOutBuf.Get(); + AscendC::LocalTensor reduceTmpHalf = reduceTmpBuf.Get(); +#else + AscendC::LocalTensor absFloat = absFloatBuf.Get(); + AscendC::LocalTensor reduceOut = reduceOutBuf.Get(); + AscendC::LocalTensor reduceTmp = reduceTmpBuf.Get(); +#endif + AscendC::LocalTensor logical = logicalBuf.Get(); + AscendC::LocalTensor logicalFloat = logical.template ReinterpretCast(); + AscendC::LocalTensor logicalInt = logical.template ReinterpretCast(); + AscendC::GlobalTensor expertOut; + expertOut.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(expertOutGM), selfSendCnt * h); + + const int64_t begin = selfSendCnt * laneId / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t end = selfSendCnt * (laneId + 1) / TileXREp::kEpUrmaCombinePackLaneCount; + const uint64_t readyValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineTxRouteReady); + const uint32_t inputAlignedElements = static_cast(inputBytes / sizeof(half)); + const uint8_t inputPadding = static_cast(inputAlignedElements - static_cast(h)); + const uint8_t repeats = static_cast(blockCount); + const AscendC::TEventID quantizeEvent = pipe.FetchEventID(); + const AscendC::TEventID scalarToVectorEvent = pipe.FetchEventID(); + bool copyPending = false; + int64_t pendingRoute = 0; + int64_t pendingSlot = 0; +#if !TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH + int64_t readyBatchFirstRoute = 0; + int64_t readyBatchCount = 0; +#endif + bool firstReadyPublished = false; + +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + AscendC::Duplicate(logical, static_cast(0), static_cast(logicalBytes)); + AscendC::Duplicate(packed0, TileXR::DATA_AS_FLAG_READY_VALUE, + static_cast(routeStride / sizeof(float))); + AscendC::Duplicate(packed1, TileXR::DATA_AS_FLAG_READY_VALUE, + static_cast(routeStride / sizeof(float))); +#endif + + for (int64_t route = begin; route < end; ++route) { + const int64_t slot = (route - begin) % kPipelineBufferCount; + const event_t eventId = PipelineEvent(slot); + EnqueuePackInput(expertOut, route, h, inputPadding, inputQueue); + const uint64_t inputWaitStart = ProfileBegin(perfTrace); + AscendC::LocalTensor input = inputQueue.DeQue(); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_INPUT_WAIT, inputWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_INPUT_WAIT, 0, + static_cast(h * static_cast(sizeof(half)))); + + const uint64_t quantizeHeadStart = ProfileBegin(perfTrace); +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + if (slot == 0) { + StartPackQuantization(input, packed0, quantFloat, absHalf, reduceOutHalf, reduceTmpHalf, + logical, h, logicalBytes, routeStride, quantizeEvent); + } else { + StartPackQuantization(input, packed1, quantFloat, absHalf, reduceOutHalf, reduceTmpHalf, + logical, h, logicalBytes, routeStride, quantizeEvent); + } +#else + if (slot == 0) { + StartPackQuantization(input, packed0, quantFloat, absFloat, reduceOut, reduceTmp, + logical, h, logicalBytes, routeStride, quantizeEvent); + } else { + StartPackQuantization(input, packed1, quantFloat, absFloat, reduceOut, reduceTmp, + logical, h, logicalBytes, routeStride, quantizeEvent); + } +#endif +#if TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_QUANTIZE, quantizeHeadStart); + if (copyPending) { + FinishAndPublishPackBatch(workspaceGM, txReadyOffset, txDataOffset, routeStride, + pendingRoute, 1, PipelineEvent(pendingSlot), readyValue, + firstReadyStart, firstReadyPublished, perfTrace, perfStats); + copyPending = false; + } + const uint64_t quantizeTailStart = ProfileBegin(perfTrace); +#endif +#if TILEXR_EP_URMA_QDC_VERSION == 1 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + if (slot == 0) { + FinishPackQuantization(input, packed0, quantFloat, reduceOutHalf, logical, logicalFloat, + logicalInt, h, blockCount, repeats, quantizeEvent, scalarToVectorEvent); + } else { + FinishPackQuantization(input, packed1, quantFloat, reduceOutHalf, logical, logicalFloat, + logicalInt, h, blockCount, repeats, quantizeEvent, scalarToVectorEvent); + } +#else + if (slot == 0) { + FinishPackQuantization(input, packed0, quantFloat, reduceOut, logical, logicalFloat, + logicalInt, h, blockCount, repeats, quantizeEvent, scalarToVectorEvent); + } else { + FinishPackQuantization(input, packed1, quantFloat, reduceOut, logical, logicalFloat, + logicalInt, h, blockCount, repeats, quantizeEvent, scalarToVectorEvent); + } +#endif + inputQueue.FreeTensor(input); +#if TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_QUANTIZE, quantizeTailStart); +#else + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_QUANTIZE, quantizeHeadStart); +#endif + ProfileAux(perfTrace, perfStats, PerfStage::PACK_QUANTIZE, 0, static_cast(h)); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_QUANTIZE, 1, static_cast(blockCount)); + + GM_ADDR txRouteAddr = workspaceGM + txDataOffset + route * routeStride; + const uint64_t dataSubmitStart = ProfileBegin(perfTrace); + if (slot == 0) { + StartPackPublish(txRouteAddr, routeStride, packed0, eventId); + } else { + StartPackPublish(txRouteAddr, routeStride, packed1, eventId); + } + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TX_DATA_SUBMIT, dataSubmitStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_DATA_SUBMIT, 0, + static_cast(routeStride)); +#if !TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH + if (copyPending) { + FinishPackCopy(workspaceGM + txDataOffset + pendingRoute * routeStride, + routeStride, PipelineEvent(pendingSlot), perfTrace, perfStats); + if (readyBatchCount == 0) { + readyBatchFirstRoute = pendingRoute; + } + ++readyBatchCount; + if (readyBatchCount == static_cast(TileXREp::kEpUrmaCombineTxReadyBatchSize)) { + const uint64_t readyPublishStart = ProfileBegin(perfTrace); + PublishTxReadyBatch(workspaceGM, txReadyOffset, txDataOffset, routeStride, + readyBatchFirstRoute, readyBatchCount, readyValue, perfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, readyPublishStart); +#if TILEXR_EP_URMA_TX_READY_SHARED_FLAG && TILEXR_EP_URMA_TX_READY_BATCH_SIZE > 1 + constexpr uint64_t publishedLineCount = 1; +#else + const uint64_t publishedLineCount = static_cast(readyBatchCount); +#endif + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 1, + publishedLineCount); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 2, + static_cast(readyBatchCount)); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 3, + publishedLineCount * TileXREp::kEpUrmaCombineCacheLineBytes); + if (!firstReadyPublished) { + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, firstReadyStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, 0, + static_cast(readyBatchFirstRoute)); + firstReadyPublished = true; + } + readyBatchCount = 0; + } + } +#endif + copyPending = true; + pendingRoute = route; + pendingSlot = slot; +#if !TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH + if (route + 1 == end) { + FinishPackCopy(txRouteAddr, routeStride, eventId, perfTrace, perfStats); + if (readyBatchCount == 0) { + readyBatchFirstRoute = route; + } + ++readyBatchCount; + copyPending = false; + const uint64_t readyPublishStart = ProfileBegin(perfTrace); + PublishTxReadyBatch(workspaceGM, txReadyOffset, txDataOffset, routeStride, + readyBatchFirstRoute, readyBatchCount, readyValue, perfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, readyPublishStart); +#if TILEXR_EP_URMA_TX_READY_SHARED_FLAG && TILEXR_EP_URMA_TX_READY_BATCH_SIZE > 1 + constexpr uint64_t publishedLineCount = 1; +#else + const uint64_t publishedLineCount = static_cast(readyBatchCount); +#endif + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 1, + publishedLineCount); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 2, + static_cast(readyBatchCount)); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_TX_PUBLISH, 3, + publishedLineCount * TileXREp::kEpUrmaCombineCacheLineBytes); + if (!firstReadyPublished) { + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, firstReadyStart); + ProfileAux(perfTrace, perfStats, PerfStage::PACK_FIRST_TX_READY, 0, + static_cast(readyBatchFirstRoute)); + firstReadyPublished = true; + } + readyBatchCount = 0; + } +#endif + } +#if TILEXR_EP_URMA_TX_READY_EARLY_PUBLISH + if (copyPending) { + FinishAndPublishPackBatch(workspaceGM, txReadyOffset, txDataOffset, routeStride, + pendingRoute, 1, PipelineEvent(pendingSlot), readyValue, + firstReadyStart, firstReadyPublished, perfTrace, perfStats); + } +#endif + pipe.ReleaseEventID(quantizeEvent); + pipe.ReleaseEventID(scalarToVectorEvent); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void StartRouteReadyCopy( + GM_ADDR routeAddr, int64_t blockCount, AscendC::LocalTensor &flags, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ +#if !TILEXR_EP_URMA_CACHELESS + const uint64_t dcciStart = ProfileBegin(perfTrace); + for (int64_t block = 0; block < blockCount; ++block) { + TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t *>(routeAddr + + block * TileXREp::kEpUrmaCombineDataBlockBytes + TileXREp::kEpUrmaCombinePayloadBytes), + TileXREp::kEpUrmaCombineFlagBytes); + } + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::RX_FLAG_POLL, dcciStart); +#endif + AscendC::GlobalTensor src; + src.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( + routeAddr + TileXREp::kEpUrmaCombinePayloadBytes)); +#if TILEXR_EP_URMA_CACHELESS + SetNoCacheRead(src); + AscendC::DataCopyExtParams copyParams { + static_cast(blockCount), static_cast(sizeof(float)), + static_cast(TileXREp::kEpUrmaCombineDataBlockBytes - sizeof(float)), 0, 0}; + AscendC::DataCopyPadExtParams padParams { + true, 0, static_cast(TileXR::DATA_AS_FLAG_FLAG_FLOATS - 1U), +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR + TileXR::DATA_AS_FLAG_READY_VALUE}; +#else + 0.0f}; +#endif +#else + AscendC::DataCopyExtParams copyParams { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombineFlagBytes), + static_cast(TileXREp::kEpUrmaCombinePayloadBytes), 0, 0}; + AscendC::DataCopyPadExtParams padParams {false, 0, 0, 0}; +#endif + AscendC::DataCopyPad(flags, src, copyParams, padParams); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool CheckRouteReady( + GM_ADDR routeAddr, int64_t blockCount, AscendC::LocalTensor &flags, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + StartRouteReadyCopy(routeAddr, blockCount, flags, perfTrace, perfStats); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +#if TILEXR_EP_URMA_CACHELESS + for (int64_t block = 0; block < blockCount; ++block) { + if (flags.GetValue(static_cast( + block * TileXR::DATA_AS_FLAG_FLAG_FLOATS)) != TileXR::DATA_AS_FLAG_READY_VALUE) { + return false; + } + } +#else + const int64_t flagCount = blockCount * TileXR::DATA_AS_FLAG_FLAG_FLOATS; + for (int64_t flag = 0; flag < flagCount; ++flag) { + if (flags.GetValue(static_cast(flag)) != TileXR::DATA_AS_FLAG_READY_VALUE) { + return false; + } + } +#endif + return true; +} + +#if TILEXR_EP_URMA_RX_READY_BATCH_MTE2 +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint32_t CheckRoutesReadyBatch( + GM_ADDR workspaceGM, int64_t rxWindowOffset, + int64_t token, int64_t topK, int64_t blockCount, int64_t routeStride, uint32_t pendingMask, + AscendC::LocalTensor &flags, +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR + AscendC::LocalTensor &delta, AscendC::LocalTensor &reduceOut, + AscendC::LocalTensor &reduceTmp, +#endif + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint32_t flagCount = static_cast( + blockCount * TileXR::DATA_AS_FLAG_FLAG_FLOATS); + uint64_t pendingRouteCount = 0; + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const uint32_t routeBit = 1U << static_cast(topKId); + if ((pendingMask & routeBit) == 0) { + continue; + } + const int64_t routeIndex = token * topK + topKId; + GM_ADDR routeAddr = workspaceGM + rxWindowOffset + routeIndex * routeStride; + AscendC::LocalTensor routeFlags = flags[static_cast(topKId) * flagCount]; + StartRouteReadyCopy(routeAddr, blockCount, routeFlags, perfTrace, perfStats); + ++pendingRouteCount; + } + if (pendingRouteCount == 0) { + return 0; + } + +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR + AscendC::SetFlag(EVENT_ID0); + const uint64_t mte2WaitStart = ProfileBegin(perfTrace); + AscendC::WaitFlag(EVENT_ID0); +#else + AscendC::SetFlag(EVENT_ID0); + const uint64_t mte2WaitStart = ProfileBegin(perfTrace); + AscendC::WaitFlag(EVENT_ID0); +#endif + ProfileEnd(perfTrace, perfStats, PerfStage::RX_READY_MTE2_WAIT, mte2WaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_READY_MTE2_WAIT, 0, pendingRouteCount); + ProfileAux(perfTrace, perfStats, PerfStage::RX_READY_MTE2_WAIT, 1, + pendingRouteCount * static_cast(blockCount)); + ProfileAux(perfTrace, perfStats, PerfStage::RX_READY_MTE2_WAIT, 2, + pendingRouteCount * static_cast(blockCount) * sizeof(float)); + + uint32_t readyMask = 0; +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR + const uint64_t vectorStart = ProfileBegin(perfTrace); + constexpr uint32_t reduceResultStride = kCursorAlignment / sizeof(float); + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const uint32_t routeBit = 1U << static_cast(topKId); + if ((pendingMask & routeBit) == 0) { + continue; + } + const uint32_t routeOffset = static_cast(topKId) * flagCount; + AscendC::LocalTensor routeFlags = flags[routeOffset]; + AscendC::LocalTensor routeDelta = delta[routeOffset]; + AscendC::LocalTensor routeReduceOut = + reduceOut[static_cast(topKId) * reduceResultStride]; + AscendC::Adds(routeDelta, routeFlags, -TileXR::DATA_AS_FLAG_READY_VALUE, flagCount); + AscendC::PipeBarrier(); + AscendC::Abs(routeFlags, routeDelta, flagCount); + AscendC::PipeBarrier(); + AscendC::ReduceMax( + routeReduceOut, routeFlags, reduceTmp, static_cast(flagCount), false); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_READY_VECTOR, vectorStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_READY_VECTOR, 0, pendingRouteCount); + ProfileAux(perfTrace, perfStats, PerfStage::RX_READY_VECTOR, 1, + pendingRouteCount * static_cast(flagCount)); + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const uint32_t routeBit = 1U << static_cast(topKId); + if ((pendingMask & routeBit) != 0 && + reduceOut.GetValue(static_cast(topKId) * reduceResultStride) == 0.0f) { + readyMask |= routeBit; + } + } +#else + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const uint32_t routeBit = 1U << static_cast(topKId); + if ((pendingMask & routeBit) == 0) { + continue; + } + const uint32_t routeOffset = static_cast(topKId) * flagCount; + bool routeReady = true; +#if TILEXR_EP_URMA_CACHELESS + for (int64_t block = 0; block < blockCount; ++block) { + if (flags.GetValue(routeOffset + static_cast( + block * TileXR::DATA_AS_FLAG_FLAG_FLOATS)) != + TileXR::DATA_AS_FLAG_READY_VALUE) { + routeReady = false; + break; + } + } +#else + for (uint32_t flag = 0; flag < flagCount; ++flag) { + if (flags.GetValue(routeOffset + flag) != TileXR::DATA_AS_FLAG_READY_VALUE) { + routeReady = false; + break; + } + } +#endif + if (routeReady) { + readyMask |= routeBit; + } + } +#endif + return readyMask; +} +#endif + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void StartUnpackRoute( + GM_ADDR routeAddr, int64_t blockCount, int64_t routeStride, + AscendC::LocalTensor &logical, event_t eventId, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ +#if !TILEXR_EP_URMA_CACHELESS + const uint64_t dcciStart = ProfileBegin(perfTrace); + TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t *>(routeAddr), routeStride); + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::RX_DATA, dcciStart); +#endif + AscendC::GlobalTensor src; + src.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(routeAddr)); + SetNoCacheRead(src); + AscendC::DataCopyExtParams copyParams { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombinePayloadBytes), + static_cast(TileXREp::kEpUrmaCombineFlagBytes), 0, 0}; + AscendC::DataCopyPadExtParams padParams {false, 0, 0, 0}; + AscendC::DataCopyPad(logical, src, copyParams, padParams); + AscendC::SetFlag(eventId); + AscendC::SetFlag(eventId); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void WaitUnpackRoute(event_t eventId) +{ + AscendC::WaitFlag(eventId); + AscendC::WaitFlag(eventId); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void StartClearRouteFlags( + GM_ADDR routeAddr, int64_t blockCount, AscendC::LocalTensor &zeroFlags) +{ + AscendC::GlobalTensor dst; + dst.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( + routeAddr + TileXREp::kEpUrmaCombinePayloadBytes)); + SetNoCacheWrite(dst); + AscendC::DataCopyExtParams copyParams { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombineFlagBytes), 0, + static_cast(TileXREp::kEpUrmaCombinePayloadBytes), 0}; + AscendC::DataCopyPad(dst, zeroFlags, copyParams); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void CleanTokenRouteFlags( + GM_ADDR firstRouteAddr, int64_t topK, int64_t blockCount, int64_t routeStride, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ +#if !TILEXR_EP_URMA_CACHELESS + const uint64_t dcciStart = ProfileBegin(perfTrace); + for (int64_t topKId = 0; topKId < topK; ++topKId) { + GM_ADDR routeAddr = firstRouteAddr + topKId * routeStride; + for (int64_t block = 0; block < blockCount; ++block) { + TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t *>(routeAddr + + block * TileXREp::kEpUrmaCombineDataBlockBytes + TileXREp::kEpUrmaCombinePayloadBytes), + TileXREp::kEpUrmaCombineFlagBytes); + } + } + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::RX_DATA, dcciStart); +#else + (void)firstRouteAddr; + (void)topK; + (void)blockCount; + (void)routeStride; + (void)perfTrace; + (void)perfStats; +#endif +} + +#if TILEXR_EP_URMA_QDC_VERSION == 0 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 +__aicore__ TILEXR_EP_LOCAL_FUNCTION void FinishTokenRouteFlagClears( + GM_ADDR firstRouteAddr, int64_t topK, int64_t blockCount, int64_t routeStride, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + CleanTokenRouteFlags(firstRouteAddr, topK, blockCount, routeStride, perfTrace, perfStats); +} +#endif + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void ReceiveTokens( + GM_ADDR topKWeightsGM, GM_ADDR yOutGM, GM_ADDR workspaceGM, int64_t bs, + int64_t h, int64_t topK, int64_t magic, int64_t rxWindowOffset, int64_t blockCount, int64_t routeStride, + int64_t rxLaneDoneOffset, int64_t errorStatusOffset, int64_t laneId, AscendC::TPipe &pipe, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const int64_t logicalBytes = blockCount * TileXREp::kEpUrmaCombinePayloadBytes; + const int64_t flagBytes = blockCount * TileXREp::kEpUrmaCombineFlagBytes; + const int64_t accumBytes = AlignUpInt64(h * static_cast(sizeof(float)), kCursorAlignment); + const int64_t outputBytes = AlignUpInt64(h * static_cast(sizeof(half)), kCursorAlignment); + const int64_t dequantBytes = AlignUpInt64(h * static_cast(sizeof(float)), kCursorAlignment); + const int64_t weightBytes = AlignUpInt64( + topK * static_cast(sizeof(float)), kCursorAlignment); + + AscendC::TBuf logicalBuf; + AscendC::TBuf flagBuf; + AscendC::TBuf clearFlagBuf; + AscendC::TBuf accumBuf; + AscendC::TBuf outputBuf; + AscendC::TBuf dequantBuf; + AscendC::TBuf weightBuf; +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR + AscendC::TBuf readyDeltaBuf; + AscendC::TBuf readyReduceOutBuf; + AscendC::TBuf readyReduceTmpBuf; +#endif + pipe.InitBuffer(logicalBuf, static_cast(kPipelineBufferCount * logicalBytes)); +#if TILEXR_EP_URMA_RX_READY_BATCH_MTE2 + pipe.InitBuffer(flagBuf, static_cast(topK * flagBytes)); +#else + pipe.InitBuffer(flagBuf, static_cast(flagBytes)); +#endif + pipe.InitBuffer(clearFlagBuf, static_cast(flagBytes)); + pipe.InitBuffer(accumBuf, static_cast(accumBytes)); + pipe.InitBuffer(outputBuf, static_cast(outputBytes)); + pipe.InitBuffer(dequantBuf, static_cast(dequantBytes)); + pipe.InitBuffer(weightBuf, static_cast(weightBytes)); +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR + pipe.InitBuffer(readyDeltaBuf, static_cast(topK * flagBytes)); + pipe.InitBuffer(readyReduceOutBuf, static_cast(topK * kCursorAlignment)); + pipe.InitBuffer(readyReduceTmpBuf, static_cast(flagBytes)); +#endif + AscendC::LocalTensor logical0 = logicalBuf.GetWithOffset( + static_cast(logicalBytes), 0); + AscendC::LocalTensor logical1 = logicalBuf.GetWithOffset( + static_cast(logicalBytes), static_cast(logicalBytes)); + AscendC::LocalTensor logicalFloat0 = logical0.template ReinterpretCast(); + AscendC::LocalTensor logicalFloat1 = logical1.template ReinterpretCast(); + AscendC::LocalTensor logicalInt0 = logical0.template ReinterpretCast(); + AscendC::LocalTensor logicalInt1 = logical1.template ReinterpretCast(); + AscendC::LocalTensor logicalInt80 = logical0.template ReinterpretCast(); + AscendC::LocalTensor logicalInt81 = logical1.template ReinterpretCast(); + AscendC::LocalTensor flags = flagBuf.Get(); + AscendC::LocalTensor zeroFlags = clearFlagBuf.Get(); + AscendC::LocalTensor accum = accumBuf.Get(); + AscendC::LocalTensor output = outputBuf.Get(); + AscendC::LocalTensor dequant = dequantBuf.Get(); + AscendC::LocalTensor weight = weightBuf.Get(); +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR + AscendC::LocalTensor readyDelta = readyDeltaBuf.Get(); + AscendC::LocalTensor readyReduceOut = readyReduceOutBuf.Get(); + AscendC::LocalTensor readyReduceTmp = readyReduceTmpBuf.Get(); +#endif + AscendC::GlobalTensor yOut; + yOut.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(yOutGM), bs * h); + + const uint32_t flagCount = static_cast( + blockCount * TileXR::DATA_AS_FLAG_FLAG_FLOATS); + AscendC::Duplicate(zeroFlags, 0.0f, flagCount); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + const int64_t tokenBegin = bs * laneId / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t tokenEnd = bs * (laneId + 1) / TileXREp::kEpUrmaCombinePackLaneCount; + const bool profileFine = TileXR::TileXRPerfTraceEnabled(perfTrace); + int64_t tokensRemaining = tokenEnd - tokenBegin; +#if TILEXR_EP_URMA_RX_SCHEDULER == 0 + int64_t nextSequentialToken = tokenBegin; +#else + int64_t activeToken[kRxTokenScheduleWindow] = {}; + bool tokenActive[kRxTokenScheduleWindow] = {}; + uint32_t routeReadyMask[kRxTokenScheduleWindow] = {}; + int64_t nextToken = tokenBegin; + for (int64_t slot = 0; slot < kRxTokenScheduleWindow && nextToken < tokenEnd; ++slot) { + activeToken[slot] = nextToken++; + routeReadyMask[slot] = 0; + tokenActive[slot] = true; + } + int64_t nextCandidate = 0; + const uint32_t allRoutesReadyMask = (1U << static_cast(topK)) - 1U; +#endif + while (tokensRemaining > 0) { +#if TILEXR_EP_URMA_RX_SCHEDULER == 0 + int64_t token = nextSequentialToken; +#else + int64_t token = 0; + int64_t selectedSlot = -1; +#endif + uint64_t bypassedTokens = 0; + const uint64_t flagPollStart = ProfileBegin(perfTrace); + uint64_t pollPasses = 0; + uint64_t routeChecks = 0; + uint64_t readyMisses = 0; +#if TILEXR_EP_URMA_RX_SCHEDULER > 0 + while (selectedSlot < 0) { + for (int64_t offset = 0; offset < kRxTokenScheduleWindow; ++offset) { + int64_t candidate = nextCandidate + offset; + if (candidate >= kRxTokenScheduleWindow) { + candidate -= kRxTokenScheduleWindow; + } + if (!tokenActive[candidate]) { + continue; + } + if (profileFine) { + ++pollPasses; + } + const int64_t candidateToken = activeToken[candidate]; + bool candidateReady = true; + uint32_t readyMask = routeReadyMask[candidate]; +#if TILEXR_EP_URMA_RX_READY_BATCH_MTE2 + uint32_t pendingMask = allRoutesReadyMask; + if constexpr (TileXREp::kEpUrmaCombineRxStickyReady) { + pendingMask &= ~readyMask; + } + if (profileFine) { + for (int64_t topKId = 0; topKId < topK; ++topKId) { + if ((pendingMask & (1U << static_cast(topKId))) != 0) { + ++routeChecks; + } + } + } + const uint32_t observedReadyMask = CheckRoutesReadyBatch( + workspaceGM, rxWindowOffset, candidateToken, topK, blockCount, routeStride, + pendingMask, flags, +#if TILEXR_EP_URMA_RX_READY_BATCH_VECTOR + readyDelta, readyReduceOut, readyReduceTmp, +#endif + perfTrace, perfStats); + if (profileFine) { + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const uint32_t routeBit = 1U << static_cast(topKId); + if ((pendingMask & routeBit) != 0 && + (observedReadyMask & routeBit) == 0) { + ++readyMisses; + } + } + } + if constexpr (TileXREp::kEpUrmaCombineRxStickyReady) { + readyMask |= observedReadyMask; + } else { + readyMask = observedReadyMask; + } + candidateReady = readyMask == allRoutesReadyMask; +#else + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const uint32_t routeBit = 1U << static_cast(topKId); + if constexpr (TileXREp::kEpUrmaCombineRxStickyReady) { + if ((readyMask & routeBit) != 0) { + continue; + } + } + const int64_t routeIndex = candidateToken * topK + topKId; + GM_ADDR routeAddr = workspaceGM + rxWindowOffset + routeIndex * routeStride; + if (profileFine) { + ++routeChecks; + } + if (CheckRouteReady(routeAddr, blockCount, flags, perfTrace, perfStats)) { + if constexpr (TileXREp::kEpUrmaCombineRxStickyReady) { + readyMask |= routeBit; + } + } else { + if (profileFine) { + ++readyMisses; + } + candidateReady = false; + if constexpr (!TileXREp::kEpUrmaCombineRxStickyReady) { + break; + } + } + } +#endif + if constexpr (TileXREp::kEpUrmaCombineRxStickyReady) { + routeReadyMask[candidate] = readyMask; + candidateReady = readyMask == allRoutesReadyMask; + } + if (candidateReady) { + selectedSlot = candidate; + break; + } + } + } + token = activeToken[selectedSlot]; + for (int64_t slot = 0; slot < kRxTokenScheduleWindow; ++slot) { + if (tokenActive[slot] && activeToken[slot] < token) { + ++bypassedTokens; + } + } + nextCandidate = selectedSlot + 1; + if (nextCandidate >= kRxTokenScheduleWindow) { + nextCandidate = 0; + } +#else + ++nextSequentialToken; + bool tokenReady = false; + while (!tokenReady) { + if (profileFine) { + ++pollPasses; + } + tokenReady = true; + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const int64_t routeIndex = token * topK + topKId; + GM_ADDR routeAddr = workspaceGM + rxWindowOffset + routeIndex * routeStride; + if (profileFine) { + ++routeChecks; + } + if (!CheckRouteReady(routeAddr, blockCount, flags, perfTrace, perfStats)) { + if (profileFine) { + ++readyMisses; + } + tokenReady = false; + break; + } + } + } +#endif + CachelessAcquireBarrier(); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, flagPollStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, 0, pollPasses); + ProfileAux(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, 1, routeChecks); + ProfileAux(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, 2, readyMisses); + ProfileAux(perfTrace, perfStats, PerfStage::RX_FLAG_POLL_WAIT, 3, + routeChecks * static_cast(blockCount)); + + LoadTopKWeights(topKWeightsGM, token * topK, topK, weight); + if (topK > 0) { + GM_ADDR firstRouteAddr = workspaceGM + rxWindowOffset + token * topK * routeStride; + StartUnpackRoute(firstRouteAddr, blockCount, routeStride, logical0, PipelineEvent(0), + perfTrace, perfStats); + } + for (int64_t topKId = 0; topKId < topK; ++topKId) { + const int64_t slot = topKId % kPipelineBufferCount; + const event_t eventId = PipelineEvent(slot); + AscendC::LocalTensor &routeLogicalFloat = slot == 0 ? logicalFloat0 : logicalFloat1; + AscendC::LocalTensor &routeLogicalInt = slot == 0 ? logicalInt0 : logicalInt1; + AscendC::LocalTensor &routeLogicalInt8 = slot == 0 ? logicalInt80 : logicalInt81; + const int64_t routeIndex = token * topK + topKId; + GM_ADDR routeAddr = workspaceGM + rxWindowOffset + routeIndex * routeStride; + const uint64_t unpackWaitStart = ProfileBegin(perfTrace); + WaitUnpackRoute(eventId); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_UNPACK_WAIT, unpackWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_UNPACK_WAIT, 0, + static_cast(blockCount * TileXREp::kEpUrmaCombinePayloadBytes)); + + const uint64_t dequantStart = ProfileBegin(perfTrace); + const float scale = routeLogicalFloat.GetValue(0); + const int32_t quantMode = routeLogicalInt.GetValue(1); + if (!(scale > 0.0f) || quantMode != TileXREp::kEpUrmaCombineQuantModeInt8PerRoute) { + StoreError(workspaceGM, errorStatusOffset, magic, + TileXREp::kEpUrmaCombineStatusInvalidQuantHeader, perfTrace, perfStats); + } + const float effectiveScale = scale > 0.0f ? scale : 1.0f; + const float topKWeight = weight.GetValue(static_cast(topKId)); + AscendC::Cast(output, + routeLogicalInt8[TileXREp::kEpUrmaCombineQuantHeaderBytes], AscendC::RoundMode::CAST_NONE, + static_cast(h)); + AscendC::PipeBarrier(); + if (topKId + 1 < topK) { + const int64_t nextSlot = (slot + 1) % kPipelineBufferCount; + GM_ADDR nextRouteAddr = routeAddr + routeStride; + AscendC::LocalTensor &nextLogical = nextSlot == 0 ? logical0 : logical1; + StartUnpackRoute(nextRouteAddr, blockCount, routeStride, + nextLogical, PipelineEvent(nextSlot), perfTrace, perfStats); + } + AscendC::Cast(dequant, output, AscendC::RoundMode::CAST_NONE, + static_cast(h)); + AscendC::PipeBarrier(); + if (topKId == 0) { + AscendC::Muls(accum, dequant, effectiveScale * topKWeight, + static_cast(h)); + } else { + AscendC::Axpy(accum, dequant, effectiveScale * topKWeight, + static_cast(h)); + } + AscendC::PipeBarrier(); + StartClearRouteFlags(routeAddr, blockCount, zeroFlags); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, dequantStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, 0, + static_cast(h)); + ProfileAux(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, 1, + static_cast(blockCount)); + } + +#if TILEXR_EP_URMA_QDC_VERSION == 0 || TILEXR_EP_URMA_QDC_VERSION == 2 || \ + TILEXR_EP_URMA_QDC_VERSION == 3 + const uint64_t clearFenceStart = ProfileBegin(perfTrace); + GM_ADDR firstRouteAddr = workspaceGM + rxWindowOffset + token * topK * routeStride; + FinishTokenRouteFlagClears(firstRouteAddr, topK, blockCount, routeStride, perfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, clearFenceStart); +#endif + + const uint64_t outputStart = ProfileBegin(perfTrace); + AscendC::Cast(output, accum, AscendC::RoundMode::CAST_ROUND, + static_cast(h)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopyExtParams outputParams { + 1, static_cast(h * static_cast(sizeof(half))), 0, 0, 0}; + AscendC::DataCopyPad(yOut[token * h], output, outputParams); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + ProfileEnd(perfTrace, perfStats, PerfStage::RX_OUTPUT, outputStart); + ProfileAux(perfTrace, perfStats, PerfStage::RX_OUTPUT, 0, + static_cast(h * static_cast(sizeof(half)))); +#if TILEXR_EP_URMA_QDC_VERSION == 1 + ProfileAux(perfTrace, perfStats, PerfStage::RX_UNPACK_DEQUANT_CLEAR, 2, + static_cast(topK * blockCount)); + GM_ADDR firstRouteAddr = workspaceGM + rxWindowOffset + token * topK * routeStride; + CleanTokenRouteFlags(firstRouteAddr, topK, blockCount, routeStride, perfTrace, perfStats); +#endif +#if TILEXR_EP_URMA_RX_SCHEDULER > 0 + ProfileAux(perfTrace, perfStats, PerfStage::RX_OUTPUT, 1, bypassedTokens); + ProfileAux(perfTrace, perfStats, PerfStage::RX_OUTPUT, 2, bypassedTokens == 0 ? 0U : 1U); +#endif + + --tokensRemaining; +#if TILEXR_EP_URMA_RX_SCHEDULER > 0 + routeReadyMask[selectedSlot] = 0; + if (nextToken < tokenEnd) { + activeToken[selectedSlot] = nextToken++; + tokenActive[selectedSlot] = true; + } else { + tokenActive[selectedSlot] = false; + } +#endif + } + + StoreControlValue(workspaceGM + rxLaneDoneOffset + + laneId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxLaneDone), perfTrace, perfStats); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void CopySelfRoute( + GM_ADDR srcAddr, GM_ADDR dstAddr, int64_t blockCount, AscendC::LocalTensor &payload, + AscendC::LocalTensor &readyFlags, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + AscendC::GlobalTensor src; + AscendC::GlobalTensor dst; + src.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(srcAddr)); + dst.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t *>(dstAddr)); +#if TILEXR_EP_URMA_CACHELESS + SetNoCacheRead(src); + SetNoCacheWrite(dst); + AscendC::DataCopyExtParams payloadIn { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombinePayloadBytes), + static_cast(TileXREp::kEpUrmaCombineFlagBytes), 0, 0}; + AscendC::DataCopyPadExtParams noPad {false, 0, 0, 0}; + AscendC::DataCopyPad(payload, src, payloadIn, noPad); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::DataCopyExtParams payloadOut { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombinePayloadBytes), + 0, static_cast(TileXREp::kEpUrmaCombineFlagBytes), 0}; + AscendC::DataCopyPad(dst, payload, payloadOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + AscendC::GlobalTensor dstFlags; + dstFlags.SetGlobalBuffer(reinterpret_cast<__gm__ float *>( + dstAddr + TileXREp::kEpUrmaCombinePayloadBytes)); + SetNoCacheWrite(dstFlags); + AscendC::DataCopyExtParams flagOut { + static_cast(blockCount), static_cast(TileXREp::kEpUrmaCombineFlagBytes), + 0, static_cast(TileXREp::kEpUrmaCombinePayloadBytes), 0}; + AscendC::DataCopyPad(dstFlags, readyFlags, flagOut); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); +#else + for (int64_t block = 0; block < blockCount; ++block) { + const int64_t offset = block * TileXREp::kEpUrmaCombineDataBlockBytes; + AscendC::DataCopy(payload, src[offset], + static_cast(TileXREp::kEpUrmaCombineDataBlockBytes)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy(dst[offset], payload, + static_cast(TileXREp::kEpUrmaCombineDataBlockBytes)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + const uint64_t dcciStart = ProfileBegin(perfTrace); + TileXR::UDMACleanCacheLines(reinterpret_cast<__gm__ uint8_t *>(dstAddr), + blockCount * TileXREp::kEpUrmaCombineDataBlockBytes); + ProfileDcciEnd(perfTrace, perfStats, DcciProfileCategory::TX_DATA, dcciStart); +#endif +} + +__aicore__ TILEXR_EP_SEND_FUNCTION bool SendRoutes(const __gm__ TileXR::CommArgs *args, GM_ADDR assistInfoGM, + GM_ADDR workspaceGM, int64_t selfSendCnt, int64_t bs, int64_t topK, int64_t magic, int64_t rxWindowOffset, + int64_t blockCount, int64_t routeStride, int64_t txReadyOffset, int64_t txDataOffset, + int64_t senderDoneOffset, int64_t errorStatusOffset, int64_t senderId, AscendC::TPipe &pipe, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const int64_t cursorBytes = AlignUpInt64( + TileXREp::kEpUrmaCombinePackLaneCount * static_cast(sizeof(uint32_t)), kCursorAlignment); + const int64_t usedPeerBytes = AlignUpInt64(args->rankSize, kCursorAlignment); + const int64_t selfCopyPayloadBytes = blockCount * TileXREp::kEpUrmaCombinePayloadBytes; + const int64_t selfCopyFlagBytes = blockCount * TileXREp::kEpUrmaCombineFlagBytes; + AscendC::TBuf cursorBuf; +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL || !TILEXR_EP_URMA_CACHELESS + AscendC::TBuf metaBuf; +#endif + AscendC::TBuf usedPeerBuf; + AscendC::TBuf selfCopyBuf; + pipe.InitBuffer(cursorBuf, static_cast(cursorBytes)); +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL + pipe.InitBuffer(metaBuf, kRouteMetaPrefetchCapacityBytes); +#elif !TILEXR_EP_URMA_CACHELESS + pipe.InitBuffer(metaBuf, kRouteMetaBufferBytes); +#endif + pipe.InitBuffer(usedPeerBuf, static_cast(usedPeerBytes)); + pipe.InitBuffer(selfCopyBuf, + static_cast(selfCopyPayloadBytes + selfCopyFlagBytes)); + AscendC::LocalTensor cursor = cursorBuf.Get(); +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL || !TILEXR_EP_URMA_CACHELESS + AscendC::LocalTensor meta = metaBuf.Get(); +#endif + AscendC::LocalTensor usedPeer = usedPeerBuf.Get(); + AscendC::LocalTensor selfCopyPayload = selfCopyBuf.GetWithOffset( + static_cast(selfCopyPayloadBytes), 0); + AscendC::LocalTensor selfCopyReady = selfCopyBuf.GetWithOffset( + static_cast(selfCopyFlagBytes / sizeof(float)), + static_cast(selfCopyPayloadBytes)); +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL + const bool useRouteMetaPrefetch = selfSendCnt > 0 && selfSendCnt <= + static_cast(kRouteMetaPrefetchCapacityBytes / sizeof(TileXREp::EpAssistTuple)); + const int64_t routeMetaBytes = useRouteMetaPrefetch ? + selfSendCnt * static_cast(sizeof(TileXREp::EpAssistTuple)) : 0; + bool routeMetaPrefetchPending = false; + if (useRouteMetaPrefetch) { + StartRouteMetaPrefetch(assistInfoGM, selfSendCnt, meta, EVENT_ID1); + routeMetaPrefetchPending = true; + ProfileAux(perfTrace, perfStats, PerfStage::TX_META_SCAN, 3, + static_cast(routeMetaBytes)); + } +#endif +#if TILEXR_EP_URMA_CACHELESS + AscendC::Duplicate(selfCopyReady, TileXR::DATA_AS_FLAG_READY_VALUE, + static_cast(selfCopyFlagBytes / sizeof(float))); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); +#endif + + for (int64_t lane = 0; lane < TileXREp::kEpUrmaCombinePackLaneCount; ++lane) { + const int64_t begin = selfSendCnt * lane / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t end = selfSendCnt * (lane + 1) / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t laneLength = end - begin; + const int64_t remainder = laneLength % TileXREp::kEpUrmaCombineSendLaneCount; + const bool reverse = (lane & 1) != 0; + + // Preserve the globally balanced residue set, but reverse odd lanes so no sender + // consistently waits for the last route produced by every Pack core. + int64_t rotation = lane % TileXREp::kEpUrmaCombineSendLaneCount; + if (remainder != 0) { + rotation = begin % TileXREp::kEpUrmaCombineSendLaneCount; + if (reverse) { + rotation = (rotation + remainder - 1) % TileXREp::kEpUrmaCombineSendLaneCount; + } + } + const int64_t signedDelta = reverse ? rotation - senderId : senderId - rotation; + const int64_t delta = (signedDelta + TileXREp::kEpUrmaCombineSendLaneCount) % + TileXREp::kEpUrmaCombineSendLaneCount; + cursor.SetValue(static_cast(lane), static_cast(begin + delta)); + } + + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + for (int32_t peer = 0; peer < rankSize; ++peer) { + usedPeer.SetValue(static_cast(peer), 0); + } + const uint64_t readyValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineTxRouteReady); +#if TILEXR_EP_URMA_DOORBELL_BATCH_SIZE > 1 + TileXR::UDMADoorbellBatchGroup doorbellBatch = {}; +#else + uint64_t remotePutCount = 0; +#endif + bool allScanned = false; + while (!allScanned) { + allScanned = true; + for (int64_t lane = 0; lane < TileXREp::kEpUrmaCombinePackLaneCount; ++lane) { + const int64_t begin = selfSendCnt * lane / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t end = selfSendCnt * (lane + 1) / TileXREp::kEpUrmaCombinePackLaneCount; + const int64_t route = cursor.GetValue(static_cast(lane)); + if (route >= end) { + continue; + } + allScanned = false; + const int64_t readyRoute = TxReadyPollRoute(route, begin); +#if TILEXR_EP_URMA_TX_READY_IN_DATA + GM_ADDR readyAddr = workspaceGM + txDataOffset + readyRoute * routeStride; +#else + GM_ADDR readyAddr = workspaceGM + txReadyOffset + + readyRoute * TileXREp::kEpUrmaCombineCacheLineBytes; +#endif + const uint64_t readyPollStart = ProfileBegin(perfTrace); +#if TILEXR_EP_URMA_TX_READY_IN_DATA + const bool routeReady = LoadTxReadyInData( + readyAddr, perfTrace, perfStats) == readyValue; +#else + const bool routeReady = LoadControlValue(readyAddr, perfTrace, perfStats) == readyValue; +#endif + ProfileEnd(perfTrace, perfStats, PerfStage::TX_READY_POLL, readyPollStart); + if (!routeReady) { + ProfileAux(perfTrace, perfStats, PerfStage::TX_READY_POLL, 0, 1); + continue; + } +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL + if (!useRouteMetaPrefetch) { + CachelessAcquireBarrier(); + } +#else + CachelessAcquireBarrier(); +#endif + ProfileAux(perfTrace, perfStats, PerfStage::TX_READY_POLL, 1, 1); + + const uint64_t metaScanStart = ProfileBegin(perfTrace); +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL + TileXREp::EpAssistTuple tuple; + if (useRouteMetaPrefetch) { + if (routeMetaPrefetchPending) { + AscendC::WaitFlag(EVENT_ID1); + routeMetaPrefetchPending = false; + } + tuple = LoadPrefetchedRouteMeta(meta, route); + } else { + tuple = LoadRouteMetaBypass(assistInfoGM, route); + } +#elif TILEXR_EP_URMA_CACHELESS + const TileXREp::EpAssistTuple tuple = LoadRouteMeta(assistInfoGM, route); +#else + const TileXREp::EpAssistTuple tuple = LoadRouteMeta(assistInfoGM, route, meta); +#endif + if (!RouteMetaValid(tuple, rankSize, bs, topK)) { + ProfileEnd(perfTrace, perfStats, PerfStage::TX_META_SCAN, metaScanStart); + ProfileAux(perfTrace, perfStats, PerfStage::TX_META_SCAN, 0, 1); + StoreError(workspaceGM, errorStatusOffset, magic, + TileXREp::kEpUrmaCombineStatusInvalidRoute, perfTrace, perfStats); + cursor.SetValue(static_cast(lane), static_cast( + route + TileXREp::kEpUrmaCombineSendLaneCount)); + continue; + } + ProfileEnd(perfTrace, perfStats, PerfStage::TX_META_SCAN, metaScanStart); + ProfileAux(perfTrace, perfStats, PerfStage::TX_META_SCAN, + tuple.srcRank == rank ? 1U : 2U, 1); +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL + if (useRouteMetaPrefetch) { + CachelessAcquireBarrier(); + } +#endif + + const int64_t routeIndex = static_cast(tuple.tokenId) * topK + tuple.topKId; + const int64_t rxRouteOffset = rxWindowOffset + routeIndex * routeStride; + GM_ADDR txRouteAddr = workspaceGM + txDataOffset + route * routeStride; + if (tuple.srcRank == rank) { + const uint64_t selfCopyStart = ProfileBegin(perfTrace); + CopySelfRoute(txRouteAddr, workspaceGM + rxRouteOffset, blockCount, selfCopyPayload, + selfCopyReady, + perfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::SELF_COPY, selfCopyStart); + ProfileAux(perfTrace, perfStats, PerfStage::SELF_COPY, 0, + static_cast(routeStride)); + } else { + const uint64_t postStart = ProfileBegin(perfTrace); +#if TILEXR_EP_URMA_DOORBELL_BATCH_SIZE == 1 + TileXREp::EpUrmaUDMAPutNbi(args, tuple.srcRank, + reinterpret_cast<__gm__ uint8_t *>(txRouteAddr), + static_cast(rxRouteOffset), static_cast(routeStride), + static_cast(senderId)); + ++remotePutCount; +#else + TileXR::UDMADoorbellBatchState *batchState = TileXR::UDMAGetDoorbellBatchState( + args, static_cast(tuple.srcRank), static_cast(senderId), + &doorbellBatch); + if (batchState == nullptr) { + TileXREp::EpUrmaUDMAPutNbi(args, tuple.srcRank, + reinterpret_cast<__gm__ uint8_t *>(txRouteAddr), + static_cast(rxRouteOffset), static_cast(routeStride), + static_cast(senderId)); + ++doorbellBatch.fallbackDoorbellCount; + } else { + TileXR::UDMAPutNbiDoorbellBatched(args, tuple.srcRank, + reinterpret_cast<__gm__ uint8_t *>(txRouteAddr), + static_cast(rxRouteOffset), static_cast(routeStride), + static_cast(senderId), TileXREp::kEpUrmaCombineDoorbellBatchSize, + batchState); + } +#endif + usedPeer.SetValue(static_cast(tuple.srcRank), 1); + ProfileEnd(perfTrace, perfStats, PerfStage::UDMA_POST, postStart); + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_POST, 0, + static_cast(routeStride)); + } + cursor.SetValue(static_cast(lane), static_cast( + route + TileXREp::kEpUrmaCombineSendLaneCount)); + } + } + +#if TILEXR_EP_URMA_TX_META_PREFETCH_FULL + if (routeMetaPrefetchPending) { + const uint64_t metaDrainStart = ProfileBegin(perfTrace); + AscendC::WaitFlag(EVENT_ID1); + ProfileEnd(perfTrace, perfStats, PerfStage::TX_META_SCAN, metaDrainStart); + routeMetaPrefetchPending = false; + } +#endif + +#if TILEXR_EP_URMA_DOORBELL_BATCH_SIZE > 1 + TileXR::UDMAFlushDoorbellBatchGroup(&doorbellBatch); + const uint64_t doorbellCommitCount = static_cast( + TileXR::UDMADoorbellBatchGroupCommitCount(&doorbellBatch)); + const uint64_t activeSendSqCount = static_cast(doorbellBatch.activeQueueCount); +#else + const uint64_t doorbellCommitCount = remotePutCount; + const uint64_t activeSendSqCount = 0; +#endif + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_POST, 1, + doorbellCommitCount); + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_POST, 2, + static_cast(TileXREp::kEpUrmaCombineDoorbellBatchSize)); + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_POST, 3, + activeSendSqCount); + + __gm__ TileXR::UDMAInfo *udmaInfo = TileXR::GetUDMAInfo(args); + for (int32_t peer = 0; peer < rankSize; ++peer) { + if (usedPeer.GetValue(static_cast(peer)) == 0) { + continue; + } + const uint64_t wqeCntAddr = TileXR::UDMAGetWQCtx( + udmaInfo, static_cast(peer), static_cast(senderId))->wqeCntAddr; + bool queueAlreadyDrained = false; + for (int32_t previous = 0; previous < peer; ++previous) { + if (usedPeer.GetValue(static_cast(previous)) != 0 && + TileXR::UDMAGetWQCtx(udmaInfo, static_cast(previous), + static_cast(senderId))->wqeCntAddr == wqeCntAddr) { + queueAlreadyDrained = true; + break; + } + } + if (queueAlreadyDrained) { + continue; + } + const uint64_t quietStart = ProfileBegin(perfTrace); + (void)TileXREp::EpUrmaUDMAQuiet(args, peer, static_cast(senderId)); + ProfileEnd(perfTrace, perfStats, PerfStage::UDMA_QUIET, quietStart); + ProfileAux(perfTrace, perfStats, PerfStage::UDMA_QUIET, 0, 1); + } + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineSenderDone), perfTrace, perfStats); + return true; +} + +#undef TILEXR_EP_SEND_FUNCTION + +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool WaitLocalLines( + GM_ADDR base, int64_t lineCount, uint64_t expected, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + for (int64_t line = 0; line < lineCount; ++line) { + GM_ADDR lineAddr = base + line * TileXREp::kEpUrmaCombineCacheLineBytes; + while (LoadControlValue(lineAddr, perfTrace, perfStats) != expected) { + } + } + CachelessAcquireBarrier(); + return true; +} + +#if TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool ReleaseGenerationReached( + uint64_t observed, uint32_t expectedGeneration) +{ + const uint32_t observedStep = static_cast(observed); + const uint32_t observedGeneration = static_cast(observed >> 32U); + return observedStep == TileXREp::kEpUrmaCombineRxBufferReleased && + static_cast(observedGeneration - expectedGeneration) >= 0; +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool WaitDeferredRoundCredit( + const __gm__ TileXR::CommArgs *args, GM_ADDR workspaceGM, int64_t magic, + int64_t roundDoneOffset, int64_t senderDoneOffset, int64_t roundCreditOffset, + int64_t senderId, + GM_ADDR perfTrace, GM_ADDR finePerfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t globalWaitStart = ProfileBegin(perfTrace); + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + const int64_t creditLaneCount = rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + rankSize : TileXREp::kEpUrmaCombineSendLaneCount; + const uint64_t expectedReady = EncodeControlValue( + magic, TileXREp::kEpUrmaCombineCreditExpectedReady); + if (senderId == 0) { + StoreControlValue(workspaceGM + roundCreditOffset, expectedReady, + finePerfTrace, perfStats); + } + uint64_t pollLoads = 0; + uint32_t expectedGeneration = 0; + if (senderId < creditLaneCount) { + if (!WaitLocalLines(workspaceGM + roundCreditOffset, 1, expectedReady, + finePerfTrace, perfStats)) { + return false; + } + const uint64_t previousRelease = LoadControlValue(workspaceGM + roundDoneOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes, finePerfTrace, perfStats); + const uint32_t previousStep = static_cast(previousRelease); + expectedGeneration = static_cast(previousRelease >> 32U); + if (previousRelease != 0 && previousStep != TileXREp::kEpUrmaCombineRxBufferReleased) { + return false; + } + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(creditLaneCount)) { + if (peer == rank || previousRelease == 0) { + continue; + } + GM_ADDR peerRelease = workspaceGM + roundDoneOffset + + peer * TileXREp::kEpUrmaCombineCacheLineBytes; + while (true) { + ++pollLoads; + if (ReleaseGenerationReached( + LoadControlValue(peerRelease, finePerfTrace, perfStats), expectedGeneration)) { + break; + } + } + } + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineCreditShardDone), + finePerfTrace, perfStats); + } + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + senderDoneOffset, creditLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineCreditShardDone), + finePerfTrace, perfStats)) { + return false; + } + StoreControlValue(workspaceGM + roundCreditOffset, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineCreditRun), + finePerfTrace, perfStats); + } + if (!WaitLocalLines(workspaceGM + roundCreditOffset, 1, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineCreditRun), + finePerfTrace, perfStats)) { + return false; + } + ProfileEnd(perfTrace, perfStats, PerfStage::GLOBAL_ROUND_WAIT, globalWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::GLOBAL_ROUND_WAIT, 0, pollLoads); + ProfileAux(perfTrace, perfStats, PerfStage::GLOBAL_ROUND_WAIT, 1, expectedGeneration); + return true; +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool WaitLocalLineShard( + GM_ADDR base, int64_t lineCount, int64_t firstLine, int64_t lineStride, + uint64_t expected, GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + for (int64_t line = firstLine; line < lineCount; line += lineStride) { + GM_ADDR lineAddr = base + line * TileXREp::kEpUrmaCombineCacheLineBytes; + while (LoadControlValue(lineAddr, perfTrace, perfStats) != expected) { + } + } + CachelessAcquireBarrier(); + return true; +} +#endif + +#if TILEXR_EP_URMA_START_GATE +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool WaitForSynchronizedStart( + const __gm__ TileXR::CommArgs *args, + GM_ADDR workspaceGM, int64_t magic, int64_t rxLaneDoneOffset, int64_t senderDoneOffset, + int64_t roundPublishOffset, int64_t startGateOffset, int64_t blockIdx, + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t gateStart = ProfileBegin(perfTrace); + const uint64_t localReady = EncodeControlValue( + magic, TileXREp::kEpUrmaCombineStartLocalReady); + const bool isPackReceive = blockIdx < TileXREp::kEpUrmaCombinePackLaneCount; + if (isPackReceive) { + StoreControlValue(workspaceGM + rxLaneDoneOffset + + blockIdx * TileXREp::kEpUrmaCombineCacheLineBytes, + localReady, nullptr, perfStats); + } else { + const int64_t senderId = blockIdx - TileXREp::kEpUrmaCombinePackLaneCount; + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + localReady, nullptr, perfStats); + } + + const uint64_t startRun = EncodeControlValue(magic, TileXREp::kEpUrmaCombineStartRun); + if (isPackReceive) { + if (!WaitLocalLines(workspaceGM + roundPublishOffset, 1, startRun, + nullptr, perfStats)) { + return false; + } + ProfileEnd(perfTrace, perfStats, PerfStage::START_GATE, gateStart); + return true; + } + + const int64_t senderId = blockIdx - TileXREp::kEpUrmaCombinePackLaneCount; + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + const uint64_t rankReady = EncodeControlValue( + magic, TileXREp::kEpUrmaCombineStartRankReady); + GM_ADDR rankReadyAddr = workspaceGM + startGateOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes; + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + rxLaneDoneOffset, TileXREp::kEpUrmaCombinePackLaneCount, + localReady, nullptr, perfStats) || + !WaitLocalLines(workspaceGM + senderDoneOffset, TileXREp::kEpUrmaCombineSendLaneCount, + localReady, nullptr, perfStats)) { + return false; + } + StoreControlValue(rankReadyAddr, rankReady, nullptr, perfStats); + StoreControlValue(workspaceGM + roundPublishOffset, rankReady, nullptr, perfStats); + } + if (!WaitLocalLines(workspaceGM + roundPublishOffset, 1, rankReady, nullptr, perfStats)) { + return false; + } + + const uint32_t qpIdx = static_cast(senderId); + const uint64_t remoteOffset = static_cast(startGateOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes); + uint64_t publishCount = 0; + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(TileXREp::kEpUrmaCombineSendLaneCount)) { + if (peer == rank) { + continue; + } + TileXREp::EpUrmaUDMAPutNbi(args, peer, + reinterpret_cast<__gm__ uint8_t *>(rankReadyAddr), remoteOffset, + static_cast(TileXREp::kEpUrmaCombineCacheLineBytes), qpIdx); + ++publishCount; + } + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(TileXREp::kEpUrmaCombineSendLaneCount)) { + if (peer == rank) { + continue; + } + (void)TileXREp::EpUrmaUDMAQuiet(args, peer, qpIdx); + } + ProfileAux(perfTrace, perfStats, PerfStage::START_GATE, 0, publishCount); + ProfileAux(perfTrace, perfStats, PerfStage::START_GATE, 1, + publishCount * TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileAux(perfTrace, perfStats, PerfStage::START_GATE, 2, publishCount); + ProfileAux(perfTrace, perfStats, PerfStage::START_GATE, 3, publishCount); + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineStartPublishDone), nullptr, perfStats); + + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + senderDoneOffset, TileXREp::kEpUrmaCombineSendLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineStartPublishDone), nullptr, perfStats) || + !WaitLocalLines(workspaceGM + startGateOffset, rankSize, rankReady, + nullptr, perfStats)) { + return false; + } + StoreControlValue(workspaceGM + roundPublishOffset, startRun, nullptr, perfStats); + } + if (!WaitLocalLines(workspaceGM + roundPublishOffset, 1, startRun, nullptr, perfStats)) { + return false; + } + ProfileEnd(perfTrace, perfStats, PerfStage::START_GATE, gateStart); + return true; +} +#endif + +__aicore__ TILEXR_EP_LOCAL_FUNCTION void RecordRoundPublishCounters( + GM_ADDR perfTrace, __ubuf__ PerfStats *perfStats, uint64_t publishCount) +{ + ProfileAux(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, 0, publishCount); + ProfileAux(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, 1, + publishCount * TileXREp::kEpUrmaCombineCacheLineBytes); + ProfileAux(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, 2, publishCount); + ProfileAux(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, 3, publishCount); +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool WaitGlobalRoundDone( + GM_ADDR workspaceGM, int32_t rankSize, + int64_t magic, int64_t roundDoneOffset, + GM_ADDR perfTrace, GM_ADDR finePerfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t globalWaitStart = ProfileBegin(perfTrace); + uint64_t pollLoads = 0; + const uint64_t expected = EncodeControlValue( + magic, TileXREp::kEpUrmaCombineRxBufferReleased); + for (int32_t peer = 0; peer < rankSize; ++peer) { + GM_ADDR lineAddr = workspaceGM + roundDoneOffset + + peer * TileXREp::kEpUrmaCombineCacheLineBytes; + while (true) { + ++pollLoads; + if (LoadControlValue(lineAddr, finePerfTrace, perfStats) == expected) { + break; + } + } + } + CachelessAcquireBarrier(); + ProfileEnd(perfTrace, perfStats, PerfStage::GLOBAL_ROUND_WAIT, globalWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::GLOBAL_ROUND_WAIT, 0, pollLoads); + return true; +} + +#if TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t StartParallelRoundPublish( + const __gm__ TileXR::CommArgs *args, + GM_ADDR workspaceGM, int64_t magic, int64_t roundDoneOffset, int64_t roundPublishOffset, + GM_ADDR perfTrace, GM_ADDR finePerfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t publishStart = ProfileBegin(perfTrace); + const uint64_t roundValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxBufferReleased); + StoreControlValue(workspaceGM + roundDoneOffset + + args->rank * TileXREp::kEpUrmaCombineCacheLineBytes, + roundValue, finePerfTrace, perfStats); +#if !TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT + // This release store is also the shared 64-byte UDMA source and the local shard start flag. + StoreControlValue(workspaceGM + roundPublishOffset, roundValue, finePerfTrace, perfStats); +#else + (void)roundPublishOffset; +#endif + return publishStart; +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t PublishRoundShard( + const __gm__ TileXR::CommArgs *args, + GM_ADDR workspaceGM, int64_t magic, int64_t roundDoneOffset, int64_t roundPublishOffset, + int64_t senderDoneOffset, int64_t senderId, + GM_ADDR perfTrace, GM_ADDR finePerfTrace, __ubuf__ PerfStats *perfStats) +{ + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + const uint64_t roundValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxBufferReleased); + GM_ADDR publishAddr = workspaceGM + roundPublishOffset; +#if TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT + publishAddr = workspaceGM + roundDoneOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes; +#endif + if (!WaitLocalLines(publishAddr, 1, roundValue, finePerfTrace, perfStats)) { + return ~0ULL; + } + + const uint64_t shardStart = senderId == 0 ? 0 : ProfileBegin(perfTrace); + const uint32_t qpIdx = static_cast(senderId); + const uint64_t remoteOffset = static_cast(roundDoneOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes); + uint64_t publishCount = 0; + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(TileXREp::kEpUrmaCombineSendLaneCount)) { + if (peer == rank) { + continue; + } + TileXREp::EpUrmaUDMAPutNbi(args, peer, + reinterpret_cast<__gm__ uint8_t *>(publishAddr), remoteOffset, + static_cast(TileXREp::kEpUrmaCombineCacheLineBytes), qpIdx); + ++publishCount; + } + for (int32_t peer = static_cast(senderId); peer < rankSize; + peer += static_cast(TileXREp::kEpUrmaCombineSendLaneCount)) { + if (peer == rank) { + continue; + } + (void)TileXREp::EpUrmaUDMAQuiet(args, peer, qpIdx); + } + + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombinePublishDone), finePerfTrace, perfStats); + if (senderId != 0) { + ProfileEnd(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, shardStart); + RecordRoundPublishCounters(perfTrace, perfStats, publishCount); + } + return publishCount; +} + +__aicore__ TILEXR_EP_LOCAL_FUNCTION bool FinishParallelRoundPublish( + GM_ADDR workspaceGM, + int64_t magic, int64_t senderDoneOffset, int64_t publisherCount, + uint64_t publishStart, uint64_t sender0PublishCount, + GM_ADDR perfTrace, GM_ADDR finePerfTrace, __ubuf__ PerfStats *perfStats) +{ + if (!WaitLocalLines(workspaceGM + senderDoneOffset, publisherCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombinePublishDone), finePerfTrace, perfStats)) { + return false; + } + ProfileEnd(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, publishStart); + RecordRoundPublishCounters(perfTrace, perfStats, sender0PublishCount); + return true; +} +#else +__aicore__ TILEXR_EP_LOCAL_FUNCTION void PublishAndWaitRoundDone( + const __gm__ TileXR::CommArgs *args, GM_ADDR workspaceGM, + int64_t magic, int64_t roundDoneOffset, int64_t roundPublishOffset, + GM_ADDR perfTrace, GM_ADDR finePerfTrace, __ubuf__ PerfStats *perfStats) +{ + const uint64_t publishStart = ProfileBegin(perfTrace); + const int32_t rank = args->rank; + const int32_t rankSize = args->rankSize; + const uint64_t roundValue = EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxBufferReleased); + GM_ADDR publishAddr = workspaceGM + roundPublishOffset; + StoreControlValue(publishAddr, roundValue, finePerfTrace, perfStats); + StoreControlValue(workspaceGM + roundDoneOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes, roundValue, finePerfTrace, perfStats); + + for (int32_t peer = 0; peer < rankSize; ++peer) { + if (peer == rank) { + continue; + } + const uint64_t remoteOffset = static_cast(roundDoneOffset + + rank * TileXREp::kEpUrmaCombineCacheLineBytes); + TileXREp::EpUrmaUDMAPutNbi(args, peer, + reinterpret_cast<__gm__ uint8_t *>(publishAddr), remoteOffset, + static_cast(TileXREp::kEpUrmaCombineCacheLineBytes)); + TileXREp::EpUrmaUDMAQuiet(args, peer); + } + ProfileEnd(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, publishStart); + RecordRoundPublishCounters(perfTrace, perfStats, static_cast(rankSize - 1)); + WaitGlobalRoundDone(workspaceGM, rankSize, magic, roundDoneOffset, + perfTrace, finePerfTrace, perfStats); +} +#endif + +} // namespace + +extern "C" __global__ __aicore__ void tilexr_ep_urma_combine_kernel(GM_ADDR commArgsGM, + GM_ADDR expertOutGM, GM_ADDR assistInfoForCombineGM, GM_ADDR topKWeightsGM, GM_ADDR yOutGM, + GM_ADDR workspaceGM, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, + int64_t workspaceBytes, int64_t magic, int64_t commBytes, int64_t blockCount, int64_t routeStride, + int64_t rxWindowBytes, int64_t rxWindowOffset0, int64_t rxWindowOffset1, int64_t roundDoneOffset0, + int64_t roundDoneOffset1, int64_t rxLaneDoneOffset, int64_t senderDoneOffset, + int64_t roundPublishOffset, int64_t roundCreditOffset, int64_t startGateOffset, + int64_t runStartGate, int64_t errorStatusOffset, + int64_t txReadyOffset, int64_t txDataOffset, GM_ADDR perfTrace, int64_t perfTraceBytes, + GM_ADDR strictKernelCycles) +{ + if constexpr (g_coreType == AscendC::AIV) { + if (commArgsGM == nullptr || topKWeightsGM == nullptr || yOutGM == nullptr || workspaceGM == nullptr || + (selfSendCnt > 0 && (expertOutGM == nullptr || assistInfoForCombineGM == nullptr)) || + selfSendCnt < 0 || bs <= 0 || h <= 0 || h > TileXREp::kEpUrmaCombineMaxHidden || topK <= 0 || + topK > TileXREp::kEpUrmaCombineMaxTopK || magic <= 0 || commBytes != + TileXREp::kEpUrmaCombineQuantHeaderBytes + AlignUpInt64(h, kCursorAlignment) || blockCount <= 0 || + blockCount != (commBytes + TileXREp::kEpUrmaCombinePayloadBytes - 1) / + TileXREp::kEpUrmaCombinePayloadBytes || + blockCount > TileXREp::kEpUrmaCombineMaxBlocksPerRoute || routeStride <= 0 || + routeStride != blockCount * TileXREp::kEpUrmaCombineDataBlockBytes || + (runStartGate != 0 && runStartGate != 1) || +#if TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT + roundCreditOffset <= 0 || + roundCreditOffset > workspaceBytes - TileXREp::kEpUrmaCombineCacheLineBytes || +#endif + rxWindowBytes <= 0 || txReadyOffset < 0 || txDataOffset < txReadyOffset || + txDataOffset > workspaceBytes || selfSendCnt > + (txDataOffset - txReadyOffset) / TileXREp::kEpUrmaCombineCacheLineBytes || + selfSendCnt > (workspaceBytes - txDataOffset) / routeStride) { + return; + } + auto args = reinterpret_cast<__gm__ TileXR::CommArgs *>(commArgsGM); + if (args->rankSize <= 0 || args->rankSize > TileXR::TILEXR_MAX_RANK_SIZE || args->rank < 0 || + args->rank >= args->rankSize || (args->rankSize > 1 && !TileXR::UDMARegistryEnabled(args))) { + return; + } + if (args->rankSize > 1) { + __gm__ TileXR::UDMAInfo *udmaInfo = TileXR::GetUDMAInfo(args); + if (udmaInfo == nullptr || udmaInfo->qpNum < + static_cast(TileXREp::kEpUrmaCombineRequiredQpCount)) { + return; + } + } +#if TILEXR_EP_URMA_START_GATE + if (startGateOffset < 0 || startGateOffset > workspaceBytes || + args->rankSize > (workspaceBytes - startGateOffset) / + TileXREp::kEpUrmaCombineCacheLineBytes) { + return; + } +#else + (void)startGateOffset; +#endif +#if !TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT + (void)roundCreditOffset; +#endif + if (strictKernelCycles != nullptr && perfTrace != nullptr) { + return; + } + if (!ProfileBufferValid(perfTrace, perfTraceBytes, args->rank, args->rankSize)) { + perfTrace = nullptr; + } + GM_ADDR kernelPerfTrace = perfTrace; + + const int64_t blockIdx = AscendC::GetBlockIdx(); + if (blockIdx < 0 || blockIdx >= TileXREp::kEpUrmaCombineAivCount) { + return; + } + const uint32_t perfRank = static_cast(args->rank); + const uint32_t perfCore = static_cast(blockIdx); + __ubuf__ PerfStats *perfStats = reinterpret_cast<__ubuf__ PerfStats *>( + TileXR::TILEXR_PERF_TRACE_LOCAL_STATS_UB_OFFSET); + TileXR::TileXRPerfLocalStatsInit(kernelPerfTrace, perfStats, perfRank, perfCore, + TileXREp::kEpUrmaCombinePerfStageCount); + GM_ADDR finePerfTrace = nullptr; + if (kernelPerfTrace != nullptr) { + const uint32_t profileDetail = + reinterpret_cast<__gm__ TileXR::TileXRPerfTraceHeader *>(kernelPerfTrace)->flags; + if (profileDetail >= 2) { + finePerfTrace = kernelPerfTrace; + } + if (profileDetail == 0) { + perfTrace = nullptr; + } + } +#if TILEXR_EP_URMA_START_GATE + if (runStartGate != 0) { + if (!WaitForSynchronizedStart(args, workspaceGM, magic, rxLaneDoneOffset, senderDoneOffset, + roundPublishOffset, startGateOffset, blockIdx, perfTrace, perfStats)) { + return; + } + } +#else + (void)runStartGate; +#endif + const uint64_t kernelStart = ProfileKernelTimingBegin(kernelPerfTrace); + const uint64_t strictKernelStart = StrictKernelTimingBegin(strictKernelCycles); + const int64_t rxBufferIndex = static_cast(magic) & 1U; + const int64_t rxWindowOffset = rxBufferIndex == 0 ? rxWindowOffset0 : rxWindowOffset1; + const int64_t roundDoneOffset = rxBufferIndex == 0 ? roundDoneOffset0 : roundDoneOffset1; + AscendC::TPipe pipe; + if (blockIdx < TileXREp::kEpUrmaCombinePackLaneCount) { + const uint64_t packStart = ProfileBegin(perfTrace); + PackRoutes(expertOutGM, workspaceGM, selfSendCnt, h, magic, blockCount, routeStride, + txReadyOffset, txDataOffset, blockIdx, pipe, finePerfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::PACK_TOTAL, packStart); + pipe.Reset(); + const uint64_t receiveStart = ProfileBegin(perfTrace); + ReceiveTokens(topKWeightsGM, yOutGM, workspaceGM, bs, h, topK, magic, rxWindowOffset, + blockCount, routeStride, rxLaneDoneOffset, errorStatusOffset, blockIdx, pipe, + finePerfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::RECEIVE_TOTAL, receiveStart); + StrictKernelTimingFinish(strictKernelCycles, perfCore, strictKernelStart); + ProfileFinish(kernelPerfTrace, perfRank, perfCore, perfStats, kernelStart); + return; + } + + const int64_t senderId = blockIdx - TileXREp::kEpUrmaCombinePackLaneCount; + if (senderId >= TileXREp::kEpUrmaCombineSendLaneCount) { + return; + } +#if TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT + if (!WaitDeferredRoundCredit(args, workspaceGM, magic, roundDoneOffset, + senderDoneOffset, roundCreditOffset, senderId, + perfTrace, finePerfTrace, perfStats)) { + return; + } +#endif + const uint64_t sendStart = ProfileBegin(perfTrace); + if (!SendRoutes(args, assistInfoForCombineGM, workspaceGM, selfSendCnt, bs, topK, magic, + rxWindowOffset, blockCount, routeStride, txReadyOffset, txDataOffset, senderDoneOffset, + errorStatusOffset, senderId, pipe, finePerfTrace, perfStats)) { + return; + } + ProfileEnd(perfTrace, perfStats, PerfStage::SEND_TOTAL, sendStart); +#if TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH +#if TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT + const uint64_t receiveWaitStart = ProfileBegin(perfTrace); + if (!WaitLocalLineShard(workspaceGM + rxLaneDoneOffset, + TileXREp::kEpUrmaCombinePackLaneCount, senderId, + TileXREp::kEpUrmaCombineSendLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxLaneDone), + finePerfTrace, perfStats)) { + return; + } + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxReleaseShardDone), + finePerfTrace, perfStats); + uint64_t publishStart = 0; + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + senderDoneOffset, + TileXREp::kEpUrmaCombineSendLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxReleaseShardDone), + finePerfTrace, perfStats)) { + return; + } + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, 0, + static_cast(TileXREp::kEpUrmaCombinePackLaneCount)); + publishStart = StartParallelRoundPublish(args, workspaceGM, magic, + roundDoneOffset, roundPublishOffset, perfTrace, finePerfTrace, perfStats); + } else { + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + } + + const int64_t publisherCount = args->rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + args->rankSize : TileXREp::kEpUrmaCombineSendLaneCount; + uint64_t publishCount = 0; + if (senderId < publisherCount) { + publishCount = PublishRoundShard(args, workspaceGM, magic, + roundDoneOffset, roundPublishOffset, senderDoneOffset, senderId, + perfTrace, finePerfTrace, perfStats); + if (publishCount == ~0ULL) { + return; + } + } + if (senderId == 0 && + !FinishParallelRoundPublish(workspaceGM, magic, senderDoneOffset, + publisherCount, publishStart, publishCount, + perfTrace, finePerfTrace, perfStats)) { + return; + } +#else + uint64_t publishStart = 0; + if (senderId == 0) { + const uint64_t senderWaitStart = ProfileBegin(perfTrace); + if (!WaitLocalLines(workspaceGM + senderDoneOffset, TileXREp::kEpUrmaCombineSendLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineSenderDone), finePerfTrace, perfStats)) { + return; + } + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_SENDER_WAIT, senderWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::LOCAL_SENDER_WAIT, 0, + static_cast(TileXREp::kEpUrmaCombineSendLaneCount)); + const uint64_t receiveWaitStart = ProfileBegin(perfTrace); + if (!WaitLocalLines(workspaceGM + rxLaneDoneOffset, TileXREp::kEpUrmaCombinePackLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxLaneDone), finePerfTrace, perfStats)) { + return; + } + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, 0, + static_cast(TileXREp::kEpUrmaCombinePackLaneCount)); + publishStart = StartParallelRoundPublish(args, workspaceGM, magic, + roundDoneOffset, roundPublishOffset, perfTrace, finePerfTrace, perfStats); + } + const uint64_t publishCount = PublishRoundShard(args, workspaceGM, magic, + roundDoneOffset, roundPublishOffset, senderDoneOffset, senderId, + perfTrace, finePerfTrace, perfStats); + if (publishCount == ~0ULL) { + return; + } + if (senderId == 0) { + if (!FinishParallelRoundPublish(workspaceGM, magic, senderDoneOffset, + TileXREp::kEpUrmaCombineSendLaneCount, publishStart, publishCount, + perfTrace, finePerfTrace, perfStats)) { + return; + } + if (!WaitGlobalRoundDone(workspaceGM, args->rankSize, magic, roundDoneOffset, + perfTrace, finePerfTrace, perfStats)) { + return; + } + } +#endif +#else + if (senderId == 0) { + const uint64_t senderWaitStart = ProfileBegin(perfTrace); + WaitLocalLines(workspaceGM + senderDoneOffset, TileXREp::kEpUrmaCombineSendLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineSenderDone), finePerfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_SENDER_WAIT, senderWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::LOCAL_SENDER_WAIT, 0, + static_cast(TileXREp::kEpUrmaCombineSendLaneCount)); + const uint64_t receiveWaitStart = ProfileBegin(perfTrace); + WaitLocalLines(workspaceGM + rxLaneDoneOffset, TileXREp::kEpUrmaCombinePackLaneCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxLaneDone), finePerfTrace, perfStats); + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, 0, + static_cast(TileXREp::kEpUrmaCombinePackLaneCount)); + PublishAndWaitRoundDone(args, workspaceGM, magic, roundDoneOffset, roundPublishOffset, + perfTrace, finePerfTrace, perfStats); + } +#endif + StrictKernelTimingFinish(strictKernelCycles, perfCore, strictKernelStart); + ProfileFinish(kernelPerfTrace, perfRank, perfCore, perfStats, kernelStart); + } +} + +#undef TILEXR_EP_LOCAL_FUNCTION + +void launch_tilexr_ep_urma_combine_kernel(uint32_t blockDim, void *stream, GM_ADDR commArgs, + GM_ADDR expertOut, GM_ADDR assistInfoForCombine, GM_ADDR topKWeights, GM_ADDR yOut, GM_ADDR workspace, + int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, int64_t workspaceBytes, int64_t magic, + int64_t commBytes, int64_t blockCount, int64_t routeStride, int64_t rxWindowBytes, int64_t rxWindowOffset0, + int64_t rxWindowOffset1, int64_t roundDoneOffset0, int64_t roundDoneOffset1, int64_t rxLaneDoneOffset, + int64_t senderDoneOffset, int64_t roundPublishOffset, int64_t roundCreditOffset, + int64_t startGateOffset, int64_t runStartGate, + int64_t errorStatusOffset, int64_t txReadyOffset, int64_t txDataOffset, GM_ADDR perfTrace, int64_t perfTraceBytes, + GM_ADDR strictKernelCycles) +{ + tilexr_ep_urma_combine_kernel<<>>(commArgs, expertOut, assistInfoForCombine, + topKWeights, yOut, workspace, selfSendCnt, bs, h, topK, workspaceBytes, magic, commBytes, + blockCount, routeStride, rxWindowBytes, rxWindowOffset0, rxWindowOffset1, + roundDoneOffset0, roundDoneOffset1, rxLaneDoneOffset, senderDoneOffset, roundPublishOffset, + roundCreditOffset, startGateOffset, runStartGate, errorStatusOffset, txReadyOffset, txDataOffset, + perfTrace, perfTraceBytes, + strictKernelCycles); +} + +#undef TILEXR_EP_URMA_CACHELESS diff --git a/src/include/tilexr_ep.h b/src/include/tilexr_ep.h index 82c79cde..83540c9a 100644 --- a/src/include/tilexr_ep.h +++ b/src/include/tilexr_ep.h @@ -25,6 +25,24 @@ int TileXRMoeEpCombineV2(void *expertOut, int32_t *assistInfoForCombine, int32_t TileXRCommPtr comm, int64_t bs, int64_t h, int64_t topK, int64_t moeExpertNum, void *yOut, void *workspace, TileXR::TileXRDataType dtype, aclrtStream stream); +// workspace must be zero-initialized before its first use and registered once +// as the communicator's UDMA region. Reuse is supported for serialized launches. +int TileXRMoeEpCombineUrmaGetWorkspaceSize(int64_t rankSize, int64_t bs, int64_t h, int64_t topK, + int64_t selfSendCapacity, int64_t *workspaceBytes); + +int TileXRMoeEpCombineUrmaGetProfileSize(int64_t rankSize, int64_t *profileBytes); + +int TileXRMoeEpCombineUrma(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, TileXR::TileXRDataType dtype, aclrtStream stream); + +// perfTrace points to a device TileXRPerfTrace buffer initialized by the caller. +// Use TileXRMoeEpCombineUrmaGetProfileSize to size the allocation. +int TileXRMoeEpCombineUrmaProfile(void *expertOut, int32_t *assistInfoForCombine, float *topKWeights, + TileXRCommPtr comm, int64_t selfSendCnt, int64_t bs, int64_t h, int64_t topK, void *yOut, + void *workspace, int64_t workspaceBytes, void *perfTrace, int64_t perfTraceBytes, + 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, diff --git a/tests/ep/CMakeLists.txt b/tests/ep/CMakeLists.txt index e505f1d8..9c6fe7ca 100644 --- a/tests/ep/CMakeLists.txt +++ b/tests/ep/CMakeLists.txt @@ -77,6 +77,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_urma_combine_layout.cpp ) add_executable(test_tilexr_ep_api_sources @@ -85,6 +86,10 @@ add_executable(test_tilexr_ep_api_sources add_executable(test_tilexr_ep_kernel_sources unit/test_tilexr_ep_kernel_sources.cpp) +add_executable(test_tilexr_ep_start_gate_window + unit/test_tilexr_ep_start_gate_window.cpp +) + add_executable(test_tilexr_ep_host_validation unit/test_tilexr_ep_host_validation.cpp ${TILEXR_ROOT}/src/ep/host/ep_layout.cpp @@ -95,6 +100,10 @@ target_include_directories(test_tilexr_ep_layout PRIVATE ${TILEXR_EP_TEST_INCLUDE_DIRS} ) +target_include_directories(test_tilexr_ep_start_gate_window PRIVATE + ${TILEXR_EP_TEST_INCLUDE_DIRS} +) + target_include_directories(test_tilexr_ep_host_validation PRIVATE ${TILEXR_EP_TEST_INCLUDE_DIRS} ) @@ -105,12 +114,14 @@ target_compile_definitions(test_tilexr_ep_kernel_sources PRIVATE TILEXR_SOURCE_R add_test(NAME test_tilexr_ep_layout COMMAND test_tilexr_ep_layout) add_test(NAME test_tilexr_ep_api_sources COMMAND test_tilexr_ep_api_sources) add_test(NAME test_tilexr_ep_kernel_sources COMMAND test_tilexr_ep_kernel_sources) +add_test(NAME test_tilexr_ep_start_gate_window COMMAND test_tilexr_ep_start_gate_window) add_test(NAME test_tilexr_ep_host_validation COMMAND test_tilexr_ep_host_validation) install(TARGETS test_tilexr_ep_layout test_tilexr_ep_api_sources test_tilexr_ep_kernel_sources + test_tilexr_ep_start_gate_window test_tilexr_ep_host_validation RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} ) diff --git a/tests/ep/unit/test_tilexr_ep_api_sources.cpp b/tests/ep/unit/test_tilexr_ep_api_sources.cpp index 7f942804..eca054ce 100644 --- a/tests/ep/unit/test_tilexr_ep_api_sources.cpp +++ b/tests/ep/unit/test_tilexr_ep_api_sources.cpp @@ -257,6 +257,35 @@ void TestDispatchDemoUsesHostBarrierBeforeValidation() CheckContains("tests/ep/demo/tilexr_ep_dispatch_demo.cpp", demo, "dispatch synchronized"); } +void TestUrmaCombineApiAndBuildIntegration() +{ + const std::string apiPath = "src/include/tilexr_ep.h"; + std::string api; + if (ReadFile(apiPath, &api)) { + CheckContains(apiPath, api, "TileXRMoeEpCombineUrmaGetWorkspaceSize"); + CheckContains(apiPath, api, "TileXRMoeEpCombineUrma("); + CheckContains(apiPath, api, "TileXRMoeEpCombineUrmaProfile("); + } + + const std::string hostPath = "src/ep/host/ep_urma_combine_host.cpp"; + std::string host; + if (ReadFile(hostPath, &host)) { + CheckContains(hostPath, host, "TileXREpPrepareUrmaCombineLaunchContext"); + CheckContains(hostPath, host, "launch_tilexr_ep_urma_combine_kernel"); + CheckNotContains(hostPath, host, "TILEXR_EP_URMA_OPERATOR_LAUNCH"); + } + + const std::string cmakePath = "src/ep/CMakeLists.txt"; + std::string cmake; + if (ReadFile(cmakePath, &cmake)) { + CheckContains(cmakePath, cmake, "tilexr_ep_urma_combine_kernel.cpp"); + CheckContains(cmakePath, cmake, "Building TileXR EP URMA combine kernel (Ascend 950 S22, O2)"); + CheckContains(cmakePath, cmake, "-O2"); + CheckNotContains(cmakePath, cmake, "DIAGNOSTIC"); + CheckNotContains(cmakePath, cmake, "KERNEL_OPT_LEVEL"); + } +} + void TestNoForbiddenDependencies() { const std::vector paths = { @@ -301,6 +330,7 @@ int main() TestDemoRunnerUsesLibAndLib64Paths(); TestDispatchDemoRegistersAlignedUdmaWorkspace(); TestDispatchDemoUsesHostBarrierBeforeValidation(); + TestUrmaCombineApiAndBuildIntegration(); TestNoForbiddenDependencies(); return g_failures == 0 ? 0 : 1; } diff --git a/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp b/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp index a8100561..e83c84eb 100644 --- a/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp +++ b/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp @@ -417,6 +417,37 @@ void TestClearLocalWindowDoesNotPreclearSlotHeaders() " }"); } +void TestUrmaCombineUsesCleanS22ProductionPath() +{ + const std::string configPath = "src/ep/common/ep_urma_combine_config.h"; + std::string config; + if (ReadFile(configPath, &config)) { + CheckContains(configPath, config, "TILEXR_EP_URMA_COMBINE_SEND_CORE_COUNT 22"); + CheckContains(configPath, config, "TILEXR_EP_URMA_RX_SCHEDULER 1"); + CheckContains(configPath, config, "TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH 1"); + CheckContains(configPath, config, "TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT 1"); + CheckContains(configPath, config, "TILEXR_EP_URMA_START_GATE 1"); + CheckContains(configPath, config, "TILEXR_EP_URMA_QDC_VERSION 3"); + CheckContains(configPath, config, "TILEXR_EP_URMA_RX_READY_STICKY_MASK 1"); + } + + const std::string kernelPath = "src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp"; + std::string kernel; + if (ReadFile(kernelPath, &kernel)) { + CheckContains(kernelPath, kernel, "const bool reverse = (lane & 1) != 0;"); + CheckContains(kernelPath, kernel, "WaitDeferredRoundCredit"); + CheckContains(kernelPath, kernel, "StartParallelRoundPublish"); + CheckContains(kernelPath, kernel, "FinishParallelRoundPublish"); + CheckNotContains(kernelPath, kernel, "TILEXR_EP_URMA_DIAGNOSTIC_"); + } + + const std::string udmaPath = "src/include/tilexr_udma.h"; + std::string udma; + if (ReadFile(udmaPath, &udma)) { + CheckNotContains(udmaPath, udma, "TILEXR_EP_URMA_DIAGNOSTIC_"); + } +} + void TestNoForbiddenDependencies() { const std::vector paths = { @@ -468,6 +499,7 @@ int main() TestKernelForwardsStaticQuantConfig(); TestKernelForwardsPerTokenDynamicQuantConfig(); TestClearLocalWindowDoesNotPreclearSlotHeaders(); + TestUrmaCombineUsesCleanS22ProductionPath(); TestNoForbiddenDependencies(); if (g_failures != 0) { std::cerr << g_failures << " TileXR EP kernel source checks failed" << std::endl; diff --git a/tests/ep/unit/test_tilexr_ep_layout.cpp b/tests/ep/unit/test_tilexr_ep_layout.cpp index 69c0e7b9..cbfbb2d4 100644 --- a/tests/ep/unit/test_tilexr_ep_layout.cpp +++ b/tests/ep/unit/test_tilexr_ep_layout.cpp @@ -1,8 +1,11 @@ #include #include +#include #include "comm_args.h" #include "ep_layout.h" +#include "ep_urma_combine.h" +#include "ep_urma_combine_layout.h" #include "tilexr_types.h" namespace { @@ -100,6 +103,147 @@ void TestRejectsInvalidConfig() TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); } +void TestUrmaCombineWorkspaceConfig() +{ + TileXREp::EpUrmaCombineWorkspaceConfig config {}; + const int ret = TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 4, 8, 2, 3, &config); + CheckInt("URMA combine valid config", ret, TileXR::TILEXR_SUCCESS); + CheckInt64("URMA combine quant bytes", config.quantDataBytes, 32); + CheckInt64("URMA combine comm bytes", config.commBytes, 64); + CheckInt64("URMA combine blocks", config.blockCount, 1); + CheckInt64("URMA combine route stride", config.routeStride, 512); + CheckInt64("URMA combine route count", config.routeCount, 8); + CheckInt64("URMA combine rx bytes", config.rxWindowBytes, 4096); + CheckInt64("URMA combine rx0", config.rxWindowOffsets[0], 512); + CheckInt64("URMA combine rx1", config.rxWindowOffsets[1], 4608); + CheckInt64("URMA combine round0", config.roundDoneOffsets[0], 8704); + CheckInt64("URMA combine round1", config.roundDoneOffsets[1], 8832); + CheckInt64("URMA combine rx lane done", config.rxLaneDoneOffset, 8960); + CheckInt64("URMA combine sender done", config.senderDoneOffset, + config.rxLaneDoneOffset + TileXREp::kEpUrmaCombinePackLaneCount * + TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA combine round publish", config.roundPublishOffset, 13056); + CheckInt64("URMA combine round credit", config.roundCreditOffset, + TileXREp::kEpUrmaCombineDeferredRoundCredit ? 13120 : 0); + CheckInt64("URMA combine start gate", config.startGateOffset, + TileXREp::kEpUrmaCombineStartGate ? + (TileXREp::kEpUrmaCombineDeferredRoundCredit ? 13184 : 13120) : 0); + CheckInt64("URMA combine error", config.errorStatusOffset, + TileXREp::kEpUrmaCombineStartGate ? + (TileXREp::kEpUrmaCombineDeferredRoundCredit ? 13312 : 13248) : + (TileXREp::kEpUrmaCombineDeferredRoundCredit ? 13184 : 13120)); + CheckInt64("URMA combine fixed bytes", config.fixedBytes, + TileXREp::kEpUrmaCombineStartGate ? + (TileXREp::kEpUrmaCombineDeferredRoundCredit ? 13376 : 13312) : + (TileXREp::kEpUrmaCombineDeferredRoundCredit ? 13248 : 13184)); + CheckInt64("URMA combine fixed control boundary", config.fixedBytes, + config.errorStatusOffset + TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA combine tx ready", config.txReadyOffset, + config.fixedBytes); + CheckInt64("URMA combine tx data", config.txDataOffset, 13824); + CheckInt64("URMA combine required bytes", config.requiredBytes, 15360); +} + +void TestUrmaCombineStartGateLayout() +{ + TileXREp::EpUrmaCombineWorkspaceConfig config {}; + CheckInt("URMA start gate config", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig( + 64, 128, 7168, 8, 16, &config), TileXR::TILEXR_SUCCESS); + if (TileXREp::kEpUrmaCombineStartGate) { + CheckInt64("URMA start gate follows publish", config.startGateOffset, + (TileXREp::kEpUrmaCombineDeferredRoundCredit ? config.roundCreditOffset : + config.roundPublishOffset) + TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA start gate rank lines", config.errorStatusOffset, + config.startGateOffset + config.rankSize * TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA start gate required QPs", TileXREp::kEpUrmaCombineRequiredQpCount, + TileXREp::kEpUrmaCombineSendLaneCount); + } else { + CheckInt64("URMA disabled start gate offset", config.startGateOffset, 0); + CheckInt64("URMA disabled start gate keeps error", config.errorStatusOffset, + (TileXREp::kEpUrmaCombineDeferredRoundCredit ? config.roundCreditOffset : + config.roundPublishOffset) + TileXREp::kEpUrmaCombineCacheLineBytes); + CheckInt64("URMA disabled start gate QPs", TileXREp::kEpUrmaCombineRequiredQpCount, + TileXREp::kEpUrmaCombineSendLaneCount); + } +} + +void TestParallelRoundPublishPeerShards() +{ + constexpr int64_t kParallelSendLanes = TileXREp::kEpUrmaCombineSendLaneCount; + const int64_t rankSizes[] = {1, 2, 8, 15, 16, 17, 64}; + for (int64_t rankSize : rankSizes) { + for (int64_t rank = 0; rank < rankSize; ++rank) { + std::vector visits(static_cast(rankSize), 0); + int64_t publishCount = 0; + for (int64_t senderId = 0; senderId < kParallelSendLanes; ++senderId) { + for (int64_t peer = senderId; peer < rankSize; peer += kParallelSendLanes) { + if (peer == rank) { + continue; + } + ++visits[static_cast(peer)]; + ++publishCount; + } + } + CheckInt64("parallel publish excludes self", visits[static_cast(rank)], 0); + CheckInt64("parallel publish peer count", publishCount, rankSize - 1); + for (int64_t peer = 0; peer < rankSize; ++peer) { + if (peer != rank) { + CheckInt("parallel publish peer owned once", + visits[static_cast(peer)], 1); + } + } + } + } + CheckInt("parallel publish control step follows release", + static_cast(TileXREp::kEpUrmaCombinePublishDone), + static_cast(TileXREp::kEpUrmaCombineRxBufferReleased + 1)); +} + +void TestUrmaCombineWorkspaceRejectsInvalidConfig() +{ + TileXREp::EpUrmaCombineWorkspaceConfig config {}; + CheckInt("URMA combine null out", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 4, 8, 2, 3, nullptr), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("URMA combine negative sends", + TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 4, 8, 2, -1, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); + CheckInt("URMA combine oversized hidden", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig( + 2, 4, TileXREp::kEpUrmaCombineMaxHidden + 1, 2, 3, &config), + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL); +} + +void TestUrmaCombineDataAsFlagPayloadBoundary() +{ + TileXREp::EpUrmaCombineWorkspaceConfig exact {}; + TileXREp::EpUrmaCombineWorkspaceConfig spill {}; + CheckInt("URMA exact payload boundary", + TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 2, 448, 1, 2, &exact), + TileXR::TILEXR_SUCCESS); + CheckInt("URMA payload spill", + TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(2, 2, 449, 1, 2, &spill), + TileXR::TILEXR_SUCCESS); + CheckInt64("URMA exact comm bytes", exact.commBytes, 480); + CheckInt64("URMA exact block count", exact.blockCount, 1); + CheckInt64("URMA spill comm bytes", spill.commBytes, 512); + CheckInt64("URMA spill block count", spill.blockCount, 2); + CheckInt64("URMA spill route stride", spill.routeStride, 1024); +} + +void TestUrmaCombineDynamicTxDoesNotMoveRemoteRegions() +{ + TileXREp::EpUrmaCombineWorkspaceConfig small {}; + TileXREp::EpUrmaCombineWorkspaceConfig large {}; + CheckInt("URMA small dynamic tx", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(4, 8, 7168, 8, 3, &small), + TileXR::TILEXR_SUCCESS); + CheckInt("URMA large dynamic tx", TileXREp::TileXREpBuildUrmaCombineWorkspaceConfig(4, 8, 7168, 8, 99, &large), + TileXR::TILEXR_SUCCESS); + CheckInt64("URMA dynamic tx keeps rx0", small.rxWindowOffsets[0], large.rxWindowOffsets[0]); + CheckInt64("URMA dynamic tx keeps rx1", small.rxWindowOffsets[1], large.rxWindowOffsets[1]); + CheckInt64("URMA dynamic tx keeps round0", small.roundDoneOffsets[0], large.roundDoneOffsets[0]); + CheckInt64("URMA dynamic tx keeps fixed bytes", small.fixedBytes, large.fixedBytes); + CheckBool("URMA dynamic tx grows required bytes", large.requiredBytes > small.requiredBytes, true); +} + } // namespace int main() @@ -108,5 +252,11 @@ int main() TestDataTypes(); TestWindowConfig(); TestRejectsInvalidConfig(); + TestUrmaCombineWorkspaceConfig(); + TestUrmaCombineStartGateLayout(); + TestParallelRoundPublishPeerShards(); + TestUrmaCombineWorkspaceRejectsInvalidConfig(); + TestUrmaCombineDataAsFlagPayloadBoundary(); + TestUrmaCombineDynamicTxDoesNotMoveRemoteRegions(); return g_failures == 0 ? 0 : 1; } diff --git a/tests/ep/unit/test_tilexr_ep_start_gate_window.cpp b/tests/ep/unit/test_tilexr_ep_start_gate_window.cpp new file mode 100644 index 00000000..4c3e07a1 --- /dev/null +++ b/tests/ep/unit/test_tilexr_ep_start_gate_window.cpp @@ -0,0 +1,45 @@ +#include "ep_urma_combine_start_gate_window.h" + +#include + +namespace { + +int g_failures = 0; + +void Check(bool condition, const char *message) +{ + if (!condition) { + std::cerr << message << std::endl; + ++g_failures; + } +} + +void TestOnlyFirstLaunchAfterSynchronizationUsesGate() +{ + TileXREp::EpUrmaCombineStartGateWindow window; + Check(window.BeginLaunch(), "the first launch must use the start gate"); + Check(!window.BeginLaunch(), "the second launch in a window must skip the start gate"); + Check(!window.BeginLaunch(), "later launches in a window must skip the start gate"); + + window.StreamSynchronized(); + Check(window.BeginLaunch(), "the first launch after synchronization must use the start gate"); + Check(!window.BeginLaunch(), "only one launch after synchronization may use the start gate"); +} + +void TestWindowOneUsesGateForEveryLaunch() +{ + TileXREp::EpUrmaCombineStartGateWindow window; + for (int launch = 0; launch < 4; ++launch) { + Check(window.BeginLaunch(), "window=1 must use the start gate for every launch"); + window.StreamSynchronized(); + } +} + +} // namespace + +int main() +{ + TestOnlyFirstLaunchAfterSynchronizationUsesGate(); + TestWindowOneUsesGateForEveryLaunch(); + return g_failures == 0 ? 0 : 1; +} From 06eaa013bc6d1b7e41f02bcd68b09921b970b82d Mon Sep 17 00:00:00 2001 From: "zhujunan (A)" Date: Thu, 30 Jul 2026 12:15:37 +0800 Subject: [PATCH 2/9] perf(ep): defer round release completion --- .../kernels/tilexr_ep_urma_combine_kernel.cpp | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp index 4948d0d7..a4a1ea71 100644 --- a/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp +++ b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp @@ -1913,19 +1913,9 @@ __aicore__ TILEXR_EP_LOCAL_FUNCTION bool WaitDeferredRoundCredit( const int32_t rankSize = args->rankSize; const int64_t creditLaneCount = rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? rankSize : TileXREp::kEpUrmaCombineSendLaneCount; - const uint64_t expectedReady = EncodeControlValue( - magic, TileXREp::kEpUrmaCombineCreditExpectedReady); - if (senderId == 0) { - StoreControlValue(workspaceGM + roundCreditOffset, expectedReady, - finePerfTrace, perfStats); - } uint64_t pollLoads = 0; uint32_t expectedGeneration = 0; if (senderId < creditLaneCount) { - if (!WaitLocalLines(workspaceGM + roundCreditOffset, 1, expectedReady, - finePerfTrace, perfStats)) { - return false; - } const uint64_t previousRelease = LoadControlValue(workspaceGM + roundDoneOffset + rank * TileXREp::kEpUrmaCombineCacheLineBytes, finePerfTrace, perfStats); const uint32_t previousStep = static_cast(previousRelease); @@ -1938,6 +1928,10 @@ __aicore__ TILEXR_EP_LOCAL_FUNCTION bool WaitDeferredRoundCredit( if (peer == rank || previousRelease == 0) { continue; } + // The same-parity roundDone line is the previous release source. Complete its + // QP before this launch can eventually overwrite that line with the next release. + (void)TileXREp::EpUrmaUDMAQuiet( + args, peer, static_cast(senderId)); GM_ADDR peerRelease = workspaceGM + roundDoneOffset + peer * TileXREp::kEpUrmaCombineCacheLineBytes; while (true) { @@ -2176,6 +2170,7 @@ __aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t PublishRoundShard( static_cast(TileXREp::kEpUrmaCombineCacheLineBytes), qpIdx); ++publishCount; } +#if !TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT for (int32_t peer = static_cast(senderId); peer < rankSize; peer += static_cast(TileXREp::kEpUrmaCombineSendLaneCount)) { if (peer == rank) { @@ -2183,10 +2178,10 @@ __aicore__ TILEXR_EP_LOCAL_FUNCTION uint64_t PublishRoundShard( } (void)TileXREp::EpUrmaUDMAQuiet(args, peer, qpIdx); } - StoreControlValue(workspaceGM + senderDoneOffset + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, EncodeControlValue(magic, TileXREp::kEpUrmaCombinePublishDone), finePerfTrace, perfStats); +#endif if (senderId != 0) { ProfileEnd(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, shardStart); RecordRoundPublishCounters(perfTrace, perfStats, publishCount); @@ -2200,10 +2195,18 @@ __aicore__ TILEXR_EP_LOCAL_FUNCTION bool FinishParallelRoundPublish( uint64_t publishStart, uint64_t sender0PublishCount, GM_ADDR perfTrace, GM_ADDR finePerfTrace, __ubuf__ PerfStats *perfStats) { +#if TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT + (void)workspaceGM; + (void)magic; + (void)senderDoneOffset; + (void)publisherCount; + (void)finePerfTrace; +#else if (!WaitLocalLines(workspaceGM + senderDoneOffset, publisherCount, EncodeControlValue(magic, TileXREp::kEpUrmaCombinePublishDone), finePerfTrace, perfStats)) { return false; } +#endif ProfileEnd(perfTrace, perfStats, PerfStage::ROUND_PUBLISH, publishStart); RecordRoundPublishCounters(perfTrace, perfStats, sender0PublishCount); return true; From ad077ec642292e278e6e3c7e3e62bc4f92b18083 Mon Sep 17 00:00:00 2001 From: "zhujunan (A)" Date: Thu, 30 Jul 2026 12:50:34 +0800 Subject: [PATCH 3/9] perf(ep): reduce deferred RX fan-in --- .../kernels/tilexr_ep_urma_combine_kernel.cpp | 52 ++++++++++--------- tests/ep/unit/test_tilexr_ep_layout.cpp | 26 +++++++++- 2 files changed, 52 insertions(+), 26 deletions(-) diff --git a/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp index a4a1ea71..1437ed1c 100644 --- a/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp +++ b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp @@ -2383,37 +2383,39 @@ extern "C" __global__ __aicore__ void tilexr_ep_urma_combine_kernel(GM_ADDR comm ProfileEnd(perfTrace, perfStats, PerfStage::SEND_TOTAL, sendStart); #if TILEXR_EP_URMA_PARALLEL_ROUND_PUBLISH #if TILEXR_EP_URMA_DEFERRED_ROUND_CREDIT - const uint64_t receiveWaitStart = ProfileBegin(perfTrace); - if (!WaitLocalLineShard(workspaceGM + rxLaneDoneOffset, - TileXREp::kEpUrmaCombinePackLaneCount, senderId, - TileXREp::kEpUrmaCombineSendLaneCount, - EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxLaneDone), - finePerfTrace, perfStats)) { - return; - } - StoreControlValue(workspaceGM + senderDoneOffset + - senderId * TileXREp::kEpUrmaCombineCacheLineBytes, - EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxReleaseShardDone), - finePerfTrace, perfStats); + const int64_t publisherCount = args->rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + args->rankSize : TileXREp::kEpUrmaCombineSendLaneCount; uint64_t publishStart = 0; - if (senderId == 0) { - if (!WaitLocalLines(workspaceGM + senderDoneOffset, - TileXREp::kEpUrmaCombineSendLaneCount, - EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxReleaseShardDone), + if (senderId < publisherCount) { + // Only release publishers participate, reducing sender0's serial fan-in while + // preserving complete coverage of every RX lane. + const uint64_t receiveWaitStart = ProfileBegin(perfTrace); + if (!WaitLocalLineShard(workspaceGM + rxLaneDoneOffset, + TileXREp::kEpUrmaCombinePackLaneCount, senderId, publisherCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxLaneDone), finePerfTrace, perfStats)) { return; } - ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); - ProfileAux(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, 0, - static_cast(TileXREp::kEpUrmaCombinePackLaneCount)); - publishStart = StartParallelRoundPublish(args, workspaceGM, magic, - roundDoneOffset, roundPublishOffset, perfTrace, finePerfTrace, perfStats); - } else { - ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + StoreControlValue(workspaceGM + senderDoneOffset + + senderId * TileXREp::kEpUrmaCombineCacheLineBytes, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxReleaseShardDone), + finePerfTrace, perfStats); + if (senderId == 0) { + if (!WaitLocalLines(workspaceGM + senderDoneOffset, publisherCount, + EncodeControlValue(magic, TileXREp::kEpUrmaCombineRxReleaseShardDone), + finePerfTrace, perfStats)) { + return; + } + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + ProfileAux(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, 0, + static_cast(TileXREp::kEpUrmaCombinePackLaneCount)); + publishStart = StartParallelRoundPublish(args, workspaceGM, magic, + roundDoneOffset, roundPublishOffset, perfTrace, finePerfTrace, perfStats); + } else { + ProfileEnd(perfTrace, perfStats, PerfStage::LOCAL_RX_WAIT, receiveWaitStart); + } } - const int64_t publisherCount = args->rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? - args->rankSize : TileXREp::kEpUrmaCombineSendLaneCount; uint64_t publishCount = 0; if (senderId < publisherCount) { publishCount = PublishRoundShard(args, workspaceGM, magic, diff --git a/tests/ep/unit/test_tilexr_ep_layout.cpp b/tests/ep/unit/test_tilexr_ep_layout.cpp index cbfbb2d4..31b51815 100644 --- a/tests/ep/unit/test_tilexr_ep_layout.cpp +++ b/tests/ep/unit/test_tilexr_ep_layout.cpp @@ -170,7 +170,7 @@ void TestUrmaCombineStartGateLayout() void TestParallelRoundPublishPeerShards() { constexpr int64_t kParallelSendLanes = TileXREp::kEpUrmaCombineSendLaneCount; - const int64_t rankSizes[] = {1, 2, 8, 15, 16, 17, 64}; + const int64_t rankSizes[] = {1, 2, 8, 15, 16, 17, 22, 23, 64, 128}; for (int64_t rankSize : rankSizes) { for (int64_t rank = 0; rank < rankSize; ++rank) { std::vector visits(static_cast(rankSize), 0); @@ -199,6 +199,29 @@ void TestParallelRoundPublishPeerShards() static_cast(TileXREp::kEpUrmaCombineRxBufferReleased + 1)); } +void TestDeferredRoundCreditRxCompletionShards() +{ + const int64_t rankSizes[] = {1, 2, 8, 15, 16, 17, 22, 23, 64, 128}; + for (int64_t rankSize : rankSizes) { + const int64_t publisherCount = rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + rankSize : TileXREp::kEpUrmaCombineSendLaneCount; + std::vector visits( + static_cast(TileXREp::kEpUrmaCombinePackLaneCount), 0); + for (int64_t senderId = 0; senderId < publisherCount; ++senderId) { + for (int64_t lane = senderId; lane < TileXREp::kEpUrmaCombinePackLaneCount; + lane += publisherCount) { + ++visits[static_cast(lane)]; + } + } + CheckInt64("RX completion publisher count", publisherCount, + rankSize < TileXREp::kEpUrmaCombineSendLaneCount ? + rankSize : TileXREp::kEpUrmaCombineSendLaneCount); + for (int visitsForLane : visits) { + CheckInt("RX completion lane owned once", visitsForLane, 1); + } + } +} + void TestUrmaCombineWorkspaceRejectsInvalidConfig() { TileXREp::EpUrmaCombineWorkspaceConfig config {}; @@ -255,6 +278,7 @@ int main() TestUrmaCombineWorkspaceConfig(); TestUrmaCombineStartGateLayout(); TestParallelRoundPublishPeerShards(); + TestDeferredRoundCreditRxCompletionShards(); TestUrmaCombineWorkspaceRejectsInvalidConfig(); TestUrmaCombineDataAsFlagPayloadBoundary(); TestUrmaCombineDynamicTxDoesNotMoveRemoteRegions(); From d1d77fd1b8834072abc2a018f8a1331dc9836222 Mon Sep 17 00:00:00 2001 From: "zhujunan (A)" Date: Thu, 30 Jul 2026 15:44:25 +0800 Subject: [PATCH 4/9] perf(ep): remove redundant TX-ready publish barrier --- src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp index 1437ed1c..d8638885 100644 --- a/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp +++ b/src/ep/kernels/tilexr_ep_urma_combine_kernel.cpp @@ -538,7 +538,6 @@ __aicore__ TILEXR_EP_LOCAL_FUNCTION void PublishTxReadyBatch( GM_ADDR firstLineAddr = workspaceGM + txReadyOffset + firstRoute * TileXREp::kEpUrmaCombineCacheLineBytes; #if TILEXR_EP_URMA_CACHELESS - AscendC::PipeBarrier(); for (int64_t index = 0; index < publishedRouteCount; ++index) { AscendC::WriteGmByPassDCache( reinterpret_cast<__gm__ uint64_t *>(firstLineAddr + From 7b8d4de68ce27c70d7d8afe83d376a7dab0f8028 Mon Sep 17 00:00:00 2001 From: "zhujunan (A)" Date: Thu, 30 Jul 2026 15:55:26 +0800 Subject: [PATCH 5/9] docs(ep): clarify profile timing boundary --- docs/EP_URMA_COMBINE_S22.md | 14 +++++++++++--- .../tilexr_ep_urma_combine_s22_bs128.html | 15 ++++++++------- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/docs/EP_URMA_COMBINE_S22.md b/docs/EP_URMA_COMBINE_S22.md index c4d7fc60..65ad5c17 100644 --- a/docs/EP_URMA_COMBINE_S22.md +++ b/docs/EP_URMA_COMBINE_S22.md @@ -54,6 +54,14 @@ launches. For BS128, H=5120, top-k=6, rank-size=8 and enqueue-window=1, S22 measured 94,889 cycles (94.889 us using 1000 cycles/us). The single retained detailed report is -[BS128 S22 profile](performance/tilexr_ep_urma_combine_s22_bs128.html). The -default proportional timeline excludes the start-gate bar and rebases the axis -to steady-state work; start-gate measurements remain in report metadata. +[BS128 S22 profile](performance/tilexr_ep_urma_combine_s22_bs128.html). Its +charts, stage maxima, heatmaps and kernel KPIs exclude the first-launch start +gate and rebase the axis to steady-state device work. The raw embedded capture +retains the gate samples only for provenance. + +This boundary explains the large host/device timing difference. The host round +starts before launch and ends after stream synchronization, so it includes the +start-gate wait, API launch/synchronization overhead and cross-rank scheduling +skew. Device `kernel_total` begins after the gate releases. The values therefore +describe different intervals and must not be subtracted to infer kernel work; +production comparisons use steady-state `strictKernelCycles`. diff --git a/docs/performance/tilexr_ep_urma_combine_s22_bs128.html b/docs/performance/tilexr_ep_urma_combine_s22_bs128.html index 771f2a66..feef464a 100644 --- a/docs/performance/tilexr_ep_urma_combine_s22_bs128.html +++ b/docs/performance/tilexr_ep_urma_combine_s22_bs128.html @@ -3,7 +3,7 @@ -TileXR EP URMA Combine - S22 P42/S22/QP22 O2 - 42 Pack/Receive + 22 Send - 22 QPs, 1-WQE doorbell, TX-ready batch 1, shared ready flag off, TX ready in data off, TX metadata full prefetch off, TX early ready off, RX sticky on, RX batched MTE2 off, RX Vector ready off, balanced send routes on, RX token round robin sticky, round publish parallel, start gate enabled: first after stream synchronize (executed in representative launch), Q/D v3 - capture +TileXR EP URMA Combine - S22 P42/S22/QP22 O2 - BS128 steady-state profile -
capture

TileXR EP URMA Combine · S22 P42/S22/QP22 O2 · 42 Pack/Receive + 22 Send · 22 QPs, 1-WQE doorbell, TX-ready batch 1, shared ready flag off, TX ready in data off, TX metadata full prefetch off, TX early ready off, RX sticky on, RX batched MTE2 off, RX Vector ready off, balanced send routes on, RX token round robin sticky, round publish parallel, start gate enabled: first after stream synchronize (executed in representative launch), Q/D v3

-
8 ranks · 64 AIV · URMA · 10 paired profile samples · representative launch8 · median cluster kernel 89.43 us (coarse attribution profile) · coarse/fine attribution medians 89.43/97.51 us · Gate executed 10/10, executed-only median 259.99 us / 56 WQEs / 56 DBs
-
bs=128, topK=6, h=5120, selfSendCnt=768, routeStride=5632B, enqueueWindow=1 · INT8 per-route quantization · deterministic randomized routes, seed=20260728 · start gate collapsed from steady-state views
+
steady-state capture

TileXR EP URMA Combine · S22 P42/S22/QP22 O2 · BS128

+
8 ranks · 64 AIV · URMA · 10 paired profile samples · representative launch8 · median cluster kernel 89.43 us (coarse attribution profile) · coarse/fine attribution medians 89.43/97.51 us
+
42 Pack/Receive + 22 Send · 22 QPs · 1-WQE doorbell · Q/D v3 · bs=128, topK=6, h=5120, selfSendCnt=768, routeStride=5632B, enqueueWindow=1
+

Host and device timing boundary

The synchronization preamble is intentionally omitted from every chart, stage maximum, heatmap, and kernel KPI in this report. The device kernel_total interval begins after that preamble has released, while the host round timer starts before launch and ends after stream synchronization. The host value therefore includes the first-launch synchronization wait, API launch and synchronization overhead, and cross-rank scheduling skew; the device value does not.

In this capture the omitted synchronization wait is the dominant reason that host latency is much larger than the approximately 89.43 us device critical path. The two values have different boundaries and must not be subtracted to infer unaccounted kernel work. Production comparisons use the steady-state device timing (strictKernelCycles), not the first-launch host round.

Coarse Stage Maxima Across AIVs

Each value is the maximum coarse sum_us for that stage across the selected rank's AIVs. Maxima can come from different cores, overlap in time, and must not be added as a critical path.

Fine Explicit DCCI Breakdown

Fine diagnostic capture; not the coarse timing launch.

-

64-Core Coarse Timeline

Outlined bars are the exact per-stage maxima listed above. Receive bar width and position are coarse; its internal colors are fine-derived composition, not chronology. The steady-state axis is rebased after the collapsed start gate.

+

64-Core Coarse Timeline

Outlined bars are the exact per-stage maxima listed above. Receive bar width and position are coarse; its internal colors are fine-derived composition, not chronology. The axis begins at the first steady-state device work; the synchronization preamble is not drawn.

Core x Fine Exposed-Time Heatmap

Fine diagnostic capture; absolute durations include instrumentation overhead.

Fine Route Transport AIV Diagnostics

Fine diagnostic capture; route counters and exposed stage time are not from the coarse launch.

Rank Critical Paths

-
The total-only capture provides the timing headline when supplied; it disables stage attribution inside the measured kernel. Coarse bars are per-core activity envelopes for attribution. Each receive_total envelope is divided into RX flag wait, unpack exposed wait, dequant+clear, and output using the same core's fine exposed wall-time. Fine profiling preserves the quant/dequant double-buffer pipeline: hidden asynchronous transfer time is not charged to an exposed-wait stage. Dequant+clear includes next-route unpack submission bookkeeping, but not that transfer's asynchronous in-flight time. Explicit DCCI spans cover combine-kernel UDMACleanCacheLines calls and exclude their surrounding barriers; SQ/CQ cache maintenance inside the UDMA helper remains part of UDMA post/quiet. DCCI category values are accumulated core-time, not additional coarse wall-time. The colored parts show composition, not chronological ordering, and must not be added to infer hardware-engine occupancy. Legacy fine captures without rx_unpack_wait predate this contract. Repeated fine stages use accumulated exposed wall-time in the heatmap, and their durations include material instrumentation overhead. The detailed timeline uses launch8, selected as the coarse sample nearest the median after maximum core per rank and maximum rank per launch. The sample CSV contains all 10 paired launches. Each rank is independently normalized; the report does not assume synchronized device clocks across NPUs. Open perfetto_trace.json for the standard TileXR drill-down.
+
The total-only capture provides the timing headline when supplied; it disables stage attribution inside the measured kernel. Coarse bars are per-core activity envelopes for attribution. The synchronization preamble remains in the embedded raw capture for provenance but is excluded from all rendered steady-state views. Each receive_total envelope is divided into RX flag wait, unpack exposed wait, dequant+clear, and output using the same core's fine exposed wall-time. Fine profiling preserves the quant/dequant double-buffer pipeline: hidden asynchronous transfer time is not charged to an exposed-wait stage. Dequant+clear includes next-route unpack submission bookkeeping, but not that transfer's asynchronous in-flight time. Explicit DCCI spans cover combine-kernel UDMACleanCacheLines calls and exclude their surrounding barriers; SQ/CQ cache maintenance inside the UDMA helper remains part of UDMA post/quiet. DCCI category values are accumulated core-time, not additional coarse wall-time. The colored parts show composition, not chronological ordering, and must not be added to infer hardware-engine occupancy. Legacy fine captures without rx_unpack_wait predate this contract. Repeated fine stages use accumulated exposed wall-time in the heatmap, and their durations include material instrumentation overhead. The detailed timeline uses launch8, selected as the coarse sample nearest the median after maximum core per rank and maximum rank per launch. The sample CSV contains all 10 paired launches. Each rank is independently normalized; the report does not assume synchronized device clocks across NPUs. Open perfetto_trace.json for the standard TileXR drill-down.
- \ No newline at end of file + From ff6b469239b6d628870d27a8f6e58eda268a797b Mon Sep 17 00:00:00 2001 From: "zhujunan (A)" Date: Thu, 30 Jul 2026 17:14:22 +0800 Subject: [PATCH 6/9] docs(ep): refresh final S22 performance evidence --- docs/EP_URMA_COMBINE_S22.md | 4 +++- docs/performance/tilexr_ep_urma_combine_s22_bs128.html | 10 +++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/EP_URMA_COMBINE_S22.md b/docs/EP_URMA_COMBINE_S22.md index 65ad5c17..6efa1520 100644 --- a/docs/EP_URMA_COMBINE_S22.md +++ b/docs/EP_URMA_COMBINE_S22.md @@ -51,7 +51,9 @@ TileXR communicator for rank information and registered-memory exchange. The production comparison metric is profiling-free `strictKernelCycles`, aggregated as max core per rank, max rank per launch, then median across 100 launches. For BS128, H=5120, top-k=6, rank-size=8 and enqueue-window=1, S22 -measured 94,889 cycles (94.889 us using 1000 cycles/us). +measured 88,343.5 cycles (88.3435 us using 1000 cycles/us). The paired baseline +was 88,602.5 cycles, a 0.29% reduction. BS32 improved from 39,753 to 38,541 +cycles (3.05%). The single retained detailed report is [BS128 S22 profile](performance/tilexr_ep_urma_combine_s22_bs128.html). Its diff --git a/docs/performance/tilexr_ep_urma_combine_s22_bs128.html b/docs/performance/tilexr_ep_urma_combine_s22_bs128.html index feef464a..8d38f484 100644 --- a/docs/performance/tilexr_ep_urma_combine_s22_bs128.html +++ b/docs/performance/tilexr_ep_urma_combine_s22_bs128.html @@ -35,22 +35,22 @@
steady-state capture

TileXR EP URMA Combine · S22 P42/S22/QP22 O2 · BS128

-
8 ranks · 64 AIV · URMA · 10 paired profile samples · representative launch8 · median cluster kernel 89.43 us (coarse attribution profile) · coarse/fine attribution medians 89.43/97.51 us
+
8 ranks · 64 AIV · URMA · 10 paired profile samples · representative launch0 · median cluster kernel 83.87 us (coarse attribution profile) · coarse/fine attribution medians 83.87/92.32 us
42 Pack/Receive + 22 Send · 22 QPs · 1-WQE doorbell · Q/D v3 · bs=128, topK=6, h=5120, selfSendCnt=768, routeStride=5632B, enqueueWindow=1
-

Host and device timing boundary

The synchronization preamble is intentionally omitted from every chart, stage maximum, heatmap, and kernel KPI in this report. The device kernel_total interval begins after that preamble has released, while the host round timer starts before launch and ends after stream synchronization. The host value therefore includes the first-launch synchronization wait, API launch and synchronization overhead, and cross-rank scheduling skew; the device value does not.

In this capture the omitted synchronization wait is the dominant reason that host latency is much larger than the approximately 89.43 us device critical path. The two values have different boundaries and must not be subtracted to infer unaccounted kernel work. Production comparisons use the steady-state device timing (strictKernelCycles), not the first-launch host round.

+

Host and device timing boundary

The synchronization preamble is intentionally omitted from every chart, stage maximum, heatmap, and kernel KPI in this report. The device kernel_total interval begins after that preamble has released, while the host round timer starts before launch and ends after stream synchronization. The host value therefore includes the first-launch synchronization wait, API launch and synchronization overhead, and cross-rank scheduling skew; the device value does not.

The omitted wait has a 227.67 us executed-only median in this capture and is the dominant reason that host latency is much larger than the 83.87 us coarse device critical path. The two values have different boundaries and must not be subtracted to infer unaccounted kernel work. Production comparisons use profiling-free strictKernelCycles; the final BS128 median is 88,343.5 cycles.

Coarse Stage Maxima Across AIVs

Each value is the maximum coarse sum_us for that stage across the selected rank's AIVs. Maxima can come from different cores, overlap in time, and must not be added as a critical path.

Fine Explicit DCCI Breakdown

Fine diagnostic capture; not the coarse timing launch.

64-Core Coarse Timeline

Outlined bars are the exact per-stage maxima listed above. Receive bar width and position are coarse; its internal colors are fine-derived composition, not chronology. The axis begins at the first steady-state device work; the synchronization preamble is not drawn.

Core x Fine Exposed-Time Heatmap

Fine diagnostic capture; absolute durations include instrumentation overhead.

Fine Route Transport AIV Diagnostics

Fine diagnostic capture; route counters and exposed stage time are not from the coarse launch.

Rank Critical Paths

-
The total-only capture provides the timing headline when supplied; it disables stage attribution inside the measured kernel. Coarse bars are per-core activity envelopes for attribution. The synchronization preamble remains in the embedded raw capture for provenance but is excluded from all rendered steady-state views. Each receive_total envelope is divided into RX flag wait, unpack exposed wait, dequant+clear, and output using the same core's fine exposed wall-time. Fine profiling preserves the quant/dequant double-buffer pipeline: hidden asynchronous transfer time is not charged to an exposed-wait stage. Dequant+clear includes next-route unpack submission bookkeeping, but not that transfer's asynchronous in-flight time. Explicit DCCI spans cover combine-kernel UDMACleanCacheLines calls and exclude their surrounding barriers; SQ/CQ cache maintenance inside the UDMA helper remains part of UDMA post/quiet. DCCI category values are accumulated core-time, not additional coarse wall-time. The colored parts show composition, not chronological ordering, and must not be added to infer hardware-engine occupancy. Legacy fine captures without rx_unpack_wait predate this contract. Repeated fine stages use accumulated exposed wall-time in the heatmap, and their durations include material instrumentation overhead. The detailed timeline uses launch8, selected as the coarse sample nearest the median after maximum core per rank and maximum rank per launch. The sample CSV contains all 10 paired launches. Each rank is independently normalized; the report does not assume synchronized device clocks across NPUs. Open perfetto_trace.json for the standard TileXR drill-down.
+
The total-only capture provides the timing headline when supplied; it disables stage attribution inside the measured kernel. Coarse bars are per-core activity envelopes for attribution. The synchronization preamble remains in the embedded raw capture for provenance but is excluded from all rendered steady-state views. Each receive_total envelope is divided into RX flag wait, unpack exposed wait, dequant+clear, and output using the same core's fine exposed wall-time. Fine profiling preserves the quant/dequant double-buffer pipeline: hidden asynchronous transfer time is not charged to an exposed-wait stage. Dequant+clear includes next-route unpack submission bookkeeping, but not that transfer's asynchronous in-flight time. Explicit DCCI spans cover combine-kernel UDMACleanCacheLines calls and exclude their surrounding barriers; SQ/CQ cache maintenance inside the UDMA helper remains part of UDMA post/quiet. DCCI category values are accumulated core-time, not additional coarse wall-time. The colored parts show composition, not chronological ordering, and must not be added to infer hardware-engine occupancy. Legacy fine captures without rx_unpack_wait predate this contract. Repeated fine stages use accumulated exposed wall-time in the heatmap, and their durations include material instrumentation overhead. The detailed timeline uses launch0, selected as the coarse sample nearest the median after maximum core per rank and maximum rank per launch. The sample CSV contains all 10 paired launches. Each rank is independently normalized; the report does not assume synchronized device clocks across NPUs. Open perfetto_trace.json for the standard TileXR drill-down.