Qualcomm NPU 编译器的逆向工程

# 逆向工程高通 NPU 编译器 来源：https://datavorous.github.io/writing/qairt/ 我的工作是最大化 NPU 的利用率，让我们想要在 NPU 上运行的任何模型在边缘部署中更快。但网上关于 NPU 的文档基本不存在，仅有的那一点也令人失望，以至于我一度想过放弃——所以我干脆逆向工程了编译器。我之前写过一篇关于 NPU 的小入门文章：它们是什么以及哪里会出问题（https://datavorous.github.io/writing/npu/），应该足以理解接下来的内容。 对于任何 SoC，高通都没有公布 VTCM 的内存容量。我该怎么知道我的张量是否到处溢出？或者量化是否真的必要？再加上我好奇他们如何在模型实际运行在硬件上之前模拟其工作，以及涉及哪些优化算法。 我（配备了 Claude Code（https://claude.com/product/claude-code））深入研究了 QNPU SDK（https://www.qualcomm.com/developer/software/neural-processing-sdk-for-ai）（v2.46.0.260424）的 `*.so` 文件，并依赖于那些在剥离过程中幸存下来的未混淆名称、用 Ghidra（https://github.com/NationalSecurityAgency/ghidra）反编译的原始机器码，以及在我的 Linux 上进行的一些经验性参数扫描。 其中一些新颖的发现是（反正没人有耐心读完整个文章）： 1. HTP 将 VTCM 放置问题视为一个 MILP，并使用 HiGHS（https://highs.dev/）（优化求解器而非启发式算法）来求解——这在公开场合完全不为人知。 2. VTCM 放置使用一个在三维坐标空间中运行的递归回溯分配器。 3. 编译器可以在放置过程中自动改变权重精度（而你不自知）以缓解内存压力。 4. 即使不同的 SoC 报告相同的 `vtcmSize`，有效适配/溢出边界也取决于目标架构。 5. HTP 包含一个隐藏的分析模拟器，称为 Heximate，从中我们恢复了屋顶线方程和争用模型。 附录包含了更多我无法塞进正文的信息。 这就是要点，本文的其余部分将涵盖我发现的三件事，我认为这些内容从未在任何互联网上公开记录过，并且对任何愿意在高通 NPU 上进行边缘部署的人都会有所帮助。 ## 内存墙 Hexagon 芯片有一个小的片上暂存内存池，称为 VTCM（向量紧耦合内存）。另一边是 DDR，即主内存，但它很慢。ML 推理的显著瓶颈在于移动数据所花费的时间。编译器在此的整个工作就是决定每个时刻哪些数据可以放在 VTCM 上，因为任何放不下的东西都必须被换出到 DDR，稍后再取回来，这代价高昂且能效低下。 你的模型中的每个张量都有一个生命周期。在执行过程中的任何时刻，都有一组张量是活跃的（可称为工作集）。如果它适合 VTCM，那么一切都将很快。如果不适合，编译器将开始插入溢出操作（将张量推送到 DDR）和填充操作（将其拉回）。这就是我想找到的悬崖。 使用相同的模型（Qwen 0.8B），在 SM8350 上，编译器报告溢出了 5.46 MB，填充了 33.9 MB，总 DDR 为 37.9 MB；而在 SM8650（V75）上，没有任何溢出，DDR 读取为 1.15 MB。仅因目标芯片不同，DDR 读取流量就跳升了 33 倍（这是意料之中的）。但出于某种特殊原因，这些芯片在编译输出中报告了相同的 VTCM 大小——一个只写着 4 的字段。现在我不知道 4 是“什么”，或者它是不是某种代码。反正行为完全不同。我没有恢复实际容量，这是我下一步想做的事情。 编译后的二进制文件携带一个名为 `spillFillBufferSize` 的元数据字段，当它为 0 时表示模型权重完全在芯片内。这可以帮助任何人快速判断推理速度慢的因果关系。 现在我可以自信地把时间花在量化我的模型上，如果我的目标芯片是 SM8350，而不是猜测它了。 还有一件事决定你是否适配，事情在那里会变得更有趣一些。 ## 调度器与时间玩俄罗斯方块 芯片执行操作的顺序决定了每个张量存活多长时间，这决定了工作集有多大，进而决定了你是否撞上悬崖。因此，编译器的调度器的任务是找到正确的顺序，使工作集最小并保持在 VTCM 的范围内。 它用一种叫做“优先级 BFS”的东西来做这件事。它遍历图，测量该顺序所需的峰值 VTCM，然后进行 `peak_tcm <= tcm_capacity` 检查。之后返回 SMALL/LARGE，分别表示无溢出/有溢出。因此，填充/溢出决策取决于它找到的最佳顺序是否将峰值工作集保持在容量以下。底层的有序度量标准是一个操作在图中的深度优先拓扑排序位置，这有一个很好的特性，即尽可能将一个值的生产者和消费者在时间上靠近。当两个操作并列时，通过一系列启发式规则来打破平局，其中一些包括操作输入和输出之间的图距离、分支内的深度、先调度最重的分支等。 一旦顺序固定，编译器必须将每个张量物理放置在 VTCM 上。它首先对寿命最长的张量进行排序，然后是寿命短的张量，以减少碎片化。顺便说一句，放置不是简单的一维地址，代码为每个块分配了一个三维坐标（https://stackoverflow.com/questions/2192087/3-dimensional-bin-packing-algorithms）`(d0, d1, d2)` 在瓦片空间中，并且包装器（这是递归的）会回溯。确切的重试条件存在于几 KB 的内联 SIMD 中，我还没有尝试解码。 那个回溯包装器是回退路径。当编译器运行其完整优化器时，放置采取了更引人注目的形式，这是一个我之前在其他地方见过的正式优化问题。 到现在为止很明显，将张量放入 VTCM 与将手提箱装进一个小后备箱是同类问题，我们允许把一些手提箱留在家里，但稍后取回它们会付出代价。当可能时，编译器不会用启发式规则来猜测，而是将整个问题写成一个正式的混合整数线性规划，并交给 HiGHS——一个知名的开源优化包。 编译器将其处理的确切问题转储为标准求解器文件格式 `.mps` 用于调试。在那之后，我从反编译的二进制文件中发现了一些琐碎的东西： 1. 它最小化的目标是总移动量（溢出到 DDR 的字节数、从 DDR 填充回来的字节数，以及多核芯片上核心之间发送和接收的字节数总和）。 2. 同时活跃的两个张量不能占用相同的片上字节。 3. 要么一个张量在某个有效地址上存在于工作台，要么它被放逐到 DDR。它是半连续的。 你的模型的数值精度可以在你不知道的情况下被重写。有一些称为 `relaxed_precision_cast` 的操作，它们在 *放置过程中* 将张量在 float、FP16 和 BF16 之间转换，以缓解内存压力。因此现在我们知道编译器可以自行插入降级操作，你的 float32 张量可以变成 FP16，如果这样能装下字节的话。（我确认这些转换是在放置过程中插入的；但它们到底是求解器内部的变量还是一个单独的 pass，从我所提取的信息中无法判断。） 优化器中还有太多内容我无法在此详述。其中一些包括： 1. 图通过最小割进行分区，以决定在哪里切割模型以便溢出。 2. 一个独立、精巧的算法通过让生产者直接写入最终布局而非复制，使常见情况下的拼接免费。 3. 卷积通过 im2col 转换为矩阵乘法。 4. 通过校验和检查其内容来发现重复操作，并将其合并。 ## Hexagon + estimate = Hextimate (？？) 所有这一切——调度器寻找最小工作集、优化器最小化移动的字节数——都依赖于编译器实际上无法拥有的东西。为了在两个调度方案之间做出选择，或者决定是否值得避免溢出，它必须知道每个选项在时间上 *有多昂贵*。但芯片并不在那里。编译器正在我的 Linux 机器上运行，对要到模型发运到手机才会存在的硬件下数千个赌注。那么它如何定价呢？ 它进行猜测，并且它有一个专门为此构建的模拟器。它叫做 Heximate（Hexagon 估计，对吧？），编译器里到处都是对它的引用。它是一个整个芯片的小型分析模型，模拟运行你的模型并逐块累加成本，分为两个桶：计算和内存移动。 它还模拟 *争用*（当芯片的两个部分想要同一个资源而其中一个必须等待时会发生什么？），并且它不返回单个数字，而是返回一个 *范围*。它恰好运行一次工作负载，假设所有事情都能完美重叠，另一次假设没有任何重叠，并报告两端的结果。这应该使之成为模拟器，而不是查找表接口。 Heximate 的大部分形状我是从其内部组件的名称中拼凑出来的。但对于内存方面，我得到了实际的公式，直接来自机器码（感谢 Sonnet 4.6），这是每个性能工程师都知道的教科书式屋顶线模型： ``` 带宽 = 通道数 * 宽度 * 效率 * 频率 时间 = 字节数 / 带宽 ``` 该模拟器为某些事物保留 *单独* 的成本，而不是将它们混在一起，我认为这些相当琐碎。不过，我能读到这些单独成本因素的名称，但无法提取它们存储的实际数值。 1. 整数和浮点计算的定价是分开的（琐碎）。 2. 一次将一块数据写入多个地方（多播）的定价与一次写入一个地方不同（也很琐碎）。 3. KV 缓存张量和权重有自己的“快速 DDR”系数。 它还包含专门的检测器，可以识别 FlashAttention、MoE、KVCache、旋转位置编码等。 ## 结论 那么这有什么用呢？ 这次逆向工程尝试的关键要点，我认为对所有从事边缘 ML 的人都有用： 1. 它告诉自定义调度器的作者，放置质量受限于求解器，而不是启发式规则，并且目标是纯粹的字节流量。 2. 自定义调度器可以镜像 Heximate 的信念来预测编译器的决策。 3. 成本模型在系数层面偏向于 LLM 工作负载（因此可识别的 KV/权重张量是一个很大的加分项）。 4. 精度是求解器无论你要求与否都会拉动的旋钮，因此质量在实际硬件上可能有所不同。 5. `spillFillBufferSize` 可以用作静态适配/溢出预言机。 它是否直接帮助任何人，或者只是满足任何人的好奇心，这是另一个问题，但我得到了我在任何地方都找不到的答案。如果你在 Hexagon 上做过边缘工作并想交流心得——或者告诉我哪里错了——[我很乐意](mailto:[email protected])。 ## 备注 问题是该信任多少。这个文件（`libHtpPrepare.so`）是我们的唯一目标（x86_64 主机构建，QAIRT 2.46.0.260424，BuildID `63e60947ee8df89fe11592a8af12a30ddedb91cd`，sha1 `8048df5fe605d743a20337a6968aebc6f1930f4b`）。我之前的逆向工程尝试都是简单/中等的 crackme，而不是专有编译器，这一次只有借助 Claude Code 才可能完成。因此请对每一个声称持保留态度。我会等到弄清所有法律问题后再发布我获得这些见解的步骤。 1. 唯一我在任何 NPU 相关文献中都无法放置的发现是本文开头的那些。“在此 SDK 版本中未公开记录”是我可以坚持的声称。 2. 这只是单个二进制文件的一个 SDK 版本。这些数字（33 倍 DDR 跳升、5.46/33.9 MB 溢出/填充、`spillFillBufferSize` 的判断）在我的机器上使用高通自己的工具作为预言机是可重现的，但不同的 QAIRT 发布版可能会重新洗牌其中的任何内容。 3. 该二进制文件中充斥着像 `dp_tcm_threshold_selector`、`TCM_THRESH_75_predictor`、`special_mlp_features` 这样的名字，看起来就像是有个训练好的模型在选择何时溢出。但没有，至少在这个二进制文件里没有。我找不到任何权重矩阵、前向传递、任何像推理的东西。`special_mlp_features` 原来是一个将特征导出到 Python 的函数，即训练数据流出。成本路径是一个普通的 CSV 查找表。所以学习部分（如果存在的话）位于高通的训练时间上游，并以表格形式交付。我无法证明一个承载权重的模型在编译时运行，我倾向于否定它。