#!/usr/bin/env python3 """Standalone inference for Blink-1. One file, no install beyond torch. pip install torch python infer.py # base, default prompts python infer.py --instruct # instruct (ChatML) variant python infer.py --prompt "the " -n 120 Blink is a *byte* model: input and output are raw UTF-8 bytes (vocab 257 = 256 bytes + EOS=256). There is no tokenizer to download. It is 1,087 parameters, so coherent words are the exception, not the rule. That's the point. The model code below is the exact Blink architecture (a looped shared transformer block with a per-loop LoRA and per-iteration embedding). It is copied here on purpose so this file runs on its own — you do not need the training repo to load a 13 KB model. """ from __future__ import annotations import argparse import math import os from dataclasses import dataclass, fields import torch from torch import Tensor, nn from torch.nn import functional as F EOS_TOKEN_ID = 256 # Blink is deployed at 8 loops through the shared block (the trained # shared_loops=2 in the config is a floor; inference uses the full loop budget, # exactly like the training repo's CLI). At 2 loops the model collapses to # emitting spaces; at 8 it actually moves. BLINK_LOOPS = 8 # -------------------------------------------------------------------------- # config (only the fields the checkpoint carries; extras are ignored on load) # -------------------------------------------------------------------------- @dataclass class ModelConfig: vocab_size: int = 257 dim: int = 3 n_heads: int = 1 prelude_layers: int = 0 coda_layers: int = 0 shared_loops: int = 2 lora_rank: int = 2 ffn_hidden: int = 6 max_seq_len: int = 4096 max_context_len: int = 8192 rope_base: float = 10_000_000.0 index_dim: int = 2 index_top_k: int = 16 local_window: int = 16 sparse_chunk_threshold: int = 4096 sparse_chunk_queries: int = 1024 ssa_block_size: int = 32 ssa_top_k_blocks: int = 2 attention_window: int = 256 fp4_weights: bool = False thinking_enabled: bool = False thinking_steps_min: int = 0 thinking_steps_max: int = 0 thinking_loss_weight: float = 0.0 # -------------------------------------------------------------------------- # model (verbatim Blink architecture) # -------------------------------------------------------------------------- def build_rope_cache(config: ModelConfig, length: int) -> tuple[Tensor, Tensor]: head_dim = config.dim // config.n_heads positions = torch.arange(length, dtype=torch.float32) inv_freq = 1.0 / (config.rope_base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) angles = torch.outer(positions, inv_freq) return torch.cos(angles), torch.sin(angles) def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: x_even = x[..., 0::2] x_odd = x[..., 1::2] rotated_even = x_even * cos - x_odd * sin rotated_odd = x_even * sin + x_odd * cos return torch.stack((rotated_even, rotated_odd), dim=-1).flatten(-2) def causal_attention(q: Tensor, k: Tensor, v: Tensor, window: int) -> Tensor: """Causal (optionally sliding-window) attention. Folds Blink's two training-time backends — padded SDPA for short sequences and flex sliding-window attention once seq_len exceeds the window — into one masked softmax. Both use scale = 1/sqrt(head_dim) on the post-RoPE query dim, so the result is identical to either backend. """ head_dim_qk = q.shape[-1] head_dim_v = v.shape[-1] scale = 1.0 / math.sqrt(head_dim_qk) scores = (q @ k.transpose(-2, -1)) * scale seq = q.shape[-2] qi = torch.arange(seq, device=q.device).view(seq, 1) ki = torch.arange(seq, device=q.device).view(1, seq) mask = ki <= qi if window > 0: mask = mask & (qi - ki < window) scores = scores.masked_fill(~mask, float("-inf")) return F.softmax(scores, dim=-1) @ v[..., :head_dim_v] class SwiGlu(nn.Module): def __init__(self, dim: int, hidden: int) -> None: super().__init__() self.gate_up = nn.Linear(dim, 2 * hidden, bias=False) self.down = nn.Linear(hidden, dim, bias=False) def forward(self, x: Tensor) -> Tensor: gate, up = self.gate_up(x).chunk(2, dim=-1) return self.down(F.silu(gate) * up) class Attention(nn.Module): def __init__(self, config: ModelConfig) -> None: super().__init__() self.config = config self.n_heads = config.n_heads self.head_dim = config.dim // config.n_heads self.qkv = nn.Linear(config.dim, 3 * config.dim, bias=False) self.out = nn.Linear(config.dim, config.dim, bias=False) def forward(self, x: Tensor, cos: Tensor, sin: Tensor, qkv_delta: Tensor | None) -> Tensor: batch, seq_len, dim = x.shape qkv = self.qkv(x) if qkv_delta is not None: qkv = qkv + qkv_delta q, k, v = qkv.split(dim, dim=-1) q = q.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2) k = k.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2) v = v.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2) q = apply_rope(q, cos, sin) k = apply_rope(k, cos, sin) window = self.config.attention_window use_window = 0 < window < seq_len attended = causal_attention(q, k, v, window if use_window else 0) merged = attended.transpose(1, 2).reshape(batch, seq_len, dim) return self.out(merged) class Block(nn.Module): def __init__(self, config: ModelConfig) -> None: super().__init__() self.attn_norm = nn.RMSNorm(config.dim) self.attn = Attention(config) self.ffn_norm = nn.RMSNorm(config.dim) self.ffn = SwiGlu(config.dim, config.ffn_hidden) def forward(self, x: Tensor, cos: Tensor, sin: Tensor, qkv_delta: Tensor | None) -> Tensor: x = x + self.attn(self.attn_norm(x), cos, sin, qkv_delta) return x + self.ffn(self.ffn_norm(x)) class LoopLora(nn.Module): def __init__(self, config: ModelConfig, max_loops: int) -> None: super().__init__() self.down = nn.ModuleList(nn.Linear(config.dim, config.lora_rank, bias=False) for _ in range(max_loops)) self.up = nn.ModuleList(nn.Linear(config.lora_rank, 3 * config.dim, bias=False) for _ in range(max_loops)) def forward(self, x: Tensor, loop_index: int) -> Tensor: clamped = min(loop_index, len(self.down) - 1) return self.up[clamped](self.down[clamped](x)) class SparseIndexer(nn.Module): """Holds the SSA gate parameter. Never fires for Blink (short sequences, tiny dim), so it is load-only — present to match the released weights.""" def __init__(self) -> None: super().__init__() self.gate = nn.Parameter(torch.tensor([0.1])) class Blink(nn.Module): def __init__(self, config: ModelConfig, max_loops: int = 8) -> None: super().__init__() self.config = config self.max_loops = max_loops self.embed = nn.Embedding(config.vocab_size, config.dim) self.indexer = SparseIndexer() self.prelude = nn.ModuleList(Block(config) for _ in range(config.prelude_layers)) self.shared = Block(config) self.loop_lora = LoopLora(config, max_loops) self.loop_embed = nn.Embedding(max_loops, config.dim) self.coda = nn.ModuleList(Block(config) for _ in range(config.coda_layers)) self.final_norm = nn.RMSNorm(config.dim) @torch.no_grad() def forward(self, tokens: Tensor, loops: int | None = None) -> Tensor: loop_count = loops if loops is not None else self.config.shared_loops seq_len = tokens.shape[1] cos, sin = build_rope_cache(self.config, seq_len) cos, sin = cos.to(tokens.device), sin.to(tokens.device) x = self.embed(tokens) for block in self.prelude: x = block(x, cos, sin, None) for loop_index in range(loop_count): clamped = min(loop_index, self.max_loops - 1) gated = x + self.loop_embed.weight[clamped] delta = self.loop_lora(gated, loop_index) x = self.shared(gated, cos, sin, delta) for block in self.coda: x = block(x, cos, sin, None) x = self.final_norm(x) return F.linear(x, self.embed.weight) # -------------------------------------------------------------------------- # loading + generation # -------------------------------------------------------------------------- def load_model(checkpoint_path: str) -> tuple[Blink, ModelConfig]: payload = torch.load(checkpoint_path, map_location="cpu", weights_only=False) known = {f.name for f in fields(ModelConfig)} config = ModelConfig(**{k: v for k, v in payload["model_config"].items() if k in known}) max_loops = payload["model"]["loop_embed.weight"].shape[0] model = Blink(config, max_loops=max_loops) model.load_state_dict(payload["model"]) model.eval() return model, config def sample_next_token(logits, temperature, top_k, repetition_penalty, recent, rng): if repetition_penalty != 1.0 and recent: logits = logits.clone() idx = torch.tensor(sorted(set(recent)), dtype=torch.long) sel = logits[idx] logits[idx] = torch.where(sel > 0, sel / repetition_penalty, sel * repetition_penalty) if temperature <= 0: return int(torch.argmax(logits)) probs = F.softmax(logits / temperature, dim=-1) if top_k and top_k > 0: v, i = torch.topk(probs, min(top_k, probs.size(-1))) return int(i[torch.multinomial(v, 1, generator=rng)].item()) return int(torch.multinomial(probs, 1, generator=rng).item()) def generate(model, prompt, max_new_tokens=80, temperature=0.5, top_k=5, repetition_penalty=1.1, repetition_window=128, seed=0): rng = torch.Generator() if seed is not None: rng.manual_seed(seed) tokens = list(prompt.encode("utf-8")) generated: list[int] = [] for _ in range(max_new_tokens): x = torch.tensor([tokens[-model.config.max_context_len:]], dtype=torch.long) logits = model(x, loops=BLINK_LOOPS)[0, -1] nxt = sample_next_token(logits, temperature, top_k, repetition_penalty, tokens[-repetition_window:], rng) if nxt == EOS_TOKEN_ID: break tokens.append(nxt) generated.append(nxt) if b"<|im_end|>" in bytes(generated[-60:]): break return bytes(generated).decode("utf-8", errors="replace").removesuffix("<|im_end|>") CHATML = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" BASE_DECODE = dict(temperature=0.5, top_k=5, repetition_penalty=1.1) INSTRUCT_DECODE = dict(temperature=0.2, top_k=10, repetition_penalty=1.5) def main() -> None: ap = argparse.ArgumentParser(description="Run Blink-1 byte-level generation.") ap.add_argument("--instruct", action="store_true", help="use the instruct (ChatML) champion") ap.add_argument("--prompt", help="single prompt; with --instruct it is wrapped in ChatML") ap.add_argument("-n", "--max-new-tokens", type=int, default=80) ap.add_argument("--seed", type=int, default=0) args = ap.parse_args() here = os.path.dirname(os.path.abspath(__file__)) ckpt = "blink-1-instruct.pt" if args.instruct else "blink-1-base.pt" model, config = load_model(os.path.join(here, ckpt)) nparams = sum(p.numel() for p in model.parameters()) print(f"loaded {ckpt}: {nparams} params, dim={config.dim}, shared_loops={config.shared_loops}\n") decode = INSTRUCT_DECODE if args.instruct else BASE_DECODE if args.prompt is not None: prompts = [CHATML.format(args.prompt) if args.instruct else args.prompt] elif args.instruct: prompts = [CHATML.format("hello"), CHATML.format("what is your name?")] else: prompts = ["the ", "I think ", "once upon "] for p in prompts: out = generate(model, p, max_new_tokens=args.max_new_tokens, seed=args.seed, **decode) print(f"prompt: {p!r}") print(f"output: {out!r}\n") if __name__ == "__main__": main()