Converting PyTorch Models to StreamIt Pipelines

With the rise of large language models, there have been efforts to optimize machine learning inference to support a large volume of queries. Currently, the two main ways to do this are running optimized kernels for computing the forward inference pass and distributing computation across multiple GPUs or different cores in a GPU. Machine learning libraries such as PyTorch produce dynamic computation graphs in order to represent the forward pass of the model. PyTorch allows conversion of these dynamic graphs into static ones through just-in-time (JIT) compilation. These graphs can then be optimized further by the compiler. We propose an alternate way of optimizing these dynamic graphs. We convert the dynamic computation graph of PyTorch to pipelines in StreamIt, a domain specific language (DSL) for streaming applications, and use the multi-stage compilation property of BuildIt to compile this pipeline in stages to inference code. We found that, while the inference latencies of models compiled in this way are slightly higher, they are still comparable to those of PyTorch models and are open to future optimizations.

Name thesis:degree_name

Master

Department dc:contributor.department

Massachusetts Institute of Technology. Department of Electrical Engineering and Computer Science

Grantor dc:publisher

Massachusetts Institute of Technology

Year dc:date.issued

2025

Author dc:creator

• Rajvee, Muhender Raj

Advisor dc:contributor.advisor

• Amarasinghe, Saman

Statement dc:rights

• In Copyright - Educational Use Permitted
• Copyright retained by author(s)

Licence dc:rights.uri

• https://rightsstatements.org/page/InC-EDU/1.0/

Handle dc:identifier.uri

https://hdl.handle.net/1721.1/163041

OAI identifier oai:identifier

oai:dspace.mit.edu:1721.1/163041