Elamur Research⌕
☰
home / papers / 2026 / Deep Kernel Fusion for Transformers

Deep Kernel Fusion for Transformers

Paper · 2026 · Zixi Zhang, Zhiwen Mo, Yiren Zhao et al.

Summary

In long-context, agentic LLM inference, decoding is limited more by memory bandwidth than by compute. The large SwiGLU MLP weights do not fit in cache, so these blocks become a costly and rarely optimized part of the Transformer. The authors present DeepFusionKernel, a CUDA operator that merges the Up and Gate GEMMs, the SiLU activation and the elementwise product into one kernel. Existing stacks use four kernels (naive PyTorch) or two (SGLang, vLLM), and they write intermediate activations to HBM; the fused kernel avoids this. It offers row-major and column-major tiling variants. A profiler-based scheduler benchmarks these variants on the target hardware before inference and keeps the fastest one. The authors tested LLaMA 3.1 70B inside SGLang on four-GPU A100 and H100 clusters. The kernel raised decoding throughput by up to 9.7% on A100 and 13.2% on H100 over SGLang, and it also beat vLLM on very long outputs. The gains require no extra FLOPs and fit into a production serving framework.

Key points

  • Fusion scope: The MLP is split into two stages. The first stage (Up/Gate GEMMs, SiLU and gating multiply) is fused into one kernel, while the Down projection remains separate. Row-major tiling favors reuse of the activations, and column-major tiling favors reuse of the weights, which suits small batches with large models.
  • Setup: LLaMA 3.1 70B in FP16 with tensor parallelism TP=4 on 80 GB SXM A100 or H100 GPUs. Tests ran in SGLang with FlashInfer and CUDA Graphs, using batch sizes 1–64 and four runs per configuration.
  • Peak speedups over SGLang: +9.7% on A100 (batch 16) and +13.2% on H100 (batch 2). Gains shrink at large batches, for example +1.3% on A100 at batch 64.
  • Long generation: Output lengths ranged from 1024 to 16,384 tokens at batch sizes 1, 4 and 16. DeepFusionKernel had the best throughput against SGLang and vLLM in the reported cases.
  • Code and limitations: The code is released at github.com/ZixiBenZhang/deepfusionkernel. The scheduler adds no dispatch cost at runtime. Different GPU interconnects and communication-induced variance were not evaluated exhaustively.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

 Fields, topics and concepts

Cite

Cite as Zixi Zhang, Zhiwen Mo, Yiren Zhao and Robert D. Mullins. 2026. Deep Kernel Fusion for Transformers. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.acl-short.15

@inproceedings{zhang2026deep,
  title = {{Deep Kernel Fusion for Transformers}},
  author = {Zixi Zhang, Zhiwen Mo, Yiren Zhao and Robert D. Mullins},
  booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)},
  year = {2026},
  doi = {10.18653/v1/2026.acl-short.15},
  url = {https://doi.org/10.18653/v1/2026.acl-short.15}
}

References

30 references, 13 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, and Yuxiong He. 2022. DeepSpeed inference: Enabling efficient inference of transformer models at unprecedented scale. Preprint, arXiv:2207.00032. find
  2. Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, Binyuan Hui, Luo Ji, Mei Li, Junyang Lin, Runji Lin, Dayiheng Liu, Gao Liu, Chengqiang Lu, Keming Lu, and 29 others. 2023. Qwen technical report. Preprint, arXiv:2309.16609.
  3. Luca Beurer-Kellner, Marc Fischer, and Martin Vechev. 2023. Prompting is programming: A query language for large language models. Proceedings of the ACM on Programming Languages, 7(PLDI):1946–1969. find
  4. Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter,… — Language Models are Few-Shot Learners
  5. Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. TVM: An automated end-to-end optimizing compiler for deep learning. Preprint, arXiv:1802.04799.
  6. DeepSeek-AI, Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, and 181 others. 2025a. DeepSeek-R1: Incentivizing reasoning capability in l…
  7. DeepSeek-AI, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, and 181 others. 2025b. DeepSeek-V3 technical report. Preprint, arXiv:… find
  8. Ofer Dekel. 2025. Blockbuster, part 1: Block-level AI operator fusion. Preprint, arXiv:2505.07829. find
  9. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, and 542 o… — The Llama 3 Herd of Models
  10. Shengyu Liu Jiashi Li. 2025. FlashMLA: Efficient MLA decoding kernels. https://github.com/ deepseek-ai/FlashMLA. ↗
  11. Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. 2023. Efficient memory management for large language model serving with PagedAttention. Preprint, arXiv:2309.06180.
  12. Microsoft. 2022. DeepSpeed Model Implementations for Inference (MII). https://github.com/ deepspeedai/DeepSpeed-MII. Accessed: 2025-05-26. ↗
  13. NVIDIA. 2018. nvidia/apex. find
  14. Apex. https://github.com/ ↗
  15. NVIDIA Corporation. 2021. NVIDIA A100 Tensor Core GPU datasheet. https://www.nvidia.com/content/dam/ en-zz/Solutions/Data-Center/a100/pdf/ nvidia-a100-datasheet-us-nvidia-1758950-r4-web. pdf. Accessed: 2025-05-26. ↗
  16. NVIDIA Corporation. 2024. NVIDIA H100 Tensor Core GPU datasheet. https://resources.nvidia.com/en-us-gpu-resources/ h100-datasheet-24306. Accessed: 2025-05-26. ↗
  17. NVIDIA Corporation. 2025a. NVIDIA DGX B200 datasheet. https://resources.nvidia.com/ en-us-dgx-systems/dgx-b200-datasheet. Accessed: 2025-05-26. ↗
  18. NVIDIA Corporation. 2025b. NVIDIA TensorRT-LLM documentation. https://nvidia.github.io/ TensorRT-LLM/. Accessed: 2025-05-26. ↗
  19. OpenAI, :, Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, Alex Iftimie, Alex Karpenko, Alex Tachard Passos, Alexander Neitz, Alexander Prokofiev, Alexander Wei, Allison Tam, and 244 others. 2024. OpenAI o1 sy…
  20. Zixuan Zhou, Xuefei Ning, Ke Hong, Tianyu Fu, Jiaming Xu, Shiyao Li, Yuming Lou, Luning Wang, Zhihang Yuan, Xiuhong Li, Shengen Yan, Guohao Dai, Xiao-Ping Zhang, Yuhan Dong, and Yu Wang. 2024. A survey on efficient inference for large language models. Preprint, arXiv:2404.14294. find
  21. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, and 2 others. 2019.… — PyTorch: An Imperative Style, High-Performance Deep Learning Library
  22. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. 2024. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. Preprint, arXiv:2407.08608.
  23. Noam Shazeer. 2020. GLU variants improve transformer. Preprint, arXiv:2002.05202. find
  24. Yining Shi, Zhi Yang, Jilong Xue, Lingxiao Ma, Yuqing Xia, Ziming Miao, Yuxiao Guo, Fan Yang, and Lidong Zhou. 2023. Welder: Scheduling deep learning memory access via tile-graph. In 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23), pages 701–718, Boston, MA. USENIX Ass… find
  25. Philippe Tillet, Hsiang-Tsung Kung, and David Cox. 2019. 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, pages 10–19.
  26. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2023. Attention is all you need. Preprint, arXiv:1706.03762.
  27. Haoran Wu, Can Xiao, Jiayi Nie, Xuan Guo, Binglei Lou, Jeffrey T. H. Wong, Zhiwen Mo, Cheng Zhang, Przemyslaw Forys, Wayne Luk, Hongxiang Fan, Jianyi Cheng, Timothy M. Jones, Rika Antonova, Robert Mullins, and Aaron Zhao. 2025. Combating the memory walls: Optimization pathways for long-context agent… find
  28. Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. 2023. React: Synergizing reasoning and acting in language models. Preprint, arXiv:2210.03629.
  29. Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, and Luis Ceze. 2025. FlashInfer: Efficient and customizable attention engine for LLM inference serving. Preprint, arXiv:2501.01005. find
  30. Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark Barrett, and Ying Sheng. 2024. SGLang: Efficient execution of structured language model programs. Preprint, arXiv:2312.07104.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhang-2026-deep-kernel-fusion-transformers/