Wall Attention（GitHub リポジトリ）
本文の状態
日本語全文を表示中
詳細モードで約5分の本文を読めます。
同じ出来事の情報源
この情報源を基点に整理
TLDR AI
開発者が公開した「Wall Attention」という技術の GitHub リポジトリが紹介された。具体的な機能や変更点は本文抜粋から読み取れないため、詳細は不明である。
Source Article
元記事を日本語で読む
本文に関係しない購読案内、埋め込み通知、サイト内プロモーションは除いています。
Wall Attention は、QK 内積に組み込まれた「チャネルごとの、タイムステップごとの乗算減衰」を特徴とするアテンションのバリアントです。標準的なアテンションがペア (i, j) を ∑_n q_{i,n} k_{j,n} でスコアリングするのに対し、Wall Attention は各チャネル n に対して、2 つの位置間で蓄積された学習済みの減衰で重み付けを行います。これにより、各クエリチャネルが独立した、コンテンツ依存性の忘却率を持つようになり、スカラーゲート（FoX）や RoPE スタイルの減衰を全チャネル次元に一般化します。g = 0 と設定すると、バニラのソフトマックスアテンションが復元されます。
詳細はブログをご覧ください：
https://blog.tilderesearch.com/blog/wall-attn
このリポジトリでは、実務で用いられる 2 つのカーネルをそれぞれ個別にパッケージ化しています。
- トレーニング / プレフィル (wall_attn): q, k, v, g に対する解析的勾配を持つ融合フォワード + バックワード Triton カーネル（FlashAttention スタイルのストリーミングソフトマックス）
- デコード (wall_attn_decode): 事前再スケーリングされた KV キャッシュを読み取る単一ステップカーネル。これにより、トークンごとの生成コストはプレフィックスの再計算ではなく、小さな GEMV 類似パス 1 つで済みます。
インストール
uv を使用（推奨）
uv sync
source .venv/bin/activate
または pip で
pip install -e .
使用方法
トレーニング / プレフィル
import torch
from wall_attn import wall_attn
B, T, H, HQ, K, V = 2, 1024, 4, 8, 64, 64 # GQA: HQ クエリヘッド、H KV ヘッド
q = torch.randn(B, T, HQ, K, device="cuda", dtype=torch.bfloat16, requires_grad=True)
k = torch.randn(B, T, H, K, device="cuda", dtype=torch.bfloat16, requires_grad=True)
v = torch.randn(B, T, H, V, device="cuda", dtype=torch.bfloat16, requires_grad=True)
g = torch.randn(B, T, HQ, K, device="cuda", dtype=torch.bfloat16, requires_grad=True) * 0.02
o = wall_attn(q, k, v, g, scale=K**-0.5) # [B, T, HQ, V]
o.sum().backward()
オプション引数: g_scalar ([B, T, HQ] FoX スタイルの加算ゲート)、sink_bias ([HQ] アテンションシンク)、window_size (スライディングウィンドウ)、cu_seqlens (可変長パッキング、B == 1 を必要とする)。
Decode (キャッシュ生成)
プリフェッチ時に再スケール済みのキャッシュを一度構築し、その後トークンを一つずつデコードする:
import torch
from fla.ops.utils.constant import RCP_LN2
from fla.ops.utils.cumsum import chunk_global_cumsum
from wall_attn import build_wall_kv_cache, wall_attn_decode
C = 64 # キャッシュチャンクサイズ (アンカー粒度)
P = chunk_global_cumsum(g, scale=RCP_LN2) # [B, T, HQ, K] プレフィックス
k_tilde, r_cache = build_wall_kv_cache(k, P, chunk_size=C)
o, _ = wall_attn_decode(
q=q[:, -1:], # 現在のクエリ [B, 1, HQ, K]
v=v, # キャッシュされた値 [B, T_kv, H, V]
p_curr=P[:, -1:], # 現在の行におけるプレフィックス
k_tilde=k_tilde, # 再スケーリング前のキー [B, T_kv, HQ, K]
r_cache=r_cache, # チャンクごとのアンカー [B, ceil(T_kv/C), HQ, K]
sink_bias=None,
scale=K**-0.5,
cache_chunk_size=C,
)
build_wall_kv_cache は、チャンクごとのアンカー R_c を用いて減衰をキーに折りたたみます (k_tilde[j] = k[j] · exp2(R_c − P[j]))。これにより、デコードカーネルはプレフィックスの再累積を行う必要がなくなります。逐次追加によるサービングループの詳細については、tests/test_decode.py::test_decode_streaming_matches_full_forward を参照してください。
コード構造
wall_attn/
├── __init__.py # パブリック API
├── training.py # フォワード/バックワード Triton カーネル + autograd Function + wall_attn()
├── decode.py # 単一ステップデコードカーネル + build_wall_kv_cache()
└── reference.py # イーガーな PyTorch リファレンス (正しさのオラクル)
tests/
├── test_training.py # パリティ + 解析的勾配 (有限差分チェック付き)
└── test_decode.py # デコード == プレフィルフォワード、ストリーミング、キャッシュ形状
機能
- GQA: クエリヘッド HQ は KV ヘッド H を超える可能性があります (HQ % H == 0)。
- 各チャネルごとの減衰 g と正確な解析的勾配に加え、オプションのスカラーゲート g_scalar をサポートします。
- アテンションシンク (sink_bias)、スライディングウィンドウ (window_size)、および可変長パッキング (cu_seqlens) をサポートします。
- 安価な自己回帰生成のための事前再スケーリングされたデコードキャッシュ、数値的に安定した長文コンテキスト（チャンクごとのアンカーにより exp2 が有界に保たれる）。
- BF16/FP32 入力；Hopper / Ampere アーキテクチャ向けに自動調整されたブロックサイズ。
テスト
pytest # CUDA GPU が必要
すべてのカーネルパスは、イージモードの wall_attn_reference と比較され、g および g_scalar の勾配は中心有限差分法に対して検証される。デコードカーネルは、ストリーミング生成ループを含むトレーニングの順方向をトークン単位で再現するかどうかも確認される。
謝辞
Triton カーネルは、flash-linear-attention (MIT) から派生した並列アテンション機構に基づいています。効率的なアテンションに関する優れた研究に尽力された FLA チームに感謝いたします。
ライセンス
MIT、詳細は LICENSE を参照してください。
原文を表示
Wall Attention is an attention variant with a per-channel, per-timestep multiplicative decay baked into the QK inner product. Where standard attention scores a pair
(
i
,
j
)
with
∑
n
q
i
,
n
,
k
j
,
n
, Wall Attention weights each channel
n
by a learned decay accumulated between the two positions. This gives each query channel an independent, content-dependent forgetting rate, generalizing scalar gating (FoX) and RoPE-style decays to the full channel dimension. Setting
g
=
0
recovers vanilla softmax attention.
See the blog for more information:
https://blog.tilderesearch.com/blog/wall-attn
This repo packages the two kernels used in practice, each on its own:
- Training / prefill (wall_attn): a fused forward + backward Triton kernel (FlashAttention-style streaming softmax) with analytic gradients for
q
,
k
,
v
,
g
.
- Decode (wall_attn_decode): a single-step kernel that reads a pre-rescaled KV cache, so per-token generation costs one small GEMV-like pass instead of recomputing the prefix.
Installation
# Using uv (recommended)
uv sync
source .venv/bin/activate
# or with pip
pip install -e .Usage
Training / prefill
import torch
from wall_attn import wall_attn
B, T, H, HQ, K, V = 2, 1024, 4, 8, 64, 64 # GQA: HQ query heads, H kv heads
q = torch.randn(B, T, HQ, K, device="cuda", dtype=torch.bfloat16, requires_grad=True)
k = torch.randn(B, T, H, K, device="cuda", dtype=torch.bfloat16, requires_grad=True)
v = torch.randn(B, T, H, V, device="cuda", dtype=torch.bfloat16, requires_grad=True)
g = torch.randn(B, T, HQ, K, device="cuda", dtype=torch.bfloat16, requires_grad=True) * 0.02
o = wall_attn(q, k, v, g, scale=K**-0.5) # [B, T, HQ, V]
o.sum().backward()Optional arguments: g_scalar ([B, T, HQ] FoX-style additive gate), sink_bias ([HQ] attention sink), window_size (sliding window), and cu_seqlens (varlen packing, requires B == 1).
Decode (cached generation)
Build the pre-rescaled cache once at prefill, then decode one token at a time:
import torch
from fla.ops.utils.constant import RCP_LN2
from fla.ops.utils.cumsum import chunk_global_cumsum
from wall_attn import build_wall_kv_cache, wall_attn_decode
C = 64 # cache chunk size (anchor granularity)
P = chunk_global_cumsum(g, scale=RCP_LN2) # [B, T, HQ, K] prefix
k_tilde, r_cache = build_wall_kv_cache(k, P, chunk_size=C)
o, _ = wall_attn_decode(
q=q[:, -1:], # current query [B, 1, HQ, K]
v=v, # cached values [B, T_kv, H, V]
p_curr=P[:, -1:], # prefix at the current row
k_tilde=k_tilde, # pre-rescaled keys [B, T_kv, HQ, K]
r_cache=r_cache, # per-chunk anchors [B, ceil(T_kv/C), HQ, K]
sink_bias=None,
scale=K**-0.5,
cache_chunk_size=C,
)build_wall_kv_cache folds the decay into the keys (k_tilde[j] = k[j] · exp2(R_c − P[j])) using a per-chunk anchor R_c, so the decode kernel never re-accumulates the prefix. See tests/test_decode.py::test_decode_streaming_matches_full_forward for the full append-as-you-go serving loop.
Code structure
wall_attn/
├── __init__.py # public API
├── training.py # forward/backward Triton kernels + autograd Function + wall_attn()
├── decode.py # single-step decode kernel + build_wall_kv_cache()
└── reference.py # eager PyTorch reference (correctness oracle)
tests/
├── test_training.py # parity + analytic gradients (finite-difference checked)
└── test_decode.py # decode == prefill forward, streaming, cache shapes
Features
- GQA: query heads HQ may exceed kv heads H (HQ % H == 0).
- Per-channel decay g with exact analytic gradient, plus an optional scalar gate g_scalar.
- Attention sink (sink_bias), sliding window (window_size), and varlen packing (cu_seqlens).
- Pre-rescaled decode cache for cheap autoregressive generation, numerically stable to long context (per-chunk anchors keep exp2 bounded).
- BF16/FP32 inputs; autotuned block sizes for Hopper / Ampere.
Testing
pytest # requires a CUDA GPUEvery kernel path is checked against the eager wall_attn_reference, and the g / g_scalar gradients are verified against central finite differences. The decode kernel is checked to reproduce the training forward token-for-token, including a streaming generation loop.
Acknowledgments
The Triton kernels build on the parallel-attention machinery from flash-linear-attention (MIT). We thank the FLA team for their excellent work on efficient attention.
License
MIT, see LICENSE.
今日のまとめ
AIデイリーブリーフで今日の重要ニュースをまとめ読み