shared_resources/tree_duplicate_identity_byteshark/INTEGRATION.md

Integration Notes

Copy duplicate_identity_probe.py next to a tree-v2 submission's sitecustomize.py and tree_v2_ext.py, then add the hooks below. The snippets match the current chiku/byteshark tree-v2 family; apply equivalent edits if the local source has drifted.

tree_v2_ext.py

Add the gated flag next to the tree env settings:

TREE_DUP_IDENTITY_PROBE = os.environ.get("TREE_DUP_IDENTITY_PROBE", "0") == "1"

In _apply_runner_tree_patch, replace the duplicate-position assignment:

t = _t.arange(K * (TREE_W - 1), device=pos.device)
pos[1 + K : R] = pos[1 + (t // (TREE_W - 1))]

with:

t = _t.arange(K * (TREE_W - 1), device=pos.device)
dup_rows = 1 + K + t
src_rows = 1 + (t // (TREE_W - 1))
pos[dup_rows] = pos[src_rows]
if TREE_DUP_IDENTITY_PROBE:
    import duplicate_identity_probe as _dup

    _dup.record_positions(pos, src_rows, dup_rows)
    _STATE["dup_identity_src_drafts"] = src_rows - 1
    _STATE["dup_identity_dup_drafts"] = dup_rows - 1

The src_rows - 1 / dup_rows - 1 values convert verify rows into draft-row indices for metadata.draft_token_ids, metadata.target_logits_indices, and target_argmax.

sitecustomize.py

In the direct rejection path, immediately after:

argmax_token_ids = logits.argmax(dim=-1)

add:

if SPEC_TREE_SPEC and os.environ.get("TREE_DUP_IDENTITY_PROBE", "0") == "1":
    try:
        import duplicate_identity_probe as _dup
        import tree_v2_ext as _T2

        src = _T2._STATE.get("dup_identity_src_drafts")
        dup = _T2._STATE.get("dup_identity_dup_drafts")
        if _T2._STATE.get("tree_step") and src is not None and dup is not None:
            _dup.record_draft_tokens(metadata.draft_token_ids, src, dup)
            _dup.record_target_indices(metadata.target_logits_indices, src, dup)
            _dup.record_argmax_rows(argmax_token_ids, src, dup)
    except Exception:
        if os.environ.get("TREE_DUP_IDENTITY_REQUIRE", "0") == "1":
            raise

Immediately after:

target_argmax = argmax_token_ids[metadata.target_logits_indices].contiguous()

add:

if SPEC_TREE_SPEC and os.environ.get("TREE_DUP_IDENTITY_PROBE", "0") == "1":
    try:
        import duplicate_identity_probe as _dup
        import tree_v2_ext as _T2

        src = _T2._STATE.get("dup_identity_src_drafts")
        dup = _T2._STATE.get("dup_identity_dup_drafts")
        if _T2._STATE.get("tree_step") and src is not None and dup is not None:
            _dup.record_target_argmax(target_argmax, src, dup)
    except Exception:
        if os.environ.get("TREE_DUP_IDENTITY_REQUIRE", "0") == "1":
            raise

Pass Condition

Run only on a duplicate-only diagnostic and require:

positions: bad=0
draft_tokens: bad=0
target_logits_indices: bad=0
argmax_rows: bad=0
target_argmax: bad=0

target_logits_indices checks index[row] == row for both source and duplicate rows on the current flat metadata path. It should not check that the source row and duplicate row have the same index; they are different rows whose outputs should match.