Hundred-billion-parameter models don't run on phones. But 300M-parameter models can—if the architecture is chosen well. LLMForge, proposed by Jiang, Luo, Qi, and team, is a hardware-aware neural architecture search (NAS) framework designed specifically for small language models on edge devices.
Three core contributions:
1. Infinite-Head Attention (IHA) — decouples the number of query heads, KV groups, and per-head dimension, expanding the search space by roughly 400x. 2. Forge-Former — an encoder-based surrogate model for fast ranking of candidate architectures. 3. Forge-DSE — an NSGA-II-based search engine combined with multi-backend hardware cost models (GPU, systolic array, ring dataflow edge accelerators).
Across four hardware backends, the search converges on clearly distinct architectures—the shapes track each backend's cost bottleneck. On a ring dataflow chip, the framework produced three 300M-parameter Pareto-front variants:
- Accuracy-first: validation loss 2.798, with fewer parameters
- Energy-efficiency-first: 40% lower energy per token
- Latency-first: 43% lower first-token and per-token latency
- How does the 300M-parameter model perform comprehensively on broad tasks (knowledge, reasoning, code)?
- Does the 400x search space from IHA actually accelerate things on real hardware—are the gains entirely from better architecture-hardware matching?
- What is the search cost—how many architectures must NSGA-II evaluate before converging?
Open questions (from the poster)
References
1. Jiang, X., Luo, J., Qi, R., et al. (2026). *LLMForge: Multi-Backend Hardware-Aware NAS with Infinite-Head Attention for Edge Language Models*. arXiv:2605.17653 [cs.LG]. 2. Shazeer, N. (2019). *Fast Transformer Decoding: One Write-Head is All You Need*. arXiv. 3. Ainslie, J., et al. (2023). *GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints*. EMNLP.