shared_resources/tree_duplicate_identity_byteshark/duplicate_identity_probe.py

"""Duplicate-row identity probe for Gemma tree-verify debug builds.

Enable only on duplicate-only/fake-request equivalence diagnostics. In normal
tree mode branch rows intentionally differ from their source rows, so this
probe is not a production correctness check.
"""

from __future__ import annotations

import os
import sys
from dataclasses import dataclass
from typing import Any


ENABLED = os.environ.get("TREE_DUP_IDENTITY_PROBE", "0") == "1"
REQUIRE = os.environ.get("TREE_DUP_IDENTITY_REQUIRE", "0") == "1"
EVERY = max(1, int(os.environ.get("TREE_DUP_IDENTITY_EVERY", "128")))
MAX_EXAMPLES = max(0, int(os.environ.get("TREE_DUP_IDENTITY_MAX_EXAMPLES", "4")))


@dataclass
class _Counter:
    checks: int = 0
    pairs: int = 0
    bad: int = 0
    first_examples: list[tuple[int, int, int, int]] | None = None


_STATE: dict[str, _Counter] = {}


def _counter(stage: str) -> _Counter:
    counter = _STATE.get(stage)
    if counter is None:
        counter = _Counter(first_examples=[])
        _STATE[stage] = counter
    return counter


def _to_long_tensor(value: Any) -> Any:
    import torch

    if isinstance(value, torch.Tensor):
        return value.to(torch.long)
    return torch.as_tensor(value, dtype=torch.long)


def _slice(values: Any, rows: Any) -> Any:
    rows_t = _to_long_tensor(rows).to(values.device)
    return values.reshape(-1)[rows_t]


def _record_mask(
    stage: str,
    bad_mask: Any,
    src_rows: Any,
    dup_rows: Any,
    lhs_values: Any,
    rhs_values: Any,
) -> int:
    if not ENABLED:
        return 0

    try:
        import torch

        src_t = _to_long_tensor(src_rows).reshape(-1)
        dup_t = _to_long_tensor(dup_rows).reshape(-1)
        lhs_v = lhs_values.reshape(-1)
        rhs_v = rhs_values.reshape(-1)
        bad_mask = bad_mask.reshape(-1)
        checked = int(bad_mask.numel())
        bad = int(bad_mask.sum().item())

        counter = _counter(stage)
        counter.checks += 1
        counter.pairs += checked
        counter.bad += bad

        if bad and counter.first_examples is not None and len(counter.first_examples) < MAX_EXAMPLES:
            bad_idx = torch.nonzero(bad_mask, as_tuple=False).reshape(-1)
            for idx in bad_idx[: MAX_EXAMPLES - len(counter.first_examples)].tolist():
                counter.first_examples.append(
                    (
                        int(src_t[idx].item()),
                        int(dup_t[idx].item()),
                        int(lhs_v[idx].item()),
                        int(rhs_v[idx].item()),
                    )
                )

        if counter.checks in (1, 2, 4, 8) or counter.checks % EVERY == 0 or bad:
            examples = counter.first_examples or []
            print(
                "[tree-dup-id] "
                f"stage={stage} checks={counter.checks} "
                f"pairs={counter.pairs} bad={counter.bad} "
                f"last_checked={checked} last_bad={bad} "
                f"examples={examples}",
                file=sys.stderr,
                flush=True,
            )

        if bad and REQUIRE:
            raise RuntimeError(
                f"[tree-dup-id] {stage} duplicate-row identity failed: "
                f"bad={bad}/{checked}; examples={counter.first_examples}"
            )
        return bad
    except Exception as exc:
        if REQUIRE:
            raise
        print(
            f"[tree-dup-id] stage={stage} disabled after error: {exc!r}",
            file=sys.stderr,
            flush=True,
        )
        return -1


def _record_equal(stage: str, lhs: Any, rhs: Any, src_rows: Any, dup_rows: Any) -> int:
    if not ENABLED:
        return 0

    try:
        lhs_v = _slice(lhs, src_rows)
        rhs_v = _slice(rhs, dup_rows)
        return _record_mask(stage, lhs_v != rhs_v, src_rows, dup_rows, lhs_v, rhs_v)
    except Exception as exc:
        if REQUIRE:
            raise
        print(
            f"[tree-dup-id] stage={stage} disabled after error: {exc!r}",
            file=sys.stderr,
            flush=True,
        )
        return -1


def record_positions(positions: Any, src_rows: Any, dup_rows: Any) -> int:
    """Assert duplicate rows inherited identical position IDs."""
    return _record_equal("positions", positions, positions, src_rows, dup_rows)


def record_draft_tokens(draft_token_ids: Any, src_drafts: Any, dup_drafts: Any) -> int:
    """Assert duplicate rows were actually fed identical draft token IDs."""
    return _record_equal(
        "draft_tokens", draft_token_ids, draft_token_ids, src_drafts, dup_drafts
    )


def record_target_indices(target_logits_indices: Any, src_drafts: Any, dup_drafts: Any) -> int:
    """Assert current flat metadata maps each draft row to its own target row."""
    if not ENABLED:
        return 0

    try:
        index_v = _to_long_tensor(target_logits_indices).reshape(-1)
        src_t = _to_long_tensor(src_drafts).reshape(-1).to(index_v.device)
        dup_t = _to_long_tensor(dup_drafts).reshape(-1).to(index_v.device)
        src_idx = index_v[src_t]
        dup_idx = index_v[dup_t]
        bad_mask = (src_idx != src_t) | (dup_idx != dup_t)
        return _record_mask(
            "target_logits_indices", bad_mask, src_t, dup_t, src_idx, dup_idx
        )
    except Exception as exc:
        if REQUIRE:
            raise
        print(
            f"[tree-dup-id] stage=target_logits_indices disabled after error: {exc!r}",
            file=sys.stderr,
            flush=True,
        )
        return -1


def record_argmax_rows(argmax_token_ids: Any, src_drafts: Any, dup_drafts: Any) -> int:
    """Assert duplicate rows produce identical argmax tokens after verify."""
    return _record_equal(
        "argmax_rows", argmax_token_ids, argmax_token_ids, src_drafts, dup_drafts
    )


def record_target_argmax(target_argmax: Any, src_drafts: Any, dup_drafts: Any) -> int:
    """Assert duplicate rows remain identical after target_logits_indices gather."""
    return _record_equal(
        "target_argmax", target_argmax, target_argmax, src_drafts, dup_drafts
    )