1. 动态张量为什么天生和"编译优化"不对付
1.1 一个真实的性能现场：变长序列把GPU拖垮了
大概半年前，我在优化一个变长序列的推理服务。那批数据每条样本长度差异非常大，短的只有十几个token，长的能到几百。为了跑batch，常规做法是pad到最长长度，于是你会在性能面板上看到一组非常刺眼的数据：显存翻了三倍，但GPU利用率只有30%出头，CPU倒是快跑满了。用火焰图往下钻，开销根本不在计算上，而全耗在张量形状检查、中间缓冲区分配和Python到C++的反复跨语言调度上。
这个现象在动态张量计算场景里非常典型。所谓"动态"，主要指三个维度：第一是形状动态，输入张量的batch、序列长度、特征维度在运行时才会确定；第二是控制流动态，比如while循环、条件分支要等实际数据出现才能走通；第三是内存布局动态，同一个逻辑张量会因为切片、转置、拼接不断改变stride和存储位置。静态编译器拿这种输入没太多办法，纯解释器倒是能应付，但性能又差得不像话。
于是我开始系统性研究一个折中方案：先做一个字节码虚拟机，把张量程序编译成字节码，跑起来之后再对热点路径做实时编译。这套东西的核心思路和LuaJIT有点像，但优化对象从标量运算换成了张量计算，要解决的问题也从"解释器太慢"变成了"动态形状导致的重编译风暴和kernel启动开销"。
1.2 静态编译在动态形状面前的挣扎
静态编译器和动态张量之间有一条很深的鸿沟。以T个时间步的循环为例，静态编译器在第一次看到输入shape是(batch=4, T=7, 512)时，可以很舒服地做特化编译：把T=7直接折叠进循环上界，allocate固定大小的中间缓冲，把该融合的算子全部融合掉。但第二次请求T变成了9，整个特化全部作废，又要重新编译一遍。
如果生产环境里T的取值分布有二十多种，结果就是二十几份特化代码互相换血，每次切换都伴随一次几百毫秒甚至秒级的编译停顿。更难受的是，编译器在编译期如果发现某个维度的值不可知，通常只能退化成非常保守的代码生成，很多融合优化当场失效，性能甚至不如写得不怎么好的eager代码。
统计下来，这类系统在动态输入下的真实编译命中率往往不到50%——也就是说一大半时间在做无效编译。我在一个用TVM做的实验里验证过，当序列长度集合是{3,4,5,6,7,8,9,10}时，XLA风格的静态编译几乎每个长度都要单独处理，中间切换的开销大于收益，最终端到端只比eager快了一点点，完全对不起编译带来的复杂度。
1.3 解释器慢而灵活，编译器快而僵硬，中间路线怎么走
纯解释执行和纯静态编译是两个极端。解释器的好处是天然支持动态控制流、动态形状，遇到什么shape就处理什么shape，缺点是每条张量指令都要走一遍大循环，做完结果还得写回内存，运算密集度一高就完全发挥不出硬件性能。我手写的块解释器跑一个简单的矩阵乘法链，比起直接用cuBLAS/oneDNN差了可能有四五倍——大量时间耗在指令分派和中间张量的读写上。
实时编译（JIT编译）正好卡在中间：程序先用字节码形式解释执行，解释器同时记录每一条指令经过时的形状分布、调用频率；某个字节码位置被发现是热点且形状模式稳定，就把它编译成高度特化的机器码。以后再来同样的形状，直接命中缓存；发现形状变了，要么编译一个新版本，要么先退回解释器兜底。
这套架构的关键词就三个：字节码VM提供灵活的执行基础，profiler识别热点，JIT编译器把热点路径变成特化代码。训练和推理都能覆盖，推理场景因为形状分布往往更集中，收益尤其明显。
2. 指令集与执行模型：先把张量字节码设计对
2.1 两层指令集：控制流与张量运算分离
设计字节码VM的第一步是定义指令集。我踩过的坑是，别把张量指令和控制流指令混在一起设计。张量指令的语义很重，每个操作都牵扯到shape推断、dtype检查、设备分配，而控制流指令是轻量级的，它们在解释器里几乎不产生实际计算。混在一起会让dispatch逻辑变得极慢。
我最终分成了两层：
| 层次 | 指令示例 | 职责 |
|---|---|---|
| Flow层 | jump、branch、call、ret、loop_head | 管理PC跳转、调用栈、循环状态 |
| Tensor层 | matmul、add、reshape、transpose、layer_norm | 真正的张量运算，带完整shape语义 |
一个动态循环累加的程序，编译出的字节码大概是这样的：
# 伪代码: sum = 0; for i in range(T): sum += x[i] load.const %0, 0 # sum = 0 load.tensor %1, x # 加载输入张量 load.const %2, 0 # i = 0 loop_head: # 动态上界判断 get.shape_axis %3, %1, 0 # 得到第0维长度 T cmp.lt %4, %2, %3 # i < T ? branch.false %4, @exit # 取 x[i]，这里x如果是连续存储可以做指针步进 tensor.index %5, %1, %2 tensor.add %6, %0, %5 move %0, %6 add.const %2, %2, 1 jump @loop_head exit: ret %0注意loop_head里的get.shape_axis，这是在动态形状下保持灵活的根基。每次进入循环体都重新读一次当前轴的长度，这样即使运行期间张量被resize过，循环也不会跑错。代价是每轮多一条shape读取指令，但相比重新编译整个循环，这点开销完全可以接受。
2.2 张量值表示：shape信息要能参与运行时计算
张量的值表示是VM设计里最容易想简单、又最致命的一环。最朴素的表示是void* data + shape + stride + dtype，但这在动态场景下不够用——你还需要一张"shape的自由度图"，告诉VM哪些维度是固定的、哪些维度是符号化的，符号与符号之间有没有约束关系。
我最后采用了带符号shape的TensorValue结构：
struct TensorValue { void* data; std::vector<int64_t> shape; std::vector<int64_t> stride; DType dtype; Device device; // 符号shape：-1表示未知，>=0表示已知 std::vector<int64_t> sym_shape; // 符号约束表，比如 sym[0] == sym[1] };sym_shape数组是关键。假设两个矩阵A、B要相乘，A的shape是(M, K)，B是(K, N)，如果M、K、N都是-1（未知），VM不会立刻报错，而是记录一条约束关系A.shape[1] == B.shape[0]。等运行到matmul指令时再拿实际shape去匹配约束。匹配成功就正常计算；匹配失败才算真正的运行时类型错误。这样设计的好处是，解释执行和JIT编译可以共享同一套shape推断逻辑。
2.3 为什么选寄存器式而不是栈式
如果你做过脚本语言VM，一定纠结过寄存器式和栈式。栈式的优点是字节码紧凑、编译简单，Java VM、Python VM都走这条路。但张量运算的特点是每个操作都涉及大块显存/内存的读写，如果每条指令都从操作数栈顶弹出再压入，会频繁产生中间TensorValue的引用计数切换，最终表现为大量的引用计数原子操作和shallow copy。
寄存器式VM天然接近于SSA形式，每条指令明确写出dst, src1, src2，解释器可以直接把src的指针传给底层kernel，dst的输出位置处在一个可预测的寄存器槽里。后续要做JIT优化时，这种形式几乎可以直接映射到LLVM IR或者等价的三地址代码，不需要再做一遍phi节点重建。
代价是字节码体积变大、编译器需要做寄存器分配。对于张量程序来说，一块显存缓冲比一条字节码指令值钱得多，这个tradeoff非常划算。实际上，TVM的Relax早期设计也走了类似路线，核心IR是函数式带显式shape标注的，只是它没有硬性地落到一个独立的VM指令集上。
3. 实时编译三件套：形状特化、指令融合与后端代码生成
3.1 形状特化与多态版本缓存
实时编译最核心的是形状特化。流程是这样的：解释器在运行每个字节码基本块时，维护一张频率表，记录该基本块被执行了多少次、进入时张量形状的签名（比如M=32,K=512,N=128）。当频率超过阈值——我常用8次——就把这个基本块标记为"可编译块"，并尝试针对那个形状签名编译一份特化版本。
编译出的特化版本放入一个形状签名到编译后代码的哈希表。后续执行到这个字节码位置时，先查一下当前的实际形状签名在不在哈希表里。命中直接跳转到特化版本；没命中就留在解释器继续跑，同时启动一次异步编译。
这里有一个很值得说的细节：特化版本并不只对应一个形状。比如MatMul的形状签名是(M,K) x (K,N)，K在矩阵乘法里是归约维度，K一样但N变化，很多kernel可以复用底层BLAS选择结果，参数化一下就行。可以把签名分成"强特化组"和"弱特化组"：强特化维度变了必须新编译，弱特化维度变了只需修改kernel参数。这个分组能把编译缓存命中率提升得非常明显，我在实验里中位数shape命中率从53%拉到了87%。
3.2 指令融合：把一串字节码变成内核
有了特化代码，接下来就是融合优化。举一个layer_norm的例子，字节码序列一般是：
tensor.mean %1, %0, axis=1 tensor.sub %2, %0, %1 tensor.square %3, %2 tensor.mean %4, %3, axis=1 tensor.add.const %5, %4, eps tensor.sqrt %6, %5 tensor.div %7, %2, %6如果每一条都单独启动一个CUDA kernel，显存里会多出五六个中间缓冲区，而且每个kernel启动都有几十微秒的固定开销。实时编译器要做的就是把这条链检测出来，融合成一个巨kernel：把mean的结果算在寄存器/smem里，不落显存；把square直接折叠进sum；最后rescale一次。
融合边界要非常小心。控制流指令绝不能跨过去，因为融合后的kernel必须是直线代码；依赖外部位点（比如某些全局参数）的指令也必须分开，否则缓存后的特化代码在参数变化时得不到正确结果。我的经验是，编译器内部先构建一张数据依赖图，只做"纯数据流子图"的融合，一旦遇到side-effect指令就立刻截断。
3.3 代码生成后端：LLVM、TinyCodeGen还是自研
实时编译的后端选择直接影响工程复杂度和编译延迟。我实验过三条路线：
- LLVM OrcJIT：优化效果好，能吃到AVX-512、SVE这类新指令集的红利。但编译延迟高，冷启动一个模块经常要几十毫秒到上百毫秒，对于几十毫秒一次的前向推理来说，等不起。
- 生成C代码再调系统编译器：用模板生成一段C++代码，丢给
cc -O2 -shared编译成.so再dlopen。有效但每一步都要起子进程，延迟同样难以控制。 - 轻量级自研代码生成器：只针对自己指令集里的几十条张量指令，直接输出内存中的机器码。编译速度快，几微秒到几百微秒搞定，缺点是优化能力有限，得靠更高层的融合逻辑兜底。
最终线上系统选的是第三条，配合一层基于指令模式匹配的融合优化器。效果是：编译延迟从30ms压到了2ms以内，融合后的内核性能比未融合的特化版本提升平均约2.3倍。如果你的系统里没有特别激进的新指令集需求，这个组合是性价比最高的。
4. 字节码VM在动态张量战场的差异化定位
4.1 和PyTorch系方案比，差在哪、好在哪里
PyTorch的TorchScript和TorchDynamo也存在解决动态形状问题，但路径本质上仍然是"图模式"，先抠出一个静态计算图，再交给后端编译。TorchScript的做法是尽量追踪出一个静态图，遇到动态控制流就抽象成loop/branch子图；TorchDynamo则靠guard机制，每次进入Python函数时检查shape是否和上次一致，不一致就重新编译。
TorchDynamo的问题在于guard机制是"全函数级"的。函数里只要有一个维度是动态的，guard就会频繁失效，导致cache miss和重编译风暴。我在一个多层GRU的动态batch训练实验里，Dynamo的cache miss率一度达到40%——每次miss都要回落到eager模式重新trace，性能反而不如干脆全程eager。字节码VM的优势在于特化粒度是"基本块级"而不是"函数级"：只有真正热点、形状稳定的基本块才会被编译，其他部分留在解释器里，互不干扰。
4.2 XLA和TVM面对动态形状的挣扎
XLA在动态shape上的核心手段是Pad和DynamicReshape，思路是把动态shape的输入pad到某个上限，让kernel逻辑变成静态的。这个方案在推理场景还勉强能用，但训练场景会遇到问题：pad出来的部分是无效计算，反向传播时还得专门mask掉，计算浪费可达30%-70%。而且一旦用户给的上限不准确，超出部分会直接报错，非常不友好。
TVM的Relax尝试在IR层面引入shape表达式，把"形状"本身变成一等公民参与编译决策。思路很优美，但工程复杂度过高。我当时评估过自己搞一套Relax-like IR要多久——排期下来一个四人小组至少四个月。字节码VM这条路就务实得多，不追求一次成功，而是"解释执行保证正确，缓存命中保证性能"，上线负担小很多。
4.3 渐进式执行：不编译也能活，编译只是加速
字节码VM最独特的地方在于它的"渐进式"特质。这意味着系统上线不存在"要么全有要么全无"的风险。如果生产环境的形状分布极端发散，特化缓存命中率上不去，系统会自动退化成纯解释执行，性能顶多比eager慢个20%，绝不会出现编译风暴把服务打挂的情况。
也正是这种特性，让我敢把这套VM直接用在线上推理服务里。它像一个带涡轮增压的引擎，涡轮不介入的时候就是普通自然吸气，一旦转速到位涡轮介入，就能获得明显性能提升。对于动态张量计算这种"大部分时间形状稳定、偶尔有长尾"的负载，没有比这更合适的形态了。
5. 上线前必须啃下的工程硬骨头
5.1 shape推断在动态场景里的隐藏大坑
动态shape下做编译优化，最脏的活永远是shape推断。难点在于，当两个符号维度看上去可能相等、但编译器无法在编译期证明时，你只有两条路：放弃优化，或者加入运行时断言。
我踩过一个特别典型的坑：计算Attention Score时，Q和K的sequence长度在绝大多数情况下相等，因此在编译期直接假设了它们相等，做了内存布局的合并。结果某天线上来了一条极端数据，Q长度和K长度不匹配，程序直接segfault。修法是给这个假设加一条运行时shape断言assert_eq(q_len, k_len)，失败就跳转到通用版本重算，而不是继续跑特化代码。从那以后，我给自己定了一条规矩：动态shape的特化编译必须携带约束表，约束表验证失败就回退，宁可不优化也不出错误结果。
5.2 副作用指令会让特化代码状态错乱
张量计算里有不少带副作用的操作：in-place update、随机数生成、显存池的显式分配释放、某些算子内部的全局状态。这些指令如果被无脑融合进特化版本，会导致编译后代码和解释器执行结果产生分歧。
举一个真例。某个模型在循环体里反复调用随机dropout，我最初把dropout和后面的matmul融合成一个kernel，看起来没问题。但训练时梯度下降需要dropout的mask是可复现的——不同batch的同一位置要产生确定性mask。融合后由于运算顺序变化，串行随机流被打乱，复现实验直接失败。最终的做法是：所有带随机状态的算子都不参与融合，保留独立的kernel调用，只在它们周围做shape特化。
5.3 编译失败后的回退要静默且稳妥
再健壮的编译器也有veto的情况：某个shape组合超出发射器的处理范围、某些指令还没移植到JIT后端、或者运行时资源不足。我见过很多项目在这些场景下直接抛异常，导致整条推理链路崩溃。我的做法是给每个编译单元一个"编译失败计数"，连续失败三次就永久标记为"不编译"，让这段字节码永远跑解释器，并输出一条warning级别的日志。
回退策略的另一个要点是：回退必须在字节码级别的同一位置进行，不能出现"前半段跑特化、后半段跑解释器"这种割裂状态。所以我在JIT入口设计了一个两阶段提交：先验证当前状态与编译时的假设完全一致，再切换PC到特化代码的入口；任何一步验证不过都留在解释模式。
5.4 编译器线程不能拖慢推理主线程
实时编译的延迟会直接影响端到端延迟，尤其是推理场景。我的调优心得是，把编译完全放到独立线程池里，主解释器只做"发布编译任务"和"查询编译结果"。解释器每到达一个热点块，先向编译线程池提交任务，然后继续解释执行；编译完成后结果被放回缓存区，下次执行到该位置时再启用。
还需要注意编译器线程池大小和CPU核心数的关系。编译过程是CPU密集的，如果线程池设太大，会和主线程抢核，导致推理延迟反而上涨。我的经验值是主服务部署在16核机器上，线程池设4，既能并行编译多个shape版本，又不至于把CPU吃满。实践中还可以结合profile阶段的采样频率来动态调整线程池大小。
6. 实测数据与调优心得
6.1 变长序列Transformer推理：从45ms降到28ms
拿一个经典的变长序列Transformer decoder推理做基准，输入batch=32，样本长度在5到60之间变化，使用RTX 3090。同一模型分别在eager模式、纯解释器VM、字节码VM+实时编译三种状态下跑，各跑5000次取中位数延迟：
| 执行模式 | 中位延迟 | 相对Eager |
|---|---|---|
| Eager（PyTorch原生） | 45ms | 1.0x |
| 纯解释器字节码VM | 58ms | 0.78x（慢） |
| 字节码VM + 实时编译 | 28ms | 1.6x（快） |
解释器慢是预期内的，它存在的意义本来就是兜底和采集profile。实时编译带来的提升主要来自两块：约65%的算子被融合成巨kernel，kernel启动次数从每层十几次降到了两三次；另外，形状特化后能提前锁定内存池里具体大小的缓冲块，避免动态分配，这部分省了约20%的时延。
6.2 动态batch的RNN训练：编译预热期的等待
训练场景的时间线值得单独记录。前8个step基本都在解释执行，profiler在后台收集基本块的形状分布。第9个step开始，第一个隐藏层基本块的特化编译完成，之后每隔几个step就多一份特化版本。从第25个step开始，所有热点基本块都编译完成，单step时间从解释期的18ms降到了8ms，加速比约2.2倍。
需要提前做好心理准备的是编译期的前30个step会有明显波动，某些step甚至会因为后台编译器线程抢占CPU而比解释期还慢。解决方法是把编译线程优先级调低，同时把"是否开启实时编译"做成运行时开关，在训练的前几个epoch先纯解释执行收集统计，等形状分布稳定了再打开开关。
6.3 几个直接影响收益的参数
最后分享几个调参经验，都是我用控制变量法试出来的：
- 热点阈值：我默认8次。阈值太高，小batch请求还没触发编译就结束了，收益为零；阈值太低，只有两三次调用也去编译，编译开销比解释执行还大。在推理场景可以降到3-4次，因为热点的形状分布通常更集中。
- 特化缓存容量：默认1024个shape版本，LRU淘汰。容量太小，长尾形状频繁挤掉核心版本；容量太大，哈希查找本身开始有可察觉开销。1024够覆盖90%以上的线上shape分布。
- 禁止编译名单：对于延迟极度敏感、且形状分布确实非常发散的场景，直接禁用实时编译走纯解释，反而更稳定。我写了一个运行时自适应开关，连续20次查询到特化缓存miss且编译完成率低于30%，就自动降级为纯解释模式。
我个人在实际操作中的体会是：这类系统最忌讳一上来就追求完美编译，应该先把字节码解释器跑通，让业务正常运作，再把profiler和JIT逐步加上去。每加一层优化都盯住端到端延迟的P99和缓存命中率，一旦发现收益为负就果断撤掉，用数据说话。动态张量的优化没有银弹，但字节码虚拟机加实时编译的组合，确实是目前我在生产环境里验证过、最能在灵活性和性能之间取得平衡的方案。
