How dots-infra/ultraep works
UltraEP is a Python/CUDA library that rebalances MoE expert load using each layer’s post-gating routing: the framework supplies routing, UltraEP plans replicas, synchronizes weights, and expands routing before expert execution, then restores weights and reduces replica gradients during backward. The README documents this workflow and broader training compatibility; sampled sources confirm the Python/native manager boundary, NVSHMEM runtime, and profiling tools. Some kernel wiring is documented rather than visible in the excerpts, so uncertain internal calls are omitted.
Main entry points: manager.py Python Manager · ultra_ep.cpp Native Manager · grad_reduce.cu Replica gradient reduction · placement.cu Placement planning
Components 17
MoE integration
- Post-gating routing
- Token dispatch
- Expert compute
Placement control
- Python Manager
ultra_ep/manager.py - Native Manager
csrc/ultra_ep.cpp
GPU operations
- Placement planning
csrc/kernels/placement.cu - Routing expansion
csrc/kernels/reroute.cu - Replica weight sync
csrc/kernels/weight_sync.cu - Replica gradient reduction
csrc/kernels/grad_reduce.cu - Load computation
csrc/kernels/load_compute.cu
Distributed runtime
- CUDA/NVSHMEM runtime
csrc/runtime.cpp - Interprocess buffers
csrc/utils/ipc_manager.cu
Load observability
- Load profiler
ultra_ep/profiling.py - Trace viewer
ultra_ep/load_viewer.py
Other components
- MoE framework
- NVSHMEM
- Trace chunks
Connections 21
- MoE framework → Python Manager: initializes
- MoE framework → Post-gating routing: produces
- MoE framework → Python Manager: registers expert pointers
- MoE framework → Python Manager: passes routing
- Python Manager → Native Manager: invokes runtime
- Python Manager → Placement planning: updates placement
- Python Manager → Replica weight sync: syncs weights
- Python Manager → Routing expansion: reroutes tokens
- Python Manager → Replica gradient reduction: reduces gradients
- Python Manager → Load computation: profiles loads
- Placement planning → Native Manager: placement plan
- Replica weight sync → Expert compute: replica weights
- Routing expansion → Token dispatch: physical routing
- Token dispatch → Expert compute: dispatches tokens
- Expert compute → Replica gradient reduction: replica gradients
- Replica gradient reduction → MoE framework: master gradients
- Native Manager → CUDA/NVSHMEM runtime: initializes
- CUDA/NVSHMEM runtime → NVSHMEM: uses
- Native Manager → Interprocess buffers: manages buffers
- Load profiler → Trace chunks: writes traces
- Trace viewer → Trace chunks: reads traces