尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

面向动态张量计算的字节码虚拟机实时编译

面向动态张量计算的字节码虚拟机实时编译 动态张量计算已成为大语言模型等现代AI工作负载的常态——输入序列长度、批量大小乃至控制流路径均在运行时方可确定。传统AI编译器依赖离线shape推导与静态kernel生成难以应对此类动态性而运行时编译虽能利用真实信息但编译链路过长对服务时延影响显著。本文提出一种基于字节码虚拟机的实时编译方案——MindSpore DVMDynamic Virtual Machine。DVM的核心思想是将算子程序在CPU侧编码为字节码在NPU侧由虚拟机解释执行避免为每种shape组合生成机器码的开销。实验表明DVM在LayerNorm上相较TorchInductor加速11.77倍Qwen3-14B微调性能提升1.07~1.21倍编译时间最高提速5个数量级。该方案已作为MindSpore的核心技术组件正式开源。关键词动态张量计算实时编译字节码虚拟机算子融合MindSpore1 引言大语言模型的推理与训练面临着日益严峻的动态性挑战。Transformer架构中的自注意力机制使得输入序列长度在运行时方可确定单个算子的shape可能呈现成千上万种组合。若为每种shape单独编译kernel编译时间与设备内存开销将不可接受若采用bucketing策略将不同shape填充至固定尺寸桶中则会引入冗余计算。更为棘手的是算子融合——大量优化机会仅存在于运行时两个张量表达式在符号层面无法判定是否可融合但实际运行时shape恰好相等此时本可减少一次全局内存访问。传统AI编译器将“编译”理解为生成接近硬件的kernel无论是AOT还是JIT编译最终产物均为机器码或设备kernel。该模式适用于静态模型——其shape、图结构、融合模式在编译前相对确定编译器可投入较多时间进行全局优化。然而动态模型的关键问题在于大量信息在编译前不可知。针对上述矛盾MindSpore团队与香港科技大学广州陈雷教授团队、湖南大学赵捷教授团队合作提出DVM方案。DVM既不试图提前准备所有可能的kernel也不在运行时走完整的机器码生成流程而是在获取真实shape与真实执行路径后仅执行生成tile级字节码的轻量级操作。设备侧不加载新kernel而是由已存在的虚拟机kernel解释执行这些字节码。2 相关工作2.1 静态图编译以TensorFlow XLA和MindSpore静态图为代表的静态图编译器在训练启动前通过整图分析完成算子融合与kernel生成。其优势在于可进行全局优化但缺陷同样明显一旦shape发生变化整图需重新编译编译时间可达秒级。在动态shape场景下这一方案与微秒级的算子下发时间尺度严重错位。2.2 动态shape的现有应对策略学术界与工业界对动态shape问题提出了三类主流应对策略策略一为每种shape单独编译kernel。运行效率较高但遇到新shape时需重新编译导致服务时延抖动且kernel缓存迅速膨胀。策略二dynamic shape kernel或bucketing。将一批shape合并处理减少编译次数但可能引入padding或保守代码牺牲部分性能。策略三预置micro-kernel或多版本kernel。运行时动态选择但依赖shape空间可控一旦融合组合或维度数量复杂化候选数量极易失控。上述策略的共同局限在于它们均在“生成机器码”的框架内寻求优化未能从根本上解决编译链路过长的问题。2.3 字节码虚拟机在AI编译中的探索字节码虚拟机的思想在通用计算领域已有成熟应用如Java JVM、Python VM但在AI编译领域尚属空白。DVM首次将该范式引入动态张量计算通过将编译产物从机器码降级为字节码将编译开销从“秒级”压缩至“微秒级”。3 DVM设计与实现3.1 设计目标DVM的设计围绕以下核心目标展开运行时亲和充分利用运行时获取的真实shape与执行路径信息而非依赖编译前推导。编译轻量化避免传统的图分析、多候选kernel benchmark等重编译流程。融合即时性在动态图执行链路中实现实时算子融合而非依赖离线整图分析。硬件透明性字节码由NPU侧虚拟机解释执行屏蔽底层硬件差异。3.2 系统架构DVM的系统架构可划分为三个层次第一层运行时捕捉与增量构图MindSpore侧。在动态图运行过程中MindSpore在FrontendTask等关键阶段对满足条件的算子进行捕捉并逐步组织成轻量级的运行时融合图。这一过程并非离线分析而是顺着真实执行中的算子流边捕捉、边判断、边决定是否继续融合。当遇到不支持融合的算子、融合边界或同步接口时触发flush将当前融合图交由后续阶段处理。第二层字节码编码CPU侧。获取运行时融合图后DVM在CPU侧将算子程序编码为tile级的字节码。与传统编译器生成机器码不同字节码编码过程无需进行复杂的硬件相关优化仅需完成算子语义到字节码指令的映射编译开销被控制在极低水平。第三层字节码解释执行NPU侧。编码后的字节码被下发至NPU由常驻的虚拟机kernel解释执行。设备侧无需加载新的kernel从而消除了传统JIT编译中kernel加载与缓存管理的开销。3.3 算子融合机制DVM支持两类融合模式Pattern融合针对特定算子组合模式如Elemwise相互融合、Reduce前向融合、MatMulElemwise等进行预定义融合。编译器在捕捉到匹配的算子模式后直接生成对应的融合字节码。Streaming融合在动态图执行过程中根据真实执行路径即时决策融合策略。例如两个张量表达式在符号层面无法判定是否可融合但运行时shape恰好相等DVM可在捕捉阶段实时决定融合。3.4 与传统AI编译器的对比DVM与传统AI编译器在五个维度上存在本质差异维度传统AI编译器DVM编译产物设备kernel或机器码tile级字节码程序编译粒度kernel级别tile级别动态信息使用编译前shape推导运行时验证运行时真实shape与执行路径缓存压力大量kernel或图编译结果需缓存仅常驻虚拟机kernel开销模型先等编译结束再执行编译开销隐藏在硬件流水线中4 实验评估4.1 实验设置DVM已在MindSpore框架中完成实现并开源。实验在Ascend NPU上进行对比基线包括PyTorch的TorchInductor编译器以及MindSpore静态图模式。4.2 单算子性能在LayerNorm算子上的测试中DVM相较TorchInductor实现了11.77倍的加速。这一提升主要源于两个因素其一DVM避免了为不同shape反复编译kernel的开销其二字节码解释执行消除了kernel加载与缓存管理的额外延迟。4.3 端到端模型性能在Qwen3-14B模型的微调任务中DVM实现了1.07~1.21倍的性能提升。提升幅度随输入动态性的增加而扩大——当batch size和序列长度在训练过程中频繁变化时DVM的优势更为显著。4.4 编译时间DVM将编译时间从传统方案的秒级压缩至微秒级最高提速5个数量级。这一改进对于需要频繁处理变长输入的服务场景具有决定性意义——传统方案中每次shape变化触发的秒级重编译将导致服务不可用而DVM的微秒级编译可完全隐藏在计算流水线中。4.5 整网端到端收益在完整的CV和NLP网络测试中DVM的动态图无图融合能力带来1%~5%的整网端到端性能收益。当前支持的融合模式包括Elemwise相互融合、Reduce前向融合及MatMulElemwise等CV融合扩展。5 讨论与展望5.1 适用场景DVM特别适合以下场景变长输入的大模型推理输入序列长度动态变化传统编译方案面临频繁重编译动态控制流的模型分支、循环等控制流路径在运行时确定在线学习与服务场景需要低延迟响应无法承受秒级编译抖动5.2 局限与未来工作当前DVM主要面向Ascend NPU后端对GPU后端的支持仍在进行中。此外当前支持的融合模式尚集中于Elemwise、Reduce和MatMul等基础算子组合未来可扩展至更复杂的融合模式如多头注意力融合、MoE专家融合等。DVM的字节码虚拟机范式为动态模型编译开辟了新的可能性。未来工作可探索将字节码优化如字节码级常量传播、死代码消除纳入虚拟机执行流程进一步降低解释执行的开销。
返回列表