Score Mods

Score mods transform attention scores before softmax. Each function returns a score_mod — a callable with signature (score, b, h, q_idx, kv_idx) -> score — that can be passed directly to flex_attention. See Concepts for the difference between score_mod and mask_mod.

from torch.nn.attention.flex_attention import flex_attention

out = flex_attention(query, key, value, score_mod=my_score_mod)

ALiBi

Attention with Linear Biases — adds a linear position-dependent bias to attention scores, removing the need for positional embeddings.

from attn_gym.mods import generate_alibi_bias

alibi = generate_alibi_bias(8)
out = flex_attention(query, key, value, score_mod=alibi)

attn_gym.mods.alibi.generate_alibi_bias(H)

Returns an alibi bias score_mod given the number of heads H

Parameters:

Name	Type	Description	Default
H	int	number of heads	required

Returns:

Name	Type	Description
alibi_bias	_score_mod_signature	alibi bias score_mod

Soft-Capping

Tanh soft-capping of attention scores, as used in Gemma-2 and Grok-1.

from attn_gym.mods import generate_tanh_softcap

softcap = generate_tanh_softcap(soft_cap=50.0)
out = flex_attention(query, key, value, score_mod=softcap)

attn_gym.mods.softcapping.generate_tanh_softcap(soft_cap, approx=False)

Returns an tanh bias score_mod given the number of heads H

Parameters:

Name	Type	Description	Default
soft_cap	int	The soft cap value to use for normalizing logits	required
approx	bool	Whether to use the tanh.approx. ptx instruction	False

Returns:

Name	Type	Description
tanh_softcap	_score_mod_signature	score_mod

Sandwich

Sandwich relative positional bias (paper) — the position-position inner product of sinusoidal embeddings with ALiBi-style per-head compression ratios, enabling length extrapolation without learned parameters.

from attn_gym.mods import generate_sandwich_bias

sandwich = generate_sandwich_bias(H=num_heads, max_seq_len=S, device=device)
out = flex_attention(query, key, value, score_mod=sandwich)

attn_gym.mods.sandwich.generate_sandwich_bias(H, max_seq_len, d_bar=128, device='cpu')

Returns a Sandwich bias score_mod.

Sandwich keeps only the position-position inner product of sinusoidal embeddings: bias[h, m, n] = sum_i cos((m - n) * w_i) / ratio_h with w_i = 1/10000^(2i/d_bar) and ALiBi-style per-head compression ratios ratio_h = 8(h+1)/H. The bias is precomputed as a 1-D relative-distance table, padded to a multiple of 8 for vector-load-friendly indexing, and the per-head ratio is applied via a small lookup table (uniform loads let the FLASH backend vectorize the score_mod).

Parameters:

Name	Type	Description	Default
H	int	number of heads.	required
max_seq_len	int	maximum sequence length the bias will be used with.	required
d_bar	int	Sandwich shape hyperparameter (paper uses 128).	128
device	str | device	device for the precomputed bias table.	'cpu'

Returns:

Name	Type	Description
sandwich_bias	_score_mod_signature	sandwich bias score_mod

Graphormer

Graphormer attention biases (paper) — two learnable terms added to attention scores between graph node pairs. Gradients flow back into the captured bias tables through flex_attention's backward, so both train end to end.

Spatial encoding — a learnable per-head bias indexed by the shortest-path distance between node pairs:

import torch
from attn_gym.mods import generate_graphormer_spatial_bias, shortest_path_distances

distances = shortest_path_distances(adjacency, max_distance=5)  # (B, N, N)
spatial_bias = torch.nn.Parameter(torch.zeros(num_heads, 5 + 2, device=device))
graphormer = generate_graphormer_spatial_bias(spatial_bias, distances)
out = flex_attention(query, key, value, score_mod=graphormer)

Edge encoding — averages a learnable scalar per (head, path position, edge type) over the edges along each pair's shortest path:

from attn_gym.mods import generate_graphormer_edge_bias, shortest_path_edge_types

path_types, path_lengths = shortest_path_edge_types(adjacency, edge_types, max_path_len=4)
edge_bias = torch.nn.Parameter(torch.zeros(num_heads, 4, num_edge_types, device=device))
graphormer_edge = generate_graphormer_edge_bias(edge_bias, path_types, path_lengths)
out = flex_attention(query, key, value, score_mod=graphormer_edge)

attn_gym.mods.graphormer.generate_graphormer_spatial_bias(spatial_bias, distances)

Returns a Graphormer spatial-encoding score_mod.

Parameters:

Name	Type	Description	Default
spatial_bias	Tensor	(H, num_buckets) bias table, one learnable scalar per head and shortest-path distance. Pass an nn.Parameter (or a tensor with requires_grad=True) to train it. num_buckets must cover every value in distances: max_distance + 2 when produced by shortest_path_distances.	required
distances	Tensor	(B, N, N) integer shortest-path distances between node pairs.	required

Returns:

Name	Type	Description
graphormer_spatial_bias	_score_mod_signature	Graphormer spatial bias score_mod

attn_gym.mods.graphormer.generate_graphormer_edge_bias(edge_bias, path_edge_types, path_lengths)

Returns a Graphormer edge-encoding score_mod.

Implements the paper's c_ij = mean_n(x_{e_n} . w_n^E) for categorical edge types, where the feature/weight dot product collapses into one learnable scalar per (head, path position, edge type). The sum over path positions is unrolled at trace time, so keep max_path_len small (the paper truncates multi-hop paths too).

The table is sliced per path position because compiled flex_attention only supports one gather per gradient-requiring captured tensor; gradients flow from each slice back to edge_bias through regular autograd. Because those slices join the autograd graph when this generator runs, call it once per forward pass when training (it is cheap).

Parameters:

Name	Type	Description	Default
edge_bias	Tensor	(H, K, num_edge_types) learnable bias table. Pass an nn.Parameter to train it.	required
path_edge_types	Tensor	(B, N, N, K) edge types from shortest_path_edge_types.	required
path_lengths	Tensor	(B, N, N) valid position counts from shortest_path_edge_types.	required

Returns:

Name	Type	Description
graphormer_edge_bias	_score_mod_signature	Graphormer edge encoding score_mod

attn_gym.mods.graphormer.shortest_path_distances(adjacency, max_distance)

Computes batched all-pairs shortest-path distances via Floyd-Warshall.

Distances beyond max_distance and unreachable pairs are bucketed together at max_distance + 1, matching Graphormer's treatment of far/disconnected nodes.

Parameters:

Name	Type	Description	Default
adjacency	Tensor	(B, N, N) boolean or 0/1 adjacency matrices (unweighted edges).	required
max_distance	int	largest distance that gets its own bias bucket.	required

Returns:

Type	Description
Tensor	(B, N, N) int32 distances with values in [0, max_distance + 1]. int32 is the
Tensor	smallest integer dtype flex_attention can gather with, and the matrix is the
Tensor	only O(N^2) state Graphormer needs: compute it once per graph in preprocessing
Tensor	and share it across every layer and head.

attn_gym.mods.graphormer.shortest_path_edge_types(adjacency, edge_types, max_path_len)

Reconstructs the edge types along each pair's shortest path for edge encoding.

Runs Floyd-Warshall with next-hop tracking, then walks each (i, j) path for up to max_path_len hops, recording the type of every traversed edge. Like the paper's multi_hop_max_dist, longer paths are truncated to their first max_path_len edges. Precompute this once per graph alongside shortest_path_distances.

Parameters:

Name	Type	Description	Default
adjacency	Tensor	(B, N, N) boolean or 0/1 adjacency matrices (unweighted edges).	required
edge_types	Tensor	(B, N, N) integer edge-type ids, read only where an edge exists.	required
max_path_len	int	number of leading path positions to record (K).	required

Returns:

Name	Type	Description
path_edge_types	Tensor	(B, N, N, K) int32 edge types; positions past the path end are 0 and must be masked with path_lengths (as the returned score_mod does).
path_lengths	Tensor	(B, N, N) int32 count of valid positions per pair: min(distance, K), with 0 for unreachable pairs and the diagonal.

Activation Score Mod

Wraps an activation function to operate in log-space on attention scores.

attn_gym.mods.activation.generate_activation_score_mod(activation=F.gelu, offset=1.0)

Returns a score_mod that replaces softmax with an arbitrary activation.

Wraps the activation as log(activation(s) + offset) so that FlexAttention's internal softmax reduces to (activation(s) + offset) / ell. Use undo_softmax() on the output to recover activation(S) @ V.

Parameters:

Name	Type	Description	Default
activation	Callable[[Tensor], Tensor]	Pointwise activation function. Must satisfy f(s) + offset > 0 for all s (e.g. gelu, relu, sigmoid).	gelu
offset	float	Additive constant to keep log argument positive. Larger values increase numerical stability but also increase the bias correction.	1.0

attn_gym.mods.activation.undo_softmax(out, lse, v_sum, offset=1.0)

Recover activation(S) @ V from FlexAttention's softmax-normalized output.

Parameters:

Name	Type	Description	Default
out	Tensor	FlexAttention output [B, H, Q_LEN, HEAD_DIM].	required
lse	Tensor	Log-sum-exp from AuxOutput.lse [B, H, Q_LEN].	required
v_sum	Tensor	Bias term sum_j(V_j) for each query's attended keys. For full (unmasked) attention: V.sum(dim=-2, keepdim=True). For causal attention: V.cumsum(dim=-2).	required
offset	float	Must match the offset used in generate_activation_score_mod.	1.0

MLA RoPE Score Mod

RoPE-based score modification for Multi-Head Latent Attention (DeepSeek-V2).

Not recommended for production use

This implementation is a demonstration of what FlexAttention can express, not a performant implementation. It works correctly but is not optimized — use it as a reference for understanding the API, not as a drop-in for real workloads.

attn_gym.mods.latent_attention.generate_mla_rope_score_mod(query_rope, key_rope, num_heads, scale=1.0)

Returns an MLA RoPE score modification function to be used w/ FlexAttention

Parameters:

Name	Type	Description	Default
query_rope	Tensor	Positional embeddings for queries [batch, num_heads, seq_len, head_dim]	required
key_rope	Tensor	Positional embeddings for keys [batch, num_heads//128, seq_len, head_dim]	required
num_heads	int	The number of query heads	required
scale	float	Scaling factor for the positional embedding contribution	1.0

Returns:

Name	Type	Description
mla_rope_score_mod	_score_mod_signature	Score modification function for FlexAttention