Runtime-Integrated Simulation of Accelerator Memory Behaviour for Distributed LLM Deployment

Developing hardware for AI is both costly and technically demanding. The software stack evolves rapidly and is often optimized around GPUs. These realities raise a fundamental challenge: how can one validate hardware that does not yet exist for emerging AI workloads? Fortunately, most AI applications are built on a relatively small set of frameworks, notably PyTorch and JAX, which provide natural integration pathways for new accelerators. Hooking into the integration layers from software frameworks to hardware offers a practical interface for accelerator simulators to assess both functional correctness and performance characteristics.

Taking inspiration from PyTorch's example accelerator integration (OpenReg), this internship effort focuses on simulating accelerator memory behaviour during distributed LLM deployment. The goal is to enable early-stage evaluation of memory bottlenecks, communication patterns, and execution efficiency before hardware is physically available, thereby reducing design risk and accelerating the development cycle for novel AI accelerators.

Objectives

• develop a simulator for faithfully modelling accelerator memory behaviour in distributed LLM execution;
• integrate the model with existing PyTorch execution paths to capture realistic memory access patterns;
• analyze memory usage, identify potential memory bottlenecks under distributed deployment scenarios;
• disseminate the findings by publishing in appropriate conference or journal.

Desirable skills

• Familiarity with Python and C++; experience with PyTorch is a plus
• Interest in computer architecture; knowledge of architectural simulators such as gem5 and onnxim is a plus
• Basic familiarity with LLM models and distribution strategies