pith. sign in

arxiv: 2510.08726 · v2 · submitted 2025-10-09 · 💻 cs.PL · cs.LG

Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs

Yifan Zhao , Egan Johnson , Prasanth Chatarasi , Vikram Adve , Sasa Misailovic This is my paper

Pith reviewed 2026-05-18 08:59 UTC · model grok-4.3

classification 💻 cs.PL cs.LG
keywords operator fusiontensor compilerattentionGPUdeep learningreduction operatorskernel optimization
0
0 comments X

The pith

Neptune enables better GPU performance for attention by fusing reduction operators through dependency breaking and algebraic corrections.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Neptune is a tensor compiler designed for advanced fusion of sequences of reduction operators in deep learning. It works by intentionally breaking some dependencies in the computation and then using algebraic correction expressions to ensure the results are correct. This allows generating high-performance kernels equivalent to FlashAttention directly from plain attention code and a scheduling template. The approach addresses limitations in existing compilers like Triton and TVM that struggle with complex reductions involving loop-carried dependencies. If the method holds, it suggests a path to more automatic optimization of key deep learning workloads on various GPU architectures.

Core claim

Neptune presents a new approach for advanced operator fusion, which intentionally breaks some existing dependencies and compensates by constructing algebraic correction expressions that allow the kernel to produce the correct result. Applying Neptune's advanced operator fusion to a plain attention operator generates operators equivalent to FlashAttention and FlashDecoding. On ten attention-based benchmarks across four GPU architectures, Neptune-generated kernels achieve an average speedup of 1.35× over the next best alternative.

What carries the argument

Algebraic correction expressions built after intentionally breaking dependencies in reduction operator sequences.

If this is right

  • Generates operators equivalent to FlashAttention and FlashDecoding from plain attention code.
  • Delivers average 1.35× speedup over Triton, TVM, and FlexAttention on attention benchmarks.
  • Achieves up to 2.65× speedup on Nvidia GPUs and up to 3.32× on AMD GPUs.
  • Applies effectively to deep learning workloads with complex reduction computations.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • This technique could potentially extend to fusion of other types of operators beyond reductions in ML models.
  • High-level scheduling templates might become a standard way to guide compilers without low-level manual tuning.
  • Such dependency-breaking with corrections might simplify the development of custom kernels for new hardware.

Load-bearing premise

The algebraic correction expressions always yield mathematically equivalent results to the original dependent computations for the reduction sequences involved.

What would settle it

Observing a discrepancy in the numerical output between the Neptune fused kernel and the original attention computation on any of the benchmarks would indicate the corrections do not preserve correctness.

Figures

Figures reproduced from arXiv: 2510.08726 by Egan Johnson, Prasanth Chatarasi, Sasa Misailovic, Vikram Adve, Yifan Zhao.

Figure 1
Figure 1. Figure 1: c. Values of the repair term and the repaired xsum are shown in green. Note,𝑠 ⟨3⟩ 𝑖 =𝑠 ⟨3⟩ 𝑖 [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: The main components of Neptune. • The template-guided optimizer automatically applies the transformations in the template to the loop nests in the program. Section 4 presents two novel advanced fusion algorithmsthat composewith otherbuilt-intransformations. At this point, we have a loop-scalar program that has un￾dergone high-level optimizations. • To facilitate tile optimizations, the loop-to-tile transla… view at source ↗
Figure 3
Figure 3. Figure 3: The result of applying split-k update to Figure 1a. Algorithm 2: SplitKUpdate(𝐿𝑡 , 𝑙) Input: 𝐿𝑡 : the target loop nest to be fused Input: 𝑙: the target loop to fuse 𝐿𝑡 under Output: A pair of loop nests: the local and global versions of fused 𝐿𝑡 1 L𝑟 = InlineMapReturnReduce(𝐿𝑡 ); 2 (𝑀local,𝑀global) = (dict(), dict()); 3 for 𝐿 ∈ L𝑟 do 4 (𝐿 local,𝐿global) = FuseAndPrivatize(𝐿, 𝑙); 5 (𝑀local[𝐿],𝑀global[𝐿]) = … view at source ↗
Figure 4
Figure 4. Figure 4: An example program in Neptune’s tile IR. Neptune translates the program in loop-scalar IR (above) to the program in tile IR (below). i.e., tile expressions express computation that the tile opti￾mizers can process. We provide the syntax of the tile IR in Appendix E. Neptune translates a program from the loop-scalar IR to the tile IR, by detecting parts of the program with computation regular enough to expr… view at source ↗
Figure 5
Figure 5. Figure 5: Performance of Neptune kernels vs. kernels by other tensor compilers. The y-axis shows relative performance of all compilers normalized to Neptune for each setup. The plots for the remaining benchmarks are in Appendix F. geomean over input sequence lengths. Out of 320 setups we evaluate, Neptune achieves better or equal performance com￾pared to all other compilers on 284 setups. Neptune shows improvement o… view at source ↗
Figure 7
Figure 7. Figure 7: Ablation studies of the Neptune kernel performance on sequence length 512 on RTX 6000 Ada. The x-axis labels mark the component of Neptune we keep and remove. not extend to non-attention operators, and it uses multiple manually optimized templates to generate kernels. Tile-Based Compilers. Many tensor compilers and pro￾gramming frameworks decompose tensor-level operations into operations over tensor tiles,… view at source ↗
Figure 6
Figure 6. Figure 6: Throughput of Neptune kernels and multiple base￾lines over increasing batch sizes, evaluated on RTX A5000 for the GQA (PF) operator. The y-axis shows throughput in TFLOPS/sec, and the x-axis shows batch size. TVM Tile Optim. No Fusion Fusion No Tuning Full Neptune 0.0 0.2 0.4 0.6 0.8 1.0 Relative Performance Global (PF) Causal (PF) [PITH_FULL_IMAGE:figures/full_fig_p011_6.png] view at source ↗
Figure 9
Figure 9. Figure 9: shows the schedule that pairs with the compute def￾inition. In 28 lines of code, it applies rolling update to fuse the attention computation into a single loop nest, and produces many of the high-performance kernels in our evaluation. 1 def create_general_attention( 2 B: int, N: int, QS: int, KVS: int, H: int, 3 mask_cond: Callable, 4 score_mod: Callable, 5 ): 6 q = placeholder((B, N, QS, H), "float16", na… view at source ↗
Figure 8
Figure 8. Figure 8: shows a flexible compute definition for attention that Neptune takes as input. It allows customizing bias and mask conditions, and in 38 lines of code it covers all operators in [PITH_FULL_IMAGE:figures/full_fig_p014_8.png] view at source ↗
Figure 10
Figure 10. Figure 10: All intermediate results of rolling update, when applied to Figure 1a. C.2 Privatization [PITH_FULL_IMAGE:figures/full_fig_p016_10.png] view at source ↗
Figure 11
Figure 11. Figure 11: presents privatization combined with rolling update in Neptune. On top of the rolling update output in Figure 1c, we privatize s_max and s_sum with a split size of 2. xmax_1p and xsump are the local arrays. 1 for i in range(2): 2 xmax_0[i] = -inf 3 for j1 in range(2): 4 for j2 in range(2): 5 # max_local 6 xmax_1p[i, j1] = max(xmax_1p[i, j1], inp[i, j1 * 2 + j2]) 7 # max_global 8 xmax_1[i] = max(xmax_1[i],… view at source ↗
Figure 13
Figure 13. Figure 13: Neptune tile IR syntax. Non-terminals that are the same as [PITH_FULL_IMAGE:figures/full_fig_p019_13.png] view at source ↗
Figure 14
Figure 14. Figure 14: The result of tensorizing the program in [PITH_FULL_IMAGE:figures/full_fig_p019_14.png] view at source ↗
Figure 15
Figure 15. Figure 15: Performance of Neptune vs. other tensor compilers on the five operators not shown in [PITH_FULL_IMAGE:figures/full_fig_p021_15.png] view at source ↗
Figure 16
Figure 16. Figure 16: Performance of Neptune vs. tensor libraries on 8 operators for which library baselines exist (we were unable to find the library implementations of SoftCap PF/DC). 22 [PITH_FULL_IMAGE:figures/full_fig_p022_16.png] view at source ↗
Figure 17
Figure 17. Figure 17: extends the scalability test of [PITH_FULL_IMAGE:figures/full_fig_p023_17.png] view at source ↗
Figure 17
Figure 17. Figure 17: (Continued) Throughput of Neptune kernels and multiple base- lines over increasing batch size, for all 10 operators on all 4 GPUs. The y-axis shows throughput in TFLOPS/sec, and the x-axis shows batch size [PITH_FULL_IMAGE:figures/full_fig_p024_17.png] view at source ↗
Figure 18
Figure 18. Figure 18: Timeline view of a profile in the Nvidia Nsight System profiler. The bottow rows show when the kernels are executing, and the “GPC Clock Frequency” row near the top shows the frequency of the GPU graphic clock over time. 24 [PITH_FULL_IMAGE:figures/full_fig_p024_18.png] view at source ↗
read the original abstract

Operator fusion has become a key optimization for deep learning, which combines multiple deep learning operators to improve data reuse and reduce global memory transfers. However, existing tensor compilers struggle to fuse complex reduction computations involving loop-carried dependencies, such as attention mechanisms. This paper introduces Neptune, a tensor compiler for advanced operator fusion for sequences of reduction operators. Neptune presents a new approach for advanced operator fusion, which intentionally breaks some existing dependencies and compensates by constructing algebraic correction expressions that allow the kernel to produce the correct result. Applying Neptune's advanced operator fusion to a plain attention operator generates operators equivalent to FlashAttention and FlashDecoding. On ten attention-based benchmarks, Neptune, starting from a plain attention code and a high-level scheduling template, outperforms existing compilers like Triton, TVM, and FlexAttention, including Triton-based implementations of FlashAttention. Across four different GPU architectures from NVIDIA and AMD, Neptune-generated kernels have an average speedup of $1.35\times$ over the next best alternative, with up to $2.65\times$ speedup on Nvidia GPUs and up to $3.32\times$ on AMD GPUs, demonstrating its effectiveness for deep learning workloads.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The paper introduces Neptune, a tensor compiler for advanced operator fusion on sequences of reduction operators with loop-carried dependencies (e.g., attention). It achieves fusion by intentionally breaking dependencies and compensating via constructed algebraic correction expressions, claiming to produce FlashAttention- and FlashDecoding-equivalent operators from plain attention code plus a high-level scheduling template. On ten attention benchmarks across four NVIDIA and AMD GPUs, Neptune reports an average 1.35× speedup over the next-best compiler (Triton, TVM, FlexAttention), with peaks of 2.65× (NVIDIA) and 3.32× (AMD).

Significance. If the algebraic corrections are shown to preserve mathematical equivalence across the targeted reduction patterns, the technique would offer a principled route to automatic fusion of complex reductions that current compilers handle poorly, with clear practical value for attention-heavy workloads. The cross-architecture speedups are a positive signal, but the absence of any derivation details, equivalence arguments, or numerical validation in the manuscript limits the strength of the contribution until those elements are supplied.

major comments (2)
  1. Abstract: the central performance claims rest on the assertion that the algebraic correction expressions 'allow the kernel to produce the correct result' and yield FlashAttention-equivalent operators. No derivation of the corrections, symbolic equivalence argument, machine-checked proof, or numerical stress-test protocol across input ranges and reduction orders is supplied, leaving the equivalence claim unverified and directly undermining the reported speedups.
  2. Experimental section (implied by benchmark results): the reported average 1.35× speedup (and per-architecture maxima) on ten benchmarks is presented without error bars, full experimental protocol, or reproducibility artifacts, making it impossible to assess whether the gains are robust or sensitive to floating-point ordering, overflow, or unhandled edge cases in the correction expressions.
minor comments (2)
  1. Abstract: the high-level scheduling template is mentioned but never characterized; a short description of its structure and how it exposes fusion opportunities would improve clarity.
  2. Consider adding a table that lists the ten benchmarks together with the exact speedup numbers against each baseline for quick reference.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive feedback. We address each major comment below and outline the revisions we will make to strengthen the manuscript.

read point-by-point responses

  1. Referee: Abstract: the central performance claims rest on the assertion that the algebraic correction expressions 'allow the kernel to produce the correct result' and yield FlashAttention-equivalent operators. No derivation of the corrections, symbolic equivalence argument, machine-checked proof, or numerical stress-test protocol across input ranges and reduction orders is supplied, leaving the equivalence claim unverified and directly undermining the reported speedups.

    Authors: The manuscript explains the construction of algebraic correction expressions that compensate for intentionally broken loop-carried dependencies in reduction sequences, enabling generation of FlashAttention-equivalent operators from plain attention code. We acknowledge that the current presentation provides only a high-level description without a full step-by-step symbolic derivation or explicit equivalence argument in the main text or appendix. In the revised manuscript we will add a dedicated subsection (or appendix) that derives the correction expressions for the targeted attention reduction patterns and presents a symbolic argument establishing mathematical equivalence to the unfused computation. We will also include a numerical validation protocol together with results across representative input ranges and reduction orders to confirm that results match within floating-point tolerance. A machine-checked proof lies outside the scope of this systems paper, but the added algebraic argument and empirical checks will directly address the verification concern. revision: yes

  2. Referee: Experimental section (implied by benchmark results): the reported average 1.35× speedup (and per-architecture maxima) on ten benchmarks is presented without error bars, full experimental protocol, or reproducibility artifacts, making it impossible to assess whether the gains are robust or sensitive to floating-point ordering, overflow, or unhandled edge cases in the correction expressions.

    Authors: We agree that the experimental reporting can be improved to allow readers to evaluate robustness. The revised manuscript will add error bars derived from multiple independent runs, a detailed experimental protocol section (including hardware configurations, software versions, benchmark construction, and measurement methodology), and explicit reproducibility artifacts such as a public code repository link and scripts. We will further include results from targeted stress tests on the correction expressions that examine sensitivity to floating-point ordering, potential overflow conditions, and edge-case inputs. revision: yes

Circularity Check

0 steps flagged

No circularity: performance claims rest on external benchmark comparisons

full rationale

The paper's central results are empirical speedups measured by executing Neptune-generated kernels on ten fixed attention benchmarks and comparing runtimes against independent external systems (Triton, TVM, FlexAttention). No equations, fitted parameters, or self-citations are shown that would make any reported speedup equivalent to an internal input by construction. The algebraic correction step is presented as a mechanism to restore equivalence after dependency breaking, but the performance numbers themselves are obtained from direct, externally verifiable execution rather than from any self-referential derivation or renaming of known results.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The central claim rests on the assumption that algebraic corrections preserve semantics after dependency breaking; no free parameters or new invented entities are mentioned in the abstract.

axioms (1)
  • domain assumption Algebraic correction expressions can be constructed that restore exact equivalence after selected dependencies are broken inside reduction loops.
    This premise is required for the fusion to remain correct while still enabling the locality and parallelism benefits described.

pith-pipeline@v0.9.0 · 5747 in / 1226 out tokens · 36757 ms · 2026-05-18T08:59:00.649551+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference

    cs.CL 2026-05 unverdicted novelty 5.0

    Ada-MK fuses LLM operators into persistent MegaKernels via MLIR DAG search and 3D shared-memory modeling, delivering up to 23.6% higher single-batch throughput than TensorRT-LLM on NVIDIA L20.

Reference graph

Works this paper leans on

47 extracted references · 47 canonical work pages · cited by 1 Pith paper · 2 internal anchors

  1. [1]

    {TensorFlow}: a system for {Large-Scale} machine learning

    Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. {TensorFlow}: a system for {Large-Scale} machine learning. In12th USENIX symposium on operating systems design and implementation (OSDI 16), pages 265–283, 2016

    work page 2016

  2. [2]

    Learning to optimize halide with tree search and random programs.ACM Transactions on Graphics (TOG), 38, 2019

    Andrew Adams, Karima Ma, Luke Anderson, Riyadh Baghdadi, Tzu-Mao Li, Michaël Gharbi, Benoit Steiner, Steven Johnson, Kayvon Fatahalian, Frédo Durand, and Jonathan Ragan-Kelley. Learning to optimize halide with tree search and random programs.ACM Transactions on Graphics (TOG), 38, 2019

    work page 2019

  3. [3]

    Pallas: a jax kernel language.https: //docs.jax.dev/en/latest/pallas/index.html, 2024

    The JAX Authors. Pallas: a jax kernel language.https: //docs.jax.dev/en/latest/pallas/index.html, 2024

    work page 2024

  4. [4]

    Bhaskaracharya, Julien Demouth, and Vinod Grover

    Somashekaracharya G. Bhaskaracharya, Julien Demouth, and Vinod Grover. Automatic kernel generation for volta tensor cores.CoRR, abs/2006.12645, 2020

    work page arXiv 2006

  5. [5]

    Blume and R

    W. Blume and R. Eigenmann. Nonlinear and symbolic data dependence testing.IEEE Transactions on Parallel and Distributed Systems, 9(12):1180–1194, 1998

    work page 1998

  6. [6]

    TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

    Tianqi Chen, Thierry Moreau, Ziheng Jiang, Haichen Shen, Eddie Q. Yan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. TVM: end-to-end optimization stack for deep learning. CoRR, abs/1802.04799, 2018

    work page internal anchor Pith review Pith/arXiv arXiv 2018

  7. [7]

    Evt: Accelerating deep learning training with epilogue visitor tree

    Zhaodong Chen, Andrew Kerr, Richard Cai, Jack Kosaian, Haicheng Wu, Yufei Ding, and Yuan Xie. Evt: Accelerating deep learning training with epilogue visitor tree. InProceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, ASPLOS ’24, page 301–316, 2024

    work page 2024

  8. [8]

    Chillee. Where do the 2000+ pytorch operators come from? (pytorch developer discussions).https://dev-discuss.pytorch.org/t/where-do- the-2000-pytorch-operators-come-from-more-than-you-wanted- to-know/373

    work page 2000

  9. [9]

    NVIDIA cuSPARSELt.https://docs.nvidia.com/ cuda/cusparselt/types.html, 2021

    NVIDIA Corporation. NVIDIA cuSPARSELt.https://docs.nvidia.com/ cuda/cusparselt/types.html, 2021

    work page 2021

  10. [10]

    NVIDIA A10 Tensor Core GPU

    NVIDIA Corporation. NVIDIA A10 Tensor Core GPU. https://www.nvidia.com/en-us/data-center/a100/, 2024

    work page 2024

  11. [11]

    NVIDIA RTX 6000 Ada-generation Graphics Card.https://www.nvidia.com/en-us/design-visualization/rtx-6000/, 2024

    NVIDIA Corporation. NVIDIA RTX 6000 Ada-generation Graphics Card.https://www.nvidia.com/en-us/design-visualization/rtx-6000/, 2024

    work page 2024

  12. [12]

    NVIDIA RTX A5000 Graphics Card

    NVIDIA Corporation. NVIDIA RTX A5000 Graphics Card. https://www.nvidia.com/en-us/design-visualization/rtx-a5000/, 2024

    work page 2024

  13. [13]

    FlashAttention-2: Faster attention with better parallelism and work partitioning

    Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. InInternational Conference on Learning Representations (ICLR), 2024

    work page 2024

  14. [14]

    Flashattention: Fast and memory-efficient exact attention with io-awareness.Advances in neural information processing systems, 35:16344–16359, 2022

    Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness.Advances in neural information processing systems, 35:16344–16359, 2022

    work page 2022

  15. [15]

    Flash-decoding for long-context inference, Oct 2023

    Tri Dao, Grigory Sizov, Francisco Massa, and Daniel Haziza. Flash-decoding for long-context inference, Oct 2023

    work page 2023

  16. [16]

    FlashAttention.https://github.com/Dao-AILab/flash- attention, 2023

    Dao-AILab. FlashAttention.https://github.com/Dao-AILab/flash- attention, 2023

    work page 2023

  17. [17]

    Flex attention: A programming model for generating optimized attention kernels, 2024

    Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kernels, 2024

    work page 2024

  18. [18]

    Tensorir: An abstraction for automatic tensorized program optimization

    Siyuan Feng, Bohan Hou, Hongyi Jin, Wuwei Lin, Junru Shao, Ruihang Lai, Zihao Ye, Lianmin Zheng, Cody Hao Yu, Yong Yu, et al. Tensorir: An abstraction for automatic tensorized program optimization. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pages 804–817, 2023

    work page 2023

  19. [19]

    Taso: optimizing deep learning computation with automatic generation of graph substitutions

    Zhihao Jia, Oded Padon, James Thomas, Todd Warszawski, Matei Za- haria, and Alex Aiken. Taso: optimizing deep learning computation with automatic generation of graph substitutions. InProceedings of the 27th ACM Symposium on Operating Systems Principles, pages 47–62, 2019

    work page 2019

  20. [20]

    Optimizing dnn computation with relaxed graph substitutions.Proceedings of Machine Learning and Systems, 12 1:27–39, 2019

    Zhihao Jia, James Thomas, Todd Warszawski, Mingyu Gao, Matei Zaharia, and Alex Aiken. Optimizing dnn computation with relaxed graph substitutions.Proceedings of Machine Learning and Systems, 12 1:27–39, 2019

    work page 2019

  21. [21]

    xformers: A modular and hackable transformer modelling library.https://github.com/facebookresearch/xformers, 2022

    Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizen- stein, and Grigory Sizov. xformers: A modular and hackable transformer modelling library.https://github.com/facebookresearch/xformers, 2022

    work page 2022

  22. [22]

    Differentiable programming for image processing and deep learning in halide.ACM Trans

    Tzu-Mao Li, Michaël Gharbi, Andrew Adams, Frédo Durand, and Jonathan Ragan-Kelley. Differentiable programming for image processing and deep learning in halide.ACM Trans. Graph., 37(4), 2018

    work page 2018

  23. [23]

    Smith, Mateusz Paprocki, Ondřej Čertík, Sergey B

    Aaron Meurer, Christopher P. Smith, Mateusz Paprocki, Ondřej Čertík, Sergey B. Kirpichev, Matthew Rocklin, AMiT Kumar, Sergiu Ivanov, Jason K. Moore, Sartaj Singh, Thilina Rathnayake, Sean Vig, Brian E. Granger, Richard P. Muller, Francesco Bonazzi, Harsh Gupta, Shivam Vats, Fredrik Johansson, Fabian Pedregosa, Matthew J. Curry, Andy R. Terrel, Štěpán Rou...

    work page 2017

  24. [24]

    Automatically scheduling halide image processing pipelines.ACM Trans

    Ravi Teja Mullapudi, Andrew Adams, Dillon Sharlet, Jonathan Ragan-Kelley, and Kayvon Fatahalian. Automatically scheduling halide image processing pipelines.ACM Trans. Graph., 35(4), jul 2016

    work page 2016

  25. [25]

    Newcomb, Andrew Adams, Steven Johnson, Rastislav Bodik, and Shoaib Kamil

    Julie L. Newcomb, Andrew Adams, Steven Johnson, Rastislav Bodik, and Shoaib Kamil. Verifying and improving halide’s term rewriting system with program synthesis.Proc. ACM Program. Lang., 4(OOPSLA), November 2020

    work page 2020

  26. [26]

    Dnnfusion: accelerating deep neural networks execution with advanced operator fusion

    Wei Niu, Jiexiong Guan, Yanzhi Wang, Gagan Agrawal, and Bin Ren. Dnnfusion: accelerating deep neural networks execution with advanced operator fusion. InProceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation, PLDI 2021, page 883–898, 2021

    work page 2021

  27. [27]

    CUTLASS: CUDA Templates for Linear Algebra Subroutines

    NVIDIA. CUTLASS: CUDA Templates for Linear Algebra Subroutines. https://github.com/NVIDIA/cutlass, 2021

    work page 2021

  28. [28]

    Fused Attention – Triton Documentation.https://triton-lang

    OpenAI. Fused Attention – Triton Documentation.https://triton-lang. org/main/getting-started/tutorials/06-fused-attention.html, 2024

    work page 2024

  29. [29]

    Xla.https://openxla.org/xla

    OpenXLA. Xla.https://openxla.org/xla

    work page

  30. [30]

    Pytorch: An imperative style, high-performance deep learn- ing library.Advances in neural information processing systems, 32, 2019

    Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learn- ing library.Advances in neural information processing systems, 32, 2019

    work page 2019

  31. [31]

    Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines

    Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines. InACM SIGPLAN Conference on Programming Language Design and Implementation, 2013

    work page 2013

  32. [32]

    Einops: Clear and versatile tensor manipulations for deep learning.https://github.com/arogozhnikov/einops, 2020

    Alex Rogozhnikov. Einops: Clear and versatile tensor manipulations for deep learning.https://github.com/arogozhnikov/einops, 2020. GitHub repository, Accessed: [Date Accessed]

    work page 2020

  33. [33]

    Tensor program optimization with probabilistic programs

    Junru Shao, Xiyou Zhou, Siyuan Feng, Bohan Hou, Ruihang Lai, Hongyi Jin, Wuwei Lin, Masahiro Masuda, Cody Hao Yu, and Tianqi Chen. Tensor program optimization with probabilistic programs. In Advances in Neural Information Processing Systems, volume 35, 2022

    work page 2022

  34. [34]

    Approxtuner: a compiler and runtime system for adaptive approximations

    Hashim Sharif, Yifan Zhao, Maria Kotsifakou, Akash Kothari, Ben Schreiber, Elizabeth Wang, Yasmin Sarita, Nathan Zhao, Keyur Joshi, Vikram S Adve, Sasa Misailovic, and Sarita V Adve. Approxtuner: a compiler and runtime system for adaptive approximations. In Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, 2021

    work page 2021

  35. [35]

    Spector, Simran Arora, Aaryan Singhal, Daniel Y

    Benjamin F. Spector, Simran Arora, Aaryan Singhal, Daniel Y. Fu, and Christopher Ré. Thunderkittens: Simple, fast, and adorable ai kernels, 2024

    work page 2024

  36. [36]

    Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, page 10–19, June 2019

    work page 2019

  37. [37]

    PET: Optimizing tensor programs with partially equivalent transformations and automated corrections

    Haojie Wang, Jidong Zhai, Mingyu Gao, Zixuan Ma, Shizhi Tang, Liyan Zheng, Yuanzhi Li, Kaiyuan Rong, Yuanyong Chen, and Zhihao Jia. PET: Optimizing tensor programs with partially equivalent transformations and automated corrections. In15th USENIX Symposium on Operating Systems Design and Implementation (OSDI 21), pages 37–54, July 2021

    work page 2021

  38. [38]

    Unit: Unifying tensorized instruction compilation

    Jian Weng, Animesh Jain, Jie Wang, Leyuan Wang, Yida Wang, and Tony Nowatzki. Unit: Unifying tensorized instruction compilation. In2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), page 77–89, February 2021

    work page 2021

  39. [39]

    Mirage: A {Multi-Level} superoptimizer for tensor programs

    Mengdi Wu, Xinhao Cheng, Shengyu Liu, Chunan Shi, Jianan Ji, Man Kit Ao, Praveen Velliengiri, Xupeng Miao, Oded Padon, and Zhihao Jia. Mirage: A {Multi-Level} superoptimizer for tensor programs. In19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25), pages 21–38, 2025

    work page 2025

  40. [40]

    Equality saturation for tensor graph superop- timization

    Yichen Yang, Phitchaya Phothilimthana, Yisu Wang, Max Willsey, Sudip Roy, and Jacques Pienaar. Equality saturation for tensor graph superop- timization. In A. Smola, A. Dimakis, and I. Stoica, editors,Proceedings of Machine Learning and Systems, volume 3, pages 255–268, 2021

    work page 2021

  41. [41]

    FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

    Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, and Luis Ceze. Flashinfer: Efficient and customiz- able attention engine for llm inference serving.arXiv preprint arXiv:2501.01005, 2025

    work page internal anchor Pith review Pith/arXiv arXiv 2025

  42. [42]

    Felix: Optimizing tensor programs with gradient descent

    Yifan Zhao, Hashim Sharif, Vikram Adve, and Sasa Misailovic. Felix: Optimizing tensor programs with gradient descent. InProceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, ASPLOS ’24, page 367–381, 2024

    work page 2024

  43. [43]

    Approxcaliper: A programmable framework for application-aware neural network optimization

    Yifan Zhao, Hashim Sharif, Peter Pao-Huang, Vatsin Ninad Shah, Arun Narenthiran Sivakumar, Mateus Valverde Gasparino, Abdulrahman Mahmoud, Nathan Zhao, Sarita Adve, Girish Chowdhary, Sasa Misailovic, and Vikram Adve. Approxcaliper: A programmable framework for application-aware neural network optimization. In Proceedings of Machine Learning and Systems 5, 2023

    work page 2023

  44. [44]

    Gonzalez, and Ion Stoica

    Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang Zhuo, Koushik Sen, Joseph E. Gonzalez, and Ion Stoica. Ansor: Generating high- performance tensor programs for deep learning. InUSENIX Conference on Operating Systems Design and Implementation, OSDI’20, 2020

    work page 2020

  45. [45]

    Amos: enabling automatic mapping for tensor computationson spatial accelerators with hardware abstraction

    Size Zheng, Renze Chen, Anjiang Wei, Yicheng Jin, Qin Han, Liqiang Lu, Bingyang Wu, Xiuhong Li, Shengen Yan, and Yun Liang. Amos: enabling automatic mapping for tensor computationson spatial accelerators with hardware abstraction. InProceedings of the 49th Annual International Symposium on Computer Architecture, ISCA ’22, page 874–887, 2022

    work page 2022

  46. [46]

    Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system

    Size Zheng, Yun Liang, Shuo Wang, Renze Chen, and Kaiwen Sheng. Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system. In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS ’20, page 859–873, 2020

    work page 2020

  47. [47]

    float16", name=

    Hongyu Zhu, Ruofan Wu, Yijia Diao, Shanbin Ke, Haoyu Li, Chen Zhang, Jilong Xue, Lingxiao Ma, Yuqing Xia, Wei Cui, Fan Yang, Mao Yang, Lidong Zhou, Asaf Cidon, and Gennady Pekhimenko. ROLLER: Fast and efficient tensor compilation for deep learning. In16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 233–248, July 2022. ...

    work page 2022