Edge/face transformer attention is O(E^2) per part with no cap, OOMs on high-edge-count B-reps

Summary

EdgeTransformerLayer.forward runs full self-attention over every edge of a part, and the topology bias in _mean_multiedge_bias builds a dense [num_nodes * num_nodes, num_heads] buffer. Both are O(E^2) in the edge count of a single part, and nothing bounds E. One part with a large edge count allocates several multi-GB tensors and OOMs the whole batch, even at --batch_size 16 on an 80 GB GPU. FaceTransformerLayer.forward has the same structure over faces.

Where

models/dual_encoder.py:

• _mean_multiedge_bias (lines 35-45): allocates aggregated = edge_bias.new_zeros((num_nodes * num_nodes, H)) and counts = edge_bias.new_zeros((num_nodes * num_nodes, 1)), i.e. dense [E^2, H] and [E^2, 1], even though the line graph is sparse (far fewer than E^2 connections).
• EdgeTransformerLayer.forward (lines 403-417): scores = einsum("ihd,jhd->hij", q, k) has shape [H, E, E] per part, then the bias is symmetrized (local_bias + local_bias.transpose(0, 1), another [E, E, H]) and permuted before being added to scores.
• FaceTransformerLayer.forward (lines 235-249): same pattern over faces.

There is no cap on edges (or faces) per part. The MAX_* constants only bound per-primitive UV sampling, not the token count fed to attention.

Memory arithmetic

For a part with E edges and H heads (float32):

• scores [H, E, E]: H * E^2 * 4 bytes
• bias [E^2, H] + its symmetrized copy + counts [E^2, 1]: roughly (2H + 1) * E^2 * 4 bytes

With E = 26,946 and H = 4 that is about 11.6 GB for scores alone, plus another ~26 GB across the bias buffers, so a single part needs well over 30 GB of transient E^2 tensors and triggers the CUDA OOM. E around 22k needs ~80 GB for scores by itself.

Context and reproducibility

We hit this while pretraining on a rebuilt corpus. The offending parts were a handful of thread-tessellation and gear bodies whose edge counts reached about 27k. The released Brep2Shape-250k data may not contain parts this large, and our tokenizer may segment curves more finely than the original, so this is partly data dependent. But the code has no upper bound and the cost is quadratic, so any sufficiently high-edge part will OOM regardless of source. Today it surfaces as a bare CUDA out of memory inside the attention einsum, which is hard to trace back to the single part that caused it.

Suggested fixes

Three options, from smallest to most involved. The first two are cheap and I am happy to send a PR.

1. Stop materializing the dense bias buffers. scores (shape [H, E, E]) already exists, so the topology bias can be scattered straight into it. The line graph is sparse, so aggregate the parallel-edge means over unique (src, dst) pairs (there are at most num_line_edges of them, far fewer than E^2) and index_add_ into a flattened view of scores, once for (src, dst) and once for (dst, src) to get the current symmetrization:

# per component, scores: [H, E, E], E = component_edges
Hh = scores.shape[0]
lb = line_bias[line_edge_start:line_edge_end]                 # [m, H], m = num_line_edges
idx = local_src * E + local_dst                               # [m] flat positions
uniq, inv = torch.unique(idx, return_inverse=True)            # compact, no E^2
s = lb.new_zeros((uniq.numel(), Hh)); s.index_add_(0, inv, lb)
c = lb.new_zeros((uniq.numel(), 1));  c.index_add_(0, inv, torch.ones_like(lb[:, :1]))
mean = (s / c).t()                                            # [H, num_uniq]
flat = scores.view(Hh, E * E)
flat.index_add_(1, uniq, mean)                                # (src, dst)
sym = (uniq % E) * E + (uniq // E)                            # swap src/dst
flat.index_add_(1, sym, mean)                                 # (dst, src)

The only E^2-sized tensor left is scores itself. This removes the [E^2, H] aggregate, the [E^2, 1] counts, the [E, E, H] symmetrized copy, and the permute. On the E = 27k part it drops peak transient from ~30+ GB to ~12 GB, and it is numerically identical to the current _mean_multiedge_bias path.

2. Add a fast, clear guard. Expose a --max_edges_per_part (and face equivalent) and raise with the offending part id and edge count instead of an opaque OOM, or filter such parts at dataset-build time. This does not change the math for normal parts and makes the failure actionable. Minimal version inside the forward:

cap = int(os.environ.get("MAX_EDGES_PER_PART", "0"))
if cap and max(line_node_counts) > cap:
    raise ValueError(
        f"A part has {max(line_node_counts)} edges (> {cap}); edge attention is "
        f"O(E^2) and will OOM. Filter this part or raise the cap."
    )

3. Memory-efficient attention (real scaling fix). Replace the dense per-part [H, E, E] scores with a tiled / flash-style attention that never materializes the full E x E matrix, applying the topology bias per tile. This is the only option that removes the quadratic memory. It is more involved because of the additive topology bias, so it is worth a separate PR.

Options 1 and 2 together keep the current behavior for normal parts, roughly halve peak memory on the pathological ones, and turn the remaining failure into a clear error. Let me know if that direction looks right and I will open a PR.

Start in models/dual_encoder.py by reading _mean_multiedge_bias, EdgeTransformerLayer.forward, and FaceTransformerLayer.forward, then reproduce the CUDA OOM with a high-edge-count part. Compare the dense bias allocations and attention scores against the sparse aggregation and cap options described in the issue. Done means pathological parts no longer fail opaquely and normal-part behavior remains unchanged.