Evaluating Cross-Architecture Performance Modeling of Distributed ML Workloads Using StableHLO

2026-04-13

AI Summary

This work addresses the challenge that existing performance modeling approaches for distributed machine learning struggle to uniformly characterize workloads across heterogeneous accelerators such as GPUs and TPUs, often relying on architecture-specific assumptions or costly real-world execution. The study presents the first systematic validation of StableHLO as a cross-architecture unified intermediate representation and introduces an MLIR-based, multi-fidelity performance modeling framework. By integrating analytical models, profiling data, and simulators, the framework enables reproducible, cross-platform performance prediction without requiring large-scale hardware deployment. Evaluated on GEMM, ResNet, and large language model training tasks, the approach accurately captures cross-architecture performance trends with prediction errors within practical bounds, while also exposing fidelity limitations in current GPU simulators.

Technology Category

• Machine Learning: Hardware-aware ML
• Humans and AI: Human-in-the-loop Machine Learning
• Natural Language Processing: (Large) Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

Abstract

Predicting the performance of large-scale distributed machine learning (ML) workloads across multiple accelerator architectures remains a central challenge in ML system design. Existing GPU and TPU focused simulators are typically architecture-specific, while distributed training simulators rely on workload-specific analytical models or costly post-execution traces, limiting portability and cross-platform comparison. This work evaluates whether MLIR's StableHLO dialect can serve as a unified workload representation for cross-architecture and cross-fidelity performance modeling of distributed ML workloads. The study establishes a StableHLO-based simulation methodology that maps a single workload representation onto multiple performance models, spanning analytical, profiling-based, and simulator-driven predictors. Using this methodology, workloads are evaluated across GPUs and TPUs without requiring access to scaled-out physical systems, enabling systematic comparison across modeling fidelities. An empirical evaluation covering distributed GEMM kernels, ResNet, and large language model training workloads demonstrates that StableHLO preserves relative performance trends across architectures and fidelities, while exposing accuracy trade-offs and simulator limitations. Across evaluated scenarios, prediction errors remain within practical bounds for early-stage design exploration, and the methodology reveals fidelity-dependent limitations in existing GPU simulators. These results indicate that StableHLO provides a viable foundation for unified, distributed ML performance modeling across accelerator architectures and simulators, supporting reusable evaluation workflows and cross-validation throughout the ML system design process.

Problem

• cross-architecture
• distributed ML
• performance modeling
• StableHLO
• workload representation

Innovation

• StableHLO
• cross-architecture modeling
• distributed ML performance
• MLIR
• performance simulation