Opt-in batched block-diagonal attention: 2x pretraining throughput, validated equivalent

[alex-crlhmmr]

Motivation

FaceTransformerLayer.forward and EdgeTransformerLayer.forward iterate over the parts of a batch in a Python loop, computing attention per part. Each step launches thousands of small kernels from one Python thread, and on fast GPUs the forward becomes kernel-launch-bound rather than compute-bound. Profiling pretraining (paper config, batch 16) on an H100: the dataloader accounts for 0.3% of wall time, the two per-part loops for ~90%, GPU utilization sits around 40%, and py-spy shows the main thread pinned inside the per-part einsum. The GPU starves while the CPU dispatches kernels.

Change

Adds _bd_attn, a batched block-diagonal attention path used by both layers:

• Parts are padded onto a batch axis. Parts cannot attend across each other, so the block-diagonal structure comes for free from batching; no cross-part masking logic is needed beyond padding.
• The topology attention bias (the _mean_multiedge_bias scatter) is computed for all parts in one index_add_ over flattened (part, src, dst) indices instead of per part.
• Pad keys are masked with -inf before the softmax.
• Attention dropout is applied to the batched weights exactly as in the loop path (this matters: attention_dropout is active at its 0.1 default during pretraining, see Expose --attention_dropout in pretraining (consistency with classification/segmentation) #2).
• The edge layer falls back to the existing loop when the largest part in the batch exceeds EDGE_CAP edges (default 2000), bounding the padded [P, Emax, Emax] allocation on edge-heavy parts (see Edge/face transformer attention is O(E^2) per part with no cap, OOMs on high-edge-count B-reps #3).

The path is opt-in via BATCHED_FACE=1 and BATCHED_EDGE=1; without the env vars, behavior is byte-identical to the current code.

Measured speedup

Pretraining, paper config (batch 16, 6-layer Dual Transformer):

setup	it/s	100-epoch wall clock (250k)
loop (stock), H100	5.39	~78 h
batched, H100	10.96	~39 h
batched, H200 + high-clock host	13.2	33 h

2.03x from the code change alone. The training math is untouched (same batch size, same LR schedule).

Equivalence evidence

• Synthetic checks: batched vs loop outputs agree within 6e-8.
• 5-epoch training comparison (seed 0, 10k subset): |batched - stock| stays flat at ~3e-4 across epochs, below the seed-to-seed floor (~1e-3 to 4e-3) at every epoch. No drift.
• 16-epoch full-corpus comparison: stock 4-seed mean final loss 7.75e-3 (spread 7.66 to 7.83e-3); batched seed-0 7.71e-3, inside the seed spread.
• End-to-end: a full 100-epoch pretraining on Brep2Shape-250k with this path, followed by fine-tuning, reproduces the paper's downstream results (FabWave 100.0 Acc, TMCAD 86.4 Acc).

Not bit-exact with the loop path: the dropout RNG stream differs because dropout is applied to one batched tensor instead of per part. Statistically the runs are indistinguishable from a seed change, as shown above.