arXiv:2605.17653cs.LGcs.AI2026-05

为边缘设备设计高效语言模型,自动适配不同硬件瓶颈。

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

论文配图:LLMForge: Multi-Backend Hardware-Aware Neural Architecture Search with Infinite-Head Attention for Edge Language Models
图 1 · 摘自论文原文
  • 提出无限头注意力机制,扩展400倍配置空间。
  • 在多硬件上搜索出三类优化模型,能耗降40%,延迟降43%。
  • 适合边缘AI部署、芯片定制与模型轻量化研究者。

参数量在十亿以下的Transformer语言模型正越来越多地部署于边缘设备,但其隐私、低延迟和低成本优势受限于严格的内存带宽、能耗与热预算,使架构选择与加速器特定开销成为高效推理的核心。我们提出LLMForge,一个硬件感知神经架构搜索框架,其三个可组合贡献共同实现边端语言模型架构搜索对硬件条件的依赖性:无限头注意力(IHA)解耦查询头数、键值组数及每头维度,使单层注意力配置空间在搜索范围内扩大约400倍;Forge-Former作为基于编码器的候选架构评分代理,优于MLP与随机森林基线;Forge-DSE采用NSGA-II设计空间探索引擎,结合跨GPU、阵列加速器与环形数据流边端加速器的多后端硬件成本模型。在四种不同硬件上,搜索结果呈现明显差异,且架构形状精准匹配各自成本瓶颈。在多芯片环形架构上,协同搜索生成三个300M规模的部署友好变体,均在FineWeb-Edu-10BT上按相同训练方案重训练,对比SmolLM2-360M与Qwen-0.5B基线:准确型验证损失最低(2.798),参数更少;能效优化型每令牌能耗降低40%;延迟优化型首字延迟(TTFT)与每令牌处理时间(TPOT)均降低43%。

原文摘要 · Abstract (English)

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%.

边缘计算架构搜索注意力机制轻量化模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。