BalancedTernaryEngine.Runtime (balanced_ternary_engine v0.1.0-dev)

Minimal runtime for executing a converted .tbin model.

Supports layer chaining, softmax, layer normalization, multi-head attention with a KV cache, and prefill/decode phases.

Summary

Types

t()

Functions

attention(q, k, v, cache \\ %{}, num_heads \\ 1)

Multi-head attention with a KV cache.

decode(rt, input)

Decode: run a single new token, reusing the cache.

forward(runtime, input)

Run the model (layer chaining) on an input tensor.

layer_norm(x, eps \\ 1.0e-5)

Layer normalization (RMSNorm-style) over the last axis.

load(path, config \\ %{})

Load a .tbin model and an optional config map.

new(layers, config \\ %{})

Create a runtime from already-loaded layers and a config map.

prefill(rt, input)

Prefill: run a full sequence through the model, populating the cache.

reset_cache(rt)

Reset the runtime cache without reloading the model.

softmax(x)

Softmax over the last axis.

validate(rt)

Validate a runtime instance.

validate_config(config)

Validate a runtime configuration map.

Types

t()

@type t() :: %BalancedTernaryEngine.Runtime{
  cache: map(),
  config: map(),
  layers: [any()]
}

Functions

attention(q, k, v, cache \\ %{}, num_heads \\ 1)

@spec attention(Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), map(), integer()) ::
  {Nx.Tensor.t(), map()}

Multi-head attention with a KV cache.

q, k, v are {seq, hidden} tensors. Returns {out, new_cache}.

decode(rt, input)

@spec decode(t(), Nx.Tensor.t()) :: {Nx.Tensor.t(), t()}

Decode: run a single new token, reusing the cache.

forward(runtime, input)

@spec forward(t(), Nx.Tensor.t()) :: Nx.Tensor.t()

Run the model (layer chaining) on an input tensor.

layer_norm(x, eps \\ 1.0e-5)

@spec layer_norm(Nx.Tensor.t(), float()) :: Nx.Tensor.t()

Layer normalization (RMSNorm-style) over the last axis.

load(path, config \\ %{})

@spec load(String.t(), map()) :: {:ok, t()} | {:error, term()}

Load a .tbin model and an optional config map.

new(layers, config \\ %{})

@spec new([any()], map()) :: {:ok, t()} | {:error, term()}

Create a runtime from already-loaded layers and a config map.

prefill(rt, input)

@spec prefill(t(), Nx.Tensor.t()) :: {Nx.Tensor.t(), t()}

Prefill: run a full sequence through the model, populating the cache.

reset_cache(rt)

@spec reset_cache(t()) :: t()

Reset the runtime cache without reloading the model.

softmax(x)

@spec softmax(Nx.Tensor.t()) :: Nx.Tensor.t()

Softmax over the last axis.

validate(rt)

@spec validate(t()) :: :ok | {:error, term()}

Validate a runtime instance.

validate_config(config)

@spec validate_config(term()) :: {:ok, map()} | {:error, term()}

Validate a runtime configuration map.

Returns {:ok, config} when every supplied value is valid, or {:error, reason} otherwise. Unknown keys are allowed so callers can carry application-specific settings.