I am a Founding Engineer at Axiom, working on the best AI model for Mathematics. I got my PhD from Rice University in Houston, Texas, developing Meta's LLM Compiler and HPCToolkit, a widely recognized open-source profiler for supercomputing systems. Throughout my career, I worked at leading research institutions such as Meta's Modern Recommendation Systems and Foundational AI Research, Berkeley Lab, and the Institute for High-Performance Microelectronics in Germany. In addition to technical expertise, I am actively involved in fostering the growth of the Serbian AI Ecosystem, serving as Ecosystem Team Lead for the Serbian AI Society.
Developing the best AI model for mathematics. I lead a 4-engineer team delivering a large-scale autoformalization pipeline, and engineered an AI system for code optimization that is formally verified in Lean. I also helped build Axiom’s first agentic system for automated theorem proving in Lean. Alongside this, I developed a system for scalable data collection and contributed to model training.
Co-designed foundation models for retrieval, ranking, and content understanding across Meta’s recommendation stack. I also optimized Facebook Feed ranking model training for GPU efficiency and scalability.
Working on use of Machine Learning in Compilers (thesis). Developed several tools and techniques that deploy Reinforcement Learning, Large Language Models, and advanced search techniques in compilers. Additionally, working on development of HPCToolkit, large-scale profiling tool used widely accross US national labs. Implemented infrastructure for scalable GPU tracing, node-level metric tracing, performance counters, and various performance analysis.
Contributed to Meta’s LLM Compiler, a suite of foundation models for compiler optimization. I fine-tuned Llama 2-based models on LLVM IR programs to solve the phase-ordering problem, and analyzed LLM capability in applying 60 LLVM optimizations. I also developed Priority Sampling, a method that reaches the performance of the test data within 30 samples.
Built LoopTune, a deep reinforcement learning framework for optimizing tensor operations. I built and tuned the cost and policy deep learning models behind it.
Profiling and analysis of power consumption on multi node GPU applications by using Nvidia NVML library.
Hands-on tutorials on cutting-edge supercomputing ATPESC 2021.
Profiling and Analysis of FFT implementation on Xtensa Platform in C and theoretical analysis window functions.
Measuring and characterization of materials on micro-identer device
| Thesis | Optimizing Compiler Heuristics with Machine Learning |
| Advisors | John Mellor-Crummey, Aleksandar Zlateski and Chris Cummins |
Thesis focus on the use of Machine Learning in Compilers. First, we developed LoopTune, a reinforcement-learning-based framework for optimizing tensor computations, a core component of ML workloads. Second, we pioneered the use of Large Language Models (LLMs) in compiler optimization by predicting the sequence of LLVM optimization flags directly from LLVM-IR in text form. Third, Finally, we developed Unique Sampling, a simple deterministic sampling technique for LLM that produces unique samples ordered by the model’s confidence and outperforms the label’s performance with 30 samples. Additionally, developed infrastructure for scalable GPU profiling over many GPU nodes. Added support for measuring performance counters and node level metrics in HPCToolkit, as well as GPU-idleness analysis, which points to the cause of serialization in GPU code.
| Grade | 10/10 |
| Thesis | Finding Shortest Path in Dynamic Large-scale Graph, based on Lambda Architecture |
| Advisors | Vladimir Dimitrieski |
Developed the system for detecting the shortest path from multiple source in large-scale dynamic graph based on Lambda Architecture. Technologies used: Spark, HDFS, Kafka, Python Dash, Docker, Python
| Grade | 9.96/10 |
| Thesis | Hardware acceleration of chess engine |
| Advisors | Vuk Vrankovic |
FPGA implementation of chess board evaluation by following RTL methodology. Technologies used: C, SystemC, VHDL, SystemVerilog
| Authors | David Kurniadi Angdinata, Evan Chen, Chris Cummins, Ben Eltschig, Dejan Grubisic, et al. |
| Publication | Indagationes Mathematicae (also on arXiv) |
Assuming the abc Conjecture, we prove that Ramanujan’s tau function misses a density-one subset of the primes, taking prime values far less often than previously known.
| Authors | Evan Chen, Chris Cummins, Ben Eltschig, Dejan Grubisic, et al. |
| Publication | Available on arXiv |
We prove that almost all primes are partially regular, with the proof automatically formalized in the Lean proof assistant, and draw implications for p-adic L-functions and algebraic K-theory.
| Authors | Evan Chen, Chris Cummins, Ben Eltschig, Dejan Grubisic, et al. |
| Publication | Available on arXiv |
We study “dead ends” in square-free digit walks, derive a closed-form expression for their density across all bases, and show the true density is far smaller than a prior 2024 prediction.
| Authors | Evan Chen, Chris Cummins, Ben Eltschig, Dejan Grubisic, et al. |
| Publication | Available on arXiv |
We prove Fel’s conjecture, giving an explicit formula for the syzygies of numerical semigroups via gap power sums, with the full proof formalized in the Lean proof assistant.
| Authors | Chris Cummins, Volker Seeker, Dejan Grubisic, Baptiste Roziere, Jonas Gehring, Gabriel Synnaeve, Hugh Leather |
| Publication | Compiler Construction (CC ‘25) |
We present LLM Compiler, a suite of LLMs built on Code Llama and trained on 546B tokens of LLVM-IR and assembly for compiler optimization, reaching 77% of the optimizing potential of an autotuning search and enabling assembly disassembly.
| Authors | Dejan Grubisic, Chris Cummins, Volker Seeker, Hugh Leather |
| Publication | Available on Arxiv |
We present Priority Sampling, a simple and deterministic sampling technique that produces unique samples ordered by the model’s confidence. Priority Sampling outperforms Nucleus Sampling for any number of samples, boosting the performance of the original model from 2.87% to 5% improvement over -Oz.
| Authors | Dejan Grubisic, Chris Cummins, Volker Seeker, Hugh Leather |
| Publication | Available on arXiv |
We introduce a system in which an LLM optimizes LLVM IR using compiler feedback: the model emits optimized IR, the best passes, and instruction counts, and the compiler validates the result to drive iterative improvement.
| Authors | Chris Cummins, Volker Seeker, Dejan Grubisic, Mostafa Elhoushi, Youwei Liang, Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Kim Hazelwood, Gabriel Synnaeve, Hugh Leather |
| Publication | Available on Arxiv |
We present a 7B-parameter LLaMa2-based model trained from scratch to optimize LLVM assembly for code size. Our approach achieves a 3.0% improvement in reducing instruction counts over the compiler with zero compilations, outperforming two state-of-the-art baselines that require thousands of compilations.
| Authors | Dejan Grubisic, Bram Wasti, Chris Cummins, Aleksandar Zlateski |
| Publication | International Conference on Compiler Construction 24’ (Under submission) |
We present LoopTune, a deep reinforcement learning framework for optimizing tensor computations in deep learning models. LoopTune consistently exceeds the performance of traditional search-based algorithms, TVM, performing at the level of hand-tuned library Numpy.
| Authors | Bram Wasti, Dejan Grubisic, Benoit Steiner, Aleksandar Zlateski |
| Publication | Neural Information Processing Systems 22’ |
We present LoopStack, a domain-specific compiler stack for tensor operations. LoopStack is orders of magnitude faster than LLVM, while resulting in equal or improved run time performance, while defining predictable optimization space suitable for tuning with reinforcement learning.
| Authors | Keren Zhou, Laksono Adhianto, Jonathon Anderson, Aaron Cherian, Dejan Grubisic, Mark Krentel, Yumeng Liu, Xiaozhu Meng, John Mellor-Crummey |
| Publication | Parallel Computing Journal |
To address the challenge of performance analysis on the US DOE forthcoming exascale supercomputers, Rice University has been extending its HPCToolkit performance tools to support measurement and analysis of GPU-accelerated applications.
| Authors | Aaron Thomas Cherian, Keren Zhou, Dejan Grubisic, Xiaozhu Meng, John Mellor-Crummey |
| Publication | ProTools |
In this paper, we describe extensions to Rice University’s HPCToolkit performance tools that support measurement and analysis of Intel’s DPC++ programming model for GPU-accelerated systems atop an implementation of the industry-standard OpenCL framework for heterogeneous parallelism on Intel GPUs
| Authors | Keren Zhou, Xiaozhu Meng, Ryuichi Sai, Dejan Grubisic, John Mellor-Crummey |
| Publication | TPDS |
In this paper, we describe GPA, a performance advisor that suggests potential code optimizations at a hierarchy of levels, including individual lines, loops, and functions. To gather the fine-grained measurements needed to produce such insights, GPA uses instruction sampling and binary instrumentation to monitor execution of GPU code. sing GPA, we obtained speedups on a Volta V100 GPU ranging from 1.01x to 3.58x, with a geometric mean of 1.22x
| Authors | Dejan Grubisic |
| Publication | Zbornik radova Fakulteta tehničkih nauka u Novom Sadu |
Finding shortest path in dynamic multi source large-scale graph, based on Lambda architecture
| Duration | 1 hour 13 minutes |
Demonstrating the use of machine learning in compilers, using reinforcement learning, large language models, and advanced techniques for sampling.
| Duration | 13 minutes |
Presentation of Priority Sampling - a simple deterministic sampling technique for LLM that achieves suprisingly high performance.
| Duration | 45 minutes |
Implementing system for finding shortest path in dynamic graph suitable for big data.
| Duration | 45 minutes |
Implemenging FPGA component for accelerating evaluation of the board in chess engine.
| Duration | 5 minutes |
Here we present the newly added features of monitoring power, temperature, and utilization on Nvidia GPUs in HPCToolkit.
Founded a summit dedicated to growing the AI ecosystem and connecting researchers, builders, and companies across Southeast Europe.
Leading ecosystem-building efforts to foster the growth of the Serbian AI community, connecting academia, industry, and the diaspora.
| Python | |
| C/C++ | |
| CudaC | |
| GNU / Linux | |
| Bash | |
| Lean | |
| OpenMP/MPI | |
| VHDL | |
| Docker | |
| Java | |
| Spark | |
| Hadoop | |
Curious, innovative, driven AI researcher with a passion for building AI for Systems.
