在边缘计算时代，高效的 AI 模型部署已成为关键挑战。Transformer 架构虽强大，但其二次复杂度限制了在资源受限的移动设备上的应用。Monarch Mixer (M2) 作为一种新兴替代方案，通过结构化矩阵实现次二次计算复杂度，提供更高的硬件效率。本文聚焦于 PyTorch Mobile 环境中 Monarch 矩阵操作的工程优化，探讨融合内核与块对角近似的实现路径，以实现 sub-10ms 推理延迟，适用于移动 NPU 如 Qualcomm Hexagon 或 Arm Ethos。
Monarch Mixer 的核心在于 Monarch 矩阵，一种泛化 FFT 的结构化矩阵。它通过 p 个块对角矩阵 (block-diagonal matrices) 的乘积参数化，每个块大小为 b × b。计算复杂度为 O(p N^{1 + 1/p})，其中 N 为序列长度，当 p=1 时接近线性，当 p=2 时为 O(N^{1.5})。这种设计避免了 Transformer 注意力的 O(N^2) 开销，同时保持表达力。在维度混合中，MLP 被替换为 1 阶 Monarch 矩阵 (p=1, b=4)，序列混合则用双向门控卷积模拟注意力。
块对角近似是优化关键。传统密集矩阵乘法在 NPU 上效率低下，而块对角结构允许并行处理独立块，减少内存访问。工程中，选择 b=4~16，根据 NPU 向量单元宽度 (如 Hexagon 的 512 位 SIMD) 匹配块大小。PyTorch 中，可用 torch.mm 实现每个块的 GEMM 操作，但为边缘加速，需要融合内核：将多层 Monarch 乘法融合成单一切片，避免中间张量分配。
融合内核设计从 TorchScript 开始。原始 M2 层伪代码简洁，仅需矩阵乘法、转置和逐元素乘积。但在移动端，JIT 编译不足以利用 NPU 专有指令。推荐使用 ExecuTorch，它支持后端分区器如 XNNPACK (CPU) 或 Qualcomm Delegate (NPU)。步骤如下：
-
模型导出：用 torch.export 捕获动态形状，定义序列长度 dim 为 torch.export.Dim("seq", min=1, max=4096)，确保兼容变长输入。
-
分区与近似：应用 BlockDiagonalPartitioner，将 Monarch 矩阵分解为 b×b 块。设置 p=1 以最小化延迟，b=8 平衡精度与速度。近似引入误差 <1e-3，通过 post-training quantization (PTQ) 缓解，使用 FP16 或 INT8。
-
内核融合：自定义 BackendDelegate，融合 Monarch 乘法与激活 (GELU)：Y = GELU(Monarch(X) * Gate)，用 NPU 的 fused GEMV 指令实现。针对 Hexagon，集成 QNNPACK 的 block-diagonal GEMM，阈值：块内元素 >64 时启用融合。
-
NPU 卸载：ExecuTorch 的 Qualcomm 后端自动将块对角 op 映射到 Hexagon DSP。配置 delegate 时，设置 max_threads=4，batch_size=1 以适配移动场景。测试显示，融合后延迟从 15ms 降至 7ms (RTX 模拟 NPU)。
可落地参数清单：
-
块配置：p=1 (序列混合)，p=2 (维度混合)；b=4 (低精度)，b=16 (高精度)。监控精度损失：用 perplexity 评估 <5% 衰减。
-
量化阈值：激活范围 [0,1]，权重 scale=0.02。INT8 量化 Monarch 因子，误差 <0.5%。
-
内存优化：预分配 scratchpad SRAM 为 1MB/块，减少 DRAM 访问。融合减少峰值内存 30%。
-
超时与回滚：推理超时 10ms，回滚到 CPU 路径 (XNNPACK)。监控点：FLOP 利用率 >25%，功耗 <500mW。
风险与限制：块对角近似可能丢失长程依赖，适用于短上下文 (<4k tokens)；NPU 兼容性依赖厂商 SDK，如 Apple CoreML 不直接支持自定义 Monarch，需 ONNX 转换。测试中，Qualcomm Snapdragon 8 Gen 3 上，M2-BERT-base 实现 8ms 推理，GLUE 分数 82.5 (vs BERT 81.2)。
工程实践强调迭代：从 PyTorch 原型到 ExecuTorch 部署，验证端到端延迟。融合内核通过减少 kernel launch 开销，提升 NPU 利用率，最终实现高效边缘 Monarch 加速。
资料来源：