# dots-infra/ultraep architecture

> Architecture diagram of the GitHub repository dots-infra/ultraep, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/dots-infra/ultraep
- Repository: https://github.com/dots-infra/ultraep

## Overview

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: [`ultra_ep/manager.py`](https://github.com/dots-infra/ultraep/blob/main/ultra_ep/manager.py) (Python Manager), [`csrc/ultra_ep.cpp`](https://github.com/dots-infra/ultraep/blob/main/csrc/ultra_ep.cpp) (Native Manager), [`csrc/kernels/grad_reduce.cu`](https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/grad_reduce.cu) (Replica gradient reduction), [`csrc/kernels/placement.cu`](https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/placement.cu) (Placement planning)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_integration["MoE integration"]
  node_routing["Post-gating routing"]
  node_dispatch["Token dispatch"]
  node_experts["Expert compute"]
end

subgraph group_control["Placement control"]
  node_manager["Python Manager<br/>[manager.py]"]
  node_native["Native Manager<br/>[ultra_ep.cpp]"]
end

subgraph group_kernels["GPU operations"]
  node_placement["Placement planning<br/>[placement.cu]"]
  node_reroute["Routing expansion<br/>[reroute.cu]"]
  node_weight_sync["Replica weight sync<br/>[weight_sync.cu]"]
  node_grad_reduce["Replica gradient reduction<br/>[grad_reduce.cu]"]
  node_load_compute["Load computation<br/>[load_compute.cu]"]
end

subgraph group_runtime["Distributed runtime"]
  node_runtime["CUDA/NVSHMEM runtime<br/>[runtime.cpp]"]
  node_ipc["Interprocess buffers<br/>[ipc_manager.cu]"]
end

subgraph group_observability["Load observability"]
  node_profiler["Load profiler<br/>[profiling.py]"]
  node_viewer["Trace viewer<br/>[load_viewer.py]"]
end

node_framework(("MoE framework"))
node_nvshmem(("NVSHMEM"))
node_trace_files[("Trace chunks")]

node_framework -->|"initializes"| node_manager
node_framework -->|"produces"| node_routing
node_framework -->|"registers expert pointers"| node_manager
node_framework -->|"passes routing"| node_manager
node_manager -->|"invokes runtime"| node_native
node_manager -->|"updates placement"| node_placement
node_manager -->|"syncs weights"| node_weight_sync
node_manager -->|"reroutes tokens"| node_reroute
node_manager -->|"reduces gradients"| node_grad_reduce
node_manager -->|"profiles loads"| node_load_compute
node_placement -->|"placement plan"| node_native
node_weight_sync -->|"replica weights"| node_experts
node_reroute -->|"physical routing"| node_dispatch
node_dispatch -->|"dispatches tokens"| node_experts
node_experts -->|"replica gradients"| node_grad_reduce
node_grad_reduce -->|"master gradients"| node_framework
node_native -->|"initializes"| node_runtime
node_runtime -->|"uses"| node_nvshmem
node_native -->|"manages buffers"| node_ipc
node_profiler -->|"writes traces"| node_trace_files
node_viewer -->|"reads traces"| node_trace_files

click node_manager "https://github.com/dots-infra/ultraep/blob/main/ultra_ep/manager.py"
click node_native "https://github.com/dots-infra/ultraep/blob/main/csrc/ultra_ep.cpp"
click node_placement "https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/placement.cu"
click node_reroute "https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/reroute.cu"
click node_weight_sync "https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/weight_sync.cu"
click node_grad_reduce "https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/grad_reduce.cu"
click node_load_compute "https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/load_compute.cu"
click node_runtime "https://github.com/dots-infra/ultraep/blob/main/csrc/runtime.cpp"
click node_ipc "https://github.com/dots-infra/ultraep/blob/main/csrc/utils/ipc_manager.cu"
click node_profiler "https://github.com/dots-infra/ultraep/blob/main/ultra_ep/profiling.py"
click node_viewer "https://github.com/dots-infra/ultraep/blob/main/ultra_ep/load_viewer.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_routing,node_dispatch,node_experts toneBlue
class node_manager,node_native,node_trace_files toneAmber
class node_placement,node_reroute,node_weight_sync,node_grad_reduce,node_load_compute toneMint
class node_runtime,node_ipc toneRose
class node_profiler,node_viewer,node_framework,node_nvshmem toneIndigo
```

## Components

### MoE integration

- **Post-gating routing**
- **Token dispatch**
- **Expert compute**

### Placement control

- [Python Manager](https://github.com/dots-infra/ultraep/blob/main/ultra_ep/manager.py): `ultra_ep/manager.py`
- [Native Manager](https://github.com/dots-infra/ultraep/blob/main/csrc/ultra_ep.cpp): `csrc/ultra_ep.cpp`

### GPU operations

- [Placement planning](https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/placement.cu): `csrc/kernels/placement.cu`
- [Routing expansion](https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/reroute.cu): `csrc/kernels/reroute.cu`
- [Replica weight sync](https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/weight_sync.cu): `csrc/kernels/weight_sync.cu`
- [Replica gradient reduction](https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/grad_reduce.cu): `csrc/kernels/grad_reduce.cu`
- [Load computation](https://github.com/dots-infra/ultraep/blob/main/csrc/kernels/load_compute.cu): `csrc/kernels/load_compute.cu`

### Distributed runtime

- [CUDA/NVSHMEM runtime](https://github.com/dots-infra/ultraep/blob/main/csrc/runtime.cpp): `csrc/runtime.cpp`
- [Interprocess buffers](https://github.com/dots-infra/ultraep/blob/main/csrc/utils/ipc_manager.cu): `csrc/utils/ipc_manager.cu`

### Load observability

- [Load profiler](https://github.com/dots-infra/ultraep/blob/main/ultra_ep/profiling.py): `ultra_ep/profiling.py`
- [Trace viewer](https://github.com/dots-infra/ultraep/blob/main/ultra_ep/load_viewer.py): `ultra_ep/load_viewer.py`

### Other components

- **MoE framework**
- **NVSHMEM**
- **Trace chunks**

## Connections

- 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