Optimize anchor-only graph transformer inference

[kmontemayor2-sc]

Summary

At batch size 1,120 on a Tesla T4, this reduces median graph-transformer transform + encoder latency from 1,298.46 ms to 781.69 ms and raises throughput from 862.6 to 1,432.8 anchors/s (+66.1%).

T4 metric (B=1,120)	GiGL main	This PR	Delta
Transform + encoder median	1,298.46 ms	781.69 ms	-39.8%
Encoder median	1,287.69 ms	762.94 ms	-40.8%
Anchor throughput	862.6/s	1,432.8/s	+66.1%
Peak allocated memory	9.158 GB	9.158 GB	0%

anchor_only readout discards every final-layer output except token zero. The specialized final-layer path keeps full-sequence keys, values, and source relation features, but computes query-side attention, relation-message targets, output projection, and feed-forward work only for the anchor.

The encoder selection logic enables this path during evaluation with anchor_only readout. Relation-aware attention uses a rectangular anchor-query bias while preserving full-sequence keys.

Relation-aware T4 measurements retain the same benefit:

B=1,120 mode	Full final layer	Anchor final layer	Throughput change
Edge-type bilinear attention	1,540.78 ms	932.41 ms	+65.25%
Edge-type HGT attention	1,592.89 ms	961.55 ms	+65.66%

Both cases use nonzero learned relation parameters and match the full path within 1.0431e-7 maximum absolute error.

Correctness

The production-count T4 fixture produced:

Output metric	Value
Cosine similarity mean	1.0
Maximum absolute error	1.0431e-7
Relative L2 error	2.2403e-7

The new real-tensor tests compare the specialized output with an independently evaluated full-sequence reference. They exercise:

• plain attention;
• learned linear relation messages;
• learned attention-weighted relation messages;
• ragged valid masks and additive attention bias;
• nonzero bilinear and HGT relation-aware attention with rectangular anchor-query bias.

Validation

• 66/66 graph-transformer unit tests pass.
• Ruff lint and format checks pass.
• Changed-file ty checks pass.
• git diff --check passes.
• Real T4 A/B used BAGL's production graph-transformer initializer with production-scale node, edge, sequence, and feature dimensions.

The benchmark graph topology, input features, and initial weights are synthetic. A checkpoint-backed L4 validation remains the next deployment gate.

Full diff: 322d351...6989558