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:

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

In `_apply_runner_tree_patch`, replace the duplicate-position assignment:

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

with:

```python
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:

```python
argmax_token_ids = logits.argmax(dim=-1)
```

add:

```python
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:

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

add:

```python
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:

```text
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.