在昇腾NPU 上做推理部署,第一个撞上的硬钉子往往是 Dynamic Shape(动态输入形状)。不管你输入的是一段 128 个 Token 的短文本还是 4096 个 Token 的长文档,CANN 在编译模型时并不知道输入序列会有多长。这个"不知道"会从图编译一路传递到 Runtime 执行,在每个层面都引发额外开销。

这篇文章不用一堆术语堆砌,而是从一个推理请求的真实旅程出发,解释动态 Shape 为什么难、昇腾是怎么处理的。


为什么动态 Shape 是难题

静态 Shape 的世界很简单。模型编译时 ATC 知道每个 Tensor 的精确尺寸——input_ids[1, 512]attention_mask[1, 512]。GE 可以精确计算每块显存要分配多少、每个 Buffer 要多大、算子融合后中间 Tensor 放在片上哪个偏移地址。这些信息在编译期全部确定。

动态 Shape 打破了这一切。

假设 input_ids 的序列长度在 128 到 4096 之间变化。编译时不知道长度,GE 就不能做精确的内存分配——不知道中间激活层的 Tensor 有多大、不知道 KV Cache 要预留多少空间。退一步说,即使把最大值 4096 传给编译期来预分配空间,实际跑 128 时显存利用率只有 3%,废了 97% 的显存。

这就是核心矛盾:编译期需要确定信息才能做优化,但推理时的输入长度无法提前锁定。


大模型推理中的动态 Shape 场景

Transformer 模型里动态 Shape 最典型的来源有三个:

序列长度变化。 对话场景中用户一次问 50 个字,下一次问 2000 个字。Attention 模块里的 Q/K/V 矩阵维度直接跟着序列长度走。预填充(Prefill)阶段和自回归解码阶段的序列长度也不同——Prefill 是满序列计算,解码是逐 Token 追加。

Batch 大小变化。 在线推理服务的 Batch 是动态的——客户端请求到达时间不确定,服务端攒够一定数量或等够一定时间才提交一次推理。dynamic_batch_size 是部署中最常用的动态 Shape 配置。

Cache 序列长度增长。 自回归解码中 KV Cache 随输出 Token 逐个增长。每次生成一个新 Token,Cache 长度 +1,下个 step 的 Attention 计算输入就变了。Runtime 需要在每个 token 生成后重新调整 KV Cache Tensor 的偏移地址。


昇腾如何处理动态输入

CANN 在三个层次上处理动态 Shape:ATC 编译层、GE 图层、Runtime 执行层。

编译层:预注册 Shape 范围。 ATC 编译时允许指定每个维度允许的取值范围。比如 --input_shape_range="input_ids:[1~8,128~4096]"。ATC 在编译时会对这个范围内的 Shape 做覆盖性优化——测试该范围内哪些算子融合策略依然有效、哪些内存分配方案可以通用。编译后的 OM 模型可以处理这个范围内的任意输入。

但覆盖性优化只能针对已知范围内的 Shape。如果推理时来了一个不在范围内的输入(比如序列长度 8192),ATC 编译时的优化全部失效,Runtime 需要触发重新编译。

图编译层:Shape 无关的图优化。 GE 把算子融合分成两类:一类依赖 Shape 信息(比如内存优化相关),一类不依赖(比如常量折叠、无用节点消除)。GE 对不依赖 Shape 的优化预编译写入 OM,对依赖 Shape 的优化延迟到运行时、拿到真实输入 Shape 后再执行。这样动态输入下至少部分优化提前做好了。

Runtime 层:动态 Tensor 管理。 Runtime 在收到推理请求时才知道真实 Shape。它需要动态调整 Tensor 的显存分配和 Buffer 偏移。CANN Runtime 的做法是维护一个动态 Tensor 表——每次推理 step 更新 Tensor 的形状、大小、显存地址。KV Cache 的管理是最典型的例子:每生成一个 Token,Runtime 更新 KV Cache 的偏移量,下一个 Attention 算子从偏移后的位置继续读写。


动态 Shape 对性能的影响

说清楚机制之后,量化一下动态 Shape 的实际代价。

场景 优化手段 推理延迟 显存利用率
固定 Shape 512 ATC 精确编译 5.2ms 100%
动态 Shape 128-4096, 最大分配 按最大值预分配 6.8ms 12-100%
动态 Shape 128-4096, 动态分配 运行时按需分配 8.5ms 60-100%

按最大值预分配延迟只多了 30%,但短序列场景下显存浪费严重。运行时按需分配显存利用率好了很多,但每次 Tensor 重分配有几十微秒的开销。

实际部署中的权衡是:大部分请求集中在某个典型长度(比如 256-1024 Token),用这个典型范围做固定 Shape 编译,超长请求走 Fallback。既保证了主流场景的极致性能,也覆盖了长尾场景。


Batch 动态化的特殊处理

动态 Batch 跟动态序列长度的处理方式不同。Batch 变化主要在 Prefill 阶段影响计算量和显存,解码阶段 Batch 不是连续变化的——服务端攒够一个 Batch 才提交。

CANN 推荐的做法是预编译几个常用 Batch 档位的 OM 模型(Batch=1,4,8),推理时根据当前排队的请求数量选择模型。这比编译一个通用动态 Batch 模型简单得多,性能也更稳定。

# 分别编译 3 个模型
atc --model=model.onnx --dynamic_batch_size="1" --output=model_bs1
atc --model=model.onnx --dynamic_batch_size="4" --output=model_bs4
atc --model=model.onnx --dynamic_batch_size="8" --output=model_bs8

推理服务根据当前队列长度在 3 个模型之间切换。切换开销只有指针替换,不需要重编译。


结语

动态 Shape 是 Transformer 推理中无法回避的问题。理解了它从 ATC 到 GE 到 Runtime 的传递路径,就能理解为什么同一个模型在固定输入和动态输入下的推理延迟可能差 2-3 倍。解决思路不是"消除动态 Shape"——这是不可能的——而是在动态 Shape 的约束下,把能预编译的优化提前做完、把只能在运行时做的优化做到最轻量。

CANN Dynamic Shape 文档

AscendCL 推理部署指南


一个具体的动态 Shape 例子

用 BERT 做文本分类,训练时每条数据 padding 到 512 Token。部署到线上一看,用户发来的句子短的十几个字、长的几千字。如果全部 padding 到 512,短文本浪费 90% 的计算量。

CANN 的 Dynamic Shape 机制允许你在不 padding 的情况下直接推理。但代价是 ATC 无法像固定 Shape 那样精确优化内存布局。比如算子融合后中间 Tensor 的位置,固定 Shape 下可以直接算偏移地址写到 OM 里;动态 Shape 下必须留一个"运行时计算偏移"的标记,每次推理多一次计算开销。

这也是为什么很多线上服务最终选择"多档静态 Shape"策略——编译 5 个模型分别在 128/256/512/1024/2048 五个长度上做精确优化,推理时根据输入长度就近选择一个。五倍编译时间,换来每档输入的最优性能。

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐