gigl.nn.packed_attention#

Attention over packed variable-length sequences: the compute half of the packed backend.

This is the attention side of the packed backend of GraphTransformerEncoder. The encoder builds one token sequence per anchor (for example, from its k-hop neighborhood with gigl.transforms.static_khop), copies the real tokens of all sequences into one [num_tokens, dim] buffer, and runs the transformer on that buffer. packed_attention is the attention step and does not depend on how the sequences were built: each sequence attends only to itself, with boundaries given by cu_seqlens.

Why this exists:

  • With uneven sequence lengths (about 90 real tokens out of 256 in our models), most of the linear-layer, FFN and attention work on a padded [num_sequences, max_seq_len, dim] batch is spent on padding. Packed tokens spend it only on real ones.

  • On CUDA this calls aten’s memory-efficient varlen attention kernel, wrapped in a custom op. The kernel is private because torch 2.8 has no public attention API that takes cu_seqlens; _varlen_attention explains why the public options do not fit. The kernel runs in fp32, which inference needs for accuracy, and aten’s own backward kernel supplies the gradients.

  • On CPU the same inputs go through SDPA on a padded grid. That is the reference the CUDA path is tested against, and it lets the packed backend run anywhere.

  • Packing needs one host read (the total token count, in pack_sequences) so the buffer holds exactly the real tokens. Everything else has fixed shapes and no syncs, so the layers can be compiled with torch.compile.

Data flow of a packed forward:

padded grid [num_sequences, max_seq_len, ...]      lengths [num_sequences]
    |                                                   |
    |      pack_sequences(lengths, max_seq_len)  <------+
    |          -> flat_index, PackedLayout(cu_seqlens, max_seq_len)
    |  grid.flatten()[flat_index]
    v
packed buffer [num_tokens, dim]
    |
    |  per layer, on real tokens only:
    |    norm -> q / k / v linears -> [num_tokens, num_heads, head_dim]
    |    packed_attention(q, k, v, layout, ...)
    |        CUDA: aten memory-efficient attention with cu_seqlens
    |        CPU:  SDPA on a padded grid (reference)
    |    output linear + residual -> FFN + residual
    v
anchor readout: rows cu_seqlens[:-1], the first token of every sequence
Constraints:

  • The varlen kernel is CUDA only; other devices use the padded SDPA path.

  • head_dim must be a multiple of 8 for the aligned kernels.

  • Attention dropout is not supported; the kernel’s varlen dropout has not been tested against the padded path.

  • cu_seqlens is int32, so the token count must fit in int32.

Example

>>> from gigl.transforms.sequence_layout import pack_sequences
>>> lengths = torch.tensor([3, 1, 2])
>>> flat_index, layout = pack_sequences(lengths, max_seq_len=3)
>>> flat_index.tolist()
[0, 1, 2, 3, 6, 7]
>>> num_tokens = int(lengths.sum())
>>> query = key = value = torch.randn(num_tokens, 2, 8)
>>> output = packed_attention(
...     query, key, value, layout, query_is_first_token=False
... )
>>> output.shape
torch.Size([6, 2, 8])
>>> first_token_only = packed_attention(
...     query[layout.cu_seqlens[:-1].long()],
...     key,
...     value,
...     layout,
...     query_is_first_token=True,
... )
>>> first_token_only.shape
torch.Size([3, 2, 8])
Results:

We have seen the following with this path. These are observations on one production model, not benchmarks of this module alone.

One L4 GPU, 3-layer graph transformer (hidden 512, 8 heads), sequences of up to 256 tokens holding about 90 real tokens on average, realistic batches, steady state:

  • Forward, no compile: 225 -> 76 ms per batch (~3.0x).

  • Forward, per-layer torch.compile: 53 ms per batch (~4.3x vs the padded path without compile).

  • Training step: 944 -> 364 ms (~2.6x).

  • Peak GPU memory: 15.4 -> 6.8 GiB training, 6.9 -> 3.8 GiB forward.

120 L4 GPUs. These runs also include torch.compile and loader changes, so the gain is not from packing alone:

  • Training step: 6.62 -> 2.66 s (2.5x), peak memory 9.8 -> 4.3 GiB.

  • Full-graph inference over ~230M nodes: 1884 -> 1014 s end to end (1.9x). The model stopped being the bottleneck; sampling was.

Functions#

packed_attention(query, key, value, layout, ...)

Attention within each packed sequence; the entry point of this module.

Module Contents#

gigl.nn.packed_attention.packed_attention(query, key, value, layout, query_is_first_token)[source]#

Attention within each packed sequence; the entry point of this module.

Uses the CUDA varlen kernel on GPU and padded SDPA elsewhere. With query_is_first_token, only each sequence’s first token is queried, which is all the anchor readout of the last layer needs.

Every sequence must contain at least one token (lengths >= 1). An empty sequence gives wrong output on the padded path and is not supported.

Args and returns are as in _padded_packed_attention.

Parameters:

  • query (jaxtyping.Float[torch.Tensor, num_queries num_heads head_dim])

  • key (jaxtyping.Float[torch.Tensor, num_tokens num_heads head_dim])

  • value (jaxtyping.Float[torch.Tensor, num_tokens num_heads head_dim])

  • layout (gigl.transforms.sequence_layout.PackedLayout)

  • query_is_first_token (bool)

Return type:

jaxtyping.Float[torch.Tensor, num_queries num_heads head_dim]