Attention and dense relations¶
The first four examples require an NVIDIA GPU and the cuda extra; GAT selects
CPU/CUDA according to availability. Core snippets omit setup; run the complete
linked scripts. Install Torch separately before Tiga.
Attention-style workloads expressed with generic relation and reducer
semantics. The mask is never code — it is the graph: a dense relation is full
attention, a triangular relation is causal attention, and a block-diagonal CSR
is varlen causal attention with cu_seqlens. Tile pruning is a separate,
approximate semantic layered on the same online-softmax reducer.
python examples/full_attention.pypython examples/causal_dense_relation.pypython examples/varlen_causal_attention.pypython examples/tile_pruned_attention.pypython examples/gat_edge_attention.py
Full attention, no mask¶
What it is. Attention: each of the N positions carries a query vector, a key vector and a value vector of width D. Every output position is a weighted average of all value vectors, where the weights are a softmax over query·key similarity scores — so the cost is the full N×N score matrix. With no mask, every query attends to every key. In graph terms each edge (j→i) pairs destination position i (the query) with source position j (the key and value); a dense relation contains all N×N edges.
This is the exact special case of everything below — the same tiled online-softmax machinery FlashAttention uses, with every tile admitted.
class FullAttention(tg.MessagePassing):
reducer = tg.online_softmax()
def edge(self, src, dst, edge, scale):
score = (src.key * dst.query).sum(dim=-1) * scale
return self.reducer(score, src.value)
program = FullAttention()
output = program( # (N, H, D)
graph=tg.Graph.dense(nodes, device="cuda"),
src={"key": key, "value": value},
dst={"query": query},
scale=width**-0.5,
)
Full source: examples/full_attention.py (runs as-is)
"""Exact full attention (no mask) as a dense Graph + online-softmax reducer."""
import torch
import tiga as tg
torch.manual_seed(0)
nodes, heads, width = 1024, 4, 64
# the streaming kernel reads lane-major storage: build (H, N, D), view as (N, H, D)
query = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
key = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
value = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
class FullAttention(tg.MessagePassing):
reducer = tg.online_softmax()
def edge(self, src, dst, edge, scale):
score = (src.key * dst.query).sum(dim=-1) * scale
return self.reducer(score, src.value)
program = FullAttention()
output = program( # (N, H, D)
graph=tg.Graph.dense(nodes, device="cuda"),
src={"key": key, "value": value},
dst={"query": query},
scale=width**-0.5,
)
# Inspect: program.explain() provider/lowering/cache summary
# program.ir("kernel_ttir") TTIR (streaming online softmax)
Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)
### FullAttentionTorchExecutor
backend: cuda
provider: triton-nvidia@3.6.0
lowering: gf-kernel-to-ttir-dense-streaming-reduction
passes: gf-verify-domain, gf-lower-domain-to-iter, gf-lower-iter-to-kernel, capture-structured-message, expand-reducer-regions, select-cartesian-coordinate-hierarchy, select-dense-query-neighbor-tile, lower-vector-contraction, lower-streaming-reducer-state, gf-kernel-to-ttir, provider-compile-serialized-ttir
provider cache key: triton-nvidia | cuda:nvidia | 3.6.0 | 2.11.0+cu128 | af81e84448f193cc | cuda:120:warp32
accepted: [machine-schedule] admitted dense-query-key-tile
unknown: [pipeline] no compiler-controlled asynchronous pipeline admitted
remark: [planning] Cartesian relation remained implicit
remark: [planning] multi-value message and reducer state were fused into tiles
remark: [planning] no pairwise score or normalized-weight tensor was materialized
variant cache: hits=0, misses=1
Causal dense relation¶
What it is. Causal attention is attention with a look-ahead ban: position i may average only over positions j ≤ i. This is the rule in language models that generate text left to right, where each position may only read the past. The example expresses the mask as the graph itself: a triangular relation contains exactly the edges (j→i) with j ≤ i, so the edge function needs no masking code at all. As above, source j supplies the key and value, destination i supplies the query.
Graph.triangular() carries the lower-inclusive topology through
Domain → Iter → Kernel, and an ordinary online reducer lowers to causal
streaming TTIR. Torch appears only as external CUDA storage.
class CausalWeightedValue(tg.MessagePassing):
reducer = tg.online_softmax()
def edge(self, src, dst, edge, scale):
score = (src.key * dst.query).sum(dim=-1) * scale
return self.reducer(score, src.value)
kernel = CausalWeightedValue()
output = kernel( # (N, H, D)
graph=tg.Graph.triangular(nodes, device="cuda"),
src={"key": key, "value": value},
dst={"query": query},
scale=width**-0.5,
)
Full source: examples/causal_dense_relation.py (runs as-is)
"""A causal reducer expressed as a triangular Graph, not a core NN operator."""
import torch
import tiga as tg
torch.manual_seed(7)
nodes, heads, width = 256, 2, 64
# the streaming kernel reads lane-major storage: (H, N, D) viewed as (N, H, D)
query = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
key = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
value = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
class CausalWeightedValue(tg.MessagePassing):
reducer = tg.online_softmax()
def edge(self, src, dst, edge, scale):
score = (src.key * dst.query).sum(dim=-1) * scale
return self.reducer(score, src.value)
kernel = CausalWeightedValue()
output = kernel( # (N, H, D)
graph=tg.Graph.triangular(nodes, device="cuda"),
src={"key": key, "value": value},
dst={"query": query},
scale=width**-0.5,
)
# Inspect: graph.explain() relation summary (triangular, edges=N(N+1)/2)
# kernel.explain() provider/lowering/cache summary
# kernel.ir("gf.kernel.ttir") generated TTIR streaming kernel
def main() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
if not torch.cuda.is_available():
raise SystemExit("this CUDA compiler example requires a CUDA device")
return output, query, key, value
if __name__ == "__main__":
main()
Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)
### CausalWeightedValueTorchExecutor
backend: cuda
provider: triton-nvidia@3.6.0
lowering: gf-kernel-to-ttir-dense-streaming-reduction
passes: gf-verify-domain, gf-lower-domain-to-iter, gf-lower-iter-to-kernel, capture-structured-message, expand-reducer-regions, select-cartesian-coordinate-hierarchy, select-dense-query-neighbor-tile, lower-vector-contraction, lower-streaming-reducer-state, gf-kernel-to-ttir, provider-compile-serialized-ttir
provider cache key: triton-nvidia | cuda:nvidia | 3.6.0 | 2.11.0+cu128 | af81e84448f193cc | cuda:120:warp32
accepted: [machine-schedule] admitted dense-query-key-tile
unknown: [pipeline] no compiler-controlled asynchronous pipeline admitted
remark: [planning] Cartesian relation remained implicit
remark: [planning] multi-value message and reducer state were fused into tiles
remark: [planning] no pairwise score or normalized-weight tensor was materialized
variant cache: hits=0, misses=1
Varlen causal attention by graph composition¶
What it is. Real training batches pack many sequences of different
lengths into one flat tensor and mark the boundaries with cumulative lengths
cu_seqlens = [0, L₁, L₁+L₂, …] (the convention of flash-attn's varlen API).
Each sequence attends causally within itself only — position i may read
position j only when both lie in the same sequence and j ≤ i. As a mask that
is a block-diagonal matrix of small triangles; as a graph it is one explicit
CSR relation with exactly those edges, Σ L(L+1)/2 of them. Boundary handling
needs no code: cross-sequence attention is an edge that does not exist.
The example composes per-sequence triangular graphs with tg.Graph.cat, then
runs the same unmasked edge UDF, including a length-1 sequence. Block k attends
only within itself. cat currently materializes CSR, not an implicit composite
relation. The compatibility constructor tg.Graph.cu_seqlens(boundaries, causal=True)
produces the same relation when cumulative boundaries begin with zero. (The tests check
all three attention examples against F.scaled_dot_product_attention.)
Execution falls back to the
exact eager oracle today: compiled streaming covers the dense/triangular
relations above and field-score CSR segment softmax, while an edge-computed
score over a general CSR is a compiler boundary, not a semantic one.
# Block-diagonal causal relation: position i attends j <= i inside its own
# sequence. cat currently materializes CSR, rather than retaining an implicit
# composite representation. No edge is added between sequences.
graph = tg.Graph.cat([
tg.Graph.triangular(length, device="cuda") for length in lengths
])
generator = torch.Generator(device="cuda").manual_seed(20260821)
query = torch.randn(
total, heads, width, device="cuda", dtype=torch.float16,
generator=generator) # (N, H, D)
key = torch.randn(query.shape, device="cuda", dtype=query.dtype, generator=generator)
value = torch.randn(query.shape, device="cuda", dtype=query.dtype, generator=generator)
class VarlenCausalAttention(tg.MessagePassing):
reducer = tg.online_softmax()
def edge(self, src, dst, edge, scale):
score = (src.key * dst.query).sum(dim=-1) * scale
return self.reducer(score, src.value)
program = VarlenCausalAttention()
output = program( # (N, H, D)
graph=graph,
src={"key": key, "value": value},
dst={"query": query},
scale=width**-0.5,
)
Full source: examples/varlen_causal_attention.py (runs as-is)
"""Varlen causal attention: compose independent triangular graph blocks."""
import torch
import tiga as tg
heads, width = 4, 64
lengths = [384, 1, 96, 640, 57]
boundaries = torch.tensor([0, *lengths], device="cuda").cumsum(0)
cu_seqlens = boundaries # Compatibility with per-sequence reference checks.
total = sum(lengths)
# Block-diagonal causal relation: position i attends j <= i inside its own
# sequence. cat currently materializes CSR, rather than retaining an implicit
# composite representation. No edge is added between sequences.
graph = tg.Graph.cat([
tg.Graph.triangular(length, device="cuda") for length in lengths
])
generator = torch.Generator(device="cuda").manual_seed(20260821)
query = torch.randn(
total, heads, width, device="cuda", dtype=torch.float16,
generator=generator) # (N, H, D)
key = torch.randn(query.shape, device="cuda", dtype=query.dtype, generator=generator)
value = torch.randn(query.shape, device="cuda", dtype=query.dtype, generator=generator)
class VarlenCausalAttention(tg.MessagePassing):
reducer = tg.online_softmax()
def edge(self, src, dst, edge, scale):
score = (src.key * dst.query).sum(dim=-1) * scale
return self.reducer(score, src.value)
program = VarlenCausalAttention()
output = program( # (N, H, D)
graph=graph,
src={"key": key, "value": value},
dst={"query": query},
scale=width**-0.5,
)
# Execution note: this edge-computed score over a general CSR runs on the
# exact eager oracle today; compiled streaming covers dense/triangular
# relations (full/causal attention) and field-score CSR segment softmax.
# Inspect: program.explain()
Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)
### VarlenCausalAttentionTorchExecutor
backend: reference
provider: torch
lowering: (unspecified)
passes: validate-domain, select-reference-csr
remark: [planning] reference evaluator selected; no performance codegen artifact exists
remark: [planning] cross-apply fusion is handled by the @tg.jit/@tg.program capture boundary rather than the single-apply reference evaluator
variant cache: hits=0, misses=1
Tile-pruned sparse attention¶
Why skip a tile?¶
This is an experimental accuracy/performance tradeoff, not a required part of attention. Exact attention scores every allowed key and accumulates values with softmax weights. If a whole key block scores very poorly, loading and accumulating its values may contribute little. A tile pairs a small query block with a small key block, matching the GPU's blockwise execution.
For example, with a running maximum score of 10 and a new tile maximum of 1, the ratio of their exponential weights is approximately:
block_prune_threshold=0.01 skips that tile. This is a query-block comparison,
not a per-query error guarantee: many omitted terms or large values can still
cause substantial output error.
| Still executed | Skipped for a pruned tile |
|---|---|
| Key loads, query/key scores, threshold comparison | Value loads and subsequent softmax/value accumulation |
Pruning does not eliminate all score computation or guarantee a speedup. Larger
thresholds prune more aggressively. Omit block_prune_threshold for exact
computation; zero is not a valid threshold. Only the specified generated-CUDA
dense forward path supports pruning; backward does not. The example's
width / nodes is an illustrative choice, not a general tuning formula.
Admission rule and limitations¶
What it is. Same attention as the first example, but approximate: tile pruning processes the keys in blocks (tiles) and skips any tile whose best score is far below the running query-block maximum. The threshold alone is not a bound on total omitted probability or output error: those also depend on omitted tile count and value magnitudes. Validate accuracy for each workload. FlashAttention is the exact special case where every tile is admitted; this example adds the admission rule.
The near-diagonal pattern is illustrative, not a fixed local window: actual admission depends on the data's scores.
The emitter compares the maximum score across the entire query/key tile with
the running maximum for that query block. In natural-log score units, the
admission rule uses the ratio rho = block_prune_threshold:
The TTIR implementation uses base-2 scores and log2(rho) equivalently.
Key loads and score computation happen before this decision; pruning skips
value loads and the subsequent softmax/value accumulation.
How the threshold works: it is a constructor parameter of the built-in
tg.online_softmax() reducer that the compiler recognizes structurally —
not a reducer subclass mechanism. The threshold rides the reducer's
specialization key into the gf.reducer IR as an optional attribute, and
only the generated-CUDA dense streaming provider acts on it, emitting a
dynamic scf.if in the tile loop that skips a low-score tile's payload
entirely. Every other path — eager oracle, native CSR, edge-nn attention,
and all backward passes — fails closed rather than silently substituting
exact softmax, so the approximation is always an explicit choice.
Subclassing tg.Reducer cannot reproduce this: the structured-dense
lowering accepts only the built-in online-softmax structure, and the tile
admission logic is written into the TTIR emitter, not derived from user
algebra. Subclassing OnlineSoftmaxReducer and setting
block_prune_threshold works — but that is the same built-in parameter,
not a new extension point.
class TilePrunedAttention(tg.MessagePassing):
def __init__(self, threshold):
super().__init__()
self.reducer = tg.online_softmax(
block_prune_threshold=threshold,
)
def edge(self, src, dst, edge, scale):
score = (src.key * dst.query).sum(dim=-1) * scale
return self.reducer(score, src.value)
program = TilePrunedAttention(width / nodes)
output = program( # (N, H, D)
graph=tg.Graph.dense(nodes, device="cuda"),
src={"key": key, "value": value},
dst={"query": query},
scale=width**-0.5,
)
Full source: examples/tile_pruned_attention.py (runs as-is)
"""Approximate tile admission as reducer semantics, not a core attention op."""
import torch
import tiga as tg
torch.manual_seed(0)
nodes, heads, width = 1024, 4, 64
# the streaming kernel reads lane-major storage: build (H, N, D), view as (N, H, D)
query = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
key = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
value = torch.randn(heads, nodes, width, device="cuda", dtype=torch.float16).permute(1, 0, 2)
class TilePrunedAttention(tg.MessagePassing):
def __init__(self, threshold):
super().__init__()
self.reducer = tg.online_softmax(
block_prune_threshold=threshold,
)
def edge(self, src, dst, edge, scale):
score = (src.key * dst.query).sum(dim=-1) * scale
return self.reducer(score, src.value)
program = TilePrunedAttention(width / nodes)
output = program( # (N, H, D)
graph=tg.Graph.dense(nodes, device="cuda"),
src={"key": key, "value": value},
dst={"query": query},
scale=width**-0.5,
)
# Inspect: program.explain() provider/lowering/cache summary
# program.ir("domain") Domain IR (block_prune_threshold attr)
# program.ir("kernel_ttir") TTIR (dynamic scf.if tile admission)
Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)
### TilePrunedAttentionTorchExecutor
backend: cuda
provider: triton-nvidia@3.6.0
lowering: gf-kernel-to-ttir-dense-streaming-reduction
passes: gf-verify-domain, gf-lower-domain-to-iter, gf-lower-iter-to-kernel, capture-structured-message, expand-reducer-regions, select-cartesian-coordinate-hierarchy, select-dense-query-neighbor-tile, lower-vector-contraction, lower-streaming-reducer-state, gf-kernel-to-ttir, provider-compile-serialized-ttir
provider cache key: triton-nvidia | cuda:nvidia | 3.6.0 | 2.11.0+cu128 | af81e84448f193cc | cuda:120:warp32
accepted: [machine-schedule] admitted dense-query-key-tile
unknown: [pipeline] no compiler-controlled asynchronous pipeline admitted
remark: [planning] Cartesian relation remained implicit
remark: [planning] multi-value message and reducer state were fused into tiles
remark: [planning] no pairwise score or normalized-weight tensor was materialized
variant cache: hits=0, misses=1
GAT edge attention with an nn score¶
What it is. Graph attention (GAT): instead of a dot product, each edge's attention score is produced by a small neural network reading the edge geometry and both endpoint features; the softmax over incoming edges then weights the source features as before. The score network is where the learnable capacity lives.
The UDF is the first example's shape with one change — tg.nn.trace(attn)
produces the score. On CUDA this compiles to a row-centric fused kernel: one
program per destination row, the score MLP evaluated in registers per edge
chunk, the online-softmax state carried in loop registers — no [E, ·] score
or attention-weight tensor exists. Grad-mode calls take the fused backward
(score-chain replay + softmax Jacobian adjoint in-tile), so training
materializes nothing per edge either. The launch geometry follows the
attention defaults (block_e=16, num_warps=1); tg.nn.trace still accepts
explicit overrides.
class GAT(tg.MessagePassing):
reducer = tg.online_softmax() # attention: softmax over incoming edges
def __init__(self, attn):
super().__init__()
self.attn = tg.nn.trace(attn) # scalar score network, per edge
def edge(self, src, dst, edge):
# s_e = attn([pos_src − pos_dst ‖ x_src ‖ x_dst]) (E, 1)
# out_i = Σ_e softmax_e(s) · x_src(e) (N, F_in)
return self.reducer(self.attn(edge.displacement, src.x, dst.x),
src.x)
graph = tg.Graph.radius(positions, cutoff=0.15) # avg degree ≈ 7
attn = nn.Sequential(
nn.Linear(3 + 2 * f_in, hidden), nn.LeakyReLU(0.2), nn.Linear(hidden, 1),
).to(device)
kernel = GAT(attn)
with torch.no_grad():
output = kernel(graph=graph, src={"x": x}, dst={"x": x}) # (N, F_in)
Full source: examples/gat_edge_attention.py (runs as-is)
"""GAT-style edge attention: an nn scores each edge, online softmax combines.
The score network runs inside a row-centric fused tile kernel — no [E, ·]
score or attention-weight tensor is materialized in either direction.
"""
import tiga as tg
import torch
from torch import nn
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(0)
nodes, f_in, hidden = 512, 8, 16
positions = torch.rand((nodes, 3), device=device) # (N, 3)
x = torch.randn((nodes, f_in), device=device) # (N, F_in)
class GAT(tg.MessagePassing):
reducer = tg.online_softmax() # attention: softmax over incoming edges
def __init__(self, attn):
super().__init__()
self.attn = tg.nn.trace(attn) # scalar score network, per edge
def edge(self, src, dst, edge):
# s_e = attn([pos_src − pos_dst ‖ x_src ‖ x_dst]) (E, 1)
# out_i = Σ_e softmax_e(s) · x_src(e) (N, F_in)
return self.reducer(self.attn(edge.displacement, src.x, dst.x),
src.x)
graph = tg.Graph.radius(positions, cutoff=0.15) # avg degree ≈ 7
attn = nn.Sequential(
nn.Linear(3 + 2 * f_in, hidden), nn.LeakyReLU(0.2), nn.Linear(hidden, 1),
).to(device)
kernel = GAT(attn)
with torch.no_grad():
output = kernel(graph=graph, src={"x": x}, dst={"x": x}) # (N, F_in)
# Training is fused as well: the forward persists only the per-row softmax
# state (m, l); the backward replays the score chain per row chunk and applies
# the softmax Jacobian adjoint in-tile, streaming gradients to the fields,
# the positions and the score-network weights.
positions_g = positions.clone().requires_grad_(True)
xg = x.clone().requires_grad_(True)
graph_g = tg.Graph.radius(positions_g, cutoff=0.15)
out_g = kernel(graph=graph_g, src={"x": xg}, dst={"x": xg})
out_g.sum().backward()
# Inspect: kernel.explain(); kernel.last_variant.artifacts["ttir"]
Measured compilation artifacts (Ryzen 7 255 · RTX 5070 Ti)
### GATTorchExecutor
backend: cuda
provider: triton-nvidia@3.6.0
lowering: gf-python-emit-edge-nn-attention-vjp
passes: capture-edge-nn-subgraph, translate-fx-to-message-dag, prove-edge-nn-attention-structure, python-emit-edge-nn-attention-ttir, symbolic-edge-nn-attention-vjp, python-emit-edge-nn-attention-vjp-ttir, provider-compile-serialized-ttir
provider cache key: triton-nvidia | cuda:nvidia | 3.6.0 | 2.11.0+cu128 | af81e84448f193cc | cuda:120:warp32
remark: [planning] training path: forward persists only the O(N) softmax state (m, l); backward replays the score chain per row chunk — no [E, ·] tensor in either direction
variant cache: hits=0, misses=2