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

[yliu2-sc]

I'm a little wary that we have some copied over logic in this separate forward, changes in the original forward might not be captured. Can the anchor only be a condition in the existing forward?

[kmontemayor2-sc]

good point - we now have a shared path :)

[yliu2-sc]

Do we need this? Should be unrelated to this PR?

[kmontemayor2-sc]

_forward_anchor_only() filters relation edges whose query position is not the anchor. During training, if that filters every edge, the full-sequence path still includes the relation parameters in autograd and yields zero gradients, while the optimized path would mark them unused (grad=None). _zero_relation_parameter_dependency() adds a zero-valued dependency only in that case, preserving DDP unused-parameter handling and optimizer semantics without changing the forward value or gradients. It is required for training parity, not inference.

[yliu2-sc]

include some comments on what and why we do this

[kmontemayor2-sc]

Do these comments help make it clear?

[zfan3-sc]

Direct method calls skip nn.Module.__call__, so for the final layer this silently bypasses forward hooks and per layer checkpoint wrappers.

Suggestion: have forward pass x.size(1) if query_seq_len is None else query_seq_len to _forward_query_prefix, and call the layer here as:

x = final_encoder_layer(
    x,
    attn_bias=attn_bias,
    valid_mask=valid_mask,
    pairwise_relation_indices=pairwise_relation_indices,
    query_seq_len=1,
)

[kmontemayor2-sc]

Good callout, done :)