Skip to content

Latest commit

 

History

History
269 lines (202 loc) · 7.2 KB

File metadata and controls

269 lines (202 loc) · 7.2 KB

Llama3.java Technical Findings

Deep dive into pure Java LLM inference and the dramatic Vector API performance improvements between JDK 21 and 25.

The 40x Performance Mystery

Benchmark results (Llama 3.2 1B, 32 tokens, Apple M1 Pro):

JDK Version Distribution Performance Ratio
JDK 25 Temurin ~13.0 tok/s Baseline
JDK 25 GraalVM CE ~13.3 tok/s +2%
JDK 21 GraalVM CE ~0.33 tok/s 40x slower!

Key finding: The Vector API implementation in JDK 25 is dramatically better than JDK 21.

Why Such a Big Difference?

Vector API Evolution

JDK 21 (Vector API - 6th Incubator):

  • Vector operations compile to scalar loops
  • Poor code generation for some patterns
  • Limited optimization
  • ARM64 NEON not fully utilized

JDK 25 (Vector API - 8th Incubator, near-final):

  • Improved JIT compilation
  • Better ARM64 NEON code generation
  • Optimized memory access patterns
  • Aligned loads/stores
  • Reduced bounds checking overhead

The difference is in the JIT compiler, not the API surface.

What Changed

Example: Vector dot product

JDK 21 generates:

; Scalar loop with poor vectorization
loop:
  ldr   w0, [x1], #4
  ldr   w2, [x3], #4
  fmul  s0, s0, s2
  fadd  s1, s1, s0
  subs  x4, x4, #1
  b.ne  loop

JDK 25 generates:

; Proper NEON vector instructions
loop:
  ld1   {v0.4s}, [x1], #16
  ld1   {v1.4s}, [x3], #16
  fmul  v0.4s, v0.4s, v1.4s
  fadd  v2.4s, v2.4s, v0.4s
  subs  x4, x4, #4
  b.ne  loop

Result: 4 elements per iteration instead of 1 = 4x speedup, plus other optimizations = 40x total.

Pure Java LLM Implementation

Architecture

Llama3.java implements:

  1. GGUF format parser - Loads llama.cpp models
  2. Tokenizer - BPE with special tokens
  3. Transformer - Attention, FFN, RMSNorm
  4. Sampling - Top-p, temperature
  5. Chat templates - Llama 3 Instruct format

All in ~3000 lines of pure Java code!

Key Techniques

1. Vector API for matrix multiplication:

FloatVector va = FloatVector.fromArray(SPECIES, a, i);
FloatVector vb = FloatVector.fromArray(SPECIES, b, j);
FloatVector result = va.mul(vb);
sum = result.reduceLanes(VectorOperators.ADD);

2. Memory-mapped files:

FileChannel fc = FileChannel.open(path, StandardOpenOption.READ);
MemorySegment segment = fc.map(READ_ONLY, 0, fc.size(), arena);

3. Unsafe for performance-critical paths:

UNSAFE.getShort(memorySegment.address() + offset)

Performance Analysis

CPU Utilization

JDK 25:

  • CPU usage: ~100% (all cores utilized)
  • Vector instructions: Fully utilized
  • Memory bandwidth: Saturated
  • Bottleneck: Memory bandwidth

JDK 21:

  • CPU usage: ~25-50% (poor utilization)
  • Vector instructions: Poorly generated
  • Memory bandwidth: Underutilized
  • Bottleneck: Scalar execution

Memory Access Patterns

Llama 3.2 1B model:

  • Parameters: 1.24 billion
  • Memory (FP16): ~2.5 GB
  • Memory bandwidth critical for performance

Why Vector API helps:

  • Aligned memory loads (128-bit at once)
  • Reduced loop overhead
  • Better instruction-level parallelism
  • Prefetching hints

Comparison with Native Solutions

Approach Speed Technology Complexity
java-llama.cpp ~50 tok/s JNI + Metal GPU Low (just use library)
Cyfra ~33 tok/s Scala + Vulkan GPU High (GPU programming)
Llama3.java (JDK 25) ~13 tok/s Pure Java + Vector API Low (single file)
CPython ~10 tok/s Python + llama.cpp CPU Low (Python script)
TornadoVM ~6 tok/s Java + OpenCL GPU Medium (TaskGraph)
Llama3.java (JDK 21) ~0.3 tok/s Pure Java + Vector API Low (single file)

Key insights:

  • Pure Java can be competitive with native CPU implementations
  • GPU acceleration makes huge difference (50 vs 13 tok/s)
  • JDK version matters more than you'd think (40x!)

Why Pure Java Matters

Advantages

1. Portability:

  • Single Java file runs everywhere
  • No platform-specific binaries
  • No GPU drivers needed
  • Works in restricted environments (no native access)

2. Simplicity:

  • Just Java code - no build complexity
  • No CMake, gcc, CUDA toolchains
  • Easy to understand and modify
  • Can be embedded in any Java app

3. Security:

  • No native code = smaller attack surface
  • Memory-safe (mostly - uses some Unsafe)
  • No library version conflicts
  • Easier to audit

Disadvantages

1. Performance ceiling:

  • CPU-only (no GPU acceleration)
  • Bounded by memory bandwidth
  • Can't match optimized native code

2. Memory overhead:

  • JVM heap + model
  • More GC pressure
  • Less memory-efficient than C++

3. JDK dependence:

  • Performance varies wildly by JDK version
  • Requires incubator modules
  • May break on future JDK updates

Vector API Status

Current state (JDK 25):

  • Still in incubator (8th preview)
  • API is stable
  • Performance is good
  • Expected to be final in JDK 26-27

Evolution:

  • JDK 16: First preview
  • JDK 17-21: Incremental improvements
  • JDK 22-25: Major performance work
  • JDK 26+: Expected to finalize

Lesson: Wait for optimizations before judging API performance!

When to Use Llama3.java

✅ Good Use Cases

  • Learning - Understand how LLMs work
  • Prototyping - Quick experiments without setup
  • Portability - Need to run everywhere
  • No GPU - CPU-only environments
  • Security - Restricted environments (no native code)
  • Embedded - Small deployments

❌ Not Recommended For

  • Production inference - Use GPU-accelerated solutions
  • High throughput - java-llama.cpp is 4x faster
  • Large models - Memory constraints
  • JDK 21 or older - Performance is terrible

The Single-File Philosophy

Llama3.java is intentionally one file:

  • Educational - Easy to read and understand
  • Self-contained - No dependency hell
  • Hackable - Modify and experiment easily
  • Portable - Just copy one file

Trade-offs:

  • Not production-grade architecture
  • Harder to maintain for large changes
  • But: Perfect for demos and learning

Lessons Learned

1. JDK Version Matters Enormously

  • Don't judge API by early implementations
  • Performance can improve 40x without API changes
  • Always test on latest JDK

2. Vector API is Production-Viable (on JDK 25+)

  • ~13 tok/s is respectable for CPU-only
  • Competitive with native CPU implementations
  • Good enough for many use cases

3. Pure Java Can Compete

  • With Vector API, Java can be surprisingly fast
  • Won't beat GPU, but beats naive implementations
  • Portability is a major win

4. Single-File Demos are Powerful

  • Great for education and experiments
  • Easy to share and modify
  • Reduces friction to near-zero

References

See Also

  • demos/java-llama-cpp/ - JNI with GPU (✅ ~50 tok/s)
  • demos/tornadovm/ - GPU via OpenCL
  • demos/graalpy/ - Python comparison
  • demos/valhalla/FINDINGS.md - Float16 Vector API research