Blame for code/offload/dspark/d2/shards.py

"""D2 shards (schema d2.shard.v1, DATASET.md): writer, reader, integrity checks and a V1 CaptureSet adapter.

Proprietary content (legal-v2a token ids, taps, target probabilities): shards live only on x86host
/opt/dsv41/x86host-data/d2/ (mode 700). Nothing here prints token ids or text.

  Shard(path)            reader: tensors (taps lazily, by slice), sidecar metadata, blocks() per DATASET.md section 4
  write_shard(...)       writer used by dataset.py build and dataset.py synth
  check_record(rec)      DATASET.md section 8 rules that one sequence can be checked against on its own
  ShardCaptureSet(...)   V1 verify.captures.CaptureSet interface over shards (rows, samples, plumbing, window,
                         window_status, arm_inputs), so V1 run_ref / analyze run on D2 data unchanged
"""

from __future__ import annotations

import collections.abc
import dataclasses
import hashlib
import json
import os
from typing import Iterator, Sequence

import torch
from safetensors import safe_open
from safetensors.torch import save_file

SCHEMA = "d2.shard.v1"
MANIFEST_SCHEMA = "d2.manifest.v1"
W = 128                      # drafter window (native config sliding_window)
K = 5                        # dspark_block_size
NOISE = 128799               # dspark_noise_token_id
TOPK = 20
H = 5120
TAP_LAYERS = (37, 38, 39)
TAPW = 3 * H
VOCAB = 129280
SEQ_FIELDS = ("tok_off", "L", "gen_start", "r0", "row_off", "n_rows", "span_off", "n_span")
SPLITS = ("train", "dev", "heldout_legal", "heldout_general", "acap")
TAPS_SOURCES = ("prefill_replay", "decode_capture", "synthetic")
MAX_SHARD_TAP_BYTES = 2 * 1024 ** 3


def tokens_sha256(tokens: Sequence[int] | torch.Tensor) -> str:
    """sha256 of the token ids as int32 little-endian (the identity used across prompt, generation and shards)."""
    t = torch.as_tensor(tokens, dtype=torch.int32) if not isinstance(tokens, torch.Tensor) else tokens.to(torch.int32)
    return hashlib.sha256(t.contiguous().numpy().astype("<i4", copy=False).tobytes()).hexdigest()


def sha256_file(path: str) -> str:
    with open(path, "rb") as fh:
        return hashlib.file_digest(fh, "sha256").hexdigest()


def r0_of(gen_start: int) -> int:
    return max(0, gen_start - W)


@dataclasses.dataclass
class SeqRecord:
    """One sequence ready to be written (DATASET.md section 3)."""
    meta: dict                 # seq_id, split, family, source, source_docs, prompt_sha256, finish_reason, segments, ...
    tokens: torch.Tensor       # int32 [L]
    gen_start: int
    taps: torch.Tensor         # bf16 [L - r0, TAPW]
    top_ids: torch.Tensor      # int32 [L - gen_start, TOPK]
    top_lp: torch.Tensor       # float32 [L - gen_start, TOPK]
    tok_lp: torch.Tensor       # float32 [L - gen_start]
    tok_rank: torch.Tensor     # int32 [L - gen_start]
    resid: torch.Tensor        # float32 [L - gen_start]

    @property
    def L(self) -> int:
        return int(self.tokens.shape[0])

    @property
    def r0(self) -> int:
        return r0_of(self.gen_start)


def check_record(rec: SeqRecord) -> list[str]:
    """DATASET.md section 8 rules 3 and 5 for one sequence; returns the violations (empty = clean)."""
    err: list[str] = []
    sid = rec.meta.get("seq_id", "?")
    L, g, r0 = rec.L, rec.gen_start, rec.r0
    if not 1 <= g <= L - 1:
        return [f"{sid}: gen_start {g} outside 1..L-1 (L={L})"]
    n_span = L - g
    shapes = {"taps": (L - r0, TAPW), "top_ids": (n_span, TOPK), "top_lp": (n_span, TOPK), "tok_lp": (n_span,),
              "tok_rank": (n_span,), "resid": (n_span,)}
    dtypes = {"taps": torch.bfloat16, "top_ids": torch.int32, "top_lp": torch.float32, "tok_lp": torch.float32,
              "tok_rank": torch.int32, "resid": torch.float32}
    for name, shape in shapes.items():
        t = getattr(rec, name)
        if tuple(t.shape) != shape or t.dtype != dtypes[name]:
            err.append(f"{sid}: {name} {tuple(t.shape)} {t.dtype}, want {shape} {dtypes[name]}")
    if rec.tokens.dtype != torch.int32 or rec.tokens.dim() != 1:
        err.append(f"{sid}: tokens {rec.tokens.dtype} dim {rec.tokens.dim()}, want int32 [L]")
    if err:
        return err
    if not bool(torch.isfinite(rec.taps).all()):
        err.append(f"{sid}: non-finite taps")
    zero = (rec.taps.view(-1, 3, TAPW // 3) == 0).all(-1).any(-1)
    if bool(zero.any()):
        err.append(f"{sid}: {int(zero.sum())} taps rows with a layer all zero (first at position "
                   f"{rec.r0 + int(zero.nonzero()[0])}; trimmed eager-prefill rows?)")
    if int(rec.tokens.min()) < 0 or int(rec.tokens.max()) >= VOCAB:
        err.append(f"{sid}: token id outside 0..{VOCAB - 1}")
    lp = rec.top_lp
    if not bool(torch.isfinite(lp).all()) or bool((lp > 1e-6).any()):
        err.append(f"{sid}: top_lp not finite or above 0")
    if bool((lp[:, 1:] > lp[:, :-1] + 1e-6).any()):
        err.append(f"{sid}: top_lp not non-increasing")
    if bool(((rec.top_ids < 0) | (rec.top_ids >= VOCAB)).any()):
        err.append(f"{sid}: top_ids outside the vocabulary")
    srt = rec.top_ids.sort(dim=-1).values
    if bool((srt[:, 1:] == srt[:, :-1]).any()):
        err.append(f"{sid}: repeated id within a top-20 row")
    mass = lp.double().exp().sum(-1)
    if bool((mass > 1 + 1e-4).any()):
        err.append(f"{sid}: top-20 mass above 1 (max {float(mass.max()):.6f})")
    want_resid = (1 - mass).clamp(0, 1).float()
    if not torch.allclose(rec.resid, want_resid, atol=1e-5):
        err.append(f"{sid}: resid != max(0, 1 - sum exp(top_lp))")
    span_tok = rec.tokens[g:].long()
    if not bool(torch.isfinite(rec.tok_lp).all()) or bool((rec.tok_lp > 1e-6).any()):
        err.append(f"{sid}: tok_lp not finite or above 0")
    hit = rec.top_ids.long() == span_tok[:, None]
    in_top = hit.any(-1)
    if not torch.equal(in_top, rec.tok_rank <= TOPK):
        err.append(f"{sid}: tok_rank <= {TOPK} disagrees with the committed token's presence in top_ids")
    if bool((rec.tok_rank < 1).any()):
        err.append(f"{sid}: tok_rank below 1")
    if bool(in_top.any()):
        at = hit.float().argmax(-1)
        idx = in_top.nonzero().flatten()
        lp_hit = lp[idx, at[idx]]
        if not torch.allclose(lp_hit, rec.tok_lp[idx], atol=1e-5):
            err.append(f"{sid}: tok_lp != top_lp at the committed token")
        rk_hit = at[idx].to(torch.int32) + 1
        # ties: a tied neighbour may sit before the committed token in the row; the rank must still point at a tie
        bad = (rk_hit != rec.tok_rank[idx]) & (rec.tok_rank[idx] <= TOPK)   # rank > TOPK here is reported above
        if bool(bad.any()):
            j = idx[bad]
            if not torch.allclose(lp[j, rec.tok_rank[j].long() - 1], rec.tok_lp[j], atol=1e-5):
                err.append(f"{sid}: tok_rank != position of the committed token in top_ids")
    out_top = ~in_top
    if bool(out_top.any()) and bool((rec.tok_lp[out_top] > lp[out_top, -1] + 1e-5).any()):
        err.append(f"{sid}: a committed token outside the top-20 has a logprob above the 20th")
    return err


def write_shard(base: str, split: str, records: list[SeqRecord], taps_source: str,
                extra_meta: dict | None = None) -> dict:
    """Write <base>.safetensors + <base>.json (DATASET.md section 6); returns the manifest file entry."""
    if split not in SPLITS:
        raise ValueError(f"split {split!r} not in {SPLITS}")
    if taps_source not in TAPS_SOURCES:
        raise ValueError(f"taps_source {taps_source!r} not in {TAPS_SOURCES}")
    if not records:
        raise ValueError("empty shard")
    errs = [e for r in records for e in check_record(r)]
    if errs:
        raise ValueError(f"{len(errs)} integrity errors, first: {errs[0]}")
    for r in records:
        if r.meta.get("split") != split:
            raise ValueError(f"{r.meta.get('seq_id')}: split {r.meta.get('split')!r} written into a {split} shard")
    idx, metas = [], []
    tok_off = row_off = span_off = 0
    for i, r in enumerate(records):
        L, g, r0 = r.L, r.gen_start, r.r0
        idx.append([tok_off, L, g, r0, row_off, L - r0, span_off, L - g])
        m = dict(r.meta)
        m.update(index=i, L=L, gen_start=g, n_gen=L - g, tokens_sha256=tokens_sha256(r.tokens))
        metas.append(m)
        tok_off += L
        row_off += L - r0
        span_off += L - g
    if row_off * TAPW * 2 > MAX_SHARD_TAP_BYTES:
        raise ValueError(f"shard taps {row_off * TAPW * 2} bytes above {MAX_SHARD_TAP_BYTES}")
    tensors = {
        "seq_index": torch.tensor(idx, dtype=torch.int64),
        "tokens": torch.cat([r.tokens for r in records]).contiguous(),
        "taps": torch.cat([r.taps for r in records]).contiguous(),
        "row_pos": torch.cat([torch.arange(r.r0, r.L, dtype=torch.int32) for r in records]),
        "top_ids": torch.cat([r.top_ids for r in records]).contiguous(),
        "top_lp": torch.cat([r.top_lp for r in records]).contiguous(),
        "tok_lp": torch.cat([r.tok_lp for r in records]).contiguous(),
        "tok_rank": torch.cat([r.tok_rank for r in records]).contiguous(),
        "resid": torch.cat([r.resid for r in records]).contiguous(),
    }
    path = base + ".safetensors"
    tmp = path + ".tmp"
    save_file(tensors, tmp, metadata={"schema": SCHEMA, "split": split})
    os.chmod(tmp, 0o600)
    os.replace(tmp, path)
    digest = sha256_file(path)
    side = {"schema": SCHEMA, "split": split, "file": os.path.basename(path), "sha256": digest, "W": W, "K": K,
            "topk": TOPK, "tap_layers": list(TAP_LAYERS), "hidden": H, "noise_token": NOISE,
            "taps_source": taps_source, **(extra_meta or {}), "sequences": metas}
    side_path = base + ".json"
    with open(side_path + ".tmp", "w") as fh:
        json.dump(side, fh, indent=1)
    os.chmod(side_path + ".tmp", 0o600)
    os.replace(side_path + ".tmp", side_path)
    return {"file": path, "sidecar": side_path, "sha256": digest, "sidecar_sha256": sha256_file(side_path),
            "bytes": os.path.getsize(path), "n_seq": len(records), "n_rows": row_off, "n_span": span_off}


class Shard:
    """Read one shard. Small tensors are loaded; taps stay in the file and are read by slice."""

    SMALL = ("seq_index", "tokens", "row_pos", "top_ids", "top_lp", "tok_lp", "tok_rank", "resid")

    def __init__(self, path: str, verify_sha256: bool = False):
        if not path.endswith(".safetensors"):
            raise ValueError(f"{path}: not a .safetensors shard")
        self.path = path
        with open(path[: -len(".safetensors")] + ".json") as fh:
            self.side = json.load(fh)
        if self.side.get("schema") != SCHEMA:
            raise ValueError(f"{path}: sidecar schema {self.side.get('schema')!r}, want {SCHEMA}")
        if verify_sha256 and sha256_file(path) != self.side["sha256"]:
            raise ValueError(f"{path}: sha256 differs from its sidecar")
        self._f = safe_open(path, framework="pt", device="cpu")
        md = self._f.metadata() or {}
        if md.get("schema") != SCHEMA or md.get("split") != self.side["split"]:
            raise ValueError(f"{path}: header metadata {md} disagrees with the sidecar")
        self.split = self.side["split"]
        self.t = {n: self._f.get_tensor(n) for n in self.SMALL}
        self.meta = self.side["sequences"]
        if self.t["seq_index"].shape != (len(self.meta), len(SEQ_FIELDS)):
            raise ValueError(f"{path}: seq_index {tuple(self.t['seq_index'].shape)} vs {len(self.meta)} sequences")
        self._taps = self._f.get_slice("taps")
        if list(self._taps.get_shape()) != [int(self.t["row_pos"].shape[0]), TAPW]:
            raise ValueError(f"{path}: taps shape {self._taps.get_shape()}")

    def __len__(self) -> int:
        return len(self.meta)

    def seq(self, i: int) -> dict:
        """seq_index fields of sequence i (ints) merged over its sidecar metadata."""
        v = self.t["seq_index"][i].tolist()
        return {**self.meta[i], **dict(zip(SEQ_FIELDS, v))}

    def tokens(self, i: int) -> torch.Tensor:
        s = self.seq(i)
        return self.t["tokens"][s["tok_off"]: s["tok_off"] + s["L"]]

    def taps(self, i: int, p0: int, p1: int) -> torch.Tensor:
        """bf16 [p1 - p0, TAPW] taps rows of positions p0 .. p1-1 of sequence i."""
        s = self.seq(i)
        if not s["r0"] <= p0 <= p1 <= s["L"]:
            raise IndexError(f"rows {p0}..{p1} outside {s['r0']}..{s['L']} of sequence {i}")
        a = s["row_off"] + p0 - s["r0"]
        return self._taps[a: a + (p1 - p0)]

    def span(self, i: int) -> dict[str, torch.Tensor]:
        """Target arrays over sequence i's span (row j = position gen_start + j)."""
        s = self.seq(i)
        a, n = s["span_off"], s["n_span"]
        return {k: self.t[k][a: a + n] for k in ("top_ids", "top_lp", "tok_lp", "tok_rank", "resid")}

    def blocks(self, i: int) -> range:
        """Every anchor-1 position P of sequence i's training blocks (DATASET.md section 4)."""
        s = self.seq(i)
        return range(s["gen_start"] - 1, s["L"] - 2)

    def block(self, i: int, P: int) -> dict:
        """One training block: window rows, anchor, and per draft position its target or None (partial block)."""
        s = self.seq(i)
        if not s["gen_start"] - 1 <= P <= s["L"] - 3:
            raise IndexError(f"P {P} outside the block range of sequence {i}")
        tok = self.tokens(i)
        lo = max(0, P - W + 1)
        sp = self.span(i)
        tgt = []
        for d in range(K):
            t = P + 2 + d
            if t > s["L"] - 1:
                tgt.append(None)
                continue
            j = t - s["gen_start"]
            tgt.append({"pos": t, "token": int(tok[t]), "prev": int(tok[t - 1]), "top_ids": sp["top_ids"][j],
                        "top_lp": sp["top_lp"][j], "resid": float(sp["resid"][j])})
        return {"P": P, "window_pos": list(range(lo, P + 1)), "window_taps": self.taps(i, lo, P + 1),
                "anchor": int(tok[P + 1]), "targets": tgt}

    def structure_errors(self) -> list[str]:
        """Cross-checks of the stored offsets (DATASET.md sections 3, 6.1); empty = clean."""
        err: list[str] = []
        tok_off = row_off = span_off = 0
        rp = self.t["row_pos"]
        for i in range(len(self)):
            s = self.seq(i)
            want = (tok_off, row_off, span_off)
            got = (s["tok_off"], s["row_off"], s["span_off"])
            if got != want:
                err.append(f"seq {i}: offsets {got}, want {want}")
            if s["r0"] != r0_of(s["gen_start"]) or s["n_rows"] != s["L"] - s["r0"] or \
                    s["n_span"] != s["L"] - s["gen_start"]:
                err.append(f"seq {i}: r0/n_rows/n_span inconsistent with L {s['L']} gen_start {s['gen_start']}")
            if not torch.equal(rp[s["row_off"]: s["row_off"] + s["n_rows"]],
                               torch.arange(s["r0"], s["L"], dtype=torch.int32)):
                err.append(f"seq {i}: row_pos is not r0..L-1")
            if tokens_sha256(self.tokens(i)) != s.get("tokens_sha256"):
                err.append(f"seq {i}: tokens_sha256 mismatch")
            tok_off += s["L"]
            row_off += s["n_rows"]
            span_off += s["n_span"]
        if (tok_off, row_off, span_off) != (self.t["tokens"].shape[0], rp.shape[0], self.t["tok_lp"].shape[0]):
            err.append("totals of seq_index differ from the tensor lengths")
        return err

    def record(self, i: int) -> SeqRecord:
        """Sequence i as a SeqRecord (taps read in full; used by re-checks and tests)."""
        s = self.seq(i)
        sp = self.span(i)
        meta = {k: v for k, v in self.meta[i].items() if k not in ("index", "L", "gen_start", "n_gen", "tokens_sha256")}
        return SeqRecord(meta, self.tokens(i).clone(), s["gen_start"], self.taps(i, s["r0"], s["L"]).clone(),
                         sp["top_ids"].clone(), sp["top_lp"].clone(), sp["tok_lp"].clone(), sp["tok_rank"].clone(),
                         sp["resid"].clone())


def shard_paths(root: str, split: str) -> list[str]:
    d = os.path.join(root, split)
    if not os.path.isdir(d):
        return []
    return sorted(os.path.join(d, f) for f in os.listdir(d) if f.endswith(".safetensors"))


# ------------------------------------------------------------------------------------------- V1 CaptureSet adapter
@dataclasses.dataclass
class Row:
    """V1 captures.Row: tok, round (= position here: any window row is older than its sample), mx, taps."""
    tok: int
    round: int
    mx: torch.Tensor | None
    taps: torch.Tensor | None


@dataclasses.dataclass
class Sample:
    """V1 captures.Sample fields. round = P and b = the sequence's index in the set, so `key` (r%06d_b%d) is unique
    and V1 analyze's dump glob (r*_b*.pt) finds it. engine_drafts holds caller-supplied drafts (e.g. the K2 worker's,
    g1.py), else NOISE placeholders (engine_drafts_placeholder True)."""
    round: int
    b: int
    req: str
    B: int
    num_tokens: int
    P: int
    anchor: int
    engine_drafts: list[int]
    engine_conf: list[float] | None
    temperature: float
    realized: list[int | None] = dataclasses.field(default_factory=list)
    engine_scheduled: int | None = None
    engine_accepted: int | None = None
    engine_drafts_placeholder: bool = False

    @property
    def key(self) -> str:
        return f"r{self.round:06d}_b{self.b}"


_ZERO_MX = torch.zeros(H, dtype=torch.bfloat16)


class _SeqRows(collections.abc.Mapping):
    """position -> Row for one sequence, built on access from the shard (taps read by slice)."""

    def __init__(self, shard: Shard, i: int):
        self.shard, self.i = shard, i
        s = shard.seq(i)
        self.r0, self.L = s["r0"], s["L"]
        self.tok = shard.tokens(i)

    def __getitem__(self, p: int) -> Row:
        if not self.r0 <= p < self.L:
            raise KeyError(p)
        # D2 has no engine main_x; V1 prepare_window only logs engine-vs-native main_x from it (a zero row reads as
        # relative error 1.0 in run_ref's mainx_engine_vs_native.json, which is meaningless for D2 sets)
        return Row(int(self.tok[p]), p, _ZERO_MX, self.shard.taps(self.i, p, p + 1)[0].clone())

    def __contains__(self, p: object) -> bool:
        return isinstance(p, int) and self.r0 <= p < self.L

    def __iter__(self) -> Iterator[int]:
        return iter(range(self.r0, self.L))

    def __len__(self) -> int:
        return self.L - self.r0


class ShardCaptureSet:
    """V1 CaptureSet interface over D2 shards. Every block of the selected sequences is a sample (stride/limit thin
    them evenly per sequence). Structural checks fill `plumbing` (a violation is never skipped quietly)."""

    WINDOW = W

    def __init__(self, paths: list[str], seq_ids: set[str] | None = None, stride: int = 1,
                 drafts: dict[str, list[int]] | None = None, verify_sha256: bool = False,
                 drafts_by_seq: dict[str, dict[int, list[int]]] | None = None):
        if not paths:
            raise FileNotFoundError("no shard paths given")
        self.shards = [Shard(p, verify_sha256=verify_sha256) for p in paths]
        self.rows: dict[str, _SeqRows] = {}
        self.samples: list[Sample] = []
        self.plumbing: list[str] = []
        self.unverified_tail: list[str] = []
        self.seq_meta: dict[str, dict] = {}
        n = 0
        for sh in self.shards:
            self.plumbing.extend(f"{os.path.basename(sh.path)} {e}" for e in sh.structure_errors())
            for i in range(len(sh)):
                s = sh.seq(i)
                sid = s["seq_id"]
                if seq_ids is not None and sid not in seq_ids:
                    continue
                if sid in self.rows:
                    self.plumbing.append(f"{sid}: present in two shards")
                    continue
                rows = _SeqRows(sh, i)
                self.rows[sid] = rows
                self.seq_meta[sid] = s
                tok = rows.tok
                sp = sh.span(i)
                hit = sp["top_ids"].long() == tok[s["gen_start"]:].long()[:, None]
                if not torch.equal(hit.any(-1), sp["tok_rank"] <= TOPK):
                    self.plumbing.append(f"{sid}: tok_rank disagrees with top_ids")
                for P in sh.blocks(i)[::stride]:
                    key_b = n
                    realized = [int(tok[t]) if t < s["L"] else None for t in range(P + 2, P + 2 + K)]
                    smp = Sample(P, key_b, sid, 1, 1, P, int(tok[P + 1]), [], None, 1.0, realized)
                    d = None if drafts is None else drafts.get(smp.key)
                    if d is None and drafts_by_seq is not None:            # g1.py k2 output: seq_id -> {P: drafts}
                        d = drafts_by_seq.get(sid, {}).get(P)
                    if d is None:
                        smp.engine_drafts = [NOISE] * K
                        smp.engine_drafts_placeholder = True
                    else:
                        if len(d) != K:
                            self.plumbing.append(f"{smp.key}: {len(d)} supplied drafts")
                        smp.engine_drafts = [int(x) for x in d]
                    self.samples.append(smp)
                n += 1

    def window(self, s: Sample) -> list[int]:
        return list(range(max(0, s.P - W + 1), s.P + 1))

    def window_status(self, s: Sample) -> dict:
        rows = self.rows[s.req]
        pos = self.window(s)
        missing = [p for p in pos if p not in rows]
        return {"n": len(pos), "missing": missing, "no_taps": []}

    def arm_inputs(self, s: Sample) -> dict:
        st = self.window_status(s)
        if st["missing"]:
            raise RuntimeError(f"{s.key}: window positions without a row: {st['missing'][:5]}")
        sh, i = self.rows[s.req].shard, self.rows[s.req].i
        lo = max(0, s.P - W + 1)
        taps = sh.taps(i, lo, s.P + 1)
        win = [(p, "taps", taps[j]) for j, p in enumerate(range(lo, s.P + 1))]
        return {"key": s.key, "req": s.req, "P": s.P, "anchor": s.anchor, "window": win,
                "block_positions": list(range(s.P + 1, s.P + 1 + K)), "noise_token": NOISE}


# ------------------------------------------------------------------------------------------------- synthetic data
def synth_record(seq_id: str, split: str, L: int, gen_start: int, seed: int, family: str = "synthetic") -> SeqRecord:
    """A schema-valid random sequence (tests; dataset.py synth). Tokens are random ids, taps N(0,1) bf16, and the
    top-20 a random distribution whose committed token sits inside the top-20 about 90 % of the time."""
    g = torch.Generator().manual_seed(seed)
    tokens = torch.randint(0, VOCAB, (L,), generator=g, dtype=torch.int32)
    r0 = r0_of(gen_start)
    taps = torch.randn(L - r0, TAPW, generator=g).to(torch.bfloat16)
    n = L - gen_start
    top_ids = torch.empty(n, TOPK, dtype=torch.int32)
    top_lp = torch.empty(n, TOPK, dtype=torch.float32)
    tok_lp = torch.empty(n, dtype=torch.float32)
    tok_rank = torch.empty(n, dtype=torch.int32)
    for j in range(n):
        t = int(tokens[gen_start + j])
        logits = torch.randn(TOPK + 1, generator=g, dtype=torch.float64) * 2.0
        lp = torch.log_softmax(logits, -1)          # 21 masses: 20 kept + the tail
        lp_top, _ = lp[:TOPK].sort(descending=True)
        cand = torch.randperm(VOCAB, generator=g)[: TOPK + 1].to(torch.int32)
        ids = cand[cand != t][:TOPK].clone()                 # 20 distinct ids, none of them the committed token
        if float(torch.rand(1, generator=g)) < 0.9:
            k = int(torch.randint(0, TOPK, (1,), generator=g))
            ids[k] = t
            tok_lp[j], tok_rank[j] = float(lp_top[k]), k + 1
        else:
            tok_lp[j], tok_rank[j] = float(lp_top[-1]) - 1.0 - float(torch.rand(1, generator=g)), TOPK + 1
        top_ids[j], top_lp[j] = ids, lp_top.float()
    resid = (1 - top_lp.double().exp().sum(-1)).clamp(0, 1).float()
    meta = {"seq_id": seq_id, "split": split, "family": family, "source": "synthetic", "source_docs": [],
            "prompt_sha256": hashlib.sha256(seq_id.encode()).hexdigest(), "finish_reason": "stop", "segments": 1,
            "replay_request_id": f"d2r.{seq_id}.{r0}.{L}"}
    return SeqRecord(meta, tokens, gen_start, taps, top_ids, top_lp, tok_lp, tok_rank, resid)