LightCode: Compiling LLM Inference for Photonic-Electronic Systems

📄 arXiv: 2509.16443v1 📥 PDF

作者: Ryan Tomich, Zhizhen Zhong, Dirk Englund

分类: physics.app-ph, cs.AI, cs.PL

发布日期: 2025-09-19

备注: 9 pages, 8 figures

💡 一句话要点

LightCode：用于光子-电子系统的LLM推理编译框架

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 光子计算 大型语言模型 异构计算 编译优化 能效优化

📋 核心要点

  1. 现有GPU在与光子张量单元等新兴加速器集成方面存在不足，无法满足LLM对低延迟、高能效推理的需求。
  2. LightCode提出了一种混合编译策略，通过堆叠图表示张量操作的多种硬件实现，并优化硬件分配。
  3. 实验表明，LightCode在GPT-2和Llama-7B上，能显著降低能耗并提升推理速度，验证了其有效性。

📝 摘要（中文）

针对大型语言模型（LLM）对低延迟、高能效推理日益增长的需求，异构架构的研究备受关注。虽然GPU仍然占据主导地位，但它们不适合与新兴的特定领域加速器（如光子张量单元（PTU））集成，后者提供低功耗、高吞吐量的线性计算。这促使人们研究结合光子和电子资源的混合编译策略。我们提出了LightCode，一个用于在混合光子-电子系统上映射LLM推理工作负载的编译器框架和模拟器。LightCode引入了堆叠图，一种中间表示，它编码了每个张量操作的多个硬件特定实现。硬件分配被表述为一个约束子图选择问题，该问题在参数化成本模型下针对延迟或能量进行优化。我们评估了LightCode在GPT-2和Llama-7B的预填充阶段的表现，结果表明，在我们的工作负载和硬件假设下，(i) 在模拟工作负载中，光子硬件在最大序列长度下降低了高达50%的能量；(ii) 多路复用和分配策略产生了超过10倍的延迟改进；(iii) 在我们的模拟中，针对延迟或能量进行优化导致了不同的硬件映射。LightCode为将LLM编译到新兴的光子加速器提供了一个模块化、基础性的框架和模拟器。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）在异构光子-电子系统上的高效推理问题。现有方法，特别是依赖GPU的方案，在能效和延迟方面存在瓶颈，无法充分利用新兴的光子计算加速器的优势。因此，如何有效地将LLM计算任务映射到混合的光子-电子硬件资源上，以实现更低的功耗和更高的吞吐量，是本研究要解决的核心问题。

核心思路：LightCode的核心思路是构建一个编译器框架，能够感知不同硬件（包括光子和电子设备）的特性，并根据给定的优化目标（例如，最小化延迟或能量消耗）自动地将LLM的计算图映射到最合适的硬件资源上。这种混合编译策略旨在充分利用光子计算在特定线性代数运算上的优势，同时利用电子计算在其他操作上的灵活性。

技术框架：LightCode的技术框架主要包含以下几个关键模块：1) 堆叠图（Stacked Graph）：作为中间表示，它将每个张量操作表示为多个硬件特定的实现，从而允许编译器在不同的硬件选项之间进行选择。2) 硬件分配器：将硬件分配问题建模为一个约束子图选择问题，并使用优化算法（例如，整数线性规划）来找到满足约束条件并最大化性能的硬件映射方案。3) 参数化成本模型：用于估计不同硬件配置下的延迟和能量消耗，为硬件分配提供依据。4) 模拟器：用于验证编译后的代码在目标硬件上的性能。

关键创新：LightCode的关键创新在于其混合编译策略和堆叠图的表示方法。传统的编译器通常针对单一类型的硬件进行优化，而LightCode能够同时考虑光子和电子硬件的特性，并根据优化目标动态地选择最合适的硬件资源。堆叠图允许编译器探索不同的硬件映射方案，从而实现更好的性能。

关键设计：LightCode的关键设计包括：1) 硬件成本模型的参数化：成本模型需要准确地反映不同硬件的性能特征，例如，光子计算的延迟和功耗与矩阵大小的关系。2) 约束子图选择问题的建模：需要定义合适的约束条件，以确保硬件映射方案的可行性，例如，硬件资源的容量限制和数据传输的带宽限制。3) 优化算法的选择：需要选择高效的优化算法，以便在合理的时间内找到最优或近似最优的硬件映射方案。

📊 实验亮点

实验结果表明，在GPT-2和Llama-7B的预填充阶段，LightCode在模拟工作负载中，光子硬件在最大序列长度下降低了高达50%的能量消耗。此外，通过多路复用和硬件分配策略，延迟改进超过10倍。针对延迟或能量进行优化会导致不同的硬件映射方案，表明LightCode能够根据不同的优化目标进行灵活的硬件资源分配。

🎯 应用场景

LightCode的研究成果可应用于各种需要低延迟、高能效LLM推理的场景，例如边缘计算设备、移动设备和数据中心。通过将LLM计算任务卸载到光子加速器上，可以显著降低功耗并提高推理速度，从而实现更高效的AI应用。该研究也为未来光子计算在人工智能领域的应用奠定了基础。

📄 摘要（原文）

The growing demand for low-latency, energy-efficient inference in large language models (LLMs) has catalyzed interest in heterogeneous architectures. While GPUs remain dominant, they are poorly suited for integration with emerging domain-specific accelerators like the Photonic Tensor Units (PTUs), which offer low-power, high-throughput linear computation. This motivates hybrid compilation strategies that combine photonic and electronic resources. We present LightCode, a compiler framework and simulator for mapping LLM inference workloads across hybrid photonic-electronic systems. LightCode introduces the Stacked Graph, an intermediate representation that encodes multiple hardware-specific realizations of each tensor operation. Hardware assignment is formulated as a constrained subgraph selection problem optimized for latency or energy under parametric cost models. We evaluate LightCode on the prefill stage of GPT-2 and Llama-7B showing that under our workload and hardware assumptions, (i) Photonic hardware reduced energy by up to 50% in our simulated workloads at maximum sequence length; (ii) multiplexing and assignment strategy yielded latency improvements exceeding 10x; and (iii) Optimizing for latency or energy resulted in distinct hardware mappings in our simulations. LightCode offers a module, foundational framework and simulator for compiling LLMs to emerging photonic accelerators.