[kmontemayor2-sc]

Summary

GraphTransformerEncoder ran sequence_construction_method="ppr" only on the padded path, so PPR configs spent attention work on padding. This change lets them run on packed tokens (one [tokens, dim] buffer and the varlen kernel) on CUDA, with the same sequences and the same outputs as the padded path.

Approach

• The PPR grid code that sat inside tensors_to_graph_transformer_input is now build_ppr_sequence_grids, returning a PprSequenceGrids tuple (node ids, valid mask, ppr_weight, ppr_relation_features, and the ppr_is_expanded token mask). The padded and packed paths call it, so both see the same tokens, features and relation filter.
• The packed path packs those grids with pack_sequences. ppr_weight and ppr_relation_features travel in the layout's token_attrs, so they work as an attention bias and as an anchor-relative input; _packed_anchor_relative_features reads any attribute in token_attrs from the layout and looks up the rest in their sparse matrices.
• Relation attention and message modes match edges to tokens with the padded path's matcher, passing the ppr_is_expanded token mask and the typed-PPR relation-channel masks so edges of tokens the anchor did not expand are dropped. The k-hop frontier-key code moved to _packed_khop_frontier_keys.
• Sinusoidal sequence positions are added from the layout's token positions.
• _packed_unsupported_reasons now rejects only k-hop with hop_distance > 4; PPR sequences do not grow with hop_distance.

User-facing changes

• PPR configs, including ppr_weight, ppr_relation_features and sinusoidal positions, run on packed tokens on CUDA; they fell back to the padded path before. CPU is unchanged. With ppr sequences the packed path reads hop_distance from the batch's sparse matrix, like the padded path; only khop sequences derive it from the layout.
• Host reads: pairwise attention biases add an .any() and a nonzero per bias, and the PPR sequence builder reads the longest candidate list and uses boolean masks. These configs and every ppr config have no fixed host-read count and do not whole-encoder compile.
• In the Notes and CHANGELOG, "training" for the padded-grid attention backend means that autograd is recording; attention dropout in training mode sends every layer, including the last layer of anchor_only, to that backend.
• CUDA runs in a dtype other than float16, bfloat16 and float32 (for example float64) use the padded path, because the varlen kernel does not support them.
• The hop_distance embedding size check applies to khop configs only (see Run supported GraphTransformerEncoder configs on packed tokens #811).
• A ppr batch with no anchors returns an empty embedding on the padded path instead of raising. Requesting ppr_relation_features with max_seq_len=1 raises an error that names max_seq_len.
• New public function gigl.transforms.graph_transformer.build_ppr_sequence_grids and the PprSequenceGrids tuple.
• CHANGELOG entry updated.

Verification

• make check_format_py, make check_format_md and make type_check pass (at 02bfa30). At the tip, ruff check, ruff format --check and a CPU run of packed_graph_transformer_test and transforms/graph_transformer_test pass (the five GT test modules: 418 ran, 82 CUDA skips).
• New tests in tests/unit/nn/packed_graph_transformer_test.py: a heterogeneous PPR packed-vs-padded parity test, a PPR config with a two-entry hop_distance embedding, a no-anchor batch, a max_seq_len=1 error test, and the dtype routing decision. Over 12 PPR configs: supported-config check; packed vs padded output (untyped and typed PPR, 3 seeds); training gradient parity; CUDA vs CPU for inference and training gradients; and a test that the ppr_is_expanded filter changes the packed output.
• At 02bfa30, CPU unit tests for the four GT test files pass (packed_graph_transformer_test 240, transforms graph_transformer_test 125, static_khop_test 17, packed_attention_test 30; CUDA cases skipped there).
• CUDA run on one L4 of the five GT 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: 418 passed, 0 skipped.
• Inference parity, padded vs packed fp32: min cosine 0.99999988 and max abs diff <= 1.1e-4 in all four configs below.

Performance

Synthetic PPR batches on one L4, torch 2.8.0+cu128, random weights, 3 layers, hidden 512, 8 heads, max_seq_len 256, anchor_only readout, hop_distance 2; each anchor has 32 to 300 PPR targets from a 20000-node pool. Inference: 512 anchors, fp32 with TF32, mean of 20 batches after 8 warmup. Training: 128 anchors, bf16 autocast, AdamW, mean of 20 steps. "Compile" is each layer under torch.compile. Cells are ms per batch or step / peak MiB, padded vs packed. Measured at 2a96c1b, the PPR code commit; the later commits are base merges and comment fixes.

Config	Inference eager	Inference compile	Training eager	Training compile
plain	180.12 / 5291 vs 116.54 / 3802	163.88 / 3483 vs 107.4 / 2554	96.78 / 2120 vs 56.96 / 1617	59.89 / 1623 vs 40.63 / 1286
ppr_weight bias	280.05 / 5561 vs 144.36 / 3805	205.11 / 3490 vs 120.84 / 2556	170.31 / 3256 vs 66.22 / 1644	91.06 / 2140 vs 47.9 / 1217
ppr_weight bias + ppr_relation_features input	298.79 / 5454 vs 145.94 / 3918	220.99 / 3564 vs 123.08 / 2669	173.32 / 3148 vs 66.68 / 1756	92.85 / 2249 vs 48.0 / 1216
sinusoidal positions	202.7 / 5704 vs 122.92 / 3912	180.02 / 3814 vs 115.99 / 2663	99.69 / 2228 vs 58.7 / 1727	62.67 / 1734 vs 42.87 / 1400

🤖 Generated with Claude Code