Latency-Aware Transformer Partitioning for Heterogeneous Edge Inference

Deploying transformer models in heterogeneous edge environments poses significant challenges due to limited memory, varied compute capabilities, and asymmetric inter-node bandwidth. To enable efficient distributed inference, it is necessary to partition the model into segments that can be collaboratively executed across multiple edge nodes. However, previous works predominantly focus on either single-point partitioning or static scheduling strategies, which fall short in capturing the joint impact of segmentation and assignment on end-to-end latency. In this work, we formulate a latency-aware transformer deployment problem that jointly optimizes model segmentation and segment-to-node assignment, explicitly accounting for memory constraints and heterogeneous communication costs. We propose efficient algorithms to explore the solution space, including a latency-aware balanced partitioning heuristic (LaBP) and a dynamic programming-based optimal strategy (DP-MCP). Extensive experiments demonstrate that our approach consistently achieves significantly lower inference latency compared to existing baselines while maintaining feasibility under stringent resource constraints.