diff --git a/docs/moonep/PREFETCH_WEIGHT_PERFORMANCE.md b/docs/moonep/PREFETCH_WEIGHT_PERFORMANCE.md new file mode 100644 index 0000000..d3345c1 --- /dev/null +++ b/docs/moonep/PREFETCH_WEIGHT_PERFORMANCE.md @@ -0,0 +1,159 @@ +# MoonEP PrefetchWeight Performance + +## Result + +The shared-QP PrefetchWeight implementation was validated on 2026-08-13 against +TileXR main and native MoonEP. The optimized implementation keeps logical slot +ownership unchanged but maps its four workers to physical QPs `{0,1,2,16}`. +This distributes equal-size slot traffic across the six-port and two-port CLOS +groups in the hardware's 3:1 port ratio. + +The test used eight physical Ascend950PR devices on `141.61.49.195`, CANN +`9.1.T560`, driver `25.1.rc1.b188`, TileXR base commit `c39c433`, and native +MoonEP commit `53e03002655d07cfc39e7e9ca2c2aa18583c6c0b`. + +| Implementation | P50 | P99 | Effective bandwidth | +| --- | ---: | ---: | ---: | +| TileXR main `c39c433` | 1272.83 us | 2662.19 us | 276.80 GB/s | +| TileXR shared-QP mapping | 986.99 us | 1044.07 us | 356.97 GB/s | +| Native MoonEP `53e0300` | 7554.13 us | 7810.42 us | 46.64 GB/s | + +Relative to TileXR main, the optimized P50 is 22.46% lower and effective +bandwidth is 28.96% higher. It is 7.65 times faster than native MoonEP by P50 +and reaches 89.24% of the 400 GB/s one-card one-way external-port ceiling. +The independent dual-card validation in +`docs/UDMA_DUAL_CARD_BANDWIDTH_VALIDATION.md` measured about 389 GB/s per card, +so the result is close to the demonstrated transport limit but does not yet +close the remaining data-plane gap. + +## Workload And Timing + +The fixed EP8 case is: + +```text +ranks: 8 physical ranks, one rank per NPU +experts: 32 total, 4 local experts per rank +remote slots: 4 per rank +dtype: BF16 +gate/up row: [7168, 2048], 28 MiB each +down row: [2048, 7168], 28 MiB +bytes per slot: 84 MiB +bytes per rank/round: 336 MiB (352321536 bytes) +samples: 3 x (5 warmup + 20 measured) +``` + +Before timing, every rank pulls the four experts owned by the previous rank and +checks every BF16 value in all twelve destination projection slots exactly. +The timed interval uses events on the current NPU stream. TileXR times one fused +gate/up/down launch; native MoonEP times its three sequential `launch_prefetch` +calls. Allocation, MR or SHMEM registration, correctness validation, the +pre-iteration distributed barrier, and post-launch status reads are excluded. + +For every iteration the report selects the maximum event time across all eight +ranks. P50 and P99 use linear interpolation over those cross-rank maxima. +Effective bandwidth is `352321536 / P50_us / 1000` GB/s. + +## Sample Quality + +The TileXR event API intermittently returned `0.044 us`, which is below any +possible 336 MiB transfer time. The benchmark preserves every raw event sample, +but classifies readings at or below 1 us as invalid and excludes only those +readings from P50, P99, and bandwidth: + +| Implementation | Raw samples | Valid samples | Invalid `0.044 us` samples | +| --- | ---: | ---: | ---: | +| TileXR main | 60 | 55 | 5 | +| TileXR shared-QP mapping | 60 | 52 | 8 | + +The optimized repeat P50 values were 975.78, 990.03, and 985.84 us. Main +contained one 4143.04 us system spike; no high sample was filtered, so its +aggregate P99 is 2662.19 us. Main's three repeat P99 values were 1384.15, +3662.42, and 1310.24 us. The optimized repeat P99 values were 1021.40, +1052.77, and 1017.52 us. P50 and bandwidth are the primary throughput comparison; +larger runs are required before treating either P99 as a production tail claim. + +## Retained Artifacts + +The complete rank JSON files, aggregate JSON, and launch logs remain on the +test host: + +```text +/tmp/TileXR-prefetch-opt-20260813-c39c433/ + artifacts/prefetch-final-main-20260813/ + artifacts/prefetch-final-opt-20260813/ + prefetch-final-main-20260813.log + prefetch-final-opt-20260813.log + +/tmp/TileXR-prefetch-baseline-20260813-c39c433/ + artifacts/prefetch-formal-native-dev/ + native-formal-dev.log + hccl-aiv-only-8r.log +``` + +Each artifact directory contains `rank_0.json` through `rank_7.json` plus +`summary.json`. The TileXR artifact labels identify the tested main and optimized +source snapshots; both are based on `c39c433`. The native aggregate records the +full native commit. The matching official HCCL AIV-only environment baseline +passed before TileXR performance investigation. + +## Implementation Boundaries + +The optimization separates logical workers from physical QPs. Logical worker +`w` still owns slots `w`, `w + workerCount`, and so on. Only physical queue +selection changes: + +| Shared-domain workers | Physical QPs | +| ---: | --- | +| 1 | `0` | +| 2 | `0,1` | +| 4 | `0,1,2,16` | +| 8 | `0,1,2,3,4,5,16,17` | + +Non-shared domains retain identity mapping. The selected QPs are packed in a +private 64-bit Kernel argument. Public MoonEP structures and Python APIs are +unchanged. Registered memory, peer-memory behavior, slot assignment, WQE +construction in UB, MTE3 SQ publication, `st_dev` doorbells, and CQ completion +semantics are unchanged. There is no second PrefetchWeight implementation path. + +## Experiments Not To Repeat + +The following experiments did not improve the large-transfer bottleneck: + +| Experiment | Result | Conclusion | +| --- | --- | --- | +| `blockDim=1/2/4` before QP remapping | 270.75 / 274.42 / 275.30 GB/s | More QPs inside the same six-port CLOS do not add useful bandwidth. | +| Reverse `TILEXR_UDMA_QP_ROUTE_SPEC` | No material change | A shared communicator constructs the fixed 32-QP profile, so this variable does not remap that profile. | + +The following setup failures were environmental or invocation mistakes, not +PrefetchWeight defects: + +- The remote `torchrun` shebang referenced a removed Conda environment. Use + `python -m torch.distributed.run`. +- `TILEXR_BUILD_EP=ON` does not build MoonEP. Use + `TILEXR_BUILD_MOONEP=ON`. +- Streaming a Windows-produced tar archive into Linux caused archive-format + problems. The successful workflow used a remote main snapshot and copied + changed files individually. + +## Next Stage + +1. Profile the remaining 32.46 GB/s gap to the independently measured + approximately 389 GB/s card transport rate. Use msprof and physical-port + counters to separate queue/WQE issue limits, the fused Kernel's scalar work, + and link utilization. Do not tune small Host overhead before this is known. +2. Increase the sample count and determine why NPU events sometimes report + `0.044 us`. Keep raw samples and correctness checks; do not silently discard + high values or quote production P99 until the event anomaly is understood. +3. Sweep large payload and slot-count distributions around the production + shape. Confirm that the 3:1 mapping remains optimal when workers own unequal + byte counts; derive a byte-aware mapping only if evidence shows imbalance. +4. Measure the transfer-size crossover against the unchanged peer-memory path. + Memory remains preferable for small transfers, while registered UDMA should + remain the large-transfer path. Do not select either transport from topology + alone. +5. After the data plane is saturated, measure fixed launch/status overhead and + consider reducing it only if it becomes a material fraction of stage time. + +These measurements prove the registered-memory UDMA data plane on this +Ascend950 topology. They do not establish UDMA performance on 910B, a simulator, +another CLOS layout, an oversubscribed rank topology, or cross-node MoonEP. diff --git a/docs/plans/2026-08-13-moonep-prefetch-weight-shared-qp.md b/docs/plans/2026-08-13-moonep-prefetch-weight-shared-qp.md new file mode 100644 index 0000000..3f2165d --- /dev/null +++ b/docs/plans/2026-08-13-moonep-prefetch-weight-shared-qp.md @@ -0,0 +1,107 @@ +# MoonEP PrefetchWeight Shared-QP Plan + +## Goal And Scope + +Implement the approved design in +`docs/specs/2026-08-13-moonep-prefetch-weight-shared-qp-design.md`, validate it +on CANN 9.1 and eight Ascend950 devices, and deliver a patch based on current +`origin/main`. + +Only PrefetchWeight layout, private launch ABI, Kernel QP selection, focused +tests, benchmark tooling, and performance documentation are in scope. Public +ABI, registration, peer-memory, and other MoonEP stages are non-goals. + +## Task 1: Implement And Test The Mapping Contract + +**Objective and role:** Separate logical worker assignment from physical QP +selection without changing slot ownership. + +**Background and prerequisites:** Use the fixed shared-domain profile from +`src/comm/udma/tilexr_udma_config.*` and the approved mappings in the design. + +**Modification scope:** + +- `src/moonep/prefetch_weight/host/prefetch_weight_layout.*` +- `tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp` + +**Constraints and non-goals:** Preserve supported worker counts and the current +block-dimension override. Non-shared QPs use identity mapping. + +**Acceptance and verification:** The Host unit test proves four-worker shared +mapping `{0,1,2,16}`, eight-worker shared mapping +`{0,1,2,3,4,5,16,17}`, identity fallback, and invalid-input behavior. + +**Artifacts and interfaces:** A packed private QP map stored in +`PrefetchWeightLayout` for Task 2. + +## Task 2: Propagate And Consume The Physical QP + +**Objective and role:** Pass the map through the registered direct-launch ABI +and use it for all PrefetchWeight WQ/CQ operations. + +**Background and prerequisites:** Depends on Task 1. The logical worker remains +the slot scheduler. + +**Modification scope:** + +- `src/moonep/prefetch_weight/host/prefetch_weight_launch.cpp` +- `src/moonep/prefetch_weight/kernels/tilexr_moonep_prefetch_weight_kernel.cpp` +- focused launch/source tests under `tests/moonep/unit/` + +**Constraints and non-goals:** Preserve WQE-in-UB, MTE3 publication, `st_dev` +doorbells, CQ accounting, status behavior, and public API. Do not introduce a +second implementation path. + +**Acceptance and verification:** Focused CTest targets pass and source checks +show mapped QPs are used for queue lookup, submit, and quiet. A target CANN 9.1 +build compiles and embeds the Kernel. + +**Artifacts and interfaces:** Installed PrefetchWeight library and integration +package used by Task 3. + +## Task 3: Run Ascend950 Correctness And Performance A/B + +**Objective and role:** Decide from hardware evidence whether the mapping is a +real improvement. + +**Background and prerequisites:** Depends on a passing Task 2 build. Reuse the +healthy HCCL AIV-only environment baseline and the existing EP8 benchmark case +on `141.61.49.195`. + +**Modification scope:** Remote files under a new `/tmp` directory and retained +benchmark artifacts. Do not modify system CANN or driver installations. + +**Constraints and non-goals:** Use eight physical ranks, BF16 exact slot +validation, NPU events, cross-rank maxima, and the same timing exclusions as +the baseline. Run a short diagnostic first; retain the production mapping only +if correctness passes and the large-transfer P50 improves materially. + +**Acceptance and verification:** Run `3 x (5 + 20)` for the optimized TileXR +build. Compare against TileXR main and native MoonEP with P50, P99, effective +GB/s, per-repeat medians, and correctness status. + +**Artifacts and interfaces:** Raw per-rank JSON, aggregate JSON, build and run +logs, and result paths for Task 4. + +## Task 4: Record Results And Deliver The Patch + +**Objective and role:** Make the optimization reproducible without the current +conversation and produce a patch based on latest main. + +**Background and prerequisites:** Depends on Tasks 1-3 and their retained raw +artifacts. + +**Modification scope:** A stable benchmark under `tools/moonep/`, focused +MoonEP documentation, this plan/spec, and the final Git diff. + +**Constraints and non-goals:** Record ineffective experiments and failed test +launches that would otherwise be repeated. Keep temporary remote paths out of +runtime code. + +**Acceptance and verification:** Documentation names commits, hardware, +dimensions, timing boundary, baseline, optimized results, validation limits, +and next-stage recommendations. Relevant tests pass from a clean build and +`git diff --check` is clean. Generate a patch against the fetched latest main. + +**Artifacts and interfaces:** A scoped commit and `.patch` file; no push or PR +unless separately requested. diff --git a/docs/specs/2026-08-13-moonep-prefetch-weight-shared-qp-design.md b/docs/specs/2026-08-13-moonep-prefetch-weight-shared-qp-design.md new file mode 100644 index 0000000..2fe417f --- /dev/null +++ b/docs/specs/2026-08-13-moonep-prefetch-weight-shared-qp-design.md @@ -0,0 +1,86 @@ +# MoonEP PrefetchWeight Shared-QP Design + +## Goal + +Increase PrefetchWeight throughput on Ascend950 shared UDMA domains by using +both external-port CLOS groups. Preserve the current fused three-projection +launch, registered-memory transport, public MoonEP ABI, and peer-memory paths. + +The target workload is EP8 with four remote expert slots per rank. Each BF16 +slot contains 28 MiB gate, 28 MiB up, and 28 MiB down rows, for 336 MiB per +rank per launch. + +## Evidence And Result + +On `141.61.49.195`, the final back-to-back run measured TileXR `c39c433` at +1272.83 us P50 and 276.80 GB/s. Native MoonEP `53e0300` achieved 7554.13 us +P50 and 7810.42 us P99, or 46.64 GB/s, for the same data and explicit plan. + +PrefetchWeight currently uses the logical worker index as the physical QP. +The fixed shared-domain profile assigns QPs 0-15 to the six-port CLOS and QPs +16-31 to the two-port CLOS. A block-dimension sweep measured 270.75, 274.42, +and 275.30 GB/s with one, two, and four workers. More QPs within the six-port +CLOS therefore provide little additional throughput, while the two-port CLOS +is unused. + +Changing `TILEXR_UDMA_QP_ROUTE_SPEC` did not alter the baseline because a +shared-QP communicator always constructs the fixed 32-QP profile. This was a +diagnostic experiment, not a supported tuning mechanism for this path. + +The implemented four-worker map `{0,1,2,16}` measured 986.99 us P50, +1044.07 us P99, and 356.97 GB/s. P50 improved by 22.46% and bandwidth by +28.96% over main, and the optimized path was 7.65 times faster than native +MoonEP by P50. Exact BF16 slot validation passed on all eight physical ranks. +`docs/moonep/PREFETCH_WEIGHT_PERFORMANCE.md` records the full measurement +method, raw artifact paths, event anomalies, ineffective experiments, and +next-stage work. + +## Design + +Keep logical work assignment unchanged: worker `w` owns slots `w`, +`w + workerCount`, and so on. Add a private logical-worker-to-physical-QP map +to `PrefetchWeightLayout` and the direct-launch Kernel ABI. + +For a fixed 32-QP shared domain: + +| Workers | Physical QPs | CLOS traffic ratio | +| ---: | --- | ---: | +| 1 | `0` | 1:0 | +| 2 | `0,1` | 2:0 | +| 4 | `0,1,2,16` | 3:1 | +| 8 | `0,1,2,3,4,5,16,17` | 6:2 | + +All non-shared domains retain identity mapping. Host validation requires every +selected QP to be below the transport QP count. The map is packed as eight +8-bit indices in one `uint64_t`; this keeps the private launch block compact +and covers the transport maximum of 32 QPs. + +The Kernel continues to use the logical worker for slot partitioning and uses +the mapped physical QP only for WQ lookup, UDMA GET submission, and CQ wait. +There is no change to registered regions, transfer sizes, WQE construction, +doorbell ordering, or completion semantics. + +## Compatibility And Non-Goals + +- Preserve C++14, CANN 9.1, and the registered direct-Kernel launch path. +- Do not change `TileXRMoonEpPrefetchWeightArgsV1` or Python-facing APIs. +- Do not change UDMA registration or the fixed shared-QP profile. +- Do not change peer-memory, Dispatch, Combine, or ReduceGrad. +- Do not optimize small fixed Host overhead until the large-transfer path is + measured after using both CLOS groups. + +## Verification + +1. Host tests cover identity mapping, 32-QP four/eight-worker mappings, and + packed-map propagation into the launch context. +2. Source/launch tests cover the private ABI and require all UDMA queue + operations to use the mapped physical QP. +3. A complete CANN 9.1 build proves Host and embedded Kernel ABI consistency. +4. Ascend950 EP8 tests validate exact BF16 slot contents and run the same + `3 x (5 warmup + 20 measured)` NPU-event benchmark as the baseline. +5. Retain the native MoonEP baseline and report P50, P99, effective bandwidth, + repeat-level results, and any ineffective experiments. + +The original performance hypothesis was 350-390 GB/s, corresponding to roughly +900-1000 us P50. The measured 356.97 GB/s and 986.99 us P50 satisfy that gate, +so the shared-QP mapping is retained. diff --git a/src/moonep/prefetch_weight/host/prefetch_weight_launch.cpp b/src/moonep/prefetch_weight/host/prefetch_weight_launch.cpp index 91be0ee..84ac47f 100644 --- a/src/moonep/prefetch_weight/host/prefetch_weight_launch.cpp +++ b/src/moonep/prefetch_weight/host/prefetch_weight_launch.cpp @@ -38,6 +38,7 @@ int TileXRMoonEpLaunchPrefetchWeightKernel( int64_t expertsPerRank; int64_t prefetchSlots; uint64_t qpNum; + uint64_t physicalQpMap; } args { context.devArgs, reinterpret_cast(const_cast(params.expertsToCopy)), @@ -50,10 +51,10 @@ int TileXRMoonEpLaunchPrefetchWeightKernel( context.layout.down.rowBytes, context.layout.rank, context.layout.rankSize, context.layout.expertsPerRank, context.layout.prefetchSlots, - context.layout.qpNum + context.layout.qpNum, context.layout.physicalQpMap }; - static_assert(sizeof(PrefetchWeightKernelArgs) == 17U * sizeof(uint64_t), + static_assert(sizeof(PrefetchWeightKernelArgs) == 18U * sizeof(uint64_t), "PrefetchWeight kernel argument ABI changed"); return LaunchRegisteredMoonEpKernel(g_prefetchWeightRegistration, diff --git a/src/moonep/prefetch_weight/host/prefetch_weight_layout.cpp b/src/moonep/prefetch_weight/host/prefetch_weight_layout.cpp index 542d64e..8d0e5ed 100644 --- a/src/moonep/prefetch_weight/host/prefetch_weight_layout.cpp +++ b/src/moonep/prefetch_weight/host/prefetch_weight_layout.cpp @@ -101,6 +101,33 @@ bool RangesOverlap(uint64_t lhsBegin, uint64_t lhsEnd, } // namespace +int TileXRMoonEpBuildPrefetchWeightQpMap(uint32_t workers, uint32_t qpNum, + bool sharedQps, uint64_t *physicalQpMap) +{ + if (physicalQpMap == nullptr || !SupportedWorkerCount(workers) || + workers > qpNum || + (sharedQps && qpNum != kPrefetchWeightSharedQpCount)) { + return TILEXR_MOONEP_ERROR_INVALID_ARGUMENT; + } + + uint64_t packed = 0; + for (uint32_t worker = 0; worker < workers; ++worker) { + uint32_t physicalQp = worker; + if (sharedQps && workers == 4U && worker == 3U) { + physicalQp = kPrefetchWeightSecondClosQpBase; + } else if (sharedQps && workers == kPrefetchWeightMaxWorkers && + worker >= 6U) { + physicalQp = kPrefetchWeightSecondClosQpBase + worker - 6U; + } + if (physicalQp >= qpNum || physicalQp > UINT8_MAX) { + return TILEXR_MOONEP_ERROR_INVALID_ARGUMENT; + } + packed |= static_cast(physicalQp) << (worker * 8U); + } + *physicalQpMap = packed; + return TILEXR_MOONEP_SUCCESS; +} + int TileXRMoonEpBuildPrefetchWeightLayout( const TileXRMoonEpPrefetchWeightArgsV1 &args, const TileXR::CommArgs &commArgs, @@ -158,6 +185,13 @@ int TileXRMoonEpBuildPrefetchWeightLayout( while (workers > static_cast(args.plan->b)) { workers >>= 1; } + const bool sharedQps = + (commArgs.extraFlag & TileXR::ExtraFlag::UDMA_SHARED_QP) != 0U; + if (TileXRMoonEpBuildPrefetchWeightQpMap( + workers, qpNum, sharedQps, &next.physicalQpMap) != + TILEXR_MOONEP_SUCCESS) { + return TILEXR_MOONEP_ERROR_INVALID_ARGUMENT; + } next.expertsPerRank = expertsPerRank; next.prefetchSlots = args.plan->b; next.rank = commArgs.rank; diff --git a/src/moonep/prefetch_weight/host/prefetch_weight_layout.h b/src/moonep/prefetch_weight/host/prefetch_weight_layout.h index 29d6c01..b047363 100644 --- a/src/moonep/prefetch_weight/host/prefetch_weight_layout.h +++ b/src/moonep/prefetch_weight/host/prefetch_weight_layout.h @@ -10,6 +10,8 @@ namespace TileXRMoonEp { constexpr uint32_t kPrefetchWeightMaxWorkers = 8; constexpr uint32_t kPrefetchWeightAlignment = 64; +constexpr uint32_t kPrefetchWeightSharedQpCount = 32; +constexpr uint32_t kPrefetchWeightSecondClosQpBase = 16; struct PrefetchWeightProjectionLayout { GM_ADDR localBase = nullptr; @@ -27,8 +29,20 @@ struct PrefetchWeightLayout { int32_t rankSize = 0; uint32_t qpNum = 0; uint32_t blockDim = 0; + uint64_t physicalQpMap = 0; }; +int TileXRMoonEpBuildPrefetchWeightQpMap(uint32_t workers, uint32_t qpNum, + bool sharedQps, uint64_t *physicalQpMap); + +inline uint32_t TileXRMoonEpPrefetchWeightPhysicalQp( + uint64_t physicalQpMap, uint32_t worker) +{ + return worker < kPrefetchWeightMaxWorkers ? + static_cast((physicalQpMap >> (worker * 8U)) & UINT64_C(0xFF)) : + UINT32_MAX; +} + int TileXRMoonEpBuildPrefetchWeightLayout( const TileXRMoonEpPrefetchWeightArgsV1 &args, const TileXR::CommArgs &commArgs, diff --git a/src/moonep/prefetch_weight/kernels/tilexr_moonep_prefetch_weight_kernel.cpp b/src/moonep/prefetch_weight/kernels/tilexr_moonep_prefetch_weight_kernel.cpp index b7fd7f5..4538c1d 100644 --- a/src/moonep/prefetch_weight/kernels/tilexr_moonep_prefetch_weight_kernel.cpp +++ b/src/moonep/prefetch_weight/kernels/tilexr_moonep_prefetch_weight_kernel.cpp @@ -19,7 +19,7 @@ class PrefetchWeightKernel { uint64_t gateOffset, uint64_t upOffset, uint64_t downOffset, uint32_t gateRowBytes, uint32_t upRowBytes, uint32_t downRowBytes, int32_t rank, int32_t rankSize, int64_t expertsPerRank, - int64_t prefetchSlots, uint32_t qpNum) + int64_t prefetchSlots, uint32_t qpNum, uint64_t physicalQpMap) { args_ = reinterpret_cast<__gm__ TileXR::CommArgs *>(commArgs); expertsToCopy_ = reinterpret_cast<__gm__ int32_t *>(expertsToCopy); @@ -42,6 +42,8 @@ class PrefetchWeightKernel { worker_ = static_cast(get_block_idx()) * subBlockCount + static_cast(get_subblockid()); workerCount_ = static_cast(get_block_num()) * subBlockCount; + physicalQp_ = static_cast( + (physicalQpMap >> (worker_ * 8U)) & UINT64_C(0xFF)); pipe_.InitBuffer(wqeBuf_, TileXR::TILEXR_UDMA_WQE_SCRATCH_BYTES); } @@ -100,7 +102,8 @@ class PrefetchWeightKernel { prefetchSlots_ > expertsPerRank_ || rank_ < 0 || rank_ >= rankSize_ || rankSize_ > static_cast(kMaxTrackedRankSize) || - workerCount_ == 0 || worker_ >= workerCount_ || worker_ >= qpNum_ || + workerCount_ == 0 || worker_ >= workerCount_ || + physicalQp_ >= qpNum_ || !TileXR::UDMARegistryEnabled(args_) || args_->rank != rank_ || args_->rankSize != rankSize_) { return kPrefetchWeightStatusInvalidRuntime; @@ -156,7 +159,8 @@ class PrefetchWeightKernel { expert % static_cast(expertsPerRank_); MarkPeer(usedPeers, owner); __gm__ TileXR::UDMAWQCtx *queue = TileXR::UDMAGetWQCtx( - TileXR::GetUDMAInfo(args_), static_cast(owner), worker_); + TileXR::GetUDMAInfo(args_), static_cast(owner), + physicalQp_); const uint32_t completionQueue = TrackCompletionQueue( queue->wqeCntAddr, owner, completionQueueIds, completionQueuePeers, completionTargets, completionQueueCount, workerStatus); @@ -170,8 +174,8 @@ class PrefetchWeightKernel { static_cast(expertsPerRank_ + slot) * rowBytes_[projection]; const uint32_t submitStatus = TileXR::UDMAGetNbiOnQp( - args_, wqeScratch, owner, worker_, destination, sourceOffset, - rowBytes_[projection]); + args_, wqeScratch, owner, physicalQp_, destination, + sourceOffset, rowBytes_[projection]); if (submitStatus != TileXR::TILEXR_UDMA_STATUS_SUCCESS && workerStatus == 0) { workerStatus = kPrefetchWeightStatusSubmitErrorBase + @@ -224,7 +228,7 @@ class PrefetchWeightKernel { continue; } const uint32_t cqStatus = TileXR::UDMAQuietStatusOnQpUntil( - args_, peer, worker_, completionTargets[queue]); + args_, peer, physicalQp_, completionTargets[queue]); if (cqStatus != 0 && workerStatus == 0) { workerStatus = kPrefetchWeightStatusCqErrorBase + (cqStatus & 0xFFU); } @@ -251,6 +255,7 @@ class PrefetchWeightKernel { uint32_t qpNum_ = 0; uint32_t worker_ = 0; uint32_t workerCount_ = 0; + uint32_t physicalQp_ = 0; AscendC::TPipe pipe_; AscendC::TBuf wqeBuf_; }; @@ -263,13 +268,14 @@ extern "C" __global__ __aicore__ void tilexr_moonep_prefetch_weight_kernel( GM_ADDR status, uint64_t gateOffset, uint64_t upOffset, uint64_t downOffset, uint64_t gateRowBytes, uint64_t upRowBytes, uint64_t downRowBytes, int64_t rank, int64_t rankSize, int64_t expertsPerRank, - int64_t prefetchSlots, uint64_t qpNum) + int64_t prefetchSlots, uint64_t qpNum, uint64_t physicalQpMap) { TileXRMoonEp::Kernel::PrefetchWeightKernel op; op.Init(commArgs, expertsToCopy, gate, up, down, status, gateOffset, upOffset, downOffset, static_cast(gateRowBytes), static_cast(upRowBytes), static_cast(downRowBytes), static_cast(rank), static_cast(rankSize), - expertsPerRank, prefetchSlots, static_cast(qpNum)); + expertsPerRank, prefetchSlots, static_cast(qpNum), + physicalQpMap); op.Process(); } diff --git a/tests/moonep/python/test_prefetch_weight_benchmark.py b/tests/moonep/python/test_prefetch_weight_benchmark.py new file mode 100644 index 0000000..b0e8d03 --- /dev/null +++ b/tests/moonep/python/test_prefetch_weight_benchmark.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import pytest + +from tools.moonep.prefetch_weight_benchmark import ( + MIN_VALID_EVENT_US, + percentile, + ring_plan, + statistics, +) + + +def test_ring_plan_pulls_remote_slots() -> None: + assert ring_plan(4, 2, 2) == [ + [6, 7], + [0, 1], + [2, 3], + [4, 5], + ] + + +def test_statistics_use_cross_rank_max_and_retain_invalid_events() -> None: + result = statistics( + ( + (0.044, 8.0, 6.0), + (0.044, 5.0, 7.0), + ) + ) + + assert result["cross_rank_max_us"] == [0.044, 8.0, 7.0] + assert result["valid_cross_rank_max_us"] == [8.0, 7.0] + assert result["invalid_event_sample_count"] == 1 + assert result["minimum_valid_event_us_exclusive"] == MIN_VALID_EVENT_US + assert result["p50_us"] == 7.5 + assert result["p99_us"] == pytest.approx(7.99) + + +def test_statistics_reject_mismatched_or_entirely_invalid_samples() -> None: + with pytest.raises(ValueError, match="same number"): + statistics(((2.0,), (2.0, 3.0))) + with pytest.raises(ValueError, match="all cross-rank NPU event samples"): + statistics(((0.044,), (0.044,))) + + +def test_percentile_rejects_empty_input() -> None: + with pytest.raises(ValueError, match="at least one"): + percentile([], 0.5) diff --git a/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp b/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp index 9f01905..a92e071 100644 --- a/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_kernel_sources.cpp @@ -207,6 +207,10 @@ int main() Contains("prefetch launch", prefetchLaunch, "LaunchRegisteredMoonEpKernel("); Excludes("prefetch launch", prefetchLaunch, "rtKernelLaunchWithFlagV2"); Contains("prefetch launch", prefetchLaunch, "kPrefetchWeightKernelSignature"); + Contains("prefetch launch", prefetchLaunch, + "context.layout.qpNum, context.layout.physicalQpMap"); + Contains("prefetch launch", prefetchLaunch, + "18U * sizeof(uint64_t)"); Contains("prefetch kernel", prefetchKernel, "extern \"C\" __global__ __aicore__ void tilexr_moonep_prefetch_weight_kernel"); Contains("prefetch kernel", prefetchKernel, "expertsToCopy"); @@ -218,6 +222,11 @@ int main() Excludes("prefetch kernel", prefetchKernel, "e_ + slot"); Contains("prefetch kernel", prefetchKernel, "UDMAGetNbiOnQp"); Contains("prefetch kernel", prefetchKernel, "UDMAQuietStatusOnQpUntil"); + Contains("prefetch kernel", prefetchKernel, "physicalQp_"); + Contains("prefetch kernel", prefetchKernel, + "owner, physicalQp_, destination"); + Contains("prefetch kernel", prefetchKernel, + "peer, physicalQp_, completionTargets[queue]"); Contains("prefetch kernel", prefetchKernel, "completionQueueIds"); Contains("prefetch kernel", prefetchKernel, "++completionTargets[completionQueue]"); diff --git a/tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp b/tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp index 1f748e7..c0115ba 100644 --- a/tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp +++ b/tests/moonep/unit/test_tilexr_moonep_prefetch_weight_host.cpp @@ -96,11 +96,13 @@ void TestLaunch() seenContext.layout.up.registryOffset == 0x1000 && seenContext.layout.down.registryOffset == 0x2000 && seenContext.layout.qpNum == 4 && seenContext.layout.blockDim == 4 && + seenContext.layout.physicalQpMap == UINT64_C(0x03020100) && seenContext.layout.expertsPerRank == 4, "prefetch UDMA layout mismatch"); Reset(); qpNum = 32; + commArgs.extraFlag |= TileXR::ExtraFlag::UDMA_SHARED_QP; plan = Plan(); gate = Weight(0x100000, 4, 8); up = Weight(0x101000, 4, 16); @@ -110,9 +112,31 @@ void TestLaunch() TileXRMoonEp::TileXRMoonEpRunPrefetchWeightV1(&args, stream), TILEXR_MOONEP_SUCCESS); Check(launchCalls == 1 && seenContext.layout.qpNum == 32 && - seenContext.layout.blockDim == 4, + seenContext.layout.blockDim == 4 && + TileXRMoonEp::TileXRMoonEpPrefetchWeightPhysicalQp( + seenContext.layout.physicalQpMap, 0) == 0 && + TileXRMoonEp::TileXRMoonEpPrefetchWeightPhysicalQp( + seenContext.layout.physicalQpMap, 1) == 1 && + TileXRMoonEp::TileXRMoonEpPrefetchWeightPhysicalQp( + seenContext.layout.physicalQpMap, 2) == 2 && + TileXRMoonEp::TileXRMoonEpPrefetchWeightPhysicalQp( + seenContext.layout.physicalQpMap, 3) == 16, "prefetch must cap workers without rejecting the shared-domain QP count"); + uint64_t physicalQpMap = 0; + Status("prefetch eight-worker shared QP map", + TileXRMoonEp::TileXRMoonEpBuildPrefetchWeightQpMap( + 8, 32, true, &physicalQpMap), TILEXR_MOONEP_SUCCESS); + const uint32_t expectedSharedQps[8] = {0, 1, 2, 3, 4, 5, 16, 17}; + for (uint32_t worker = 0; worker < 8; ++worker) { + Check(TileXRMoonEp::TileXRMoonEpPrefetchWeightPhysicalQp( + physicalQpMap, worker) == expectedSharedQps[worker], + "prefetch eight-worker shared QP mapping mismatch"); + } + Status("prefetch malformed shared QP domain", + TileXRMoonEp::TileXRMoonEpBuildPrefetchWeightQpMap( + 4, 4, true, &physicalQpMap), TILEXR_MOONEP_ERROR_INVALID_ARGUMENT); + Reset(); qpNum = 3; plan = Plan(); gate = Weight(0x100000, 4, 8); diff --git a/tools/moonep/README.md b/tools/moonep/README.md index e8e52c4..afa4023 100644 --- a/tools/moonep/README.md +++ b/tools/moonep/README.md @@ -74,6 +74,56 @@ Run the CPU/fake-backend suite without NPU hardware: python -m pytest tests/moonep/python -q ``` +## Standalone PrefetchWeight Benchmark + +`prefetch_weight_benchmark.py` compares only the PrefetchWeight stage. It uses +eight physical ranks by default in the validated EP8 case: 32 experts, four +experts and four remote slots per rank, BF16 `H=7168`, and `Hf=2048`. Each slot +copies three 28 MiB projections, so one rank transfers 336 MiB per measured +iteration. + +Build and install the MoonEP targets, then run TileXR with the Python module +launcher. The `torchrun` executable can contain a stale Conda shebang on shared +hosts, so it is intentionally not used here. + +```bash +source scripts/common_env.sh +export TILEXR_INSTALL_PREFIX="$PWD/install" +export PYTHONPATH="$PWD/integrations/moonep_torch:$PWD:${PYTHONPATH:-}" +export TILEXR_COMM_ID=127.0.0.1:10067 + +python -m torch.distributed.run --standalone --nproc-per-node=8 \ + tools/moonep/prefetch_weight_benchmark.py \ + --backend tilexr \ + --install-prefix "$TILEXR_INSTALL_PREFIX" \ + --output-dir output/prefetch-weight-tilexr \ + --warmup 5 --iterations 20 --repeats 3 \ + --tilexr-commit "$(git rev-parse HEAD)" +``` + +Run the native MoonEP comparison with the same case and iteration counts: + +```bash +python -m torch.distributed.run --standalone --nproc-per-node=8 \ + tools/moonep/prefetch_weight_benchmark.py \ + --backend native \ + --native-root /path/to/ascend-moonep \ + --output-dir output/prefetch-weight-native \ + --warmup 5 --iterations 20 --repeats 3 \ + --native-commit "$(git -C /path/to/ascend-moonep rev-parse HEAD)" +``` + +Both backends initialize and validate exact BF16 slot contents before timing. +Each sample uses NPU events around only the stage launch, and rank 0 aggregates +the slowest rank for every iteration. Allocation, memory registration, barriers, +correctness checks, and status reads are excluded. Raw per-rank samples and the +aggregate are retained in `rank_.json` and `summary.json`. Event readings +at or below 1 us remain in `cross_rank_max_us` for diagnosis but are excluded +from `valid_cross_rank_max_us`, P50, P99, and effective bandwidth. + +The validated Ascend950 results and known experimental dead ends are recorded in +`docs/moonep/PREFETCH_WEIGHT_PERFORMANCE.md`. + ## Public API NPU E2E Tool `test_npu_e2e.py` is a standalone base test tool for the upstream-compatible diff --git a/tools/moonep/prefetch_weight_benchmark.py b/tools/moonep/prefetch_weight_benchmark.py new file mode 100644 index 0000000..c7c6454 --- /dev/null +++ b/tools/moonep/prefetch_weight_benchmark.py @@ -0,0 +1,548 @@ +from __future__ import annotations + +import argparse +import json +import math +import os +import platform +import socket +import subprocess +import sys +import time +from pathlib import Path +from types import SimpleNamespace + + +PROJECTION_NAMES = ("gate", "up", "down") +TILEXR_SUCCESS_STATUS = 4000 +MIN_VALID_EVENT_US = 1.0 + + +def projection_shapes(args): + return ( + (args.hidden, args.projection), + (args.hidden, args.projection), + (args.projection, args.hidden), + ) + + +def percentile(values, quantile): + if not values: + raise ValueError("percentile requires at least one value") + ordered = sorted(float(value) for value in values) + position = (len(ordered) - 1) * quantile + lower = int(math.floor(position)) + upper = int(math.ceil(position)) + if lower == upper: + return ordered[lower] + fraction = position - lower + return ordered[lower] + (ordered[upper] - ordered[lower]) * fraction + + +def statistics(samples_by_rank): + if not samples_by_rank or not samples_by_rank[0]: + raise ValueError("cross-rank statistics require samples") + iterations = len(samples_by_rank[0]) + if any(len(samples) != iterations for samples in samples_by_rank): + raise ValueError("all ranks must report the same number of samples") + maxima = [ + max(float(rank_samples[index]) for rank_samples in samples_by_rank) + for index in range(iterations) + ] + valid = [value for value in maxima if value > MIN_VALID_EVENT_US] + if not valid: + raise ValueError("all cross-rank NPU event samples are invalid") + return { + "cross_rank_max_us": maxima, + "valid_cross_rank_max_us": valid, + "invalid_event_sample_count": len(maxima) - len(valid), + "minimum_valid_event_us_exclusive": MIN_VALID_EVENT_US, + "p50_us": percentile(valid, 0.50), + "p99_us": percentile(valid, 0.99), + "min_us": min(valid), + "max_us": max(valid), + "mean_us": sum(valid) / len(valid), + } + + +def constant_value(owner, local_expert, projection): + return float(owner * 16 + local_expert + projection * 0.25) + + +def ring_plan(ranks, experts_per_rank, slots): + if slots > experts_per_rank: + raise ValueError("slots must not exceed experts_per_rank") + return [ + [((rank - 1) % ranks) * experts_per_rank + slot for slot in range(slots)] + for rank in range(ranks) + ] + + +def current_stream_ptr(torch_module): + stream = torch_module.npu.current_stream() + value = getattr(stream, "npu_stream", getattr(stream, "stream", None)) + if value is None: + raise RuntimeError("current NPU stream exposes no native pointer") + return int(value) + + +def command_output(command, cwd=None): + try: + completed = subprocess.run( + command, cwd=cwd, check=False, capture_output=True, text=True, timeout=10 + ) + except (OSError, subprocess.SubprocessError): + return None + return (completed.stdout.strip() or completed.stderr.strip()) or None + + +class TileXRRunner: + def __init__(self, torch_module, args, rank, ranks, device, plan_host): + from tilexr_moonep import MoonEPPlan, ProjectionBuffers, TileXRMoonEPRuntime + + self.torch = torch_module + self.rank = rank + self.ranks = ranks + self.device = device + self.args = args + self.shapes = projection_shapes(args) + os.environ["TILEXR_UDMA_QP_ROUTE_SPEC"] = args.qp_route_spec + if args.num_sms is not None: + os.environ["TILEXR_MOONEP_PREFETCH_BLOCK_DIM"] = str(args.num_sms) + self.runtime = TileXRMoonEPRuntime( + rank, ranks, install_prefix=args.install_prefix + ) + self.context = SimpleNamespace( + planner_group_size=ranks, + expert_count=ranks * args.experts_per_rank, + prefetch_slots=args.slots, + nv_s=1, + topk=1, + ) + plan_tensor = torch_module.tensor( + plan_host, dtype=torch_module.int32, device=device + ) + self.plan = MoonEPPlan( + dst=torch_module.zeros((1,), dtype=torch_module.int32, device=device), + experts_to_copy=plan_tensor, + zero_fill_ranges=torch_module.zeros( + (self.context.expert_count + args.slots, 2), + dtype=torch_module.int32, + device=device, + ), + remote_stats=torch_module.tensor( + [args.slots, args.slots], dtype=torch_module.int32, device=device + ), + dup_groups=torch_module.zeros( + (1, 3), dtype=torch_module.int32, device=device + ), + dup_loffs=torch_module.zeros((1,), dtype=torch_module.int32, device=device), + dup_counts=torch_module.zeros((2,), dtype=torch_module.int32, device=device), + status=torch_module.zeros((1,), dtype=torch_module.int32, device=device), + reduce_grad_status=torch_module.zeros( + (1,), dtype=torch_module.int32, device=device + ), + workspace=torch_module.empty((1,), dtype=torch_module.uint8, device=device), + n=1, + tokens_per_rank=1, + topk=1, + expert_count=self.context.expert_count, + rank_size=ranks, + prefetch_slots=args.slots, + nv_s=1, + token_padding=1, + epoch=1, + backend="tilexr", + runtime=self.runtime, + ) + local_weights = [] + for projection, shape in enumerate(self.shapes): + tensor = torch_module.empty( + (args.experts_per_rank, *shape), + dtype=torch_module.bfloat16, + device=device, + ) + for expert in range(args.experts_per_rank): + tensor[expert].fill_(constant_value(rank, expert, projection)) + local_weights.append(tensor) + self.projections = ProjectionBuffers.from_local_weights( + SimpleNamespace( + experts_per_rank=args.experts_per_rank, + dtype=torch_module.bfloat16, + device_index=int(os.environ["LOCAL_RANK"]), + ), + *local_weights, + slot_fill_value=-7.0, + torch_module=torch_module, + ) + del local_weights + torch_module.npu.synchronize() + self.handle = self.runtime.udma_register(self.projections.backing) + self.stream_ptr = current_stream_ptr(torch_module) + + def launch(self): + self.runtime.prefetch_weight( + self.context, self.plan, self.projections, self.stream_ptr + ) + + def synchronize(self): + self.torch.npu.synchronize() + actual = int(self.plan.status.item()) + if actual != TILEXR_SUCCESS_STATUS: + raise RuntimeError( + f"TileXR PrefetchWeight status {actual}, expected {TILEXR_SUCCESS_STATUS}" + ) + + def validate(self, plan_host): + owner = (self.rank - 1) % self.ranks + checks = {} + for projection, name in enumerate(PROJECTION_NAMES): + tensor = getattr(self.projections, name) + for slot, expert in enumerate(plan_host[self.rank]): + local_expert = int(expert) % self.args.experts_per_rank + expected = constant_value(owner, local_expert, projection) + actual = tensor[self.args.experts_per_rank + slot] + passed = bool(self.torch.all(actual == expected).item()) + checks[f"{name}_{slot}"] = passed + if not passed: + raise RuntimeError( + f"TileXR {name} slot {slot} differs from expert {expert}" + ) + return checks + + def layout(self): + return { + "launch_count": 1, + "transport": "registered-memory UDMA GET", + "qp_count": self.runtime.udma_qp_count, + "block_dim_override": self.args.num_sms, + } + + def close(self): + self.torch.npu.synchronize() + self.runtime.udma_unregister(self.handle) + self.runtime.close() + + +class NativeRunner: + def __init__(self, torch_module, dist, args, rank, ranks, device, plan_host): + native_root = Path(args.native_root).resolve() + sys.path.insert(0, str(native_root)) + from ascend_moonep import ShmemRuntime, launch_prefetch + from ascend_moonep.buffer_c import ( + create_sym_tensor_from_ptr, + create_vmm_physical, + map_to_sym_ptr, + reset_symmetric_descriptors, + ) + + self.torch = torch_module + self.dist = dist + self.args = args + self.shapes = projection_shapes(args) + self.rank = rank + self.ranks = ranks + self.device = device + self.launch_prefetch = launch_prefetch + self.ShmemRuntime = ShmemRuntime + self.reset_symmetric_descriptors = reset_symmetric_descriptors + self.sources = [] + for projection, shape in enumerate(self.shapes): + source = create_vmm_physical( + (args.experts_per_rank + args.slots, *shape), + torch_module.bfloat16, + int(os.environ["LOCAL_RANK"]), + )[0] + for expert in range(args.experts_per_rank): + source[expert].fill_(constant_value(rank, expert, projection)) + source[args.experts_per_rank :].fill_(-7.0) + self.sources.append(source) + torch_module.npu.synchronize() + dist.barrier() + ShmemRuntime.init_with_buffer(group=None) + self.plan = torch_module.tensor( + plan_host[rank], dtype=torch_module.int32, device=device + ).contiguous() + self.remote_experts = [] + self.slots = [] + for projection, (source, shape) in enumerate( + zip(self.sources, self.shapes) + ): + sym_base = map_to_sym_ptr(source.data_ptr(), f"projection_{projection}") + row_bytes = math.prod(shape) * 2 + self.remote_experts.append( + create_sym_tensor_from_ptr( + sym_base, + (ranks * args.experts_per_rank, *shape), + torch_module.bfloat16, + device, + ) + ) + self.slots.append( + create_sym_tensor_from_ptr( + sym_base + args.experts_per_rank * row_bytes, + (args.slots, *shape), + torch_module.bfloat16, + device, + ) + ) + + def launch(self): + num_sms = self.args.num_sms or 32 + for remote_expert, slots in zip(self.remote_experts, self.slots): + self.launch_prefetch(remote_expert, slots, self.plan, num_sms=num_sms) + + def synchronize(self): + self.torch.npu.synchronize() + + def validate(self, plan_host): + owner = (self.rank - 1) % self.ranks + checks = {} + for projection, (name, tensor) in enumerate( + zip(PROJECTION_NAMES, self.slots) + ): + for slot, expert in enumerate(plan_host[self.rank]): + local_expert = int(expert) % self.args.experts_per_rank + expected = constant_value(owner, local_expert, projection) + passed = bool(self.torch.all(tensor[slot] == expected).item()) + checks[f"{name}_{slot}"] = passed + if not passed: + raise RuntimeError( + f"native {name} slot {slot} differs from expert {expert}" + ) + return checks + + def layout(self): + return { + "launch_count": 3, + "transport": "CANN SHMEM MTE remote-to-UB-to-GM GET", + "block_dim_requested": self.args.num_sms or 32, + } + + def close(self): + self.torch.npu.synchronize() + self.dist.barrier() + if self.ShmemRuntime.is_initialized(): + self.ShmemRuntime.finalize() + for source in self.sources: + allocation = getattr(source, "_vmm_allocation", None) + if allocation is not None: + allocation.destroy() + self.reset_symmetric_descriptors() + + +def time_once(torch_module, dist, runner): + torch_module.npu.synchronize() + dist.barrier() + start = torch_module.npu.Event(enable_timing=True) + end = torch_module.npu.Event(enable_timing=True) + start.record() + runner.launch() + end.record() + end.synchronize() + elapsed_us = float(start.elapsed_time(end)) * 1000.0 + runner.synchronize() + return elapsed_us + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--backend", choices=("tilexr", "native"), required=True) + parser.add_argument("--install-prefix") + parser.add_argument("--native-root") + parser.add_argument("--output-dir", required=True) + parser.add_argument("--experts-per-rank", type=int, default=4) + parser.add_argument("--slots", type=int, default=4) + parser.add_argument("--hidden", type=int, default=7168) + parser.add_argument("--projection", type=int, default=2048) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iterations", type=int, default=20) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--num-sms", type=int) + parser.add_argument( + "--qp-route-spec", default="port_count:6,port_count:2" + ) + parser.add_argument("--tilexr-commit", default="unknown") + parser.add_argument("--tilexr-source-sha256", default="unknown") + parser.add_argument("--native-commit", default="unknown") + return parser.parse_args() + + +def main(): + args = parse_args() + if min( + args.experts_per_rank, + args.slots, + args.hidden, + args.projection, + args.iterations, + args.repeats, + ) <= 0 or args.warmup < 0: + raise ValueError("invalid dimensions or iteration count") + if args.backend == "native" and not args.native_root: + raise ValueError("--native-root is required for the native backend") + + import torch + import torch.distributed as dist + import torch_npu + + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + ranks = int(os.environ["WORLD_SIZE"]) + torch.npu.set_device(local_rank) + coordination_backend = "gloo" if args.backend == "tilexr" else "hccl" + init_kwargs = {"backend": coordination_backend} + if coordination_backend == "hccl": + init_kwargs["device_id"] = torch.device(f"npu:{local_rank}") + dist.init_process_group(**init_kwargs) + device = f"npu:{local_rank}" + plan_host = ring_plan(ranks, args.experts_per_rank, args.slots) + runner = None + result = { + "schema_version": 1, + "status": "failed", + "backend": args.backend, + "rank": rank, + "world_size": ranks, + "coordination_backend": coordination_backend, + "case": { + "experts_per_rank": args.experts_per_rank, + "expert_count": ranks * args.experts_per_rank, + "slots": args.slots, + "projection_shapes": projection_shapes(args), + "dtype": "bfloat16", + "warmup": args.warmup, + "iterations": args.iterations, + "repeats": args.repeats, + }, + "plan": plan_host, + "timing_boundary": { + "clock": "NPU events on the current stream", + "tilexr": "one fused gate/up/down PrefetchWeight launch", + "native": "three sequential launch_prefetch calls", + "excluded": [ + "allocation", + "MR or SHMEM registration", + "correctness validation", + "pre-iteration cross-rank barrier", + "post-launch status reads", + ], + }, + } + failure = None + try: + setup_start = time.perf_counter() + if args.backend == "tilexr": + runner = TileXRRunner(torch, args, rank, ranks, device, plan_host) + else: + runner = NativeRunner(torch, dist, args, rank, ranks, device, plan_host) + result["setup_ms"] = (time.perf_counter() - setup_start) * 1000.0 + result["layout"] = runner.layout() + dist.barrier() + runner.launch() + runner.synchronize() + result["correctness"] = { + "passed": True, + "checks": runner.validate(plan_host), + } + all_repeats = [] + gathered_repeats = [] + for repeat in range(args.repeats): + for _ in range(args.warmup): + time_once(torch, dist, runner) + local_samples = [ + time_once(torch, dist, runner) for _ in range(args.iterations) + ] + gathered = [None for _ in range(ranks)] + dist.all_gather_object(gathered, local_samples) + all_repeats.append(local_samples) + if rank == 0: + gathered_repeats.append(gathered) + result["local_samples_us"] = all_repeats + result["status"] = "passed" + if rank == 0: + projection_row_bytes = [ + math.prod(shape) * 2 for shape in projection_shapes(args) + ] + bytes_per_rank = args.slots * sum(projection_row_bytes) + per_repeat = [statistics(samples) for samples in gathered_repeats] + all_cross_rank = [ + value + for repeat in per_repeat + for value in repeat["cross_rank_max_us"] + ] + all_valid_cross_rank = [ + value + for repeat in per_repeat + for value in repeat["valid_cross_rank_max_us"] + ] + aggregate = { + "cross_rank_max_us": all_cross_rank, + "valid_cross_rank_max_us": all_valid_cross_rank, + "invalid_event_sample_count": + len(all_cross_rank) - len(all_valid_cross_rank), + "minimum_valid_event_us_exclusive": MIN_VALID_EVENT_US, + "p50_us": percentile(all_valid_cross_rank, 0.50), + "p99_us": percentile(all_valid_cross_rank, 0.99), + "min_us": min(all_valid_cross_rank), + "max_us": max(all_valid_cross_rank), + "mean_us": sum(all_valid_cross_rank) / len(all_valid_cross_rank), + } + aggregate["effective_GBps_at_p50"] = ( + bytes_per_rank / aggregate["p50_us"] / 1000.0 + ) + result["bytes"] = { + "per_projection_row": projection_row_bytes, + "per_slot": sum(projection_row_bytes), + "per_rank_per_iteration": bytes_per_rank, + } + result["samples_by_rank_us"] = gathered_repeats + result["statistics_by_repeat"] = per_repeat + result["statistics"] = aggregate + result["commits"] = { + "tilexr": args.tilexr_commit, + "tilexr_source_snapshot_sha256": args.tilexr_source_sha256, + "native_moonep": args.native_commit, + } + result["environment"] = { + "hostname": socket.gethostname(), + "platform": platform.platform(), + "python": sys.version, + "torch": str(torch.__version__), + "torch_npu": str(torch_npu.__version__), + "cann_home": os.environ.get("ASCEND_HOME_PATH"), + "soc": str(torch.npu.get_device_name()), + "npu_smi": command_output(("npu-smi", "info")), + } + except Exception as exc: + result["failure_reason"] = f"{type(exc).__name__}: {exc}" + failure = (exc, exc.__traceback__) + finally: + try: + if runner is not None: + runner.close() + dist.barrier() + except Exception as cleanup_error: + result["cleanup_error"] = ( + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + if failure is None: + failure = (cleanup_error, cleanup_error.__traceback__) + result["status"] = "failed" + output_dir = Path(args.output_dir).resolve() + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / f"rank_{rank}.json").write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + if rank == 0: + (output_dir / "summary.json").write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + dist.destroy_process_group() + if failure is not None: + raise failure[0].with_traceback(failure[1]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())