[kmontemayor2-sc]

PR 3/3 for Packed Attention in GT

1. Add a fixed-shape k-hop sequence layout for packed graph transformer tokens #814
2. Add varlen attention over packed sequences #815
3. Run supported GraphTransformerEncoder configs on packed tokens #811

Summary

• GraphTransformerEncoder pads every sequence to max_seq_len. When sampled batches fill only a fraction of the
  slots, most of the attention and feed-forward work goes to padding.
• This PR adds a second backend that runs the layers on only the real tokens ("packed" tokens, with variable-length
  attention). The encoder picks it automatically on CUDA when the configuration supports it. There are no new arguments.
• Measured on one L4 (3 layers, hidden 512, 8 heads, max_seq_len 256, 2-hop sampled batches, 512 anchors), the packed
  backend is about 4x faster than origin/main in both inference and training, and its peak memory is about a third.
  Outputs match the padded path.
• Third of three PRs. It builds on the k-hop layout PR (kmonte/gt-static-khop: khop_layout, pack_sequences,
  PackedSequences) and the packed attention PR (kmonte/gt-packed-attention: packed_attention), which hold the
  layout and attention details. This PR is the encoder integration. Next, the neighbour-attention readout PR
  (kmonte/gt-packed-neighbor-readout) stacks on this one.
• It also builds on the tensor-dict inputs PR (tensors_to_graph_transformer_input, Build GT sequence inputs from tensor dicts instead of an internal HeteroData #809), and supersedes the earlier
  stacked PRs (length bucketing, layout-based hop distance, packed forward, packed training).

Approach

• Path selection: the encoder decides at construction whether the config can run packed (see "When packed is used"). A
  batch whose layout is over budget falls back to the padded path, and the encoder counts these batches and warns.
• The packed backend sizes its buffers from the batch. It reads two values back to the host (2 host syncs per batch):
  the largest node degree and the total token count. It then allocates exactly that many tokens.
• GraphTransformerEncoderLayer.forward_packed is the packed version of the layer forward. It shares weights with the
  padded forward, so state dicts are unchanged.
• Anchor-only readout means the last layer only needs the anchor rows as queries, so it computes attention for those
  alone (query_is_first_token in the attention PR).
• Hop distance for the hop_distance token input is derived from the layout (the first-reach hop), so the batch
  attribute is not needed on the packed path.
• Training backpropagates through aten's own autograd formula for the variable-length kernel, so the backward is also
  packed.
• The constructor validates the packed-path constraints (see "Path selection").
• Helpers in gigl/transforms/graph_transformer.py that the encoder needs are now public: get_node_type_offsets,
  concat_tensors, to_homogeneous_edge_index and HOP_DISTANCE_FEATURE_NAME.

When packed is used

The encoder uses the packed backend only when all of these hold. Otherwise it runs the padded path as before.

• The input is on CUDA.
• The sequence builder is khop (prioritize_hop_order is supported).
• The readout is anchor_only.
• The relation mode is none.
• There are no attention biases (no pairwise, PPR or relative-encoding biases).
• Every anchor-based input is a hop-distance input (a subset of hop_distance).
• The attention head dimension is a multiple of 8.
• hop_distance is at most 4.
• attention_dropout_rate is 0.

The hop_distance <= 4 limit comes from memory. The packed layout grows as anchors x width^hops, so past four hops
nearly every batch would exceed the budget below. Per batch, the layout is also skipped (the batch runs padded; the
encoder counts these batches and warns on the 1st, 2nd, 4th, 8th, ... one) when
anchors x (1 + width + ... + width^hop_distance) + (nodes + 1) x width exceeds 2^25 cells, where width is the
largest node degree over all edge types. The second term is the neighbour table, and the check runs before the table is
allocated, so a hub node cannot blow up memory either way.

Why attention dropout is excluded:

• The variable-length kernel's dropout has not been tested against the padded path's per-head, per-sequence masks, so it
  stays off until it has been.

Path selection:

• The check is made at construction, so a model takes the same path in training and eval. This matters because the two
  paths define the hop_distance token input slightly differently when the batch attribute was built with another
  h_max or direction. When the attribute is present, the transform's h_max must be at least the encoder's
  hop_distance and its direction must match sampling_direction.
• A hop_distance embedding on a config that can run packed must have more than hop_distance entries; the constructor
  raises ValueError otherwise.
• When a config is not packed, the reasons are logged once at construction.

User-facing changes

• New public API, all additive:
  • GraphTransformerEncoderLayer.forward_packed
  • get_node_type_offsets, concat_tensors, to_homogeneous_edge_index and HOP_DISTANCE_FEATURE_NAME in
    gigl/transforms/graph_transformer.py, previously private helpers
  • khop_layout, pack_sequences and PackedSequences come from the k-hop layout PR, and packed_attention from the
    packed attention PR.
• No API is removed or renamed.
• Behavior changes:
  • A supported CUDA configuration switches to the packed backend automatically. Outputs match the padded path (see
    Verification), but the computation is different.
  • Training dropout draws differ from the padded path. The distribution is the same, so seeded runs will not reproduce
    earlier ones bit for bit.
  • On the packed path the batch's hop_distance attribute is not read and need not exist. Hop distance is computed
    from the layout as the directed first-reach hop. A batch that falls back to the padded path (layout over budget)
    does need the attribute: without it the encoder raises a ValueError that names AddHeteroHopDistanceEncoding.
    Keep the transform in the pipeline unless every batch is sure to fit the budget.
  • The device argument of forward now accepts a string such as "cuda" as well as a torch.device.
  • GraphTransformerEncoder raises ValueError for max_seq_len < 1 or hop_distance < 0.
• Unsupported configurations and CPU runs are unchanged.

Verification

• CPU, final tip (dev box, tests run as scripts): all pass, ruff and ty clean.
  • packed_graph_transformer_test: 49 tests (7 CUDA-only skipped).
  • packed_attention_test: 3 (2 skipped).
  • nn/graph_transformer_test: 76.
  • transforms/graph_transformer_test: 46 (1 skipped).
  • static_khop_test: 7.
  • The _encode_packed doctest passes.
• The attention and pack_sequences unit tests moved to the PRs that add that code (the packed attention PR and the
  k-hop layout PR). This PR's tests cover the encoder.
• The packed tests force the packed path on CPU, so the same code is checked against the padded path without a GPU.
  CUDA-only tests cover the real kernels:
  • Packed variable-length attention matches the padded reference in forward and backward, in fp32 and bf16 (these tests
    now live in the packed attention PR).
  • The encoder forward on CUDA matches the CPU padded forward (max abs diff about 4e-5, cosine similarity 1.0000).
  • Training gradients on CUDA match the CPU padded gradients.
  • The CUDA forward does exactly 2 host synchronizations per batch, checked with torch.cuda sync debug mode. The L4
    run printed SYNC_COUNT 2.
• CUDA, final tree: job gt-split-ut-20261023-165324, one L4 (g2-standard-48, cuda:0), about 80 s of test time: 181
  tests pass, none skipped, so every CUDA test ran. packed_graph_transformer_test 49, packed_attention_test 3,
  nn/graph_transformer_test 76, transforms/graph_transformer_test 46, static_khop_test 7.
• Earlier L4 run, from the pre-split tip (same code; job gt-adv-a-20261022-185022, g2-standard-48, cuda:0, tests
  ran 45 s): the packed test module, static_khop_test and the other unit tests it imports pass in full, 59 tests, none
  skipped, including every CUDA test (the 2-sync assertion, the 1- and 2-anchor training steps, CUDA training gradients
  at hop_distance=3 and with prioritize_hop_order, the varlen forward with and without grad, and the Inductor
  whole-encoder compile).
• Parity in the benchmark (same weights, all variants against origin/main eager):
  • Inference: minimum cosine similarity 0.9999999, max abs diff at most 4.0e-5.
  • Training: gradient relative L2 difference 7.1e-6 for packed eager and about 2.1e-4 for compiled variants (the
    compiled origin/main whole-encoder run shows the same 2.2e-4 against its own eager).

Whole-encoder torch.compile

torch.compile(encoder) on the packed backend had two problems, both fixed in this PR:

• It recompiled on every batch. forward passed self._encode_packed(...) inline to self._project_output(...).
  _encode_packed has graph breaks (data-dependent sizes), and a graph break inside a call leaves the bound method on
  Dynamo's stack. torch.compile guards on that fresh object, so the guard never matches the next batch. The output
  projection now runs at the end of _encode_packed, so nothing is pending across the break (a CPU test checks the
  graph count stays flat across batches, with and without a fallback batch).
• The attention call is plain aten. An earlier revision wrapped the varlen kernels in torch.library custom ops
  with fakes and a hand-written autograd Function, on the belief that Inductor's layout rules for the aten op fail on
  the [1, tokens, ...] views. They do not here: the token count is a backed dynamic dimension, because the layout's
  host reads break the graph first. The encoder now calls torch.ops.aten._efficient_attention_forward directly (aten
  supplies the autograd formula and the meta function) and asks for the logsumexp only when a gradient is needed. That
  removed about 280 lines. The Inductor compile tests, the compile-trace matrix below and every CUDA test pass.

Result on one L4 (job gt-adv-a-20261022-185022, about 31 min program wall):

• Setup: same model as the Performance section, but with bf16 autocast, 200 distinct 512-anchor inference batches and
  100 distinct training steps (256 queries plus 512 random negatives), then a replay of the first batches.
• Whole-encoder compile compiles on the first 2 calls and never again, including on the replay.
• Nothing hit the 30-call cap.

inference, 512 anchors	ms/batch	peak GiB	compiles (last at)	graphs / breaks	compile total	min cosine vs eager
eager	76.0	3.80	0	0 / 0	0	1
per-layer compile	53.3	3.06	1 (at 0)	3 / 0	31.4 s	0.9999998
whole encoder	53.7	3.01	2 (at 1)	18 / 7	63.4 s	0.9999998
whole, layout eager	52.7	3.01	2 (at 1)	7 / 1	17.7 s	0.9999998

training step	ms/step	peak GiB	compiles (last at)	graphs / breaks	compile total	grad rel. L2 vs eager
eager	353.9	6.81	0	0 / 0	0	0
per-layer	256.1	5.26	1 (at 0)	38 / 1	92.9 s	4.7e-3
whole encoder	245.7	5.27	2 (at 1)	92 / 8	215.4 s	4.8e-3
whole, layout eager	255.7	5.27	1 (at 0)	59 / 2	77.3 s	4.7e-3

• Pick: per-layer compile for inference (53.3 vs 53.7 ms/batch, and 31 s vs 63 s one-time compile). Whole-encoder
  compile for training (245.7 vs 256.1 ms/step, about 4% faster, at 215 s vs 93 s one-time compile). All variants
  settle.
• "Whole, layout eager" runs the k-hop layout and token packing under torch.compiler.disable; it is a benchmark
  variant, not part of the PR.
• fp32 check on the first training batch pair (no autocast, TF32 off): gradients differ from fp32 eager by 1.6e-7 for
  whole-encoder compile and 1.6e-7 for per-layer compile. Eager in bf16 autocast differs from fp32 eager by 1.25e-2, so
  the bf16 gradient differences above are precision noise, not a compile difference.
• These batches are fuller (mean 91 of 256 slots) than the 68 used in the tables below, so eager is slower here.
• The tables below were measured on the base before the review-fix commits; the numbers in this section were measured on
  the pre-split tip, which has the same code as this PR.

Performance

One L4, fp32 with TF32, inference and a full train step (forward + backward).

• Setup: 3 layers, hidden 512, 8 heads, max_seq_len 256, 2-hop sampled batches with fanout 15, 512 anchors, about 68
  of 256 token slots filled, same weights in every cell.
• The tables below are from the earlier run of this revision.
• A re-run after the review fixes measured packed eager at 44.8 ms/batch inference and 148.5 ms/step training on the
  same hardware, with 2 host syncs and exact parity. The earlier harness run of the unfixed tip read 47.7 ms/batch and
  152.5 ms/step.

Inference, ms per batch (peak GiB):

mode	origin/main	#809 only	this PR
eager	184.7 (5.15)	191.6 (5.15)	48.2 (1.50)
per-layer torch.compile	160.3 (3.23)	165.5 (3.23)	45.0 (0.98)
torch.compile(encoder)	152.1 (4.43)	151.5 (4.43)	42.9 (1.33)

Train step, ms per step (peak GiB):

mode	origin/main	#809 only	this PR
eager	623.4 (11.19)	629.6 (11.23)	155.9 (3.35)
per-layer torch.compile	fails	fails	136.6 (3.01)
torch.compile(encoder)	488.2 (10.71)	486.7 (10.71)	138.5 (3.01)

• origin/main and Build GT sequence inputs from tensor dicts instead of an internal HeteroData #809 per-layer compile fails in training with an Inductor AssertionError in the padded gather
  (shown as "fails").
• Compiled graphs and first-call time (inference / training):
  • Per-layer compile: 2 graphs in every version. First call 9.8 s inference on origin/main (training fails), 2.4 s on
    Build GT sequence inputs from tensor dicts instead of an internal HeteroData #809 (training fails), and 6.2 / 13.4 s on this PR.
  • Whole-encoder compile: origin/main 20 graphs, 16.0 s inference and 18 graphs, 20.3 s training. Build GT sequence inputs from tensor dicts instead of an internal HeteroData #809 14 graphs, 2.6
    s inference and 14 graphs, 4.5 s training. This PR 23 graphs, 25.6 s inference and 21 graphs, 15.5 s training.
  • Eager first-call is under a second everywhere.

Headlines:

• Code changes only (origin/main eager to this PR eager):
  • Inference 184.7 to 48.2 ms/batch (3.8x), peak 5.15 to 1.50 GiB.
  • Training 623.4 to 155.9 ms/step (4.0x), peak 11.19 to 3.35 GiB.
• Code plus torch.compile (origin/main eager to this PR's best compiled variant):
  • Inference 184.7 to 42.9 ms/batch with torch.compile(encoder) (4.3x), peak 1.33 GiB.
  • Training 623.4 to 136.6 ms/step with per-layer compile (4.6x), peak 3.01 GiB.

🤖 Generated with Claude Code

[kmontemayor2-sc]

/all_test

[github-actions]

GiGL Automation

@ 18:50:26UTC : 🔄 E2E Test started.

@ 20:57:10UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:50:26UTC : 🔄 Lint Test started.

@ 19:00:13UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:50:26UTC : 🔄 Integration Test started.

@ 20:23:03UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:50:27UTC : 🔄 Scala Unit Test started.

@ 19:00:59UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:50:27UTC : 🔄 Python Unit Test started.

@ 20:31:17UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 18:50:28UTC : 🔄 C++ Unit Test started.

@ 18:52:31UTC : ✅ Workflow completed successfully.

[yliu2-sc]

can it?

[yliu2-sc]

robot comment: _encode_packed() calls encoder_layer.forward_packed(...), but torch.compile(layer) only compiles the module’s ordinary forward. Custom methods delegate to the original module and execute eagerly. BAGL’s active inference helper replaces every encoder layer with torch.compile(layer) (line 88), so supported CUDA configurations will automatically switch to packed execution while enable_torch_compile silently stops compiling the layers. In an exact-head reproduction, compiled.forward_packed(...) left Dynamo counters empty; compiled(...) produced one captured graph. Route packed inputs through the compiled callable, explicitly wrap forward_packed, or change BAGL to whole-encoder compilation, with a regression covering the actual OptimizedModule wrapper.