动态Shape支持机制：计算平台元数据定义与编译优化实战

做推理引擎这些年，我最大的感受是：静态 shape 的优化已经卷到头了，真正拉开差距的反而是动态 shape 的支持能力。前阵子我们服务里一个模型输入尺寸从固定 512 改成允许 256 到 1024 动态变化，原本编译好的计算图直接报废，离线编译的二十多个 shape 变体全都不匹配，线上流量一上来就卡在 shape check 上。后来我们把精力砸在“计算平台元数据定义”这一层，重新设计了 shape 表达式的表示和编译期的约束传播，才算把动态 shape 的坑填平。

这篇文章就聊聊我在这过程中的理解：计算平台里的元数据到底在定义什么，计算图编译为什么怕动态 Shape，以及一套能落地运行的动态 Shape 支持机制应该怎么设计。内容偏底层，适合正在做推理引擎、编译器优化，或者被动态 shape 折磨过的小伙伴，看完可以直接对照自己项目里的数据结构和技术选型。

1. 先从元数据定义说起

1.1 一个 Tensor 的身份证到底包含什么

很多人一听到“元数据定义”，第一反应就是 shape 和 dtype。但在真正的计算图编译流程里，Tensor 的元数据远不止这两个字段。我常用“身份证”来打比方：shape 和 dtype 只是姓名和性别，真正决定你能不能坐高铁、住酒店的还有更多字段。

一份完整的 Tensor 元数据，在编译器和运行时之间传递时，至少要包含以下几类信息：

别小看 strides 和 memory format 这两个字段。一个x[:, ::2]的切片结果，shape 看着没问题，但 strides 是乱的；如果编译器没捕获这个信息，后续直接把该 tensor 当连续内存传给 kernel，数据全错。动态 Shape 问题也是一样：它不只是把 shape 从常量变成变量，而是会让上面所有字段的推导从“编译期可确定”变成“运行期才可确定”。

所以在设计元数据的时候，我强烈建议把 shape、dtype、strides、memory format 看成一张完整的表，而不是单独处理。动态 Shape 支持机制，本质上要保证这张表在“未知维度”存在的情况下，依然可以被推导、被约束、被验证。

1.2 为什么静态 Shape 这么好用

静态 Shape 指计算图中每个 Tensor 的每一个维度，在编译期都是确定的整数。这种情况下，编译器的日子非常好过，因为所有优化 pass 拿到的都是精确数字，可以做很多激进的重构。

举个例子，给定输入[1, 3, 224, 224]，编译器可以直接算出卷积输出是[1, 64, 112, 112]，然后提前分配一块固定大小的内存，所有中间结果都在这块 arena 里复用。kernel 选择也简单：224 的尺寸配 3x3 卷积，用滑窗实现还是用 implicit GEMM，按 shape 大小查表就能定。还有 layout assignment，NCHW 还是 NHWC，编译器可以根据后续算子的访问偏好统一决定。

这种模式在目标检测、分类这种固定分辨率的模型上非常有效，很多部署栈也因此把“静态 shape”当成默认配置。但现实场景里，动态 shape 是躲不掉的：

• NLP 模型里 batch size 和序列长度几乎不可能固定，一条样本 32 个 token，另一条 512 个 token，模型结构一样，shape 完全不同。
• 推荐系统里用户特征数量、候选 item 数量都是动态的。
• 目标检测模型经常要接收不同分辨率的输入，或者在一个 batch 内做 multi-scale 推理。

静态 shape 编译那一套在动态输入面前很尴尬：你没法提前知道内存多大，没法固定 kernel 参数，连计算图里的中间 tensor 是否存在都会变。于是就有了“动态 Shape 支持机制”的需求。

2. 动态 Shape 为什么让编译器头疼

2.1 动态 Shape 和动态图不是一回事

先把概念理清，很多人会把“动态 Shape”和“动态图”混在一起。

动态图是 PyTorch 默认那种 eager 模式：每次执行一行 Python 代码，就现场建一个算子节点，现场执行。这种情况下 shape 是天然动态的，因为根本没有一张先编译后执行的完整图，也没有所谓的“编译期优化”。

而计算图编译模式下，我们希望先把整张模型编译成一个可执行文件或可部署的二进制包，再在服务端反复执行。这时候输入 shape 如果是动态的，编译器面对的问题就是：如何在不知道具体维度数值的情况下，生成正确且高效的代码。

所以动态 Shape 支持机制，是静态图和动态语义之间的一个妥协：静态编译的底子还在，但要给“未知量”留出表达空间。这个“未知量”就是符号 dimension，所有的难点都围绕它展开。

2.2 一个动态维度引发的连锁反应

表面上看，动态 shape 不就是一个维度写“None”嘛，运行时拿到具体值填进去不就行了？但真正做一次 shape 推断就会发现，一个动态维会像多米诺骨牌一样，把下游所有 Tensor 的元数据全部变成表达式。

假设输入是[N, L]，其中 N 是 batch，L 是序列长度，二者都是动态的。经过 embedding 后变成[N, L, 128]，再经过reshape变成[N * L, 128]，然后做矩阵乘。下游输出 shape 写成[N * L, 64]还算简单，可一旦中间夹着切片、concat、split 或者带 mask 的操作，shape 表达式就会越来越复杂。比如：

# 伪代码 y = x[:, :seq_len] # [N, seq_len, 128] z = concat([y, padding], axis=1) # [N, max_len, 128] w = z.reshape([N * max_len, 128]) # [N*max_len, 128]

这里的seq_len是一个动态值，max_len又可能是另一个约束下的常量。编译器要回答的问题变成：

• 我怎么知道z的第二个维度等于seq_len + padding_len？
• 我怎么证明 reshape 操作的元素数量是可整除的？
• 我在内存上应该给z预留多少空间？

这些问题如果全靠运行时去碰运气，那编译优化基本无从谈起。所以动态 shape 难的不是“报错”，而是如何在编译期留下足够多的元数据约束，让编译器即使不知道具体值，也能做出安全、不退回解释执行的决定。

2.3 静态优化手段集体失效

一旦 shape 变成动态，以下几类优化会遇到大麻烦：

• 布局优化：NCHW 和 NHWC 的转换往往依赖固定维度顺序。动态 shape 下，维度的顺序本身是固定的，但每个维度的“语义”在运行期才绑定到具体值，布局 pass 不能随手改内存 order，否则运行时校验会很麻烦。
• 内存复用：arena 方式要求所有 tensor 的 size 编译期可算。动态 shape 下每个 tensor 的 size 是一个表达式，两个 tensor 能不能复用同一块内存，需要做表达式等价性判定，这比整数比较难得多。
• 算子融合：某些融合要求中间 tensor 的 shape 是静态的，比如卷积+ReLU 可以融合，但卷积+动态 reshape 就不能乱来，因为 reshape 可能在处理动态维时引入数据拷贝。
• kernel 选择：GEMM 有专门为小 shape 和大 shape 优化过的实现，编译期没法确定走哪条分支，运行时再做 dispatch 又可能带来额外开销。

动态 Shape 支持机制，本质上就是为这些“失效点”逐一提供替代方案：内存上做动态规划，kernel 上做多版本派发，融合上尽量收缩动态区域。

3. 元数据定义层怎么设计才能扛住动态 Shape

3.1 从 int list 升级到 shape expression

既然动态 shape 无法用固定整数表示，第一步就是改造 Tensor 元数据里的 shape 字段。我见过最偷懒的做法是把某个维设成0或者-1来表示动态，然后在运行时去填。这个方案能做 demo，但完全无法支撑编译优化，因为-1不携带任何约束，编译器对它的了解为零。

更好的做法是把 shape 设计成“表达式”，而不是一个简单整数。每个维度可以是下面三种之一：

• 常量整数，例如128
• 符号变量，例如N、seq_len
• 符号变量的表达式，例如N * 2 + 4

在代码里大致长这样：

from dataclasses import dataclass from typing import Optional, Union, Tuple @dataclass(frozen=True) class SymDim: name: str lower_bound: Optional[int] = None upper_bound: Optional[int] = None @dataclass(frozen=True) class TensorMetadata: shape: Tuple[Union[int, SymDim, tuple], ...] # 简化版，tuple代表表达式 dtype: str strides: Optional[Tuple[Union[int, str], ...]] = None memory_format: str = "dense" quant_scale: Optional[float] = None quant_zero_point: Optional[int] = None

这里我解释一下为什么用符号变量而不是直接写死表达式节点。符号变量的意义在于：它可以在整个计算图里共享。同一个 batch size 出现在输入、中间结果、输出的多个维度里，它们应该引用同一个符号，这样编译器才能知道x1.shape[0]和x2.shape[0]是同一个值，可以做维度绑定。

表达式节点则可以表达缩放和偏移，比如某种 padding 后的序列长度是seq_len + 2，这比新造一个符号更准确，也更利于约束传播。

3.2 在 IR 级别表达动态 Shape

光在宿主语言里定义 TensorMetadata 还不够，计算图编译的各个 pass 需要在中间表示（IR）里直接看到动态 shape。否则每个 pass 都要绕回到外部数据结构去查，工程上很容易失控。

我常用的做法是在 IR 的 type 体系里直接支持“符号 shape”。如果你们用的是 MLIR 风格，可以写成类似这样：

func.func @inference(%arg0: tensor<?x?xf32> {shape_sym = ["N", "L"]}) -> tensor<?xf32> { %c128 = arith.constant 128 : i64 %slice = tensor.extract_slice %arg0[0, 0] [%N, %L] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32> %pad = tensor.pad %slice low[0, 0] high[%c128, %c128] : tensor<?x?xf32> to tensor<?x?xf32> // ... }

这里?表示动态维，而shape_sym属性把动态维和具体符号名绑定起来。这样 IR 里的每个 value 都自带一份完整的“shape 表达式”信息，编译器 pass 可以做统一处理。

同时，还要在 IR 里维护一个约束集合，例如：

#shape_constraint = #shape.constraint< N >= 1, N <= 64, L >= 8, L <= 512, N * L % 8 == 0 >

我把这个约束集合挂在包含 dynamic shape 的 func 或 region 上，所有涉及 shape 的 pass 都先读取它。这不是锦上添花，而是动态 shape 支持机制的核心：光有表达式没有约束，优化 pass 依然不敢做任何假设。

3.3 约束集合是动态 Shape 的“隐藏裁判”

很多框架实现动态 shape 时，只做到“表达”，没做到“约束”。结果就是编译期能画出 shape 表达式，但无法回答“这个维度能不能被 32 整除”“这两个维度是不是一定相等”这类问题。没有约束，就无法安全生成高效的向量化 kernel，也无法做内存复用。

约束集合需要支持几种常见关系：

• 相等关系：batch == input_0.dim0，即两个符号代表同一个值，可以 union-find 合并。
• 区间关系：1 <= seq_len <= 1024，提供 range。有 range 之后，编译器可以估算最大内存，准备 buffer pool。
• 对齐关系：seq_len % 8 == 0，这个信息对 SIMD 优化至关重要。
• 算术关系：full_len == seq_len + 2 * pad_len，用于推导 reshape/concat 后的维度。

约束集合的维护不是一次性完成的，而是随着 shape inference pass 的处理不断叠加。每个 op 的 shape function 在生成输出 shape 时，也会生成一些新的约束条件。编译器要做的，是把这些约束收编进一个全局 solver/checker，在涉及到分支判断时向它提问。

我在工程实践中还发现一个窍门：约束集合要同时支持“编译期可满足性检查”和“运行期校验”。编译期我们只需要知道某个假设是否有可能成立，比如“如果 seq_len 能被 8 整除，就可以走 vectorized kernel”，这时候 solver 做一些区间推理就够了。运行期则要真正校验实际 shape 是否满足这些约束，如果不满足，就 fallback 到通用路径。

4. 动态 Shape 支持机制在编译管线中的落地

4.1 基于形状表达式的 shape inference

动态 shape 的 shape inference 不能再用简单的常量传播，每个算子都要有一个“符号 shape 推导规则”。举个例子，矩阵乘的推导规则：

def infer_matmul(a: TensorMetadata, b: TensorMetadata) -> TensorMetadata: # a: [M, K], b: [K, N] m = a.shape[0] k_a = a.shape[1] k_b = b.shape[0] n = b.shape[1] assert_symbolic_eq(k_a, k_b, "matmul contract dim mismatch") return TensorMetadata( shape=(m, n), dtype=infer_matmul_dtype(a.dtype, b.dtype), ) def assert_symbolic_eq(dim_a, dim_b, message): # 如果两个维度都是符号，加入相等约束； # 如果一个是常量，另一个是符号，尝试用区间约束证明。 solver.add_eq(dim_a, dim_b)

注意这里的assert_symbolic_eq不能像静态 shape 那样直接在编译期报错。动态 shape 下，两个符号维度是否相等，可能要留到运行期才能完全确认。所以它实际上是“记录约束 + 运行期检查”，而不是“编译期死磕”。

真实项目里，每个算子都要写类似的 shape function。Pointwise 算子最简单，输出 shape 是对齐后的输入 shape；Reduce、Concat、Reshape 各自有不同的推导逻辑。Complicated 的是 Reshape：输入的总元素数是一个符号表达式，输出 shape 必须能让这些元素数保持一致，编译器只能生成一个运行期 assertion，然后继续往下。

4.2 形状特化：bucketing 和 guard

纯符号推导解决了“正确性”，却解决不了“性能”。动态 shape 下如果所有算子都按最通用实现来编译，性能基本没法看。所以在实际部署里，更常见的策略是“形状特化”，也就是把动态 shape 值域划分成若干桶（bucket），每个桶单独编译一套静态 shape 的变体。

桶的划分一般结合线上 shape 分布。比如：

划分原则是：桶的数量不要太多，一般控制在 20~50 之间；每个桶内的 shape 差异对性能影响不大；极端 shape 不单独开桶，走 fallback。

编译期，IR 会被复制成多个变体，每个变体的动态符号被绑定到一个具体整数。这个绑定发生在 shape inference 之后、所有静态优化之前。每个变体前面插入一个 guard 节点，运行时判断输入 shape 是否匹配当前变体。我再强调一遍，这个 guard 是在运行时执行的，不是编译期比较，所以速度必须快，通常用简单的整数比较就能完成。

4.3 静态子图动态子图混合执行

如果整张图所有地方都动态，那编译优化空间就很小。实际模型通常只有一部分 shape 是动态的，比如输入层和某些动态控制流，其他大部分算子都是形状确定的。

所以我会把计算图划分成静态子图和动态子图。划分规则很直接：

• 一个动态 shape 算子产生的 tensor，如果被后续所有算子使用，而后续算子又要求 shape 精确匹配，那这些算子都要算作动态子图。
• 如果某个动态 tensor 经过 padding、clip 或者特化后 shape 变成常量，那此后的算子可以划回静态子图。
• 边界上可能需要插入 layout conversion 或 memory copy，成本要算在动态子图里。

这种混合执行的好处是，静态子图可以继续享受算子融合、常量折叠、内存复用；动态子图则单独走动态 kernel 和动态分配。我在实际项目里见过一个模型，动态输入进来后先做一轮pad_to_fixed, 之后几十个算子全部恢复成静态 shape，速度几乎追平纯静态编译。

4.4 内存计划和动态缓冲池

动态 shape 对内存规划的影响最大。静态 shape 时，arena 可以一次性把整张图的中间张量排布好；动态 shape 时，每个中间张量 size 都是运行期才知道，传统 arena 方案失效。

我的处理办法是两层分配策略：

1. 静态区：编译期能确定 size 的 tensor，继续走静态 arena，这部分内存可以复用。
2. 动态区：size 只有在运行期确定的 tensor，走一块独立的动态 buffer pool。每个 bucket 会预分配一块最大 size 的 buffer，实际 shape 小于等于 bucket 上限时反复使用。

动态区的复用策略比静态区宽松一些，因为同 bucket 内的 tensor size 不一定完全相等，但 buffer 容量是固定的，可以复用。加上内存池扩容机制，避免每次都走系统 malloc。

另外还要注意对齐问题。动态 shape 取到的值经常是 3、5、17 这种不是 16 的倍数，直接分配会导致 SIMD 访存效率下降。所以我一般会把动态维度向上对齐到 16 或 32，计算时用 mask 或者 pad 处理边界。这个对齐信息也要写进 TensorMetadata，不然运行时不记得 buffer 的有效长度和分配长度之间的差距。

4.5 运行时的形状 dispatch

最后落到运行时，shape dispatch 逻辑其实是一个match_shape(inputs) -> executable的过程。我会为每个编译变体生成一个 shape descriptor，运行时用输入的 shape 去匹配。

匹配顺序一般是：

• 完全相等匹配，命中则直接执行。
• 如果多个变体都能覆盖当前 shape，选最近的一个 bucket，避免频繁重编译。
• 如果都匹配不上，走 fallback：解释执行或者调用一个通用动态 kernel。
• 对于一些高频出现的未匹配 shape，可以触发 on-the-fly JIT 编译，编译结果缓存起来。

运行时 dispatch 的开销必须控制住。形状比较不要每次都用 tuple 哈希，最好把 shape 值打包成几个 64 位整数快速比较。否则动态 shape 的收益会被 dispatch 本身吃掉。

5. 实操中的坑与排查实录

5.1 符号碰撞：所有输出 shape 突然变错

动态 shape 最常见的问题，是约束集合里两个本该独立的符号被错误等价了。比如模型有两个动态输入，一个 batch size，一个 beam size，如果 shape function 里粗心地把它们 merge 到同一个符号，那后面所有 shape 推导都会带上这个错误假设。

我们当时排查了一个很诡异的 bug：模型在 batch=16、beam=4 时正确，batch=8、beam=4 时输出 shape 完全乱掉。最后打印符号绑定表才发现，编译器把batch * beam当成一个新的符号，又和某个中间维度的batch符号 union 了，导致约束关系错误。

排查手段就是“符号绑定表”。在 shape inference 之后，打印每个符号对应的实际维度来源、区间约束、相等类。维护起来不复杂，但对定位这类问题帮助极大。

5.2 编译时间爆炸

动态 shape + bucketing 很容易产生笛卡尔积式的编译变体。比如 batch 有 6 个桶，seq_len 有 6 个桶，图像尺寸有 5 个桶，三者的组合就是 180 个变体，每个变体都要跑一遍完整编译优化，几个小时都编不完。

我最后的解法是：

• 用线上 shape 分布的 top-K 和覆盖率来决定桶，而不是穷举所有组合。
• 提供“编译预算”参数，比如最多编译 32 个变体，超过后直接使用通用 kernel。
• 对模板特化按需 JIT，而不是离线全量编译。
• 在 kernel 层面，尽量把 shape 作为运行时参数而不是编译期模板参数，这样同一个 kernel 可以服务多个桶，大幅减少编译次数。

5.3 动态 shape 下算不对：一个 reshape 引发的血案

动态 shape 下最容易出错的是 reshape。静态 shape 时，元素总数直接算好，reshape 不可能非法。动态 shape 时，[N, 128]reshape 成[N // 2, -1]看着没问题，但如果N是奇数，运行期就会崩。更隐蔽的是，某些实现里-1推导会得到小数，比如N * 128 / 24，如果不整除，结果完全错误。

处理方式必须是：编译期生成动态断言，运行期检查整除性。不要尝试在编译期把所有情况证明完。我给 shape function 加了一个检查点：所有 reshape 操作都会生成assert(total_elements % new_shape_prod == 0)，这样错误能提前暴露在推理阶段，而不是等到 kernel 越界。

5.4 调试技巧小结

我把这几年常用的动态 shape 调试手段整理成一张速查表：

还有一个通用调试开关，我建议所有动态 shape 框架都支持：打开后每个算子的输入和输出会打印两套 shape，一套是编译期的符号表达式，一套是运行期的真实 shape。两者不一致时直接报 warning。这个开关几乎能定位 90% 的动态 shape 问题。

6. 几句题外话：动态 Shape 支持的分寸感

最后聊点个人体会。动态 Shape 支持机制不是越复杂越好。如果你的场景只是 batch size 变一变，直接 padding 到固定 batch 或者做几个 bucket，就能解决绝大多数问题。符号 shape 表达式和约束求解器这套东西，工程成本不低，如果不是要支持那种 shape 变化非常自由的模型，不建议一上来就全上。

我踩过的坑是：一开始想着“做大而全的动态 shape 系统”，结果编译器 pass 越写越多，约束 solver 性能跟不上，最后反而拖慢了上线进度。后来改成“静态为主、动态度量、bucket 兜底”的务实方案，效果反而好很多。

真正要把动态 Shape 做好，元数据定义层的设计是关键。shape 表达式、约束集合、运行期 guard 这三样东西，是整套机制的骨架。先把它们在 IR 里表达清楚，后续任何 pass 都好做；如果这一步偷懒，后面只会越补越乱。

做计算平台优化，很多时候拼的不是技巧，而是对“元数据契约”的尊重。动态 Shape 只是其中一个缩影。希望这篇内容对正在被 shape 困恼的你有帮助，也欢迎在实践中多尝试不同的约束表示方式，找到最适合你业务形态的那一套。