[kmontemayor2-sc]

Summary

Packed-token attention with a per-pair bias (pairwise bias or relation attention) always used the padded-grid SDPA backend, even when nothing needs a gradient. This change sends those calls to the varlen kernel at inference, and folds the relation score into a per-key bias in anchor-only last layers.

Approach

• _varlen_attention takes an optional per-sequence attn_bias. The kernel indexes the bias of sequence i at i * stride(0), so the grid is passed as a strided view. When one head and full sequences make that view reach past the grid, the grid is copied into a zero buffer that holds the view. The kernel has no backward for a bias with cu_seqlens, so a call with a bias raises ValueError when the bias, the query, the key or the value requires grad; packed_attention routes those calls to the padded path.
• packed_attention dispatch: dropout goes to the padded path. A per-pair bias with a gradient and all tokens queried goes to the padded path. Otherwise the kernel runs. With query_is_first_token, row 0 of the per-pair bias is added to the per-key bias and folded into q and k, so it runs on the kernel with or without a gradient.
• With first_token_output, every relation edge targets the anchor, so the relation score depends only on the key. The layer adds it to bias.per_key instead of building a per-pair grid. The shared score computation moved to _relation_attention_scores.
• The per-pair bias is made contiguous before the strided view.

User-facing changes

No API changes. Packed inference with pairwise_attention_bias_attr_names or relation_attention_mode is faster and uses less memory; training with all tokens queried is unchanged.

Verification

• ruff check, ruff format --check and dprint check pass at this tip.
• New CUDA tests: kernel vs CPU padded reference (fp32 and bf16, per-key on and off, max_seq_len not a multiple of 8, one head with full sequences, permuted-view bias); dispatch routing; the bias-requires-grad error; encoder-level CUDA vs CPU for pairwise, pairwise plus key bias, and relation bilinear/hgt across readout modes; whole-encoder compile with no recompiles.
• 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: 340 passed, 0 skipped.
• Inference parity: min cosine >= 0.9999997 vs padded fp32 eager; 0 recompiles.

Performance

One L4, torch 2.8.0+cu128, random weights, 3 layers, hidden 512, 8 heads (head dim 64), anchor_only readout, max_seq_len 256, pairwise bias, 512 anchors per batch, fp32 with TF32, mean of 40 batches. Both columns are on packed tokens: grid backend forced vs the new varlen dispatch.

Inference	Grid backend forced	Varlen dispatch
Eager	185.7 ms / 5606 MiB	123.8 ms / 4895 MiB
Per-layer compile	145.4 ms / 5410 MiB	100.0 ms / 4224 MiB

The training anchor-only fold was measured earlier at 1% to 7% faster, because only the last of three layers changes. The kernel still holds the dense per-pair bias; it skips the scores and padded tokens.

🤖 Generated with Claude Code