
第30题自研推理框架、动态维度、模型导出、算子与计算图1. 核心回答如果让我设计一个自研推理框架我会把整个系统拆成四层模型导出 → 统一计算图 IR → 图优化与算子 Lowering → 后端 Kernel 执行核心流程可以表示为Framework Model→Exported Graph→Unified IR→Graph Optimization→Backend Lowering→Kernel Execution \text{Framework Model} \rightarrow \text{Exported Graph} \rightarrow \text{Unified IR} \rightarrow \text{Graph Optimization} \rightarrow \text{Backend Lowering} \rightarrow \text{Kernel Execution}Framework Model→Exported Graph→Unified IR→Graph Optimization→Backend Lowering→Kernel Execution其中需要重点解决四个工程问题如何把 PyTorch 等训练框架中的模型稳定导出如何用统一 IR 表达算子、Tensor 和计算依赖如何处理 Batch Size、Sequence Length 等动态维度如何处理后端暂时不支持的算子。整个设计的基本原则是先保证模型语义正确再进行图级和算子级性能优化。2. 整体推理框架怎么设计2.1 前端模型导出训练模型通常来自PyTorchTensorFlowONNX其他模型框架。自研推理框架首先需要把这些不同模型转换成自己的统一表示。例如 PyTorch 模型可以通过torch.exportONNX Export获得显式计算图。导出阶段需要保存的不只有算子名称还包括TensorTensor dtypeTensor shape算子类型算子属性常量输入输出关系数据依赖动态维度约束。例如x ↓ MatMul ↓ Add ↓ GELU ↓ y可以在内部表示为Tensor x ↓ MatMul(x, weight) ↓ Add(bias) ↓ GELU ↓ Tensor y这样后续优化阶段可以直接分析整张计算图。3. 什么是计算图和 IR3.1 计算图计算图可以抽象成G(V,E) G(V,E)G(V,E)其中VVV表示算子节点EEE表示 Tensor 数据依赖。例如YReLU(XWb) Y\operatorname{ReLU}(XWb)YReLU(XWb)可以转换成X ----\ MatMul → Add → ReLU → Y W ----/ ↑ b计算图让推理框架能够分析哪些算子可以融合哪些常量可以提前计算Tensor 生命周期是什么哪些内存可以复用哪些算子可以并行。3.2 IR 的作用IR 即 Intermediate Representation中间表示。它位于PyTorch / ONNX ↓ Unified IR ↓ CUDA / CPU / NPU统一 IR 的价值在于把前端模型格式与底层硬件解耦。例如一个 MatMulmatmul(A, B)在不同后端可以 Lower 成CUDA → cuBLAS GEMM CPU → oneDNN GEMM NPU → Vendor MatMul Kernel这样增加新硬件时主要增加新的 Backend 和 Kernel 映射不需要重新设计整个模型前端。4. 图优化阶段做什么模型成功转换到 IR 后可以执行一系列 Graph Pass。4.1 常量折叠例如yx(2×3) yx(2\times3)yx(2×3)其中2×36 2\times362×36可以在编译阶段直接计算。最终图变成yx6 yx6yx6这样运行时减少一次计算。4.2 Dead Code Elimination如果某个计算结果最终没有参与模型输出A → B → C \ D其中 D 永远不会影响最终结果就可以删除 D。4.3 算子融合例如MatMul ↓ Bias Add ↓ GELU可以尝试融合成Fused MatMul Bias GELU减少Kernel Launch中间 Tensor显存读写内存带宽开销。很多模型推理优化的收益来自减少内存访问而不只是减少 FLOPs。5. 动态维度怎么处理动态 Shape 是自研推理框架中比较困难的一部分。例如语言模型输入X∈RB×L×d X\in\mathbb{R}^{B\times L\times d}X∈RB×L×d其中BBBBatch SizeLLLSequence LengthdddHidden Dimension。通常ddd固定但B, L B,\ LB,L可能在运行时变化。例如[1, 128, 4096] [8, 512, 4096] [16, 2048, 4096]都可能输入同一个模型。6. Symbolic Shape 怎么设计6.1 使用符号表示动态维度内部可以把 shape[batch, seq, 4096]表示成[B,L,4096] [B,L,4096][B,L,4096]其中BBB和LLL是 Symbolic Dimension。例如矩阵X:[B,L,4096] X:[B,L,4096]X:[B,L,4096]经过 LinearW:[4096,11008] W:[4096,11008]W:[4096,11008]则输出可以推导为Y:[B,L,11008] Y:[B,L,11008]Y:[B,L,11008]这里无需提前知道BBB和LLL的具体值。这就是 Shape Inference。6.2 Shape Constraint动态维度通常仍然需要约束。例如1≤B≤64 1\leq B\leq641≤B≤641≤L≤8192 1\leq L\leq81921≤L≤8192还可能存在维度关系A.shape[0]B.shape[0] A.shape[0]B.shape[0]A.shape[0]B.shape[0]这些约束需要记录在 IR 或编译环境中。PyTorchtorch.export当前也采用 Symbolic Dimension 和 Range Constraint 表达动态 Shape。7. 为什么动态 Shape 不能完全无限制如果所有 Shape 都完全未知很多编译期优化就无法完成。例如 GPU Kernel 的最优实现可能依赖Batch SizeSequence LengthMatrix ShapeAlignmentWorkspace Size。因此实际工程中通常会使用有限范围动态。例如MIN [1, 1, 4096] OPT [8, 512, 4096] MAX [32, 4096, 4096]TensorRT 的 Optimization Profile 就采用min / opt / max描述动态输入范围。optShape 用于重点选择高性能执行策略。因此可以把动态 Shape 分成Symbolic Shape ↓ Shape Constraint ↓ Optimization Profile ↓ Runtime Shape8. Runtime 怎么处理动态 Shape运行时收到Input Shape [6, 700, 4096]框架首先检查shape∈valid range shape\in valid\ rangeshape∈validrange如果合法绑定 Symbolic Shape完成剩余 Shape 推导选择对应执行计划分配 Workspace调用对应 Kernel。如果输入 Shape 超过已有优化范围可以根据框架设计选择其他 Optimization Profile触发新的 specialization重新编译回退到通用执行路径返回明确的 Shape Error。具体策略取决于实时性和部署要求。9. 动态 Shape 的难点是什么9.1 内存规划困难静态 Shape 下Tensor A 10 MB Tensor B 20 MB Tensor C 5 MB编译阶段就能准确完成 Memory Planning。动态 Shape 下size(A)f(B,L) size(A)f(B,L)size(A)f(B,L)实际大小只能在运行时确定。因此需要最大 Shape 预分配Runtime AllocationMemory PoolWorkspace Reuse等策略。9.2 Kernel 选择困难某个 Kernel 可能在M128时最快。另一个 Kernel 可能在M4096时最快。因此通常需要shape→kernel shape\rightarrow kernelshape→kernel的映射策略。可以通过heuristicautotuningprofilekernel cache选择执行方案。10. 不支持的算子怎么处理模型导出后经常会遇到Unsupported Operator例如模型中存在自定义算子CustomAttention后端没有对应 Kernel。通常有三种处理方式。10.1 算子分解优先检查高层算子能否拆成已有基础算子。例如CustomLinear可以分解成MatMul Add如果基础算子已经受到后端支持就无需新增 Kernel。这种过程可以理解为Ophigh→{Oplow1,Oplow2,...} Op_{high} \rightarrow \{Op_{low}^{1},Op_{low}^{2},...\}Ophigh→{Oplow1,Oplow2,...}MLIR 的 Dialect Conversion 也提供类似的 Lowering 和 Rewrite 机制。10.2 自定义 Kernel如果算子无法有效分解可以实现自己的 Kernel。例如FlashAttentionOp对应CUDA FlashAttention Kernel然后注册IR Operator ↓ Backend Operator ↓ Kernel需要同时定义输入输出类型Shape Inferencedtype 支持Kernel 参数Workspace错误检查。10.3 Backend Fallback如果某个算子暂时没有优化实现可以回退到已有运行时。例如自研 CUDA Backend ↓ Unsupported Op ↓ PyTorch / ONNX Runtime / CPU FallbackFallback 可以提高模型覆盖率。代价可能包括Host/Device CopyTensor Layout 转换额外同步延迟增加。因此生产环境需要记录Fallback Rate如果大量算子都发生 fallback自研推理框架通常无法获得预期性能收益。11. Backend 和 Kernel 怎么组织可以设计统一 Backend InterfaceBackend ├── CPUBackend ├── CUDABackend └── NPUBackend每个 Backend 注册自己支持的 Operator。例如MatMul ├── CPU → oneDNN ├── CUDA → cuBLAS └── NPU → Vendor GEMM编译阶段根据Devicedtypeshapelayoutoperator attributes选择 Kernel。最终形成OperatorShapeDTypeDevice→Kernel Operator Shape DType Device \rightarrow KernelOperatorShapeDTypeDevice→Kernel12. 内存规划怎么做推理性能不仅由计算速度决定。GPU 推理中频繁cudaMalloc cudaFree会产生明显开销。因此编译阶段可以分析 Tensor Lifetime。例如Tensor A: Node1 → Node3 Tensor B: Node4 → Node6如果 A 在 Node3 后已经死亡而 B 在 Node4 才创建就可以让Memory(A)Memory(B) Memory(A)Memory(B)Memory(A)Memory(B)复用同一块显存。这一过程叫Memory Planning / Memory Reuse最终可以减少峰值显存Allocation 次数内存碎片。13. 怎么验证自研推理框架是正确的我会从四个维度验证。13.1 数值正确性比较Yreference Y_{reference}Yreference与Yengine Y_{engine}Yengine例如计算Errormax∣Yreference−Yengine∣ Error \max|Y_{reference}-Y_{engine}|Errormax∣Yreference−Yengine∣针对 FP32、FP16、BF16、INT8 分别设置合理容差。13.2 模型覆盖率统计CoverageSupported ModelsTotal Models Coverage \frac{Supported\ Models} {Total\ Models}CoverageTotalModelsSupportedModels同时报告Unsupported Op 数量Fallback RateDynamic Shape 覆盖范围自定义算子数量。13.3 性能需要分别测LatencyP50 / P95 / P99 LatencyThroughputTime to First TokenTokens/sCompilation Time。性能比较必须保持相同硬件相同模型相同精度相同输入 Shape相同 Batch Size。13.4 资源消耗还需要测Peak GPU MemoryWorkspaceCPU MemoryGPU UtilizationKernel Launch 数量。这样才能判断优化收益来自什么地方。14. 如果面试官问“你真正做了哪些东西”回答时需要明确区分模型导出 IR设计 Shape推导 Graph Pass Operator Lowering Kernel实现 Memory Planning Runtime Benchmark然后指出自己真正负责的模块。例如我主要负责 Dynamic Shape 和 Operator Lowering。模型导出后我在 IR 中使用 Symbolic Dimension 表示动态 Batch 和 Sequence Length并维护维度范围及等式约束。编译阶段先做 Shape Inference再根据 Shape Profile 生成不同执行策略。对于后端不支持的高层算子我优先通过 Rewrite Pattern 分解成基础算子无法有效分解时再注册自定义 CUDA Kernel。最后通过跨 Shape 数值对齐测试、模型覆盖率、P95 延迟和峰值显存验证实现。这种回答能够清楚体现实际工程边界。15. 面试时可以压缩成下面这段如果让我设计一个自研推理框架我会采用“模型导出、统一 IR、图优化、Backend Lowering、Kernel Runtime”这条链路。前端通过torch.export或 ONNX 获得计算图然后转换成内部统一 IR。IR 中保存 Operator、Tensor、dtype、shape、attribute、常量以及数据依赖。静态 Shape 可以在编译阶段完成 Shape Inference 和内存规划。对于动态 Batch 或 Sequence Length我会使用 Symbolic Shape 表示并维护 Range Constraint。编译时根据常见 Shape 建立 Optimization Profile运行时绑定真实 Shape再进行 Kernel 选择和 Workspace 分配。图优化阶段执行常量折叠、无效节点删除、算子融合和内存复用。后端不支持的算子优先做 Operator Decomposition 或 Lowering确实需要时实现 Custom Kernel暂时无法支持时提供 Fallback并监控 Fallback Rate。最后从四个维度验收数值正确性、模型和算子覆盖率、延迟与吞吐、显存和 Workspace。整个过程首先保证与原模型语义一致然后再判断优化是否真正带来性能收益。16. 来源PyTorch Documentation,torch.export模型导出、Dynamic Shapes、Symbolic Dimension 与 Range Constraint。ONNX Documentation,Shape Inference静态维度、符号维度及计算图 Shape 推导。NVIDIA TensorRT Documentation,Dynamic ShapesRuntime Dimension 与 Optimization Profile 的min / opt / max机制。NVIDIA TensorRT Documentation,Dynamic Shapes: Advanced Topics动态 Shape 的部分算子约束和 Shape Tensor。MLIR Documentation,Dialect Conversion通过 Conversion Target、Rewrite Pattern 和 Type Converter 实现 Operator Lowering。