From a845d8459612c3f634e3e89fd57000a999a294e6 Mon Sep 17 00:00:00 2001 From: chaowick Date: Thu, 13 Aug 2026 20:34:05 +0800 Subject: [PATCH] feat(moonep): fuse dispatch payload epoch --- docs/moonep/MINDSPEED_DEBUGGING_EXPERIENCE.md | 34 + .../2026-08-13-moonep-dispatch-fused-epoch.md | 49 ++ ...8-13-moonep-dispatch-fused-epoch-design.md | 93 +++ .../dispatch/urma/common/dispatch_profile.h | 1 + .../dispatch/urma/common/dispatch_wqe_batch.h | 45 +- .../dispatch/urma/host/dispatch_host.cpp | 34 +- src/moonep/dispatch/urma/host/dispatch_host.h | 3 + .../dispatch/urma/host/dispatch_launch.cpp | 70 +- .../dispatch/urma/host/dispatch_launch.h | 46 +- .../dispatch/urma/host/dispatch_layout.cpp | 24 +- .../kernels/tilexr_moonep_dispatch_kernel.cpp | 720 ++++++++++++------ tests/moonep/CMakeLists.txt | 31 +- .../test_dispatch_hot_loop_diagnostics.py | 289 +++++++ tests/moonep/python/test_moonep_modes.py | 6 +- .../test_tilexr_moonep_dispatch_layout.cpp | 61 +- .../test_tilexr_moonep_dispatch_schedule.cpp | 24 + .../test_tilexr_moonep_dispatch_urma_host.cpp | 263 +++++++ ...est_tilexr_moonep_dispatch_urma_launch.cpp | 215 ++++++ .../test_tilexr_moonep_kernel_sources.cpp | 49 ++ tools/moonep/benchmark.py | 17 +- tools/moonep/dispatch_hot_loop.py | 188 +++-- tools/moonep/launcher.py | 6 +- 22 files changed, 1891 insertions(+), 377 deletions(-) create mode 100644 docs/plans/2026-08-13-moonep-dispatch-fused-epoch.md create mode 100644 docs/specs/2026-08-13-moonep-dispatch-fused-epoch-design.md create mode 100644 tests/moonep/python/test_dispatch_hot_loop_diagnostics.py create mode 100644 tests/moonep/unit/test_tilexr_moonep_dispatch_urma_host.cpp create mode 100644 tests/moonep/unit/test_tilexr_moonep_dispatch_urma_launch.cpp diff --git a/docs/moonep/MINDSPEED_DEBUGGING_EXPERIENCE.md b/docs/moonep/MINDSPEED_DEBUGGING_EXPERIENCE.md index 3932778..fd4a72a 100644 --- a/docs/moonep/MINDSPEED_DEBUGGING_EXPERIENCE.md +++ b/docs/moonep/MINDSPEED_DEBUGGING_EXPERIENCE.md @@ -182,6 +182,40 @@ reset 等单一变量。一次同时修改 Kernel、Host、timeout 和路由， - 不要因某次补丁通过完整模型就跳过最小 reproducer；最小 reproducer 才能证明因果。 - 不要删除被推翻的假设记录。保留否定证据可以防止后续重复猜测。 +## Dispatch V2 fused epoch 约束与验证边界 + +一次 paired `TileXRMoonEpDispatchV2` 应只有一个 magic、一次 AICore launch 和 +一个 UDMA epoch。Hidden 与 RouteWeight 必须使用互不重叠的 source 和双 scratch； +同一 route 在同一逻辑 QP 上按 Hidden、Weight 顺序发 WQE，ordered completion 必须 +排在该 QP 的全部 payload WQE 之后。completion、group credit、CQ reclaim 和 final +quiet 每个 fused epoch 只执行一轮。发生上游或设备错误后可以停止 payload work，但 +不能跳过 signal-only completion、incoming wait/credit 和最终 status/quiet 收敛，否则 +健康 rank 会退化为超时发现错误。 + +诊断仍保留独立 Hidden/Weight Profile 与 DFX，但共享阶段只能有一个 owner。paired +以 Weight record 承载 route scan、flag wait、credit、CQ 和 quiet；Hidden record 的共享 +耗时为零，并由 kernel status 的 fused feature bit 明确标识。profiling OFF 的延迟运行与 +profiling ON 的阶段分析必须分开，不能把两个 payload record 伪装成两轮独立通信。 + +实机验证时还要注意以下边界： + +- Host 环境值是 `group_credit`，不是 `group-credit`；错误拼写会在 launch 前返回 `-3`。 +- grouped/group-credit 只适用于 vector route selection。小于 64 routes、非 2 的幂 + `NvS` 或 UB/vector 不满足条件的 shape 应使用 legacy scalar-tiled 路径；同步返回 + `-6` 不是 UDMA 数据面失败。 +- padded zero-fill 要验证完整输出 tensor，不能只比较有效 slot。若 shape 不满足 grouped + 条件，用 legacy 路径证明 Hidden 和 Weight 的空 slot 都为零。 +- shared-QP 不能仅凭配置名推断。应同时确认 runtime 调用 shared-QP-domain 初始化、日志 + 显示 domain 启用，且 communicator `extraFlag` 含 UDMA 与 `UDMA_SHARED_QP`；本次 + Ascend950PR 单机 8 rank 证据对应固定 32 QP shared domain。 +- 参考形状 `S=128,K=16,H=3584,NvS=2048` 的不重叠布局仍为 30 MiB；其他 shape + 必须使用 checked add/multiply 计算真实容量，并在扩大 workspace 绑定后重新检查全部 + active region 与 common tail 边界。 + +这些 Host/mock/source guard 只能证明 ABI、布局和源码协议不变量。只有相同 CANN、设备、 +拓扑上的 HCCL baseline 和 Ascend950 实机逐元素多轮结果，才能证明 UDMA 数据面；性能结论 +还必须使用同 shape、同 pair 模式、同 warmup/迭代和独立 profiling-off 构建做 A/B。 + ## 当前实现状态说明 本文记录的是已验证经验，不代表所有修复都已经进入 `main`。截至 2026-08-12： diff --git a/docs/plans/2026-08-13-moonep-dispatch-fused-epoch.md b/docs/plans/2026-08-13-moonep-dispatch-fused-epoch.md new file mode 100644 index 0000000..6a8c54d --- /dev/null +++ b/docs/plans/2026-08-13-moonep-dispatch-fused-epoch.md @@ -0,0 +1,49 @@ +# MoonEP Dispatch Fused Epoch Implementation Plan + +## Scope + +Implement the approved design in +`docs/specs/2026-08-13-moonep-dispatch-fused-epoch-design.md`. Preserve the +public MoonEP and Python APIs. Do not modify comparison-only `reference/` code or +add Weight-only Dispatch. + +## Tasks + +### 1. Layout And Host Contract + +Update the URMA layout to append Hidden and Weight active regions rather than +overlay them. Extend launch parameters and the registered direct-Kernel ABI to +carry both pointer pairs and active offsets. Refactor `RunDispatchUrma` so a +paired call validates both descriptors but invokes the launcher once. Add layout +and Host/launch tests that prove one magic and one launch. + +### 2. Fused WQE Accounting + +Extend the common batch helpers and UB WQE builder so one selected route emits +one Hidden WQE and an optional immediately following Weight WQE on the same +logical QP. Preserve completion WQE ordering, SQ capacity reserve, ring wrap, +CQE final-BB accounting, and grouped staged-doorbell behavior. Add Host-testable +helpers for WQE counts and selected indices. + +### 3. One-Epoch Kernel + +Refactor the current single-payload Kernel body into one fused execution path. +Stage both sources before the first local barrier, scan routes once, service +both local and remote payloads, wait once, copy both outputs, and quiet once. +Keep communication convergence on device-detected errors and preserve sticky +status. Ensure UB resources are reset only after their producers/consumers have +completed. + +### 4. Diagnostics And Tools + +Version or feature-mark fused diagnostics without breaking existing record-size +consumers. Preserve separate per-payload records and final paired Kernel status. +Update `dispatch_hot_loop.py`, reporting tests, source guards, and API tests. + +### 5. Validation And Documentation + +Run focused C++ and Python tests, target CANN 9.1 Host/Kernel build, and inspect +the final diff. On `141.61.49.195`, establish NPU/CANN/source/binary provenance, +run matching HCCL Test, deploy to a task-specific directory, and execute the +approved correctness/stability/performance ladder. Record reusable fused-epoch +workspace and completion-order lessons in the maintained Dispatch design. diff --git a/docs/specs/2026-08-13-moonep-dispatch-fused-epoch-design.md b/docs/specs/2026-08-13-moonep-dispatch-fused-epoch-design.md new file mode 100644 index 0000000..584bad4 --- /dev/null +++ b/docs/specs/2026-08-13-moonep-dispatch-fused-epoch-design.md @@ -0,0 +1,93 @@ +# MoonEP Dispatch Fused Epoch Design + +## Goal + +Execute the Hidden payload and optional FP32 route-weight payload of one +`TileXRMoonEpDispatchV2` call in one registered AICore Kernel launch and one UDMA +communication epoch. Hidden-only calls remain one launch. The public C, Torch, +MindSpeed, asynchronous-event, and zero-copy contracts do not change. + +## Direct-Launch Contract + +The Host validates the mandatory Hidden descriptors and the optional paired +route-weight descriptors, builds one disjoint registered-workspace layout, +obtains one communicator magic, and calls `rtKernelLaunchWithFlagV2` once. The +Kernel ABI carries both payload pointer pairs, both active layouts, separate +diagnostic offsets, one plan status, and one communication configuration. + +Standalone route-weight Dispatch is not added. The mandatory primary payload +remains FP16 or BF16 Hidden `[S,H] -> [NvS,H]`; the optional pair remains FP32 +`[S,K] -> [NvS]`. + +## Workspace + +The registered region contains, in this order: + +1. Hidden source `[S,H]` and two Hidden receive scratch slots `[NvS,H]`; +2. Weight source `[S,K]` and two Weight receive scratch slots `[NvS]`; +3. shared completion flags and signal source; +4. separate Hidden and Weight Profile arrays; +5. separate Hidden and Weight DFX arrays; +6. shared Kernel status. + +All offsets use checked 64-bit arithmetic and 64-byte internal alignment. The +whole registration is rounded to 2 MiB. Binding a larger registered allocation +moves the common tail while preserving disjoint active regions. For +`S=128,K=16,H=3584,NvS=2048`, the result remains 30 MiB after registration +alignment. + +## Kernel Flow + +1. Validate uniform scalar, pointer, layout, route, transport, and core-count + arguments without partially entering the communication protocol. +2. Cooperatively stage Hidden source and, when present, Weight source into their + disjoint registered regions; converge through one `SyncAll`. +3. Load and select the route plan once. For every selected route and logical QP, + append the Hidden WQE followed by the optional Weight WQE. Hidden uses + `sourceRow=route/K` and `H*2` bytes; Weight uses `sourceRow=route` and four + bytes. Both target the same decoded rank and slot in different scratch + regions. +4. Append the ordered completion WQE after both payload WQEs for that peer/QP. + Build every WQE in UB, publish complete batches to SQ through MTE3, then ring + doorbells only with `st_dev` after MTE3 completion. +5. Use one completion-flag exchange, grouped credit progression, CQ recovery, + and final quiet for the fused epoch. Errors after protocol entry continue the + bounded convergence path so peers are not stranded. +6. After all receive completions and local-core convergence, zero-fill and copy + Hidden output, then optional Weight output, from their respective scratch + slot. Output-copy UB is reused only after explicit pipeline completion/reset. +7. Preserve sticky first-error publication in `plan.status` and write complete + per-payload diagnostics plus shared Kernel status. + +## Diagnostics + +Hidden and Weight retain separate Profile and DFX records for compatibility and +payload-specific byte/output counters. A new diagnostic feature bit identifies +a fused epoch. Shared route-selection, flag-wait, credit, CQ, and quiet work is +owned by the final active payload record: Hidden for hidden-only calls and +RouteWeight for paired calls. The non-owning Hidden paired record reports zero +for shared-stage durations instead of duplicating time. Both records carry the +same magic and failure context. Paired shared Kernel status keeps +`payloadMode=RouteWeight`; hidden-only keeps `payloadMode=Hidden`. + +## Non-Goals + +- Weight-only public Dispatch. +- `WRITE_WITH_NOTIFY` as a weight carrier. +- Multi-SGE or packed Hidden/Weight WQEs. +- Public ABI or MindSpeed adapter changes. +- Removing required local barriers or weakening timeout/error convergence. + +## Verification + +Host and layout tests prove disjoint ranges, checked arithmetic, one magic, and +one launch. WQE helper tests prove paired counts, address/length selection, QP +split, signal ordering, batching, wrap, and CQ accounting. Source guards retain +UB-only WQE construction, MTE3 SQ publication, and `st_dev` doorbells. Python +tests preserve the one-FFI-call contract and parse fused diagnostics. + +The target CANN 9.1 build must compile Host and Kernel. Hardware validation on +`141.61.49.195` starts with the matching official HCCL Test and then covers +single-rank, two-rank, and full-host Hidden/paired exactness, repeated rounds, +alternating plans, grouped/group-credit/shared-QP configurations where +supported, and profiling-off/on `pair` A/B evidence. diff --git a/src/moonep/dispatch/urma/common/dispatch_profile.h b/src/moonep/dispatch/urma/common/dispatch_profile.h index 1723371..0eaad22 100644 --- a/src/moonep/dispatch/urma/common/dispatch_profile.h +++ b/src/moonep/dispatch/urma/common/dispatch_profile.h @@ -11,6 +11,7 @@ constexpr uint32_t kDispatchKernelStatusMarker = 0x54584453U; // TXDS constexpr uint16_t kDispatchDiagnosticVersion = 3U; constexpr uint64_t kDispatchKernelStatusFeatureDfxEnabled = 1U << 0; constexpr uint64_t kDispatchKernelStatusFeatureProfilingEnabled = 1U << 1; +constexpr uint64_t kDispatchKernelStatusFeatureFusedEpoch = 1U << 2; enum DispatchSelectMode : uint32_t { kDispatchSelectScalarTiled = 0, diff --git a/src/moonep/dispatch/urma/common/dispatch_wqe_batch.h b/src/moonep/dispatch/urma/common/dispatch_wqe_batch.h index 671fdce..5bb5034 100644 --- a/src/moonep/dispatch/urma/common/dispatch_wqe_batch.h +++ b/src/moonep/dispatch/urma/common/dispatch_wqe_batch.h @@ -23,6 +23,44 @@ constexpr uint32_t kDispatchSharedQpCoreCount = 16U; constexpr uint32_t kDispatchSharedQpCount = kDispatchQpCount * kDispatchSharedQpCoreCount; +TILEXR_MOONEP_WQE_BATCH_INLINE uint32_t DispatchPayloadWqesPerRoute( + bool hasWeight) +{ + return hasWeight ? 2U : 1U; +} + +TILEXR_MOONEP_WQE_BATCH_INLINE bool DispatchDataWqeCount( + uint64_t routeCount, bool hasWeight, uint64_t &wqeCount) +{ + const uint32_t perRoute = DispatchPayloadWqesPerRoute(hasWeight); + if (routeCount > UINT64_MAX / perRoute) { + wqeCount = 0U; + return false; + } + wqeCount = routeCount * perRoute; + return true; +} + +TILEXR_MOONEP_WQE_BATCH_INLINE uint32_t DispatchDataTaskRouteIndex( + uint32_t dataTask, bool hasWeight) +{ + return hasWeight ? dataTask / 2U : dataTask; +} + +TILEXR_MOONEP_WQE_BATCH_INLINE bool DispatchDataTaskIsWeight( + uint32_t dataTask, bool hasWeight) +{ + return hasWeight && (dataTask & 1U) != 0U; +} + +TILEXR_MOONEP_WQE_BATCH_INLINE bool DispatchSignalFitsAfterData( + uint64_t remainingRoutes, bool hasWeight, uint32_t availableWqes) +{ + uint64_t remainingDataWqes = 0U; + return DispatchDataWqeCount(remainingRoutes, hasWeight, + remainingDataWqes) && remainingDataWqes + 1U <= availableWqes; +} + TILEXR_MOONEP_WQE_BATCH_INLINE bool DispatchQpCountSupported( uint32_t availableQpCount, bool sharedQp = false) { @@ -159,9 +197,12 @@ TILEXR_MOONEP_WQE_BATCH_INLINE uint32_t DispatchQpSelectedIndex( TILEXR_MOONEP_WQE_BATCH_INLINE bool DispatchPeerWqesStreamable( uint64_t routeCount, uint32_t sqEntryCount, - uint32_t reserve = kDispatchSqPollReserve) + uint32_t reserve = kDispatchSqPollReserve, bool hasWeight = false) { - if (routeCount > UINT32_MAX || sqEntryCount <= reserve) { + uint64_t dataWqeCount = 0U; + if (routeCount > UINT32_MAX || + !DispatchDataWqeCount(routeCount, hasWeight, dataWqeCount) || + sqEntryCount <= reserve) { return false; } return sqEntryCount - reserve >= kDispatchWqeBatchCapacity; diff --git a/src/moonep/dispatch/urma/host/dispatch_host.cpp b/src/moonep/dispatch/urma/host/dispatch_host.cpp index 6e1440e..9b05d9f 100644 --- a/src/moonep/dispatch/urma/host/dispatch_host.cpp +++ b/src/moonep/dispatch/urma/host/dispatch_host.cpp @@ -55,7 +55,7 @@ bool ResolveDispatchPeerConfig(DispatchPeerConfig &config) } bool DispatchVectorBatchShapeSupported(uint64_t routeCount, - uint64_t destinationCapacity) + uint64_t destinationCapacity, bool hasWeight) { constexpr uint64_t vectorCompareMinElements = 256U / sizeof(int32_t); return routeCount >= vectorCompareMinElements && @@ -65,17 +65,19 @@ bool DispatchVectorBatchShapeSupported(uint64_t routeCount, destinationCapacity <= UINT32_MAX && (destinationCapacity & (destinationCapacity - 1U)) == 0U && DispatchPeerWqesStreamable(routeCount, - TileXR::TILEXR_UDMA_SQ_BB_COUNT); + TileXR::TILEXR_UDMA_SQ_BB_COUNT, + kDispatchSqPollReserve, hasWeight); } int ValidateDispatchPeerConfig(const DispatchPeerConfig &config, const TileXR::CommArgs &commArgs, uint64_t routeCount, - uint64_t destinationCapacity) + uint64_t destinationCapacity, bool hasWeight) { if (!DispatchPeerModeUsesGroups(static_cast(config.mode))) { return TILEXR_MOONEP_SUCCESS; } - if (!DispatchVectorBatchShapeSupported(routeCount, destinationCapacity)) { + if (!DispatchVectorBatchShapeSupported( + routeCount, destinationCapacity, hasWeight)) { return TILEXR_MOONEP_ERROR_NOT_SUPPORTED; } if (commArgs.rankSize > 1 && @@ -378,7 +380,8 @@ static int RunDispatchUrma(const TileXRMoonEpDispatchArgsV1 *args, } ret = ValidateDispatchPeerConfig(peerConfig, *commArgs, static_cast(layout.routeCount), - static_cast(layout.destinationCapacity)); + static_cast(layout.destinationCapacity), + args->routeWeightsSk != nullptr); if (ret != TILEXR_MOONEP_SUCCESS) { return ret; } @@ -454,6 +457,12 @@ static int RunDispatchUrma(const TileXRMoonEpDispatchArgsV1 *args, params.groupWidth = peerConfig.groupWidth; params.zeroFillRangeCount = args->plan->e + args->plan->b; params.layout = layout; + params.hiddenInput = args->hiddenSh->data; + params.hiddenOutput = args->hiddenNvsh->data; + params.weightInput = args->routeWeightsSk == nullptr ? nullptr : + args->routeWeightsSk->data; + params.weightOutput = args->routeWeightsNvs == nullptr ? nullptr : + args->routeWeightsNvs->data; const bool statusResetEnqueued = resetStatus || (args->flags & TILEXR_MOONEP_FLAG_RESET_STATUS) != 0; @@ -462,21 +471,6 @@ static int RunDispatchUrma(const TileXRMoonEpDispatchArgsV1 *args, return TILEXR_MOONEP_ERROR_INTERNAL; } - params.input = args->hiddenSh->data; - params.output = args->hiddenNvsh->data; - params.mode = DispatchPayloadMode::Hidden; - ret = MapLaunchStatus(TileXRMoonEpLaunchDispatchUrmaKernel(params)); - if (ret != TILEXR_MOONEP_SUCCESS || args->routeWeightsSk == nullptr) { - if (ret != TILEXR_MOONEP_SUCCESS && statusResetEnqueued && - aclrtSynchronizeStream(stream) != ACL_SUCCESS) { - return TILEXR_MOONEP_ERROR_INTERNAL; - } - return ret; - } - - params.input = args->routeWeightsSk->data; - params.output = args->routeWeightsNvs->data; - params.mode = DispatchPayloadMode::RouteWeight; ret = MapLaunchStatus(TileXRMoonEpLaunchDispatchUrmaKernel(params)); if (ret != TILEXR_MOONEP_SUCCESS && statusResetEnqueued && aclrtSynchronizeStream(stream) != ACL_SUCCESS) { diff --git a/src/moonep/dispatch/urma/host/dispatch_host.h b/src/moonep/dispatch/urma/host/dispatch_host.h index 92aebbf..804ed43 100644 --- a/src/moonep/dispatch/urma/host/dispatch_host.h +++ b/src/moonep/dispatch/urma/host/dispatch_host.h @@ -12,6 +12,9 @@ int TileXRMoonEpQueryDispatchUrmaWorkspace(TileXRCommPtr comm, int64_t s, int TileXRMoonEpRunDispatchUrmaV1(const TileXRMoonEpDispatchArgsV1 *args, aclrtStream stream); +int TileXRMoonEpRunDispatchUrmaV2(const TileXRMoonEpDispatchArgsV2 *args, + aclrtStream stream); + } // namespace TileXRMoonEp #endif // TILEXR_MOONEP_DISPATCH_URMA_HOST_H diff --git a/src/moonep/dispatch/urma/host/dispatch_launch.cpp b/src/moonep/dispatch/urma/host/dispatch_launch.cpp index e28944d..9c4d0f1 100644 --- a/src/moonep/dispatch/urma/host/dispatch_launch.cpp +++ b/src/moonep/dispatch/urma/host/dispatch_launch.cpp @@ -39,37 +39,6 @@ int gDispatchRegistrationStatus = TileXR::TILEXR_ERROR_NOT_INITIALIZED; void *gDispatchBinaryHandle = nullptr; uint8_t gDispatchKernelStub = 0; -struct DispatchKernelArgs { - GM_ADDR commArgs; - GM_ADDR input; - GM_ADDR dst; - GM_ADDR zeroFillRanges; - GM_ADDR workspace; - GM_ADDR output; - GM_ADDR planStatus; - uint64_t profileOffset; - uint64_t scratchOffset; - uint64_t completionFlagsOffset; - uint64_t signalOffset; - uint64_t dfxOffset; - uint64_t kernelStatusOffset; - int64_t s; - int64_t k; - int64_t h; - int64_t routeCount; - int64_t destinationCapacity; - int64_t zeroFillRangeCount; - uint64_t rowBytes; - uint64_t payloadMode; - int64_t magic; - uint64_t completionTimeoutTicks; - uint64_t peerMode; - uint64_t groupWidth; -}; - -static_assert(sizeof(DispatchKernelArgs) == 25U * sizeof(uint64_t), - "MoonEP Dispatch Host/Kernel ABI changed"); - int EnsureDispatchKernelRegistered() { std::lock_guard guard(gDispatchRegistrationMutex); @@ -152,24 +121,24 @@ uint64_t TileXRMoonEpDispatchCompletionTimeoutTicks() int TileXRMoonEpLaunchDispatchUrmaKernel(const DispatchUrmaLaunchParams ¶ms) { - const DispatchUrmaActiveLayout *active = TileXRMoonEpGetActiveDispatchUrmaLayout( - params.layout, params.mode); - if (active == nullptr) { + const bool hasWeight = params.weightInput != nullptr; + if (params.hiddenInput == nullptr || params.hiddenOutput == nullptr || + (hasWeight != (params.weightOutput != nullptr))) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } const uint32_t aivCoreCount = ResolveDispatchAivCoreCount(); if (aivCoreCount == 0U) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } + if (!DispatchPeerModeValid(static_cast(params.peerMode)) || + !DispatchGroupWidthValid(params.groupWidth)) { + return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; + } int64_t magic = 0; int ret = TileXRCommNextMagic(params.comm, &magic); if (ret != TileXR::TILEXR_SUCCESS) { return ret; } - if (!DispatchPeerModeValid(static_cast(params.peerMode)) || - !DispatchGroupWidthValid(params.groupWidth)) { - return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; - } if (DispatchPeerModeUsesCredit(static_cast(params.peerMode))) { const uint32_t groupCount = DispatchGroupedGroupCount( params.layout.rankSize, params.groupWidth); @@ -184,29 +153,36 @@ int TileXRMoonEpLaunchDispatchUrmaKernel(const DispatchUrmaLaunchParams ¶ms) return ret; } - const bool hidden = params.mode == DispatchPayloadMode::Hidden; DispatchKernelArgs args { params.commArgs, - reinterpret_cast(const_cast(params.input)), + reinterpret_cast(const_cast(params.hiddenInput)), + reinterpret_cast(const_cast(params.weightInput)), reinterpret_cast(const_cast(params.dst)), reinterpret_cast(const_cast(params.zeroFillRanges)), static_cast(params.workspace), - static_cast(params.output), + static_cast(params.hiddenOutput), + static_cast(params.weightOutput), reinterpret_cast(params.planStatus), - hidden ? params.layout.hiddenProfileOffset : params.layout.weightProfileOffset, - active->scratchOffset, + params.layout.hidden.sourceOffset, + params.layout.hidden.scratchOffset, + params.layout.hidden.rowBytes, + params.layout.weight.sourceOffset, + params.layout.weight.scratchOffset, + params.layout.weight.rowBytes, params.layout.completionFlagsOffset, params.layout.signalOffset, - hidden ? params.layout.hiddenDfxOffset : params.layout.weightDfxOffset, + params.layout.hiddenProfileOffset, + params.layout.weightProfileOffset, + params.layout.hiddenDfxOffset, + params.layout.weightDfxOffset, params.layout.kernelStatusOffset, params.layout.s, params.layout.k, - hidden ? params.layout.h : 1, + params.layout.h, params.layout.routeCount, params.layout.destinationCapacity, params.zeroFillRangeCount, - active->rowBytes, - static_cast(params.mode), + hasWeight ? 1U : 0U, magic, TileXRMoonEpDispatchCompletionTimeoutTicks(), static_cast(params.peerMode), diff --git a/src/moonep/dispatch/urma/host/dispatch_launch.h b/src/moonep/dispatch/urma/host/dispatch_launch.h index 5a42627..6eabba2 100644 --- a/src/moonep/dispatch/urma/host/dispatch_launch.h +++ b/src/moonep/dispatch/urma/host/dispatch_launch.h @@ -10,21 +10,61 @@ namespace TileXRMoonEp { struct DispatchUrmaLaunchParams { GM_ADDR commArgs = nullptr; - const void *input = nullptr; + const void *hiddenInput = nullptr; + const void *weightInput = nullptr; const int32_t *dst = nullptr; const int32_t *zeroFillRanges = nullptr; void *workspace = nullptr; - void *output = nullptr; + void *hiddenOutput = nullptr; + void *weightOutput = nullptr; int32_t *planStatus = nullptr; TileXRCommPtr comm = nullptr; aclrtStream stream = nullptr; - DispatchPayloadMode mode = DispatchPayloadMode::Hidden; DispatchPeerMode peerMode = DispatchPeerMode::Legacy; uint32_t groupWidth = kDispatchDefaultGroupWidth; int64_t zeroFillRangeCount = 0; MoonEpDispatchUrmaLayout layout {}; }; +struct DispatchKernelArgs { + GM_ADDR commArgs; + GM_ADDR hiddenInput; + GM_ADDR weightInput; + GM_ADDR dst; + GM_ADDR zeroFillRanges; + GM_ADDR workspace; + GM_ADDR hiddenOutput; + GM_ADDR weightOutput; + GM_ADDR planStatus; + uint64_t hiddenSourceOffset; + uint64_t hiddenScratchOffset; + uint64_t hiddenRowBytes; + uint64_t weightSourceOffset; + uint64_t weightScratchOffset; + uint64_t weightRowBytes; + uint64_t completionFlagsOffset; + uint64_t signalOffset; + uint64_t hiddenProfileOffset; + uint64_t weightProfileOffset; + uint64_t hiddenDfxOffset; + uint64_t weightDfxOffset; + uint64_t kernelStatusOffset; + int64_t s; + int64_t k; + int64_t h; + int64_t routeCount; + int64_t destinationCapacity; + int64_t zeroFillRangeCount; + uint64_t hasWeight; + int64_t magic; + uint64_t completionTimeoutTicks; + uint64_t peerMode; + uint64_t groupWidth; +}; + +static_assert(sizeof(DispatchKernelArgs) == 33U * sizeof(uint64_t), + "MoonEP Dispatch Host/Kernel ABI changed"); + int TileXRMoonEpLaunchDispatchUrmaKernel(const DispatchUrmaLaunchParams ¶ms); uint64_t TileXRMoonEpDispatchCompletionTimeoutTicks(); diff --git a/src/moonep/dispatch/urma/host/dispatch_layout.cpp b/src/moonep/dispatch/urma/host/dispatch_layout.cpp index 26a76e0..fad8492 100644 --- a/src/moonep/dispatch/urma/host/dispatch_layout.cpp +++ b/src/moonep/dispatch/urma/host/dispatch_layout.cpp @@ -48,23 +48,24 @@ bool AppendBytes(uint64_t bytes, uint64_t *cursor, uint64_t *offset) } bool BuildActiveLayout(uint64_t sourceRows, uint64_t destinationCapacity, - uint64_t rowBytes, DispatchUrmaActiveLayout *out) + uint64_t rowBytes, uint64_t *cursor, DispatchUrmaActiveLayout *out) { - if (out == nullptr || sourceRows == 0 || destinationCapacity == 0 || rowBytes == 0) { + if (cursor == nullptr || out == nullptr || sourceRows == 0 || + destinationCapacity == 0 || rowBytes == 0) { return false; } DispatchUrmaActiveLayout next {}; - uint64_t cursor = 0; + const uint64_t begin = *cursor; next.rowBytes = rowBytes; if (!CheckedMul(sourceRows, rowBytes, &next.sourceBytes) || - !AppendBytes(next.sourceBytes, &cursor, &next.sourceOffset) || + !AppendBytes(next.sourceBytes, cursor, &next.sourceOffset) || !CheckedMul(destinationCapacity, rowBytes, &next.scratchSlotBytes) || !CheckedMul(next.scratchSlotBytes, kDispatchScratchBufferCount, &next.scratchBytes) || - !AppendBytes(next.scratchBytes, &cursor, &next.scratchOffset)) { + !AppendBytes(next.scratchBytes, cursor, &next.scratchOffset)) { return false; } - next.activeDataBytes = cursor; + next.activeDataBytes = *cursor - begin; *out = next; return true; } @@ -103,16 +104,15 @@ int TileXRMoonEpBuildDispatchUrmaLayout(int64_t rankSize, int64_t s, int64_t k, } MoonEpDispatchUrmaLayout next {}; + uint64_t cursor = 0; if (!BuildActiveLayout(static_cast(s), static_cast(destinationCapacity), hiddenRowBytes, - &next.hidden) || + &cursor, &next.hidden) || !BuildActiveLayout(routeCount, static_cast(destinationCapacity), - sizeof(float), &next.weight)) { + sizeof(float), &cursor, &next.weight)) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } - uint64_t cursor = next.hidden.activeDataBytes > next.weight.activeDataBytes ? - next.hidden.activeDataBytes : next.weight.activeDataBytes; next.commonOffset = TileXRMoonEpDispatchUrmaAlignUp( cursor, kDispatchInternalAlignmentBytes); cursor = next.commonOffset; @@ -177,9 +177,7 @@ int TileXRMoonEpBindDispatchUrmaWorkspace(uint64_t workspaceBytes, } const uint64_t commonBytes = layout->totalBytes - layout->commonOffset; const uint64_t nextCommonOffset = workspaceBytes - commonBytes; - const uint64_t activeBytes = layout->hidden.activeDataBytes > layout->weight.activeDataBytes ? - layout->hidden.activeDataBytes : layout->weight.activeDataBytes; - if (nextCommonOffset < activeBytes) { + if (nextCommonOffset < layout->commonOffset) { return TileXR::TILEXR_ERROR_PARA_CHECK_FAIL; } const uint64_t shift = nextCommonOffset - layout->commonOffset; diff --git a/src/moonep/dispatch/urma/kernels/tilexr_moonep_dispatch_kernel.cpp b/src/moonep/dispatch/urma/kernels/tilexr_moonep_dispatch_kernel.cpp index 2c7e3ac..2639839 100644 --- a/src/moonep/dispatch/urma/kernels/tilexr_moonep_dispatch_kernel.cpp +++ b/src/moonep/dispatch/urma/kernels/tilexr_moonep_dispatch_kernel.cpp @@ -38,9 +38,12 @@ constexpr uint32_t kDispatchWqeBatchBytes = TileXRMoonEp::kDispatchWqeBatchCapacity * kDispatchUdmaWqeBytes; struct alignas(32) DispatchWqeBatchContext { - uint64_t localSourceBase; - uint64_t remoteScratchBase; - uint64_t rowBytes; + uint64_t hiddenLocalSourceBase; + uint64_t hiddenRemoteScratchBase; + uint64_t hiddenRowBytes; + uint64_t weightLocalSourceBase; + uint64_t weightRemoteScratchBase; + uint64_t weightRowBytes; uint64_t routeCountMask; uint64_t signalLocalAddr; uint64_t signalRemoteAddr; @@ -48,18 +51,18 @@ struct alignas(32) DispatchWqeBatchContext { uint64_t rmtEidH; uint32_t batchHead; uint32_t batchOutputOffset; - uint32_t tokenCount; + uint32_t dataTaskCount; uint32_t appendSignal; uint32_t topKMagic; uint32_t topKShift; - uint32_t hiddenMode; + uint32_t hasWeight; uint32_t tokenEn; uint32_t rmtJettyType; uint32_t targetHint; uint32_t tpId; uint32_t rmtJettyOrSegId; uint32_t rmtTokenValue; - uint32_t selectedStart; + uint32_t dataTaskStart; uint32_t qpSelection; uint32_t routePlanStart; }; @@ -76,7 +79,7 @@ static_assert(sizeof(TileXR::UDMACqeCtx) == 64U, "MoonEP Dispatch UDMA CQE must occupy one cache line"); static_assert(kDispatchWqeBatchBytes == 8192U, "MoonEP Dispatch WQE batch must occupy 8 KiB of UB"); -static_assert(sizeof(DispatchWqeBatchContext) == 128U, +static_assert(sizeof(DispatchWqeBatchContext) == 160U, "MoonEP Dispatch WQE batch context ABI changed"); static_assert(kDispatchUdmaIssueUbBytes % kUbAlignBytes == 0U, "MoonEP Dispatch issue UB must be 32-byte aligned"); @@ -90,7 +93,7 @@ inline void DispatchBuildWriteWqeBatchVf(__ubuf__ uint8_t *wqeBytes, __ubuf__ const int32_t *dstValues, __ubuf__ const DispatchWqeBatchContext *context) { - const uint32_t taskCount = context->tokenCount + context->appendSignal; + const uint32_t taskCount = context->dataTaskCount + context->appendSignal; for (uint32_t task = static_cast(threadIdx.x); task < taskCount; task += kDispatchWqeBuildThreads) { const uint32_t outputIndex = context->batchOutputOffset + task; @@ -107,25 +110,36 @@ inline void DispatchBuildWriteWqeBatchVf(__ubuf__ uint8_t *wqeBytes, uint64_t remoteAddr = context->signalRemoteAddr; uint32_t sqeFlag = TileXR::TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION; - if (task < context->tokenCount) { + uint64_t rowBytes = sizeof(uint64_t); + if (task < context->dataTaskCount) { + const uint32_t dataTask = context->dataTaskStart + task; + const bool weightTask = + TileXRMoonEp::DispatchDataTaskIsWeight( + dataTask, context->hasWeight != 0U); + const uint32_t qpRouteIndex = + TileXRMoonEp::DispatchDataTaskRouteIndex( + dataTask, context->hasWeight != 0U); const uint32_t qpIdx = context->qpSelection >> 2U; const uint32_t sequencePhase = context->qpSelection & 3U; const uint32_t selectedIndex = TileXRMoonEp::DispatchQpSelectedIndex( - context->selectedStart + task, sequencePhase, qpIdx); + qpRouteIndex, sequencePhase, qpIdx); const uint32_t routeId = static_cast( selectedRouteIndices[selectedIndex]); const uint64_t targetSlot = static_cast( static_cast( dstValues[routeId - context->routePlanStart])) & context->routeCountMask; - const uint32_t sourceRow = context->hiddenMode != 0U ? + const uint32_t sourceRow = weightTask ? routeId : AscendC::Simt::UintDiv(routeId, context->topKMagic, - context->topKShift) : routeId; - localAddr = context->localSourceBase + - static_cast(sourceRow) * context->rowBytes; - remoteAddr = context->remoteScratchBase + - targetSlot * context->rowBytes; + context->topKShift); + rowBytes = weightTask ? context->weightRowBytes : + context->hiddenRowBytes; + localAddr = (weightTask ? context->weightLocalSourceBase : + context->hiddenLocalSourceBase) + + static_cast(sourceRow) * rowBytes; + remoteAddr = (weightTask ? context->weightRemoteScratchBase : + context->hiddenRemoteScratchBase) + targetSlot * rowBytes; sqeFlag = 0U; } @@ -156,9 +170,7 @@ inline void DispatchBuildWriteWqeBatchVf(__ubuf__ uint8_t *wqeBytes, __ubuf__ TileXR::UDMASgeCtx *sge = reinterpret_cast<__ubuf__ TileXR::UDMASgeCtx *>( wqe + sizeof(TileXR::UDMASqeCtx)); - sge->len = task < context->tokenCount ? - static_cast(context->rowBytes) : - static_cast(sizeof(uint64_t)); + sge->len = static_cast(rowBytes); sge->tokenId = 0U; sge->va = localAddr; } @@ -239,18 +251,22 @@ struct DispatchWqeBatchInitContext { const __gm__ TileXR::CommArgs *args; __gm__ TileXR::UDMAInfo *udmaInfo; __gm__ TileXR::TileXRUDMARegistry *registry; - uint64_t remoteScratchOffset; - uint64_t scratchBytes; + uint64_t hiddenRemoteScratchOffset; + uint64_t hiddenScratchBytes; + uint64_t weightRemoteScratchOffset; + uint64_t weightScratchBytes; uint64_t remoteFlagBase; uint32_t coreIdx; bool sharedQp; + bool hasWeight; }; struct DispatchWqeBatchState { const __gm__ TileXR::CommArgs *args; __gm__ TileXR::UDMAWQCtx *qpCtxEntry; __gm__ TileXR::UDMACQCtx *cqCtxEntry; - __gm__ uint8_t *remoteScratchBase; + __gm__ uint8_t *hiddenRemoteScratchBase; + __gm__ uint8_t *weightRemoteScratchBase; __gm__ uint8_t *remoteSignalAddr; uint64_t rmtEidL; uint64_t rmtEidH; @@ -284,9 +300,10 @@ struct DispatchPreparedPeer { }; __aicore__ inline bool InitDispatchWqeBatchInitContext( - const __gm__ TileXR::CommArgs *args, uint64_t remoteScratchOffset, - uint64_t scratchBytes, uint64_t remoteFlagBase, uint32_t coreIdx, - DispatchWqeBatchInitContext &context) + const __gm__ TileXR::CommArgs *args, uint64_t hiddenRemoteScratchOffset, + uint64_t hiddenScratchBytes, uint64_t weightRemoteScratchOffset, + uint64_t weightScratchBytes, uint64_t remoteFlagBase, uint32_t coreIdx, + bool hasWeight, DispatchWqeBatchInitContext &context) { if (!TileXR::UDMARegistryEnabled(args)) { return false; @@ -306,17 +323,21 @@ __aicore__ inline bool InitDispatchWqeBatchInitContext( context.args = args; context.udmaInfo = udmaInfo; context.registry = registry; - context.remoteScratchOffset = remoteScratchOffset; - context.scratchBytes = scratchBytes; + context.hiddenRemoteScratchOffset = hiddenRemoteScratchOffset; + context.hiddenScratchBytes = hiddenScratchBytes; + context.weightRemoteScratchOffset = weightRemoteScratchOffset; + context.weightScratchBytes = weightScratchBytes; context.remoteFlagBase = remoteFlagBase; context.coreIdx = coreIdx; context.sharedQp = sharedQp; + context.hasWeight = hasWeight; return true; } __aicore__ inline bool InitDispatchWqeBatchState( const DispatchWqeBatchInitContext &context, int32_t targetRank, - uint32_t qpIdx, __gm__ uint8_t *remoteScratchBase, + uint32_t qpIdx, __gm__ uint8_t *hiddenRemoteScratchBase, + __gm__ uint8_t *weightRemoteScratchBase, __gm__ uint8_t *remoteSignalAddr, DispatchWqeBatchState &state) { if (qpIdx >= TileXRMoonEp::kDispatchQpCount) { @@ -352,7 +373,8 @@ __aicore__ inline bool InitDispatchWqeBatchState( state.args = context.args; state.qpCtxEntry = qpCtxEntry; state.cqCtxEntry = cqCtxEntry; - state.remoteScratchBase = remoteScratchBase; + state.hiddenRemoteScratchBase = hiddenRemoteScratchBase; + state.weightRemoteScratchBase = weightRemoteScratchBase; state.remoteSignalAddr = remoteSignalAddr; state.rmtEidL = remoteEid[0]; state.rmtEidH = remoteEid[1]; @@ -395,19 +417,25 @@ __aicore__ inline bool InitDispatchPreparedPeer( preparedPeer.issuePhase = issuePhase; preparedPeer.initialized = false; if (!TileXR::UDMARegisteredRangeValid(context.registry, targetRank, - context.remoteScratchOffset, context.scratchBytes) || + context.hiddenRemoteScratchOffset, context.hiddenScratchBytes) || + (context.hasWeight && !TileXR::UDMARegisteredRangeValid( + context.registry, targetRank, context.weightRemoteScratchOffset, + context.weightScratchBytes)) || !TileXR::UDMARegisteredRangeValid(context.registry, targetRank, context.remoteFlagBase, remoteFlagBytes)) { return false; } - __gm__ uint8_t *remoteScratchBase = TileXR::UDMARegisteredRemoteAddr( - context.registry, targetRank, context.remoteScratchOffset); + __gm__ uint8_t *hiddenRemoteScratchBase = TileXR::UDMARegisteredRemoteAddr( + context.registry, targetRank, context.hiddenRemoteScratchOffset); + __gm__ uint8_t *weightRemoteScratchBase = context.hasWeight ? + TileXR::UDMARegisteredRemoteAddr(context.registry, targetRank, + context.weightRemoteScratchOffset) : nullptr; __gm__ uint8_t *remoteFlagBase = TileXR::UDMARegisteredRemoteAddr( context.registry, targetRank, context.remoteFlagBase); for (uint32_t qpIdx = 0U; qpIdx < TileXRMoonEp::kDispatchQpCount; ++qpIdx) { if (!InitDispatchWqeBatchState(context, targetRank, qpIdx, - remoteScratchBase, + hiddenRemoteScratchBase, weightRemoteScratchBase, remoteFlagBase + qpIdx * sizeof(uint64_t), preparedPeer.qpState[qpIdx])) { return false; @@ -502,10 +530,11 @@ __aicore__ inline bool BuildDispatchWriteWqeBatch( AscendC::LocalTensor issueLocal, AscendC::LocalTensor selectedRouteIndices, AscendC::LocalTensor dstValues, - DispatchWqeBatchState &state, uint64_t localSourceBase, - uint64_t rowBytes, uint64_t routeCountMask, uint32_t topKMagic, - uint32_t topKShift, bool hiddenMode, uint32_t selectedStart, uint32_t tokenCount, - bool appendSignal, uint64_t signalLocalAddr, + DispatchWqeBatchState &state, uint64_t hiddenLocalSourceBase, + uint64_t hiddenRowBytes, uint64_t weightLocalSourceBase, + uint64_t weightRowBytes, bool hasWeight, uint64_t routeCountMask, + uint32_t topKMagic, uint32_t topKShift, uint32_t dataTaskStart, + uint32_t dataTaskCount, bool appendSignal, uint64_t signalLocalAddr, uint32_t sequencePhase, uint32_t routePlanStart) { __ubuf__ uint8_t *issueAddr = reinterpret_cast<__ubuf__ uint8_t *>( @@ -513,10 +542,14 @@ __aicore__ inline bool BuildDispatchWriteWqeBatch( __ubuf__ DispatchWqeBatchContext *context = reinterpret_cast<__ubuf__ DispatchWqeBatchContext *>( issueAddr + kDispatchWqeBatchContextOffset); - context->localSourceBase = localSourceBase; - context->remoteScratchBase = - reinterpret_cast(state.remoteScratchBase); - context->rowBytes = rowBytes; + context->hiddenLocalSourceBase = hiddenLocalSourceBase; + context->hiddenRemoteScratchBase = + reinterpret_cast(state.hiddenRemoteScratchBase); + context->hiddenRowBytes = hiddenRowBytes; + context->weightLocalSourceBase = weightLocalSourceBase; + context->weightRemoteScratchBase = + reinterpret_cast(state.weightRemoteScratchBase); + context->weightRowBytes = weightRowBytes; context->routeCountMask = routeCountMask; context->signalLocalAddr = signalLocalAddr; context->signalRemoteAddr = @@ -525,18 +558,18 @@ __aicore__ inline bool BuildDispatchWriteWqeBatch( context->rmtEidH = state.rmtEidH; context->batchHead = state.head; context->batchOutputOffset = state.batchCount; - context->tokenCount = tokenCount; + context->dataTaskCount = dataTaskCount; context->appendSignal = appendSignal ? 1U : 0U; context->topKMagic = topKMagic; context->topKShift = topKShift; - context->hiddenMode = hiddenMode ? 1U : 0U; + context->hasWeight = hasWeight ? 1U : 0U; context->tokenEn = state.tokenEn; context->rmtJettyType = state.rmtJettyType; context->targetHint = state.targetHint; context->tpId = state.tpId; context->rmtJettyOrSegId = state.rmtJettyOrSegId; context->rmtTokenValue = state.rmtTokenValue; - context->selectedStart = selectedStart; + context->dataTaskStart = dataTaskStart; context->qpSelection = (state.qpIdx << 2U) | (sequencePhase & 3U); context->routePlanStart = routePlanStart; @@ -717,16 +750,24 @@ __aicore__ inline bool AppendDispatchWqes(DispatchWqeBatchState &state, AscendC::LocalTensor issueLocal, AscendC::LocalTensor cqeLocal, AscendC::LocalTensor selectedRouteIndices, - AscendC::LocalTensor dstValues, uint64_t localSourceBase, - uint64_t rowBytes, uint64_t routeCountMask, uint32_t topKMagic, - uint32_t topKShift, bool hiddenMode, uint32_t selectedRouteCount, bool appendSignal, + AscendC::LocalTensor dstValues, uint64_t hiddenLocalSourceBase, + uint64_t hiddenRowBytes, uint64_t weightLocalSourceBase, + uint64_t weightRowBytes, bool hasWeight, uint64_t routeCountMask, + uint32_t topKMagic, uint32_t topKShift, uint32_t selectedRouteCount, + bool appendSignal, uint64_t signalLocalAddr, uint32_t phase, uint32_t &dfxFlags, uint32_t &firstQuietStatus, uint32_t &firstQuietPhase, uint32_t sequencePhase, uint32_t routePlanStart) { - uint32_t selectedStart = 0U; + uint64_t dataTaskCount64 = 0U; + if (!TileXRMoonEp::DispatchDataWqeCount(selectedRouteCount, hasWeight, + dataTaskCount64) || dataTaskCount64 > UINT32_MAX) { + return false; + } + const uint32_t dataTaskCount = static_cast(dataTaskCount64); + uint32_t dataTaskStart = 0U; bool signalPending = appendSignal; - while (selectedStart < selectedRouteCount || signalPending) { + while (dataTaskStart < dataTaskCount || signalPending) { if (state.batchCount == state.batchLimit && !SubmitDispatchWqeBatch(state, issueLocal, cqeLocal, phase, dfxFlags, @@ -734,23 +775,25 @@ __aicore__ inline bool AppendDispatchWqes(DispatchWqeBatchState &state, return false; } const uint32_t available = state.batchLimit - state.batchCount; - const uint32_t selectedRemaining = selectedRouteCount - selectedStart; + const uint32_t dataTaskRemaining = dataTaskCount - dataTaskStart; const bool appendSignalNow = signalPending && - static_cast(selectedRemaining) + 1U <= available; - const uint32_t tokenCapacity = available - + static_cast(dataTaskRemaining) + 1U <= available; + const uint32_t dataTaskCapacity = available - (appendSignalNow ? 1U : 0U); - const uint32_t tokenCount = selectedRemaining < tokenCapacity ? - selectedRemaining : tokenCapacity; + const uint32_t batchDataTaskCount = + dataTaskRemaining < dataTaskCapacity ? + dataTaskRemaining : dataTaskCapacity; if (!BuildDispatchWriteWqeBatch(issueLocal, - selectedRouteIndices, dstValues, state, localSourceBase, - rowBytes, routeCountMask, topKMagic, topKShift, hiddenMode, - selectedStart, tokenCount, + selectedRouteIndices, dstValues, state, hiddenLocalSourceBase, + hiddenRowBytes, weightLocalSourceBase, weightRowBytes, + hasWeight, routeCountMask, topKMagic, topKShift, + dataTaskStart, batchDataTaskCount, appendSignalNow, signalLocalAddr, sequencePhase, routePlanStart)) { return false; } - state.batchCount += tokenCount + (appendSignalNow ? 1U : 0U); - selectedStart += tokenCount; + state.batchCount += batchDataTaskCount + (appendSignalNow ? 1U : 0U); + dataTaskStart += batchDataTaskCount; if (appendSignalNow) { signalPending = false; } @@ -831,39 +874,48 @@ __aicore__ inline bool DispatchDrainHistoricalCq( __aicore__ inline bool DispatchBuildGroupedQpBatch( DispatchWqeBatchState &state, AscendC::LocalTensor issueLocal, AscendC::LocalTensor selectedRouteIndices, - AscendC::LocalTensor dstValues, uint64_t localSourceBase, - uint64_t rowBytes, uint64_t routeCountMask, uint32_t topKMagic, - uint32_t topKShift, bool hiddenMode, uint32_t selectedRouteCount, - uint32_t &selectedStart, bool &signalPending, uint64_t signalLocalAddr, + AscendC::LocalTensor dstValues, uint64_t hiddenLocalSourceBase, + uint64_t hiddenRowBytes, uint64_t weightLocalSourceBase, + uint64_t weightRowBytes, bool hasWeight, uint64_t routeCountMask, + uint32_t topKMagic, uint32_t topKShift, uint32_t selectedRouteCount, + uint32_t &dataTaskStart, bool &signalPending, uint64_t signalLocalAddr, uint32_t sequencePhase, uint32_t routePlanStart, bool &finalBatch) { finalBatch = false; state.batchCount = 0U; - if (selectedStart >= selectedRouteCount && !signalPending) { + uint64_t dataTaskCount64 = 0U; + if (!TileXRMoonEp::DispatchDataWqeCount(selectedRouteCount, hasWeight, + dataTaskCount64) || dataTaskCount64 > UINT32_MAX) { + return false; + } + const uint32_t dataTaskCount = static_cast(dataTaskCount64); + if (dataTaskStart >= dataTaskCount && !signalPending) { return true; } const uint32_t available = state.batchLimit; if (available == 0U) { return false; } - const uint32_t selectedRemaining = selectedRouteCount - selectedStart; + const uint32_t dataTaskRemaining = dataTaskCount - dataTaskStart; const bool appendSignalNow = signalPending && - static_cast(selectedRemaining) + 1U <= available; - const uint32_t tokenCapacity = available - (appendSignalNow ? 1U : 0U); - const uint32_t tokenCount = selectedRemaining < tokenCapacity ? - selectedRemaining : tokenCapacity; - if (tokenCount == 0U && !appendSignalNow) { + static_cast(dataTaskRemaining) + 1U <= available; + const uint32_t dataTaskCapacity = available - + (appendSignalNow ? 1U : 0U); + const uint32_t batchDataTaskCount = dataTaskRemaining < dataTaskCapacity ? + dataTaskRemaining : dataTaskCapacity; + if (batchDataTaskCount == 0U && !appendSignalNow) { return false; } if (!BuildDispatchWriteWqeBatch(issueLocal, selectedRouteIndices, - dstValues, state, localSourceBase, rowBytes, routeCountMask, - topKMagic, topKShift, hiddenMode, selectedStart, tokenCount, + dstValues, state, hiddenLocalSourceBase, hiddenRowBytes, + weightLocalSourceBase, weightRowBytes, hasWeight, routeCountMask, + topKMagic, topKShift, dataTaskStart, batchDataTaskCount, appendSignalNow, signalLocalAddr, sequencePhase, routePlanStart)) { return false; } - state.batchCount = tokenCount + (appendSignalNow ? 1U : 0U); - selectedStart += tokenCount; + state.batchCount = batchDataTaskCount + (appendSignalNow ? 1U : 0U); + dataTaskStart += batchDataTaskCount; if (appendSignalNow) { signalPending = false; finalBatch = true; @@ -1511,35 +1563,45 @@ __aicore__ inline bool ClearDispatchZeroFillRanges( } // namespace extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( - GM_ADDR commArgsGM, GM_ADDR inputGM, GM_ADDR dstGM, - GM_ADDR zeroFillRangesGM, GM_ADDR workspaceGM, GM_ADDR outputGM, - GM_ADDR planStatusGM, uint64_t profileOffset, - uint64_t scratchOffset, uint64_t completionFlagsOffset, - uint64_t signalOffset, uint64_t dfxOffset, uint64_t kernelStatusOffset, + GM_ADDR commArgsGM, GM_ADDR hiddenInputGM, GM_ADDR weightInputGM, + GM_ADDR dstGM, GM_ADDR zeroFillRangesGM, GM_ADDR workspaceGM, + GM_ADDR hiddenOutputGM, GM_ADDR weightOutputGM, GM_ADDR planStatusGM, + uint64_t hiddenSourceOffset, uint64_t hiddenScratchOffset, + uint64_t hiddenRowBytes, uint64_t weightSourceOffset, + uint64_t weightScratchOffset, uint64_t weightRowBytes, + uint64_t completionFlagsOffset, uint64_t signalOffset, + uint64_t hiddenProfileOffset, uint64_t weightProfileOffset, + uint64_t hiddenDfxOffset, uint64_t weightDfxOffset, + uint64_t kernelStatusOffset, int64_t s, int64_t k, int64_t h, int64_t routeCountArg, int64_t destinationCapacityArg, int64_t zeroFillRangeCountArg, - uint64_t rowBytes, - uint64_t payloadMode, int64_t magic, + uint64_t hasWeightArg, int64_t magic, uint64_t completionTimeoutTicks, uint64_t peerMode, uint64_t groupWidthArg) { if constexpr (g_coreType == AscendC::AIV) { auto args = reinterpret_cast<__gm__ TileXR::CommArgs *>(commArgsGM); - auto input = reinterpret_cast<__gm__ uint8_t *>(inputGM); + auto hiddenInput = reinterpret_cast<__gm__ uint8_t *>(hiddenInputGM); + auto weightInput = reinterpret_cast<__gm__ uint8_t *>(weightInputGM); auto dst = reinterpret_cast<__gm__ int32_t *>(dstGM); auto zeroFillRanges = reinterpret_cast<__gm__ int32_t *>(zeroFillRangesGM); auto workspace = reinterpret_cast<__gm__ uint8_t *>(workspaceGM); - auto output = reinterpret_cast<__gm__ uint8_t *>(outputGM); + auto hiddenOutput = reinterpret_cast<__gm__ uint8_t *>(hiddenOutputGM); + auto weightOutput = reinterpret_cast<__gm__ uint8_t *>(weightOutputGM); auto planStatus = reinterpret_cast<__gm__ int32_t *>(planStatusGM); - if (args == nullptr || input == nullptr || dst == nullptr || - zeroFillRanges == nullptr || workspace == nullptr || output == nullptr || + const bool hasWeight = hasWeightArg != 0U; + if (args == nullptr || hiddenInput == nullptr || dst == nullptr || + zeroFillRanges == nullptr || workspace == nullptr || + hiddenOutput == nullptr || + hasWeightArg > 1U || hasWeight != (weightInput != nullptr) || + hasWeight != (weightOutput != nullptr) || planStatus == nullptr || s <= 0 || k <= 0 || h <= 0 || routeCountArg <= 0 || destinationCapacityArg < routeCountArg || zeroFillRangeCountArg <= 0 || zeroFillRangeCountArg > UINT32_MAX || - rowBytes == 0U || rowBytes > UINT32_MAX || magic <= 0 || + hiddenRowBytes == 0U || hiddenRowBytes > UINT32_MAX || + weightRowBytes != sizeof(float) || magic <= 0 || completionTimeoutTicks == 0U || peerMode > UINT32_MAX || groupWidthArg > UINT32_MAX || - !TileXRMoonEp::DispatchPayloadModeValid(static_cast(payloadMode)) || !TileXRMoonEp::DispatchPeerModeValid(static_cast(peerMode)) || !TileXRMoonEp::DispatchGroupWidthValid( static_cast(groupWidthArg))) { @@ -1595,16 +1657,16 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( return; } - const bool hiddenMode = payloadMode == static_cast( - TileXRMoonEp::DispatchPayloadMode::Hidden); uint32_t topKMagic = 0U; uint32_t topKShift = 0U; AscendC::GetUintDivMagicAndShift(topKMagic, topKShift, static_cast(k)); - const uint64_t sourceRows = hiddenMode ? static_cast(s) : routeCount; - const uint64_t scratchSlotBytes = MultiplyU32ToU64( + const uint64_t hiddenScratchSlotBytes = MultiplyU32ToU64( static_cast(destinationCapacity), - static_cast(rowBytes)); + static_cast(hiddenRowBytes)); + const uint64_t weightScratchSlotBytes = MultiplyU32ToU64( + static_cast(destinationCapacity), + static_cast(weightRowBytes)); const uint64_t expectedFlag = static_cast(magic); const uint64_t scratchIndex = expectedFlag % TileXRMoonEp::kDispatchScratchBufferCount; @@ -1735,9 +1797,21 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( const uint64_t kernelStartCycle = static_cast(AscendC::GetSystemCycle()); const uint64_t stagingStartCycle = kernelStartCycle; #endif - for (uint64_t sourceRow = blockIdx; sourceRow < sourceRows; sourceRow += blockNum) { - CopyBytesGmToGm(workspace + sourceRow * rowBytes, - input + sourceRow * rowBytes, static_cast(rowBytes), relayLocal); + for (uint64_t sourceRow = blockIdx; sourceRow < static_cast(s); + sourceRow += blockNum) { + CopyBytesGmToGm(workspace + hiddenSourceOffset + + sourceRow * hiddenRowBytes, + hiddenInput + sourceRow * hiddenRowBytes, + static_cast(hiddenRowBytes), relayLocal); + } + if (hasWeight) { + for (uint64_t sourceRow = blockIdx; sourceRow < routeCount; + sourceRow += blockNum) { + CopyBytesGmToGm(workspace + weightSourceOffset + + sourceRow * weightRowBytes, + weightInput + sourceRow * weightRowBytes, + static_cast(weightRowBytes), relayLocal); + } } AscendC::SyncAll(); #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_PROFILING) @@ -1746,7 +1820,10 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( #endif const int32_t upstreamStatus = LoadDispatchPlanStatus(planStatus, relayLocal); - auto currentScratch = workspace + scratchOffset + scratchIndex * scratchSlotBytes; + auto currentHiddenScratch = workspace + hiddenScratchOffset + + scratchIndex * hiddenScratchSlotBytes; + auto currentWeightScratch = workspace + weightScratchOffset + + scratchIndex * weightScratchSlotBytes; auto receiveFlags = reinterpret_cast<__gm__ uint64_t *>( workspace + completionFlagsOffset); auto signalSource = reinterpret_cast<__gm__ uint64_t *>( @@ -1772,6 +1849,7 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( uint64_t matchedRouteCount = 0U; uint64_t selectedRouteCount = 0U; uint64_t processedRouteCount = 0U; + uint64_t issuedRouteCount = 0U; uint64_t issuedPutCount = 0U; uint64_t issuedPutBytes = 0U; uint64_t visitedPeerCount = 0U; @@ -1801,16 +1879,20 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( static_cast(AscendC::GetSystemCycle()); #endif if (useVectorSlotSelect && groupedPeerMode) { - const uint64_t remoteScratchOffset = scratchOffset + - scratchIndex * scratchSlotBytes; + const uint64_t remoteHiddenScratchOffset = hiddenScratchOffset + + scratchIndex * hiddenScratchSlotBytes; + const uint64_t remoteWeightScratchOffset = weightScratchOffset + + scratchIndex * weightScratchSlotBytes; const uint64_t remoteFlagBase = completionFlagsOffset + static_cast(rank) * TileXRMoonEp::kDispatchQpCount * sizeof(uint64_t); DispatchWqeBatchInitContext initContext {}; const bool initContextValid = localOnly || - InitDispatchWqeBatchInitContext(args, remoteScratchOffset, - scratchSlotBytes, remoteFlagBase, - static_cast(blockIdx), initContext); + InitDispatchWqeBatchInitContext(args, + remoteHiddenScratchOffset, hiddenScratchSlotBytes, + remoteWeightScratchOffset, weightScratchSlotBytes, + remoteFlagBase, static_cast(blockIdx), + hasWeight, initContext); DispatchPreparedPeer previousPeer {}; bool previousPeerValid = false; @@ -1824,28 +1906,31 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( continue; } ++visitedPeerCount; - if (peerValue == rank || upstreamStatus != - TileXRMoonEp::kDispatchStatusSuccess) { + if (peerValue == rank) { continue; } const int32_t peer = static_cast(peerValue); + const bool payloadReady = upstreamStatus == + TileXRMoonEp::kDispatchStatusSuccess && dfxFlags == 0U; #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_PROFILING) const uint64_t putIssueStartCycle = static_cast(AscendC::GetSystemCycle()); #endif DispatchPreparedPeer preparedPeer {}; - if (!initContextValid || !InitDispatchPreparedPeer(initContext, - peer, group, preparedPeer) || - !DispatchDrainHistoricalCq(preparedPeer, relayLocal, + bool sendOk = initContextValid && + InitDispatchPreparedPeer(initContext, peer, group, + preparedPeer); + if (sendOk) { + sendOk = DispatchDrainHistoricalCq(preparedPeer, relayLocal, completionTimeoutTicks, dfxFlags, firstQuietStatus, firstQuietPhase, timeoutPeer, timeoutPhase, - timeoutObservedFlag)) { + timeoutObservedFlag); + } + if (!sendOk) { dfxFlags |= TileXRMoonEp::kDispatchDfxCqError; - break; } bool firstLogicalBatch = true; - bool sendOk = true; uint64_t peerSelectedCount = 0U; for (uint32_t routeTileStart = 0U; sendOk && routeTileStart < routeCount;) { @@ -1853,22 +1938,26 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( TileXRMoonEp::DispatchRouteTileCount( static_cast(routeCount), routeTileStart, kRouteTileElements); - if (!PrepareDispatchRouteTile(dstGlobal, routePlanLocal, - routeRankLocal, routeIndexLocal, routeTileStart, - routeTileCount, routeShift)) { - sendOk = false; - break; + uint32_t selectedCount = 0U; + if (payloadReady) { + if (!PrepareDispatchRouteTile(dstGlobal, routePlanLocal, + routeRankLocal, routeIndexLocal, routeTileStart, + routeTileCount, routeShift)) { + sendOk = false; + break; + } + selectedCount = SelectDispatchPeerRoutes( + compareMaskLocal, routeRankLocal, routeIndexLocal, + selectedRouteIndexLocal, peer, routeTileCount); + scannedRouteCount += routeTileCount; + matchedRouteCount += selectedCount; + selectedRouteCount += selectedCount; + peerSelectedCount += selectedCount; } - const uint32_t selectedCount = SelectDispatchPeerRoutes( - compareMaskLocal, routeRankLocal, routeIndexLocal, - selectedRouteIndexLocal, peer, routeTileCount); - scannedRouteCount += routeTileCount; - matchedRouteCount += selectedCount; - selectedRouteCount += selectedCount; - peerSelectedCount += selectedCount; uint32_t qpSelectedCount[TileXRMoonEp::kDispatchQpCount] = {}; - uint32_t qpSelectedStart[TileXRMoonEp::kDispatchQpCount] = {}; + uint32_t qpDataTaskCount[TileXRMoonEp::kDispatchQpCount] = {}; + uint32_t qpDataTaskStart[TileXRMoonEp::kDispatchQpCount] = {}; const bool finalRouteTile = routeTileStart + routeTileCount == routeCount; bool qpSignalPending[TileXRMoonEp::kDispatchQpCount] = { @@ -1878,11 +1967,13 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( qpSelectedCount[qpIdx] = TileXRMoonEp::DispatchQpRouteCount( selectedCount, 0U, qpIdx); + qpDataTaskCount[qpIdx] = qpSelectedCount[qpIdx] * + TileXRMoonEp::DispatchPayloadWqesPerRoute(hasWeight); } while (sendOk && - (qpSelectedStart[0] < qpSelectedCount[0] || - qpSelectedStart[1] < qpSelectedCount[1] || + (qpDataTaskStart[0] < qpDataTaskCount[0] || + qpDataTaskStart[1] < qpDataTaskCount[1] || qpSignalPending[0] || qpSignalPending[1])) { bool finalBatch[TileXRMoonEp::kDispatchQpCount] = {}; for (uint32_t qpIdx = 0U; @@ -1896,10 +1987,14 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( sendOk = DispatchBuildGroupedQpBatch(qpState, issueLocal, selectedRouteIndexLocal, routePlanLocal, - reinterpret_cast(workspace), rowBytes, - destinationCapacity - 1U, topKMagic, topKShift, - hiddenMode, qpSelectedCount[qpIdx], - qpSelectedStart[qpIdx], qpSignalPending[qpIdx], + reinterpret_cast(workspace + + hiddenSourceOffset), hiddenRowBytes, + reinterpret_cast(workspace + + weightSourceOffset), weightRowBytes, + hasWeight, destinationCapacity - 1U, + topKMagic, topKShift, + qpSelectedCount[qpIdx], + qpDataTaskStart[qpIdx], qpSignalPending[qpIdx], reinterpret_cast(signalSource + qpIdx), 0U, routeTileStart, finalBatch[qpIdx]); } @@ -1947,30 +2042,67 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( routeTileStart += routeTileCount; } - if (!sendOk || preparedPeer.qpState[0].finalStaged == 0U || - preparedPeer.qpState[1].finalStaged == 0U) { + const bool optimizedSignalSubmitted = sendOk && + preparedPeer.qpState[0].finalStaged != 0U && + preparedPeer.qpState[1].finalStaged != 0U; + bool completionSubmitted = optimizedSignalSubmitted; + if (!optimizedSignalSubmitted) { if ((dfxFlags & (TileXRMoonEp::kDispatchDfxCreditTimeout | TileXRMoonEp::kDispatchDfxCqError)) == 0U) { dfxFlags |= TileXRMoonEp::kDispatchDfxInvalidConfig; } - break; + requiresFinalQuiet = true; + completionSubmitted = true; + for (uint32_t qpIdx = 0U; + qpIdx < TileXRMoonEp::kDispatchQpCount; ++qpIdx) { + const uint32_t signalStatus = + TileXR::UDMAPutNbiOnQpWithFlagDeferred( + args, qpIdx == 0U ? udmaIssueQp0Local : + udmaIssueQp1Local, + peer, physicalQp[qpIdx], signalSource + qpIdx, + remoteFlagBase + qpIdx * sizeof(uint64_t), + static_cast(sizeof(uint64_t)), + TileXR::TILEXR_UDMA_SQE_FLAG_ORDERED_COMPLETION); + const uint32_t flushStatus = + TileXR::UDMAFlushQpDoorbell( + args, peer, physicalQp[qpIdx]); + if (signalStatus != TileXR::TILEXR_UDMA_STATUS_SUCCESS || + flushStatus != TileXR::TILEXR_UDMA_STATUS_SUCCESS) { + completionSubmitted = false; + dfxFlags |= TileXRMoonEp::kDispatchDfxQuietError; + if (firstQuietStatus == 0U) { + firstQuietStatus = signalStatus != + TileXR::TILEXR_UDMA_STATUS_SUCCESS ? + signalStatus : flushStatus; + firstQuietPhase = group; + } + } + } } - issuedPutCount += peerSelectedCount; - issuedPutBytes += peerSelectedCount * rowBytes; - processedRouteCount += peerSelectedCount; - completionFlagCount += TileXRMoonEp::kDispatchQpCount; - previousPeer = preparedPeer; - previousPeerValid = true; + if (optimizedSignalSubmitted) { + issuedPutCount += peerSelectedCount * + TileXRMoonEp::DispatchPayloadWqesPerRoute(hasWeight); + issuedPutBytes += peerSelectedCount * + (hiddenRowBytes + (hasWeight ? weightRowBytes : 0U)); + issuedRouteCount += peerSelectedCount; + processedRouteCount += peerSelectedCount; + previousPeer = preparedPeer; + previousPeerValid = true; + } + if (completionSubmitted) { + completionFlagCount += TileXRMoonEp::kDispatchQpCount; + } if (!WaitDispatchIncomingPeerAndPublishCredit(args, receiveFlags, rank, rankSize, peer, group, lane, groupWidth, magic, creditPeerMode, completionTimeoutTicks, relayLocal, dfxFlags, timeoutPeer, timeoutPhase, timeoutObservedFlag)) { - ProbeDispatchPeerFinalCq(previousPeer, relayLocal, - outgoingCqStatuses, outgoingRemainingSqEntries); - break; + if (optimizedSignalSubmitted) { + ProbeDispatchPeerFinalCq(preparedPeer, relayLocal, + outgoingCqStatuses, outgoingRemainingSqEntries); + } } #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_PROFILING) putIssueCycles += @@ -1985,15 +2117,19 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( dfxFlags |= TileXRMoonEp::kDispatchDfxCqError; } } else if (useVectorSlotSelect) { - const uint64_t remoteScratchOffset = scratchOffset + - scratchIndex * scratchSlotBytes; + const uint64_t remoteHiddenScratchOffset = hiddenScratchOffset + + scratchIndex * hiddenScratchSlotBytes; + const uint64_t remoteWeightScratchOffset = weightScratchOffset + + scratchIndex * weightScratchSlotBytes; const uint64_t remoteFlagBase = completionFlagsOffset + static_cast(rank) * TileXRMoonEp::kDispatchQpCount * sizeof(uint64_t); DispatchWqeBatchInitContext initContext {}; const bool initContextValid = InitDispatchWqeBatchInitContext( - args, remoteScratchOffset, scratchSlotBytes, - remoteFlagBase, static_cast(blockIdx), initContext); + args, remoteHiddenScratchOffset, hiddenScratchSlotBytes, + remoteWeightScratchOffset, weightScratchSlotBytes, + remoteFlagBase, static_cast(blockIdx), hasWeight, + initContext); DispatchPreparedPeer peerBatch[kDispatchPreparedPeerCapacity] {}; uint64_t peerCursor = 0U; const uint64_t totalPeerAssignments = @@ -2024,7 +2160,9 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( matchedRouteCount += selectedCount; selectedRouteCount += selectedCount; } - const uint64_t peerWqeCount = selectedCount; + const uint64_t peerWqeCount = static_cast( + selectedCount) * + TileXRMoonEp::DispatchPayloadWqesPerRoute(hasWeight); uint64_t issuedPeerWqeCount = 0U; bool signalSubmitted = false; if (upstreamStatus == TileXRMoonEp::kDispatchStatusSuccess) { @@ -2046,9 +2184,12 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( udmaIssueQp1Local, relayLocal, selectedRouteIndexLocal, routePlanLocal, - reinterpret_cast(workspace), - rowBytes, destinationCapacity - 1U, - topKMagic, topKShift, hiddenMode, + reinterpret_cast(workspace + + hiddenSourceOffset), hiddenRowBytes, + reinterpret_cast(workspace + + weightSourceOffset), weightRowBytes, + hasWeight, destinationCapacity - 1U, + topKMagic, topKShift, qpRouteCount, true, reinterpret_cast( signalSource + qpIdx), @@ -2056,9 +2197,12 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( firstQuietPhase, sequencePhase, 0U); } if (batchOk) { - issuedPutCount += selectedCount; - issuedPutBytes += selectedCount * rowBytes; + issuedPutCount += peerWqeCount; + issuedPutBytes += selectedCount * + (hiddenRowBytes + + (hasWeight ? weightRowBytes : 0U)); processedRouteCount += selectedCount; + issuedRouteCount += selectedCount; signalSubmitted = true; } else { dfxFlags |= TileXRMoonEp::kDispatchDfxQuietError; @@ -2090,23 +2234,43 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( firstInvalidRawDst); continue; } - const uint64_t sourceRow = hiddenMode ? - routeId / static_cast(k) : routeId; TileXR::UDMAPutNbiOnQpWithFlagDeferred( args, udmaIssueQp0Local, targetRank, physicalQp[0], - workspace + sourceRow * rowBytes, - remoteScratchOffset + targetSlot * rowBytes, - static_cast(rowBytes), + workspace + hiddenSourceOffset + + routeId / static_cast(k) * + hiddenRowBytes, + remoteHiddenScratchOffset + + targetSlot * hiddenRowBytes, + static_cast(hiddenRowBytes), TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION); ++issuedPeerWqeCount; ++issuedPutCount; - issuedPutBytes += rowBytes; - ++processedRouteCount; + issuedPutBytes += hiddenRowBytes; ReclaimDeferredSegment(args, targetRank, physicalQp[0], issuedPeerWqeCount, issuePhase, dfxFlags, firstQuietStatus, firstQuietPhase); + if (hasWeight) { + TileXR::UDMAPutNbiOnQpWithFlagDeferred( + args, udmaIssueQp0Local, targetRank, + physicalQp[0], + workspace + weightSourceOffset + + routeId * weightRowBytes, + remoteWeightScratchOffset + + targetSlot * weightRowBytes, + static_cast(weightRowBytes), + TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION); + ++issuedPeerWqeCount; + ++issuedPutCount; + issuedPutBytes += weightRowBytes; + ReclaimDeferredSegment(args, targetRank, + physicalQp[0], issuedPeerWqeCount, + issuePhase, dfxFlags, + firstQuietStatus, firstQuietPhase); + } + ++processedRouteCount; + ++issuedRouteCount; if (ShouldFlushPartialDoorbell( issuedPeerWqeCount, peerWqeCount)) { TileXR::UDMAFlushQpDoorbell( @@ -2158,7 +2322,6 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( const uint64_t putIssueStartCycle = static_cast(AscendC::GetSystemCycle()); #endif - uint64_t peerWqeCount = 0U; uint64_t issuedPeerWqeCount = 0U; if (upstreamStatus == TileXRMoonEp::kDispatchStatusSuccess) { @@ -2184,56 +2347,55 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( firstInvalidRawDst); continue; } - if (targetRank == peer) { - ++peerWqeCount; - ++matchedRouteCount; - } - } - } - selectedRouteCount += peerWqeCount; - for (uint64_t tileStart = 0U; - tileStart < routeCount; - tileStart += kRouteTileElements) { - const uint32_t tileElements = LoadRouteTile( - dstGlobal, tileStart, routeCount, - routePlanLocal); - scannedRouteCount += tileElements; - for (uint32_t tileRoute = 0U; - tileRoute < tileElements; ++tileRoute) { - const uint64_t routeId = tileStart + tileRoute; - const int32_t encoded = - routePlanLocal.GetValue(tileRoute); - int32_t targetRank = -1; - uint64_t targetSlot = 0U; - if (!DecodeSendDst(encoded, destinationCapacity, - rankSize, targetRank, targetSlot) || - targetRank != peer) { + if (targetRank != peer) { continue; } - const uint64_t sourceRow = hiddenMode ? - routeId / static_cast(k) : - routeId; + ++matchedRouteCount; + ++selectedRouteCount; TileXR::UDMAPutNbiOnQpWithFlagDeferred( args, udmaIssueQp0Local, targetRank, physicalQp[0], - workspace + sourceRow * rowBytes, - scratchOffset + - scratchIndex * scratchSlotBytes + - targetSlot * rowBytes, - static_cast(rowBytes), + workspace + hiddenSourceOffset + + routeId / static_cast(k) * + hiddenRowBytes, + hiddenScratchOffset + + scratchIndex * hiddenScratchSlotBytes + + targetSlot * hiddenRowBytes, + static_cast(hiddenRowBytes), TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION); ++issuedPeerWqeCount; ++issuedPutCount; - issuedPutBytes += rowBytes; - ++processedRouteCount; + issuedPutBytes += hiddenRowBytes; ReclaimDeferredSegment(args, targetRank, physicalQp[0], issuedPeerWqeCount, static_cast(issuePhase), dfxFlags, firstQuietStatus, firstQuietPhase); - if (ShouldFlushPartialDoorbell( - issuedPeerWqeCount, peerWqeCount)) { + if (hasWeight) { + TileXR::UDMAPutNbiOnQpWithFlagDeferred( + args, udmaIssueQp0Local, targetRank, + physicalQp[0], + workspace + weightSourceOffset + + routeId * weightRowBytes, + weightScratchOffset + + scratchIndex * weightScratchSlotBytes + + targetSlot * weightRowBytes, + static_cast(weightRowBytes), + TileXR::TILEXR_UDMA_SQE_FLAG_COMPLETION); + ++issuedPeerWqeCount; + ++issuedPutCount; + issuedPutBytes += weightRowBytes; + ReclaimDeferredSegment(args, targetRank, + physicalQp[0], issuedPeerWqeCount, + static_cast(issuePhase), + dfxFlags, firstQuietStatus, + firstQuietPhase); + } + ++processedRouteCount; + ++issuedRouteCount; + if (issuedPeerWqeCount % + TileXRMoonEp::kDispatchWqeBatchCapacity == 0U) { TileXR::UDMAFlushQpDoorbell( args, targetRank, physicalQp[0]); } @@ -2303,11 +2465,18 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( firstInvalidRouteId, firstInvalidRawDst); continue; } - const uint64_t sourceRow = hiddenMode ? - routeId / static_cast(k) : routeId; - CopyBytesGmToGm(currentScratch + targetSlot * rowBytes, - workspace + sourceRow * rowBytes, - static_cast(rowBytes), relayLocal); + CopyBytesGmToGm( + currentHiddenScratch + targetSlot * hiddenRowBytes, + workspace + hiddenSourceOffset + + routeId / static_cast(k) * hiddenRowBytes, + static_cast(hiddenRowBytes), relayLocal); + if (hasWeight) { + CopyBytesGmToGm( + currentWeightScratch + targetSlot * weightRowBytes, + workspace + weightSourceOffset + + routeId * weightRowBytes, + static_cast(weightRowBytes), relayLocal); + } ++processedRouteCount; } } else { @@ -2339,12 +2508,18 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( } ++matchedRouteCount; ++selectedRouteCount; - const uint64_t sourceRow = hiddenMode ? - routeId / static_cast(k) : routeId; CopyBytesGmToGm( - currentScratch + targetSlot * rowBytes, - workspace + sourceRow * rowBytes, - static_cast(rowBytes), relayLocal); + currentHiddenScratch + targetSlot * hiddenRowBytes, + workspace + hiddenSourceOffset + + routeId / static_cast(k) * hiddenRowBytes, + static_cast(hiddenRowBytes), relayLocal); + if (hasWeight) { + CopyBytesGmToGm( + currentWeightScratch + targetSlot * weightRowBytes, + workspace + weightSourceOffset + + routeId * weightRowBytes, + static_cast(weightRowBytes), relayLocal); + } ++processedRouteCount; } } @@ -2415,25 +2590,42 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( dfxFlags |= TileXRMoonEp::kDispatchDfxRouteCountMismatch; } if (upstreamStatus == TileXRMoonEp::kDispatchStatusSuccess && - !ClearDispatchZeroFillRanges(currentScratch, zeroFillRanges, - zeroFillRangeCount, destinationCapacity, rowBytes, - blockIdx, blockNum, relayLocal)) { + (!ClearDispatchZeroFillRanges(currentHiddenScratch, zeroFillRanges, + zeroFillRangeCount, destinationCapacity, hiddenRowBytes, + blockIdx, blockNum, relayLocal) || + (hasWeight && !ClearDispatchZeroFillRanges(currentWeightScratch, + zeroFillRanges, zeroFillRangeCount, destinationCapacity, + weightRowBytes, blockIdx, blockNum, relayLocal)))) { dfxFlags |= TileXRMoonEp::kDispatchDfxInvalidConfig; } const int32_t localExecutionStatus = StatusFromDfxFlags(dfxFlags); DispatchPublishFirstStatus(planStatus, localExecutionStatus); + const uint32_t ownerPayloadMode = static_cast(hasWeight ? + TileXRMoonEp::DispatchPayloadMode::RouteWeight : + TileXRMoonEp::DispatchPayloadMode::Hidden); #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_DFX) uint64_t localSignalObserved = expectedFlag; if (!localOnly && dfxFlags != 0U) { localSignalObserved = LoadCompletionFlag(signalSource); } - WriteDfxRecord(workspace + dfxOffset, static_cast(payloadMode), + WriteDfxRecord(workspace + hiddenDfxOffset, + static_cast(TileXRMoonEp::DispatchPayloadMode::Hidden), static_cast(rank), static_cast(blockIdx), dfxFlags, firstInvalidRouteId, firstInvalidRawDst, firstQuietStatus, firstQuietPhase, timeoutPeer, timeoutPhase, routeCount, processedRouteCount, expectedFlag, expectedFlag, timeoutObservedFlag, localSignalObserved, completionFlagCount, outgoingCqStatuses, outgoingRemainingSqEntries, diagnosticLocal); + if (hasWeight) { + WriteDfxRecord(workspace + weightDfxOffset, ownerPayloadMode, + static_cast(rank), static_cast(blockIdx), + dfxFlags, firstInvalidRouteId, firstInvalidRawDst, + firstQuietStatus, firstQuietPhase, timeoutPeer, timeoutPhase, + routeCount, processedRouteCount, expectedFlag, expectedFlag, + timeoutObservedFlag, localSignalObserved, completionFlagCount, + outgoingCqStatuses, outgoingRemainingSqEntries, + diagnosticLocal); + } #endif #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_PROFILING) const uint64_t dfxWriteEndCycle = @@ -2449,7 +2641,7 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_DFX) uint32_t globalDfxFlags = 0U; auto allDfx = reinterpret_cast<__gm__ TileXRMoonEp::DispatchDfxRecord *>( - workspace + dfxOffset); + workspace + (hasWeight ? weightDfxOffset : hiddenDfxOffset)); for (uint32_t core = 0U; core < static_cast(blockNum); ++core) { globalDfxFlags |= allDfx[core].flags; } @@ -2464,7 +2656,7 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( kernelStatus->version = TileXRMoonEp::kDispatchDiagnosticVersion; kernelStatus->recordBytes = sizeof(TileXRMoonEp::DispatchKernelStatus); kernelStatus->status = executionStatus; - kernelStatus->payloadMode = static_cast(payloadMode); + kernelStatus->payloadMode = ownerPayloadMode; kernelStatus->magic = expectedFlag; for (uint32_t index = 0U; index < 5U; ++index) { kernelStatus->reserved[index] = 0U; @@ -2477,6 +2669,10 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( kernelStatus->reserved[0] |= TileXRMoonEp::kDispatchKernelStatusFeatureProfilingEnabled; #endif + if (hasWeight) { + kernelStatus->reserved[0] |= + TileXRMoonEp::kDispatchKernelStatusFeatureFusedEpoch; + } } #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_PROFILING) @@ -2493,16 +2689,26 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( const uint64_t outputSlotCount = outputEndSlot - outputStartSlot; if (outputCopyTileBytes != 0U) { CopyContiguousBytesGmToGmPipelined( - output + outputStartSlot * rowBytes, - currentScratch + outputStartSlot * rowBytes, - outputSlotCount * rowBytes, outputCopyTileBytes, + hiddenOutput + outputStartSlot * hiddenRowBytes, + currentHiddenScratch + outputStartSlot * hiddenRowBytes, + outputSlotCount * hiddenRowBytes, outputCopyTileBytes, outputCopyQueue); } else { for (uint64_t targetSlot = outputStartSlot; targetSlot < outputEndSlot; ++targetSlot) { - CopyBytesGmToGm(output + targetSlot * rowBytes, - currentScratch + targetSlot * rowBytes, - static_cast(rowBytes), relayLocal); + CopyBytesGmToGm( + hiddenOutput + targetSlot * hiddenRowBytes, + currentHiddenScratch + targetSlot * hiddenRowBytes, + static_cast(hiddenRowBytes), relayLocal); + } + } + if (hasWeight) { + for (uint64_t targetSlot = outputStartSlot; + targetSlot < outputEndSlot; ++targetSlot) { + CopyBytesGmToGm( + weightOutput + targetSlot * weightRowBytes, + currentWeightScratch + targetSlot * weightRowBytes, + static_cast(weightRowBytes), relayLocal); } } } @@ -2546,40 +2752,88 @@ extern "C" __global__ __aicore__ void tilexr_moonep_dispatch_urma_kernel( const uint64_t quietEndCycle = static_cast(AscendC::GetSystemCycle()); quietCycles = quietEndCycle - quietStartCycle; #endif + AscendC::SyncAll(); + if (blockIdx == 0U) { + kernelStatus->status = LoadDispatchPlanStatus(planStatus, relayLocal); + } #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_DFX) - WriteDfxRecord(workspace + dfxOffset, static_cast(payloadMode), + WriteDfxRecord(workspace + hiddenDfxOffset, + static_cast(TileXRMoonEp::DispatchPayloadMode::Hidden), static_cast(rank), static_cast(blockIdx), dfxFlags, firstInvalidRouteId, firstInvalidRawDst, firstQuietStatus, firstQuietPhase, timeoutPeer, timeoutPhase, routeCount, processedRouteCount, expectedFlag, expectedFlag, timeoutObservedFlag, localSignalObserved, completionFlagCount, outgoingCqStatuses, outgoingRemainingSqEntries, diagnosticLocal); + if (hasWeight) { + WriteDfxRecord(workspace + weightDfxOffset, ownerPayloadMode, + static_cast(rank), static_cast(blockIdx), + dfxFlags, firstInvalidRouteId, firstInvalidRawDst, + firstQuietStatus, firstQuietPhase, timeoutPeer, timeoutPhase, + routeCount, processedRouteCount, expectedFlag, expectedFlag, + timeoutObservedFlag, localSignalObserved, completionFlagCount, + outgoingCqStatuses, outgoingRemainingSqEntries, + diagnosticLocal); + } #endif #if defined(TILEXR_MOONEP_DISPATCH_ENABLE_PROFILING) const uint64_t kernelCycles = static_cast(AscendC::GetSystemCycle()) - kernelStartCycle; - WriteProfileRecord(workspace + profileOffset, static_cast(payloadMode), + const uint64_t payloadPutCount = issuedRouteCount; + WriteProfileRecord(workspace + hiddenProfileOffset, + static_cast(TileXRMoonEp::DispatchPayloadMode::Hidden), static_cast(rank), static_cast(blockIdx), static_cast(blockNum), dfxFlags, expectedFlag, static_cast(scratchIndex), static_cast(groupCount), useVectorSlotSelect ? TileXRMoonEp::kDispatchSelectVector : TileXRMoonEp::kDispatchSelectScalarTiled, - fallbackReason, scannedRouteCount, matchedRouteCount, selectedRouteCount, - processedRouteCount, issuedPutCount, issuedPutBytes, visitedPeerCount, - completionFlagCount, kernelCycles, stagingCycles, putIssueCycles, - flagWaitCycles, outputCopyCycles, quietCycles, - stagingEndCycle - kernelStartCycle, - issueWindowStartCycle - kernelStartCycle, - remoteIssueEndCycle - kernelStartCycle, - issueWindowEndCycle - kernelStartCycle, - flagWaitStartCycle - kernelStartCycle, - flagWaitEndCycle - kernelStartCycle, - dfxWriteEndCycle - kernelStartCycle, - syncAllEndCycle - kernelStartCycle, - outputCopyStartCycle - kernelStartCycle, - outputCopyEndCycle - kernelStartCycle, - quietEndCycle - kernelStartCycle, diagnosticLocal); + fallbackReason, hasWeight ? 0U : scannedRouteCount, + hasWeight ? 0U : matchedRouteCount, + hasWeight ? 0U : selectedRouteCount, processedRouteCount, + payloadPutCount, payloadPutCount * hiddenRowBytes, + hasWeight ? 0U : visitedPeerCount, + hasWeight ? 0U : completionFlagCount, + hasWeight ? 0U : kernelCycles, hasWeight ? 0U : stagingCycles, + hasWeight ? 0U : putIssueCycles, hasWeight ? 0U : flagWaitCycles, + hasWeight ? 0U : outputCopyCycles, hasWeight ? 0U : quietCycles, + hasWeight ? 0U : stagingEndCycle - kernelStartCycle, + hasWeight ? 0U : issueWindowStartCycle - kernelStartCycle, + hasWeight ? 0U : remoteIssueEndCycle - kernelStartCycle, + hasWeight ? 0U : issueWindowEndCycle - kernelStartCycle, + hasWeight ? 0U : flagWaitStartCycle - kernelStartCycle, + hasWeight ? 0U : flagWaitEndCycle - kernelStartCycle, + hasWeight ? 0U : dfxWriteEndCycle - kernelStartCycle, + hasWeight ? 0U : syncAllEndCycle - kernelStartCycle, + hasWeight ? 0U : outputCopyStartCycle - kernelStartCycle, + hasWeight ? 0U : outputCopyEndCycle - kernelStartCycle, + hasWeight ? 0U : quietEndCycle - kernelStartCycle, + diagnosticLocal); + if (hasWeight) { + WriteProfileRecord(workspace + weightProfileOffset, + ownerPayloadMode, static_cast(rank), + static_cast(blockIdx), static_cast(blockNum), + dfxFlags, expectedFlag, static_cast(scratchIndex), + static_cast(groupCount), useVectorSlotSelect ? + TileXRMoonEp::kDispatchSelectVector : + TileXRMoonEp::kDispatchSelectScalarTiled, + fallbackReason, scannedRouteCount, matchedRouteCount, + selectedRouteCount, processedRouteCount, payloadPutCount, + payloadPutCount * weightRowBytes, visitedPeerCount, + completionFlagCount, kernelCycles, stagingCycles, + putIssueCycles, flagWaitCycles, outputCopyCycles, quietCycles, + stagingEndCycle - kernelStartCycle, + issueWindowStartCycle - kernelStartCycle, + remoteIssueEndCycle - kernelStartCycle, + issueWindowEndCycle - kernelStartCycle, + flagWaitStartCycle - kernelStartCycle, + flagWaitEndCycle - kernelStartCycle, + dfxWriteEndCycle - kernelStartCycle, + syncAllEndCycle - kernelStartCycle, + outputCopyStartCycle - kernelStartCycle, + outputCopyEndCycle - kernelStartCycle, + quietEndCycle - kernelStartCycle, diagnosticLocal); + } #endif } } diff --git a/tests/moonep/CMakeLists.txt b/tests/moonep/CMakeLists.txt index 4c5bdc8..69f1688 100644 --- a/tests/moonep/CMakeLists.txt +++ b/tests/moonep/CMakeLists.txt @@ -32,10 +32,13 @@ file(WRITE "${TILEXR_MOONEP_FAKE_INCLUDE_DIR}/acl/acl_rt.h" "#include \n" "#include \"acl_base.h\"\n" "typedef enum aclrtMemcpyKind { ACL_MEMCPY_DEVICE_TO_DEVICE = 3 } aclrtMemcpyKind;\n" +"typedef enum aclrtDevAttr { ACL_DEV_ATTR_VECTOR_CORE_NUM = 0, ACL_DEV_ATTR_UBUF_PER_VECTOR_CORE = 1 } aclrtDevAttr;\n" "#ifdef __cplusplus\nextern \"C\" {\n#endif\n" "aclError aclrtMemsetAsync(void *devPtr, size_t maxCount, int32_t value, size_t count, aclrtStream stream);\n" "aclError aclrtMemcpyAsync(void *dst, size_t destMax, const void *src, size_t count, aclrtMemcpyKind kind, aclrtStream stream);\n" "aclError aclrtSynchronizeStream(aclrtStream stream);\n" +"aclError aclrtGetDevice(int32_t *deviceId);\n" +"aclError aclrtGetDeviceInfo(uint32_t deviceId, aclrtDevAttr attr, int64_t *value);\n" "#ifdef __cplusplus\n}\n#endif\n" "#endif\n") @@ -50,7 +53,7 @@ file(WRITE "${TILEXR_MOONEP_FAKE_INCLUDE_DIR}/runtime/kernel.h" "typedef struct rtDevBinary { uint32_t magic; uint32_t version; const void *data; uint64_t length; } rtDevBinary_t;\n" "typedef struct rtArgsEx { void *args; size_t argsSize; } rtArgsEx_t;\n" "typedef struct rtSmDesc { uint32_t reserved; } rtSmDesc_t;\n" -"typedef struct rtTaskCfgInfo { uint32_t schemMode; } rtTaskCfgInfo_t;\n" +"typedef struct rtTaskCfgInfo { uint32_t schemMode; uint32_t localMemorySize; } rtTaskCfgInfo_t;\n" "#define RT_ERROR_NONE 0\n" "#define RT_DEV_BINARY_MAGIC_ELF_AIVEC 0x41415246U\n" "#ifdef __cplusplus\nextern \"C\" {\n#endif\n" @@ -165,6 +168,28 @@ target_include_directories(test_tilexr_moonep_dispatch_urma_layout PRIVATE add_executable(test_tilexr_moonep_dispatch_urma_schedule unit/test_tilexr_moonep_dispatch_schedule.cpp ) + +add_executable(test_tilexr_moonep_dispatch_urma_host + unit/test_tilexr_moonep_dispatch_urma_host.cpp + ${TILEXR_ROOT}/src/moonep/dispatch/urma/host/dispatch_layout.cpp + ${TILEXR_ROOT}/src/moonep/dispatch/urma/host/dispatch_host.cpp +) +target_include_directories(test_tilexr_moonep_dispatch_urma_host PRIVATE + ${TILEXR_MOONEP_TEST_INCLUDE_DIRS} + ${TILEXR_ROOT}/src/moonep/dispatch/urma/common + ${TILEXR_ROOT}/src/moonep/dispatch/urma/host +) + +add_executable(test_tilexr_moonep_dispatch_urma_launch + unit/test_tilexr_moonep_dispatch_urma_launch.cpp + ${TILEXR_ROOT}/src/moonep/dispatch/urma/host/dispatch_layout.cpp + ${TILEXR_ROOT}/src/moonep/dispatch/urma/host/dispatch_launch.cpp +) +target_include_directories(test_tilexr_moonep_dispatch_urma_launch PRIVATE + ${TILEXR_MOONEP_TEST_INCLUDE_DIRS} + ${TILEXR_ROOT}/src/moonep/dispatch/urma/common + ${TILEXR_ROOT}/src/moonep/dispatch/urma/host +) target_include_directories(test_tilexr_moonep_dispatch_urma_schedule PRIVATE ${TILEXR_MOONEP_TEST_INCLUDE_DIRS} ${TILEXR_ROOT}/src/moonep/dispatch/urma/common @@ -248,6 +273,10 @@ add_test(NAME test_tilexr_moonep_dispatch_urma_layout COMMAND test_tilexr_moonep_dispatch_urma_layout) add_test(NAME test_tilexr_moonep_dispatch_urma_schedule COMMAND test_tilexr_moonep_dispatch_urma_schedule) +add_test(NAME test_tilexr_moonep_dispatch_urma_host + COMMAND test_tilexr_moonep_dispatch_urma_host) +add_test(NAME test_tilexr_moonep_dispatch_urma_launch + COMMAND test_tilexr_moonep_dispatch_urma_launch) add_test(NAME test_tilexr_moonep_dispatch_host COMMAND test_tilexr_moonep_dispatch_host) add_test(NAME test_tilexr_moonep_combine_host COMMAND test_tilexr_moonep_combine_host) add_test(NAME test_tilexr_moonep_prefetch_weight_host COMMAND test_tilexr_moonep_prefetch_weight_host) diff --git a/tests/moonep/python/test_dispatch_hot_loop_diagnostics.py b/tests/moonep/python/test_dispatch_hot_loop_diagnostics.py new file mode 100644 index 0000000..877d945 --- /dev/null +++ b/tests/moonep/python/test_dispatch_hot_loop_diagnostics.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +import ctypes +from types import SimpleNamespace + +from tools.moonep import dispatch_hot_loop + + +def _profile(payload_mode: int, *, shared_owner: bool): + record = dispatch_hot_loop.DispatchProfileRecord() + record.marker = dispatch_hot_loop.PROFILE_MARKER + record.version = dispatch_hot_loop.DIAGNOSTIC_VERSION + record.record_bytes = ctypes.sizeof(dispatch_hot_loop.DispatchProfileRecord) + record.payload_mode = payload_mode + record.rank = 0 + record.core = 0 + record.block_dim = 1 + record.select_mode = 1 + record.fallback_reason = 0 + record.magic = 17 + record.processed = 8 + record.put_count = 8 + if shared_owner: + record.scanned = 16 + record.matched = 8 + record.selected = 8 + record.visited_peers = 1 + record.completion_flags = 2 + record.kernel_cycles = 110 + record.staging_cycles = 10 + record.put_issue_cycles = 20 + record.flag_wait_cycles = 30 + record.output_copy_cycles = 40 + record.quiet_cycles = 10 + for index in range(11): + record.reserved[index] = (index + 1) * 10 + return record + + +def _dfx(payload_mode: int): + record = dispatch_hot_loop.DispatchDfxRecord() + record.marker = dispatch_hot_loop.DFX_MARKER + record.version = dispatch_hot_loop.DIAGNOSTIC_VERSION + record.record_bytes = ctypes.sizeof(dispatch_hot_loop.DispatchDfxRecord) + record.payload_mode = payload_mode + record.rank = 0 + record.core = 0 + record.magic = 17 + record.expected_routes = 16 + record.processed_routes = 8 + return record + + +def test_paired_fused_diagnostics_marks_one_shared_epoch_owner(monkeypatch) -> None: + monkeypatch.setenv("TILEXR_MOONEP_DISPATCH_AIV_CORE_COUNT", "1") + context = SimpleNamespace( + planner_group_rank=0, + _dispatch_workspace_bytes=( + dispatch_hot_loop.COMPLETION_BYTES + + dispatch_hot_loop.SIGNAL_BYTES + + 2 * dispatch_hot_loop.PROFILE_COUNT + * ctypes.sizeof(dispatch_hot_loop.DispatchProfileRecord) + + 2 * dispatch_hot_loop.DFX_COUNT + * ctypes.sizeof(dispatch_hot_loop.DispatchDfxRecord) + + ctypes.sizeof(dispatch_hot_loop.DispatchKernelStatus) + ), + ) + offsets = dispatch_hot_loop._diagnostic_layout(context) + + hidden_profile = bytearray( + dispatch_hot_loop.PROFILE_COUNT + * ctypes.sizeof(dispatch_hot_loop.DispatchProfileRecord) + ) + weight_profile = bytearray(hidden_profile) + hidden_profile[: ctypes.sizeof(dispatch_hot_loop.DispatchProfileRecord)] = bytes( + _profile(0, shared_owner=False) + ) + weight_profile[: ctypes.sizeof(dispatch_hot_loop.DispatchProfileRecord)] = bytes( + _profile(1, shared_owner=True) + ) + hidden_dfx = bytearray( + dispatch_hot_loop.DFX_COUNT + * ctypes.sizeof(dispatch_hot_loop.DispatchDfxRecord) + ) + weight_dfx = bytearray(hidden_dfx) + hidden_dfx[: ctypes.sizeof(dispatch_hot_loop.DispatchDfxRecord)] = bytes( + _dfx(0) + ) + weight_dfx[: ctypes.sizeof(dispatch_hot_loop.DispatchDfxRecord)] = bytes( + _dfx(1) + ) + status = dispatch_hot_loop.DispatchKernelStatus() + status.marker = dispatch_hot_loop.KERNEL_STATUS_MARKER + status.version = dispatch_hot_loop.DIAGNOSTIC_VERSION + status.record_bytes = ctypes.sizeof(dispatch_hot_loop.DispatchKernelStatus) + status.payload_mode = 1 + status.magic = 17 + status.reserved[0] = ( + dispatch_hot_loop.KERNEL_STATUS_FEATURE_DFX_ENABLED + | dispatch_hot_loop.KERNEL_STATUS_FEATURE_PROFILING_ENABLED + | dispatch_hot_loop.KERNEL_STATUS_FEATURE_FUSED_EPOCH + ) + blobs = { + offsets["hidden_profile"]: bytes(hidden_profile), + offsets["weight_profile"]: bytes(weight_profile), + offsets["hidden_dfx"]: bytes(hidden_dfx), + offsets["weight_dfx"]: bytes(weight_dfx), + offsets["kernel_status"]: bytes(status), + } + + monkeypatch.setattr( + dispatch_hot_loop, + "_workspace_blob", + lambda unused_context, byte_offset, byte_count: blobs[byte_offset][ + :byte_count + ], + ) + + result = dispatch_hot_loop._diagnostics(context) + + assert result["fused_epoch"] is True + assert result["shared_owner_mode"] == "weight" + assert result["kernel_status"]["payload_mode"] == 1 + assert result["hidden"]["shared_epoch_owner"] is False + assert result["weight"]["shared_epoch_owner"] is True + assert result["hidden"]["profile"][0]["kernel_cycles"] == 0 + assert result["hidden"]["profile"][0]["reserved"] == [0] * 11 + assert result["weight"]["profile"][0]["kernel_cycles"] == 110 + assert len(result["hidden"]["dfx"]) == 1 + assert len(result["weight"]["dfx"]) == 1 + + +def test_non_fused_diagnostics_do_not_claim_a_shared_owner(monkeypatch) -> None: + monkeypatch.setenv("TILEXR_MOONEP_DISPATCH_AIV_CORE_COUNT", "1") + context = SimpleNamespace( + planner_group_rank=0, + _dispatch_workspace_bytes=( + dispatch_hot_loop.COMPLETION_BYTES + + dispatch_hot_loop.SIGNAL_BYTES + + 2 * dispatch_hot_loop.PROFILE_COUNT + * ctypes.sizeof(dispatch_hot_loop.DispatchProfileRecord) + + 2 * dispatch_hot_loop.DFX_COUNT + * ctypes.sizeof(dispatch_hot_loop.DispatchDfxRecord) + + ctypes.sizeof(dispatch_hot_loop.DispatchKernelStatus) + ), + ) + offsets = dispatch_hot_loop._diagnostic_layout(context) + status = dispatch_hot_loop.DispatchKernelStatus() + status.marker = dispatch_hot_loop.KERNEL_STATUS_MARKER + status.version = dispatch_hot_loop.DIAGNOSTIC_VERSION + status.record_bytes = ctypes.sizeof(dispatch_hot_loop.DispatchKernelStatus) + status.payload_mode = 0 + status.magic = 17 + status.reserved[0] = dispatch_hot_loop.KERNEL_STATUS_FEATURE_DFX_ENABLED + empty_profile = bytes( + dispatch_hot_loop.PROFILE_COUNT + * ctypes.sizeof(dispatch_hot_loop.DispatchProfileRecord) + ) + empty_dfx = bytes( + dispatch_hot_loop.DFX_COUNT + * ctypes.sizeof(dispatch_hot_loop.DispatchDfxRecord) + ) + blobs = { + offsets["hidden_profile"]: empty_profile, + offsets["weight_profile"]: empty_profile, + offsets["hidden_dfx"]: empty_dfx, + offsets["weight_dfx"]: empty_dfx, + offsets["kernel_status"]: bytes(status), + } + monkeypatch.setattr( + dispatch_hot_loop, + "_workspace_blob", + lambda unused_context, byte_offset, byte_count: blobs[byte_offset][ + :byte_count + ], + ) + + result = dispatch_hot_loop._diagnostics( + context, required_modes=() + ) + + assert result["fused_epoch"] is False + assert result["shared_owner_mode"] is None + assert result["hidden"]["shared_epoch_owner"] is False + assert result["weight"]["shared_epoch_owner"] is False + + +def test_skewed_hot_loop_routes_concentrate_on_the_same_experts() -> None: + case = SimpleNamespace( + tokens_per_rank=3, + topk=2, + expert_count=8, + routing_pattern="skewed", + route_distribution="rank_shifted_uniform", + ) + + rank_zero = dispatch_hot_loop._case_rank_topk(case, 0, 4) + rank_three = dispatch_hot_loop._case_rank_topk(case, 3, 4) + + assert rank_zero == (0, 1, 0, 1, 0, 1) + assert rank_three == rank_zero + assert dispatch_hot_loop._all_case_topk(case, 4) == rank_zero * 4 + + +def test_reference_slots_leave_padding_for_zero_fill() -> None: + case = SimpleNamespace( + tokens_per_rank=2, + topk=1, + expert_count=2, + routing_pattern="balanced", + route_distribution="moonep_combine_balanced", + ) + context = SimpleNamespace( + planner_group_rank=0, + planner_group_size=2, + nv_s=8, + prefetch_slots=1, + token_padding=4, + ) + + sources, tokens, routes = dispatch_hot_loop._reference_slot_assignments( + case, context + ) + + assert len(sources) == len(tokens) == len(routes) == context.nv_s + assert sum(source >= 0 for source in sources) == 2 + assert sum(source < 0 for source in sources) == 6 + + +class _FakePlan: + def __init__(self) -> None: + import torch + + self.dst = torch.tensor([0, 1], dtype=torch.int32) + + +class _FakeBuffer: + def __init__(self) -> None: + self.runtime = SimpleNamespace() + self.context = SimpleNamespace(dispatch_workspace=(object(), 4096)) + self.quiesce_calls = 0 + + def _stream_ptr(self): + return object() + + def quiesce(self) -> None: + self.quiesce_calls += 1 + + +def test_repeated_exact_check_compares_every_round_and_alternates_plan( + monkeypatch, +) -> None: + import torch + + case = SimpleNamespace(topk=1) + plan = _FakePlan() + buffer = _FakeBuffer() + hidden_out = torch.zeros((2, 1), dtype=torch.bfloat16) + weights_out = torch.zeros((2,), dtype=torch.float32) + observed = [] + + monkeypatch.setattr(torch, "npu", SimpleNamespace(synchronize=lambda: None), + raising=False) + monkeypatch.setattr( + dispatch_hot_loop, + "_launch_dispatch", + lambda unused_runtime, unused_context, active_plan, *unused_args: ( + observed.append(tuple(active_plan.dst.tolist())) + ), + ) + monkeypatch.setattr( + dispatch_hot_loop, + "_expected", + lambda *unused_args, destination_roll=0: ( + hidden_out.clone(), weights_out.clone(), [0, 0] + ), + ) + + result = dispatch_hot_loop._repeated_exact_check( + torch, case, buffer, plan, object(), hidden_out, object(), weights_out, + "pair", 4 + ) + + assert result["rounds"] == 4 + assert result["hidden_exact"] is True + assert result["weight_exact"] is True + assert observed == [(0, 1), (1, 0), (0, 1), (1, 0)] + assert buffer.quiesce_calls == 4 + assert plan.dst.tolist() == [0, 1] diff --git a/tests/moonep/python/test_moonep_modes.py b/tests/moonep/python/test_moonep_modes.py index ace4a18..b13a0ab 100644 --- a/tests/moonep/python/test_moonep_modes.py +++ b/tests/moonep/python/test_moonep_modes.py @@ -420,13 +420,17 @@ def test_single_node_launcher_builds_hidden_dispatch_hot_loop_command(tmp_path) "dispatch_hot_loop", "--dispatch-modes", "hidden", + "--exact-rounds", + "100", ] ) command = _process_command(args) assert "tools.moonep.dispatch_hot_loop" in command assert "tools.moonep.benchmark" not in command - assert command[command.index("--dispatch-modes") + 1 :] == ["hidden"] + assert command[command.index("--dispatch-modes") + 1 :] == [ + "hidden", "--exact-rounds", "100" + ] assert "--mode" not in command diff --git a/tests/moonep/unit/test_tilexr_moonep_dispatch_layout.cpp b/tests/moonep/unit/test_tilexr_moonep_dispatch_layout.cpp index 2d68d72..4b0debd 100644 --- a/tests/moonep/unit/test_tilexr_moonep_dispatch_layout.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_dispatch_layout.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include "dispatch_layout.h" #include "dispatch_profile.h" @@ -16,6 +17,39 @@ int g_failures = 0; if (lhsValue != rhsValue) { std::cerr << "CHECK_EQ line " << __LINE__ \ << ": " #lhs " != " #rhs << std::endl; ++g_failures; } } while (0) +struct Range { + uint64_t offset; + uint64_t bytes; +}; + +void CheckDisjointLayout(const TileXRMoonEp::MoonEpDispatchUrmaLayout &layout) +{ + const std::vector ranges { + {layout.hidden.sourceOffset, layout.hidden.sourceBytes}, + {layout.hidden.scratchOffset, layout.hidden.scratchBytes}, + {layout.weight.sourceOffset, layout.weight.sourceBytes}, + {layout.weight.scratchOffset, layout.weight.scratchBytes}, + {layout.completionFlagsOffset, layout.completionFlagsBytes}, + {layout.signalOffset, layout.signalBytes}, + {layout.hiddenProfileOffset, layout.profileBytes}, + {layout.weightProfileOffset, layout.profileBytes}, + {layout.hiddenDfxOffset, layout.dfxBytes}, + {layout.weightDfxOffset, layout.dfxBytes}, + {layout.kernelStatusOffset, layout.kernelStatusBytes}, + }; + for (std::size_t index = 0; index < ranges.size(); ++index) { + const Range ¤t = ranges[index]; + CHECK_TRUE(current.bytes > 0U); + CHECK_TRUE(current.offset <= layout.totalBytes); + CHECK_TRUE(current.bytes <= layout.totalBytes - current.offset); + for (std::size_t other = index + 1; other < ranges.size(); ++other) { + const Range &candidate = ranges[other]; + CHECK_TRUE(current.offset + current.bytes <= candidate.offset || + candidate.offset + candidate.bytes <= current.offset); + } + } +} + void TestReferenceShape() { TileXRMoonEp::MoonEpDispatchUrmaLayout layout {}; @@ -30,6 +64,15 @@ void TestReferenceShape() CHECK_EQ(layout.weight.rowBytes, UINT64_C(4)); CHECK_EQ(layout.weight.sourceBytes, UINT64_C(8192)); CHECK_EQ(layout.weight.scratchSlotBytes, UINT64_C(8192)); + CHECK_EQ(layout.hidden.sourceOffset, UINT64_C(0)); + CHECK_TRUE(layout.hidden.scratchOffset >= + layout.hidden.sourceOffset + layout.hidden.sourceBytes); + CHECK_TRUE(layout.weight.sourceOffset >= + layout.hidden.scratchOffset + layout.hidden.scratchBytes); + CHECK_TRUE(layout.weight.scratchOffset >= + layout.weight.sourceOffset + layout.weight.sourceBytes); + CHECK_TRUE(layout.commonOffset >= + layout.weight.scratchOffset + layout.weight.scratchBytes); CHECK_EQ(layout.completionFlagsBytes, UINT64_C(8192)); CHECK_TRUE(layout.signalOffset >= layout.completionFlagsOffset + layout.completionFlagsBytes); @@ -37,11 +80,11 @@ void TestReferenceShape() UINT64_C(64) * sizeof(TileXRMoonEp::DispatchProfileRecord)); CHECK_EQ(layout.dfxBytes, UINT64_C(64) * sizeof(TileXRMoonEp::DispatchDfxRecord)); - CHECK_TRUE(layout.commonOffset >= layout.hidden.activeDataBytes); - CHECK_TRUE(layout.commonOffset >= layout.weight.activeDataBytes); + CHECK_EQ(layout.totalBytes, UINT64_C(30) * 1024U * 1024U); CHECK_TRUE(layout.requiredBytes <= layout.totalBytes); CHECK_EQ(layout.totalBytes % TileXRMoonEp::kDispatchRegistrationAlignmentBytes, UINT64_C(0)); + CheckDisjointLayout(layout); TileXRMoonEp::MoonEpDispatchUrmaLayout expanded = layout; const uint64_t oldCommonOffset = expanded.commonOffset; @@ -50,6 +93,7 @@ void TestReferenceShape() &expanded), TileXR::TILEXR_SUCCESS); CHECK_EQ(expanded.commonOffset, oldCommonOffset + TileXRMoonEp::kDispatchRegistrationAlignmentBytes); + CheckDisjointLayout(expanded); CHECK_TRUE(TileXRMoonEp::TileXRMoonEpBindDispatchUrmaWorkspace( layout.totalBytes - TileXRMoonEp::kDispatchRegistrationAlignmentBytes, &layout) != TileXR::TILEXR_SUCCESS); @@ -82,6 +126,19 @@ void TestFailuresAndModes() layout, TileXRMoonEp::DispatchPayloadMode::Hidden) == &layout.hidden); CHECK_TRUE(TileXRMoonEp::TileXRMoonEpGetActiveDispatchUrmaLayout( layout, TileXRMoonEp::DispatchPayloadMode::RouteWeight) == &layout.weight); + + CHECK_EQ(TileXRMoonEp::TileXRMoonEpBuildDispatchUrmaLayout( + 1, 1, 1, 1, 1, &layout), TileXR::TILEXR_SUCCESS); + CheckDisjointLayout(layout); + CHECK_EQ(layout.hidden.rowBytes, UINT64_C(2)); + CHECK_EQ(layout.weight.rowBytes, UINT64_C(4)); + + CHECK_EQ(TileXRMoonEp::TileXRMoonEpBuildDispatchUrmaLayout( + 3, 3, 2, 3, 7, &layout), TileXR::TILEXR_SUCCESS); + CheckDisjointLayout(layout); + CHECK_EQ(layout.routeCount, 6); + CHECK_EQ(layout.destinationCapacity, 7); + CHECK_EQ(layout.hidden.rowBytes, UINT64_C(6)); } } // namespace diff --git a/tests/moonep/unit/test_tilexr_moonep_dispatch_schedule.cpp b/tests/moonep/unit/test_tilexr_moonep_dispatch_schedule.cpp index 6c77e2e..722abb3 100644 --- a/tests/moonep/unit/test_tilexr_moonep_dispatch_schedule.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_dispatch_schedule.cpp @@ -266,6 +266,30 @@ void TestWqeBatchBoundaries() CHECK_EQ(TileXRMoonEp::DispatchRouteTileCount(1000U, 0U, 1024U), 1000U); CHECK_EQ(TileXRMoonEp::DispatchRouteTileCount(1000U, 0U, 0U), 0U); + uint64_t dataWqes = UINT64_MAX; + CHECK_TRUE(TileXRMoonEp::DispatchDataWqeCount(17U, false, dataWqes)); + CHECK_EQ(dataWqes, UINT64_C(17)); + CHECK_TRUE(TileXRMoonEp::DispatchDataWqeCount(17U, true, dataWqes)); + CHECK_EQ(dataWqes, UINT64_C(34)); + CHECK_TRUE(!TileXRMoonEp::DispatchDataWqeCount( + UINT64_MAX, true, dataWqes)); + CHECK_EQ(dataWqes, UINT64_C(0)); + + for (uint32_t route = 0U; route < 129U; ++route) { + const uint32_t hiddenTask = route * 2U; + const uint32_t weightTask = hiddenTask + 1U; + CHECK_EQ(TileXRMoonEp::DispatchDataTaskRouteIndex(hiddenTask, true), + route); + CHECK_EQ(TileXRMoonEp::DispatchDataTaskRouteIndex(weightTask, true), + route); + CHECK_TRUE(!TileXRMoonEp::DispatchDataTaskIsWeight(hiddenTask, true)); + CHECK_TRUE(TileXRMoonEp::DispatchDataTaskIsWeight(weightTask, true)); + } + CHECK_TRUE(!TileXRMoonEp::DispatchSignalFitsAfterData(64U, true, 128U)); + CHECK_TRUE(TileXRMoonEp::DispatchSignalFitsAfterData(63U, true, 128U)); + CHECK_TRUE(TileXRMoonEp::DispatchSignalFitsAfterData(127U, false, 128U)); + CHECK_TRUE(!TileXRMoonEp::DispatchSignalFitsAfterData(128U, false, 128U)); + } void TestSparseCqBatchAndOwnerGeneration() diff --git a/tests/moonep/unit/test_tilexr_moonep_dispatch_urma_host.cpp b/tests/moonep/unit/test_tilexr_moonep_dispatch_urma_host.cpp new file mode 100644 index 0000000..ded0176 --- /dev/null +++ b/tests/moonep/unit/test_tilexr_moonep_dispatch_urma_host.cpp @@ -0,0 +1,263 @@ +#include +#include +#include +#include + +#include "acl/acl_rt.h" +#include "dispatch_host.h" +#include "dispatch_launch.h" +#include "tilexr_types.h" + +namespace { + +int failures = 0; +int hostArgsCalls = 0; +int devArgsCalls = 0; +int launchCalls = 0; +int memsetCalls = 0; +int synchronizeCalls = 0; +int launchReturn = TileXR::TILEXR_SUCCESS; +aclError memsetReturn = ACL_SUCCESS; +aclError synchronizeReturn = ACL_SUCCESS; +TileXR::CommArgs commArgs {}; +GM_ADDR devArgs = reinterpret_cast(uintptr_t {0x9000}); +TileXRMoonEp::DispatchUrmaLaunchParams launchedParams {}; + +void Check(bool condition, const std::string &message) +{ + if (!condition) { + std::cerr << message << '\n'; + ++failures; + } +} + +void CheckStatus(const char *label, int actual, int expected) +{ + if (actual != expected) { + std::cerr << label << ": expected " << expected << ", got " << actual << '\n'; + ++failures; + } +} + +void Reset() +{ + hostArgsCalls = devArgsCalls = launchCalls = 0; + memsetCalls = synchronizeCalls = 0; + launchReturn = TileXR::TILEXR_SUCCESS; + memsetReturn = synchronizeReturn = ACL_SUCCESS; + commArgs = TileXR::CommArgs {}; + commArgs.rank = 0; + commArgs.localRank = 0; + commArgs.rankSize = 1; + commArgs.localRankSize = 1; + devArgs = reinterpret_cast(uintptr_t {0x9000}); + launchedParams = TileXRMoonEp::DispatchUrmaLaunchParams {}; +} + +TileXRMoonEpPlanV1 ValidPlan() +{ + TileXRMoonEpPlanV1 plan {}; + plan.structSize = sizeof(plan); + plan.abiVersion = TILEXR_MOONEP_ABI_VERSION_V1; + plan.n = 4; + plan.r = 1; + plan.e = 8; + plan.b = 2; + plan.nvS = 4; + plan.k = 2; + plan.dst = reinterpret_cast(uintptr_t {0x3000}); + plan.expertsToCopy = reinterpret_cast(uintptr_t {0x3100}); + plan.zeroFillRanges = reinterpret_cast(uintptr_t {0x3200}); + plan.remoteStats = reinterpret_cast(uintptr_t {0x3300}); + plan.dupGroups = reinterpret_cast(uintptr_t {0x3400}); + plan.dupLoffs = reinterpret_cast(uintptr_t {0x3500}); + plan.dupCounts = reinterpret_cast(uintptr_t {0x3600}); + plan.status = reinterpret_cast(uintptr_t {0x3700}); + return plan; +} + +TileXRMoonEpTensorV1 Tensor(void *data, uint32_t dtype, uint32_t rank, + int64_t dim0, int64_t dim1, uint64_t elements) +{ + TileXRMoonEpTensorV1 tensor {}; + tensor.structSize = sizeof(tensor); + tensor.abiVersion = TILEXR_MOONEP_ABI_VERSION_V1; + tensor.data = data; + tensor.elementCount = elements; + tensor.dtype = dtype; + tensor.rank = rank; + tensor.shape[0] = dim0; + tensor.shape[1] = dim1; + return tensor; +} + +TileXRMoonEpDispatchArgsV2 Args(const TileXRMoonEpPlanV1 *plan, + const TileXRMoonEpTensorV1 *hiddenInput, + const TileXRMoonEpTensorV1 *weightInput, + TileXRMoonEpTensorV1 *hiddenOutput, + TileXRMoonEpTensorV1 *weightOutput) +{ + TileXRMoonEpDispatchArgsV2 args {}; + args.structSize = sizeof(args); + args.abiVersion = TILEXR_MOONEP_ABI_VERSION_V2; + args.comm = reinterpret_cast(uintptr_t {0x1000}); + args.plan = plan; + args.hiddenSh = hiddenInput; + args.routeWeightsSk = weightInput; + args.hiddenNvsh = hiddenOutput; + args.routeWeightsNvs = weightOutput; + args.flags = TILEXR_MOONEP_FLAG_RESET_STATUS; + args.registeredWorkspace = reinterpret_cast(uintptr_t {0x200000}); + args.registeredWorkspaceBytes = UINT64_C(2) * 1024U * 1024U; + return args; +} + +void TestPairedSingleLaunch() +{ + Reset(); + TileXRMoonEpPlanV1 plan = ValidPlan(); + TileXRMoonEpTensorV1 hiddenInput = Tensor( + reinterpret_cast(uintptr_t {0x4000}), + TILEXR_MOONEP_DTYPE_BFLOAT16, 2, 2, 17, 34); + TileXRMoonEpTensorV1 hiddenOutput = Tensor( + reinterpret_cast(uintptr_t {0x5000}), + TILEXR_MOONEP_DTYPE_BFLOAT16, 2, 4, 17, 68); + TileXRMoonEpTensorV1 weightInput = Tensor( + reinterpret_cast(uintptr_t {0x6000}), + TILEXR_MOONEP_DTYPE_FLOAT32, 2, 2, 2, 4); + TileXRMoonEpTensorV1 weightOutput = Tensor( + reinterpret_cast(uintptr_t {0x7000}), + TILEXR_MOONEP_DTYPE_FLOAT32, 1, 4, 0, 4); + TileXRMoonEpDispatchArgsV2 args = Args(&plan, &hiddenInput, &weightInput, + &hiddenOutput, &weightOutput); + aclrtStream stream = reinterpret_cast(uintptr_t {0x8000}); + + CheckStatus("paired", TileXRMoonEp::TileXRMoonEpRunDispatchUrmaV2( + &args, stream), TILEXR_MOONEP_SUCCESS); + Check(hostArgsCalls == 1 && devArgsCalls == 1 && launchCalls == 1, + "paired Dispatch must make exactly one internal launch"); + Check(memsetCalls == 1 && synchronizeCalls == 0, + "paired success must enqueue one reset without synchronizing"); + Check(launchedParams.hiddenInput == hiddenInput.data && + launchedParams.hiddenOutput == hiddenOutput.data && + launchedParams.weightInput == weightInput.data && + launchedParams.weightOutput == weightOutput.data, + "paired launch did not carry all four payload pointers"); + Check(launchedParams.layout.weight.sourceOffset >= + launchedParams.layout.hidden.scratchOffset + + launchedParams.layout.hidden.scratchBytes, + "paired launch layout overlaps Hidden and Weight active regions"); +} + +void TestHiddenOnlyAndFailureBoundaries() +{ + Reset(); + TileXRMoonEpPlanV1 plan = ValidPlan(); + TileXRMoonEpTensorV1 hiddenInput = Tensor( + reinterpret_cast(uintptr_t {0x4000}), + TILEXR_MOONEP_DTYPE_FLOAT16, 2, 2, 16, 32); + TileXRMoonEpTensorV1 hiddenOutput = Tensor( + reinterpret_cast(uintptr_t {0x5000}), + TILEXR_MOONEP_DTYPE_FLOAT16, 2, 4, 16, 64); + TileXRMoonEpDispatchArgsV2 args = Args( + &plan, &hiddenInput, nullptr, &hiddenOutput, nullptr); + aclrtStream stream = reinterpret_cast(uintptr_t {0x8000}); + + CheckStatus("hidden", TileXRMoonEp::TileXRMoonEpRunDispatchUrmaV2( + &args, stream), TILEXR_MOONEP_SUCCESS); + Check(launchCalls == 1 && launchedParams.weightInput == nullptr && + launchedParams.weightOutput == nullptr, + "hidden-only Dispatch must launch once without Weight pointers"); + + Reset(); + launchReturn = TileXR::TILEXR_ERROR_MKIRT; + CheckStatus("launch failure", TileXRMoonEp::TileXRMoonEpRunDispatchUrmaV2( + &args, stream), TILEXR_MOONEP_ERROR_INTERNAL); + Check(launchCalls == 1 && synchronizeCalls == 1, + "same-stream reset must be synchronized after launch failure"); + + Reset(); + memsetReturn = 1; + CheckStatus("reset failure", TileXRMoonEp::TileXRMoonEpRunDispatchUrmaV2( + &args, stream), TILEXR_MOONEP_ERROR_INTERNAL); + Check(launchCalls == 0, "reset failure must prevent launch"); +} + +} // namespace + +extern "C" int TileXRGetCommArgsHost(TileXRCommPtr, TileXR::CommArgs *&args) +{ + ++hostArgsCalls; + args = &commArgs; + return TileXR::TILEXR_SUCCESS; +} + +extern "C" int TileXRGetCommArgsDev(TileXRCommPtr, GM_ADDR &args) +{ + ++devArgsCalls; + args = devArgs; + return TileXR::TILEXR_SUCCESS; +} + +extern "C" int TileXRUDMAGetQpCount(TileXRCommPtr, uint32_t *qpCount) +{ + if (qpCount != nullptr) { + *qpCount = 2U; + } + return TileXR::TILEXR_SUCCESS; +} + +extern "C" int TileXRGetUDMARegistryHost( + TileXRCommPtr, const TileXR::TileXRUDMARegistry **registry) +{ + if (registry != nullptr) { + *registry = nullptr; + } + return TileXR::TILEXR_SUCCESS; +} + +extern "C" aclError aclrtMemsetAsync( + void *, size_t, int32_t, size_t, aclrtStream) +{ + ++memsetCalls; + return memsetReturn; +} + +extern "C" aclError aclrtSynchronizeStream(aclrtStream) +{ + ++synchronizeCalls; + return synchronizeReturn; +} + +extern "C" aclError aclrtGetDevice(int32_t *deviceId) +{ + if (deviceId != nullptr) { + *deviceId = 0; + } + return ACL_SUCCESS; +} + +extern "C" aclError aclrtGetDeviceInfo( + uint32_t, aclrtDevAttr, int64_t *value) +{ + if (value != nullptr) { + *value = TileXRMoonEp::kDispatchAivCoreCount; + } + return ACL_SUCCESS; +} + +namespace TileXRMoonEp { +int TileXRMoonEpLaunchDispatchUrmaKernel(const DispatchUrmaLaunchParams ¶ms) +{ + ++launchCalls; + launchedParams = params; + return launchReturn; +} +} // namespace TileXRMoonEp + +int main() +{ + TestPairedSingleLaunch(); + TestHiddenOnlyAndFailureBoundaries(); + return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/tests/moonep/unit/test_tilexr_moonep_dispatch_urma_launch.cpp b/tests/moonep/unit/test_tilexr_moonep_dispatch_urma_launch.cpp new file mode 100644 index 0000000..61e09a1 --- /dev/null +++ b/tests/moonep/unit/test_tilexr_moonep_dispatch_urma_launch.cpp @@ -0,0 +1,215 @@ +#include +#include +#include +#include + +#include "acl/acl_rt.h" +#include "dispatch_launch.h" +#include "dispatch_layout.h" +#include "runtime/kernel.h" +#include "tilexr_types.h" + +extern "C" { +extern const unsigned char TileXRMoonEpDispatchUrmaKernelBinaryData[] = { + 0x7f, 'E', 'L', 'F'}; +extern const std::size_t TileXRMoonEpDispatchUrmaKernelBinarySize = + sizeof(TileXRMoonEpDispatchUrmaKernelBinaryData); +} + +namespace { + +int failures = 0; +int magicCalls = 0; +int binaryRegisterCalls = 0; +int functionRegisterCalls = 0; +int launchCalls = 0; +int magicReturn = TileXR::TILEXR_SUCCESS; +int64_t nextMagic = 41; +rtError_t launchReturn = RT_ERROR_NONE; +uint32_t launchedBlockDim = 0U; +size_t launchedArgsSize = 0U; +uint32_t launchedLocalMemorySize = 0U; +TileXRMoonEp::DispatchKernelArgs launchedArgs {}; + +void Check(bool condition, const char *message) +{ + if (!condition) { + std::cerr << message << std::endl; + ++failures; + } +} + +TileXRMoonEp::DispatchUrmaLaunchParams Params(bool paired) +{ + TileXRMoonEp::DispatchUrmaLaunchParams params {}; + params.commArgs = reinterpret_cast(uintptr_t {0x1000}); + params.hiddenInput = reinterpret_cast(uintptr_t {0x2000}); + params.weightInput = paired ? + reinterpret_cast(uintptr_t {0x3000}) : nullptr; + params.dst = reinterpret_cast(uintptr_t {0x4000}); + params.zeroFillRanges = reinterpret_cast(uintptr_t {0x5000}); + params.workspace = reinterpret_cast(uintptr_t {0x6000}); + params.hiddenOutput = reinterpret_cast(uintptr_t {0x7000}); + params.weightOutput = paired ? + reinterpret_cast(uintptr_t {0x8000}) : nullptr; + params.planStatus = reinterpret_cast(uintptr_t {0x9000}); + params.comm = reinterpret_cast(uintptr_t {0xa000}); + params.stream = reinterpret_cast(uintptr_t {0xb000}); + params.peerMode = TileXRMoonEp::DispatchPeerMode::Legacy; + params.groupWidth = TileXRMoonEp::kDispatchDefaultGroupWidth; + params.zeroFillRangeCount = 10; + Check(TileXRMoonEp::TileXRMoonEpBuildDispatchUrmaLayout( + 4, 2, 2, 17, 4, ¶ms.layout) == TileXR::TILEXR_SUCCESS, + "failed to build launcher test layout"); + return params; +} + +void TestPairedAndHiddenOnly() +{ + TileXRMoonEp::DispatchUrmaLaunchParams paired = Params(true); + Check(TileXRMoonEp::TileXRMoonEpLaunchDispatchUrmaKernel(paired) == + TileXR::TILEXR_SUCCESS, "paired launcher failed"); + Check(magicCalls == 1 && launchCalls == 1, + "paired launcher must take one magic and launch once"); + Check(binaryRegisterCalls == 1 && functionRegisterCalls == 1, + "paired launcher did not register the embedded kernel once"); + Check(launchedArgsSize == sizeof(TileXRMoonEp::DispatchKernelArgs) && + launchedBlockDim == TileXRMoonEp::kDispatchAivCoreCount, + "paired launch metadata mismatch"); + Check(launchedArgs.hiddenInput == paired.hiddenInput && + launchedArgs.weightInput == paired.weightInput && + launchedArgs.hiddenOutput == paired.hiddenOutput && + launchedArgs.weightOutput == paired.weightOutput, + "paired Kernel ABI pointer order mismatch"); + Check(launchedArgs.hiddenSourceOffset == paired.layout.hidden.sourceOffset && + launchedArgs.hiddenScratchOffset == paired.layout.hidden.scratchOffset && + launchedArgs.hiddenRowBytes == paired.layout.hidden.rowBytes && + launchedArgs.weightSourceOffset == paired.layout.weight.sourceOffset && + launchedArgs.weightScratchOffset == paired.layout.weight.scratchOffset && + launchedArgs.weightRowBytes == paired.layout.weight.rowBytes, + "paired Kernel ABI active layout order mismatch"); + Check(launchedArgs.hiddenProfileOffset == paired.layout.hiddenProfileOffset && + launchedArgs.weightProfileOffset == paired.layout.weightProfileOffset && + launchedArgs.hiddenDfxOffset == paired.layout.hiddenDfxOffset && + launchedArgs.weightDfxOffset == paired.layout.weightDfxOffset && + launchedArgs.hasWeight == 1U && launchedArgs.magic == nextMagic, + "paired Kernel ABI diagnostics or epoch fields mismatch"); + Check(launchedLocalMemorySize >= 190U * 1024U, + "paired launch did not configure the required dynamic UB"); + + TileXRMoonEp::DispatchUrmaLaunchParams hidden = Params(false); + nextMagic = 42; + Check(TileXRMoonEp::TileXRMoonEpLaunchDispatchUrmaKernel(hidden) == + TileXR::TILEXR_SUCCESS, "hidden-only launcher failed"); + Check(magicCalls == 2 && launchCalls == 2 && + binaryRegisterCalls == 1 && functionRegisterCalls == 1, + "hidden-only launch did not reuse one registered function and one epoch"); + Check(launchedArgs.weightInput == nullptr && + launchedArgs.weightOutput == nullptr && launchedArgs.hasWeight == 0U && + launchedArgs.magic == nextMagic, + "hidden-only Kernel ABI optional Weight contract mismatch"); +} + +void TestFailureBoundaries() +{ + TileXRMoonEp::DispatchUrmaLaunchParams params = Params(true); + params.groupWidth = 7U; + Check(TileXRMoonEp::TileXRMoonEpLaunchDispatchUrmaKernel(params) == + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL, + "invalid group width was accepted"); + Check(magicCalls == 2 && launchCalls == 2, + "Host-detectable configuration error consumed an epoch"); + + params.groupWidth = TileXRMoonEp::kDispatchDefaultGroupWidth; + magicReturn = -91; + Check(TileXRMoonEp::TileXRMoonEpLaunchDispatchUrmaKernel(params) == -91, + "magic failure was not propagated"); + Check(magicCalls == 3 && launchCalls == 2, + "magic failure must stop before runtime launch"); + + magicReturn = TileXR::TILEXR_SUCCESS; + launchReturn = -92; + Check(TileXRMoonEp::TileXRMoonEpLaunchDispatchUrmaKernel(params) == + TileXR::TILEXR_ERROR_MKIRT, "runtime launch failure was not mapped"); + Check(magicCalls == 4 && launchCalls == 3, + "runtime launch failure call counts mismatch"); + + params.weightOutput = nullptr; + Check(TileXRMoonEp::TileXRMoonEpLaunchDispatchUrmaKernel(params) == + TileXR::TILEXR_ERROR_PARA_CHECK_FAIL, + "unpaired Weight pointers were accepted"); + Check(magicCalls == 4 && launchCalls == 3, + "pointer validation must happen before the epoch starts"); +} + +} // namespace + +extern "C" int TileXRCommNextMagic(TileXRCommPtr, int64_t *magic) +{ + ++magicCalls; + if (magicReturn == TileXR::TILEXR_SUCCESS && magic != nullptr) { + *magic = nextMagic; + } + return magicReturn; +} + +extern "C" aclError aclrtGetDevice(int32_t *deviceId) +{ + if (deviceId != nullptr) { + *deviceId = 0; + } + return ACL_SUCCESS; +} + +extern "C" aclError aclrtGetDeviceInfo(uint32_t, aclrtDevAttr attr, int64_t *value) +{ + if (value != nullptr) { + *value = attr == ACL_DEV_ATTR_UBUF_PER_VECTOR_CORE ? + 192U * 1024U : TileXRMoonEp::kDispatchAivCoreCount; + } + return ACL_SUCCESS; +} + +extern "C" rtError_t rtDevBinaryRegister(const rtDevBinary_t *, void **handle) +{ + ++binaryRegisterCalls; + if (handle != nullptr) { + *handle = reinterpret_cast(uintptr_t {0xc000}); + } + return RT_ERROR_NONE; +} + +extern "C" rtError_t rtDevBinaryUnRegister(void *) +{ + return RT_ERROR_NONE; +} + +extern "C" rtError_t rtFunctionRegister(void *, const void *, const char_t *, + const void *, uint32_t) +{ + ++functionRegisterCalls; + return RT_ERROR_NONE; +} + +extern "C" rtError_t rtKernelLaunchWithFlagV2(const void *, uint32_t blockDim, + rtArgsEx_t *argsInfo, rtSmDesc_t *, rtStream_t, uint32_t, + const rtTaskCfgInfo_t *cfgInfo) +{ + ++launchCalls; + launchedBlockDim = blockDim; + launchedArgsSize = argsInfo == nullptr ? 0U : argsInfo->argsSize; + launchedLocalMemorySize = cfgInfo == nullptr ? 0U : cfgInfo->localMemorySize; + if (argsInfo != nullptr && argsInfo->args != nullptr && + argsInfo->argsSize == sizeof(TileXRMoonEp::DispatchKernelArgs)) { + launchedArgs = *static_cast( + argsInfo->args); + } + return launchReturn; +} + +int main() +{ + TestPairedAndHiddenOnly(); + TestFailureBoundaries(); + return failures == 0 ? EXIT_SUCCESS : EXIT_FAILURE; +} diff --git a/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp b/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp index a92e071..18bb21a 100644 --- a/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp @@ -52,6 +52,21 @@ void Excludes(const char *label, const std::string &contents, const char *needle } } +void CountEquals(const char *label, const std::string &contents, + const char *needle, size_t expected) +{ + size_t count = 0U; + for (size_t position = contents.find(needle); position != std::string::npos; + position = contents.find(needle, position + 1U)) { + ++count; + } + if (count != expected) { + std::cerr << label << " expected " << expected << " occurrences of " + << needle << ", got " << count << '\n'; + ++failures; + } +} + } // namespace int main() @@ -68,6 +83,10 @@ int main() ReadFile("src/moonep/dispatch/host/dispatch_launch.cpp"); const std::string dispatchKernel = ReadFile("src/moonep/dispatch/kernels/tilexr_moonep_dispatch_kernel.cpp"); + const std::string dispatchUrmaLaunch = + ReadFile("src/moonep/dispatch/urma/host/dispatch_launch.cpp"); + const std::string dispatchUrmaKernel = ReadFile( + "src/moonep/dispatch/urma/kernels/tilexr_moonep_dispatch_kernel.cpp"); const std::string dispatchCmake = ReadFile("src/moonep/dispatch/CMakeLists.txt"); const std::string combineCommon = ReadFile("src/moonep/combine/common/combine_common.h"); @@ -156,6 +175,36 @@ int main() Excludes("dispatch kernel", dispatchKernel, "WaitRankInnerFlag"); Excludes("dispatch kernel", dispatchKernel, "aclrtSynchronizeStream"); Excludes("dispatch kernel", Lower(dispatchKernel), "udma"); + CountEquals("URMA dispatch launcher", dispatchUrmaLaunch, + "TileXRCommNextMagic(", 1U); + CountEquals("URMA dispatch launcher", dispatchUrmaLaunch, + "rtKernelLaunchWithFlagV2(", 1U); + Contains("URMA Dispatch WQE", dispatchUrmaKernel, + "__ubuf__ TileXR::UDMASqeCtx *sqe"); + Contains("URMA Dispatch WQE", dispatchUrmaKernel, + "__ubuf__ TileXR::UDMASgeCtx *sge"); + Contains("URMA Dispatch SQ publish", dispatchUrmaKernel, + "AscendC::DataCopyPad(wqeGlobal, issueLocal, copyParams)"); + Contains("URMA Dispatch SQ publish", dispatchUrmaKernel, + "SyncFunc()"); + Contains("URMA Dispatch doorbell", dispatchUrmaKernel, + "state.qpCtxEntry->dbAddr), 0)"); + Contains("URMA Dispatch fused payload", dispatchUrmaKernel, + "DispatchDataTaskIsWeight("); + Contains("URMA Dispatch fused payload", dispatchUrmaKernel, + "DispatchPayloadWqesPerRoute(hasWeight)"); + Contains("URMA Dispatch grouped convergence", dispatchUrmaKernel, + "const bool payloadReady = upstreamStatus =="); + Contains("URMA Dispatch grouped convergence", dispatchUrmaKernel, + "bool completionSubmitted = optimizedSignalSubmitted"); + Contains("URMA Dispatch grouped convergence", dispatchUrmaKernel, + "WaitDispatchIncomingPeerAndPublishCredit(args"); + Excludes("URMA Dispatch grouped convergence", dispatchUrmaKernel, + "peerValue == rank || upstreamStatus !="); + Excludes("URMA Dispatch WQE", dispatchUrmaKernel, "WRITE_WITH_NOTIFY"); + Excludes("URMA Dispatch WQE", dispatchUrmaKernel, "sgeNum = 2U"); + Excludes("URMA Dispatch WQE", dispatchUrmaKernel, + "__gm__ TileXR::UDMASqeCtx *"); Contains("combine common", combineCommon, "moonep_peer_window.h"); Contains("combine launch", combineLaunch, "#include \"moonep_kernel_launch.h\""); Contains("combine launch", combineLaunch, "LaunchRegisteredMoonEpKernel("); diff --git a/tools/moonep/benchmark.py b/tools/moonep/benchmark.py index 861a59e..af409f4 100644 --- a/tools/moonep/benchmark.py +++ b/tools/moonep/benchmark.py @@ -614,12 +614,21 @@ def _require_exact(name: str, actual: list[int], expected) -> None: def validate_plan( plan, context, cu_seqlens=None, *, expected_status: int = 0, route_distribution: str = DEFAULT_ROUTE_DISTRIBUTION, + all_topk=None, ) -> dict[str, object]: status = int(plan.status.item()) if status != int(expected_status): raise RuntimeError( f"MoonEP device status is {status}, expected {int(expected_status)}" ) + if all_topk is None: + all_topk = deterministic_all_topk( + context.planner_group_size, + context.tokens_per_rank, + context.topk, + context.expert_count, + route_distribution, + ) reference = build_reference_plan( rank=context.planner_group_rank, rank_size=context.planner_group_size, @@ -628,13 +637,7 @@ def validate_plan( expert_count=context.expert_count, prefetch_slots=context.prefetch_slots, token_padding=context.token_padding, - all_topk=deterministic_all_topk( - context.planner_group_size, - context.tokens_per_rank, - context.topk, - context.expert_count, - route_distribution, - ), + all_topk=all_topk, ) _require_exact("dst", _tensor_values(plan.dst), reference.dst) if cu_seqlens is None: diff --git a/tools/moonep/dispatch_hot_loop.py b/tools/moonep/dispatch_hot_loop.py index 770619d..1318cc3 100644 --- a/tools/moonep/dispatch_hot_loop.py +++ b/tools/moonep/dispatch_hot_loop.py @@ -11,7 +11,6 @@ from .config import apply_overrides, build_case_parser, load_cases, select_cases from .planner_reference import ( build_reference_plan, - deterministic_all_topk, deterministic_rank_topk, ) from .rendezvous import completion_barrier_from_env @@ -30,6 +29,7 @@ DIAGNOSTIC_VERSION = 3 KERNEL_STATUS_FEATURE_DFX_ENABLED = 1 << 0 KERNEL_STATUS_FEATURE_PROFILING_ENABLED = 1 << 1 +KERNEL_STATUS_FEATURE_FUSED_EPOCH = 1 << 2 PROFILE_COUNT = 64 DFX_COUNT = 64 COMPLETION_BYTES = 512 * 2 * 8 @@ -219,6 +219,28 @@ def _dtype(torch_module, name: str): return {"bfloat16": torch_module.bfloat16, "float16": torch_module.float16}[name] +def _case_rank_topk(case, rank: int, rank_size: int) -> tuple[int, ...]: + route_count = case.tokens_per_rank * case.topk + if case.routing_pattern == "skewed": + return tuple(route % case.topk for route in range(route_count)) + return deterministic_rank_topk( + rank, + rank_size, + case.tokens_per_rank, + case.topk, + case.expert_count, + case.route_distribution, + ) + + +def _all_case_topk(case, rank_size: int) -> tuple[int, ...]: + return tuple( + expert + for rank in range(rank_size) + for expert in _case_rank_topk(case, rank, rank_size) + ) + + def _inputs(torch_module, case, context): device = f"npu:{context.device_index}" route_count = case.tokens_per_rank * case.topk @@ -226,13 +248,8 @@ def _inputs(torch_module, case, context): route_count, dtype=torch_module.int32, device=device ) topk = torch_module.tensor( - deterministic_rank_topk( - context.planner_group_rank, - context.planner_group_size, - case.tokens_per_rank, - case.topk, - case.expert_count, - case.route_distribution, + _case_rank_topk( + case, context.planner_group_rank, context.planner_group_size ), dtype=torch_module.int32, device=device, @@ -253,10 +270,8 @@ def _inputs(torch_module, case, context): return topk, tpe, hidden.contiguous(), weights.contiguous() -def _expected(torch_module, case, context, destination_roll: int = 0): - all_topk = deterministic_all_topk(context.planner_group_size, - case.tokens_per_rank, case.topk, case.expert_count, - case.route_distribution) +def _reference_slot_assignments(case, context, destination_roll: int = 0): + all_topk = _all_case_topk(case, context.planner_group_size) route_count = case.tokens_per_rank * case.topk destination_capacity = int(context.nv_s) source_by_slot = [-1] * destination_capacity @@ -279,29 +294,68 @@ def _expected(torch_module, case, context, destination_roll: int = 0): raw = encoded target, slot = divmod(raw, destination_capacity) if target == context.planner_group_rank: + if source_by_slot[slot] >= 0: + raise RuntimeError(f"reference destination collision at slot {slot}") source_by_slot[slot] = source token_by_slot[slot] = route // case.topk route_by_slot[slot] = route - if min(source_by_slot) < 0: - raise RuntimeError("reference destinations do not fill every local slot") + return source_by_slot, token_by_slot, route_by_slot + + +def _expected(torch_module, case, context, destination_roll: int = 0): + source_by_slot, token_by_slot, route_by_slot = _reference_slot_assignments( + case, context, destination_roll + ) device = f"npu:{context.device_index}" - source = torch_module.tensor(source_by_slot, dtype=torch_module.int32, + valid = torch_module.tensor( + [source >= 0 for source in source_by_slot], + dtype=torch_module.bool, + device=device, + ) + source = torch_module.tensor([max(0, value) for value in source_by_slot], + dtype=torch_module.int32, device=device).reshape(-1, 1) - token = torch_module.tensor(token_by_slot, dtype=torch_module.int32, + token = torch_module.tensor([max(0, value) for value in token_by_slot], + dtype=torch_module.int32, device=device).reshape(-1, 1) - route = torch_module.tensor(route_by_slot, dtype=torch_module.int32, + route = torch_module.tensor([max(0, value) for value in route_by_slot], + dtype=torch_module.int32, device=device) cols = torch_module.arange(case.hidden_size, dtype=torch_module.int32, device=device).reshape(1, -1) - hidden = ((token * 13 + cols + source * 7) % 64).to( + hidden_values = ((token * 13 + cols + source * 7) % 64).to( dtype=_dtype(torch_module, case.dtype) ) - weights = ((route * 5 + source.reshape(-1) * 11) % 97).to( + weight_values = ((route * 5 + source.reshape(-1) * 11) % 97).to( dtype=torch_module.float32 ) / 128.0 + hidden = torch_module.where( + valid.reshape(-1, 1), hidden_values, torch_module.zeros_like(hidden_values) + ) + weights = torch_module.where( + valid, weight_values, torch_module.zeros_like(weight_values) + ) return hidden, weights, source_by_slot +def _launch_dispatch(runtime, context, plan, hidden, hidden_out, stream, + weights, weights_out, mode: str, workspace, workspace_bytes) -> None: + if mode == "hidden": + runtime.dispatch(context, plan, hidden, hidden_out, stream, + registered_workspace=workspace, + registered_workspace_bytes=workspace_bytes) + elif mode == "weight": + runtime.dispatch(context, plan, weights, weights_out, stream, + registered_workspace=workspace, + registered_workspace_bytes=workspace_bytes) + elif mode == "pair": + runtime.dispatch(context, plan, hidden, hidden_out, stream, weights, + weights_out, registered_workspace=workspace, + registered_workspace_bytes=workspace_bytes) + else: + raise ValueError(f"unsupported Dispatch mode: {mode}") + + def _alternating_plan_check(torch_module, case, buffer, plan, hidden, hidden_out, weights, weights_out, mode: str) -> dict[str, object]: runtime = buffer.runtime @@ -312,20 +366,8 @@ def _alternating_plan_check(torch_module, case, buffer, plan, hidden, destination_roll = case.topk def launch() -> None: - if mode == "hidden": - runtime.dispatch(context, plan, hidden, hidden_out, stream, - registered_workspace=workspace, - registered_workspace_bytes=workspace_bytes) - elif mode == "weight": - runtime.dispatch(context, plan, weights, weights_out, stream, - registered_workspace=workspace, - registered_workspace_bytes=workspace_bytes) - elif mode == "pair": - runtime.dispatch(context, plan, hidden, hidden_out, stream, weights, - weights_out, registered_workspace=workspace, - registered_workspace_bytes=workspace_bytes) - else: - raise ValueError(f"unsupported Dispatch mode: {mode}") + _launch_dispatch(runtime, context, plan, hidden, hidden_out, stream, + weights, weights_out, mode, workspace, workspace_bytes) try: launch() @@ -361,6 +403,52 @@ def launch() -> None: torch_module.npu.synchronize() +def _repeated_exact_check(torch_module, case, buffer, plan, hidden, + hidden_out, weights, weights_out, mode: str, rounds: int) -> dict[str, object]: + if rounds <= 0: + return {"enabled": False, "rounds": 0} + runtime = buffer.runtime + context = buffer.context + workspace, workspace_bytes = context.dispatch_workspace + stream = buffer._stream_ptr() + original_dst = plan.dst.clone() + destination_roll = case.topk + try: + for round_index in range(rounds): + roll = destination_roll if round_index % 2 else 0 + if roll: + plan.dst.copy_(torch_module.roll(original_dst, shifts=roll)) + else: + plan.dst.copy_(original_dst) + torch_module.npu.synchronize() + _launch_dispatch(runtime, context, plan, hidden, hidden_out, stream, + weights, weights_out, mode, workspace, workspace_bytes) + buffer.quiesce() + expected_hidden, expected_weights, _ = _expected( + torch_module, case, context, destination_roll=roll) + hidden_ok = (bool(torch_module.equal(hidden_out, expected_hidden)) + if mode in ("hidden", "pair") else None) + weight_ok = (bool(torch_module.equal(weights_out, expected_weights)) + if mode in ("weight", "pair") else None) + if hidden_ok is False or weight_ok is False: + raise RuntimeError( + "repeated Dispatch mismatch " + f"round={round_index} hidden={hidden_ok} weight={weight_ok} " + f"roll={roll}" + ) + return { + "enabled": True, + "rounds": rounds, + "dispatch_mode": mode, + "alternating_plan": True, + "hidden_exact": mode in ("hidden", "pair"), + "weight_exact": mode in ("weight", "pair"), + } + finally: + plan.dst.copy_(original_dst) + torch_module.npu.synchronize() + + def _workspace_blob(context, byte_offset: int, byte_count: int) -> bytes: raw = context._dispatch_workspace_owner aligned_offset = context._dispatch_workspace_ptr - int(raw.data_ptr()) @@ -458,7 +546,13 @@ def _diagnostics(context, *, require_complete: bool = True, KERNEL_STATUS_FEATURE_DFX_ENABLED) dfx_enabled = bool(features & KERNEL_STATUS_FEATURE_DFX_ENABLED) profiling_enabled = bool(features & KERNEL_STATUS_FEATURE_PROFILING_ENABLED) - result = {"kernel_status": kernel_status} + fused_epoch = bool(features & KERNEL_STATUS_FEATURE_FUSED_EPOCH) + shared_owner_mode = "weight" if fused_epoch else None + result = { + "kernel_status": kernel_status, + "fused_epoch": fused_epoch, + "shared_owner_mode": shared_owner_mode, + } rank = int(context.planner_group_rank) for payload_mode, mode in enumerate(("hidden", "weight")): profile_blob = _workspace_blob(context, offsets[f"{mode}_profile"], profile_bytes) @@ -524,6 +618,7 @@ def _diagnostics(context, *, require_complete: bool = True, result[mode] = { "profile_available": profiling_enabled, "dfx_available": dfx_enabled, + "shared_epoch_owner": fused_epoch and mode == shared_owner_mode, "profile": profiles, "dfx": dfx, } @@ -586,18 +681,8 @@ def _measure(torch_module, case, buffer, plan, hidden, hidden_out, weights, stream = buffer._stream_ptr() def launch(): - if mode == "hidden": - runtime.dispatch(context, plan, hidden, hidden_out, stream, - registered_workspace=workspace, - registered_workspace_bytes=workspace_bytes) - elif mode == "weight": - runtime.dispatch(context, plan, weights, weights_out, stream, - registered_workspace=workspace, - registered_workspace_bytes=workspace_bytes) - else: - runtime.dispatch(context, plan, hidden, hidden_out, stream, weights, - weights_out, registered_workspace=workspace, - registered_workspace_bytes=workspace_bytes) + _launch_dispatch(runtime, context, plan, hidden, hidden_out, stream, + weights, weights_out, mode, workspace, workspace_bytes) for _ in range(case.warmup): launch() @@ -662,7 +747,8 @@ def run_case(torch_module, case, args) -> None: buffer.synchronize() validation = validate_plan( plan, context, cu_seqlens, - route_distribution=case.route_distribution) + route_distribution=case.route_distribution, + all_topk=_all_case_topk(case, context.planner_group_size)) route_count = context.dispatched_capacity hidden_out = torch_module.empty((route_count, case.hidden_size), dtype=hidden.dtype, device=hidden.device) @@ -673,6 +759,10 @@ def run_case(torch_module, case, args) -> None: alternating_mode = "pair" if "pair" in dispatch_modes else dispatch_modes[0] alternating_plan = _alternating_plan_check(torch_module, case, buffer, plan, hidden, hidden_out, weights, weights_out, alternating_mode) + exact_mode = "pair" if "pair" in dispatch_modes else dispatch_modes[0] + repeated_exact = _repeated_exact_check(torch_module, case, buffer, plan, + hidden, hidden_out, weights, weights_out, exact_mode, + args.exact_rounds) modes = {} for mode in dispatch_modes: modes[mode] = _measure(torch_module, case, buffer, plan, hidden, @@ -719,6 +809,7 @@ def run_case(torch_module, case, args) -> None: if validate_weight: validation["dispatch_weight_exact"] = weight_ok validation["alternating_plan"] = alternating_plan + validation["repeated_exact"] = repeated_exact validation["passed"] = True validation["mode"] = "planner_and_dispatch_bit_exact" result["validation"] = validation @@ -773,9 +864,12 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--output-dir", required=True) parser.add_argument("--install-prefix", default=None) parser.add_argument("--wait-iterations", type=int, default=1_000_000) + parser.add_argument("--exact-rounds", type=int, default=0) parser.add_argument("--dispatch-modes", nargs="+", choices=DISPATCH_MODES, default=DISPATCH_MODES) args = parser.parse_args(argv) + if args.exact_rounds < 0: + parser.error("--exact-rounds must be non-negative") if os.environ.get("TILEXR_UDMA_QP_ROUTE_SPEC") != "port_count:6,port_count:2": raise RuntimeError( "TILEXR_UDMA_QP_ROUTE_SPEC=port_count:6,port_count:2 is required before init") diff --git a/tools/moonep/launcher.py b/tools/moonep/launcher.py index 01868fe..618ee20 100644 --- a/tools/moonep/launcher.py +++ b/tools/moonep/launcher.py @@ -136,6 +136,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--comm-id", default=None) parser.add_argument("--python", default=sys.executable) parser.add_argument("--wait-iterations", type=int, default=1_000_000) + parser.add_argument("--exact-rounds", type=int, default=0) parser.add_argument("--timeout-sec", type=float, default=1800.0) parser.add_argument( "--mode", choices=("benchmark", "reference", "correctness"), default="benchmark" @@ -184,6 +185,7 @@ def _process_command(args: argparse.Namespace) -> list[str]: if dispatch_hot_loop: command.append("--dispatch-modes") command.extend(args.dispatch_modes) + command.extend(("--exact-rounds", str(getattr(args, "exact_rounds", 0)))) else: if args.candidate_backend: command.extend(("--candidate-backend", args.candidate_backend)) @@ -199,11 +201,13 @@ def main(argv: list[str] | None = None) -> int: validate_iteration_overrides(args.warmup, args.iterations) if ( args.wait_iterations <= 0 + or args.exact_rounds < 0 or args.timeout_sec <= 0 or args.tensor_preview_elements <= 0 ): raise ValueError( - "wait_iterations, timeout_sec, and tensor_preview_elements must be positive" + "wait_iterations, timeout_sec, and tensor_preview_elements must be positive; " + "exact_rounds must be non-negative" ) if args.mode == "benchmark" and args.dump_stage_tensors: raise ValueError("--dump-stage-tensors is only valid in reference/correctness mode")