动态 Shape 与多尺寸分档在推理编译器中的实现在大语言模型LLM与多模态扩散模型Diffusion的在线推理服务中输入张量的维度是高度动态的用户输入的 Prompt 长度可能从 1 个 Token极短提问到 8192 个 Token超长文档问答之间剧烈波动并发请求的 Batch Size 随着在线流量洪峰在 1 到 64 之间动态调整。传统的 AI 编译优化器如 TensorRT 早期的静态图优化、TVM 早期模式通常假设张量形状是严格编译期固定的静态常量。如果每次收到一个新的动态维度都触发一次运行时 JIT 重新编译编译耗时会直接导致在线服务不可用而如果完全放弃编译优化、退化为纯解释执行又会白白浪费 GPU 50% 以上的硬件峰值算力。多尺寸分档Bucket Tiling / Bucketing Strategy结合符号推导动态 Shape 计算图Symbolic Shape Graph是工业级 AI 编译器在灵活性与性能之间达成最佳平衡的核心工程范式。-------------------------------------------------------------------------- | 动态 Shape 多尺寸分档 (Bucketing) 路由架构 | -------------------------------------------------------------------------- | 用户动态输入 Prompt (实际长度: SeqLen 105 个 Token) | | | | v | [智能分档路由器 (Dynamic Bucket Router)] | | 预定义分档阶梯: [32, 64, 128, 256, 512, 1024, 2048, 4096] | | - 向上对齐路由至最近档位: Bucket 128 (仅需 Padding 23 个零) | -------------------------------------------------------------------------- | v 命中预编译好的静态专用 Kernel | [Bucket 128 编译优化执行引擎] | | - 享受 100% 静态 Tile 优化、指令级并行展开与最大 Tensor Core 算力利用率! | | - 在最终输出层自动截断 (Slice 0..105)对上层用户呈现完全透明! | --------------------------------------------------------------------------1. 为什么纯动态 Shape 难以榨干 GPU 性能在 GPU 体系结构中算子之所以能跑出极致性能关键在于编译器对矩阵乘法GEMM做出的精确分块Tiling如果已知矩阵尺寸 $M128, N4096, K4096$编译器可以精确把线程块划分为 $128 \times 128$ 的网格让每个线程刚好计算 8 个元素寄存器利用率达到 100%Shared Memory 没有任何 Bank Conflict 冲突但如果 $M$ 是一个未知变量Dynamic Shape编译器无法在编译期确定循环迭代次数无法进行激进的循环展开Loop Unrolling无法静态分配固定大小的 Shared Memory 缓冲区只能采用保守的通用分块策略实测表明未分档的纯动态泛型 Kernel 性能通常比针对特定尺寸专门编译的静态 Kernel 慢 30% ~ 50%2. 多尺寸分档Bucketing的工程落地策略Bucketing 的核心思想是将连续无限的动态输入空间离散化为有限个经过高度优化的几何阶梯分档阶梯设计Bucket Scale Design通常采用非线性对数阶梯例如对于序列长度[64, 128, 256, 512, 1024, 2048, 4096, 8192]在短文本区间 512分档较为密集减少 Padding 比例在长文本区间 2048分档适度稀疏控制总分档数量离线 AOT 批量调优Multi-Profile Auto-Tuning在离线编译阶段编译器针对这 8 个固定的 Bucket Profile 分别执行激进的 Auto-tuning 搜索生成 8 组极致优化的 Cubin 二进制 Kernel运行时零拷贝路由与截断Routing Slicing当实际长度为 105 的请求到达时系统将其路由至Bucket 128在输入 Tensor 末尾追加 23 个零Padding运行针对 128 优化的极速 Kernel在最终获取输出 Logits 时直接在显存中取前 105 个元素的切片视图零拷贝丢弃无用的 Padding 部分。3. 符号形状推导Symbolic Shape Inference进阶对于多层复合算子网络编译器前端还需构建符号推理引擎Symbolic Shape Evaluator#[derive(Debug, Clone, PartialEq, Eq)] pub enum Dim { Static(usize), Dynamic(String), // 符号变量如 batch_size, seq_len } pub struct SymbolicTensor { pub shape: VecDim, } impl SymbolicTensor { /// 编译期符号推导矩阵乘输出维度: [B, M, K] x [B, K, N] - [B, M, N] pub fn matmul_shape(self, rhs: SymbolicTensor) - ResultSymbolicTensor, static str { if self.shape.len() ! 3 || rhs.shape.len() ! 3 { return Err(Rank mismatch); } // 校验 Batch 维与 K 维符号一致性 assert_eq!(self.shape[0], rhs.shape[0]); assert_eq!(self.shape[2], rhs.shape[1]); Ok(SymbolicTensor { shape: vec![self.shape[0].clone(), self.shape[1].clone(), rhs.shape[2].clone()], }) } }通过将动态符号推导与运行期多尺寸分档完美融合现代 AI 编译器彻底摆脱了动态输入的性能枷锁在多变复杂的生产流量中始终保持着如钢铁般的极致计算效率。