Practical LLM inference in modern Java.pptx
Download as PPTX, PDF0 likes35 views
The document discusses the implementation of a fast LLM inference engine in modern Java, including performance optimization techniques using Java vector APIs and GraalVM. It highlights the advantages of local inference for various models, the importance of memory bandwidth in inference performance, and the development of a lightweight Java library for LLMs with no native dependencies. Additionally, it emphasizes the educational value and accessibility of this approach for developers looking to work with large language models.
![LLM engines
[Token][ize][r]
Prompt Format
Inference
Sampler
GGUF File
Parser](https://image.slidesharecdn.com/practicalllminferenceinmodernjava-241218135810-f50500a5/85/Practical-LLM-inference-in-modern-Java-pptx-17-320.jpg)
![Transformer architecture implemented in
Java
*Low Level Technicals of LLMs: Daniel
Han
class FloatTensor {
float getFloat(int index);
void setFloat(int index, float value);
// ...
void matmul(FloatTensor that, FloatTensor out, int dim0, int dim1) {
for (int i = 0; i < dim0; ++i) {
float result = 0f;
for (int j = 0; j < dim1; j++) {
result += this.getFloat(i * dim1 + j) * that.getFloat(j);
}
out.setFloat(i, result);
}
}
// ...
}
// qkv matmuls for this position
weights.wq[l].matmul(state.xb, state.q, dim, dim);
weights.wk[l].matmul(state.xb, state.k, kvDim, dim);
weights.wv[l].matmul(state.xb, state.v, kvDim, dim);](https://image.slidesharecdn.com/practicalllminferenceinmodernjava-241218135810-f50500a5/85/Practical-LLM-inference-in-modern-Java-pptx-18-320.jpg)
Practical LLM inference in modern Java.pptx
- 1. Practical LLM inference in modern Java Alina Yurenko, Alfonso² Peterssen Oracle Labs
- 2. Who we are • Alina Yurenko, developer advocate @ GraalVM • Alfonso² Peterssen, Java on Java @ GraalVM alina-yur mukel Questions, comments, pull requests?:)
- 3. What we will cover today • Implementing a fast LLM inference engine in modern Java • Executing such Java LLM engine locally • Optimizing for the best CPU-based performance • Performance optimizations: Java Vector API, GraalVM, AOT techniques • Integrating with LangChain4j
- 4. Once upon a tweet...
- 5. Llama3.java Based on llama2.c by Andrej Karpathy and his excellent educational videos, and llama.cpp github.com/mukel/ llama3.java
- 6. (Not so) Large Language Models 16B+ 4B+ < 4B
- 7. Large Language Models are everywhere
- 8. (Not so) Large Language Models • Meta Llama 3+ (1B & 3B & 8B) • Mistral (7B) • Microsoft Phi-3 (3.5B & 7B) • Google Gemma 2 (2.6B & 9B) • Alibaba Qwen2.5 (0.5B & 1.5B & 3B & 7B & 14B) • … or fine-tune your own!
- 9. (Not so) Large Language Models are everywhere Microsoft shipping RWKV.cpp
- 10. (Not so) Large Language Models are everywhere Apple On-Device model (~3B)
- 11. (Not so) Large Language Models are everywhere Gemini Nano, a powerful 3.25B parameter LLM, 100% locally i n your browser!
- 12. (Not so) Large Language Models are everywhere JetBrains Line Compl etion (100M)
- 13. (Not so) Large Language Models are everywhere DevoxxGenie IDEA Plugin
- 15. Local LLM inference Cost Privacy Control
- 16. Local LLM inference in Java No native dependencies Cost Privacy Control Developer productivity Performance
- 17. LLM engines [Token][ize][r] Prompt Format Inference Sampler GGUF File Parser
- 18. Transformer architecture implemented in Java *Low Level Technicals of LLMs: Daniel Han class FloatTensor { float getFloat(int index); void setFloat(int index, float value); // ... void matmul(FloatTensor that, FloatTensor out, int dim0, int dim1) { for (int i = 0; i < dim0; ++i) { float result = 0f; for (int j = 0; j < dim1; j++) { result += this.getFloat(i * dim1 + j) * that.getFloat(j); } out.setFloat(i, result); } } // ... } // qkv matmuls for this position weights.wq[l].matmul(state.xb, state.q, dim, dim); weights.wk[l].matmul(state.xb, state.k, kvDim, dim); weights.wv[l].matmul(state.xb, state.v, kvDim, dim);
- 19. Transformer architecture implemented in Java *Low Level Technicals of LLMs: Daniel Han // RoPE relative positional encoding: complex-valued rotate q and k in each head for (int i = 0; i < dim; i += 2) { int head_dim = i % headSize; float fcr = weights.freq_cis_real.get(position * (headSize / 2) + (head_dim / 2)); float fci = weights.freq_cis_imag.get(position * (headSize / 2) + (head_dim / 2)); int rotn = i < kvDim ? 2 : 1; // how many vectors? 2 = q & k, 1 = q only for (int v = 0; v < rotn; v++) { FloatTensor vec = v == 0 ? state.q : state.k; // the vector to rotate (query or key float v0 = vec.getFloat(i); float v1 = vec.getFloat(i + 1); vec.setFloat(i, v0 * fcr - v1 * fci); vec.setFloat(i + 1, v0 * fci + v1 * fcr); } }
- 20. Transformer architecture implemented in Java *Low Level Technicals of LLMs: Daniel Han NOT a single token, but a vector of probabilities
- 21. Anatomy of LLM weights • Mostly matrices, quantized … • During inference, every weight is read exactly once https://bbycroft.net/llm Is LLM inference the new Bitcoin mining?
- 22. What is the limiting factor of inference performance? Memory bandwidth *and the ability to fully utilize it Project Panama JEP 489: Vector API (Ninth Incubator) Introduce an API to express vector computations that reliably compile at runtime to optimal vector instructions on supported CPU architectures, thus achieving performance superior to equivalent scalar computations.
- 23. Memory bandwidth 50 GB/s 10 TB/s 1 TB/s 400 GB/s
- 24. AI and Memory Wall • Inference is memory bound • 90% of inference is spent on matrix × vector operations "AI and Memory Wall" void matrixVectorMul(FloatTensor m, FloatTensor v, FloatTensor out, int dim0, int dim1) { for (int i = 0; i < dim0; ++i) { float result = 0f; for (int j = 0; j < dim1; j++) { result += m.getFloat(i * dim1 + j) * v.getFloat(j); } out.setFloat(i, result); } }
- 25. Quantizations 16 Colors Grayscale 256 Colors 24bit RGB
- 26. Smaller weights, lower accuracy, faster inference Lossy Quantization Theoretical throughput (tokens/s) of Llama 3.1 8B @ 50 GB/s Approximate (!) comfortable human perception rate
- 27. LLM Inference Engine in Java • High performance • Modern APIs, such as Vector API and FFM API • Rich ecosystem • More control over the models and the inference process (caching, etc) Why Java?
- 28. 🤝 Latest Java APIs Foreign Function and Memory API • Interface for interop between Java code and native code • Just works tm on Graal JIT • Experimental support in Native Image with -H:+ForeignAPISupport (upcalls, downcalls, foreign memory) Vector API • Enables fast vector computations • Initial support in GraalVM for JDK 21, more coming in GraalVM for JDK 24 • Works with Native Image! github.com/ gergo- Thank you Gergö! 🏆
- 29. Faster LLM inference on GraalVM 🚀 • Local LLM inference, powered by modern Java APIs and GraalVM's performance optimizations • ~15% faster inference on Oracle GraalVM (across several models, such as Llama3+) 🔥 • Updates in GraalVM for JDK 24 – get the latest EA build (JDK 24 EA build 15+): • sdk install java 24.ea.15-graal
- 30. Faster LLM inference on GraalVM • AOT at the speed of JIT 🚀 • inference code is AOT-friendly • Combining the power of Java with native performance: you can pre-parse GGUF meta-data and cache prompts at build time for instant inference • <25ms time to first token Native Image
- 31. Demo – Native llama3 & AOT optimizations
- 32. Demo – LangChain4j Integration
- 34. • Performance of llama3.java running on GraalVM Native Image is comparable to llama.cpp, and gets even closer as the model size increases. More optimizations coming soon:) • Different approaches: idiomatic Java vs hand-tuned tensor operations • Free performance boost: quantization helps, and larger models are resilient to it llama3 native image Oracle GraalVM 24 EA 15 Linux 3950X 64GB@3800
- 35. Other models? • Meta Llama 3+ (1B & 3B & 8B) … but also 30+ other models, such as: • Mistral (7B) • Microsoft Phi-3 (3.5B & 7B) • Google Gemma 2 (2.6B & 9B) • Alibaba Qwen2.5 (0.5B & 1.5B & 3B & 7B & 14B) • … and specialized models for math, programming, and more Llama3.java can be easily adapted to different models and vendors huggingface.co/mukel
- 36. LLM Inference Engine • Performance testing and tuning (ARM, Apple Silicon, AVX512...) • Implement additional features (prompt caching, quantizations, YARN...) • Support for GPUs (HAT, TornadoVM?) • Further integrations with the Java libraries and frameworks (such as LangChain4j) • Audio • Vision Help wanted ♂️ ️ ♂️ ♂️ ♂️ ♂️ ♂️ ♂️ ♂️ ♂️ ♂️ ♂️ ♂️ ♂️ ♂️
- 37. Wrap up • FAST LLM inference in pure, modern Java • No dependencies (73kB jar file) • GraalVM makes it even faster 🚀 • Native Image support with AOT model pre-loading, for instant time-to-first-token • Simple, accessible and with a high educational value for learning about LLMs • Use a lightweight local inference engine or your inspiration for your next Java project • Fun to hack and fun to use! github.com/mukel/llama3.java
- 38. Please rate this session :)