[kmontemayor2-sc]

Summary

GraphTransformerEncoder on CUDA now keeps configs with relation_attention_mode="edge_type_bilinear" or "edge_type_hgt" on the packed-token path instead of falling back to padded sequences. Together with the earlier packed-path PRs in this stack, this leaves only ppr_weight, ppr_relation_features, non-khop sequences, hop_distance > 4, CPU, and oversized k-hop candidate sets on the padded path.

The PR also removes a per-batch cost from the packed pairwise attention bias (the bias is now made contiguous once per forward) and adds a training test that compiles the whole encoder.

Approach

• Each layer turns its per-edge relation scores into a per-layer pair bias on top of any incoming pairwise bias, and attention runs on the existing padded-grid backend. The incoming bias is copied, never written in place, so layers do not see each other's scores.
• The score computation is shared with the padded path. It scores each relation segment with batched matmuls, and one scatter writes all scores into the grid. This removes the per-relation .item() reads that made the padded path graph-break under torch.compile, and the packed path does not recompile when per-relation edge counts change between batches.
• The pairwise bias projection returns a permuted view; the encoder now calls .contiguous() once per forward so each layer's attention reads a dense tensor.

User-facing changes

• relation_attention_mode values edge_type_bilinear and edge_type_hgt are accepted on the packed path (any relation_message_mode, both readout modes, with or without pairwise biases). Embeddings and gradients match the padded path up to float error. No config or API changes.
• GraphTransformerEncoderLayer.forward_packed now raises ValueError only when layout.relation_edges is missing for a relation attention or message mode; it no longer rejects relation attention.
• The packed-tokens CHANGELOG.md entry now lists every config the packed path supports at this point in the stack, including the neighbor readout, attention dropout and pairwise biases from the earlier PRs.
• Known limitation: with relation attention under per-layer torch.compile, the packed path builds 35 graphs in inference (3 after the first batch) and 60 in training, with long warmup (about 108 s in inference). Timing is stable after warmup. This is not fixed here.

Verification

• ruff check, ruff format --check and dprint check pass at this tip.
• make unit_test_py (including type_check, CPU) at an earlier tip of this branch: 1563 tests, OK, 45 skipped.
• New CPU tests: packed vs padded outputs for bilinear and HGT (with message modes, out sampling direction, neighbor readout, pairwise bias), training gradients, filtered-edge gradients, and no recompiles across batches.
• New CUDA tests: whole-encoder inductor compile in training against eager gradients (8 configs, including relation attention, pairwise bias and anchor bias), and relation training gradients against CPU.
• CUDA run on one L4 of the five graph-transformer test modules (tests/unit/nn/packed_attention_test.py, tests/unit/nn/packed_graph_transformer_test.py, tests/unit/nn/graph_transformer_test.py, tests/unit/transforms/graph_transformer_test.py, tests/unit/transforms/static_khop_test.py) at this tip: 322 passed, 0 skipped.
• Benchmark parity: min cosine >= 0.9999997 against padded fp32 eager in every inference cell below; 0 recompiles in every cell.

Performance

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

Relation attention edge_type_bilinear, padded vs packed:

Mode	Padded	Packed
Inference, eager	280.8 ms / 6982 MiB	206.6 ms / 5604 MiB
Inference, per-layer compile	246.5 ms / 6059 MiB	126.4 ms / 5152 MiB
Training, eager (quarter-size batch)	1095.4 ms / 18386 MiB	676.8 ms / 12195 MiB
Training, per-layer compile (quarter-size batch)	901.0 ms / 14413 MiB	452.9 ms / 10619 MiB

The padded per-layer compile has 52 (inference) and 62 (training) graph breaks in _get_relation_segments and .item() reads; the packed path has 0 and 1. The compiled ratios therefore flatter the packed path; the eager ratios (1.4x inference, 1.6x training) are the conservative ones.

Pairwise bias, padded vs packed, at this PR's tip:

Mode	Padded	Packed
Inference, eager	359.5 ms / 8135 MiB	187.6 ms / 5606 MiB
Inference, per-layer compile	274.8 ms / 6172 MiB	146.4 ms / 5410 MiB
Training, eager	854.8 ms / 11893 MiB	371.1 ms / 6266 MiB
Training, per-layer compile	517.5 ms / 8477 MiB	278.6 ms / 5455 MiB

The base branch's packed pairwise-bias inference was 293.2 ms eager and 184.9 ms per-layer compile (same config, measured at the base tip). In a same-process A/B, the contiguous bias takes packed inference eager from 290.9 to 188.6 ms (fp32 output max abs diff 1.3e-7, gradient relative L2 5.3e-7). Per-layer-compile inference peak rises 256 MiB (5154 to 5410 MiB).

🤖 Generated with Claude Code