AI Infra 随笔

Portfolio avatar
上海交通大学

主要关注大模型推理优化、计算机系统以及并行计算。

67
Total posts
33
Active days
0
Day streak
  • 27 May

    GPU 通信拓扑实战

    5 min

    在 现代 AI 集群通信体系一文中我们分析了不同通信硬件、互连拓扑和网络协议栈,本文通过 nvidia-smi topo -m 展示的是服务器内部 GPU、NIC、CPU Socket、NUMA 与 PCIe/NVLink 的物理通信拓扑:本机中 GPU0-7 两两之间均为 NV18,说明 8 张 GPU 通过 NVLink/NVSwitch 形成了全互连高速 Fabric,因此节点内 GPU-GPU 通信通常直接走 NVLink,而不经过 CPU 与 NUMA hierarchy,非常适合 TP、EP、AllReduce 等高通信量并行;而跨节点通信时则需要关注 GPU 与 NIC 的 PCIe/NUMA 距离,此时 PIX > NODE > SYS 分别对应越来越远的 PCIe/NUMA 路径与更高通信代价,因此训练和推理框架(如 NCCL、Megatron、vLLM)通常会优先将 GPU 绑定到最近的 NIC,以减少 PCIe hierarchy 与跨 NUMA 带来的 latency 和 bandwidth loss。

  • 26 May

    现代 AI 集群通信体系

    这篇文章介绍了现代 AI / HPC 集群中的 GPU 通信体系,以及为什么“大模型训练”本质上越来越像一个网络系统问题。文章从单机内 GPU 互连开始,介绍 PCIe、NVLink 与 NVSwitch 的拓扑与数据流;再扩展到多机场景中的 RDMA、InfiniBand 与 RoCE,解释为什么跨机器通信会成为分布式训练与推理的核心瓶颈。

  • 19 May

    Python ZMQ:消息传递库

    2 min

    在 ZMQ 中,每个客户端(进程)都需要先加入 ZMQ 的消息系统,然后才能进行通信。对于每个客户端来说,通常需要完成三件事:创建自己的 ZMQ runtime,可以将其理解为一个轻量级的消息通信操作...

  • 19 May

    客户端间通信模式

    6 min

    客户端之间的消息通信,通常围绕三种经典模式构建:Request-Reply、Pub-Sub 与 Push-Pull。模式 核心思想 通信方式 典型关系 Request-Reply 我问,你答 同步 1...

  • 16 May

    图论基础

    8 min

    在未来将会被重写 图是离散数学中描述关系结构的核心模型,记作:G=(V,E)G = (V, E) 其中,VV 表示节点(Vertex)集合,EE 表示边(Edge)集合。节点描述对象,边描述对象之间...

  • 14 May

    信息论及其应用

    36 min

    信息论(information theory)研究的核心问题是消息在不确定性意义上带来了多少信息。如果一个事件本来几乎必然发生,那么它发生时不会带来多少新信息;如果一个事件本来极不可能发生,那么它一旦发生,就显著减少了我们对世界的无知。围绕这个直觉,信息论建立了一套用概率刻画信息、压缩、加密与通信可靠性的体系。本文对信息论的基本概念、压缩与信道编码、线性分组码、有限域、BCH 码与 Reed–Solomon 码进行总结。

  • 14 May

    估计与统计识别

    20 min

    估计与统计识别研究的核心问题是:当我们只能观察到有限的随机样本时,如何从这些样本中推断隐藏在背后的未知参数。换言之,我们并不知道真实世界的概率分布,但我们假设它属于某个参数化模型族,然后试图用数据把其...

  • 14 May

    CUDA Graph:合并多个算子

    2 min

    在 CUDA 编程模型中,CPU 不断向 GPU 提交 CUDA kernels,而一个 Transformer forward 则会进一步拆解成大量这样的 kernel 调用。即使单个 kernel 很快,CPU 发起 kernel launch 仍然存在固定开销,包括 runtime 调度、参数准备、driver 调用以及 GPU command queue 提交等。 在大模型推理中,模型结构通常是固定的,尤其是在 decode 阶段,执行路径会高度重复:每一层做相同的计算,只是输入数据不同。因此,大量 iteration 实际上是在重复提交几乎相同的一串 kernels。 CUDA Graph 的核心思想,就是把这一整段 GPU 执行流程提前录制下来并复用。录制过程中,CUDA runtime 会捕获 kernel 的执行顺序、依赖关系以及相关 memory operations,形成一个 executable graph。之后再次执行时,CPU 不再需要逐个 launch kernel,而是只需一次 graph launch,就可以让 GPU 按照预先记录好的 execution graph 执行整段计算。

  • 10 May

    博弈论(2):标准形博弈

    17 min

    在很多现实问题中,一个人的决策结果不仅取决于自己的选择,还取决于其他人的行为。例如,企业定价会受到竞争对手影响,国家政策会影响其他国家的回应,求职者是否接受 offer 也会受到公司策略变化的影响。在这类问题中,决策主体之间形成了相互依赖:每个人都在根据别人可能采取的行动来调整自己的选择。 传统的优化问题通常只研究给定环境下如何选择最优行动,而博弈论研究的问题则更进一步:环境本身也是由其他理性决策者共同构成的。因此,一个玩家面对的并不是固定世界,而是一个会随着所有人策略变化而不断变化的互动系统。 有限标准形博弈(finite normal-form game)正是对这类策略相互依赖问题的一种最基本抽象。它用参与者、可选策略以及不同策略组合对应的收益,来描述一个完整的战略互动结构。

  • 06 May

    大模型推理批处理和吞吐优化

    6 min

    通过 Roofline 模型 可以发现:Transformer 在不同阶段的性能瓶颈并不相同。对于 Prefill 阶段,由于 sequence length 较长,Attention 和 FFN 通常已经能够形成较大的矩阵乘法,因此 batch size 对 arithmetic intensity 的提升有限;而在 Decode 阶段,由于每次仅生成一个 token,FFN / Linear 层会退化为小规模 GEMV,算术强度显著下降,更容易受到 memory bandwidth 限制。此时增大 batch size 可以将 GEMV 重新转化为更高效的 GEMM,从而显著提升 GPU 利用率与整体吞吐。 与此同时,Attention 层在 Decode 阶段的 arithmetic intensity 与 batch size 基本无关,因此批处理对其性能提升有限,甚至可能因为 BMM 工作量增加而带来额外延迟。整体来看,LLM 推理中的 batching 收益主要来自 Decode 阶段 Dense Layer 的吞吐提升,而不是 Attention 本身。

  • 04 May

    Punica:Multi-LoRA 推理优化

    3 min

    Punica 将 Multi-LoRA 推理中按请求执行的 for-loop 重写为分段聚合的矩阵向量计算(SGMV),在一个融合 kernel 中同时处理多个 LoRA Adapter,从而避免小规模 GEMV 的频繁调度,显著提升 GPU 利用率。

  • 04 May

    LoRA 和 Multi-LoRA:模型参数微调

    4 min

    LoRA 在原模型权重上 WW 引入一个低秩结构的增量项(ΔW=BA\Delta W = BA),在不改变原参数的情况下,仅训练该低秩参数,从而在低维子空间中实现高效微调。

  • 04 May

    EPLB:MoE Expert 负载均衡

    4 min

    在 Mixture-of-Experts (MoE) 架构中,token 通过 gating 网络被动态路由到不同专家,导致 expert-level 的输入分布呈现明显的偏斜。当这些专家被分布到多个 GPU 上执行时,这种不均匀的 token 分布会直接转化为设备级的计算负载不均衡。EPLB 用于在推理阶段使不同设备/节点上的 MoE Expert 负载均衡。

  • 03 May

    Paged Attention:高效管理 KV Cache

    1 min

    TBD:未来将会增加更多细节 Paged Attention 使用类似虚拟内存分页的管理方式,使得 KV Cache 可以按需分配与回收,从而有效降低显存碎片化并提高内存利用率。PagedAtten...

  • 03 May

    Pipeline Parallel:模型并行

    6 min

    模型并行的核心思想是:不再让每张 GPU 都保存完整模型,而是把模型本身切分到不同 GPU 上。不同 GPU 负责模型的不同部分,并通过通信把中间结果传递给下一部分。换句话说,数据并行切的是数据,模型...

  • 03 May

    ZeRO 和 FSDP:DP优化——将模型参数和中间状态分片到多卡

    11 min

    在传统数据并行中,每个设备都维护完整的模型参数、梯度以及优化器状态,导致显著的内存冗余。 ZeRO(Zero Redundancy Optimizer)通过对这些状态进行分布式切分来消除冗余存储,在可控的通信开销下显著降低单设备的内存占用。 ZeRO 分为三个阶段(ZeRO-1 至 ZeRO-3),逐步对优化器状态、梯度以及模型参数进行切分。

  • 03 May

    C++ 基础(5):String

    2 min

    本文是对 C++ 基础(3):Streams 的应用。在工程里,如果需求只是不断把不同类型的数据拼接起来,最后得到一个完整字符串,那么最推荐使用 std::ostringstream。它的语义很明确...

  • 01 May

    基本数据结构:二叉堆和优先级队列

    5 min

    优先级队列通常基于二叉堆实现,支持在 O(1)O(1) 时间内访问当前最大/最小元素,并在 O(logn)O(\log n) 时间内完成插入和删除操作,相比于每次线性扫描 O(n)O(n),在需要频繁动态维护最值的场景中具有明显优势。

  • 01 May

    C++ 基础(4):函数和 Lambda

    2 min

    本文介绍了如何将函数作为可调用对象,并通过在创建可调用对象时捕获其上下文环境,更好地支持泛型编程。Represent functions as variables. Predicate: boolea...

  • 01 May

    Continuous Batching

    8 min

    Continuous batching 是一种面向在线推理的调度机制,在每个 decode step 动态组织当前活跃请求,从而避免 padding 和长尾带来的资源浪费,提高系统吞吐、提高 GPU 利用率.

  • 01 May

    在无外网远程机器上使用 HuggingFace 或 Codex

    4 min

    在某些环境中远程机器无法直接访问外网,导致无法在远程机器上下载 HuggingFace 模型、无法调用 OpenAI / Codex API,pip install 或 git clone 失败等等情况。我们希望能够让远程机器借用自己本地电脑的网络来对这些网站和服务器进行访问。 本文将介绍基于这个思路的解决方案和实现方式。本质上,我们通过 SSH 隧道把远程服务器的请求转运到本地,再由本地去访问互联网,最后把结果再返回给远程。

  • 01 May

    Roofline Model:从算术强度出发理解性能瓶颈

    6 min

    Roofline Model 用于衡量一个运算 / kernel 在特定计算平台(例如 GPU 或 TPU)上是计算受限(compute-bound)还是内存受限(memory-bound)。任何一个计算过程都同时受到计算能力内存带宽两方面的约束,最终性能由两者中较小值决定。Roofline 模型不仅可以判断算子/模型的 Arithmetic Intensity 是否过低(导致 memory-bound),更重要的是用于定位性能瓶颈,指导优化方向。

  • 28 Apr

    基本数据结构:栈和队列

    2 min

    查询:队列某一位置元素:[i] 队列头/尾元素:.front(), .back() 插入:在队列末尾插入(推荐!):.push_back(element) 在队列头部插入(推荐!):.push_fr...

  • 28 Apr

    基本数据结构:动态数组

    1 min

    内存分配和指针:查询:数组某一位置元素:[i] 数组头/尾元素:.front(), .back() 插入:在数组末尾插入(推荐!):.push_back(element) 在数组末尾插入并构造:.e...

  • 24 Apr

    C++ 常用 API 一图流

    1 min

    本章全部内容均来自:Cheat Sheets & Infographics std::string interface overview cheat sheet standard library se...

  • 24 Apr

    C++ 基础(3):Streams

    3 min

    在 C++ 中,数据可能来源于、或者想写入屏幕、文件或者字符串中。如何统一面对来源、去向以及数据类型完全不同的数据?本文简单介绍了 stream 作为统一的 I/O 抽象,stringstream 能...

  • 23 Apr

    基本数据结构:键值对和哈希表

    5 min

    哈希表 是用来高效存储和查找键值 (key, value) 对的数据结构,使得任意 dict 结构的 get, put, remove 方法平均时间复杂度都是 O(1)O(1). 哈希表的底层数据结构是数组。由于不同 key 可能映射到相同位置(哈希冲突),通常通过拉链法或开放寻址法来处理冲突,并通过优化哈希函数和控制负载因子(动态扩容)来降低冲突概率。为了保证哈希值的稳定性,key 通常要求是不可变对象。

  • 23 Apr

    基本数据结构:循环数组

    2 min

    在保持数组这种连续内存、高效访问结构的前提下,循环数组希望实现对队首和队尾的插入删除都能达到 O(1)O(1) 的效率。普通数组在头部删除或插入时,会在数组中留下“空位”,如果试图通过数据搬移来填补这些空...

  • 23 Apr

    数组操作技巧:差分数组

    3 min

    在某些情况下,我们需要频繁对原始数组的某个区间的元素进行同样数量的增减,并且我们希望最终获得数组中各元素的值。如果使用暴力方法,复杂度是 O(n)O(n),但是如果我们只记录元素之间的差(差分),那我们只...

  • 22 Apr

    Flash Decoding

    4 min

    FlashAttention 在训练和 prefill 阶段可以借助 batch 和 sequence length 提供充足的并行度,从而高效利用 GPU;但在 decode 阶段,每一步仅涉及单个 query,导致可并行的计算任务数量有限,从而难以填满 GPU,单 step 的 SM 利用率较低。 Flash Decoding利用 online softmax 中间状态 (m,l,O)(m,l,O) 的可结合性,将原本的串行递推过程重写为并行分块计算 + 规约合并:显著提升了 decode 阶段的 SM 利用率。

  • 21 Apr

    Flash Attention 2 (FA2)

    8 min

    Flash Attention V2 在 Flash Attention V1 的基础上,从计算模式与并行策略两个层面进行了共四点系统性优化。在算法层面,FA2 减少了非矩阵乘法运算(例如不必要的 rescale 和 normalize 等逐元素操作),使整体计算过程更加接近连续的 GEMM。同时,对于 causal masking,FA2 采用了更高效的处理方式,尽量避免在被 mask 的区域上进行无效计算,进一步提升了计算效率。 在并行化方面,FA2 沿 sequence length(即 QQ 维度)引入了更细粒度的并行划分,将原本较粗粒度的计算任务拆分并分配给更多的 thread block,从而显著提高了硬件并行度与资源利用率。在 thread block 内部,FA2 重新设计了 warp 的工作分配方式,使不同 warp 分别负责不同的 QQ 子块(对应不同输出行),从而避免多个 warp 同时对同一输出进行并行写入且避免了对中间结果的合并需求,显著降低了 warp 间的同步与通信开销。

  • 21 Apr

    Ring Attention:Attention Block 的序列并行

    3 min

    在长上下文场景下,序列并行(Sequence Parallelism)通常用于切分 seq_len 维度。然而 Attention 计算具有全局依赖性:每个 Query 都需要访问完整的 KV 序列,这使得序列维度上的直接切分变得困难。同时,在长序列场景下,单个设备往往无法容纳完整的 KV Cache,因此需要对 KV 进行分布式存储。在序列并行下,需要先 AllGather 从所有设备收集完整 KV Cache,再在本地算局部 Q shard 的 attention 计算。 为了解决这一问题,Ring Attention 在 Attention 计算过程中将 AllGather 操作替代为使用 ring 通信机制 在设备之间轮转 KV 块,采用 blockwise / online softmax 的方式完成 Attention 计算。

  • 20 Apr

    Online Softmax 推导

    7 min

    Online softmax 通过维护可增量更新的中间状态 (m,l)(m,l),将原本需要全局两次归约(max 与 sum)的计算转化为可流式处理的形式,从而避免额外遍历,并使 softmax 支持分块计算与归并。

  • 19 Apr

    Chunked Prefill

    2 min

    Chunked Prefill 通过将原本规模为 O(S2)O(S^2) 的 Causal Attention 计算重排为若干较小的分块计算,从而降低 Prefill 阶段的峰值显存占用,以及使得调度单元变得更加细粒度。

  • 19 Apr

    DeepSpeed Ulysses:Attention Block 的序列并行

    2 min

    SP 因为 Attention 计算在序列维度 ss 上的数据依赖而变得困难,DeepSpeed Ulysses 通过在 Attention 计算之前引入 All-to-All 通信来将原本按 sequence 切分的数据重排为按 head 切分,使得每个设备能够在本地对完整序列、但仅部分 attention heads 执行计算,从而实现高效的分布式并行。

  • 18 Apr

    C++ 基础(2):STL Containers

    3 min

    本文首先介绍了 C++ STL 常用的容器,分别介绍了对应 Python 中的 list(), dict() 的容器,并且对于各容器的优缺点进行了总结;然后介绍了遍历容器的迭代器 Iterator,最后略微提及了可以指向所有类型的指针。

  • 18 Apr

    C++ 基础(1):Types, Structs, Structured Bindings

    4 min

    在很多情况下(例如函数需要返回多个结果时),我们希望将多个相关的值组合为一个整体进行返回。C++ 提供了多种方式来对多个值进行建模与打包,例如使用 structstd::pairstd::tuple 将多个值组合成一个类型。在调用端,可以通过 structured binding 对这些聚合类型进行高效而清晰的解包,从而方便地获取各个返回值。

  • 24 Mar

    Git 基础 (1):Git 对象和引用

    8 min

    在软件开发中,一个项目从最初的几行代码,到逐渐演化成复杂系统,整个开发过程本质上是一系列变化的积累。但问题在于,变化本身是很难管理的。我们如何描述项目状态的变化?如果没有工具,我们只能:手动复制文件做...

  • 24 Mar

    MoE 模型基础

    5 min

    Mixture of Experts (MoE) 模型:将传统 FFN 替代为路由器 + 很多稀疏的 FFN,也叫 Experts 每次推理过程中选择 Experts 的子集 优势:在相同的 FLO...

  • 10 Mar

    集合通信操作和其实现

    10 min

    集合通信(Collective Communication) 描述的是多个设备之间如何协同交换和聚合数据,其核心操作包括:Broadcast(一个设备向所有设备扩散数据)、Gather/Scatter(数据的收集与分发)、Reduce(对多个设备上的数据执行求和等归约计算)以及它们的 All 版本(所有设备最终都获得结果)。这些操作本质上是在回答两个问题:数据该流向哪里(一个 / 所有),以及 数据在流动过程中是否需要被计算(Reduce)。在实际系统中,这些逻辑语义通常通过 Ring 通信实现:所有设备首尾相连形成环,每一步只与左右邻居通信,从而形成高带宽、可流水化的数据传输。基于这种思想,AllReduce 可以进一步分解为 ReduceScatter + AllGather:先让不同设备分别完成不同数据分片的归约计算,再将结果重新同步给所有设备。这种分解避免了单点聚合带来的瓶颈,也是现代分布式训练(如 DDP、ZeRO、Megatron 等)中的核心通信基础。对于 pp 个设备和大小为 NN 的 tensor,Ring AllReduce 的总通信量约为 2p1pN2N2\frac{p-1}{p}N \approx 2N,这也是为什么 AllReduce 常被认为需要“约两倍 tensor 大小”的通信开销。

  • 04 Mar

    Flash Attention (FA1)

    10 min

    FlashAttention 1 相较于标准 Attention 计算其更快的主要原因在于它是 IO-aware 的。其将注意力计算划分为能够放入 GPU 片上 SRAM(共享内存)的小块(tile),并通过 online softmax 逐步计算 softmax。中间结果始终保存在更快的 SRAM 中,而不会被写回 HBM。每个块的最终输出只在计算完成后写回一次 HBM。HBM 的数据访问量从原来的 O(n2)O(n^2) 降低到 O(nd)O(n \cdot d), 其中 dd 是每个 head 的维度,达到了读取输入和写入输出所需的理论最小值。在 NVIDIA A100 上,对于长序列,FlashAttention 相比标准注意力通常可以获得 2–4 倍的加速。

  • 04 Mar

    CUDA Grid, Block, Thread 和寻址

    4 min

    本文从数据并行计算的角度出发,介绍 CUDA 中 Grid、Block 和 Thread 的分层组织方式,并解释如何通过线程索引将计算任务映射到具体数据。