perf(moonep): balance prefetch over shared QPs

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.

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.

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.

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<GM_ADDR>(const_cast<int32_t *>(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,