AI Optimization Stack

导言

“模型优化”“框架优化”“算子优化”和“软硬协同”经常出现在同一段讨论里，却不一定在回答同一个问题。本文建立一张七层诊断地图：先确认被控制的对象、真正变化的物理量和最后采用的验收指标，再判断优化属于哪一层。这个分层不是唯一的行业标准，而是一套避免跨层归因错误的工作方法。

同一次大模型训练或推理，可能同时涉及模型结构、并行策略、请求调度、计算图、Kernel、通信库和加速器。它们最终共同决定成本与性能，因此都能被称为“AI 模型优化”；但如果不先分清层次，很容易得到几个看似合理、实际上不可验证的结论：

Kernel 快了，模型就一定同比例变快。
用了 GPU、NPU 或 RDMA，就是软硬协同。
训练优化只看算力，推理优化只看带宽。
框架机制可以不考虑具体硬件。

更可靠的提问方式是：

对象是什么？ 是模型语义、参数与梯度、请求与 KV Cache、计算图、张量 Tile、通信消息，还是硬件资源？
物理量变了什么？ 是 FLOPs、HBM 读写、峰值显存、网络字节、Kernel 启动、排队时间，还是设备能力？
用什么验收？ 是精度、收敛、tokens/s、TTFT、TPOT、P99、单算子时延，还是能效？

AI 优化层次的七层诊断地图 — 七层诊断地图。每层都有自己的控制对象、瓶颈与验收闭环；右侧的软硬协同贯穿多层，而不是额外的第八层。图为本文自绘。

七个优化层次¶

这里将广义的 AI 模型优化拆成七层。层与层之间不是互斥的组织边界，而是定位问题时的主要责任边界：一个改动可以跨层实现，但必须能说明每一层分别改了什么。

层次	主要控制对象	常见物理瓶颈	主要验收指标
模型与算法	目标函数、网络结构、数据与数值语义	有效计算量、样本效率、表达能力	质量、收敛速度、鲁棒性
训练系统	参数、梯度、优化器状态、激活、并行 Rank	设备利用率、显存峰值、同步等待	tokens/s、step time、可训练规模
推理服务	请求、Batch、KV Cache、Prefill 与 Decode	排队、KV 容量、尾时延、空泡	TTFT、TPOT/ITL、吞吐、P99、成本
编译与运行时	计算图、形状特化、内存计划、执行流	Graph Break、重编译、调度与启动开销	编译时间、命中率、端到端时延
算子与 Kernel	单个算子、Tile、Layout、融合边界	HBM 流量、算术强度、并行度	单算子时延、带宽或算力利用率
通信与系统软件	Collective、Buffer、Transport、Topology	链路带宽、消息时延、同步与拥塞	有效带宽、Collective 时延、Overlap
硬件	计算单元、存储层级、互连与指令能力	峰值能力、容量、功耗、拓扑约束	峰值规格、能效、可实现上限

模型与算法¶

模型与算法层决定“算什么”。它可以改变模型结构、训练目标、数据使用方式和数值表示，也可能改变最终可达到的质量。例如，用稀疏结构降低每个 Token 激活的参数量，或用蒸馏将教师模型的能力迁移到更小模型，首先都是算法层问题。

这一层最容易出现不公平比较。若新方法同时用了更多数据、更长训练时间或不同质量阈值，就不能只比较最终吞吐或参数量。至少要固定任务、数据预算、计算预算与质量门槛中的关键条件。

性能不等于质量

端到端变快并不自动证明模型更优；模型质量提高也不自动证明系统效率提高。算法实验与系统实验可以共享一次运行，但必须保留两套验收指标。

训练系统¶

训练系统层决定“怎么把前向、反向和参数更新组织起来”。它管理参数、梯度、优化器状态和激活的切分与生命周期，并安排数据并行、张量并行、流水线并行、上下文并行或专家并行之间的计算与通信。

以 PyTorch FSDP 的 FULL_SHARD 为例，参数在计算前按需 AllGather，梯度在反向后 ReduceScatter，并在空闲阶段重新分片；优化对象不是注意力公式，而是分布式状态的驻留方式和通信时机。¹

训练系统优化通常关心：

可训练规模： 模型与状态是否能放进给定设备；
稳态吞吐： tokens/s、samples/s 或 step time；
扩展效率： 增加设备后，有效吞吐是否接近理想比例；
语义等价： 换并行策略、重计算或混合精度后，损失曲线和最终质量是否仍在可接受范围内。

训练并非天然“计算受限”。小 Batch、细粒度流水线、稀疏专家路由或跨节点并行都可能让通信、内存容量和调度空泡成为主导瓶颈。

推理服务¶

推理服务层决定“请求如何进入设备并持续生成 Token”。它没有训练中的反向和优化器状态，却多了请求到达、排队、Prefill、Decode、抢占和 KV Cache 生命周期。

其核心不是固定 Batch 的一次前向，而是在不断变化的请求集合中平衡：

首 Token 时延 TTFT： 用户多久看到第一次响应；
Token 间时延 TPOT/ITL： 连续生成是否流畅；
吞吐： 单位时间完成多少请求或生成多少 Token；
尾时延与成本： P99 是否满足 SLO，每百万 Token 需要多少资源。

vLLM 的调度配置就同时涉及 Token Budget、KV Cache 容量、Prefill/Decode 策略和异步调度；官方指标也把 TTFT、Inter-Token Latency、队列和吞吐分开暴露。² ³

维度	训练系统	推理服务
主要阶段	Forward、Backward、Optimizer Step	Prefill、Decode
长驻状态	参数、梯度、优化器、激活	参数、KV Cache、请求状态
调度核心	并行 Rank、Micro-batch、通信依赖	请求、Token Budget、抢占与 KV 配额
主要指标	step time、tokens/s、扩展效率、峰值显存	TTFT、TPOT/ITL、吞吐、P99、成本
常见风险	收敛漂移、同步等待、流水线空泡	头阻塞、尾时延、KV 碎片或容量耗尽

因此，训练框架与推理框架都在提高设备利用率，但它们控制的状态机和目标函数并不相同。

编译与运行时¶

编译与运行时层决定“高层程序如何变成可执行图和执行序列”。典型工作包括图捕获、算子分解、形状特化、融合、内存规划、代码生成和执行流调度。

例如 torch.compile 会捕获可编译区域，交给后端生成并缓存执行产物；Guard 失败、Graph Break 或动态形状变化可能触发回退或重新编译。⁴ 这类问题不能简单归为“某个 Kernel 慢”，因为瓶颈可能发生在图边界、编译缓存或运行时调度上。

验收时既要看稳态执行，也要看：

首次编译成本与冷启动；
不同 Shape 下的缓存命中和重编译次数；
图覆盖率、Graph Break 与回退路径；
编译后峰值内存和端到端时延。

算子与 Kernel¶

算子与 Kernel 层决定“一个局部计算如何映射到设备”。这里会调整 Tile 大小、线程或任务分工、内存布局、流水化、向量化与融合边界。

同一个矩阵乘可以是计算受限，也可以是内存受限。判断取决于算术强度、矩阵形状和具体硬件的 Roofline，而不是“矩阵乘一定吃算力”这样的标签。NVIDIA 的矩阵乘性能指南也把形状、Tile 和算术强度作为判断性能上限的关键因素。⁵

算子优化必须先通过精度和边界条件测试，再比较单算子性能；最后还要回到端到端关键路径。若该算子只占总时间的 5%，即使局部快 2 倍，理想端到端收益也不到 3%：

\[ S_{\text{end-to-end}} = \frac{1} {(1-f)+\frac{f}{S_{\text{kernel}}}} \]

其中 \(f\) 是原始运行时间中该算子的占比，\(S_{\text{kernel}}\) 是单算子加速比。

通信与系统软件¶

通信与系统软件层决定“数据如何跨设备和节点移动”。它控制 Collective 算法、Buffer、Transport、拓扑路径、完成通知和通信流。

NCCL 提供 AllReduce、AllGather、ReduceScatter 等拓扑感知 Collective，并面向多 GPU、PCIe、NVLink 和网络互连优化通信。⁶ 但通信库不是完整训练框架：它不知道损失函数、优化器状态或请求 SLO，框架需要决定何时、对谁、以多大粒度调用它。

这一层的验收不能只看链路标称带宽，还应观察有效带宽、小消息时延、拓扑差异、拥塞、同步等待和计算通信 Overlap。通信变快但关键路径不变，端到端也可能没有收益。

硬件¶

硬件层决定“软件可用的物理能力与约束”，包括矩阵或向量计算单元、HBM、片上 SRAM、互连、DMA Engine、指令集和功耗上限。

硬件规格给出的是上限，不是应用性能。软件只有选择合适的数据类型、布局、并行度和数据搬运方式，才能逼近峰值。CUDA 编程指南对线程层次、存储层次和异步执行的描述，本质上就是软件映射硬件资源时需要遵守的执行模型。⁷

先定位层次，再拧动正确的瓶颈螺丝 — 概念插图：优化不是从栈顶一路“调参”，而是先找到真正振动的瓶颈螺丝，再沿上下游确认影响。图为本文生成的手绘概念图，不作为性能证据。

软硬协同¶

软硬协同不是与七层并列的第八层，而是一种跨层优化范式。 它要求软件设计显式利用硬件特性，或让后续硬件能力根据工作负载共同演进。

可以把它看成一条连续谱：

硬件感知的软件优化： 硬件不变，软件根据片上存储、Tensor Core/Cube、DMA Engine 或网络拓扑重新选择 Tile、Layout 和流水线。
接口协同： 编译器、运行时与通信库围绕硬件能力建立稳定抽象，使上层框架能可靠使用异步拷贝、低精度计算或设备直连。
共同设计： 软件工作负载反过来影响指令、缓存、互连或加速单元设计，软件也同步改变数据流和算法表达。

因此，FlashAttention 可以称为典型的硬件感知算法与 Kernel 设计：它没有修改注意力的数学结果，也没有要求先制造一种新 GPU，而是根据 GPU 的 HBM 与片上 SRAM 层级重新安排计算。只有当硬件本身也因这种数据流而发生设计变化时，才是更严格意义上的共同设计。

训练框架优化是否属于软硬协同

取决于改动的证据链。FSDP 的“状态分片”概念可以跨硬件复用；具体的 Bucket 大小、Overlap 时机、Stream 编排和通信路径则可能高度依赖设备与网络。不能仅凭代码运行在 GPU/NPU 上，就把整个改动归为软硬协同。

现实中的“软硬协同优化岗位”往往跨越训练或推理框架、编译器、Kernel、通信库和硬件性能分析。岗位名称描述的是协作半径，不是单一的软件层。

FlashAttention 机制卡¶

FlashAttention 适合用来观察跨层优化：数学上仍是精确注意力；算法上重新排列计算；Kernel 上利用片上存储分块；框架或编译器再负责选择和调用具体后端。下面以 FlashAttention-2 论文中的前向 Algorithm 1 与附录反向推导为依据，建立一张教学机制卡。⁸

起因与直觉¶

标准缩放点积注意力为：

\[ S = \frac{QK^\mathsf{T}}{\sqrt d} + M,\qquad P = \operatorname{softmax}(S),\qquad O = PV \]

其中：

\(Q,K,V\in\mathbb{R}^{B\times H\times N\times d}\)；
\(B\) 是 Batch，\(H\) 是注意力头数，\(N\) 是序列长度，\(d\) 是每个头的维度；
\(M\) 是广播到 \([B,H,N,N]\) 的 Mask；
语义上的 \(S,P\in\mathbb{R}^{B\times H\times N\times N}\)，输出 \(O\in\mathbb{R}^{B\times H\times N\times d}\)。

朴素实现若把完整 \(S\) 和 \(P\) 写回 HBM，会产生随 \(N^2\) 增长的中间读写。FlashAttention 的核心直觉不是“少算一个近似注意力”，而是：

把 \(Q,K,V\) 分块搬到片上快存储，在块内完成得分、Softmax 和输出累加，只把在线归一化状态与最终输出写回 HBM。

Softmax 不能简单按块独立计算，因为后来的块可能出现更大值。对第 \(i\) 个 Query 行块，算法维护：

\(m_i\)：目前见过的行最大值；
\(\ell_i\)：在统一最大值基准下的指数和；
\(O_i\)：已经归一化的部分输出。

当新块的局部最大值 \(\widetilde m_i\) 到来时，旧结果通过 \(\exp(m_i-m_i^{\text{new}})\) 重缩放，新块通过 \(\exp(\widetilde m_i-m_i^{\text{new}})\) 对齐，从而得到与一次性 Softmax 相同的结果。

三种不同粒度

概念： IO-aware exact attention，即不改数学语义地减少外部存储读写。
可复用机制： 分块、在线 Softmax、状态重缩放与反向重计算。
框架补丁： 某个 PyTorch、Triton、CUDA 或 NPU 后端的调用、Shape 限制与调度。本文伪代码解释前两者，不宣称逐行对应某个当前后端源码。

对象与生命周期¶

对象	逻辑形状	生命周期	是否完整写入 HBM
\(Q,K,V\)	\([B,H,N,d]\)	整个算子	是，作为输入
\(Q_i\)	\([B,H,B_r,d]\)	一个 Query Tile	否，片上暂存
\(K_j,V_j\)	\([B,H,B_c,d]\)	一个 Key/Value Tile	否，片上暂存
\(S_{ij}\)	\([B,H,B_r,B_c]\)	一次 Tile 计算	否，用后释放
\(\widetilde P_{ij}\)	\([B,H,B_r,B_c]\)	一次 Tile 计算	否，用后释放
\(m_i,\ell_i\)	\([B,H,B_r]\)	跨 Key/Value Tile	是，小状态
\(O_i\)	\([B,H,B_r,d]\)	跨 Key/Value Tile	是，滚动更新
\(D_i=\operatorname{rowsum}(dO_i\odot O_i)\)	\([B,H,B_r]\)	反向的一行统计	可预计算保存
\(dQ,dK,dV\)	同对应输入	整个反向	是，最终梯度

这里最重要的区分是：完整 \(S\) 和 \(P\) 在数学语义中存在，但在物理执行中不必被完整物化。 这也是分析编译器与 Kernel 时常用的“语义对象—物理对象”分离。

完整伪代码¶

下面的伪代码刻意写出 Mask、在线归一化、反向重计算与临时对象释放，便于检查输入、输出和生命周期。真实实现还会根据设备加入线程映射、流水线、向量化和融合。

def flash_attention_forward(Q, K, V, mask, Br, Bc):
    # Q, K, V: [B, H, N, d]
    # O:       [B, H, N, d]
    # m, l:    [B, H, N]
    O = zeros_like(Q)
    m = full([B, H, N], -inf)
    l = zeros([B, H, N])

    for j in range(ceil_div(N, Bc)):
        K_j = load_tile(K, j, Bc)
        V_j = load_tile(V, j, Bc)

        for i in range(ceil_div(N, Br)):
            Q_i = load_tile(Q, i, Br)
            O_i, m_i, l_i = load_state(O, m, l, i, Br)

            S_ij = matmul(Q_i, transpose(K_j)) / sqrt(d)
            S_ij = apply_mask(S_ij, mask, i, j, Br, Bc)

            m_tile = rowmax(S_ij)
            P_tilde = exp(S_ij - m_tile[:, None])
            l_tile = rowsum(P_tilde)

            m_new = maximum(m_i, m_tile)
            alpha = exp(m_i - m_new)
            beta = exp(m_tile - m_new)
            l_new = alpha * l_i + beta * l_tile

            O_numerator = (
                (alpha * l_i)[:, None] * O_i
                + beta[:, None] * matmul(P_tilde, V_j)
            )
            O_new = O_numerator / l_new[:, None]

            store_state(O, m, l, i, O_new, m_new, l_new)
            release(S_ij, P_tilde)

    return O, m, l

反向不保存完整 \(P\)，而是用 \(Q_i,K_j,m_i,\ell_i\) 重建当前 Tile 的概率，再累计梯度：

def flash_attention_backward(Q, K, V, O, dO, mask, m, l, Br, Bc):
    # 行统计用于 softmax 反向：
    # dS = P * (dP - rowsum(dP * P))
    D = rowsum(dO * O)
    dQ = zeros_like(Q)
    dK = zeros_like(K)
    dV = zeros_like(V)

    for j in range(ceil_div(N, Bc)):
        K_j = load_tile(K, j, Bc)
        V_j = load_tile(V, j, Bc)
        dK_j = zeros_like(K_j)
        dV_j = zeros_like(V_j)

        for i in range(ceil_div(N, Br)):
            Q_i = load_tile(Q, i, Br)
            dO_i = load_tile(dO, i, Br)
            m_i = load_tile(m, i, Br)
            l_i = load_tile(l, i, Br)
            D_i = load_tile(D, i, Br)

            S_ij = matmul(Q_i, transpose(K_j)) / sqrt(d)
            S_ij = apply_mask(S_ij, mask, i, j, Br, Bc)
            P_ij = exp(S_ij - m_i[:, None]) / l_i[:, None]

            dV_j += matmul(transpose(P_ij), dO_i)
            dP_ij = matmul(dO_i, transpose(V_j))
            dS_ij = P_ij * (dP_ij - D_i[:, None])

            add_tile(dQ, i, matmul(dS_ij, K_j) / sqrt(d))
            dK_j += matmul(transpose(dS_ij), Q_i) / sqrt(d)
            release(S_ij, P_ij, dP_ij, dS_ij)

        store_tile(dK, j, dK_j)
        store_tile(dV, j, dV_j)

    return dQ, dK, dV

这段伪代码给出了机制闭环，但不是性能实现：真实 Kernel 会重新选择循环顺序、让多个 Query Tile 共享已加载的 Key/Value Tile，并处理因果 Mask、Dropout、GQA/MQA、变长序列和数值精度等边界。

五视图¶

五张视图回答的是不同问题：

物理结构图： 哪些对象在 HBM，哪些 Tile 只在片上短暂存在；
因果逻辑图： 为什么减少 HBM 往返可能比单纯减少 FLOPs 更有效；
过程图： 前向如何在线归一化，反向如何重计算概率；
时序图： 临时对象何时创建、使用和释放；
数据流图： 输入、滚动状态、输出与梯度之间如何传递。

把它们混成一张普通“流程图”，通常会隐藏最关键的存储位置和生命周期。

论文证据与边界¶

FlashAttention-2 论文中的 GPU 存储层次与 GPT-2 训练对比 — FlashAttention-2 论文 Figure 1：左侧给出 GPU 存储层次，右侧展示一个 GPT-2 训练配置中的端到端吞吐对比。它支持“IO-aware Attention 可带来训练收益”的论点，但不是所有模型、后端和硬件的通用倍数。来源：Dao, 2023，arXiv:2205.14135v2。

论文的概览图把存储层次与端到端训练结果放在一起，说明机制收益可以穿过 Kernel 边界进入模型训练；但右侧柱状图仍然只是特定 GPT-2 配置、软件版本与 A100 环境下的一次结果，不能直接外推到 NPU、不同 Mask、超长序列或推理服务。

FlashAttention-2 论文 Figure 2 的 FLOPs、HBM 读写与时间对比 — FlashAttention-2 论文 Figure 2 的受控示例：A100、GPT-2 medium 的一个注意力前向与反向，序列长度 1024、头维 64、16 头、Batch 64。标准注意力为 66.6 GFLOPs、40.3 GB HBM 读写、41.7 ms；FlashAttention 为 75.2 GFLOPs、4.4 GB、7.3 ms。它说明“多算一点但少搬很多数据”可以更快，不代表任意 Shape 都获得相同收益。来源：Dao, 2023，arXiv:2205.14135v2。

这个对比特别重要，因为 FlashAttention 在示例中 FLOPs 反而更多，但 HBM 读写更少、执行时间更短。它反驳了“优化就是减少计算量”的单一视角，也给出完整的因果链：

\[ \text{分块与在线 Softmax} \rightarrow \text{不物化完整 }S,P \rightarrow \text{HBM 读写下降} \rightarrow \text{受 IO 限制的 Shape 加速} \]

这条链的最后一步是条件结论。若 Shape 很小、Kernel 并行度不足、实现存在额外同步，或者运行已由其他算子主导，收益会缩小。

显存影响契约¶

分析“省显存”时，不能只写一句“复杂度从 \(O(N^2)\) 降为 \(O(N)\)”。端到端峰值更接近以下账本：

\[ \begin{aligned} M_{\text{peak}}={}& M_{\text{model state}} +M_{\text{saved activations}} +M_{\text{operator workspace}}\\ &+M_{\text{communication buffers}} +M_{\text{live inputs/outputs}}\\ &+M_{\text{allocator reserved unused}} +M_{\text{runtime margin}} \end{aligned} \]

FlashAttention 直接改变的主要是：

算子工作区： 不再长期物化完整 \(S\) 与 \(P\)；
保存激活： 可以保存较小的归一化统计，在反向重计算局部概率；
生命周期： \(S_{ij}\)、\(P_{ij}\)、\(dS_{ij}\) 只在当前 Tile 内短暂存在。

它不会直接减少模型参数、优化器状态、KV Cache、通信 Buffer 或内存分配器碎片。因此，算子级 \(O(N^2)\) 中间量消失，不等于整个模型峰值严格按相同比例下降。

一个可执行的效果契约应写成：

目标项： \(M_{\text{operator workspace}}\) 与相关 \(M_{\text{saved activations}}\)，以及 HBM 读写量。
动作： 通过分块、在线归一化和反向重计算，减少完整中间张量的物化并缩短生命周期。
代价： 额外 FLOPs、Tile 工作区、Kernel Shape 约束，以及 Mask、Dropout、精度和反向实现复杂度。
同场测量： 固定模型、Batch、序列长度、头维、精度、Mask/Dropout、硬件、软件版本和分配器配置，再比较单算子与端到端结果。

如何定位优化层¶

面对一个新的“性能优化”需求，可以按下面的顺序工作。

固定语义与指标¶

先写下不可变化的条件。训练任务可能要求损失曲线与最终质量等价；在线推理可能要求输出一致、P99 小于阈值；算子替换则需要明确误差容限。

然后只选一个主目标，例如：

在质量不下降的条件下提高训练 tokens/s；
在 P99 不超过 2 s 的条件下降低单 Token 成本；
在固定 Shape 与精度下减少注意力 HBM 流量。

没有语义约束和主指标，后面的“加速”很容易只是换了任务。

建立关键路径与内存账本¶

使用 Profiler 把时间拆成计算、通信、排队、编译、Kernel Launch 和空闲；把峰值内存拆成模型状态、激活、工作区、通信 Buffer、活跃输入输出、分配器保留与运行时余量。

先找活跃瓶颈

静态规格只告诉你哪里“可能慢”。需要在真实 Shape、Batch、拓扑与请求分布下观察关键路径，才能知道哪个瓶颈正在支配当前运行。

选择最窄责任层¶

找到占主导的对象和生命周期后，选择能够直接控制它的最窄层次：

质量或有效计算量不合适，回到模型与算法层；
参数、梯度或激活驻留不合适，优先看训练系统；
排队、KV 配额或 Prefill/Decode 竞争，优先看推理服务；
Graph Break、重编译或内存计划，优先看编译与运行时；
局部 HBM、Tile 或融合问题，优先看 Kernel；
Collective、链路或拓扑等待，优先看通信；
现有能力无法表达目标数据流，再考虑硬件变化。

“最窄”不意味着只改一层，而是先建立清晰的主因，再检查上下游接口。

做两级验证¶

局部优化至少需要两级实验：

Microbenchmark： 验证被控制的局部物理量确实改变，例如 Kernel 时延、HBM 字节或 Collective 带宽。
End-to-end： 验证改动进入了真实关键路径，没有被调度、编译、通信或其他算子抵消。

可以把流程写成下面的通用伪代码：

baseline = measure(workload, fixed_semantics, fixed_environment)
bottleneck = locate_critical_path(baseline.trace, baseline.memory_ledger)
layer = owner_of(bottleneck.object, bottleneck.lifetime)

candidate = change_only_the_controllable_factor(layer)
local_result = microbenchmark(candidate, same_shape_and_precision)
e2e_result = measure(candidate, fixed_semantics, fixed_environment)

accept = (
    correctness_within_tolerance(candidate)
    and local_result.changes_expected_physical_quantity
    and e2e_result.meets_primary_metric
    and e2e_result.meets_slo_and_memory_budget
)

最后报告条件、代价和未覆盖边界，而不是只报告最快的一次数字。

常见误区¶

框架优化就是算子优化： 框架可能调用更快的 Kernel，但它还负责状态、调度和跨算子依赖。反过来，单个 Kernel 可以在完全不理解训练并行或请求 SLO 的情况下被独立优化。二者通过接口连接，不应在归因时合并。
Kernel 快等于模型快： 端到端收益受该 Kernel 在关键路径中的占比限制，还可能被 CPU 调度、编译、通信或同步抵消。必须同时报告 Microbenchmark 和 End-to-end。
训练吃算力，推理吃带宽： 这只能作为部分场景的经验起点。训练也可能受通信、显存或流水线空泡限制；长 Prompt 的 Prefill 可能具有较高算术强度，Decode 才更常受权重与 KV 读写影响。应针对具体阶段与 Shape 判断。
使用硬件 API 就是软硬协同： 调用 CUDA、CANN、NCCL 或 HCCL 只说明使用了某个接口。只有软件设计显式围绕硬件约束重构数据流，或软硬件共同改变接口与能力，才形成有意义的协同证据。
换更强硬件就不需要软件优化： 新硬件会改变算力、带宽、容量和拓扑的比例，也会移动瓶颈。旧软件的数据布局、并行策略或调度可能无法利用新能力，甚至让更高峰值长期闲置。

总结¶

AI 模型优化不是一条从算法到芯片的单向流水线，而是七个相互反馈的控制层：

模型与算法决定算什么；
训练系统组织状态、前后向和分布式并行；
推理服务组织请求、Token 与 KV Cache；
编译与运行时把程序变成稳定的执行图；
算子与 Kernel把局部计算映射到设备；
通信与系统软件搬运跨设备数据；
硬件提供物理能力与约束。

软硬协同贯穿这些层。判断一个优化属于哪里，不看它使用了多少流行名词，而看控制对象、物理变化、验收指标和证据边界是否闭环。FlashAttention 的价值也正在于此：它把数学语义、IO 复杂度、片上存储、Kernel 数据流和端到端训练收益连成了一条可验证的因果链。

可复用提示词¶

当需要分析新的 AI 优化方案时，可以复用下面的提示词：

请不要先按项目名称归类，而是分析这个优化的：

1. 固定不变的模型语义与质量约束；
2. 主要控制对象、逻辑形状和生命周期；
3. 真正改变的物理量：FLOPs、HBM 字节、峰值内存、
   网络字节、Kernel 启动、排队时间或硬件能力；
4. 主要责任层：模型算法、训练系统、推理服务、编译运行时、
   Kernel、通信系统软件或硬件；
5. 是否跨层，若是，每层分别修改了什么；
6. Microbenchmark 与 End-to-end 的验收指标；
7. 同场比较所需的模型、Shape、Batch、精度、拓扑、软件版本；
8. 已知代价、失败边界和不能由现有证据推出的结论。

最后用“对象—动作—物理量—指标—边界”的格式给出结论。

参考资料¶

PyTorch, Fully Sharded Data Parallel. ↩
vLLM, Scheduler Configuration. ↩
vLLM, Metrics. ↩
PyTorch, torch.compile. ↩
NVIDIA, Matrix Multiplication Background User's Guide. ↩
NVIDIA, NCCL Documentation. ↩
NVIDIA, CUDA Programming Guide. ↩
Tri Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, arXiv:2205.14135v2, 2023. ↩