张芷铭的个人博客

端到端吞吐优化三项核心技术：算子融合、多模态GEMM调度、显存/激活管理，分别从不同维度突破瓶颈。

Adaln Varlen Fuse（自适应层归一化 + 变长序列融合）

核心原理

组件作用
AdalnLayerNorm + 可学习缩放/偏移，适配多模态输入分布差异
Varlen变长序列支持，避免padding计算浪费
Fuse合并Adaln + 激活 + 残差 + 变长mask为单内核，消除访存开销

收益：E2E吞吐+10%

实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

class AdaLayerNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(1, 1, dim))
        self.shift = nn.Parameter(torch.zeros(1, 1, dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x = (x - mean) / torch.sqrt(var + self.eps)
        return x * self.scale + self.shift
  • 变长处理：pack_padded_sequence 或 masked计算
  • 算子融合：torch.compile、FlashAttention-2、Triton自动融合

Modality-Aware Group GEMM（模态感知分组矩阵乘法）

核心原理

技术说明
Group GEMM多个小GEMM打包为大GEMM并行执行，解决小矩阵效率低问题
Modality-Aware按模态分组，每组使用专属分块策略、张量核心利用率

收益：E2E吞吐+20%（多模态MoE场景显著）

实现

1
2
3
4
5
6
7

# 传统：N次小GEMM（低效）
outputs = [torch.matmul(x, e.weight) for e in experts]

# Group GEMM：1次大GEMM（高效）
group_weights = torch.cat([e.weight for e in experts], dim=1)
group_out = torch.matmul(x, group_weights)
outputs = group_out.split(d_out, dim=-1)
  • 算子库：DeepGEMM、CUTLASS GroupGEMM、Megatron-LM
  • 配置：--moe-grouped-gemm，按模态设置block_m、block_n

Multi-Stream Activation Offload（多流激活卸载）

核心原理

技术作用
Gradient Checkpoint只保存检查点，反向重计算激活，显存节省30%-60%
Activation Offload激活值卸载到CPU内存，释放GPU显存
Multi-Stream计算流 + 传输流并行，抵消Offload时间开销

收益：E2E吞吐+33%

实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def forward(self, x):
        return checkpoint(self._forward, x, use_reentrant=False)

    def _forward(self, x):
        x = x + self.attn(self.adaln1(x))
        x = x + self.ffn(self.adaln2(x))
        return x

1
2
3
4
5
6
7
8
9

# 多流重叠
stream_compute = torch.cuda.Stream()
stream_transfer = torch.cuda.Stream()

with torch.cuda.stream(stream_compute):
    out = self.block(x)
with torch.cuda.stream(stream_transfer):
    checkpoint_cpu = out.to('cpu', non_blocking=True)
torch.cuda.synchronize()

技术对比

技术优化点收益适用场景难度
Adaln Varlen Fuse算子融合 + 变长计算+10%多模态Transformer中
Modality-Aware Group GEMM分组GEMM + 模态调度+20%MoE、多分支模型中高
Multi-Stream Activation Offload重计算 + 卸载 + 多流+33%大模型显存受限中

落地优先级

  1. 先开Activation Offload：框架原生支持，见效最快
  2. 再做Group GEMM：多模态/MoE必选
  3. 最后优化Adaln Varlen Fuse：算子级极致优化

避坑要点

技术注意事项
Adaln Fuse融合后校准数值稳定性
Group GEMM分组过大导致显存峰值，按GPU调整group_size
Activation Offload设置non_blocking=True，平衡重计算与Batch

Comments