[kmontemayor2-sc]

Summary

GraphTransformerEncoder can now run relation_message_mode="edge_type_linear" and "edge_type_attention" on packed tokens. Before this PR, any relation message mode sent the batch down the padded path, so relation-aware models paid padding FLOPs and memory for every sequence. relation_attention_mode still has to be "none" on the packed path.

Approach

• PackedRelationEdges carries each relation edge as packed rows (target, source, relation) plus host-side (relation, start, end) segments. GraphTransformerEncoderLayer.forward_packed takes it as relation_edges, which is required when a message mode is set.
• _aggregate_relation_messages is the new layout-independent core. The padded _compute_relation_messages and _compute_relation_attention_messages now call it too, so both paths share one implementation of per-relation transforms, mean or softmax weighting, and scatter-add.
• lookup_pairwise_relation_indices sends all relation edges through the token-pair matcher in one call instead of one call per relation. Host syncs per forward therefore no longer grow with the number of relations or layers. _match_directed_edges_to_token_pairs also returns the matched edge index, and repeat_interleave calls pass output_size to avoid extra syncs.
• With anchor_only readout, the last layer keeps only edges into each sequence's first token.
• When anchor filtering drops every edge in training, the layer adds a zero dependency on the relation parameters so they still receive gradients, as the padded path does.

User-facing changes

• Configs with relation_message_mode in edge_type_linear or edge_type_attention (and relation_attention_mode="none") now run on packed tokens on CUDA. Embeddings and gradients match the padded path up to float error.
• CHANGELOG entry updated. No API signature changes for users of GraphTransformerEncoder.

Verification

• ruff check, ruff format --check and dprint check pass at this tip.
• At an earlier tip of this branch, make unit_test_py (CPU) for packed_attention_test.py (4 ran, 3 CUDA skips), packed_graph_transformer_test.py (93 ran, 12 CUDA skips), graph_transformer_test.py (125 ran, 1 skip), static_khop_test.py (14 ran): all pass, including type_check.
• New PackedRelationMessageTest compares packed against padded outputs and gradients (truncation, edgeless batches, filtered edges, hetero graphs, checkpointing). The no_edges_into_anchors gradient case has no anchor self-loops and asserts that the packed path takes the all-filtered branch. CUDA-only tests check parity with CPU and that the number of host syncs does not depend on layers or relations.
• CUDA run on one L4 of the five graph-transformer test modules at this tip: 237 passed, 0 skipped.

Performance

One L4, torch 2.8.0+cu128, random weights, 3 layers, hidden 512, 8 heads, anchor_only readout, edge_type_linear, max_seq_len 256. Padded vs packed. Inference: 512 anchors, fp32 with TF32, mean of 40 batches. Training: bf16 autocast, mean of 24 steps. Times are ms per batch or step.

Case	Padded	Packed
Inference, eager	249.0 ms / 6984 MiB	97.1 ms / 3973 MiB
Inference, per-layer compile	214.7 ms	75.4 ms
Training, eager	1051.4 ms / 15821 MiB	413.0 ms / 7004 MiB
Training, per-layer compile	802.7 ms	315.8 ms

These were measured before later PRs in the stack, so numbers at this tip may differ slightly.

🤖 Generated with Claude Code