AI 时代下的 Java 算力突围:TornadoVM & Project Babylon 初探

人工智能与大模型技术近年来迅猛发展，计算范式也正在经历从"CPU 中心化"向"异构并行化"的剧变。AI 大模型的训练与推理通常涉及海量矩阵运算与数据并行计算，GPU、FPGA 等异构硬件在算力效率上已远超传统 CPU。然而，Java 作为企业级应用的主导语言，其生态体系长期面向 CPU 优化，缺乏对异构硬件的原生支持能力，这成为 Java 在 AI 时代面临的核心挑战之一。幸运的是，一些业界的先驱者也敏锐地察觉到了这一点，本文就以 TornadoVM 与 Project Babylon 两个关键项目为切入点，简要前瞻一下未来的 Java 生态如何通过异构加速技术来突破算力瓶颈。

特别感谢 Juan Fumero，他是 TornadoVM 项目的核心成员，本文接下来的很多内容和图示就参考自他的博客来进行总结，有兴趣的读者可以阅读文末标注的参考资料来获取更细节的一手材料。

1 TornadoVM

TornadoVM 提供了一种 Java 应用的通用异构加速方式，释放了 AI 和 LLM 所需要的 GPU、FPGA 等异构并行算力，使得 Java 生态集成 AI 负载更快更简单。

开源项目地址：https://github.com/beehive-lab/TornadoVM

1.1  什么是 TornadoVM

TornadoVM 由曼彻斯特大学 Beehive Lab 研发，虽然叫 VM，但并不是一个完整虚拟机。它是一个开源插件，可以集成于常见的业界各种 JDK  (e.g., GraalVM, Red Hat Mandrel, Amazon Corretto, OpenJDK, Azul Zulu) 上。

诞生的动机?

AI 和 LLM 应用发展迅速，但它们在工业领域需要的算力是非常巨大的。AI 大模型的本质是一个"基于数据模式匹配的概率生成系统"，即通过海量数据的预训练形成一个从输入 token 到输出 token 的映射关系模型，内部的计算流程通俗地讲就是"参数"，这需要海量迭代 + 大规模并行计算来生成一个较好的模型。根据有关文章资料显示，为了达到类似人类的水平，大模型至少需要 11 万亿参数、228 万亿 token 数、1.55×10²⁸ 次的浮点运算，这带来的硬件算力成本是昂贵的。计算过程涉及大规模矩阵乘法和卷积等数据并行计算任务，不论是训练还是推理，AI 大模型应用对计算算力的需求方面，GPU 等异构硬件在大规模推理方面的表现比传统 CPU 要更好些。然而目前 Java 生态体系主要面向 CPU，对于异构硬件的算力调度能力没有直接手段，需要通过 JNI 封装底层 C/C++ 代码，不能够像处理 CPU 那样处理异构设备——检测频繁执行的代码，并生成优化的机器码。TornadoVM 就是为了解决这一问题而诞生的。

1.2 TornadoVM 架构和实现原理

从架构层面而言，TornadoVM 表现为一个 JVM 插件，通过 SDK 安装、maven pom API 集成的方式添加到项目中，配合各 JDK 发行版来支持 Java 程序不仅能在 CPU 上运行，也能在其他异构硬件上运行，并能够针对每一种架构优化代码而无需修改源代码，发挥异构硬件在不同并行化类型下（任务并行化->CPU、数据并行化->GPU、管道并行化->FPGA）的性能优势。它有三个代码生成后端：OpenCL、CUDA/PTX 和 SPIR-V。开发人员可以选择使用其中一个。

使用 TornadoVM 编程有它自己的一套软件栈，是一个层级和微内核的软件架构。

首先在 API 层面，TornadoVM 暴露了一些 API 组件帮助程序员进行开发，根据这些 API 组件背后的思想，TornadoVM 的上层设计包含以下三方面：

（1）表示并行性

Java 并非主要为并行性和现代异构硬件设计的编程语言，TornadoVM 也不具备自动检测并行化的能力，而且并非所有类型的 Java 计算都适合卸载到 TornadoVM 运行在异构硬件上，具有 for 循环且没有依赖关系的负载比较适合，特别是基于矩阵的机器学习等比较合适，因此 TornadoVM 仍然需要一个关于如何并行化代码的提示。TornadoVM 提供了两种方式表达并行性，一是 Loop Parallel API，通过 @Parallel 和 @Reduce 注解标注并行化循环，这将提示 TornadoVM 的 JIT 编译器自动将其转化为 OpenCL 等后端并行编程语言，二是 Kernel API 进行显式的 GPU 风格编程，带有线程 ID、局部内存、可用的屏障等概念，适合专业的异构编程开发者使用，相对复杂。本下文仅作简单举例，如下示例便对矩阵乘法的两个 for 循环进行了并行化标注，表示这是可以并行化的计算。

class Compute {    public static void matrixMultiplication(final float[] A, final float[] B, final float[] C, final int size) { for (@Parallel int i = 0; i < size; i++) { for (@Parallel int j = 0; j < size; j++) { float sum = 0.0f; for (int k = 0; k < size; k++)                    sum += A[(i * size) + k] * B[(k * size) + j];               C[(i * size) + j] = sum;           }       }    } }

（2）识别要加速的 Java 方法

TornadoVM 在方法层面卸载代码，类似于 Hotspot 中的 JIT 编译器，为了定义和识别需要加速的 Java 代码，TornadoVM 提供了 TaskGraph API（在 TornadoVM 0.15 版本以前是 TaskSchedule API，职能差不多），这是一个基于任务的编程 API，每个任务对应一个现有的 Java 方法，用于创建和定义计算数据和任务，一般使用其中的 transferToDevice、task、transferToHost 方法定义输入、输出和执行的计算任务，定义输入和输出对 TornadoVM 运行时在主机与异构加速器之间进行数据迁移非常有用。继续上面的例子：

TaskGraph t = new TaskGraph("s0")       .transferToDevice(DataTransferMode.EVERY_EXECUTION, matrixA, matrixB)       .task("t0", Compute::matrixMultiplication, matrixA, matrixB, result, size)        .transferToHost(DataTransferMode.EVERY_EXECUTION, result);

（3） 构建执行计划

通过 Task Graph，开发者可以在加速器上加速该图指示的所有方法，最后要将任务图实例化为 TornadoExecutionPlan 对象并开启执行任务，构建执行计划时可以进行一些高层级优化，例如选择特定 GPU、启用分析器等来优化执行计划。

TornadoExecutionPlan executionPlan = new TornadoExecutionPlan(t.snapshot());executionPlan.withProfiler(ProfilerMode.SILENT) // Enable Profiling    .withWarmUp() //  Perform a warmup (compile the code and install it in a code-cache).    .withDevice(device) // Select a specific device    .execute();

在内部，TornadoVM 运行时系统负责在异构硬件上优化、编译和调度应用。具体技术上它扩展了 GraalVM JIT 编译器，将 Java 代码卸载为低等级的 GPU 友好代码，即上述提到的三种后端语言。此外，Tornado 引擎还负责目标硬件的数据迁移、数据处理和任务计划的优化与执行等，是一个面向 Java 异构硬件编程、运行和优化的全包解决方案，有关运行时的细节要深入代码理解，笔者就不作继续展开了，有机会再详细了解。

1.3  TornadoVM 应用

（1）异构性能优势

TornadoVM 利用异构硬件性能，在一些并行计算任务上展现出比原有 Java 标准并行流高得多的性能。

在模糊滤镜例子中，在 NVIDIA GPU 2060 上运行应用程序，可以获得高达 340x 的性能加速，相比原有 Java 标准并行流提升 30 多倍。类似的例子在金融科技、数学物理模拟、信号处理等方面都有涉及，最高可达 4000x 的加速比。

凭借这些优势，业界也有一些企业实践表明了 TornadoVM 的应用潜力。其中一个应用场景来自卢森堡的 Neurocom 公司，用它运行一种自然语言处理算法。通过在 GPU 上运行分层聚类算法实现了 30 倍的性能提升。另一个应用场景来自 Spark Works 公司，使用强大的 GPU 处理来自物联网设备的信息。与原生 Java 相比，可以获得 460 倍的性能提升。更多案例可参考官网：https://www.tornadovm.org/use-cases

（2）与 AI 的结合

回到本文的话题，AI 是需要算力的，而这些算力由 GPU 等异构硬件提供更为合适，TornadoVM 就成为了助力 Java 在这一领域发展的很好土壤。甚至 TornadoVM 自己提供了一个 Java 应用向我们展示了它应用于 AI 应用的可能性。

项目地址：https://github.com/beehive-lab/GPULlama3.java?tab=readme-ov-file

GPULlama3.java 是一个基于纯 Java 实现的 Llama3 模型，并在此基础上引入 TornadoVM 进行 GPU 加速推理，并已经正式作为 Java LLM 集成框架 LangChain4j 的一个模型提供者，业已集成进 Quarkus 全栈框架中方便开发者使用。有了 TornadoVM 的加持，以 Java 为基础编写 AI 应用我们就可以利用到异构硬件的算力，填补了传统 JVM 仅能利用 CPU 的空白，实现算力调度和异构编译优化，在训练和推理过程中减少 AI 大模型的时间和金钱成本。

2. Project Babylon

开源项目地址：https://github.com/openjdk/babylon

Babylon 项目是 Oracle 官方推进的异构编程框架，利用代码反射技术将原生 Java 集成扩展到非常规编程模型。有关这一项目的细节资料目前公之于众的还比较有限，下面很多内容依然参考了 Juan Fumero 的博客，他还对 Babylon 和 TornadoVM 也做了一些对比。

2.1 什么是 Babylon 项目

Project Babylon 是 Oracle 官方孵化中的一个新的 OpenJDK 项目，2023 年由 Oracle Java 架构师 Paul Sandoz 发起，旨在将 Java 编程模型扩展到非常规编程模型，支持 SQL、可微分编程、机器学习模型和 GPU 等外部编程模型。Babylon 通过代码反射（Code Reflection） 技术进行实现，增强 Java 的反射编程，使 Java 能够以标准形式访问、分析和转换代码，一旦实现，该增强功能将允许任何外部编程模型实现为 Java 库。

Sandoz 举了一个例子，其中开发人员希望用 Java 编写 GPU 内核，然后在 GPU 上运行它。开发人员的代码需要经过分析并转化为可行的 GPU 内核。虽然 Java 库可以管理此操作，但它需要以符号形式访问 Java 代码。本系统在程序生命周期的各个阶段（例如编译时或运行时）限制对非标准API或约定的这种访问。此外，可用的符号形式（字节码或抽象语法树）通常不支持正确的分析和转换。

2.2 Babylon 技术方案

(1) 代码模型（Code model）和 代码反射（Code Reflection）

概念介绍参考：https://openjdk.org/projects/babylon/articles/code-models

现有的两种对 Java 代码建模的方法是：AST 和字节码。源代码通过 javac 建模成 AST（抽象语法树），并根据 AST 生成包含字节码指令的 class 文件，字节码是一种平台无关的中间表示，交由 JVM 解释/编译成可执行的机器码。Babylon 认为这两种表示都不适合分析和转换——AST 过于接近源代码，包含太多细节；而字节码则过于接近可执行程序，做了类型擦除，结构平铺等。

Babylon 定义了一个术语——代码模型，这是第三种建模方式。它识别方法和 lambda 表达式，由 javac 编译器生成并存储于类文件中，并且可以通过反射 API 访问和操作它。笔者理解为是一种集两家之长、避两家之短的中庸建模方式，专门适配异构计算所需的分析与转化信息。

Babylon 通过@Reflect/ @CodeReflection （前者是原有的 Java 反射 API、后者是 Babylon 新增的，虽没有明确资料表明，但理解上后者就是对前者的用户友好型包装）注解标注需要被编译器编译为代码模型的方法，然后经过 Babylon 增强的反射 API，就能够在运行时访问和操作这些代码模型，例如下面的方法代码（摘自https://github.com/openjdk/babylon/blob/code-reflection/test/jdk/java/lang/reflect/code/TestExpressionGraphs.java）

@Reflectstatic double sub(double a, double b) {return a - b;}

Method m = TestExpressionGraphs.class.getDeclaredMethod("sub", double.class, double.class);        // Get the code model for method sub        Optional<CoreOp.FuncOp> oModel = Op.ofMethod(m);        CoreOp.FuncOp model = oModel.orElseThrow();

输出 model.toText()，可以获得下面的输出：

func @loc="46:5:file:/.../TestExpressionGraphs.java" @"sub"(%0 : java.type:"double", %1 : java.type:"double")java.type:"double" -> {    %2 : Var<java.type:"double"> = var %0 @loc="46:5" @"a";    %3 : Var<java.type:"double"> = var %1 @loc="46:5" @"b";    %4 : java.type:"double" = var.load %2 @loc="48:16";    %5 : java.type:"double" = var.load %3 @loc="48:20";    %6 : java.type:"double" = sub %4 %5 @loc="48:16";return %6 @loc="48:9";};

它易于人类阅读，并且完成了一定程度的标准化抽象设计，无源码依赖，又保持了无损语义，不会像字节码一样丢失一些并行语义，包含类型与控制流信息，便于分析和转换。是 Babylon 完成符号形式访问 Java 代码的技术核心和基石。

有了 Code Model，就可以通过某种"转换器"分析和转化 Java 代码为 GPU 等异构代码，从而直接在 Java 中实现非常规编程。

（2）HAT（Heterogeneous Accelerator Toolkit，异构加速器工具包）

HAT 子项目地址：https://github.com/openjdk/babylon/tree/code-reflection/hat

HAT 是 Babylon 内部的子项目，它通俗地说就是负责对接 GPU 等异构硬件的 Java 代码转换器，也是 Babylon 的实现载体。基于 Code Reflection 技术，HAT 可以动态生成和操作 Java 程序，为各种硬件加速器生成定制的 GPU 代码。HAT 接口分为三类，我们可以从中发现与 TornadoVM 非常类似的一些概念：

• NDRange Kernel API

  NDRange API 是 HAT 用于让 Java 开发者创建并行 Kernel 的接口，类似于 TornadoVM 的第二类表示并行性的 API。例如下面的 method

  @CodeReflectionpublic void vectorAddition(F32Array a, F32Array b, F32Array c, KernelContext context) {      int idx = context.x;float sum = a.array(idx) + b.array(idx);      c.array(idx, sum);}

• 通过 @CodeReflection 注解指示生成代码模型；
• 参数类型变化以用于表示与 GPU 兼容的数据结构；
• 引入 KernelContext 用于访问包括线程 ID 等 GPU 执行参数，并且不再需要 for 循环，取而代之使用线程 ID 来访问数据，这是一种标准的 GPU 编程范式——从而写 Java 代码就像是在写异构代码。

• 内存映射

  算是 HAT 的一个亮点，定义了名为 iFaceMapper 的系列接口逻辑用于表示数据，数据通过 Panama 内存 API 存储在堆外，从而摆脱 JVM 的 GC 管理，便于 GPU 计算。通过扩展该接口，开发人员也可以创建与 GPU 或其他硬件加速器兼容的自定义数据对象，这是一种非常有前景的集成方法，具有很强的灵活性和可控制性。

• Accelerator 和 Compute Context

  类似于 TornadoVM 的任务图和执行计划，这两个对象定义了要使用的加速后端和分载的 Kernel 任务列表。

  var accelerator = new Accelerator(lookup, Backend.FIRST);accelerator.compute(cc ->       MyClass.methodToOffload(cc, matrixA, matrixB, matrixC, size));@CodeReflectionpublic static void methodToOffload(ComputeContext cc, MyCustomArray matrixA) {   cc.dispatchKernel(size, kc -> myGPUKernel(kc, data));}

（3）内部工作机制

上图抽象地展示了 Babylon/HAT 从 Java 代码生成 GPU 代码过程的工作：

1. Java 代码输入。使用 NDRange API 和 HAT 的 Accelerator 和 Compute Context API 等，并进行 @CodeReflection 注解，这些 API 帮助定义了用户期望的 Java 式非常规编程；
2. 代码模型生成。然后 javac 编译器会为其生成代码模型并存储在 class 中；
3. 深度优化。图中绘制的"降阶"，我理解这应该属于一个优化步骤，HAT 基于反射 API 对代码模型进行深度优化，就像 TornadoVM 的 Graal JIT Compiler 那样，做一些循环优化、内存访问分析、无用代码剔除、降阶等，将原始代码模型分析优化并转换为一种类似于 LLVM IR 的底层中间表示；
4. 异构代码生成。接着 HAT 会根据底层中间表示生成对应的 OpenCL、PTX 等异构代码。

2.3 Babylon 的性能与应用

Babylon 整体还处于单独项目开发孵化中，因此还未得到太多场景验证，预计可能会在未来 1~2 年里提交 JEP 提案，并逐渐公之于众成为主干特性之一。

Juan 做了一些性能对比试验，可以看到，即使在包含 JIT 编译与执行的首轮运行中，TornadoVM 的运行速度依然显著快于 Babylon。与 Babylon 相比，TornadoVM 1D 版本快 2.3 倍，2D 版本快 9.3 倍。目前的数据表明 TornadoVM 还是强于 Babylon 的，但作者仍然谦虚地表示这并不是绝对的结论，因为 TornadoVM 有一个 JIT 编译器和优化器，可以在某些应用程序中带来性能优势。

3. TornadoVM vs Babylon

目前，两个项目都在积极开发中，它们的核心目标是大体相同的，都是为了提供一个便于对 Java 代码进行操作和转换的接口，帮助开发人员在 Java 中定义和运行 GPU 等异构编程并行框架，我们也可以从它们的框架中发现相似的影子。

TornadoVM 相对成熟，已有场景应用，而 Babylon 还在孵化中，并且一定程度上也借鉴了 TornadoVM 的一些东西，但是它是社区官方项目，可能会在未来成为一种主流范式，并能够与其他一些先进项目例如 Panama 等融合。

TornadoVM 的优势：

• 它是一个更加完整和成熟的解决方案，有工程化实践，进度上比 Babylon 领先
• 提供多设备选择和动态任务迁移能力。不仅限于 GPU，还涵盖各类现代硬件加速器和运行时调度
• 编译时优化。目前 Babylon 尚未提供优化编译器，而 TornadoVM 扩展了 Graal JIT 编译器，提供了针对不同架构和供应商的编译器优化管道。
• 多维并行和 API 的复杂性。Babylon 尚未原生支持 2D 和 3D 并行计算，并且采用 Range 编程模型，这就要求开发者必须掌握 GPU 等编程模型的专业知识。相对地，TornadoVM 多维问题支持更完善，并且为新手开发者提供基于注解的并行 API，和与 Range 类似的低级 Kernel API，双重策略更照顾开发者；

Babylon 的优势：