sampler.py

"""Generation loop of the adaptive hybrid: AR control, in-block denoising, AR answer.

Numpy port of the inference paths in ``diffusion_lm.hybrid`` and the playground's
streaming loop. Sampling stays in fp64 (Gumbel argmax, the reference semantics of the
torch backend); the RNG is a ``numpy.random.Generator``, so sampled transcripts match the
torch serving path statistically, not token-for-token.
"""

from __future__ import annotations

import time

import numpy as np

from masks import adaptive_block_mask, prefix_causal_blocked
from model import Denoiser

FLOAT_MIN = np.finfo(np.float32).min


def sample_categorical(logits: np.ndarray, temperature: float, rng: np.random.Generator):
    """Sample one token per row with fp64 Gumbel noise; argmax at temperature zero."""

    if temperature < 0:
        raise ValueError('temperature must be non-negative')
    if temperature == 0:
        return logits.argmax(axis=-1)
    scaled = logits.astype(np.float64) / temperature
    uniform = rng.random(scaled.shape).clip(1e-12, 1.0 - 1e-12)
    return (scaled - np.log(-np.log(uniform))).argmax(axis=-1)


def apply_repetition_penalty(
    logits: np.ndarray, token_ids: list[int], penalty: float
) -> np.ndarray:
    """Divide logits of already-emitted tokens by ``penalty`` (CTRL convention)."""

    if penalty == 1.0 or not token_ids:
        return logits
    index = np.asarray(sorted(set(token_ids)))
    adjusted = logits.copy()
    selected = logits[index]
    adjusted[index] = np.where(selected > 0, selected / penalty, selected * penalty)
    return adjusted


def apply_top_p(logits: np.ndarray, top_p: float) -> np.ndarray:
    """Restrict sampling to the smallest set of tokens whose mass reaches ``top_p``."""

    if top_p >= 1.0:
        return logits
    order = np.argsort(-logits, kind='stable')
    ordered = logits[order].astype(np.float64)
    shifted = np.exp(ordered - ordered.max())
    probabilities = shifted / shifted.sum()
    remove = probabilities.cumsum() - probabilities >= top_p
    ordered = np.where(remove, FLOAT_MIN, ordered)
    restricted = np.empty_like(logits)
    restricted[order] = ordered.astype(logits.dtype)
    return restricted


def predict_control(
    denoiser: Denoiser,
    sequence: list[int],
    control_ids: list[int],
    *,
    prefix_len: int,
    causal_prefix: bool,
    temperature: float,
    rng: np.random.Generator,
) -> int:
    """Predict the next control token, restricted to the size/stop menu."""

    seq_len = len(sequence)
    blocked = prefix_causal_blocked(prefix_len, seq_len, causal_prefix=causal_prefix)
    logits = denoiser.logits(sequence, blocked, [seq_len - 1])[0]
    restricted = np.full_like(logits, FLOAT_MIN)
    index = np.asarray(control_ids)
    restricted[index] = logits[index]
    return int(sample_categorical(restricted[None, :], temperature, rng)[0])


def denoise_block_steps(
    denoiser: Denoiser,
    sequence: list[int],
    size: int,
    *,
    problem_len: int,
    steps: int,
    temperature: float,
    causal_prefix: bool,
    rng: np.random.Generator,
    cache_prefix: bool = False,
):
    """Yield ``(step, block_ids)`` while ancestrally denoising one ``<szN>`` block.

    The ancestral reverse transition reveals each still-masked position with probability
    ``1 / steps_left``, so the final step always completes the block. ``cache_prefix``
    encodes the context once per block and re-feeds the changing rows each step, trimming
    the cache back in between. The cache boundary sits BEFORE the ``<szN>`` token: the
    trained mask groups that token with its own block, attended bidirectionally, so its
    keys depend on the block's current state and must be recomputed every step — priming
    it with the context would freeze a key that never saw the block. Everything earlier
    never attends forward into the new block, which is what makes its cache sound.
    """

    mdlm = denoiser.config['mdlm']
    mask_id = mdlm['mask_token_id']
    size_ids = tuple(int(i) for i in mdlm['size_ids'].values())
    window = np.asarray([*sequence, *([mask_id] * size)], dtype=np.int64)
    window_prefix = len(sequence)
    prime_len = window_prefix - 1
    blocked = adaptive_block_mask(
        window, problem_len, window.shape[0], size_ids, mdlm['end_think_id'],
        causal_prefix=causal_prefix,
    )
    caches = None
    if cache_prefix:
        caches = denoiser.new_caches()
        denoiser.logits(window[:prime_len], blocked[:prime_len, :prime_len], None, caches)

    for step in range(steps):
        masked = window == mask_id
        if not masked.any():
            break
        if caches is not None:
            denoiser.trim_caches(caches, prime_len)
            positions = np.flatnonzero(masked[prime_len:])
            logits = denoiser.logits(
                window[prime_len:], blocked[prime_len:, :], positions, caches
            )
        else:
            logits = denoiser.logits(window, blocked, np.flatnonzero(masked))
        predictions = sample_categorical(logits, temperature, rng)
        proposed = window.copy()
        proposed[masked] = predictions
        reveal_probability = 1.0 / (steps - step)
        reveal = (rng.random(window.shape[0]) < reveal_probability) & masked
        window = np.where(reveal, proposed, window)
        yield step + 1, [int(t) for t in window[window_prefix:]]

    if (window == mask_id).any():
        raise RuntimeError('sampler finished with masked positions')


def decode_answer_steps(
    denoiser: Denoiser,
    sequence: list[int],
    *,
    prefix_len: int,
    stop_ids: tuple[int, ...],
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
    causal_prefix: bool,
    rng: np.random.Generator,
):
    """Yield answer tokens decoded autoregressively with the prefix in a key/value cache."""

    caches = denoiser.new_caches()
    generated: list[int] = []
    step_ids = list(sequence)
    for _ in range(max(0, max_new_tokens)):
        seen = caches[0].offset
        length = seen + len(step_ids)
        blocked = prefix_causal_blocked(prefix_len, length, causal_prefix=causal_prefix)
        logits = denoiser.logits(step_ids, blocked[seen:, :], [len(step_ids) - 1], caches)[0]
        logits = apply_repetition_penalty(logits, generated, repetition_penalty)
        logits = apply_top_p(logits, top_p)
        token = int(sample_categorical(logits[None, :], temperature, rng)[0])
        generated.append(token)
        yield token
        if token in stop_ids:
            break
        step_ids = [token]


def block_contents(
    think_ids: list[int], *, size_by_id: dict[int, int], pad_id: int, end_id: int
) -> list[tuple[int, list[int]]]:
    """Split the think stream into ``(size, content ids)`` blocks by walking ``<szN>``."""

    blocks = []
    cut = 0
    while cut < len(think_ids):
        size = size_by_id.get(think_ids[cut])
        if size is None:
            cut += 1
            continue
        block = think_ids[cut + 1 : cut + 1 + size]
        blocks.append((size, [t for t in block if t not in (pad_id, end_id)]))
        cut += 1 + size
    return blocks


def stream_turn(
    denoiser: Denoiser,
    prompt_ids: list[int],
    *,
    rng: np.random.Generator,
    temperature: float = 0.8,
    top_p: float = 0.95,
    repetition_penalty: float = 1.0,
    steps_per_block: int | None = None,
    max_answer_tokens: int | None = None,
    max_blocks: int = 48,
    # Prefix caching during denoising measured 1.7x faster per block on M1 Pro (2026-08-14,
    # 2.64s vs 4.44s for sz32 at prefix 500) with cached==uncached logits at fp32 noise level.
    cache_blocks: bool = True,
):
    """Stream one assistant turn: control loop, per-block denoising, AR answer.

    Yields ``('block', index, size, step, total, block_ids)`` during thinking,
    ``('answer', token_id)`` per answer token, and one final ``('done', result)`` where
    ``result`` carries the think/answer ids, per-block notes, and wall times.
    """

    mdlm = denoiser.config['mdlm']
    max_seq_len = mdlm['max_seq_len']
    end_think_id = mdlm['end_think_id']
    size_by_id = {int(i): int(size) for size, i in mdlm['size_ids'].items()}
    control_ids = [*size_by_id.keys(), end_think_id]
    causal_prefix = mdlm['causal_prefix']
    steps = steps_per_block or mdlm['steps_per_block']
    answer_limit = max_answer_tokens or mdlm['max_answer_tokens']

    sequence = [*prompt_ids, mdlm['think_id']]
    prefix_len = len(prompt_ids) + 1
    terminated = False
    think_started = time.perf_counter()
    for block_index in range(max_blocks):
        control = predict_control(
            denoiser, sequence, control_ids,
            prefix_len=prefix_len, causal_prefix=causal_prefix, temperature=0.0, rng=rng,
        )
        if control == end_think_id:
            terminated = True
            break
        size = size_by_id[control]
        if len(sequence) + 1 + size > max_seq_len - 8:
            break
        sequence.append(control)
        block: list[int] = []
        for step, block_ids in denoise_block_steps(
            denoiser, sequence, size,
            problem_len=len(prompt_ids), steps=steps, temperature=temperature,
            causal_prefix=causal_prefix, rng=rng, cache_prefix=cache_blocks,
        ):
            block = block_ids
            yield 'block', block_index, size, step, steps, block_ids
        sequence.extend(block)
    sequence.append(end_think_id)
    think_seconds = time.perf_counter() - think_started
    think_ids = sequence[prefix_len:]

    answer_ids: list[int] = []
    answer_started = time.perf_counter()
    for token in decode_answer_steps(
        denoiser, sequence,
        prefix_len=prefix_len, stop_ids=(mdlm['im_end_id'], mdlm['eos_id']),
        max_new_tokens=min(answer_limit, max_seq_len - len(sequence)),
        temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty,
        causal_prefix=causal_prefix, rng=rng,
    ):
        answer_ids.append(token)
        yield 'answer', token

    yield 'done', {
        'think_ids': think_ids,
        'answer_ids': answer_ids,
        'blocks': block_contents(
            think_ids, size_by_id=size_by_id,
            pad_id=mdlm['thought_pad_id'], end_id=end_think_id,
        ),
        'terminated': terminated,
        'think_seconds': think_seconds,
        'answer_seconds': time.perf_counter() - answer_started,
    }