一片自留地 我理解的 ai 性能优化

magicyang · 2025年04月23日 · 566 次阅读

1. 理论:

从本质上来说就是并行化。
将一个相同的任务拆解成 n 个子任务,而 n 个子任务相互不影响的时候,那么就有并行化的意义。

并行化通常是硬件、软件、算法相互结合的结果。
其中硬件包括:nvidia 的 simt,cpu 的 simd 指令。
软件则包括:框架调度、程序调度等。
算法则是:如果将串行任务进行拆解。

为什么并行计算能起来?
因为发现神经网络的相关任务,是一个非常容易通过并行化拆解,从而提升执行速度的任务。
随着 ai 的发展,计算量的越来越大,相关的行业也迅速发展了起来。

2. 瓶颈

在目前冯诺伊曼的计算机体系结构下,计算通常会遇到三个瓶颈:
(1)计算量
(2)io
(3)通讯
在 ai 计算的过程中,前两个瓶颈会贯穿在所有的计算当中。
这里通常会引入计算密度 ci 以及 roofline 模型来描述算法的计算密度和 io 在不同硬件体系下的适配情况。

对于通讯,尤其是在 llm 的背景下,受制于模型大小,单算子或者是切片模型(tp)只依靠单卡已经解决不了问题,那么就会涉及到卡间通讯和机器间通讯,这也是目前 nvidia 发力的重点,这部分内容因为我受制于资源规模,实际没有相关经验。这里只做简单描述。

目前对于计算、io 都是如何解决的呢?

2.1 计算

首先看计算,nvidia 从 pascal 架构后,专门引入了 tensorcore 的硬件结构,专门解决卷积、矩阵乘这类高 ci 的问题。
底层提供了 ptx mma/wmma 的矩阵乘指令,对于 fp16,bf16,int8,fp8 等提供了 10x 的硬件加速。
我个人理解,从 pascal 架构到现在的架构,这一块虽然一直在进化,允许更大、更快的尺寸指令去执行,但是底层逻辑并没有大的改变。

2.2 io

实际上,我们真正能掌控的就是 io。
如果熟悉体系结构的话,会知道目前我们的内存体系,分为了 l1,l2,gpu smem,gpu gmem, cpu mem 多层结构。
不同的内存结构间,其传输效率,带宽都不同。

实际操作中,cpu->global 的操作目前通常只会在初始化中进行。这里 nvidia 使用了 hbm 的内存,带宽会比 ddr 的高不少,这部分通常是可操控较少的部分。
shared mem 的性能相较于 global mem 有一个数量级的优势。
因此大部分优化都是基于 shared mem 进行的优化。
比如常用的:

  1. flash-attn 就是充分利用单个 sm 的 shared mem,进行的 tensor 切片运算。
  2. 比如计算过程中的 bank conflict,也是基于 shared mem,尽量降低 l2 cache miss。

除了这两部分,系统层面的优化还包括:

  1. 通过 fuse-ops,去降低 tensor->tensor 之间的 global mem 吞吐,在 kernel 内部,reg->reg 的访问是最快速的。(有代价,当 kernel reg 不够用的时候,会利用 global mem 临时申请,会导致更糟糕的结果。)
  2. 通过 graph 的模式,去降低频繁 kernel launch 加载所带来的影响。
  3. 通过 cp 和计算的异步流水化,将 cp 带来的同步影响消除。 cp 异步也是和硬件强相关的。
  4. 量化 int8/int4/fp8 的量化。其中 int4 的代价有些大,量化一定是会影响精度的。 训练/感知量化相对影响小,但是需要数据集和训练成本,这条路在 llm 基本走不通。后量化的代价也些大,如何取舍也是个问题。

3. llm 优化

这块我没有切实做个大集群的优化。
从我的理解看,dp,tp,pp 的部分已经是相当成熟。ep 我不太熟悉。
我主要做推理。
在系统缓存层面,如 kv cache,page attention,prefix caching 都已经集成到了开源框架中。
在小模型 + 大模型优化中,诸如 spec decoding 的方法也已合入。
想超越开源,提前开源,我暂时还做不到。
没有集群经验,不想写算子的我,也一直还没有找到合适的机会。

写在最后;
所有的优化都会有代价,为了速度快,一定会牺牲很大程度上的代码复杂度。
你让业务去根据你的要求去写那种复杂代码,排斥的心理会非常的大。
所以需要一个合理且重要的需求逼着你必须这么去做,这一直都是性能优化面临的最大问题。

暂无回复。
需要 登录 后方可回复, 如果你还没有账号请点击这里 注册