jspace/scripts/carry_common.py

"""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):
    """Sequential scan. step_updates: list of (row_idx, pos) index tensors."""
    for rows, pos in step_updates:
        if rows.numel() == 0:
            continue
        seed = 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))
    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 carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens,
                 k, use_checkpoint=False, feedforward=False):
    """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))
        return looper.suffix_logits(S, calls)
    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)
    S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
                       use_checkpoint=use_checkpoint)
    return looper.suffix_logits(S, calls)


@torch.no_grad()
def generate_carry_c(looper, adapter, tok, input_ids, attention_mask,
                     k, p, max_new_tokens=10, feedforward=False):
    """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)]
            S, X = carry_steps(looper, adapter, e, calls, S, X, updates)
        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