Relay Buffer Independent Communication over Pooled HBM for Efficient MoE Inference on Ascend

2026-05-07

AI Summary

This work addresses the performance bottlenecks in Mixture-of-Experts (MoE) inference caused by large-scale cross-device token communication, where the dispatch and combine phases suffer from significant overhead due to conventional relay buffers inducing costly layout transformations. To overcome this, the authors propose a relay-buffer-free communication mechanism on the Ascend platform, leveraging globally pooled HBM, symmetric memory allocation, and direct expert window access. This design enables tokens to be written directly into target expert memory regions and read remotely, relying solely on lightweight synchronization metadata. The approach substantially reduces dispatch and combine latency, accelerates time-to-first-token (TTFT), maintains competitive time-per-output-token (TPOT), and expands the feasible scheduling space under latency constraints.

Abstract

Mixture-of-Experts (MoE) inference requires large-scale token exchange across devices, making dispatch and combine major bottlenecks in both prefill and decode. Beyond network transfer, routing-driven layout transformation, temporary relay, and output restoration can add substantial overhead. Existing MoE communication paths are often buffer-centric, using explicit inter-process relay and reordering buffers around collective transfer. This report presents a relay-buffer-free communication design for MoE inference acceleration on Ascend systems. The design reorganizes dispatch and combine around direct placement into destination expert windows and direct reading from remote expert windows. Built on globally pooled high-bandwidth memory and symmetric-memory allocation, it removes most intermediate relay and reordering buffers while retaining only lightweight control state, including counts, offsets, and synchronization metadata. We instantiate the design as two schedules for the main phases of MoE inference: a prefill schedule with richer planning state for throughput-oriented execution, and a compact decode schedule for latency-sensitive execution. Experiments on Ascend-based MoE workloads show reduced dispatch and combine latency in both settings. At the serving level, the implementation improves time to first token (TTFT), preserves competitive time per output token (TPOT), and enlarges the feasible scheduling space under practical latency constraints. These results indicate that, on platforms with globally addressable device memory, reducing intermediate buffering and output restoration around expert execution is an effective direction for accelerating MoE inference.

Problem

• Mixture-of-Experts
• MoE inference
• dispatch and combine
• relay buffer
• communication overhead

Innovation

• relay-buffer-free
• pooled HBM
• MoE inference
• direct placement
• symmetric-memory allocation