3、动态形状的底层原理：Tensor的内存布局、动态分配机制、CPU与GPU上的差异
好，咱们今天聊点硬核的。动态形状这个东西，很多同学在框架层面用得挺溜，但一问到底层怎么实现的，就有点懵了。我当年刚接触动态形状时也踩过不少坑，尤其是内存分配那部分，真的是「一顿操作猛如虎，一看显存二百五」。
说白了，动态形状的核心挑战就一个：张量的大小在运行时才能确定。这跟静态形状完全不同——静态形状下，编译器可以提前把内存安排得明明白白；动态形状下，你得在推理过程中「随机应变」。嗯，这里面的门道，咱们一层层剥开。
3.1 Tensor的内存布局：连续性与步长
先说说Tensor在内存里到底怎么放的。你想想看，一个多维数组，在物理内存里其实是一维的。怎么映射？靠的就是步长（stride）。
举个例子，一个形状为 [batch, 3, 224, 224] 的张量，如果它是连续（contiguous）的，那它的步长就是 [3*224*224, 224*224, 224, 1]。访问 [i, j, k, l] 这个元素时，内存地址就是 base + i*stride[0] + j*stride[1] + k*stride[2] + l*stride[3]。
但动态形状下，事情就复杂了。比如你做了个切片操作：
# 假设 x 形状是 [8, 3, 224, 224]
y = x[2:5, :, :, :] # 形状变成 [3, 3, 224, 224]
这时候 y 的内存布局还是跟 x 共享的，只是步长变了。我建议你记住一个原则：视图操作不改变内存布局，只改变步长和形状。但一旦你做了 y.contiguous() 或者 y.reshape()，那就可能触发内存拷贝。
关键点：动态形状下，频繁的 reshape 和 transpose 会导致内存碎片化。我在项目中遇到过，一个模型推理时显存占用忽高忽低，查了半天发现是某个动态形状的切片操作没做 contiguous()，导致后续算子访问内存时 cache miss 严重。
3.2 动态分配机制：谁在背后帮你打理内存？
动态形状的内存分配，说白了就是「用多少，申请多少」。但这里有个性能陷阱——频繁的内存分配和释放。
在CPU上，malloc/free 的开销其实还好，因为操作系统有内存池。但在GPU上，cudaMalloc/cudaFree 可是重量级操作，一次调用可能就要几十微秒。你想想看，如果每个动态形状的张量都单独分配，那推理延迟直接爆炸。
那怎么办？业界主流做法是内存池（Memory Pool）。以PyTorch的 CUDAAllocator 为例：
- 它会预先申请一大块显存（比如2GB）
- 然后内部维护一个空闲块列表
- 当需要分配时，从空闲块中找大小合适的
- 释放时，不真正还给系统，而是放回空闲列表
我曾经调试过一个线上问题：某个动态形状模型跑着跑着就OOM了。后来发现是内存池里的块大小不匹配，导致大量碎片。解决方案？嗯，我建议你设置 PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128，限制最大分裂块大小，能有效缓解碎片问题。
个人经验：动态形状下，尽量复用张量。比如用 torch.empty() 代替 torch.zeros()，减少不必要的清零操作。另外，torch.cuda.empty_cache() 这玩意儿少用，它会清空整个内存池，反而拖慢后续分配。
3.3 CPU与GPU上的差异：一个天一个地
CPU和GPU在动态形状处理上，差异真的很大。我列个表，你一看就明白：
| 维度 | CPU | GPU |
|---|---|---|
| 内存分配速度 | 快（微秒级） | 慢（几十微秒到毫秒级） |
| 内存碎片影响 | 较小（OS有虚拟内存） | 严重（物理显存连续） |
| 动态形状开销 | 主要在访存模式 | 主要在kernel launch和内存分配 |
| 典型优化手段 | 数据对齐、预取 | 内存池、kernel fusion |
为什么会这样？核心原因在于GPU的并行执行模型。GPU上成千上万个线程同时跑，如果每个线程访问的动态形状张量地址不同，那内存访问模式就是「随机散落」的，这会导致严重的bank conflict和全局内存合并失败。
我记得有一次优化一个动态形状的Transformer模型，在CPU上跑得好好的，一上GPU就慢得离谱。后来用NVIDIA Nsight分析，发现是动态的序列长度导致attention score矩阵的维度变化，GPU的warp调度完全乱掉了。解决方案？把动态的序列长度padding到8的倍数，让warp size对齐。
避坑指南：我曾经在GPU上直接用 torch.where() 做动态形状的mask操作，结果显存直接爆了。原因是 torch.where() 会为两个分支都分配内存，即使其中一个分支不会被用到。后来改用 torch.masked_fill() 才解决。
3.4 实战建议：如何写出动态形状友好的代码？
说了这么多原理，最后给几条实操建议：
- 尽量保持连续性：动态形状操作后，如果后续有密集计算，记得调用
.contiguous()。虽然会多一次拷贝，但能换来更好的访存性能。 - 预分配内存：如果知道动态形状的上限，可以预先分配最大尺寸的张量，然后通过切片来使用。比如
buffer = torch.empty(max_batch, ...)，每次只用前actual_batch行。 - 避免在GPU上频繁分配：把动态形状的分配逻辑尽量放到CPU上，或者用
torch.cuda.Stream异步处理。 - 监控内存碎片：用
torch.cuda.memory_summary()查看显存使用情况，如果碎片率超过30%，就要考虑调整内存池参数了。
嗯，动态形状的底层原理就聊到这儿。说白了，就是「内存布局 + 分配策略 + 硬件特性」这三者的博弈。你只要记住：CPU上动态形状的代价是访存模式，GPU上动态形状的代价是分配和调度。下次遇到动态形状的性能问题，先定位瓶颈在哪一端，再对症下药。
一句话总结：动态形状不可怕，可怕的是你不知道它底层怎么折腾内存。搞懂了Tensor的步长、内存池的机制、CPU和GPU的差异，你就能写出又快又稳的动态形状推理代码。