TraceDiff: Post WF name disambiguation

[kyle-hoffmeyer]

Summary

This PR adds name disambiguation to TraceDiff's Wagner-Fischer alignment along with other minor changes.

---

Problem

Wagner-Fischer aligns child node lists by normalized name. When multiple children share the same name, WF treats them as interchangeable (all have edit cost 0) and picks matches based on positional alignment. This can pair structurally unrelated nodes.

Consider a transformer attention block where both traces have multiple aten::bmm children, but one trace has an extra wrapping layer that shifts positions:

trace1 children: [aten::bmm_A, aten::softmax, aten::bmm_B, aten::linear]
trace2 children: [aten::dropout, aten::bmm_A, aten::softmax, aten::bmm_B]

Here bmm_A computes Q×K^T (children: aten::transpose, aten::expand) and bmm_B computes Attn×V (children: aten::expand, aten::reshape). They have completely different subtrees but identical names.

WF aligns by position and produces:

match:  trace1[0] aten::bmm_A  ↔  trace2[1] aten::bmm_A   ✓ correct
match:  trace1[1] aten::softmax ↔ trace2[2] aten::softmax  ✓ correct
match:  trace1[2] aten::bmm_B  ↔  trace2[3] aten::bmm_B   ✓ correct
insert: trace2[0] aten::dropout                             ✓ correct
delete: trace1[3] aten::linear                              ✓ correct

But if the position shift is different — say trace2 has the aten::bmm nodes in a different order relative to other ops — WF may instead produce:

match:  trace1[0] aten::bmm_A  ↔  trace2[3] aten::bmm_B   ✗ wrong!

WF sees both as aten::bmm with cost 0 and has no way to distinguish them. The resulting diff falsely reports that bmm_A's children changed (transpose/expand → expand/reshape) when in reality the operations are identical across traces — they were just matched with the wrong counterpart.

---

Solution: Post-WF Disambiguation

A new post-processing pass (_disambiguate_same_name_candidates) runs after every WF alignment. It identifies ambiguous matches and resolves them using subtree structure.

How it works

Step 1 — Index unmatched nodes by name.
Partition WF ops into match/delete/insert. Group deleted indices (trace1 unmatched) and inserted indices (trace2 unmatched) by name.

Step 2 — Detect ambiguity.
For each WF match (i, j) with name N, check if there are also unmatched nodes named N on either side. If not, the match is unambiguous — skip it.

Step 3 — Disambiguate using BFS child-name levels.
When ambiguity exists:

• Anchor: the matched node on the fixed side (e.g., children1[i] when disambiguating trace2 candidates).
• Candidates: the current WF match plus all same-named unmatched nodes on the other side.

_bfs_child_name_levels computes a fingerprint for each node: a list of sorted tuples, one per BFS depth level, containing the normalized names of all GPU-path children at that depth. Example:

aten::linear
├── aten::mm          ← level 1
├── aten::add         ← level 1
│   ├── aten::copy_   ← level 2
│   └── aten::mul     ← level 2
└── aten::reshape     ← level 1

fingerprint = [
    ("aten::add", "aten::mm", "aten::reshape"),
    ("aten::copy_", "aten::mul"),
]

_find_bfs_survivors compares each candidate's fingerprint against the anchor's. A candidate is evicted if any level differs. If exactly one candidate survives and it differs from WF's original choice, a reassignment is recorded.

Step 4 — Apply reassignments.
Substitution maps (swap_i, swap_j) rewrite the ops list: demoted match sides become delete/insert, promoted nodes are consumed from the delete/insert pool, and the new match is emitted in place.

Integration

The disambiguation pass is called at the end of wagner_fischer, so all callers (Phase 2, Phase 3, Phase 4 in traverse_and_merge) get disambiguated results automatically.

This new post-WF disambiguation does not change the number of matches/inserts/deletes. It only changes existing matches.

---

Other Changes

cuda_driver category support

cuda_driver events (e.g., cuLaunchKernel) are now treated equivalently to cuda_runtime events throughout the TraceDiff pipeline:

Kernel launch name equivalence

A new mapping (_KERNEL_LAUNCH_EQUIVALENTS) normalizes vendor-specific kernel launch API names to a common name during comparison:

{
    "hipModuleLaunchKernel": "__kernel_launch__",
    "cuLaunchKernel": "__kernel_launch__",
}

This prevents spurious diffs when comparing HIP and CUDA traces that use different launch APIs for the same operation.

collapse_single_gpu_child stops at kernel-dispatching cpu_ops

Previously, collapse_single_gpu_child only stopped at kernel dispatch leaf nodes or nodes with != 1 GPU-path child. Now it also stops at cpu_op nodes that have a kernel dispatch child. This prevents collapsing through the last meaningful cpu_op before a kernel launch, preserving it as a comparison point in the diff.

Phase 3 reconciliation preserves Phase 2 matches

When reconcile_unmatched expands wrapper nodes, the resulting children lists are re-indexed. Previously, all children were re-aligned with a fresh WF pass, which could undo correct Phase 2 matches (after the new post-WF disambiguation). Now, Phase 2 match ops are pinned by UID pair, and WF only runs on the unpinned (newly expanded) children.