- The paper proposes a measurement-backed, greedy dispatch policy that combines exact-shape operator costs with backend-transition penalties across PyTorch CPU, PyTorch CUDA, and ONNX Runtime CPU.
- Transition-aware replay improves latency by 17.4%, energy by 14.4%, and energy-delay product by 28.5% over the best static backend policy across 9,584 operator instances.
- The method reduces backend switching by roughly 14–15% and generalizes to most held-out models, but ModernBERT latency worsens by 15.49%, highlighting the need for workload-matched training or fallback mechanisms.
Overview
This paper addresses backend selection for transformer inference on edge platforms, where the relative performance of individual operators varies with tensor shape and where switching between execution backends incurs measurable overhead. The authors propose a transition-aware dispatch policy that augments per-operator feature-based selection with the previously selected backend state, so that a candidate backend is chosen by combining its measured operator cost with a measured transition penalty. The evaluation is conducted through measurement-backed trace replay on an NVIDIA Jetson Orin Nano across seven transformer models, three eager-mode backends (PyTorch CPU, PyTorch CUDA, and ONNX Runtime CPU), and four operator classes (Linear, MLP activation, RMSNorm, LayerNorm). Across 9,584 ordered operator instances and 278 exact shape groups, transition-aware dispatch reduces replayed latency, energy, and energy-delay product (EDP) relative to the best static policy by 17.4%, 14.4%, and 28.5% on average, respectively.
The paper identifies two failure modes in existing deployment strategies. Static assignment of one backend to an entire model cannot exploit shape-dependent performance differences: a backend that minimizes latency for large matrix multiplications may be inefficient for small normalization or activation operators, and the latency-optimal backend may differ from the energy- or EDP-optimal one. Operator-local selection, which independently assigns each supported operator to its lowest-cost backend, captures these differences but can generate frequent device and framework switches whose synchronization, data movement, and framework-conversion costs offset local gains. Prior heterogeneous inference systems such as Neurosurgeon, LaLaRAND, AxoNN, and Band schedule at the granularity of models, layers, or graph partitions; this work instead targets repeated exact-shape operators within eager-mode transformer execution, considering transitions across both devices and runtime frameworks.
The dispatch scope is deliberately restricted to four operator classes that dominate transformer compute; all other operators retain a static assignment common to every evaluated policy. This scoping isolates the contribution of selective dispatch but also bounds the achievable gains, since attention and KV-cache operations are excluded from dynamic selection.
Methodology
The pipeline proceeds in five stages. First, ordered operator traces are collected from full-model inference runs of Qwen2.5-0.5B, Qwen3-0.6B, TinyLlama-1.1B, SmolLM2-135M, SmolLM2-360M, OLMo2-1B, and ModernBERT-base, using synthetic inputs at batch size 1, FP32 precision, and prefill lengths of 16–128 tokens. Notably, decode traces contain a single length-1 forward call rather than a complete autoregressive KV-cache loop, so decode-phase behavior is only partially represented. Second, instances are grouped by exact shape and context, yielding 278 unique groups. Third, each shape-backend pair is benchmarked directly on Jetson, producing 834 measurement rows with latency, idle-corrected power, energy, EDP, and correctness validation against PyTorch CPU references. Fourth, selectors are trained separately for latency, energy, and EDP. The transition-aware cost model is additive:
Ci(bi)=Ciop(bi)+Citrans(bi−1,bi)
with directed transition penalties estimated from source-to-target backend-pair measurements. Expanding every shape across three previous-backend states yields 2,502 candidate-cost rows and 834 decision contexts. Histogram-based gradient boosting (HistGBM) is selected via GroupKFold validation grouped by shape identifier, preventing leakage between training and validation folds. Fifth, policies are evaluated by analytical replay over the model-derived traces, retrieving measured costs rather than executing an integrated mixed-framework runtime.
Two methodological caveats bear directly on interpretation. The replay is causal and greedy rather than globally sequence-optimal, so reported gains are lower bounds on what a sequence-optimal policy could achieve. More importantly, integrated dispatcher and model-level framework overheads are outside the replay scope, meaning the reported improvements quantify potential under a measured-cost model rather than end-to-end system measurements.
Selector quality
Both selector formulations maintain strong prediction quality. The operator-local selector achieves accuracies of 0.8741, 0.8705, and 0.8597 for latency, energy, and EDP, with regret below 0.35% for every objective and as low as 0.0052% for EDP. The transition-aware selector achieves higher accuracies (0.8921, 0.8897, 0.9005) over its expanded previous-backend decision contexts, with regret remaining below 0.41%. Because the two selectors operate over different context spaces and target costs, their accuracies are not treated as a head-to-head comparison; the relevant question is whether conditioning on previous-backend state improves end-to-end replay outcomes, which the ablation addresses.
Main results
Positive gains over the best static policy are obtained for all seven traces and all three objectives. Best static policies average 4,782.0 ms and 39.10 J per trace; transition-aware replay reduces these to 4,012.4 ms and 34.07 J. Per-trace gains vary substantially with trace structure: SmolLM2-135M shows the strongest result at 37.1% latency, 32.0% energy, and 57.3% EDP improvement, while ModernBERT-base and OLMo2-1B show more modest EDP gains of 14.2% and 14.1%. Gains are computed per trace before averaging, preventing larger traces from dominating the mean.
The ablation separates the contributions of the two design elements. Most of the improvement over static assignment comes from shape-dependent selection itself: operator-local selection alone achieves average gains of 16.8%, 13.7%, and 27.4% for latency, energy, and EDP. Transition awareness adds modestly to these figures (17.4%, 14.4%, 28.5%), but its primary effect is stability: it reduces backend switches by roughly 14–15% relative to operator-local selection, with device switches decreasing by up to 19.7% and framework switches by up to 17.6%. This decomposition indicates that previous-backend context functions mainly as a switch-suppression mechanism rather than a major source of additional cost reduction.
Generalization
Leave-one-model-out evaluation tests transfer beyond the training distribution. Energy improves for all seven held-out models, and latency and EDP improve for six, with mean held-out gains of 12.51%, 12.96%, and 23.68% respectively. Six decoder-oriented traces generalize consistently, with SmolLM2-135M again strongest at 31.66% latency and 49.75% EDP gain.
ModernBERT-base is the clear stress case and the paper reports it plainly: while its energy improves by 2.61%, latency and EDP degrade by −15.49% and −6.10%. This negative result indicates that the learned policy does not transfer reliably when a structurally different encoder-style trace presents an operator-shape distribution absent from decoder-dominated training data. The implication is practical: deployment of such a selector requires either training coverage matching the target workload family or a fallback mechanism for out-of-distribution shapes, neither of which the current system provides.
Limitations and open questions
Several limitations are acknowledged or evident. The evaluation relies on trace replay rather than an integrated mixed-backend runtime, so dispatcher overheads, framework conversion costs during connected execution, and any interaction effects between consecutive operators are not captured; validating replay predictions end-to-end remains open. The workload coverage excludes attention, KV-cache operations, quantized operators, fused kernels, and optimized backends such as ONNX Runtime CUDA and TensorRT, leaving the generality of the approach across these configurations untested. Decode-phase modeling is limited to single length-1 forward calls rather than autoregressive generation loops. The greedy, causal policy leaves open how much additional benefit a globally sequence-optimal formulation would provide. Finally, the ModernBERT-base regression raises the specific question of how much encoder-style training data is needed for reliable cross-architecture transfer.
Conclusion
This paper demonstrates that exact-shape operator costs and backend-transition context are useful signals for selective backend dispatch on heterogeneous edge platforms. Transition-aware dispatch improves replayed latency, energy, and EDP by 17.4%, 14.4%, and 28.5% on average over the best static policy, reduces switching relative to operator-local selection, and transfers to six of seven held-out models—while honestly exposing its boundary on structurally dissimilar workloads. The principal open issue is whether these replay-measured gains survive integration into a real mixed-backend inference runtime.