TinyLoop

TinyLoop is a C++/CUDA inference framework for weight-shared looped transformers.

It is built around a narrow assumption that mainstream runtimes do not make: if a model already reuses the same transformer block many times, the runtime should preserve that structure instead of flattening it into a conventional deep stack and paying the full parameter-memory cost.

📖 Full documentation: ethanyangtw.github.io/tinyloop — PyTorch-style API reference, runtime modes, KV cache modes deep-dive, and measured benchmarks across 407M / 1B-effective on H100. Available in English, 繁體中文, 한국어, and 日本語.

Status

TinyLoop is no longer just a benchmark binary. The repository currently includes:

• an installable CMake library target
• a CLI for inspect, benchmark, generate, and speculate
• an optional Python module for evaluation and scripting
• a custom .tinyloop model format
• CUDA kernels for quantized GEMM, attention, and fused ops
• Marlin INT4 tensor-core path on Ampere / Ada (SM 8.0–8.9) — opt-in via TINYLOOP_USE_MARLIN=1 plus tl.pack_model_for_marlin(model). Routes every INT4 loop / pre linear through upstream Marlin at ~86 % of the 4090's FP16 tensor-core peak, with a companion correction kernel that folds TinyLoop's GPTQ zero points into Marlin's symmetric-only layout. End-to-end validated: Model.score() matches native INT4 within 0.67 % PPL drift on a 2 B-class model. Full details: Marlin INT4 on Ada
• Rotary Position Embeddings (RoPE) — fused in-place Q/K rotation across all forward paths, with configurable theta. Enables Llama-family model support.
• Grouped-Query Attention (GQA) — fewer KV heads than Q heads, shrinking both the shared loop block's attn_qkv weight and per-layer KV cache. In a looped transformer, the L2 residency and VRAM savings multiply across L iterations.
• default-on KV-cache decode with parity coverage
• self-hosted CUDA CI wiring through CTest and GitHub Actions

What it is not yet:

• a generic transformer framework
• a continuous-batching serving stack
• an OpenAI-compatible hosted inference server
• a finished production platform for arbitrary checkpoint families

Why TinyLoop Exists

TinyLoop targets runtimes shaped like:

tokens
  -> embed
  -> pre blocks
  -> shared loop block x L
  -> output norm
  -> head

That changes deployment economics in three ways:

1. Weight memory is tied to unique blocks, not effective depth.
2. Runtime loop count becomes a compute knob.
3. A self-speculative decode path can reuse the same weights for draft and verify passes.

TinyLoop is therefore best suited to:

• looped or recurrent transformer research
• low-bit deployment work
• teams building a specialized runtime for a controlled model family

Current Runtime Snapshot

Validated H100 numbers for the current target artifact:

Path	Shape	Result	Notes
Default low-bit benchmark	seq_len=128, loops=8, no logits	30.48 ms	Current default INT2 runtime path
FP16-body benchmark	seq_len=128, loops=8, no logits	2.82 ms	TINYLOOP_EXPERIMENTAL_FP16_BODY=1 plus the safer tiled prefill path
Cached decode attention	K=128, D=2048, heads=16	0.092 ms	Direct single-query decode kernel
Full prefill attention	T=128, D=2048, heads=16	0.143 ms -> 0.087 ms	New tiled prefill path vs reference

Quality statement for those speedups:

• CUDA unit tests pass on the validated environment.
• Cached-vs-uncached parity is covered by decode, CLI, tokenizer-aware, and small eval-slice regressions.
• The recent prefill attention path matched the reference path on the checked raw-byte and tokenizer-backed prompts.

That supports no regression detected in validated paths. It does not prove universal equivalence across every model artifact or prompt distribution.

Documentation

Published site: ethanyangtw.github.io/tinyloop (also available as wiki/docs/ source in this repository).

Map

• Getting Started — installation · quickstart · model conversion
• Framework — current status · CLI reference · config reference · C++ API · Python API · model format · runtime modes · KV cache modes
• Architecture — runtime architecture · production roadmap
• Operations — performance and memory · testing and CI · troubleshooting · GitHub Pages

Python API reference (PyTorch-style)

Split into focused subpages:

• Model — lifecycle, config, VRAM, thread safety
• Scoring — score, score_last, score_logit_lens, score_trajectory, uncertainty & consistency
• Generation — generate, generate_stream, generate_speculative
• Prefix cache — 2.4–3.0× throughput on shared-prefix workloads
• Warm-start mid-loop — −32.5 % wall-clock, break-even at N=1 follow-up
• KV cache modes — all five storage modes (FP16 KV / INT8 KV / FP16-h / INT8-h / INT4-h)

Shortest route

1. Read installation.
2. Run the quickstart.
3. Jump to CLI reference or Python API depending on your integration.

Repository Layout

tinyloop/
├── include/        public headers
├── src/            model loading, inference, generation, tokenizer, CLI
├── cuda/           CUDA kernels
├── tests/          CUDA tests and Python regressions
├── tools/          conversion and evaluation helpers
├── wiki/           Docusaurus documentation site
└── .github/        GitHub Actions workflows

Important entry points:

• include/tinyloop.h: public API
• src/main.cpp: CLI behavior
• src/inference.cpp: core execution orchestration
• src/generate.cpp: generation and speculative decoding
• cuda/attention.cu: prefill and cached decode attention kernels (with GQA head mapping)
• cuda/rope.cu: rotary position embeddings

Install

From PyPI

pip install loopformer

This installs the Python bindings and CLI. Requires a working CUDA toolkit (nvcc on PATH) and an NVIDIA GPU — the source distribution compiles CUDA kernels locally via scikit-build-core. H100 (sm_90) is the default target; override with CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=80" for A100.

From source

For a full source build (framework library, CLI, tests), see below.

Build

Requirements

• CMake 3.18+
• CUDA toolkit with nvcc
• C++17 compiler
• Python 3 for tools and optional bindings
• pybind11 if you want tinyloop_py

Configure

cmake -S tinyloop -B tinyloop/build -DCMAKE_CUDA_ARCHITECTURES=90

Examples:

• 80 for A100
• 90 for H100

Build

cmake --build tinyloop/build -j

Install

cmake --install tinyloop/build --prefix /tmp/tinyloop-install

Installed artifacts include:

• TinyLoop::framework
• tinyloop CLI
• public headers under include/tinyloop
• exported CMake package files

Quick CLI Tour

Inspect

./tinyloop/build/tinyloop model.tinyloop inspect

Use this first to validate:

• format version
• dimensions and head count
• embedding mode
• estimated weight, buffer, and KV-cache memory

Benchmark

./tinyloop/build/tinyloop model.tinyloop benchmark --loops 8 --seq-len 128 --repeat 10

For phase-level timings:

TINYLOOP_CUDA_PROFILE=1 \
./tinyloop/build/tinyloop model.tinyloop benchmark --loops 8 --seq-len 128 --repeat 1

Generate

./tinyloop/build/tinyloop model.tinyloop generate \
  --prompt "Looped transformers are" \
  --loops 8 \
  --max-tokens 64 \
  --temperature 0.8 \
  --top-k 50

Generation uses cached decode by default. For the uncached reference path:

TINYLOOP_DISABLE_KV_CACHE=1 \
./tinyloop/build/tinyloop model.tinyloop generate --prompt "Looped transformers are"

Speculate

./tinyloop/build/tinyloop model.tinyloop speculate \
  --prompt "Looped transformers are" \
  --draft-loops 2 \
  --verify-loops 8 \
  --draft-ahead 4 \
  --max-tokens 64

Current speculative decoding status:

• core accept-or-resample logic exists
• cache reuse exists
• regression coverage exists
• the path is still best treated as advanced or experimental runtime behavior rather than a final serving contract

Python Quick Start

import numpy as np
import tinyloop_py

model = tinyloop_py.Model("model.tinyloop", max_seq_len=2048)
tokens = np.asarray([15496, 995], dtype=np.int32)

last_logits = model.score_last(tokens, loops=8)
generated = model.generate(tokens, max_tokens=32, loops=8, temperature=0.8, top_k=50)

Prefix caching is also available:

prefix = model.build_prefix_cache(tokens, loops=8, cache_window=0)
continuation = model.generate_from_prefix_cache(prefix, max_tokens=32, loops=8)

C++ Quick Start

#include <tinyloop/tinyloop.h>
#include <vector>

int main() {
    tinyloop::Model* model = tinyloop::load_model("model.tinyloop", 2048);

    std::vector<int32_t> tokens = {15496, 995};
    std::vector<float> logits(model->config.vocab_size);
    tinyloop::score_last_token(model, tokens.data(), (int)tokens.size(), logits.data(), 8);

    tinyloop::GenerateConfig cfg;
    cfg.n_loops = 8;
    cfg.max_tokens = 32;
    auto out = tinyloop::generate(model, tokens.data(), (int)tokens.size(), cfg);

    tinyloop::free_model(model);
}

Runtime Flags

Common environment variables:

Variable	Effect
TINYLOOP_CUDA_PROFILE=1	Print phase timings during benchmark
TINYLOOP_DISABLE_KV_CACHE=1	Force uncached generation path
TINYLOOP_EXPERIMENTAL_FP16_BODY=1	Dequantize body weights to FP16 caches at load time
TINYLOOP_DISABLE_FLASH2_PREFILL=1	Disable the new tiled prefill attention path
TINYLOOP_DISABLE_CUBLAS_FP16=1	Disable cuBLAS-backed FP16 GEMM path
TINYLOOP_TOKEN_EXIT=N	Enable token-stable early-exit probing in forward scoring
TINYLOOP_TOKEN_EXIT_VERBOSE=1	Print token-stable early-exit diagnostics

Test-only variables:

Variable	Effect
TINYLOOP_TEST_MODEL_PATH	Enable model-dependent regressions in CTest
TINYLOOP_TEST_BINARY	Point direct CLI regression scripts at a built binary
TINYLOOP_TEST_BUILD_DIR	Point Python regressions at a build directory containing tinyloop_py

Validation

Current validation layers include:

• CUDA correctness tests for kernels and fused ops
• cache and decode parity tests
• raw-byte CLI generation regressions
• tokenizer-aware Python regressions
• eval-slice regressions
• hot-op microbenchmarks

Run the main suite with:

ctest --test-dir tinyloop/build --output-on-failure

For model-dependent coverage:

TINYLOOP_TEST_MODEL_PATH=/absolute/path/to/model.tinyloop \
ctest --test-dir tinyloop/build --output-on-failure

Current Limits

Be explicit about these:

• the CLI still uses raw-byte prompt tokenization
• batching and request scheduling are not built yet
• there is no HTTP server surface yet
• the runtime is specialized to a narrow looped-transformer family
• FlashAttention-2 is still an open roadmap item even though the current tiled prefill kernel is a real step toward it

Docs Site

The docs site is published at ethanyangtw.github.io/tinyloop via the Publish Wiki GitHub Actions workflow on every main push.

Run locally

cd tinyloop/wiki
npm ci
npm run start

Build the static site

cd tinyloop/wiki
npm run build

All four locales (en, zh-TW, ko, ja) build from the same content; translations are under wiki/i18n/.