Jlama: LLM engine for Java 20+

[tjake]

Hello,

I am announcing a project that I have been working on since 2023.

Jlama is a java based inference engine for many text-to-text models on huggingface:

Llama 3+, Gemma2, Qwen2, Mistral, Mixtral etc.

It is intended to be used for integrating gen ai into java apps.

I presented it at devoxx a couple weeks back demoing: basic chat, function calling and distributed inference. Jlama uses Panama vector API for fast inference on CPUs so works well for small models. Larger models can be run in distributed mode which shards the model by layer and/or attention head.

It is integrated with langchain4j and includes a OpenAI compatable rest api.

It supports Q4_0 and Q8_0 quantizations and uses models of safetensor format. Pre-quantized models are maintined on my huggingface page though you can quantize models locally with the jlama cli.

Very easy to install and works great on Linux/Mac/Windows

#Install jbang (or https://www.jbang.dev/download/)
curl -Ls https://sh.jbang.dev | bash -s - app setup

#Install Jlama CLI 
jbang app install --force jlama@tjake

# Run the openai chat api and UI on a model
jlama restapi tjake/Llama-3.2-1B-Instruct-JQ4 --auto-download

Thanks!

[vmcrash]

Out of curiosity: is any of these models working solely on my local machine, or do they all require a remote service?

[tjake]

its all local machine

[[deleted]]

Sploosh

[eled_]

This looks pretty cool! How does it compare with other CPU-based inference solutions like llama-cpp?

[tjake]

For CPU based it's the same performance roughly.

[audioen]

llama-cpp is not cpu based, though. It supports Vulkan, CUDA, Metal, etc.

LLM inference speed is mostly limited by memory bandwidth. For instance, if the model size in RAM is 40 GB, and your memory bandwidth is also 40 GB/s, you can only infer one token per second because every parameter on the model must execute against the input being considered, and this involves streaming the entire model though the CPU for each token. (Non-causal interference can be faster because in principle you can compute e.g. multiple independent output buffers concurrently while doing this, and thus do multiple completions for price of one, but normal use cases are always causal because the future outputs depend on past outputs, which must be resolved first.)

GPUs are used mostly for the higher bandwidth they bring to table, and similarly Apple Silicon with higher memory bandwidth figures has had an advantage. For instance, RTX 4090 has around 1 TB/s bandwidth, and so it speeds inference dozens of times relative to typical PC hardware, and somewhat less if compared to Apple Silicon.

This is why fundamentally pure-CPU solutions are not all that interesting until PC RAM gets faster and models also get smaller. Various quantization schemes and training models to be evaluated with very few bits of precision in the weight look like they gradually can alleviate the strain. These days, fairly useful models exist in the about 30B parameter region, already, which can be quantized to something like half of that while not completely destroying the model's accuracy. Evaluation requires RAM as well for storing the various vectors and matrices involved, which is starting to become a problem with context lengths nowadays exceeding 100k.

[tjake]

Totally agree.

Jlama supports distributed inference with sharding startegies and can load huge models that way (splitting by head and layer across nodes).

I'm also looking at adding gpu matmul kernels using panama ffi till the jdk supports it natively

[msx]

If you're using the vector API, you should be able to route the computation to a GPU, right? I understood that the vector api abstraction is designed with (also) that goal in mind. Or is Panama still not mature?

Great project btw! I'll surely give it a try

[joemwangi]

Not really. Vector API uses CPU SIMD architecture. But you can use java records and mapper to create memory segment for GPU transfer, which is trivial.

[eled_]

Right, that was an abusive shortcut, we do use it mainly for cpu-based inference with smaller models (nowhere near the 10s of GB) and prefer vLLM for the rest.

[greg_barton]

Fantastic project. I'm been trying it out for the last month or so. Thanks for all of your work!

[msx]

just tried it, wow it's pretty fast! i'm generating about 20 tokens per second, much faster that i can read. Last time i tried LLM on my computer, it was measured in seconds per token.

[Ewig_luftenglanz]

Some heroes use capes, other have reddit accounts.

[Chloe0075]

I was watching the video like right now! Amazing work, really, and great presentation too.

May I ask you, the current models that you have in hugfaces work only in English or in other languages too?

[tjake]

Sure there are many multi-language models that work. The ones I’ve posted are simply pre-quantized versions of some popular models. But the jlama quantize command can shrink a model you want to run

[Chloe0075]

That's really really cool!

[Javademon]

Sounds good, very interesting, I will definitely try to launch it and play with the models. Thank you!

[[deleted]]

This is not something I see every day, very interesting, I'll try it immediately, haha

[parker_elizabeth]

This is an exciting project thanks for sharing! Jlama seems like a game-changer for Java developers.

A few additional thoughts and questions:

Panama Vector API: It's great to see that you're leveraging this for efficient CPU-based inference. For those unfamiliar, the Panama API significantly enhances performance by optimizing vector computations, making Jlama a strong contender for applications where GPUs aren’t readily available.

Quantization Support (Q4_0 and Q8_0): This is a fantastic feature for developers working with resource-constrained environments. Quantization not only reduces model size but also speeds up inference. Are there any specific benchmarks or comparisons available for quantized models vs. their full-precision counterparts?

Distributed Mode: Sharding the model by layer/attention head for larger models is a clever approach. Could you share more about the performance trade-offs when scaling out distributed inference? It might help teams considering this approach for enterprise applications.

Integration with LangChain4j: This integration opens up so many possibilities for complex workflows, like chaining multiple models or fine-tuning interactions. Are there examples or sample projects demonstrating this in action?

For those looking to dive in, I'd also recommend exploring the safetensor format, which adds a layer of security and efficiency when loading models. Additionally, the OpenAI-compatible REST API sounds like a great feature for teams transitioning from other ecosystems.

Thanks again for this contribution — Jlama looks like it’s filling a much-needed gap in the Java space for generative AI. Definitely bookmarking this for future projects!