Add owner-sharded activation objects and direct publication

[foraxe]

Draft PR: Add owner-sharded activation objects and direct publication

Summary

This change adds opt-in object-centric activation paths to MoonEP while preserving ordinary dispatch as the default and universal fallback.

Baseline

producer local output
    |
    v
rank-to-rank payload dispatch
    |
    v
consumer-local routed materialization
    |
    v
consumer

This PR

producer
    |
    +--> direct-out to final owner placement --> publish --> consumer
    |
    `--> one-copy owner-sharded source --> peer-indexed view
                                          |
                                          `--> optional local materialization

• prepare_dispatch exposes a planned final-owner target for a compatible producer epilogue.
• prepare_dispatch_publication overlaps route metadata, padding, and duplicate-map preparation with that producer.
• publish_dispatch publishes producer writes and expands same-rank duplicates without recopying representative hidden payloads.
• source_storage and prepare_source_dispatch expose one physical source-activation shard per rank.
• materialize_source_dispatch is the contiguous-consumer fallback: each rank explicitly pulls representative peer rows and fans the loaded vector directly into repeated final slots.

No existing call changes behavior unless the new APIs or enable_source_object=True are selected.

Motivation

The ordinary path materializes routed payloads through a rank-to-rank dispatch even when the producer or consumer can operate on a shared logical activation object. The new contracts move placement, peer addressing, publication, and duplicate handling behind MoonEP APIs:

producer natural output
-> final owner placement or one-copy source object
-> visibility publication
-> consumer execution view

Physical movement remains available when a contiguous local view is faster. The runtime does not require generic CUDA operations to dereference peer VMM mappings.

Public contracts

Producer direct-out

prepare_dispatch returns a MoonEPDispatchTarget whose plan.dst maps each source token/top-k position to a representative final-owner row. A compatible producer writes non-negative destinations to target.mapped_hidden. publish_dispatch completes MoonEP-owned metadata, cross-rank visibility, padding, and duplicate expansion.

The optional two-phase path starts metadata work with prepare_dispatch_publication, runs the producer on the caller stream, then consumes the one-shot prepared object in publish_dispatch.

Owner-sharded source

Buffer(enable_source_object=True) allocates one [S, H] source shard and one [S, K] route-weight shard per rank. prepare_source_dispatch publishes peer-indexed provenance without routed hidden copies. materialize_source_dispatch pulls only representative rows from their physical owners and produces the ordinary contiguous expert layout.

Performance evidence

Hardware: 4 x NVIDIA GB200. Shape unless noted: S=8192, E=896, routing skew 1.0. Timing uses CUDA graphs, 10 warmups, and 100 measured iterations.

Owner-pull materialization

Each row is the mean of five fresh processes.

H	K	Skew	Ordinary ready view	Owner-pull	Absolute saving	Improvement	Candidate SD
3584	8	0.1	633.86 us	285.50 us	348.36 us	54.96%	2.02 us
3584	8	1.0	626.64 us	286.02 us	340.62 us	54.36%	1.16 us
3584	8	5.0	603.78 us	362.96 us	240.82 us	39.89%	5.69 us
3584	16	0.1	967.70 us	338.72 us	628.98 us	65.00%	0.16 us
3584	16	1.0	964.76 us	338.34 us	626.42 us	64.93%	0.40 us
3584	16	5.0	950.50 us	404.12 us	546.38 us	57.48%	0.67 us
7168	8	0.1	799.42 us	559.52 us	239.90 us	30.01%	6.65 us
7168	8	1.0	790.80 us	554.22 us	236.58 us	29.92%	3.34 us
7168	8	5.0	761.72 us	722.88 us	38.84 us	5.10%	5.92 us
7168	16	0.1	1240.40 us	667.34 us	573.06 us	46.20%	0.34 us
7168	16	1.0	1227.20 us	669.74 us	557.46 us	45.43%	2.57 us
7168	16	5.0	1196.40 us	784.88 us	411.52 us	34.40%	0.58 us

Producer direct-out with overlapped publication

The benchmark uses a synthetic tensor-core producer with equivalent local-output and direct-out epilogues. It measures the producer-to-ready-dispatch boundary, not a model step. Each row is the mean of five fresh processes.

H	K	Skew	Local producer + dispatch	Overlapped direct-out	Absolute saving	Improvement	Candidate SD
3584	8	0.1	588.26 us	327.96 us	260.30 us	44.25%	0.22 us
3584	8	1.0	586.13 us	327.04 us	259.09 us	44.20%	0.78 us
3584	8	5.0	572.82 us	335.13 us	237.70 us	41.50%	0.31 us
3584	16	0.1	807.57 us	533.08 us	274.49 us	33.99%	0.14 us
3584	16	1.0	802.60 us	531.94 us	270.66 us	33.72%	0.15 us
3584	16	5.0	786.99 us	528.81 us	258.18 us	32.81%	0.76 us
7168	8	0.1	863.49 us	612.54 us	250.95 us	29.06%	0.90 us
7168	8	1.0	860.31 us	611.89 us	248.42 us	28.88%	0.49 us
7168	8	5.0	849.06 us	633.65 us	215.41 us	25.37%	0.15 us
7168	16	0.1	1120.93 us	1011.13 us	109.79 us	9.79%	0.57 us
7168	16	1.0	1109.75 us	1012.44 us	97.31 us	8.77%	0.71 us
7168	16	5.0	1086.70 us	1000.49 us	86.21 us	7.93%	0.45 us

The non-overlapped direct-out path regressed by 37.04-42.02 us for all three H=7168, K=16 cells. Those cells require the prepared-publication overlap; ordinary dispatch remains the fallback when the framework cannot schedule it.

Completed coverage

The final EP4 gate H={3584,7168} x K={8,16} x skew={0.1,1,5} completed with five fresh processes per point for both paths. This proves the explicit benchmarked boundaries, not universal performance outside the measured shapes or full-model latency.

Representative commands are:

torchrun --nproc_per_node=4 --module benchmarks.bench_comm \
  --ep 4 --enable-source-object --warmup 10 --iters 100

torchrun --nproc_per_node=4 --module benchmarks.bench_producer_direct_out \
  --ep 4 --seq-len 8192 --experts 896 --hidden 3584 --topk 8 \
  --unbalance-ratio 1.0 --warmup 10 --iters 100 \
  --out results/direct_out.csv

The submitter retained all raw logs, CSVs, the failed initial harness-order attempt, and the audited summary.

Correctness validation

Completed on four NVIDIA GB200 ranks:

• CPU object-contract and planning checks: 5 passed.
• Test collection: 77 object-contract, planning, dispatch, combine, e2e, prefetch, and gradient-reduction cases.
• Distributed kernel/API suite: 75 passed and 1 intentional skip per rank.
• Isolated public end-to-end suite: 1 passed per rank.
• All 60 producer direct-out benchmark runs verified the public two-phase result bit-for-bit against ordinary dispatch.
• CUDA/C++ extension rebuild for SM100: passed.
• Extension import and both benchmark module entry points: passed.

Safety and compatibility

• All new storage is opt-in; enable_source_object=False preserves the original allocation size.
• Ordinary dispatch remains the fallback for unsupported producers, unsupported consumers, and unmeasured selector cells.
• Targets, views, and prepared publications are bound to one Buffer and generation.
• Publication is one-shot; stale or mismatched objects fail before launch.
• Peer mappings carry explicit capability strings. Generic tensor consumers are not assumed peer-safe.
• cuda.tile is used only by a benchmark producer and is not imported by the runtime package.

Community coordination

No open MoonEP issue or PR currently implements owner-sharded activation objects or producer direct publication. PRs #10-#13 remain open with no reported checks. This branch does not absorb their broad fixes: #11 overlaps public API validation, #12 overlaps plan provenance, and #13 overlaps VMM/C++ error handling. Rebase or stack after maintainers choose their order, and preserve the memory-semantic diff as a separate feature review.

AI assistance

AI assistance was used for implementation, test design, benchmark orchestration, analysis, and drafting. The human submitter must review and understand every changed line and reproduce the reported validation before publication.