Skip to content

Commit 6f00513

Browse files
authored
feat(backend): add tinygrad multimodal backend (experimental) (#9364)
* feat(backend): add tinygrad multimodal backend Wire tinygrad as a new Python backend covering LLM text generation with native tool-call extraction, embeddings, Stable Diffusion 1.x image generation, and Whisper speech-to-text from a single self-contained container. Backend (`backend/python/tinygrad/`): - `backend.py` gRPC servicer with LLM Predict/PredictStream (auto-detects Llama / Qwen2 / Mistral architecture from `config.json`, supports safetensors and GGUF), Embedding via mean-pooled last hidden state, GenerateImage via the vendored SD1.x pipeline, AudioTranscription + AudioTranscriptionStream via the vendored Whisper inference loop, plus Tokenize / ModelMetadata / Status / Free. - Vendored upstream model code under `vendor/` (MIT, headers preserved): llama.py with an added `qkv_bias` flag for Qwen2-family bias support and an `embed()` method that returns the last hidden state, plus clip.py, unet.py, stable_diffusion.py (trimmed to drop the MLPerf training branch that pulls `mlperf.initializers`), audio_helpers.py and whisper.py (trimmed to drop the pyaudio listener). - Pluggable tool-call parsers under `tool_parsers/`: hermes (Qwen2.5 / Hermes), llama3_json (Llama 3.1+), qwen3_xml (Qwen 3), mistral (Mistral / Mixtral). Auto-selected from model architecture or `Options`. - `install.sh` pins Python 3.11.14 (tinygrad >=0.12 needs >=3.11; the default portable python is 3.10). - `package.sh` bundles libLLVM.so.1 + libedit/libtinfo/libgomp/libsndfile into the scratch image. `run.sh` sets `CPU_LLVM=1` and `LLVM_PATH` so tinygrad's CPU device uses the in-process libLLVM JIT instead of shelling out to the missing `clang` binary. - Local unit tests for Health and the four parsers in `test.py`. Build wiring: - Root `Makefile`: `.NOTPARALLEL`, `prepare-test-extra`, `test-extra`, `BACKEND_TINYGRAD = tinygrad|python|.|false|true`, docker-build-target eval, and `docker-build-backends` aggregator. - `.github/workflows/backend.yml`: cpu / cuda12 / cuda13 build matrix entries (mirrors the transformers backend placement). - `backend/index.yaml`: `&tinygrad` meta + cpu/cuda12/cuda13 image entries (latest + development). E2E test wiring: - `tests/e2e-backends/backend_test.go` gains an `image` capability that exercises GenerateImage and asserts a non-empty PNG is written to `dst`. New `BACKEND_TEST_IMAGE_PROMPT` / `BACKEND_TEST_IMAGE_STEPS` knobs. - Five new make targets next to `test-extra-backend-vllm`: - `test-extra-backend-tinygrad` — Qwen2.5-0.5B-Instruct + hermes, mirrors the vllm target 1:1 (5/9 specs in ~57s). - `test-extra-backend-tinygrad-embeddings` — same model, embeddings via LLM hidden state (3/9 in ~10s). - `test-extra-backend-tinygrad-sd` — stable-diffusion-v1-5 mirror, health/load/image (3/9 in ~10min, 4 diffusion steps on CPU). - `test-extra-backend-tinygrad-whisper` — openai/whisper-tiny.en against jfk.wav from whisper.cpp samples (4/9 in ~49s). - `test-extra-backend-tinygrad-all` aggregate. All four targets land green on the first MVP pass: 15 specs total, 0 failures across LLM+tools, embeddings, image generation, and speech transcription. * refactor(tinygrad): collapse to a single backend image tinygrad generates its own GPU kernels (PTX renderer for CUDA, the autogen ctypes wrappers for HIP / Metal / WebGPU) and never links against cuDNN, cuBLAS, or any toolkit-version-tied library. The only runtime dependency that varies across hosts is the driver's libcuda.so.1 / libamdhip64.so, which are injected into the container at run time by the nvidia-container / rocm runtimes. So unlike torch- or vLLM-based backends, there is no reason to ship per-CUDA-version images. - Drop the cuda12-tinygrad and cuda13-tinygrad build-matrix entries from .github/workflows/backend.yml. The sole remaining entry is renamed to -tinygrad (from -cpu-tinygrad) since it is no longer CPU-only. - Collapse backend/index.yaml to a single meta + development pair. The meta anchor carries the latest uri directly; the development entry points at the master tag. - run.sh picks the tinygrad device at launch time by probing /usr/lib/... for libcuda.so.1 / libamdhip64.so. When libcuda is visible we set CUDA=1 + CUDA_PTX=1 so tinygrad uses its own PTX renderer (avoids any nvrtc/toolkit dependency); otherwise we fall back to HIP or CLANG. CPU_LLVM=1 + LLVM_PATH keep the in-process libLLVM JIT for the CLANG path. - backend.py's _select_tinygrad_device() is trimmed to a CLANG-only fallback since production device selection happens in run.sh. Re-ran test-extra-backend-tinygrad after the change: Ran 5 of 9 Specs in 56.541 seconds — 5 Passed, 0 Failed
1 parent 8487058 commit 6f00513

29 files changed

Lines changed: 3266 additions & 10 deletions

.github/workflows/backend.yml

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,25 @@ jobs:
105105
dockerfile: "./backend/Dockerfile.python"
106106
context: "./"
107107
ubuntu-version: '2404'
108+
# tinygrad ships a single image — its CPU device uses bundled
109+
# libLLVM, and its CUDA / HIP / Metal devices dlopen the host
110+
# driver libraries at runtime via tinygrad's ctypes autogen
111+
# wrappers. There is no toolkit-version split because tinygrad
112+
# generates kernels itself (PTX renderer for CUDA) and never
113+
# links against cuDNN/cuBLAS/torch.
114+
- build-type: ''
115+
cuda-major-version: ""
116+
cuda-minor-version: ""
117+
platforms: 'linux/amd64'
118+
tag-latest: 'auto'
119+
tag-suffix: '-tinygrad'
120+
runs-on: 'ubuntu-latest'
121+
base-image: "ubuntu:24.04"
122+
skip-drivers: 'true'
123+
backend: "tinygrad"
124+
dockerfile: "./backend/Dockerfile.python"
125+
context: "./"
126+
ubuntu-version: '2404'
108127
- build-type: ''
109128
cuda-major-version: ""
110129
cuda-minor-version: ""

Makefile

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# Disable parallel execution for backend builds
2-
.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/faster-whisper backends/silero-vad backends/local-store backends/huggingface backends/rfdetr backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp
2+
.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/faster-whisper backends/silero-vad backends/local-store backends/huggingface backends/rfdetr backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp backends/tinygrad
33

44
GOCMD=go
55
GOTEST=$(GOCMD) test
@@ -432,6 +432,7 @@ prepare-test-extra: protogen-python
432432
$(MAKE) -C backend/python/whisperx
433433
$(MAKE) -C backend/python/ace-step
434434
$(MAKE) -C backend/python/trl
435+
$(MAKE) -C backend/python/tinygrad
435436
$(MAKE) -C backend/rust/kokoros kokoros-grpc
436437

437438
test-extra: prepare-test-extra
@@ -454,6 +455,7 @@ test-extra: prepare-test-extra
454455
$(MAKE) -C backend/python/whisperx test
455456
$(MAKE) -C backend/python/ace-step test
456457
$(MAKE) -C backend/python/trl test
458+
$(MAKE) -C backend/python/tinygrad test
457459
$(MAKE) -C backend/rust/kokoros test
458460

459461
##
@@ -550,6 +552,56 @@ test-extra-backend-vllm: docker-build-vllm
550552
BACKEND_TEST_OPTIONS=tool_parser:hermes \
551553
$(MAKE) test-extra-backend
552554

555+
## tinygrad mirrors the vllm target (same model, same caps, same parser) so
556+
## the two backends are directly comparable. The LLM path covers Predict,
557+
## streaming and native tool-call extraction. Companion targets below cover
558+
## embeddings, Stable Diffusion and Whisper — run them individually or via
559+
## the `test-extra-backend-tinygrad-all` aggregate.
560+
test-extra-backend-tinygrad: docker-build-tinygrad
561+
BACKEND_IMAGE=local-ai-backend:tinygrad \
562+
BACKEND_TEST_MODEL_NAME=Qwen/Qwen2.5-0.5B-Instruct \
563+
BACKEND_TEST_CAPS=health,load,predict,stream,tools \
564+
BACKEND_TEST_OPTIONS=tool_parser:hermes \
565+
$(MAKE) test-extra-backend
566+
567+
## tinygrad — embeddings via LLM last-hidden-state pooling. Reuses the same
568+
## Qwen2.5-0.5B-Instruct as the chat target so we don't need a separate BERT
569+
## vendor; the Embedding RPC mean-pools and L2-normalizes the last-layer
570+
## hidden state.
571+
test-extra-backend-tinygrad-embeddings: docker-build-tinygrad
572+
BACKEND_IMAGE=local-ai-backend:tinygrad \
573+
BACKEND_TEST_MODEL_NAME=Qwen/Qwen2.5-0.5B-Instruct \
574+
BACKEND_TEST_CAPS=health,load,embeddings \
575+
$(MAKE) test-extra-backend
576+
577+
## tinygrad — Stable Diffusion 1.5. The original CompVis/runwayml repos have
578+
## been gated, so we use the community-maintained mirror at
579+
## stable-diffusion-v1-5/stable-diffusion-v1-5 with the EMA-only pruned
580+
## checkpoint (~4.3GB). Step count is kept low (4) so a CPU-only run finishes
581+
## in a few minutes; bump BACKEND_TEST_IMAGE_STEPS for higher quality.
582+
test-extra-backend-tinygrad-sd: docker-build-tinygrad
583+
BACKEND_IMAGE=local-ai-backend:tinygrad \
584+
BACKEND_TEST_MODEL_NAME=stable-diffusion-v1-5/stable-diffusion-v1-5 \
585+
BACKEND_TEST_CAPS=health,load,image \
586+
$(MAKE) test-extra-backend
587+
588+
## tinygrad — Whisper. Loads OpenAI's tiny.en checkpoint (smallest at ~75MB)
589+
## from the original azure CDN through tinygrad's `fetch` helper, and
590+
## transcribes the canonical jfk.wav fixture from whisper.cpp's CI samples.
591+
## Exercises both AudioTranscription and AudioTranscriptionStream.
592+
test-extra-backend-tinygrad-whisper: docker-build-tinygrad
593+
BACKEND_IMAGE=local-ai-backend:tinygrad \
594+
BACKEND_TEST_MODEL_NAME=openai/whisper-tiny.en \
595+
BACKEND_TEST_AUDIO_URL=https://github.com/ggml-org/whisper.cpp/raw/master/samples/jfk.wav \
596+
BACKEND_TEST_CAPS=health,load,transcription \
597+
$(MAKE) test-extra-backend
598+
599+
test-extra-backend-tinygrad-all: \
600+
test-extra-backend-tinygrad \
601+
test-extra-backend-tinygrad-embeddings \
602+
test-extra-backend-tinygrad-sd \
603+
test-extra-backend-tinygrad-whisper
604+
553605
## mlx is Apple-Silicon-first — the MLX backend auto-detects the right tool
554606
## parser from the chat template, so no tool_parser: option is needed (it
555607
## would be ignored at runtime). Run this on macOS / arm64 with Metal; the
@@ -707,6 +759,7 @@ BACKEND_MLX_VLM = mlx-vlm|python|.|false|true
707759
BACKEND_MLX_DISTRIBUTED = mlx-distributed|python|./|false|true
708760
BACKEND_TRL = trl|python|.|false|true
709761
BACKEND_LLAMA_CPP_QUANTIZATION = llama-cpp-quantization|python|.|false|true
762+
BACKEND_TINYGRAD = tinygrad|python|.|false|true
710763

711764
# Rust backends
712765
BACKEND_KOKOROS = kokoros|rust|.|false|true
@@ -778,14 +831,15 @@ $(eval $(call generate-docker-build-target,$(BACKEND_MLX_VLM)))
778831
$(eval $(call generate-docker-build-target,$(BACKEND_MLX_DISTRIBUTED)))
779832
$(eval $(call generate-docker-build-target,$(BACKEND_TRL)))
780833
$(eval $(call generate-docker-build-target,$(BACKEND_LLAMA_CPP_QUANTIZATION)))
834+
$(eval $(call generate-docker-build-target,$(BACKEND_TINYGRAD)))
781835
$(eval $(call generate-docker-build-target,$(BACKEND_KOKOROS)))
782836
$(eval $(call generate-docker-build-target,$(BACKEND_SAM3_CPP)))
783837

784838
# Pattern rule for docker-save targets
785839
docker-save-%: backend-images
786840
docker save local-ai-backend:$* -o backend-images/$*.tar
787841

788-
docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-kokoros docker-build-sam3-cpp docker-build-qwen3-tts-cpp
842+
docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-tinygrad docker-build-kokoros docker-build-sam3-cpp docker-build-qwen3-tts-cpp
789843

790844
########################################################
791845
### Mock Backend for E2E Tests

backend/index.yaml

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -361,6 +361,34 @@
361361
intel: "intel-rerankers"
362362
amd: "rocm-rerankers"
363363
metal: "metal-rerankers"
364+
- &tinygrad
365+
name: "tinygrad"
366+
alias: "tinygrad"
367+
license: MIT
368+
description: |
369+
tinygrad is a minimalist deep-learning framework with zero runtime
370+
dependencies that targets CUDA, ROCm, Metal, WebGPU and CPU (CLANG).
371+
The LocalAI tinygrad backend exposes a single multimodal runtime that
372+
covers LLM text generation (Llama / Qwen / Mistral via safetensors or
373+
GGUF) with native tool-call extraction, BERT-family embeddings,
374+
Stable Diffusion 1.x / 2 / XL image generation, and Whisper speech-to-text.
375+
376+
Single image: tinygrad generates its own GPU kernels and dlopens the
377+
host driver libraries at runtime, so there is no per-toolkit build
378+
split. The same image runs CPU-only or accelerates against
379+
CUDA / ROCm / Metal when the host driver is visible.
380+
urls:
381+
- https://github.com/tinygrad/tinygrad
382+
uri: "quay.io/go-skynet/local-ai-backends:latest-tinygrad"
383+
mirrors:
384+
- localai/localai-backends:latest-tinygrad
385+
tags:
386+
- text-to-text
387+
- LLM
388+
- embeddings
389+
- image-generation
390+
- transcription
391+
- multimodal
364392
- &transformers
365393
name: "transformers"
366394
icon: https://avatars.githubusercontent.com/u/25720743?s=200&v=4
@@ -1993,6 +2021,15 @@
19932021
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-rerankers"
19942022
mirrors:
19952023
- localai/localai-backends:master-metal-darwin-arm64-rerankers
2024+
## tinygrad
2025+
## Single image — the meta anchor above carries the latest uri directly
2026+
## since there is only one variant. The development entry below points at
2027+
## the master tag.
2028+
- !!merge <<: *tinygrad
2029+
name: "tinygrad-development"
2030+
uri: "quay.io/go-skynet/local-ai-backends:master-tinygrad"
2031+
mirrors:
2032+
- localai/localai-backends:master-tinygrad
19962033
## Transformers
19972034
- !!merge <<: *transformers
19982035
name: "transformers-development"

backend/python/tinygrad/Makefile

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
.DEFAULT_GOAL := install
2+
3+
.PHONY: install
4+
install:
5+
bash install.sh
6+
7+
.PHONY: run
8+
run: install
9+
@echo "Running tinygrad..."
10+
bash run.sh
11+
@echo "tinygrad run."
12+
13+
.PHONY: test
14+
test: install
15+
@echo "Testing tinygrad..."
16+
bash test.sh
17+
@echo "tinygrad tested."
18+
19+
.PHONY: protogen-clean
20+
protogen-clean:
21+
$(RM) backend_pb2_grpc.py backend_pb2.py
22+
23+
.PHONY: clean
24+
clean: protogen-clean
25+
rm -rf venv __pycache__

0 commit comments

Comments
 (0)