gistfile1.txt

AI serving infrastructure 

Nx + Bumblebee + Nx.Serving is a **three-layer ML stack for Elixir** that goes from raw tensor math all the way up to distributed model inference, without leaving the BEAM ecosystem. 
Lots of features that cost thousands of lines of code, or are never delivered at all, come for free here.

***

## Layer 1: Nx (Numerical Elixir)

`Nx` is the foundation — a multi-dimensional array (tensor) library for Elixir, announced by José Valim in early 2021. [1] Its central design decision is a **backend/compiler split**:

- **Backends** hold tensors in memory and dispatch operations eagerly, like NumPy. The default is a pure-Elixir binary backend, but production code uses `EXLA.Backend` (Google XLA) or `Torchx.Backend` (LibTorch / PyTorch C++ core). [2]
- **Compilers** work with `defn` — a special macro that defines a *numerical definition*, a subset of Elixir that Nx converts into a portable computation graph. When you call a `defn` function, the selected compiler JIT-compiles a hardware-specialized binary for the exact tensor shape and dtype. [3]

```elixir
defn softmax(tensor) do
  Nx.exp(tensor) / Nx.sum(Nx.exp(tensor))
end
```

Calling `softmax` with EXLA configured triggers XLA to compile a CUDA or ROCm kernel and cache it; the next call with the same shape hits the cache directly. [1] The same `defn` body runs on CPU, NVIDIA GPU (`client: :cuda`), or AMD GPU (`client: :rocm`) without modification. [1]

**EXLA** and **Torchx** are the two main backends worth knowing:

| | EXLA | Torchx |
|---|---|---|
| Underlying engine | Google XLA (same as JAX/TensorFlow) | LibTorch (PyTorch C++ core) |
| Strength | Maximum optimization via XLA fusion | Broad operator support, rapid prototyping |
| Compilation style | AOT + JIT via `defn` | Eager execution by default |
| GPU support | CUDA, ROCm | CUDA, ROCm |

***

## Layer 2: Axon

Before Bumblebee, there is **Axon** — a functional neural network library built directly on Nx, comparable conceptually to Flax (JAX) or PyTorch Ignite. [4] Models are expressed as pure functional graphs:

```elixir
model =
  Axon.input("input", shape: {nil, 784})
  |> Axon.dense(128, activation: :relu)
  |> Axon.dense(10, activation: :softmax)
```

Axon handles forward passes, gradient computation via `Nx.Defn`'s automatic differentiation, and training loops. Bumblebee's transformer models are all implemented as Axon graphs, which means they are entirely introspectable, modifiable, and transferable to other Axon pipelines. [5]

***

## Layer 3: Bumblebee

Bumblebee is the **pre-trained transformer model library** for Elixir — the equivalent of `transformers` (HuggingFace) in Python. [4] Every model in Bumblebee is written in 100% pure Elixir as an Axon graph. [5]

The key contribution is **automatic HuggingFace Hub integration**: Bumblebee can download model weights directly from the Hub and convert PyTorch checkpoints (`.safetensors`, `.bin`) into Axon-compatible parameter maps entirely in Elixir, with no Python subprocess. [5]

```elixir
{:ok, model_info} = Bumblebee.load_model({:hf, "meta-llama/Llama-3.2-1B"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "meta-llama/Llama-3.2-1B"})
{:ok, generation_config} = Bumblebee.load_generation_config({:hf, "meta-llama/Llama-3.2-1B"})
```

Supported model families include BERT, BART, GPT-2, Llama, Whisper, CLIP, ResNet, ViT, and Stable Diffusion. [4] A notable gap as of 2026 is that GGUF quantized models are not yet supported — loading a full Llama 3.1 70B requires full-precision VRAM since quantized formats like AWQ/GGUF aren't implemented yet. [6]

***

## Layer 4: Nx.Serving — The Serving Primitive

`Nx.Serving` is the runtime abstraction that sits on top of the model and makes it production-safe. [7] It is best understood as a **process that wraps a model** and handles all the operational complexity:

**Dynamic batching** is the core mechanism. [8] Multiple Elixir processes can call `Nx.Serving.batched_run/2` concurrently; the serving process collects incoming requests for up to a configurable timeout (e.g., 100ms) or until a batch size limit is reached, then stacks them into a single tensor batch, runs one GPU forward pass, and distributes the sliced results back to each waiting caller. [8] A single GPU forward pass over 16 requests is enormously more efficient than 16 sequential calls.

```elixir
# Start the serving as an OTP process in your supervision tree
serving = Bumblebee.Text.generation(model_info, tokenizer, generation_config)

{:ok, _} = Nx.Serving.start_link(serving: serving, name: MyServing, batch_size: 16, batch_timeout: 100)

# Any process, on any node, can call this:
Nx.Serving.batched_run(MyServing, "What is MLIR?")
```

**Distributed by default** is what separates `Nx.Serving` from anything in Python. [9] When you call `Nx.Serving.batched_run(MyServing, input)`, the BEAM automatically routes the request to *any node in the cluster* running a serving registered under that name. You can add a dedicated GPU node to your cluster, supervise a named serving on it, and every other application node immediately benefits — zero configuration changes in the calling code. [9] Scaling out means adding nodes and starting the same serving process on each.

**Pre/post-processing hooks** run on the calling process (client side), while the GPU forward pass runs on the serving process (server side), keeping expensive tokenization work off the GPU node. [8]

***

## The Full Stack as a Supervision Tree

In a real Phoenix application, the entire inference pipeline integrates into standard OTP:

```elixir
defmodule MyApp.Application do
  use Application

  def start(_type, _args) do
    {:ok, model} = Bumblebee.load_model({:hf, "meta-llama/Llama-3.2-1B"})
    {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "meta-llama/Llama-3.2-1B"})
    {:ok, gen_config} = Bumblebee.load_generation_config({:hf, "meta-llama/Llama-3.2-1B"})

    serving = Bumblebee.Text.generation(model, tokenizer, gen_config,
      compile: [batch_size: 4, sequence_length: 512],
      defn_options: [compiler: EXLA]
    )

    children = [
      {Nx.Serving, serving: serving, name: LLMServing, batch_timeout: 100},
      MyAppWeb.Endpoint
    ]

    Supervisor.start_link(children, strategy: :one_for_one)
  end
end
```

The serving is just another child in the supervision tree. [7] If it crashes, the supervisor restarts it — no separate sidecar process, no health-check endpoint, no Kubernetes liveness probe complexity. This is the BEAM's core value proposition applied directly to ML serving: a model crash is handled exactly like any other process crash in the system.