JVM in the Age of AI: Babylon, Valhalla, TornadoVM and friends

Are you tired of all the hype around yet another LLM-as-a-Service? That's why I want to talk about something more interesting than just another tool for prompt engineers.

While the entire industry is discussing it, my goal is to explain what needs to happen within the virtual machine for the JVM to become a good platform for Machine Learning and AI. We will discuss hardware and the challenges its evolution poses for the JVM, projects like Valhalla and Babylon, as well as standardization efforts like the JSR381 Visual Recognition API. We will also look at initiatives like TornadoVM.

This will be an overall birds-eye view to understand how the JVM can meet the demands of contemporary artificial intelligence and machine learning.

1. Artur Skowroński
JVM in the Age of AI
Babylon, Valhalla, TornadoVM and friends

4. Artur Skowronski
Head of Java / Kotlin Development

7. What we won't be talking about today...

8. What we won't be talking about today...
•Using LLM through API (Langchain4j,
Semantic Kernel, Spring AI etc.)
•Data Engineering (data collection, data
cleaning, etc.)
•Model Development (MLOps, observability
etc.)

9. What we will be talking about today!
•The usage and inference from existing
models
•Mainly about GPU programming !

10. What we will be talking about today!
•The usage and inference from existing
models
•Mainly about GPU programming !
•…but if I had told you about it,
then probably no one would have shown up "

11. Helicopter View

12. 152 Slides / 45 Minutes

14. A bit of theory

15. Model Inference
Training Data
Trained Model (Fixed
Parameters)
New Data
Inference
Prediction

16. Model Inference
Inference
Prediction
Layer1
Layer2
Layer3
Next Token of Text
Classification of Image
Inference

17. Model Inference
Trained Model
File
New Data
(Query)
Inference
Prediction

18. Model Inference
Trained Model
File
New Data
(Query)
Inference
Prediction

19. Different Model Formats
.h5 (weights & architecture)
.pkl (Pickle)
.pt
.pb

20. All these frameworks are Python based
.h5 (weights & architecture)
.pkl (Pickle)
.pt
.pb

21. Use Python Inference tooling Libraries with JVMUse Python Inference tooling Libraries with JVM

22. GraalVM

23. GraalVM

26. GraalVM

27. GraalVM
Kubernetes

28. Truffle & GraalPy

29. GraalVM in Truffle Mode
GraalVM
Byte Code Interpreter
Compiled Code #
Graal Compiler
Language Interpreter

30. GraalVM in Truffle Mode
GraalPy
GraalJS
Espresso
Language Interpreter
GraalRuby

31. GraalVM in Truffle Mode

32. Truffle & GraalPy

34. Go a Bit More Complex: GraalPy

35. Go a Bit More Complex: GraalPy

36. Go a Bit More Complex: GraalPy

37. Shortcuts don’t help here
Infer
model in
Java
reusing python
components

38. Yesterday, 15:10

39. Inference of model from Java

40. Model Inference
Inference
Response to User with
Predictions
Query from API
App Logic
Trained Model
File

41. TensorFlow
TensorFlow is an open-source
machine learning framework
for building and deploying ML
models, especially neural
networks.

42. Deeplearning4j (DL4J)
DL4J is a popular ML library in
Java that allows for easy
loading of models trained in
other environments (e.g.,
Keras, TensorFlow).

43. Tribuo
Tribuo is a relatively new machine
learning library for Java, designed
by Oracle Labs.
Its main goal is to simplify the
creation, training, and deployment
of machine learning models in
Java applications.

45. There are
compatible
implementations
of the common
standards

46. API
JPA
(Java Persistence API)
JDBC
(Java Database Connectivity)
JAX-RS
(Java API for RESTful Web
Services)
Servlet API
JSR (Java Specification Request)

48. Interesting JSR
JSR 381 – Visual Recognition (VisRec) API
JSR 376 – Java Platform Module System (Project Jigsaw)
JSR 383 – Java SE 18.3 (Java 10)
JSR 315 – Java Servlet 3.0
JSR 380 – Bean Validation 2.0
JSR 244 – Java EE 5
JSR 303 – Bean Validation 1.0
JSR 907 – JDBC 4.0
JSR 133 – Java Memory Model and Thread Specification
JSR 292 – Dynamically Typed Languages Support (invokedynamic)

49. DeepJavaLibrary

50. Visual Recognition (VisRec) JSR #381

51. Visual Recognition (VisRec) JSR #381

52. Different Model Formats
.h5 (weights & architecture)
.pkl (Pickle)
.pt
.pb

53. What about performance of the process?

54. Model Inference
Inference
Prediction
Layer1
Layer2
Layer3
Next Token of Text
Classification of Image
Computational Complexity
Each layer of requires performing a vast
number of mathematical operations,
such as matrix multiplications. The more
layers a model has, the greater the
number of such operations.

55. Model Inference
Inference
Prediction
Layer1
Layer2
Layer3
Next Token of Text
Classification of ImageComputational Complexity
Each layer of requires performing a vast
number of mathematical operations,
such as matrix multiplications. The more
layers a model has, the greater the
number of such operations.
Model Size
Models with many layers typically have
millions, or even billions, of parameters.
Storing and processing such a large
number of weights requires significant
memory resources.

56. Optimization: Half Floats and Project Valhalla

57. Half Floats
In standard floating-point numbers (commonly
known as floats), we have 32 bits, with some
allocated for the sign of the number, some for the
exponent, and some for the mantissa, which
determines precision.
1 bit for the sign (positive or negative)
8 bits for the exponent (range)
23 bits for the mantissa (precision)

58. Half Floats
In half-floats, we only have 16 bits, meaning they
can store smaller and less precise numbers.
Narrower range with less numbers in given range
1 bit for the sign (positive or negative)
5 bits for the exponent (range)
10 bits for the mantissa (precision)

59. 21 September 2024

60. Not a silver bullet - access to memory is more important

61. Native Calls

64. Start Simple - JNI and Calling C Libraries

65. Start Simple - JNI and Calling C Libraries

66. Quick Intro to JNI
JNI stands for Java Native Interface, a framework that allows
Java code running in the Java Virtual Machine (JVM) to
interact with native applications and libraries written in other
programming languages like C or C++.
JNI is (was?) typically used when Java applications need to
perform tasks that cannot be efficiently accomplished with
pure Java, such as using system-specific libraries or interacting
with low-level hardware.

67. Java 1.1, 1997

68. Quick Intro to JNI
Declare the native method in Java
In your Java code, declare a method as native
without providing an implementation…
…and compile the code, which generates a .class file.

69. Quick Intro to JNI
Using the javah tool (available until JDK 8), generate a C/C++
header file. This file contains the native method signatures in
C/C++…
…which you need to manually write the native implementation
in C/C++

70. Quick Intro to JNI
Compile the C/C++ code into a native library (e.g., .dll on
Windows, .so on Linux, .dylib on macOS).
Java will load this library at runtime using System.loadLibrary().

71. Quick Intro to JNI

72. Project Panama

73. Jextract - Java 17 (2021)
jextract - introduced in Project Panama - automates the
process of generating Java bindings from C/C++ header files.
It eliminates the need to manually write JNI code.
Changes to the native library are handled simply by re-running
jextract.
The Java code generated by jextract maps directly to the
native functions and data structures, allowing seamless
interaction between Java and native code.

74. Project Panama
1. Foreign Function & Memory API (FFM API)
2. Linker API
3. Foreign Data Layout API
4. Interconnect with Foreign Languages
5. Improved Performance for Native Calls
6. Vector API

75. Model Inference run on GPU
Inference
Response to User with
Predictions
Query from API
App Logic
Trained Model
File
GPU

76. Nvidia Stock

77. Can it be efficiently run on CPU without frameworks?
Inference
Response to User with
Predictions
Query from API
App Logic
Trained Model
File
CPU?

79. Model Inference run on GPU
Inference
Response to User with
Predictions
Query from API
App Logic
Trained Model
File
GPU

80. Panama Vector API and Vector Databases
llama.cpp is a 700-line, C++-
based inference
implementation optimized
for CPU inference.

It bypasses the traditional
GPU-centric frameworks like
PyTorch or TensorFlow.

81. Panama Vector API and Vector Databases

82. Model Inference
Inference
Prediction
Layer1
Layer2
Layer3
Next Token of Text
Classification of Image
Inference

83. Can it be efficiently run on CPU without frameworks?
Inference
Response to User with
Predictions
Query from API
App Logic
Trained Model
File
CPU?

84. https://github.com/kherud/java-llama.cpp - Java Bindings

85. Llama3.java

86. Today, 13:50

87. JLama

88. Today, 15:00

89. Panama Vector API and Vector Databases

90. Project Valhalla & Panama Together: SIMD

91. Panama Vector API and Vector Databases
Instruction Result 1Data1
Data2
Data3
SISD
Data1
Instruction Data2
Instruction Data3
Result 1
Result 1
Instruction
Result 1
Result 2
Result 3
Data1
Data2
Data3
SIMD
Data1
Data2
Data3

92. Autovectorization
arrayA[i] + arrayB[I];
result[]
arrayA[]
arrayB[]
SIMD
arrayA[0] + arrayB[0] arrayA[1] + arrayB[1]
arrayA[2] + arrayB[2]
JIT

93. Panama Vector API and Vector Databases

94. Panama Vector API and Vector Databases

95. Panama Vector API and Vector Databases

96. Model Inference
Inference
Prediction
Layer1
Layer2
Layer3
Next Token of Text
Classification of Image
Computational Complexity
Each layer of requires performing a vast
number of mathematical operations,
such as matrix multiplications. The more
layers a model has, the greater the
number of such operations.Model Size
Models with many layers typically have
millions, or even billions, of parameters.
Storing and processing such a large
number of weights requires signifi

97. Special Hardware!

101. OVERLY SIMPLIFIED EXAMPLE ALERT!

102. Model Inference run on GPU
CPU
Instruction 1
Instruction 2
Instruction 3
Instruction 4
Instruction 5
Instruction 6
Instruction 7
Instruction 8
Instruction 9
Instruction 10

103. Model Inference run on GPU
CPU
Instruction 5
Instruction 6
Instruction 7
Instruction 8
Instruction 9
Instruction 10
Instruction 1
Instruction 2
Instruction 3
Instruction 4

104. Model Inference run on GPU
CPU
Instruction 7
Instruction 8
Instruction 9
Instruction 10
Result 1
Result 3
Instruction 4
Result 2 Instruction 5
Instruction 6

105. Model Inference run on GPU
GPU
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7
Kernel 8
Kernel 9
Kernel 10

106. Model Inference run on GPU
GPU
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7
Kernel 8
Kernel 9
Kernel 10
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7

107. Model Inference run on GPU
GPU
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7
Kernel 8
Kernel 9
Kernel 10

109. Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
Kernel
NDGrid (n-dimensional grid)
refers to a multi-dimensional
grid of threads used for parallel
computation on a GPU.
The GPU runs threads in grids
with multiple dimensions,
allowing kernel processing in
different dimensions (e.g., 1D, 2D,
3D) simultaneously.
What is Kernel and NDGrid?

110. What is Kernel and NDGrid?
Kernel
Kernel is a piece of code that is
executed by hundreds or
thousands of threads
simultaneously.
Each thread runs the same code
but operates on different data,
enabling the processing of large
amounts of data in a highly
parallel manner.

111. C and JNI
Kernels are embedded as strings in the code, and programmers had to
manually manage data movement and task execution.
Calculating the square of each element in an array⬤
The first steps in GPU programming
in Java began with projects that
exposed CUDA and OpenCL APIs
through JNI.

Projects such as JCuda, and
JOpenCL allowes the use of native
GPU libraries from within Java, but
require the programmer to manage
GPU tasks in a rather low-level
manner.

112. Kernels expressed in Java
Aparapi and Rootbeer are projects
that allowed the expression of
kernels directly in Java code,
eliminating the need to manually
write CUDA or OpenCL code.

They don’t provide a direct API for
defining ndgrid in the way CUDA or
OpenCL does

113. Kernels expressed in Java
1.Lack of support for proper data
types
2.Lack of support for diverse
hardware
3.In-efficient transpiled code…
4.…or manual hacks

114. Project Sumatra
GPU
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7
Kernel 8
Kernel 9
Kernel 10
Parallel processing
?

115. Project Sumatra

116. Project Sumatra
GPU
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7
Kernel 8
Kernel 9
Kernel 10
Streams.parallel()

117. Project Sumatra
GPU
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7
Kernel 8
Kernel 9
Kernel 10
Streams.parallel()

118. Project Sumatra
GPU
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7
Kernel 8
Kernel 9
Kernel 10
Streams.parallel()

119. Project Sumatra
GPU
Kernel 1
Kernel 2
Kernel 3
Kernel 4
Kernel 5
Kernel 6
Kernel 7
Kernel 8
Kernel 9
Kernel 10
Streams.parallel()
JDK Sumatra’s approach to GPU
acceleration has been largely
superseded by GraalVM’s
capabilities, which provide broader
optimization features (e.g., native
image and polyglot execution) and
better performance across various
platforms
Data transfer between JVM and
GPU can reduce performance
benefits, especially for small tasks.

120. Next Level: TornadoVM

121. Next Level: TornadoVM
API Task Task GraphAnnotations

122. Next Level: TornadoVM
API Task Task GraphAnnotations

123. Next Level: TornadoVM
API Task Annotations Task Graph

124. Next Level: TornadoVM
API Task Annotations
Graph Optimizer
Task Graph

125. Next Level: TornadoVM
API Task Annotations Task Graph
Graph Optimizer

126. Next Level: TornadoVM
API Task Annotations
Graph Optimizer
Execution Engine
Task Graph

127. Next Level: TornadoVM
API Task Annotations
Graph Optimizer
Execution Engine
CPU GPU FPGA
Task Graph

128. Tornado JIT Compiler
Next Level: TornadoVM
API Task Annotations
Graph Optimizer
Execution Engine
CPU GPU FPGA
Task Graph

129. GraalVM

130. Tornado JIT Compiler
Next Level: TornadoVM
API Task Annotations Task Graph
Graph Optimizer
Execution Engine
CPU GPU FPGA
CUDA OpenCL PTX

131. Next Level: TornadoVM
Task Graphs - Direct abstraction
for NDGrid
A lot of different backends (not
only CUDA)
Managed Memory
JIT Support through GraalVM

133. Future: HAT & Project Babylon

134. Heterogenous Accelerator Toolkit (HAT)

135. Heterogenous Accelerator Toolkit (HAT)
NDRange API
FFM data wrapping patterns
Support for Code Reflection from
Project Babylon
Support for FFM API
Application
JDK
Vendor
Native
runtime
CPU GPU
HAT
Code Reflection Panama FFM API

136. Code Reflection
Code reflection is an extension of
Java reflection that allows access to
symbolic representations of Java
code, such as method bodies and
lambda expressions, at runtime.
This makes it possible to
programmatically manipulate Java
code.
Code Reflection

137. Code Models
bytecode
Type Flattening
Abstract Syntax Tree
Full Syntactic Details

138. Code Models
bytecode
Type Flattening
Abstract Syntax Tree
Full Syntactic Details
Code Models
Information about types and Structures

139. JDK
GPU
Code Models

Code Reflection
Code Models
Information about types and Structures
Vendor Native Runtime Code

140. Code Models

141. Code Models
Babylon can use code reflection to
dynamically generate GPU code (e.g.,
OpenCL, CUDA) by analyzing the code
model and converting fragments of
Java code into corresponding GPU
instructions.
Panama FFM API and off-heap
MemorySegments eliminate the
need for manual memory
management between the JVM and
GPU, speeding up data exchange and
increasing performance.
Application
JDK with Babylon
Vendor
Native
runtime
CPU GPU
HAT
Code Reflection Panama FFM API

142. Heterogeneous Accelerator Toolkit (HAT) Update #JVMLS

143. Conclusions: What the future holds

144. Where are we?
Panama & Valhalla SIMD:

145. Where are we?
Panama & Valhalla SIMD:
Valhalla Half-Float:

146. Where are we?
Panama & Valhalla SIMD:
Valhalla Half-Float:
GraalPy:

147. Where are we?
Panama & Valhalla SIMD:
Valhalla Half-Float:
GraalPy:
JCuda:

148. Where are we?
Panama & Valhalla SIMD:
Valhalla Half-Float:
GraalPy:
JCuda:
Sumatra:

149. Where are we?
Panama & Valhalla SIMD:
Valhalla Half-Float:
GraalPy:
JCuda:
Sumatra:
TornadoVM:

150. Where are we?
Panama & Valhalla SIMD:
Valhalla Half-Float:
GraalPy:
JCuda:
Sumatra:
TornadoVM:
HAT & Babylon:

152. Thank you $
@ArturSkowronski