Wall Attention (GitHub Repo)｜Wall Attention（GitHub リポジトリ）

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](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](https://github.com/fla-org/flash-linear-attention) (MIT) から派生した並列アテンション機構に基づいています。効率的なアテンションに関する優れた研究に尽力された FLA チームに感謝いたします。 ## ライセンス MIT、詳細は [LICENSE](https://github.com/tilde-research/wall-attention-release/blob/main/LICENSE) を参照してください。

Wall Attention（GitHub リポジトリ）

背景や根拠まで確認しますか？

調べる

選ぶ

サイト