# LLMForge: Multi-Backend Hardware-Aware Neural Architecture Search with Infinite-Head Attention for Edge Language Models

arXiv: https://arxiv.org/abs/2605.17653 · HTML: https://mlsanity.com/paper/2605.17653/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 29 of its 37 references that are in the map.
- First posted on arXiv 17 May 2026.

## Builds on (most cited first, 29 of 29)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [Crowdsourcing Multiple Choice Science Questions](https://mlsanity.com/paper/1707.06209/) (2017)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Model](https://mlsanity.com/paper/2502.02737/) (2025)
- [MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases](https://mlsanity.com/paper/2402.14905/) (2024)
- [HAT: Hardware-Aware Transformers for Efficient Natural Language Processing](https://mlsanity.com/paper/2005.14187/) (2020)
- [Zeus: Understanding and Optimizing GPU Energy Consumption of DNN Training](https://mlsanity.com/paper/2208.06102/) (2022)
- [The MiniPile Challenge for Data-Efficient Language Models](https://mlsanity.com/paper/2304.08442/) (2023)
- [OpenELM: An Efficient Language Model Family with Open Training and Inference Framework](https://mlsanity.com/paper/2404.14619/) (2024)
- [Gemmini: Enabling Systematic Deep-Learning Architecture Evaluation via Full-Stack Integration](https://mlsanity.com/paper/1911.09925/) (2019)
- [Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search](https://mlsanity.com/paper/2508.15884/) (2025)
- [FLAT: An Optimized Dataflow for Mitigating Attention Bottlenecks](https://mlsanity.com/paper/2107.06419/) (2021)
- [Transformers in Speech Processing: A Survey](https://mlsanity.com/paper/2303.11607/) (2023)
- [MELTing point: Mobile Evaluation of Language Transformers](https://mlsanity.com/paper/2403.12844/) (2024)
- [HW-GPT-Bench: Hardware-Aware Architecture Benchmark for Language Models](https://mlsanity.com/paper/2405.10299/) (2024)
- [STAR: Synthesis of Tailored Architectures](https://mlsanity.com/paper/2411.17800/) (2024)
- [Composer: A Search Framework for Hybrid Neural Architecture Design](https://mlsanity.com/paper/2510.00379/) (2025)
- [Conformer-Based Speech Recognition On Extreme Edge-Computing Devices](https://mlsanity.com/paper/2312.10359/) (2023)
- [Dynamic Transformer for Efficient Machine Translation on Embedded Devices](https://mlsanity.com/paper/2107.08199/) (2021)
- [TransCODE: Co-design of Transformers and Accelerators for Efficient Training and Inference](https://mlsanity.com/paper/2303.14882/) (2023)

## Abstract (arXiv)

Sub-billion-parameter Transformer language models are increasingly deployed on edge devices, where the privacy, latency, and operating-cost advantages of on-device inference are constrained by tight memory-bandwidth, energy, and thermal budgets that make architectural choice and accelerator-specific cost central to efficient inference. We present LLMForge, a hardware-aware neural architecture search (NAS) framework whose three composable contributions together make edge-LM architecture search hardware-conditioned, since different substrates impose different hardware cost bottlenecks. Infinite-Head Attention (IHA) decouples the number of query heads, KV groups, and per-head query/key and value dimensions, expanding the feasible per-layer attention configuration space by approximately 400x over grouped-query attention within our search-space ranges. Forge-Former, an encoder-based surrogate for ranking architectural candidates, outperforms MLP and random-forest baselines. Forge-DSE, an NSGA-II-based design-space-exploration engine, pairs Forge-Former with a multi-backend hardware cost model spanning GPUs, systolic accelerators, and ring-dataflow edge accelerators. Across four different hardware substrates, the searches converge to visibly different architectures whose shapes track each substrate's cost bottleneck. On the multi-chip ring substrate, our co-search returns three 300M-scale deployment-aware variants on the Pareto front. Each is re-trained on FineWeb-Edu-10BT under matched recipe against SmolLM2-360M and Qwen-0.5B architecture baselines. The accurate variant has the lowest validation loss 2.798 and competitive benchmark performance with fewer parameters, the energy-optimized variant lowers energy per token by 40%, and the latency-optimized variant lowers TTFT and TPOT by 43%.