LLMForge：多后端硬件感知神经架构搜索，具有边缘语言模型的无限头注意力

数十亿参数的 Transformer 语言模型越来越多地部署在边缘设备上，其中设备上推理的隐私、延迟和运营成本优势受到紧张的内存带宽、能源和热预算的限制，这使得架构选择和加速器特定的成本成为高效推理的核心。我们提出了 LLMForge，一种硬件感知的神经架构搜索 (NAS) 框架，其三个可组合的贡献共同使边缘 LM 架构搜索受到硬件限制，因为不同的基底会带来不同的硬件成本瓶颈。无限头注意力（IHA）解耦了查询头的数量、KV 组以及每头查询/键和值维度，在我们的搜索空间范围内将可行的每层注意力配置空间比分组查询注意力扩展了大约 400 倍。 Forge-Former 是一种基于编码器的代理，用于对候选架构进行排名，其性能优于 MLP 和随机森林基线。 Forge-DSE 是一种基于 NSGA-II 的设计空间探索引擎，将 Forge-Former 与涵盖 GPU、脉动加速器和环数据流边缘加速器的多后端硬件成本模型配对。在四种不同的硬件基板上，搜索集中到明显不同的架构，其形状跟踪每个基板的成本瓶颈。在多芯片环形基板上，我们的联合搜索返回了 Pareto 前沿的三个 300M 规模的部署感知变体。每个模型都在 FineWeb-Edu-10BT 上根据 SmolLM2-360M 和 Qwen-0.5B 架构基线的匹配配方进行了重新训练。准确的变体具有最低的验证损失 2.798 和具有竞争力的基准性能，参数较少，能量优化的变体将每个词元的能量降低了 40%，延迟优化的变体将 TTFT 和 TPOT 降低了 43%。