注意力与稠密关系¶
前四个示例需要 NVIDIA GPU 和 cuda extra;GAT 根据设备可用性选择 CPU/CUDA。
核心片段不包含全部初始化,运行命令对应完整脚本。Torch 需在 Tiga 之前单独安装。
用通用的关系(relation)与 reducer 语义来表达注意力类工作负载。
掩码(mask)从来不是代码——它就是图本身:稠密关系是全量注意力,三角关系是
因果注意力,而块对角 CSR 则是带 cu_seqlens 的 varlen 因果注意力。tile 剪枝
是叠加在同一个 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
全量注意力,无掩码¶
它是什么。 注意力:N 个位置中每个位置都携带一个 query 向量、一个 key 向量 和一个宽度为 D 的 value 向量。每个输出位置是所有 value 向量的加权平均,权重 是 query·key 相似度分数上的 softmax——因此代价是完整的 N×N 分数矩阵。没有 掩码时,每个 query 都会关注每个 key。用图的术语来说,每条边 (j→i) 把目的位置 i(query)与源位置 j(key 和 value)配对;稠密关系包含全部 N×N 条边。
这是下面所有内容的精确特例——与 FlashAttention 相同的分块 online softmax 机制,只是每个 tile 都被接纳。
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,
)
完整源码:examples/full_attention.py(可直接运行)
"""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)
本机编译产物(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
因果稠密关系¶
它是什么。 因果注意力是带“禁止向前看”约束的注意力:位置 i 只能对 j ≤ i 的位置求平均。这是从左到右生成文本的语言模型的规则,其中每个位置 只能读取过去。该示例把掩码表达为图本身:三角关系恰好包含所有满足 j ≤ i 的边 (j→i),因此边函数完全不需要任何掩码代码。同上,源 j 提供 key 和 value,目的 i 提供 query。
Graph.triangular() 把下三角(含对角)拓扑一路携带穿过
Domain → Iter → Kernel,普通的 online reducer 经 lowering 得到因果流式 TTIR。
Torch 只作为外部 CUDA 存储出现。
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,
)
完整源码:examples/causal_dense_relation.py(可直接运行)
"""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()
本机编译产物(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 因果注意力¶
它是什么。 真实的训练批次会把许多长度各异的序列打包进一个扁平张量,
并用累积长度 cu_seqlens = [0, L₁, L₁+L₂, …] 标记边界(即 flash-attn 的
varlen API 约定)。每个序列仅在自身内部做因果注意力——位置 i 只有当
j 与 i 位于同一序列且 j ≤ i 时才能读取位置 j。作为掩码,它是一个由小三角
组成的块对角矩阵;作为图,它是一个恰好包含这些边的显式 CSR 关系,共
Σ L(L+1)/2 条。边界处理无需任何代码:跨序列注意力就是一条不存在的边。
该示例用 tg.Graph.cat 组合每个序列的三角图,运行同一个无掩码 edge UDF,
包括长度为 1 的序列。块 k 只在内部互相注意。当前 cat 会物化 CSR,
尚非隐式组合关系。兼容构造器 tg.Graph.cu_seqlens(boundaries, causal=True)
在累计边界以零开头时产生相同关系。(测试会将全部三个注意力示例与
F.scaled_dot_product_attention 的结果逐一比对。)目前执行会回退到精确的
eager oracle:编译后的流式路径覆盖上面的稠密/三角关系以及基于字段分数的 CSR
段 softmax,而在通用 CSR 上由边计算分数属于编译器边界,而非语义边界。
# 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,
)
完整源码:examples/varlen_causal_attention.py(可直接运行)
"""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()
本机编译产物(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 剪枝的稀疏注意力¶
为什么要跳过一块¶
这是一个用精度换取潜在性能的实验性选项,不是使用 attention 的必需步骤。 普通 attention 对每个允许的 key 计算 score,再用 softmax 权重累加 value。 当某一整块 key 的 score 都很低时,继续加载该块 value 并累加,可能贡献很小。 tile 是一小块 query 与一小块 key 的配对;批量跳过是为了适合 GPU 的分块执行。
例如,已见最高 score 为 10,下一块最高 score 只有 1,两者指数权重的比例约为:
若 block_prune_threshold=0.01,这块会被跳过。这个比较针对整个 query block,
不是每个 query 各自的误差保证;很多低分项相加,或 value 很大时,仍可能造成明显误差。
| 仍然执行 | 被剪 tile 跳过 |
|---|---|
| 读取 key、计算 query/key score、判断阈值 | 读取 value、后续 softmax 与 value 累加 |
因此它不减少全部 score 计算,也不保证加速。阈值越大越激进;
精确计算应省略 block_prune_threshold,不是把它设为 0。
当前仅支持指定 generated-CUDA dense 前向路径,不支持 backward。
示例中的 width / nodes 只是演示参数,不是通用调参公式。
具体规则与边界¶
它是什么。 与第一个示例相同的注意力,但是近似的:tile 剪枝把 key 按块 (tile)处理,跳过最高分数远低于当前 query block 最大值的 tile。阈值本身 不保证被省略的总概率或输出误差上界;这些还取决于被剪 tile 的数量与 value 的幅度,需要按工作负载验证精度。FlashAttention 正是 每个 tile 都被接纳的精确特例;本示例添加了接纳规则。
图中对角附近被保留仅是示意,不是固定局部窗口;实际保留哪些块由数据的 score 决定。
emitter 将整个 query/key tile 的最大 score 与当前 query block 的滚动
最大值比较。以自然对数 score 为单位,准入规则使用比例
rho = block_prune_threshold:
TTIR 等价地使用以 2 为底的 score 和 log2(rho)。key 加载和 score 计算发生
在准入判断之前;剪枝跳过的是 value 加载及后续 softmax/value 累加。
阈值的工作机制:这是内置 tg.online_softmax() reducer 的一个构造参数,
由编译器按结构识别——不是 subclass reducer 的扩展机制。阈值随 reducer 的
特化键进入 gf.reducer IR,成为一个可选属性;只有 generated-CUDA dense
streaming provider 会消费该属性,在 tile 主循环里生成一个动态的 scf.if,把
低分 tile 的载荷整个跳过。其余所有路径——eager oracle、原生 CSR、edge-nn
attention、以及全部反向 pass——都 fail closed,而不是静默退回精确 softmax,
因此这个近似永远是显式选择。
subclass tg.Reducer 无法复刻这个机制:structured-dense lowering 只接受内置
online-softmax 结构,tile 接纳逻辑写在 TTIR emitter 里,而不是从用户代数
推导出来的。subclass OnlineSoftmaxReducer 再设置 block_prune_threshold
是可行的——但那用的还是同一个内置参数,不是新的扩展点。
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,
)
完整源码:examples/tile_pruned_attention.py(可直接运行)
"""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)
本机编译产物(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
用 nn 打分的 GAT 边注意力¶
它是什么。 图注意力(GAT):不用点积,每条边的注意力分数由一个小型 神经网络产生,该网络读取边几何信息和两个端点的特征;随后对入边的 softmax 像之前一样对源特征加权。可学习的参数就集中在这个分数网络中。
UDF 沿用第一个示例的结构,只有一处变化——tg.nn.trace(attn) 产生分数。
在 CUDA 上编译为一个以行为中心的融合 kernel:每个目的行一个程序,分数 MLP
在寄存器中按边块求值,online softmax 状态保存在循环寄存器中——不存在任何
[E, ·] 形状的分数或注意力权重张量。梯度模式的调用走融合的反向路径(分数链
重放 + tile 内 softmax Jacobian 伴随),因此训练时同样不会按边物化任何数据。
launch geometry 采用注意力示例的默认值(block_e=16, num_warps=1);
tg.nn.trace 仍接受显式覆盖值。
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)
完整源码:examples/gat_edge_attention.py(可直接运行)
"""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"]
本机编译产物(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