This paper addresses the challenges of deploying Large Language Models on specialized hardware by introducing an MLIR-based compilation pipeline. It utilizes two dialects: TopOp for high-level, framework-agnostic model semantics, and TpuOp for target hardware-specific operations. The method aims to improve the import of trained models into compiler-friendly representations and optimize scheduling for autoregressive inference under memory constraints.
- Introduces TopOp dialect for framework-independent model representation
- Uses TpuOp dialect to map operations to specific AI accelerator hardware
- Addresses memory limitations during autoregressive inference scheduling
- Provides a structured path from trained models to hardware execution