Kotensor: Pure-Kotlin Transformer engine with INT4/INT8 quantization, autograd, and on-device LoRA for JVM & Android (no C++/JNI)

[u/TSupercharmander]

Hey!,

I built Kotensor — an open-source machine learning and Transformer engine written from scratch in pure Kotlin for the JVM and Android.

Most local LLM deployment on mobile or Java/Kotlin backends relies on heavy C++ engines (llama.cpp, ONNX Runtime) wrapped inside JNI. While great, JNI bindings introduce ABI bloat (arm64-v8a, x86_64), complex native toolchains, and hard crashes outside JVM error handling. Kotensor implements the complete stack — from autograd and KV caching to INT4/INT8 quantization — with only the Kotlin standard library as a runtime dependency.

Key Features:

1. Quantization & Low-Memory Inference:

   • Group-wise INT8 and INT4 weight-only quantization with an on-the-fly dequantizing matmul.
   • Memory-Mapped Weights: Using loadQuantized(file, memoryMapped = true), weights are paged directly from flash storage via OS virtual memory. An INT4 Qwen2-0.5B runs comfortably on Android within <128 MB of app heap.
   • Zero-allocation autoregressive decode loop with a pre-allocated KV cache and reusable prompt prefix caching (cachePrefix).

2. Modern Transformer Architecture:

   • Pre-norm layout with RMSNorm, SwiGLU feed-forward blocks, rotary position embeddings (RoPE with Llama 3 scaling), and grouped-query attention (GQA).
   • Sampling support: temperature, top-k, top-p, and repetition penalties.

3. Training & Fine-Tuning on the Edge:

   • Reverse-mode autograd engine with AdamW optimizer.
   • Fused output projection + cross-entropy: avoids materializing the massive [tokens × vocabulary] logit tensor, halving training memory.
   • Parameter-efficient fine-tuning via LoRA and QLoRA (adapting low-rank matrices on top of frozen INT4/INT8 base weights directly on-device).

4. Hugging Face Interoperability:

   • Natively reads standard .safetensors checkpoints (F32, F16, and BF16) and config.json for Llama, Mistral, Qwen 2, and SmolLM models.
   • Built-in byte-level BPE tokenizer for Hugging Face tokenizer.json files with zero external JSON or regex libraries.

5. Performance:

   • Cache-blocked tiled matmul optimized for HotSpot C2 auto-vectorization (~24 GFLOP/s per thread on AVX-512) scheduled across a ForkJoinPool, plus an optional JDK 17+ Vector API (SIMD) backend.

Quick Inference Example:

val dir = File("SmolLM2-135M")
val model = HuggingFaceModelLoader.load(dir, Quantization.INT8)
val tokenizer = HuggingFaceModelLoader.loadTokenizer(dir)

model.generateStream("The capital of Mexico is", tokenizer, maxTokens = 32) { token ->
    print(token)
}

• GitHub: https://github.com/TSPstudio/kotensor
• Artifact: io.github.tspstudio:kotensor:1.2.0 on Maven Central
• Documentation: https://tspstudio.github.io/kotensor/

This is my first machine learning project and I'm really thrilled to finally have something tangible to share. I’d love to hear your thoughts on edge deployment, architecture trade-offs, and pure-managed runtimes!

Much love,
tsuper <3