loopgraft

An engine for loop transformers on GB10 (DGX Spark / SM121).

Grafts a damped Runge–Kutta loop over a window of decoder layers of a standard HuggingFace decoder-only transformer — no retraining, no architecture change — and makes it work under KV-cached autoregressive generation, not just single-forward scoring.

Implements the block-mode recipe from Training-Free Looped Transformers (arXiv:2605.23872); all credit for the method to the original authors. loopgraft adds the generation-side machinery: per-pass replay KV caches, a fidelity gate, and a decode benchmark.

capture:  g(x0) = window(x0)            # runs normally, writes the main cache
substeps: x <- x + (window(x) - x)/K    # K damped steps, per-pass replay caches
output:   beta * g(x0) + (1-beta) * x_K

Per-pass replay caches keep cached decoding exactly equivalent to the paper's full-sequence formulation: at step t, substep p attends to past tokens' pass-p key/values, written when those tokens took their own substep p.

Quickstart: gemma-4-12B-it on GB10

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
from loopgraft import GraftConfig, graft, ungraft

MODEL = "google/gemma-4-12B-it"  # or a local snapshot
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL, dtype=torch.bfloat16, device_map="cuda", attn_implementation="sdpa")
model.eval()

# paper recipe: 4-layer window at ~0.5 fractional depth, K=3, beta=0.5
handle = graft(model, GraftConfig(window=(22, 26), K=3, beta=0.5))

ids = tok.apply_chat_template([{"role": "user", "content": "hi"}],
                              add_generation_prompt=True, return_tensors="pt").to("cuda")
handle.new_sequence()  # fresh replay caches — required before EVERY generate()
out = model.generate(ids, past_key_values=DynamicCache(config=model.config.get_text_config()),
                     max_new_tokens=256)
print(tok.decode(out[0, ids.shape[1]:]))

ungraft(handle)  # model restored in place

Fidelity gate

Run before trusting generations. The cache-less full-sequence forward is the exact reference; verify_fidelity teacher-forces continuations and compares next-token logits between cached decode and that reference, calibrated against the ungrafted noise floor (plain-model cached-vs-full disagreement is pure attention-kernel numerics).

from loopgraft import verify_fidelity
cfg = GraftConfig(window=(22, 26), K=3, beta=0.5)
verify_fidelity(model, prompt_ids, {"floor": None, "rk": cfg}, n_steps=48)
verify_fidelity(model, prompt_ids, {"floor": None, "rk": cfg},
                n_steps=1300, cmp_from=1000)  # long gate: cross the sliding-window size

Decode benchmark

python -m loopgraft.bench --model /path/to/model --window 22 26 \
    --batches 1 8 16 24 --ctx 1024 4096 --mem-bw-gbs 273

Reports plain vs grafted decode steps/s across batch sizes and context depths, against two printed references: the weight-bandwidth ceiling (mem_bandwidth / weight_bytes — the right mental model on unified-memory boxes like GB10) and the graft's ideal overhead 1 + K × window/num_layers (its irreducible extra weight traffic; 1.25× for the default 4/48 @ K=3). vs ideal ≈ 1.0 means the loop costs only what it must; well above 1.0 is host/cache overhead — batch amortizes it.

Perf notes: window placement

Looping a global (full-attention) layer is disproportionately expensive: its replay caches grow with full context (× K passes) instead of staying an O(sliding_window) ring, its cache KV is re-read K extra times per token at full length, and on hardware with no fused kernel for the global head geometry (e.g. gemma-4's head_dim-512 MQA globals on GB10) the loop multiplies the slowest attention in the model. The paper-recipe window (22, 26) on gemma-4-12B straddles global layer 23 — gemma-4 places globals every 6th layer (5, 11, …, 47).

sliding_only_window(model) picks the global-free window nearest 0.5 fractional depth instead — (24, 28) on gemma-4-12B — keeping every replay cache a fixed ring:

from loopgraft import sliding_only_window
handle = graft(model, GraftConfig(window=sliding_only_window(model), K=3, beta=0.5))

Perf-wise this is strictly better; quality is recipe-dependent, so re-run the fidelity gate and your evals after moving the window.

Tests

python tests/test_cpu_fidelity.py   # CPU-only, no downloads, ~2 min

Builds a tiny random gemma-4-family model (mixed sliding/global layers, MQA globals, attention_k_eq_v) and asserts cached grafted decode matches the cache-less reference at the ungrafted noise floor, across sliding-ring rollover, for straddling, sliding-only, and naive-mode windows.

Assumptions

• Decoder layers live in one nn.ModuleList and take the shared Cache as a past_key_values kwarg (llama/gemma-family, transformers ≥ 5.x; override cache_kwarg or pass layers= otherwise).
• Replay caches are DynamicCache(config=text_config) so per-layer sliding-ring vs full-length geometry matches the main cache.
• Models with cross-layer KV sharing: audit donor/reader layer positions before grafting a window that straddles a donor.

Attribution

Method: Training-Free Looped Transformers, arXiv:2605.23872. This repo is an independent implementation focused on GB10 inference; please cite the original paper for the technique.

License

Apache-2.0