1. 张量运算融合技术概述

在深度学习和大规模数值计算领域，张量运算（Tensor Operations）构成了计算图的核心骨架。随着模型规模的不断扩大，如何高效执行这些运算成为系统优化的关键挑战。传统方法中，每个运算独立执行，导致大量中间结果需要在内存和计算单元之间频繁搬运，这种数据移动往往成为性能瓶颈。

Einsum（爱因斯坦求和约定）作为一种通用的张量运算表示法，能够简洁地描述矩阵乘法、转置、收缩等各种线性代数操作。例如，矩阵乘法可以表示为 ij,jk->ik ，向量点积则是 i,i-> 。这种表示法的优势在于：

• 统一性：可以表达绝大多数张量运算
• 显式性：清晰地展示了输入输出张量维度的变换关系
• 可组合性：多个Einsum可以串联形成计算链

在实际应用中，一个典型的Transformer层可能包含数十个Einsum操作，这些操作之间往往存在数据依赖关系，形成了所谓的"计算瀑布"（Compute Cascade）。

2. 贪婪融合策略的核心思想

2.1 迭代空间分析基础

贪婪融合策略（Greedy Stitching）的核心在于分析相邻Einsum操作的迭代空间（Iteration Space）关系。迭代空间是指执行某个运算时需要遍历的所有维度组合，例如矩阵乘法 [M,N] x [N,P] -> [M,P] 的迭代空间是 [M,N,P] 。

两个Einsum能否融合的关键指标是它们的迭代空间交集：

• 完全一致（Rank-Isomorphic, RI）：可以完美融合
• 子集关系（Rank-Subsetted, RSb）：部分融合可能
• 超集关系（Rank-Supersetted, RSp）：特定条件下可融合
• 无交集（Rank-Disjoint, RD）：难以直接融合

2.2 融合算法实现细节

算法1展示了贪婪融合的伪代码实现，其核心流程如下：

1. 初始化当前融合组为空，将前两个Einsum加入组中
2. 计算它们的迭代空间交集I_prev
3. 对于后续每个Einsum：
   • 计算其与当前组最后一个Einsum的迭代空间交集I_curr
   • 如果I_curr与I_prev满足RI/RSb/RSp关系，则加入当前组
   • 否则结束当前组，从该Einsum开始新组
4. 最终输出所有融合组列表

# 简化版的贪婪融合实现示例
def greedy_stitching(cascade):
    fusion_groups = []
    current_group = [cascade[0], cascade[1]]
    I_prev = intersect(cascade[0].iter_space, cascade[1].iter_space)
    
    for i in range(2, len(cascade)):
        curr_einsum = cascade[i]
        I_curr = intersect(current_group[-1].iter_space, curr_einsum.iter_space)
        
        if is_ri(I_prev, I_curr) or is_rsb(I_prev, I_curr) or is_rsp(I_prev, I_curr):
            current_group.append(curr_einsum)
            I_prev = I_curr
        else:
            fusion_groups.append(current_group)
            current_group = [curr_einsum]
            I_prev = curr_einsum.iter_space
    
    if current_group:
        fusion_groups.append(current_group)
    return fusion_groups

2.3 融合类型的具体表现

图8展示了一个五Einsum级联的融合实例：

• E1-E3形成第一个融合组（蓝色）
  • E1->E2: 保留[M,N]
  • E2->E3: 保留[M,N,P]
• E4-E5形成第二个融合组（黄色）
  • E4->E5: 仅保留[N]

这种分组使得：

1. 组内张量可以保持在芯片上，避免频繁访存
2. 循环结构可以共享，减少控制开销
3. 数据局部性得到最大化利用

3. Mamba加速器中的融合实践

3.1 Mamba计算图特点

Mamba模型的计算图（图1）具有以下特征：

• 高比例的低强度操作（如SiLU、exp等）
• 复杂的张量形状变化
• 长距离的数据依赖链
• 迭代性状态更新（SSM部分）

这些特点使得传统的逐操作执行方式效率低下，而融合技术可以带来显著提升。

3.2 四级融合策略演进

Mamba加速器实现了渐进式的融合策略：

3.2.1 RI-only融合

仅融合迭代空间完全相同的Einsum：

• 将24个独立Einsum减少到12个融合组
• 特别适合SSM部分（E16-E21）
• 在token生成阶段表现最佳（2.23×加速）

3.2.2 RI+RSb融合

增加子集关系的融合：

• 融合组减少到8个
• 实现GEMM后接element-wise的融合（E14-E15）
• 预填充阶段提升至2.99×

3.2.3 RI+RSb+RSp融合

进一步加入超集关系融合：

• 融合组降至3个
• 解决NEX到TX的融合（E5-E6）
• 预填充阶段达3.35×

3.2.4 完全融合

通过分块技术实现近似RD融合：

• 形成单一融合组
• 采用tile-based流水线执行
• 预填充阶段最高4.9×加速

3.3 架构支持关键设计

Mambalaya加速器的创新架构设计使能了上述融合策略：

1. 可重构PE阵列 ：

   • 2D模式（256×256 PE）：处理GEMM等高强度运算
   • 1D模式（8192 PE）：优化低强度element-wise操作

2. 混合计算单元 ：

   • 每个PE包含MAC和特殊函数单元（log/max/SiLU/exp）
   • 6级流水线设计，每周期完成一个操作

3. 层次化存储 ：

   • 32MB全局缓存
   • 分布式寄存器文件（总计4.25MB）
   • 2039GB/s内存带宽（匹配H100）

表3对比了Mambalaya与H100的关键参数，在相同工艺下实现面积等效。

4. 实现细节与优化技巧

4.1 共享输入张量合并

在融合前对计算图进行代数变换：

1. 识别共享输入张量（如X用于生成B、C和TTΔ）
2. 将多个消费者合并为单个多输出Einsum
3. 示例：

   # 原始分开计算
   B = einsum('...', X, Wb)
   C = einsum('...', X, Wc)
   TTΔ = einsum('...', X, Wt)
   
   # 合并后
   B, C, TTΔ = fused_einsum('...', X, [Wb, Wc, Wt])
   

4.2 分块策略选择

对于包含迭代秩（如H状态张量）的Einsum：

• 沿I秩分块：平衡片上存储需求
• B-D-N静态数据流：最小化中间状态
• 混合策略：
  • 预填充阶段：优先I分块
  • Token生成：B-D-N静态更优

4.3 流水线调度

完全融合模式下的关键优化：

1. 上游Einsum产生部分结果后立即触发下游
2. 避免等待完整张量就绪
3. 示例（图8c）：
   • X的Q-fiber就绪后立即开始V计算
   • 重叠计算与数据传输

5. 性能评估与对比

5.1 实验设置

• 模型：mamba-370m和mamba-2.8b
• Batch size：64
• 上下文长度：1（生成）到220（预填充）
• 对比基线：
  • MARCA-like：仅SSM部分RI融合
  • Geens-like：细粒度分块融合

5.2 关键结果

图12展示了不同场景下的加速比：

• 小上下文长生成：RI最优（2.23×）
• 大上下文短生成：完全融合最优（4.9×）
• 中等情况：RI+RSb+RSp平衡（3.35×）

图13显示相比SOTA：

• 超越MARCA-like 4.9×
• 超越Geens-like 1.5×

5.3 流量分析

图14揭示：

• 完全融合减少inter-Einsum流量34×
• 但intra-Einsum流量增加（部分乘积）
• 其他策略接近算法最小访问量

5.4 计算利用率

图15的Roofline分析表明：

• 完全融合使传统内存受限操作（如SSM）达到峰值算力
• RI+RSb+RSp在多数场景接近理想利用率
• 基线方案存在明显"屋顶"距离

6. 实际应用建议

6.1 策略选择指南

场景特征	推荐策略	预期加速
长生成（I=1）	RI-only	2-2.5×
短生成中上下文	RI+RSb	3×
长上下文短生成	完全融合	4-5×
内存受限系统	RI+RSb+RSp	3-3.5×

6.2 常见问题解决

1. 融合组过大导致寄存器压力 ：

   • 采用更激进的分块
   • 部分中间结果写回内存
   • 示例：RX（E8）显式offload

2. 控制流复杂化 ：

   • 使用predicated execution
   • 维护显式依赖标记
   • 静态调度与动态触发结合

3. 特殊函数单元冲突 ：

   • 交错安排计算类型
   • 使用PE阵列子集执行低强度操作

6.3 扩展应用方向

1. 生成式场景：

   • 动态调整融合策略（KV缓存增长时）

2. 多模态模型：

   • 跨模态Einsum的融合规则

3. 稀疏张量：

   • 结合COO/CSR格式的融合条件判断

在真实部署中，建议通过Timeloop等工具预先评估不同融合策略在目标硬件上的表现，建立策略选择器。对于Mamba类模型，通常可以观察到：

• 前50%层更适合激进融合
• 后50%层需要保守策略
• 注意力部分需要特殊处理