Technical Summary: UBEP – Re-architecting Expert Parallelism Communication Library for Production Superpods

1. Problem Statement

The deployment of Mixture-of-Experts (MoE) models on production-grade "superpods" (e.g., NVIDIA NVL72/576, Huawei CloudMatrix384) faces critical performance bottlenecks that raw interconnect bandwidth alone cannot resolve. While these systems offer unified global address spaces (UGAS) and high-bandwidth fabrics, existing Expert Parallelism Communication Libraries (EPCLs) fail to fully utilize these capabilities due to three fundamental issues:

1. Strict BSP Serialization: Current libraries rely on the Bulk Synchronous Parallel (BSP) model, which enforces coarse-grained global barriers between interdependent communication phases (e.g., routing, token dispatch, reordering). In low-latency superpods, this "stop-and-wait" behavior serializes execution, preventing the overlap of independent phases and leaving high-bandwidth interconnects underutilized.
2. Synchronization Tax: As link bandwidth scales, the relative cost of synchronization primitives (flags, barriers, kernel launches) becomes dominant. Traditional control-data decoupled mechanisms introduce non-negligible overhead that limits scalability in ultra-low-latency environments.
3. Topology-Agnostic Scheduling: Existing EPCLs treat the superpod fabric as a flat, uniform space, scheduling workloads based solely on token counts. They ignore the physical latency non-uniformity inherent in multi-tier switching fabrics (e.g., one-hop vs. two-hop access), leading to severe load imbalance and stragglers that degrade tail latency.

2. Methodology

The authors introduce UBEP (Unified-Bus Expert Parallelism), a production-ready communication library designed specifically for modern multi-tier superpods. UBEP rethinks the All-to-All primitive through three core architectural innovations:

A. Kernel Decomposition (Dependency-Driven Execution)

Departing from the monolithic BSP model, UBEP decomposes the All-to-All primitive into fine-grained tasks scheduled based on data availability rather than global barriers.

• Mechanism: The communication workflow is partitioned so that distinct groups of AI Cores (AIVs) execute independent tasks in parallel. For instance, while most AIVs handle token transmission, a separate subset calculates and sends token counts (metadata).
• Synchronization: Global barriers are replaced with lightweight point-to-point synchronization. AIVs responsible for metadata (e.g., prefix sums) write results to global shared memory, and other AIVs poll for these updates to proceed immediately, enabling aggressive overlap of metadata exchange, token dispatch, and reordering.

B. Hierarchical Token-Level Scheduling

To address physical latency non-uniformity, UBEP employs a scheduler that jointly optimizes token-to-core mapping based on both load balance and physical fabric distance.

• Problem Formulation: The scheduling problem is modeled to minimize the maximum completion time across AIVs, considering transmission latency and network Round-Trip Time (RTT) variance caused by multi-hop traversals.
• Heuristic Solution: Recognizing the NP-hard nature of the problem, UBEP uses a "latency homogenization" heuristic. It equalizes the distribution of hop-counts (e.g., one-hop vs. two-hop accesses) across AIVs.
• Implementation: A hardware-accelerated mapper uses logical matrix transposition and a quaternary search algorithm (Algorithm 1) to assign tokens. This ensures each AIV handles a balanced mix of local and remote experts, minimizing stragglers caused by the high latency of two-hop accesses (which can be 11.5× slower than local access).

C. Data-as-Flag Synchronization

To eliminate the "synchronization tax" of explicit barriers, UBEP proposes embedding synchronization signals directly into data payloads using hardware-native atomic operations.

• Mechanism: Leveraging the 512-byte atomic load/store capability of the Ascend NPU and the Unified-Bus (UB) fabric, UBEP encapsulates control flags within the data block itself.
• Variants:
  • Token-Flag Fusion (TFF): Encapsulates a 32-byte flag and 480-byte payload in a single 512-byte atomic write.
  • Data-for-Checksum (DC): Uses the data itself as a verification marker, transmitting a checksum separately.
  • Sentinel Polling (SP): Receivers initialize buffers with a sentinel value; arrival is detected when the atomic write overwrites this value.
• Benefit: This approach achieves implicit, near-zero-overhead coordination at the token level, removing the need for separate control-plane messages and global barriers.

3. Key Contributions

1. Bottleneck Characterization: The paper identifies and quantifies how BSP-style serialization, synchronization overhead, and topology-agnostic scheduling limit bandwidth utilization in modern superpods.
2. UBEP Design & Implementation: The authors design and implement UBEP, featuring:
   • Token-level kernel decomposition to maximize parallelism.
   • A hierarchical scheduler to mitigate tail latency from fabric heterogeneity.
   • The Data-as-Flag mechanism to minimize synchronization overhead.
3. Production-Scale Evaluation: UBEP is evaluated on a production Huawei CM384 superpod (256 NPU dies), demonstrating significant latency reductions and end-to-end performance gains for MoE inference.

4. Results

Experiments were conducted on a 256-NPU CM384 superpod using models ranging from 30B to 685B parameters (including Qwen3, GLM-4.7, and DeepSeek-R1).

• Operator-Level Performance: UBEP reduces All-to-All latency by up to 52.4% compared to the baseline CANN EP library. This improvement is driven by a 35.3%–40.8% increase in effective bandwidth and a reduction in synchronization overhead from ~70% of runtime (in baselines) to ~30%.
• End-to-End Performance: For MoE inference, UBEP improves the Time Per Output Token (TPOT) by up to 11.1%. The paper notes that while MoE communication accounts for ~50% of per-token latency, it represents only ~20% of actual hardware execution time due to dependency stalls; thus, the operator-level speedup translates to a smaller but significant end-to-end gain.
• Scalability: UBEP maintains consistent performance improvements across varying batch sizes, cluster scales (up to 256 ranks), and expert counts, effectively mitigating straggler effects in multi-tier topologies.

5. Significance and Claims

The paper claims that UBEP represents a necessary evolution in communication software for next-generation AI superpods. Its significance lies in shifting the paradigm from coarse-grained orchestration (BSP) to fine-grained parallelism driven by data dependencies.

• Architectural Shift: The work demonstrates that to fully unleash the potential of unified-memory, high-bandwidth superpods, software must adapt to the physical realities of the hardware, specifically the exposure of software overhead in low-latency fabrics and the non-uniformity of multi-tier topologies.
• Production Readiness: Unlike theoretical proposals, UBEP is implemented on the CANN stack and validated on a production-scale superpod, proving that dependency-driven execution and hardware-atomic synchronization are viable for large-scale MoE inference.
• Future Direction: The authors suggest that while UBEP leverages current hardware features (like 512B atomics), the ultimate evolution involves offloading control signals to dedicated hardware units (e.g., Scatter-Gather Engines) to further reduce programming complexity and latency.

The paper concludes that re-architecting communication libraries to align with the specific constraints and capabilities of modern superpods is critical for scaling MoE models efficiently.