用虚拟 Tensor 消除 DNN 数据搬运(OSDI 2026)

原题:VTC: DNN Compilation with Virtual Tensors for Data Movement Elimination

一句话总结:Transpose、Split、ScatterND 等纯数据搬运算子只改变元素位置,却会把中间结果完整写回再读出;VTC 用“物理 tensor 指针 + index mapping”表示不落地的 virtual tensor,再通过 profile-guided VTOG 搜索选择值得虚拟化的路径,在五类模型组件、A100/H100、batch 1/16 上相对每项最快 compiler 最多快 1.93 倍、论文报告平均快 1.28 倍,但完整模型训练、动态 shape、多 GPU 和 greedy 最优性都没有验证。

问题与动机

GPU 的计算能力增长快于显存带宽。DNN compiler 已能让 MatMul、attention 等计算算子高效使用 Tensor Core,但 Transpose、Split、ScatterND、Slice、Expand 等算子不做数值计算,仍要在 global memory 中读出、重排、写回整块 tensor。随着模型越来越 memory-bound,这些搬运可能比相邻计算更慢(图 1)。

现有 layout optimization 通常只探索 Reshape、Transpose 等少数 layout;operator fusion 又依赖能否把相邻计算合成一个 kernel,或依赖手写 pattern。Llama 3 8B decode 是一个具体例子:TensorRT 已把 QKV projection 合成一个 MatMul,也把 attention 做成 FlashDecoding,却仍用 6 个 kernel 完成中间的 Split、KV-cache update、Expand 和 layout change。这个区段要 0.9083 ms,比前后两个 compute kernel 的 0.7926 ms 总和还长(图 2)。

VTC 的核心问题是:中间 tensor 是否一定要有一块完整、连续的 global-memory storage?作者观察到,Tensor Core 等计算单元要求进入 on-chip buffer 的 tile 连续,却不要求整个 global tensor 都连续。只要 load/store stage 能根据 index mapping 访问多个物理区域,计算主体可以保持不变,中间搬运算子也可以消失。

关键观察 / 隐含假设

  • 观察 1:纯数据搬运算子可以写成 output index 到唯一 input element 的映射。 它改变位置而不改变 value,因此可保存 mapping,而不保存一份新 tensor(定义 2–3、§2.2、§4.1)。
    • 依赖假设:算子语义可由 VTC 支持的 mapping rule 表达,不涉及 reduction、数值计算、不可见 side effect 或必须 materialize 的 graph input/output。
    • 可能失效场景:custom operator、复杂 alias/in-place update、动态 value-dependent layout,或外部 library 必须接收普通 contiguous pointer。
  • 观察 2:global memory 只需“局部连续”也能 coalesce。 Llama 的 head dimension 为 128,大于 warp size 32;QKV output 分写到 Q 与 KV-Cache 时,每个连续 chunk 仍足够宽,不必先生成完整 QKV tensor(图 3、定义 4–5)。
    • 依赖假设:连续 chunk 大于最小 memory transaction,mapping branch 不发散,额外 address arithmetic 小于省掉的 read/write。
    • 可能失效场景:chunk 很小或 irregular、映射跨许多物理 tensor,或 TMA/CUTLASS 对 tile layout 有更强约束。
  • 观察 3:虚拟化不是越多越好。 非连续访问会让 compute kernel 变慢;图 2 的 QKV MatMul 就从 TensorRT 的 0.1034 ms 变成 VTC 的 0.1403 ms。只有省掉的搬运大于这个回归时,整体才更快。
    • 依赖假设:目标 GPU 上的 profile 能代表 deployment configuration;compiler 可以在收益为负时保留 materialization。
  • 观察 4:数据搬运机会在新模型结构中更大。 A100 breakdown 中,TensorRT 的 data-movement 占比在 Gemma batch 16 达 72%,而部分 CNN 只有约 13%–28%;论文相应看到 Transformer 的平均收益更高(图 10)。
    • 依赖假设:所选 decoder layer/attention block 能代表完整应用的热点,而不是只代表一个有利子图。
  • 假设 1:greedy profiling 能找到足够好的全局组合。
    • 证据强度:中。五类组件上结果良好,但一般 latency set function 不是 submodular,论文没有 approximation bound 或与 exhaustive optimum 的对比(§5.2)。

核心方法

1. Virtual tensor 只保留逻辑 index space。 一个 virtual tensor 写成 (F, P1, ..., Pn)P 是实际物理 tensor 指针,F(x) 告诉访问逻辑位置 x 时应去哪个 Pj[y]。mapping 可以组合,因此连续的 Reshape、Split、ScatterND 不必每步都 materialize。computation-graph 的 input/output 保持物理 tensor,VTC 只虚拟化中间结果(定义 3、算法 1)。

2. 只改 kernel 的 global-memory I/O。 普通 kernel 可分为 global→on-chip load、on-chip compute、on-chip→global store 三段。VTC overload tl.load/tl.store,根据已知 mapping 生成专用 address code;MatMul 或 attention 的计算主体不变。producer 可以直接把 merged QKV output 写入 Q 与 KV cache,consumer 也可以从分散的 cache 直接读入连续 on-chip tile(图 3、§4.2)。

3. 用连续性判断明显有利的机会。 fully contiguous mapping 被归为 Type I,论文定理 1 称它总是有利;partial-contiguous mapping 在 chunk 足够宽时通常有利;其他 Type II 情况交给 profile。VTC 为 Split、Expand、ScatterND 等算子手写 mapping 与双向虚拟化规则。ONNX 中数据搬运算子数量有限,论文举例为 18 个,但这也意味着新算子仍要开发者补规则(§4.3–4.4)。

4. Virtual Tensor Opportunity Graph(VTOG)枚举机会。 fusion 后的 computation graph 中,每个 tensor 成为 VTOG node;若删掉一个 data-movement operator 后 u 可由 v 表示,就添加 u→v edge。某些 outgoing edge 会让同一逻辑 index 指向互斥的 storage,系统将它们记为 conflict set。删除冲突 edge 后得到 points-to graph,它一一对应一套全局 virtual/materialized 策略(图 6–8、§5.1)。

5. Global greedy 选择实际策略。 全局最优需要对一般、只能实测的 latency function 做指数搜索,论文指出问题为 NP-hard。VTC 从必须物理化的 anchor 开始,每轮选择当前非冲突 marginal saving 总和最大的 node/edge set;加入新 virtual tensor 后,重新 profile 指向它的 edge。最大 saving 为负时停止,因此系统可以完全跳过一个不利机会。算法控制逻辑约为二次复杂度,但硬件 profiling 才是主要编译成本(算法 2、§5.2)。

6. 编译器集成。 VTC 在 TorchInductor 加 analysis pass 和 transformation pass:前者构造 VTOG/points-to graph,后者删除搬运 node、修改 IR 的 load/store,再交给 Triton 生成专用 Python/CUDA code。单次被 profile 的 configuration 少于 10 s,所测完整 graph 的总编译时间少于 10 min。作者声称与原 compiler 保持端到端数值等价,但正文没有给误差阈值或 differential-test 规模(§6)。

设计取舍

  • 少 materialization 换地址计算。 节省 global read/write 和 allocation,却可能让 producer/consumer 的 memory access 更散,并增加 mapping、branch 与 register 压力。
  • 通用抽象换 operator rule 维护。 virtual tensor 能组合多种搬运,但每种新 data-movement operator 仍需开发者写正确的映射和冲突规则。
  • profile robustness 换编译时间。 不必准确建模不同 GPU,代价是每个 shape/batch/hardware configuration 都可能需要数分钟重新搜索。
  • greedy 换可实现性。 二次搜索比指数枚举简单,但一般情形没有最优保证,也没有证明局部 edge saving 可可靠预测组合收益。
  • 接在 fusion 后换低侵入。 VTC 可复用现有 compute kernel,但上游 fusion/layout 已经作出的决定也可能提前关闭更好的 virtual strategy。
  • 适用边界。 大中间 tensor、多段纯搬运、连续 chunk 足够宽时收益大;domain-specific kernel 已手工消除搬运、compute kernel quality 主导或 graph boundary 频繁 materialize 时收益小。

实验设置

  • 硬件是单张 A100 80 GB PCIe 或 H100 NVL;软件为 Ubuntu 22.04、CUDA 12.1。没有 multi-GPU、B200/TMA 或训练实验(§7.1)。
  • 五类 workload 是 Llama 3 8B decoder layer、Gemma 2 9B decoder layer、EfficientViT attention block、YOLOv11n C3K2 block 和 ShuffleNet ShuffleUnit,batch 为 1 或 16(表 1)。论文称“end-to-end inference”,但评测单位实际是这些模型组件,不是完整模型服务。
  • baseline 是 PyTorch 2.6 torch.compile、ONNX Runtime 1.21.1、XLA 和 TensorRT 10;每个 case 与四者中最快的一个比较。Transformer 使用 BF16,vision component 使用 TF32。
  • 主延迟一共有 5 组件×2 batch×2 GPU 的 20 个 case;显存只在 A100 上相对 PyTorch 测 10 个 case。没有真实 arrival、tail latency、能耗或 strategy-search optimum。

实验与结果

  • 组件级端到端延迟:20 个 model-component/batch/GPU case 中,VTC 相对每项最快的 PyTorch、ONNX Runtime、XLA、TensorRT baseline 最多快 1.93 倍,论文报告平均快 1.28 倍;最大值来自 A100、Llama decoder layer、batch 16,相对 TensorRT(图 9、§7.2)。
  • 模型与硬件差异:Transformer 三类组件平均快 1.36 倍,YOLOv11/ShuffleNet 两类 CNN 平均快 1.15 倍;同一组件与 batch 的 10 对 A100/H100 比较中,有 7 对在 H100 上收益更大,但 H100 并非所有 case 都单调更好(图 9、§7.2)。
  • 搬运消除与计算回归:A100 的 10 个 component/batch case 中有 7 个把显式 data-movement operator 全部消除;VTC 的 compute 部分在多数 case 反而比 TensorRT 慢。Llama batch 16 子图中,搬运由 6 kernels、0.9083 ms 降到 0,attention 从 0.6892 降到 0.2812 ms,但 QKV MatMul 从 0.1034 增到 0.1403 ms,总子图约快 4 倍(图 2、10、§7.4)。
  • 峰值显存:仅在 A100、相对 PyTorch built-in monitor 的结果中,VTC 最多节省 60%、平均 17.5%;最大值是 Llama batch 16,从 1,787.3 MB 降到 714.5 MB。ShuffleNet batch 1 只省 0.9%,说明收益取决于是否能虚拟化大中间 tensor(表 2、§7.3)。
  • 强 vLLM 对照vLLM V1、Llama 3 8B、4096-token input、BF16 下,A100 的 QKV-to-attention 子图由 1.972 降到 1.890 ms,仅快 1.043 倍;整个 decoder layer 由 9.152 降到 9.055 ms,只快 1.011 倍。H100 默认 profile 选择不优化,与 vLLM 相同;强制虚拟化会让 decoder layer 从 4.230 增到 4.600 ms,退化 8%(表 3–4、§7.5)。
  • 搜索与边界 case:每个 profile configuration 少于 10 s、所测 graph 总编译少于 10 min(§6)。EfficientViT 因现有 fusion 已覆盖多数机会,只获得约 1.1 倍加速和约 4.4% 显存节省;YOLOv11 则能把大 Split/Concat 中间 tensor 都指向最终 Y(图 11–12、§7.6)。

论断—证据表

论断证据评测边界置信度
virtual tensor 能显著减少 compiler 未消除的数据搬运图 9–10:最多 1.93 倍;A100 10 项中 7 项无显式搬运五种模型组件、batch 1/16、单 A100/H100
虚拟化可同时降低峰值显存表 2:最多 60%、平均 17.5%仅 A100、相对 PyTorch、组件级 graph
profiling 能避开负优化表 3:H100 默认保持 1.000 倍,强制版本为 0.920 倍单个 vLLM/Llama decoder 配置
VTC 对强 domain-specific serving 的端到端收益很小或没有表 3:A100 decoder layer 仅 1.011 倍,H100 为 1.000 倍vLLM V1、Llama 3 8B、4096-token input
greedy strategy 接近全局最优§5.2:一般问题 NP-hard,算法按 marginal profile 贪心无 exhaustive、ILP 或 approximation-gap 实验

批判性分析

论证链条

论文先用纯数据搬运的 index semantics 建立 virtual tensor,再把 kernel 改写限制在 I/O stage,最后用 VTOG 和 profile 解决“哪些值得做”,从语义到实现再到性能的主链是闭合的。图 2 很好地说明了关键取舍:VTC 允许 compute kernel 变慢,只要删除搬运后的总时间更短。作者的“覆盖所有不必要数据搬运”仍应收窄为:在已写 mapping rule、满足一对一表示、允许改写 load/store 且 graph boundary 不要求 materialize 的范围内搜索;它不是任意 tensor program 的完备证明。

假设压力测试

收益取决于搬运占比、连续 chunk、backend kernel quality 和 graph boundary。H100 vLLM 强制版本退化 8% 已证明“compute/memory gap 变大”不会自动让 VTC 更好:cuBLAS MatMul 比 Triton 快得多时,删除一个小 copy 也不划算。动态 sequence、training backward、activation alias、custom CUDA op 和 collective 会扩大 mapping/conflict 空间,并迫使更多 tensor materialize。新 TMA bulk-transfer 还可能让原本昂贵的搬运变便宜,改变 profile 结论。

实验可信度

每项取四个 compiler 中最快者,包含 TensorRT 的手工 attention pattern,baseline 较强;两代 GPU、Transformer/CNN、batch 1/16、显存、breakdown 和 vLLM 负结果也让证据比较诚实。主要外部边界是评测单位:表 1 实际只跑 decoder layer 或 block,不是完整 LLM/CNN、scheduler 和 data pipeline。五类组件、两个 batch 也不足以说明 arbitrary composition;没有重复次数、方差、tail、数值误差、training、multi-GPU 和真实 serving trace。

系统性缺陷

virtual tensor 没有独立 storage,debugger、profiler、checkpoint 和外部 custom kernel 若期待普通 tensor,就必须理解 mapping 或临时 materialize。错误 mapping 可能静默读写错误地址;正文只声明 zero precision loss,没有展示 systematic differential test。VTOG 规则、conflict detection、specialized codegen 和 profile cache 都增加 compiler correctness surface。数分钟编译还会影响 autoscaling,论文也没有讨论 cache key、版本升级和 OOM/failure fallback。

局限与后续工作

  • 局限 1:只测 single-GPU inference component,没有完整模型服务、training/backward、dynamic shape 或 multi-GPU collective。
  • 局限 2:数据搬运规则需按 operator 编写;custom op、alias/in-place 和外部 kernel compatibility 未覆盖。
  • 局限 3:一般 latency function 下 greedy 没有最优保证,profile 搜索又可接近 10 min。
  • 局限 4:实现依赖 Triton,H100 vLLM case 已显示其 compute kernel 可输给 cuBLAS,Blackwell/TMA 尚未验证。
  • 后续工作 1:对小 VTOG 用 exhaustive 或 ILP 求 optimum,按 graph size、conflict density 报 greedy 的 latency gap、profile 数和 compile time。
  • 后续工作 2:扩展到完整 training graph,对 forward/backward alias、optimizer、activation checkpoint 做随机 shape differential test,并报告数值误差与峰值显存。
  • 后续工作 3:在真实 LLM arrival trace 上集成完整 vLLM,测 TTFT、TPOT、P99、goodput、KV-cache capacity 与策略 recompile 次数,而不是只测一个 decoder layer。
  • 后续工作 4:给 CUTLASS/TMA 增加 virtual I/O,在 H100/B200 上比较 memory transaction、compute regression 和 profile 选择是否翻转。
  • 后续工作 5:定义跨 custom kernel 的 materialization ABI,注入错误 mapping、unsupported op 和 OOM,验证安全 fallback、debug trace 与 cache invalidation。

相关