Projects/SOUP/Documentation

[ SOUP / MLBRICKS KIT 1.0.0B1 ]

SOUP Quick Start

Build uniform and per-layer SOUP models and use recurrent generation.

DISTRIBUTIONmlbricks-kitDOCS VERSION1.0.0b1

QUICK START

Build a basic SOUP model.

PYTHON
import torch
from mlbricks import SOUP

model = SOUP(
    dim=512,
    width=1116,
    depth=2,
    mixer="esa",
    ffn="saffn",
    backend="auto",
)

x = torch.randn(2, 128, 512)
y = model(x)

PER-LAYER CHOICES

PYTHON
model = SOUP(
    dim=512,
    width=[1116, 1116],
    depth=2,
    mixer=["esa", "bolt"],
    ffn=["saffn", "saffn"],
    backend="auto",
)

RECURRENT GENERATION

PYTHON
model = model.eval()
model.prepare_generation()

prompt_out, cache = model.prefill(prompt_hidden)
next_out, cache = model.decode_step(next_hidden, cache)
← InstallationArchitecture →