238 lines
11 KiB
Python
238 lines
11 KiB
Python
"""Design C machinery: k-loop the prompt, then per-token carry generation.
|
|
|
|
Structure per sequence: [prompt] [p pause tokens] [answer]
|
|
1. prefill: prompt positions looped k times through the merge (settled plan)
|
|
2. carry scan: each pause/answer position's band input is
|
|
x_t = merge(e_t, s_{t-1})
|
|
where s_{t-1} is the *previous position's* final band output — the
|
|
latent thought chain runs through the pause positions' workspace states
|
|
(token content of pauses is inert: <unused0>).
|
|
3. suffix -> logits; CE on answer tokens only.
|
|
|
|
The scan is sequential per position (RNN-style through the band); each step
|
|
re-runs the band on the full sequence — causality keeps settled positions
|
|
stable, provisional later positions are recomputed next step anyway.
|
|
"""
|
|
|
|
import torch
|
|
from torch.utils.checkpoint import checkpoint
|
|
|
|
PAUSE_ID = 6 # <unused0>
|
|
|
|
|
|
def _pos_ids(mask):
|
|
return (mask.cumsum(-1) - 1).clamp(min=0)
|
|
|
|
|
|
def prompt_prefill(looper, adapter, e, calls, prompt_mask, k):
|
|
"""k merge->band loops over prompt positions. Returns (S, X)."""
|
|
with torch.no_grad():
|
|
s = looper.band(e, calls)
|
|
x = e
|
|
for _ in range(k):
|
|
x = torch.where(prompt_mask[..., None], adapter(e, s), e)
|
|
s = looper.band(x, calls)
|
|
return s, x
|
|
|
|
|
|
def carry_steps(looper, adapter, e, calls, S, X, step_updates,
|
|
use_checkpoint=False, iter_states=None):
|
|
"""Sequential scan. step_updates: list of (row_idx, pos, [inplace])
|
|
index tensors. inplace=True re-iterates the SAME position, seeding
|
|
from its own previous band output (internal band looping) instead of
|
|
its left neighbor. iter_states: optional list — every inplace
|
|
update's fresh band output at pos is appended (for lens losses)."""
|
|
for upd in step_updates:
|
|
rows, pos = upd[0], upd[1]
|
|
inplace = len(upd) > 2 and upd[2]
|
|
if rows.numel() == 0:
|
|
continue
|
|
seed = S[rows, pos] if inplace else S[rows, pos - 1]
|
|
x_new = adapter(e[rows, pos], seed)
|
|
X = X.clone()
|
|
X[rows, pos] = x_new.to(X.dtype)
|
|
S = (checkpoint(lambda X_: looper.band(X_, calls), X,
|
|
use_reentrant=False) if use_checkpoint
|
|
else looper.band(X, calls))
|
|
if inplace and iter_states is not None:
|
|
iter_states.append(S[rows, pos])
|
|
return S, X
|
|
|
|
|
|
def build_step_updates(prompt_lens, total_lens, device):
|
|
"""For right-padded batches: step j updates row b at prompt_lens[b]+j."""
|
|
max_span = int((total_lens - prompt_lens).max())
|
|
updates = []
|
|
for j in range(max_span):
|
|
pos = prompt_lens + j
|
|
sel = pos < total_lens
|
|
rows = torch.nonzero(sel, as_tuple=False).squeeze(-1)
|
|
updates.append((rows.to(device), pos[sel].to(device)))
|
|
return updates
|
|
|
|
|
|
def splice_inner_iters(updates, inner_iters, inner_at, prompt_lens, dev, B):
|
|
"""Insert inner_iters in-place band iterations at the (batch-uniform
|
|
offset) anchor position, right after the scan first settles it."""
|
|
j_anchor = int((inner_at - prompt_lens).max())
|
|
rows = torch.arange(B, device=dev)
|
|
inner = [(rows, inner_at.to(dev), True)] * inner_iters
|
|
if j_anchor < 0:
|
|
# anchor is the last PROMPT position (no pauses at all):
|
|
# iterate there before the scan enters the visible tokens
|
|
return inner + updates
|
|
out = []
|
|
for j, u in enumerate(updates):
|
|
out.append(u)
|
|
if j == j_anchor:
|
|
out += inner
|
|
return out
|
|
|
|
|
|
def sym_iterate(looper, adapter, proj, e, calls, S, X, rows, anchor, m,
|
|
lens_fn, embed_w, sym_tf=None, start_id=None, topk=32,
|
|
use_checkpoint=False, iter_states=None):
|
|
"""Item 32: discrete latent chain at the anchor. Each tick reads the
|
|
previous anchor state through the lens, snaps it to a token
|
|
(straight-through over top-k) or takes the teacher token (sym_tf:
|
|
(B, m) ids, teacher forcing), and feeds that token's embedding back
|
|
through a zero-init projector ALONGSIDE the analog carry:
|
|
x_i = merge(e, s_{i-1}) + proj(E(sym))
|
|
Tick 0 uses start_id (a newline: 'a step begins')."""
|
|
for i in range(m):
|
|
s_prev = S[rows, anchor]
|
|
if sym_tf is not None:
|
|
symb = embed_w[sym_tf[:, i]]
|
|
elif i == 0:
|
|
symb = embed_w[torch.full((rows.shape[0],), start_id,
|
|
device=e.device)]
|
|
else:
|
|
logits = lens_fn(s_prev).float()
|
|
p, idx = torch.softmax(logits, -1).topk(topk, dim=-1)
|
|
p = p / p.sum(-1, keepdim=True)
|
|
soft = (p.unsqueeze(-1) * embed_w[idx].float()).sum(-2)
|
|
hard = embed_w[idx[:, 0]].float()
|
|
symb = hard + soft - soft.detach()
|
|
x_new = (adapter(e[rows, anchor], s_prev).float()
|
|
+ proj(symb.float()))
|
|
X = X.clone()
|
|
X[rows, anchor] = x_new.to(X.dtype)
|
|
S = (checkpoint(lambda X_: looper.band(X_, calls), X,
|
|
use_reentrant=False) if use_checkpoint
|
|
else looper.band(X, calls))
|
|
if iter_states is not None:
|
|
iter_states.append(S[rows, anchor])
|
|
return S, X
|
|
|
|
|
|
def carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens,
|
|
k, use_checkpoint=False, feedforward=False,
|
|
return_states=False, inner_iters=0, inner_at=None,
|
|
iter_states=None):
|
|
"""Teacher-forced design-C forward (right-padded batch).
|
|
|
|
feedforward=True: pause-token control — same positions get the adapter as
|
|
x_t = merge(e_t, e_t) in ONE band pass; no state propagates between
|
|
positions. Isolates 'pause compute + weights' from the recurrence."""
|
|
dev = input_ids.device
|
|
calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1)
|
|
e = looper._hin[looper.l0].detach()
|
|
ar = torch.arange(input_ids.shape[1], device=dev)
|
|
prompt_mask = ar[None, :] < prompt_lens[:, None].to(dev)
|
|
if feedforward:
|
|
touched = attention_mask.bool()
|
|
x = torch.where(touched[..., None], adapter(e, e), e)
|
|
S = (checkpoint(lambda x_: looper.band(x_, calls), x,
|
|
use_reentrant=False) if use_checkpoint
|
|
else looper.band(x, calls))
|
|
out = looper.suffix_logits(S, calls)
|
|
return (out, S) if return_states else out
|
|
S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k)
|
|
total_lens = attention_mask.sum(-1)
|
|
updates = build_step_updates(prompt_lens.to(dev), total_lens.to(dev), dev)
|
|
if inner_iters and inner_at is not None:
|
|
updates = splice_inner_iters(updates, inner_iters, inner_at,
|
|
prompt_lens.to(dev), dev,
|
|
input_ids.shape[0])
|
|
S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
|
|
use_checkpoint=use_checkpoint,
|
|
iter_states=iter_states)
|
|
out = looper.suffix_logits(S, calls)
|
|
return (out, S) if return_states else out
|
|
|
|
|
|
@torch.no_grad()
|
|
def generate_carry_c(looper, adapter, tok, input_ids, attention_mask,
|
|
k, p, max_new_tokens=10, feedforward=False,
|
|
inner_iters=0, kvmem=None, symchain=None):
|
|
"""Greedy design-C generation (left-padded batch, uniform positions).
|
|
|
|
Appends p pause tokens, prefill-loops the prompt, carries through the
|
|
pauses, then generates with per-token carry."""
|
|
B = input_ids.shape[0]
|
|
dev = input_ids.device
|
|
pauses = torch.full((B, p), PAUSE_ID, dtype=torch.long, device=dev)
|
|
ids = torch.cat([input_ids, pauses], 1)
|
|
mask = torch.cat([attention_mask,
|
|
torch.ones_like(pauses)], 1)
|
|
n_prompt = input_ids.shape[1]
|
|
|
|
eos = {tok.eos_token_id, tok.convert_tokens_to_ids("<end_of_turn>")}
|
|
done = torch.zeros(B, dtype=torch.bool, device=dev)
|
|
X_store = None # merged band inputs for settled positions
|
|
for step in range(max_new_tokens + 1): # step 0 = prefill + pause scan
|
|
calls, _ = looper.capture(ids, mask, logits_to_keep=1,
|
|
position_ids=_pos_ids(mask))
|
|
e = looper._hin[looper.l0]
|
|
if feedforward:
|
|
x = torch.where(mask.bool()[..., None], adapter(e, e), e)
|
|
S = looper.band(x, calls)
|
|
X = x
|
|
elif X_store is None:
|
|
ar = torch.arange(ids.shape[1], device=dev)
|
|
prompt_mask = (ar[None, :] < n_prompt) & mask.bool()
|
|
S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k)
|
|
updates = [(torch.arange(B, device=dev),
|
|
torch.full((B,), n_prompt + j, device=dev,
|
|
dtype=torch.long)) for j in range(p)]
|
|
if symchain is not None and inner_iters:
|
|
anchor_sc = torch.full((B,), n_prompt + p - 1, device=dev,
|
|
dtype=torch.long)
|
|
S, X = sym_iterate(
|
|
looper, adapter, symchain["proj"], e, calls, S, X,
|
|
torch.arange(B, device=dev), anchor_sc, inner_iters,
|
|
symchain["lens_fn"], symchain["embed_w"],
|
|
start_id=symchain["start_id"])
|
|
updates = updates # pauses (if any) already handled above
|
|
elif inner_iters:
|
|
anchor = torch.full((B,), n_prompt + p - 1, device=dev,
|
|
dtype=torch.long)
|
|
inner = [(torch.arange(B, device=dev), anchor, True)
|
|
for _ in range(inner_iters)]
|
|
# p=0: iterate at the last prompt position, before any
|
|
# visible token — no pause tokens involved
|
|
updates = (inner + updates) if p == 0 else (updates + inner)
|
|
itst = [] if (kvmem is not None and inner_iters) else None
|
|
S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
|
|
iter_states=itst)
|
|
if kvmem is not None and itst:
|
|
from kv_memory import arm_memory
|
|
arm_memory(kvmem(torch.stack(itst, 1)))
|
|
else:
|
|
X = torch.cat([X_store, e[:, X_store.shape[1]:]], 1)
|
|
rows = torch.arange(B, device=dev)
|
|
pos = torch.full((B,), ids.shape[1] - 1, device=dev,
|
|
dtype=torch.long)
|
|
S = looper.band(X, calls) # settled prefix + provisional new pos
|
|
S, X = carry_steps(looper, adapter, e, calls, S, X, [(rows, pos)])
|
|
X_store = X
|
|
logits = looper.suffix_logits(S, calls, last_only=True)
|
|
nxt = logits[:, -1].argmax(-1)
|
|
nxt = torch.where(done, torch.full_like(nxt, list(eos)[0]), nxt)
|
|
ids = torch.cat([ids, nxt[:, None]], 1)
|
|
mask = torch.cat([mask, (~done)[:, None].long()], 1)
|
|
done |= torch.tensor([t.item() in eos for t in nxt], device=dev)
|
|
if done.all():
|
|
break
|
|
return ids
|