Search papers, labs, and topics across Lattice.
This paper investigates StableHLO as a unified intermediate representation for cross-architecture performance modeling of distributed ML workloads. They map StableHLO representations of workloads to analytical models, profiling-based predictors, and simulators for GPUs and TPUs. Results show that StableHLO preserves relative performance trends across architectures and modeling fidelities, enabling cross-platform comparison without requiring scaled-out physical systems.
StableHLO unlocks cross-architecture performance prediction for distributed ML, revealing accuracy trade-offs and simulator limitations without needing access to expensive, scaled-out hardware.
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.
