jeeves_src/inference/engine.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, NamedTuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file
from torch.nn.attention import sdpa_kernel

from export import load_export
from inference.types import PRECISIONS, Options, Result
from model.config import LINEAR
from model import metal
from model.model import (EFFICIENT, FrozenRMSNorm, RMSNorm, apply_rotary, conv_step, decode_attention, delta_gates,
                         eager_conv_step, gated_delta_rule_advance_inplace, gated_delta_rule_chunk, gated_delta_rule_step, metal_enabled,
                         rotate_into_cache)
from drafter.view import DrafterView
from prep.format import DataFormat, Question

# SDPA's efficient kernel pads a copy of any mask whose row stride is not aligned, in every call.
MASK_ROW_ALIGNMENT = 16
# Longer passes run eagerly at their exact length: there the GPU work outweighs the launch cost a graph removes, and padding would add to it.
EXTEND_GRAPH_LENGTHS = (32, 64, 128, 256, 512, 1024, 1536, 2048)
# Passes up to this length round their row count up like decode cycles, so fewer graphs are captured; on longer passes the extra rows would cost real GPU time.
ROUNDED_ROWS_MAX_LENGTH = 256
MANY_QUESTIONS = 3


@dataclass(frozen=True)
class DecodeBlock:
    K: int
    J: int
    M: int
    ar_k: torch.Tensor
    ar_j: torch.Tensor
    mask_offset: torch.Tensor
    limit_offset: torch.Tensor
    allow_mm: torch.Tensor


def decode_block(K: int, device: torch.device) -> DecodeBlock:
    J = K - 1
    ar_k, ar_j = torch.arange(K, device=device), torch.arange(J, device=device)
    block_of_mask = ar_k.repeat_interleave(J)
    row_block = torch.cat((torch.full((K,), -1, device=device), block_of_mask))
    allow_mm = (row_block[:, None] == block_of_mask[None, :]) & (row_block[:, None] >= 0)
    return DecodeBlock(K, J, K * J, ar_k, ar_j, block_of_mask + (ar_j + 1).repeat(K), torch.cat((ar_k, block_of_mask)), allow_mm)


def store_transposed(linear: nn.Linear) -> None:
    linear.weight = nn.Parameter(linear.weight.t().contiguous().t(), requires_grad=False)


def merge_linears_in_place(linears: tuple[nn.Module, ...]) -> nn.Module | None:
    if any(linear.bias is not None for linear in linears):
        return None
    if all(type(linear) is nn.Linear for linear in linears):
        sizes = [linear.out_features for linear in linears]
        weight = torch.cat([linear.weight.detach() for linear in linears])
        merged = nn.Linear(linears[0].in_features, sum(sizes), bias=False, device="meta")
        merged.weight = nn.Parameter(weight, requires_grad=False)
        for linear, part in zip(linears, weight.split(sizes)):
            linear.weight = nn.Parameter(part, requires_grad=False)
        return merged
    from inference.fp8 import FP8Linear
    return FP8Linear.concatenated(list(linears)) if all(type(linear) is FP8Linear for linear in linears) else None


def freeze_rms_norms(module: nn.Module) -> None:
    for name, child in module.named_children():
        if type(child) is RMSNorm:
            setattr(module, name, FrozenRMSNorm(child))
        else:
            freeze_rms_norms(child)


def additive_mask(allowed: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
    padded = -(-allowed.shape[-1] // MASK_ROW_ALIGNMENT) * MASK_ROW_ALIGNMENT
    mask = torch.full((*allowed.shape[:-1], padded), float("-inf"), dtype=dtype, device=allowed.device)[..., :allowed.shape[-1]]
    return mask.masked_fill_(allowed, 0.0)


class GroupRow(NamedTuple):
    record: DataFormat
    question: Question
    options: Options


def pack_requests(sizes: list[int], capacity: int) -> list[list[tuple[int, int]]]:
    groups: list[list[tuple[int, int]]] = []
    for request, size in enumerate(sizes):
        questions = [(request, i) for i in range(size)]
        for start in range(0, size, capacity):
            part = questions[start:start + capacity]
            group = next((g for g in groups if len(g) + len(part) <= capacity), None)
            if group is None:
                groups.append(part)
            else:
                group.extend(part)
    return groups


def common_prefix(seqs: list[list[int]]) -> int:
    n = min(len(s) for s in seqs) - 1
    for i in range(n):
        t = seqs[0][i]
        if any(s[i] != t for s in seqs):
            return i
    return n


class Engine:
    def __init__(self, model: str, drafter: str, block: int = 4, precision: str = "bf16", max_rows: int = 8, max_len: int = 8192,
                 window_step: int = 512, sync_every: int = 1, device: str | None = None):
        self.device = dev = torch.device(device) if device else torch.accelerator.current_accelerator()
        self.mps, self.cuda = dev.type == "mps", dev.type == "cuda"
        if precision not in PRECISIONS:
            raise ValueError(f"precision must be one of {PRECISIONS}, got {precision!r}")
        self.precision = precision
        fp8 = precision == "fp8"
        fp8_supported = self.metal if self.mps else (self.cuda and torch.cuda.get_device_capability(dev) >= (8, 9))
        if fp8 and not fp8_supported:
            raise ValueError("precision='fp8' needs a CUDA GPU with compute capability 8.9 or higher, or MPS with the Metal kernels enabled "
                             "(unset QWEN35_KERNELS=0)")
        if fp8:
            if self.mps:
                from inference.fp8_metal import load_view
            else:
                from inference.fp8 import load_view
            view, head, encoder = load_view(model, drafter, block, dev)
            base = view.base
            self.dtype = base.model.embed_tokens.weight.dtype
        else:
            base, head, encoder = load_export(model, device=dev)
            view = DrafterView(base, block=block).to(dev)
            loaded = view.load_state_dict(load_file(drafter, device=str(dev)), strict=False)
            missing = [k for k in loaded.missing_keys if not k.startswith("base.")]
            if loaded.unexpected_keys or missing:
                raise ValueError(f"drafter keys do not match: unexpected {loaded.unexpected_keys[:3]}, missing {missing[:3]}")
            view.requires_grad_(False)
            self.dtype = base.lm_head.weight.dtype
            for proj in view.projections():
                proj.to(self.dtype)
        if self.mps and not fp8 and torch.backends.mps.is_macos_or_newer(15, 0):
            # On torch 2.14, MPS computes x @ W.T faster, and bitwise equal, with W.T contiguous once x has 10 or more rows, but slower with fewer; so only projections fed the K + M cycle rows switch, except attn_k/attn_v, which are slower even at M rows.
            for layer in base.model.layers:
                mixer = (layer.linear_attn.in_proj_z, layer.linear_attn.out_proj) if layer.layer_type == LINEAR else (layer.self_attn.o_proj,)
                for proj in (*mixer, layer.mlp.gate_proj, layer.mlp.up_proj, layer.mlp.down_proj):
                    store_transposed(proj)
            for proj in (*view.delta_q.values(), *view.delta_k.values(), *view.delta_v.values(), *view.attn_q.values()):
                store_transposed(proj)
        # On CUDA, projections that read the same input run as one GEMM, which saves launches and reads the input once. In fp8, in_proj_b and in_proj_a stay bf16, so a delta layer's group falls back to merging just that pair.
        self.merged_projections: dict[tuple[nn.Module, ...], nn.Module] = {}
        if self.cuda:
            groups = [(layer.linear_attn.in_proj_qkv, layer.linear_attn.in_proj_b, layer.linear_attn.in_proj_a) if layer.layer_type == LINEAR
                      else (layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj) for layer in base.model.layers]
            groups += [(view.delta_q[key], view.delta_k[key], view.delta_v[key]) for key in view.delta_q]
            groups += [(view.attn_q[key], view.attn_k[key], view.attn_v[key]) for key in view.attn_q]
            for group in groups:
                for candidate in (group, group[1:]):
                    if (merged := merge_linears_in_place(candidate)) is not None:
                        self.merged_projections[candidate] = merged
                        break
            freeze_rms_norms(base)
            if fp8:
                from inference.fp8 import warm
                warm(nn.ModuleList(self.merged_projections.values()))
        self.log_decay_rates = {i: -layer.linear_attn.A_log.float().exp() for i, layer in enumerate(base.model.layers) if layer.layer_type == LINEAR}
        torch.accelerator.empty_cache()
        self.base, self.head, self.encoder, self.view = base, head, encoder, view
        self.cfg = cfg = base.cfg
        self.K = K = block
        self.J = K - 1
        self.M = K * self.J
        self.R, self.L = max_rows, max_len
        self.trash = max_len + self.M
        slots = max_len + self.M + 1
        if sync_every < 1:
            raise ValueError(f"sync_every must be at least 1, got {sync_every}")
        # A host sync leaves the GPU idle for a short gap in every cycle, while syncing every n cycles runs up to n - 1 full cycles after the last row finishes, so syncing every cycle wins unless a decode runs more cycles than about twice a cycle's time divided by a sync gap.
        self.window_step, self.sync_every = window_step, sync_every
        self.eos = encoder.think_end_id
        self.pad = encoder.pad_id
        self.opt_end = encoder.opt_end_id
        self.empty_think = list(encoder.empty_think)
        self.bias = torch.zeros(cfg.vocab_size, device=dev, dtype=self.dtype)
        self.bias[torch.tensor(encoder.banned_ids, device=dev)] = -1e4
        self.mask_embed = view.mask_embed.detach().to(self.dtype)
        self.layers = base.model.layers
        self.k, self.v, self.dk, self.dv, self.conv, self.rec = {}, {}, {}, {}, {}, {}
        delta = [layer.linear_attn for layer in self.layers if layer.layer_type == LINEAR]
        conv_dim, conv_kernel = delta[0].conv_dim, delta[0].conv_kernel
        # Every delta layer's conv state is a view of one buffer, and so is the conv input of a decode cycle, so that one gather commits them all.
        self.conv_states = torch.zeros(len(delta), max_rows, conv_dim, conv_kernel, device=dev, dtype=self.dtype)
        self.conv_inputs: dict[int, torch.Tensor] = {}
        self.conv_slot: dict[int, int] = {}
        for i, layer in enumerate(self.layers):
            if layer.layer_type == LINEAR:
                lin = layer.linear_attn
                self.dk[i] = torch.zeros(max_rows, slots, lin.num_v_heads, lin.head_v_dim, device=dev, dtype=self.dtype)
                self.dv[i] = torch.zeros_like(self.dk[i])
                self.conv_slot[i] = len(self.conv_slot)
                self.conv[i] = self.conv_states[self.conv_slot[i]]
                self.rec[i] = torch.zeros(max_rows, lin.num_v_heads, lin.head_k_dim, lin.head_v_dim, device=dev, dtype=torch.float32)
            else:
                att = layer.self_attn
                self.k[i] = torch.zeros(max_rows, slots, att.num_kv_heads, att.head_dim, device=dev, dtype=self.dtype)
                self.v[i] = torch.zeros_like(self.k[i])
        self.n = torch.zeros(max_rows, dtype=torch.long, device=dev)
        self.n_out = torch.zeros(max_rows, dtype=torch.long, device=dev)
        self.cap = torch.zeros(max_rows, dtype=torch.long, device=dev)
        self.done = torch.ones(max_rows, dtype=torch.bool, device=dev)
        self.cand = torch.zeros(max_rows, K, dtype=torch.long, device=dev)
        self.out = torch.zeros(max_rows, max_len + K, dtype=torch.long, device=dev)
        self.ar_l = torch.arange(max_len + 1, device=dev)
        self.ar_conv_kernel = torch.arange(cfg.linear_conv_kernel_dim, device=dev)
        self.block = decode_block(K, dev)
        # Tuned on an M4 Pro with torch 2.14: in FP8, a cycle of MANY_QUESTIONS or more is compute-bound, so block 2's 4 rows per question
        # beat block 4's 16 rows even though block 2 accepts fewer tokens. In bf16 they do not, because MPS's bf16 matmul is slow at 9 to 15 rows.
        self.many_questions_block = decode_block(2, dev) if self.mps and fp8 and K > 2 else self.block
        for block_K in {self.block.K, self.many_questions_block.K}:
            self.conv_inputs[block_K] = torch.empty(len(delta), max_rows, conv_dim, conv_kernel + block_K, device=dev, dtype=self.dtype)
        self.cycle_graphs: dict[tuple[int, int, int], torch.cuda.CUDAGraph] = {}
        self.extend_graphs: dict[tuple[int, int, int, bool], torch.cuda.CUDAGraph] = {}
        self.pool = torch.cuda.graph_pool_handle() if self.cuda else None
        if self.cuda:
            self.capture_stream = torch.cuda.Stream()
            longest = min(EXTEND_GRAPH_LENGTHS[-1], max_len)
            self.extend_ids = torch.full((max_rows, longest), self.pad, dtype=torch.long, device=dev)
            self.extend_valid = torch.zeros(max_rows, longest, dtype=torch.bool, device=dev)
            self.extend_lens = torch.zeros(max_rows, dtype=torch.long, device=dev)
            self.extend_starts = torch.zeros(max_rows, dtype=torch.long, device=dev)
            self.extend_out = torch.empty(max_rows, longest, cfg.hidden_size, dtype=self.dtype, device=dev)

    @property
    def metal(self) -> bool:
        return self.mps and metal_enabled()

    def delta_extend(self, q, k, v, g, beta, state, commit: bool) -> torch.Tensor:
        if self.metal:
            return metal.gated_delta_rule_outputs(q, k, v, g, beta, state, advance_state=commit)
        o, final = gated_delta_rule_chunk(q, k, v, g, beta, initial_state=state, output_final_state=True)
        if commit:
            state.copy_(final)
        return o

    def delta_outputs(self, q, k, v, g, beta, state) -> torch.Tensor:
        if self.metal:
            return metal.gated_delta_rule_outputs(q, k, v, g, beta, state)
        return gated_delta_rule_step(q, k, v, g, beta, initial_state=state, output_final_state=False)[0]

    def delta_commit(self, k, v, g, beta, state, accepted) -> None:
        if self.metal:
            metal.gated_delta_rule_advance_inplace(k, v, g, beta, state, accepted)
        else:
            gated_delta_rule_advance_inplace(k, v, g, beta, state, accepted)

    def attention_mask(self, allowed: torch.Tensor) -> torch.Tensor:
        return additive_mask(allowed, self.dtype) if self.cuda else allowed

    def cached_attention(self, q, k_cache, v_cache, mask, scale: float, used_end: torch.Tensor | None, tail_start: int) -> torch.Tensor:
        B, gqa = q.shape[0], q.shape[2] != k_cache.shape[2]
        if self.metal:
            # Tuned on an M4 Pro with torch 2.14: the chunked kernel wins from 2 rows for the delta layers' mask queries, and from 4 rows for attention.
            if B >= (4 if gqa else 2):
                return metal.chunked_cached_attention(q, k_cache, v_cache, mask[:, 0], scale, used_end, tail_start)
            return metal.cached_attention(q, k_cache, v_cache, mask[:, 0], scale, used_end, tail_start)
        return decode_attention(q, k_cache, v_cache, mask, scale)

    def project(self, x: torch.Tensor, *linears: nn.Module) -> tuple[torch.Tensor, ...]:
        if not self.merged_projections:
            return tuple(linear(x) for linear in linears)
        outputs, start = [], 0
        while start < len(linears):
            end = next((end for end in range(len(linears), start + 1, -1) if linears[start:end] in self.merged_projections), start + 1)
            if end == start + 1:
                outputs.append(linears[start](x))
            else:
                merged = self.merged_projections[linears[start:end]]
                outputs += merged(x).split([linear.out_features for linear in linears[start:end]], dim=-1)
            start = end
        return tuple(outputs)

    def split_candidate_rows(self, h: torch.Tensor, K: int) -> tuple[torch.Tensor, torch.Tensor]:
        hc, hm = h[:, :K], h[:, K:]
        # With more than one row both halves are strided: CUDA then runs a batched GEMM that reads the weights once per row, and MPS on torch 2.14 multiplies a strided input by the untransposed candidate weights about 2x slower.
        if self.cuda:
            return hc.contiguous(), hm.contiguous()
        return (hc.contiguous() if self.mps else hc), hm

    def vocab_logits(self, h: torch.Tensor) -> torch.Tensor:
        if not self.mps or not isinstance(self.base.lm_head, nn.Linear):
            return self.base.lm_head(h)
        # MPS in torch 2.14 picks a slow kernel for the 248k-row matmul; row chunks of the weight give bitwise-equal logits much faster.
        return torch.cat([F.linear(h, w) for w in self.base.lm_head.weight.chunk(8)], dim=-1)

    def decode_rows(self, n: int) -> int:
        if not self.cuda:
            return n
        # CUDA captures a graph per row count, so rounding up to a power of two keeps the number of graphs small.
        b = 1
        while b < n:
            b *= 2
        return min(b, self.R)

    @torch.no_grad()
    def extend(self, seqs: list[list[int]], starts: list[int], commit: bool) -> torch.Tensor:
        B, T = len(seqs), max(len(s) for s in seqs)
        Lk = max(st + len(s) for st, s in zip(starts, seqs))
        padded_T = next((n for n in EXTEND_GRAPH_LENGTHS if T <= n <= self.L), None) if self.cuda else None
        graph_rows = self.decode_rows(B) if padded_T is not None and padded_T <= ROUNDED_ROWS_MAX_LENGTH else B
        width = T if padded_T is None else padded_T
        ids = torch.full((B, width), self.pad, dtype=torch.long)
        valid = torch.zeros(B, width, dtype=torch.bool)
        for b, s in enumerate(seqs):
            ids[b, :len(s)] = torch.tensor(s, dtype=torch.long)
            valid[b, :len(s)] = True
        lens, start_positions = torch.tensor([len(s) for s in seqs]), torch.tensor(starts)
        if padded_T is None:
            dev = self.device
            return self.extend_pass(ids.to(dev), valid.to(dev), lens.to(dev), start_positions.to(dev), Lk, commit)
        padded_Lk = padded_T if max(starts) == 0 else min(max(self.window_step, 1 << (Lk - 1).bit_length()), self.L)
        self.extend_ids[:graph_rows, :padded_T].fill_(self.pad)
        self.extend_valid[:graph_rows, :padded_T].zero_()
        self.extend_lens[:graph_rows].zero_()
        self.extend_starts[:graph_rows].zero_()
        self.extend_ids[:B, :padded_T].copy_(ids)
        self.extend_valid[:B, :padded_T].copy_(valid)
        self.extend_lens[:B].copy_(lens)
        self.extend_starts[:B].copy_(start_positions)
        self.extend_graph(graph_rows, padded_T, padded_Lk, commit).replay()
        # A view of the shared output buffer: the next extend overwrites it.
        return self.extend_out[:B, :T]

    def extend_graph(self, B: int, T: int, Lk: int, commit: bool) -> torch.cuda.CUDAGraph:
        graph = self.extend_graphs.get((B, T, Lk, commit))
        if graph is not None:
            return graph
        inputs = (self.extend_ids[:B, :T], self.extend_valid[:B, :T], self.extend_lens[:B], self.extend_starts[:B], Lk, commit)
        # A committing warm-up advances conv and rec, so they are restored before the caller's replay.
        states = [self.conv_states[:, :B], *(t[:B] for t in self.rec.values())] if commit else []
        saved = [t.clone() for t in states]
        graph = self.captured(lambda: self.extend_out[:B, :T].copy_(self.extend_pass(*inputs)), warmups=1)
        for t, s in zip(states, saved):
            t.copy_(s)
        self.extend_graphs[(B, T, Lk, commit)] = graph
        return graph

    def captured(self, run: Callable[[], object], warmups: int) -> torch.cuda.CUDAGraph:
        self.capture_stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(self.capture_stream):
            for _ in range(warmups):
                run()
        torch.cuda.current_stream().wait_stream(self.capture_stream)
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph, pool=self.pool, stream=self.capture_stream):
            run()
        return graph

    def extend_pass(self, ids: torch.Tensor, valid: torch.Tensor, lens: torch.Tensor, starts: torch.Tensor, Lk: int, commit: bool) -> torch.Tensor:
        dev, base = self.device, self.base
        B, T = ids.shape
        pos = starts[:, None] + torch.arange(T, device=dev)[None]
        wpos = torch.where(valid, pos, self.trash)
        rows = torch.arange(B, device=dev)[:, None].expand(B, T)
        mask = self.attention_mask((self.ar_l[:Lk][None, None, :] <= pos[:, :, None])[:, None])
        cos, sin = base.model.rotary_emb(pos)
        conv_columns = (lens[:, None] + self.ar_conv_kernel[None])[:, None, :]
        x = base.model.embed_tokens(ids)
        for i, layer in enumerate(self.layers):
            h = layer.input_layernorm(x)
            if layer.layer_type == LINEAR:
                lin = layer.linear_attn
                qkv, b_in, a_in = self.project(h, lin.in_proj_qkv, lin.in_proj_b, lin.in_proj_a)
                ext, q, k, v = eager_conv_step(self.conv[i][:B], qkv, lin.conv1d.weight, lin.key_dim, lin.value_dim, lin.head_k_dim, lin.head_v_dim)
                beta, g = delta_gates(b_in, a_in, lin.dt_bias, self.log_decay_rates[i], valid)
                o = self.delta_extend(q, k, v, g, beta, self.rec[i][:B], commit)
                if commit:
                    torch.gather(ext, 2, conv_columns.expand(B, lin.conv_dim, -1), out=self.conv[i][:B])
                self.dk[i][rows, wpos] = apply_rotary(k, cos, sin)
                self.dv[i][rows, wpos] = v
                z = lin.in_proj_z(h).view(B, T, lin.num_v_heads, lin.head_v_dim)
                o = lin.norm(o.reshape(-1, lin.head_v_dim), z.reshape(-1, lin.head_v_dim)).view(B, T, lin.value_dim)
                x = x + lin.out_proj(o)
            else:
                att = layer.self_attn
                H, Hkv, D = att.num_heads, att.num_kv_heads, att.head_dim
                q_in, k_in, v_in = self.project(h, att.q_proj, att.k_proj, att.v_proj)
                q, gate = q_in.view(B, T, H, 2 * D).chunk(2, dim=-1)
                q = apply_rotary(att.q_norm(q), cos, sin)
                self.k[i][rows, wpos] = apply_rotary(att.k_norm(k_in.view(B, T, Hkv, D)), cos, sin)
                self.v[i][rows, wpos] = v_in.view(B, T, Hkv, D)
                keys, vals = self.k[i][:B, :Lk], self.v[i][:B, :Lk]
                with sdpa_kernel(EFFICIENT):
                    # Space patch: expand KV heads instead of enable_gqa (the efficient kernel rejects GQA on this torch build).
                    rep = H // Hkv
                    o = F.scaled_dot_product_attention(q.transpose(1, 2), keys.repeat_interleave(rep, dim=2).transpose(1, 2),
                                                       vals.repeat_interleave(rep, dim=2).transpose(1, 2), attn_mask=mask,
                                                       scale=att.scaling).transpose(1, 2)
                o = o.reshape(B, T, H * D) * torch.sigmoid(gate.reshape(B, T, H * D))
                x = x + att.o_proj(o)
            x = x + layer.mlp(layer.post_attention_layernorm(x))
        return base.model.norm(x)

    def cycle(self, B: int, Lw: int, block: DecodeBlock) -> None:
        base, view = self.base, self.view
        K, J, M = block.K, block.J, block.M
        n, done, cand = self.n[:B], self.done[:B], self.cand[:B, :K]
        pos_c = n[:, None] + block.ar_k[None]
        pos = torch.cat((pos_c, n[:, None] + block.mask_offset[None]), dim=1)
        cos, sin = base.model.rotary_emb(pos)
        cos_c, sin_c, cos_m, sin_m = (t.contiguous() for t in (cos[:, :K], sin[:, :K], cos[:, K:], sin[:, K:]))
        used_end = n + K if self.metal else None
        rows = torch.arange(B, device=self.device)[:, None].expand(B, K)
        limit = n[:, None] + block.limit_offset[None]
        allow_cache = self.ar_l[:Lw][None, None, :] <= limit[:, :, None]
        mask = self.attention_mask(torch.cat((allow_cache, block.allow_mm[None].expand(B, -1, -1)), dim=2)[:, None])
        x = torch.cat((base.model.embed_tokens(cand), self.mask_embed.view(1, 1, -1).expand(B, M, -1)), dim=1)
        pending = {}
        for i, layer in enumerate(self.layers):
            h = layer.input_layernorm(x)
            key = str(i)
            if layer.layer_type == LINEAR:
                lin = layer.linear_attn
                Hv, Dv = lin.num_v_heads, lin.head_v_dim
                hc, hm = self.split_candidate_rows(h, K)
                qkv, b_in, a_in = self.project(hc, lin.in_proj_qkv, lin.in_proj_b, lin.in_proj_a)
                _, q, k, v = conv_step(self.conv[i][:B], qkv, lin.conv1d.weight, key_dim=lin.key_dim, value_dim=lin.value_dim,
                                       head_k_dim=lin.head_k_dim, head_v_dim=Dv, ext_out=self.conv_inputs[K][self.conv_slot[i], :B])
                beta, g = delta_gates(b_in, a_in, lin.dt_bias, self.log_decay_rates[i])
                o_c = self.delta_outputs(q, k, v, g, beta, self.rec[i][:B])
                pending[i] = (k, v, g, beta)
                qm, km, vm = (t.view(B, M, Hv, Dv) for t in self.project(hm, view.delta_q[key], view.delta_k[key], view.delta_v[key]))
                qm = rotate_into_cache(k, v, qm, km, vm, cos_c, sin_c, cos_m, sin_m, k_cache=self.dk[i], v_cache=self.dv[i], rows=rows,
                                       positions=pos_c, tail_start=Lw)
                o_m = self.cached_attention(qm, self.dk[i], self.dv[i], mask[:, :, K:], Dv ** -0.5, used_end, Lw)
                o = torch.cat((o_c, o_m), dim=1)
                z = lin.in_proj_z(h).view(B, K + M, Hv, Dv)
                o = lin.norm(o.reshape(-1, Dv), z.reshape(-1, Dv)).view(B, K + M, lin.value_dim)
                x = x + lin.out_proj(o)
            else:
                att = layer.self_attn
                H, Hkv, D = att.num_heads, att.num_kv_heads, att.head_dim
                hc, hm = self.split_candidate_rows(h, K)
                qc_in, kc_in, vc_in = self.project(hc, att.q_proj, att.k_proj, att.v_proj)
                qm_in, km_in, vm_in = self.project(hm, view.attn_q[key], view.attn_k[key], view.attn_v[key])
                qc, gc = qc_in.view(B, K, H, 2 * D).chunk(2, dim=-1)
                qm, gm = qm_in.view(B, M, H, 2 * D).chunk(2, dim=-1)
                q = apply_rotary(att.q_norm(torch.cat((qc, qm), 1)), cos, sin)
                kc = att.k_norm(kc_in.view(B, K, Hkv, D))
                km = att.k_norm(km_in.view(B, M, Hkv, D))
                k = apply_rotary(torch.cat((kc, km), 1), cos, sin)
                self.k[i][rows, pos_c] = k[:, :K]
                self.v[i][rows, pos_c] = vc_in.view(B, K, Hkv, D)
                self.k[i][:B, Lw:Lw + M] = k[:, K:]
                self.v[i][:B, Lw:Lw + M] = vm_in.view(B, M, Hkv, D)
                o = self.cached_attention(q, self.k[i], self.v[i], mask, att.scaling, used_end, Lw)
                o = o.reshape(B, K + M, H * D) * torch.sigmoid(torch.cat((gc, gm), 1).reshape(B, K + M, H * D))
                x = x + att.o_proj(o)
            x = x + layer.mlp(layer.post_attention_layernorm(x))
        am = (self.vocab_logits(base.model.norm(x)) + self.bias).argmax(-1)
        pred = am[:, :K]
        match = (cand[:, 1:] == pred[:, :-1]).long()
        j = torch.cumprod(match, 1).sum(1) + 1
        first_eos = torch.where(cand == self.eos, block.ar_k[None], K).amin(1)
        j = torch.minimum(j, first_eos + 1)
        j = torch.minimum(j, self.cap[:B] - self.n_out[:B])
        j = torch.where(done, 0, j)
        jm1 = (j - 1).clamp(min=0)
        nxt = torch.cat((pred.gather(1, jm1[:, None]), am.gather(1, K + jm1[:, None] * J + block.ar_j[None])), dim=1)
        self.out[:B].scatter_(1, self.n_out[:B, None] + block.ar_k[None], cand)
        for i, (k, v, g, beta) in pending.items():
            self.delta_commit(k, v, g, beta, self.rec[i][:B], j)
        conv_columns = (j[:, None] + self.ar_conv_kernel[None])[None, :, None, :]
        torch.gather(self.conv_inputs[K][:, :B], 3, conv_columns.expand(*self.conv_states[:, :B].shape), out=self.conv_states[:, :B])
        cand.copy_(torch.where(done[:, None], cand, nxt))
        done.logical_or_((first_eos < j) | (self.n_out[:B] + j >= self.cap[:B]))
        n.add_(j)
        self.n_out[:B].add_(j)

    def cycle_graph(self, B: int, Lw: int, block: DecodeBlock) -> torch.cuda.CUDAGraph:
        graph = self.cycle_graphs.get((B, Lw, block.K))
        if graph is not None:
            return graph
        saved = self.done[:B].clone()
        self.done[:B].fill_(True)
        graph = self.captured(lambda: self.cycle(B, Lw, block), warmups=2)
        self.done[:B].copy_(saved)
        self.cycle_graphs[(B, Lw, block.K)] = graph
        return graph

    def cycle_runner(self, B: int, Lw: int, block: DecodeBlock) -> Callable[[], None]:
        if self.cuda:
            return self.cycle_graph(B, Lw, block).replay
        return lambda: self.cycle(B, Lw, block)

    @torch.no_grad()
    def decode(self, B: int, prompts: list[list[int]], first: torch.Tensor, caps: list[int]) -> tuple[list[list[int]], list[bool]]:
        Bp = self.decode_rows(B)
        dev = self.device
        self.n[:Bp].zero_()
        self.n[:B].copy_(torch.tensor([len(p) for p in prompts], device=dev))
        self.n_out[:Bp].zero_()
        self.cap[:Bp].zero_()
        self.cap[:B].copy_(torch.tensor(caps, device=dev))
        self.done[:Bp].fill_(True)
        self.done[:B].copy_(self.cap[:B] == 0)
        self.cand[:Bp].zero_()
        self.cand[:B, 0].copy_(first)
        top = max(len(p) for p in prompts)
        block = self.many_questions_block if B >= MANY_QUESTIONS else self.block
        while True:
            need = top + (self.sync_every + 1) * block.K + 1
            run_cycle = self.cycle_runner(Bp, min(self.L, -(-need // self.window_step) * self.window_step), block)
            for _ in range(self.sync_every):
                run_cycle()
            finished, top = torch.stack((self.done[:Bp].all().long(), self.n[:B].max())).tolist()
            if finished:
                break
        n_out = self.n_out[:B].tolist()
        out = self.out[:B].tolist()
        chains = [out[b][:n_out[b]] for b in range(B)]
        return chains, [bool(c) and c[-1] == self.eos for c in chains]

    def readout(self, h: torch.Tensor, seq: list[int]) -> list[float]:
        opts = [t for t, tok in enumerate(seq) if tok == self.opt_end]
        logits = self.head(h[opts][None], h[len(seq) - 1][None]).float()
        return F.softmax(logits[0], dim=-1).tolist()

    def reset(self, B: int) -> None:
        self.conv_states[:, :B].zero_()
        for t in self.rec.values():
            t[:B].zero_()

    def broadcast(self, B: int, P: int) -> None:
        if B == 1:
            return
        for d in (self.k, self.v, self.dk, self.dv):
            for t in d.values():
                t[1:B, :P].copy_(t[:1, :P].expand(B - 1, -1, -1, -1))
        self.conv_states[:, 1:B].copy_(self.conv_states[:, :1].expand(-1, B - 1, -1, -1))
        for t in self.rec.values():
            t[1:B].copy_(t[:1].expand(B - 1, *t.shape[1:]))

    @torch.no_grad()
    def answer(self, record: DataFormat, opts: Options) -> list[Result]:
        return self.answer_batch([(record, opts)])[0]

    @torch.no_grad()
    def answer_batch(self, jobs: list[tuple[DataFormat, Options]]) -> list[list[Result]]:
        results: list[list[Result | None]] = [[None] * len(record.questions) for record, _ in jobs]
        # A row that cannot think would hold a decode row in a thinking group. group() shares only the prompt prefix common to all its rows, so a request's text is shared only in a group that holds no other request.
        for thinking in (False, True):
            requests = [job for job, (_, opts) in enumerate(jobs) if (opts.think and opts.max_think > 0) == thinking]
            for group in pack_requests([len(jobs[job][0].questions) for job in requests], self.R):
                rows = [GroupRow(jobs[requests[r]][0], jobs[requests[r]][0].questions[i], jobs[requests[r]][1]) for r, i in group]
                for (r, i), result in zip(group, self.group(rows)):
                    results[requests[r]][i] = result
        return results

    def group(self, rows: list[GroupRow]) -> list[Result]:
        enc = self.encoder
        prompts = [enc.encode(enc.prompt_text(row.record, row.question)) for row in rows]
        rems = [enc.encode(enc.remainder_text(row.question)) for row in rows]
        B = len(rows)
        slack = self.K + self.J + 2
        for p, r in zip(prompts, rems):
            if len(p) + len(r) + len(self.empty_think) + slack > self.L:
                raise ValueError(f"input of {len(p) + len(r)} tokens exceeds the {self.L}-token limit")
        budgets = [row.options.max_think if row.options.think else 0 for row in rows]
        thresholds = [row.options.nothink_threshold for row in rows]
        caps = [max(0, min(budget, self.L - len(p) - len(r) - slack)) for budget, p, r in zip(budgets, prompts, rems)]
        wants_nothink = [t is not None or not c for t, c in zip(thresholds, caps)]
        P = common_prefix(prompts) if B > 1 else 0
        self.reset(self.decode_rows(B))
        if P > 0:
            self.extend([prompts[0][:P]], [0], commit=True)
            self.broadcast(B, P)
        tails = [self.empty_think + r for r in rems]
        first = nothink = None
        if not any(caps):
            hn = self.extend([p[P:] + t for p, t in zip(prompts, tails)], [P] * B, commit=False)
            nothink = [self.readout(hn[b, len(p) - P:], t) for b, (p, t) in enumerate(zip(prompts, tails))]
        else:
            h = self.extend([p[P:] for p in prompts], [P] * B, commit=True)
            last = h[torch.arange(B, device=self.device), torch.tensor([len(p) - P - 1 for p in prompts], device=self.device)]
            first = (self.vocab_logits(last) + self.bias).argmax(-1)
            if any(t is not None for t in thresholds) or not all(caps):
                hn = self.extend(tails, [len(p) for p in prompts], commit=False)
                nothink = [self.readout(hn[b], tails[b]) for b in range(B)]
                caps = [c if t is None or max(nothink[b]) < t else 0 for b, (c, t) in enumerate(zip(caps, thresholds))]
        chains, closed = [[] for _ in range(B)], [False] * B
        probs = list(nothink) if nothink is not None else [[] for _ in range(B)]
        if any(caps):
            chains, closed = self.decode(B, prompts, first, caps)
            seqs = [r[1:] if c else r for r, c in zip(rems, closed)]
            hr = self.extend(seqs, [len(p) + len(ch) for p, ch in zip(prompts, chains)], commit=False)
            for b in range(B):
                if caps[b]:
                    probs[b] = self.readout(hr[b], seqs[b])
        return [Result(probs=probs[b], thought=bool(caps[b]), chain=chains[b], closed=closed[b],
                       nothink_probs=nothink[b] if nothink is not None and wants_nothink[b] else None,
                       prompt_tokens=len(prompts[b]) - (P if b else 0))
                for b in range(B)]