[kmontemayor2-sc]

PR 1/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

• The packed GraphTransformerEncoder backend needs, for every anchor, the list of neighbour tokens within k hops and
  the hop distance of each token. The existing sparse k-hop builders in gigl/transforms/graph_transformer.py produce
  this with data-dependent shapes (sparse matmul, coalesce, unique, nonzero, bool-mask indexing). That costs many
  host syncs and blocks torch.compile.
• Hop distance on that path also comes from an all-pairs hop matrix, which is the memory-heavy part.
• This PR adds gigl/transforms/static_khop.py, which builds the same layout with fixed shapes and per-token hop
  distance, and no all-pairs matrix.
• First of three PRs: this one (k-hop layout), then the packed attention PR (kmonte/gt-packed-attention), then the
  encoder integration (Run supported GraphTransformerEncoder configs on packed tokens #811, kmonte/gt-packed-backend). Nothing uses this module yet, so behavior is unchanged.

Approach

• A dense neighbour table ([num_nodes + 1, width], -1 padded, width = largest node degree) turns each hop into a
  gather. Only sort, gather and scatter are used, so output shapes are known up front.
• Each token is one int64, (node << hop_bits) | hop. The layout scatters and gathers one tensor, and consumers decode
  it with a shift and a mask.
• The order matches the existing sparse builders: anchor first, then neighbours in ascending node id (or by hop and then
  node id with prioritize_hop_order), cut at max_seq_len. The hop is the first-reach hop.
• Memory is O(num_anchors * width ** hop_distance), and one hub node sets width for the whole batch. khop_layout
  counts the cells (candidates plus the neighbour table) before allocating them and raises KhopLayoutTooLargeError
  past a 2^25-cell budget (MAX_LAYOUT_CANDIDATES). The caller can then fall back to the padded path.
• The whole path makes two host reads: the table width, and in pack_sequences the total token count.
• pack_sequences turns the per-sequence lengths into flat gather indices and PackedSequences (int32 cu_seqlens
  offsets, max_seq_len), which the attention PR consumes.

User-facing changes

• New public API in gigl/transforms/static_khop.py, all additive and unused so far:
  • khop_layout
  • pack_sequences
  • PackedSequences
  • KhopLayoutTooLargeError
  • MAX_LAYOUT_CANDIDATES
• khop_layout accepts int32 or int64 edge indices and anchors.
• No existing behavior changes. No CHANGELOG entry; the integration PR's entry covers the feature.

Verification

• tests/unit/transforms/static_khop_test.py, 7 tests, pass (CPU):
  • The layout equals the existing sparse k-hop builders across configurations.
  • Hop order truncation drops the farthest hops first, and hop is first-reach.
  • Over-budget layouts and huge hub tables raise before allocating.
  • pack_sequences offsets are correct.
• static_khop doctests: 9 examples pass.
• ruff and ty are clean.
• CUDA: job gt-split-ut-20261023-165324, one L4 (g2-standard-48, cuda:0), run on the final integration tree, which
  contains this PR's code unchanged: static_khop_test 7 tests pass. The same job ran the whole stack's tests, 181 in
  total, all passing and none skipped.

🤖 Generated with Claude Code

[yliu2-sc]

we'll have separate PR for PPR?

[kmontemayor2-sc]

we'll have separate PR for PPR?

Yes, and we will support many more GT configs.