# アテンションのバックエンド

すべてのアテンション実装は同じ計算を実行します。各トークンは他のすべてのトークンと比較されます。違いは、計算が*どのように*実行されるかという点です。基本的なアテンションは、メモリ上に完全なアテンション行列をマテリアライズするため、スケーリングが不良であり、推論を遅くするボトルネックを生み出します。最適化された実装は、計算を再構成してメモリトラフィックを削減し、より高速で低コストな推論を実現します。

[AttentionInterface](/docs/transformers/main/ja/internal/modeling_utils#transformers.AttentionInterface) は、最適化されたアテンションの実装を提供します。アテンションの実装とモデルの実装を分離することで、異なる関数を用いた実験を簡素化します。この一貫したインターフェースにより、新しいバックエンドを簡単に追加できます。

Attention バックエンドDescription "flash_attention_3"は、FlashAttention-2 を改善し、操作のオーバーラップと、フォワードパスとバックワードパスのより密接なフュージョンを実現します "flash_attention_2"は、計算をより小さなブロックに分割し、高速なオンチップメモリを使用します "flex_attention"は、低レベルのカーネルを手作業で書くことなく、カスタムのアテンションパターン（スパース、ブロックローカル、スライディングウィンドウ）を指定するためのフレームワークです "sdpa"は、スケーリングドットプロダクトアテンションの PyTorch 組み込み実装です "paged|flash_attention_3"は、FlashAttention-3 のページング版です "paged|flash_attention_2"は、FlashAttention-2 のページング版です "paged|sdpa"は、SDPA のページング版です "paged|eager"は、eager のページング版です

## アテンションバックエンドを設定する

`attn_implementation` 引数を [from_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.from_pretrained) に渡すことで、特定の注意機構関数を用いたモデルをインスタンス化できます。

```py
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="flash_attention_2"
)
```

モデルを再読み込みせずに、[set_attn_implementation()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.set_attn_implementation) を使用してランタイムでアテンションのバックエンドを切り替えます。

```py
model.set_attn_implementation("sdpa")
```

### カーネル

[Kernels](https://huggingface.co/docs/kernels/index) ライブラリを使用すると、実行時にコンパイル済みの計算カーネルを [Hub](https://huggingface.co/models?other=kernels) から直接ダウンロードして読み込むことができます。これにより、PyTorch や CUDA のバージョン不一致によるパッケージングの問題を回避できます。

Kernels は検出時に自動的に [AttentionInterface](/docs/transformers/main/ja/internal/modeling_utils#transformers.AttentionInterface) に登録されます。FlashAttention パッケージを明示的にインストールする必要はありません。FlashAttention を名前で指定した場合も Hub kernel にフォールバックします。詳細は [FlashAttention fallback](./kernel_doc/loading_kernels#flashattention-fallback) を参照してください。

```py
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="kernels-community/flash-attn2"
)
```

### SDPAコンテキストマネージャー

PyTorchのスケーリングドットプロダクトアテンション（SDPA）は、CUDAバックエンドに対して最速のアテンション関数を自動的に選択します。その他のバックエンドでは、デフォルトでPyTorchのC++実装を使用します。

[torch.nn.attention.sdpa_kernel](https://pytorch.org/docs/stable/generated/torch.nn.attention.sdpa_kernel.html) コンテキストマネージャーを使用して、SDPAに特定の実装を使用させることができます。

```py
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="sdpa"
)

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    outputs = model.generate(**inputs)
```

## バックボーン固有の注意

マルチモーダルモデルでは、各モダリティごとに異なるバックボーンを使用します。各バックボーンに特定の注意機構を割り当てることで、パフォーマンスを最適化できます。例えば、一部のビジョンバックボーンはfp32でより良いパフォーマンスを発揮しますが、FlashAttentionはこれをサポートしていません。

dict を使用してビジョンのバックボーンを異なるアテンション関数にマッピングし、テキストのバックボーンは引き続き FlashAttention を使用します。アテンション実装のキーは、サブ設定の名前と一致する必要があります。

```py
from transformers import AutoModelForImageTextToText

attention_implementation_per_backbone = {"vision_config": "sdpa", "text_config": "flash_attention_2"}

for key in attention_implementation_per_backbone:
    assert key in model.config.sub_configs, f"Invalid key in `attention_implementation`"

model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation=attention_implementation_per_backbone
)
```

dictから特定のbackboneを省略すると、デフォルトのattention関数（SDPA）が使用されます。

```py
model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation={"text_config": "flash_attention_2"}
)
```

単一の文字列で、すべてのバックボーンに同じアテンション関数を設定します。

```py
model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation="eager"
)
```

空のキーでアテンション関数をグローバルに設定します。

```py
model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation={"": "eager"}
)
```

## 新しいアテンション関数を作成する

[AttentionInterface.register()](/docs/transformers/main/ja/internal/modeling_utils#transformers.AttentionInterface.register) を使用してアテンションレジストリに追加することで、アテンション関数をカスタマイズしたり、新しいアテンション関数を作成したりできます。モデルは `attn_implementation` 引数を通じてこれらの関数を使用します。

> [!WARNING]  
> カスタムアテンション関数を登録する際は、対応するアテンションマスク関数も登録してください。カスタムの`attn_implementation`名が[AttentionMaskInterface](/docs/transformers/main/ja/internal/modeling_utils#transformers.AttentionMaskInterface)に登録されていない場合、Transformersはマスクの作成をスキップし、アテンションレイヤーに`attention_mask=None`を渡します。そのため、アテンション関数内で因果性、パディング、パッキング、またはスライディングウィンドウの制約を自分で処理する必要があります。そうしないと、これらの制約が黙って無視されてしまいます。

この例では、アテンション関数をカスタマイズして、各レイヤーごとにステートメントを出力します。元の実装のマスクを保持するために、`masking_utils.sdpa_mask`をアテンションマスク関数として登録しています。

```python
import torch
from transformers import AutoModelForCausalLM, AttentionInterface, AttentionMaskInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import sdpa_mask

def my_new_sdpa(*args, **kwargs):
    print("I just entered the attention computation")
    return sdpa_attention_forward(*args, **kwargs)

AttentionInterface.register("my_new_sdpa", my_new_sdpa)
AttentionMaskInterface.register("my_new_sdpa", sdpa_mask)  # must have the same name as the registered attention function

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", attn_implementation="my_new_sdpa")
model(torch.ones(1, 5, dtype=int))
```

アテンション関数に新しい引数を追加することもできます。[AttentionInterface](/docs/transformers/main/ja/internal/modeling_utils#transformers.AttentionInterface)をサポートするモデルは、kwargsをアテンション層およびアテンション関数に伝播します。モデルのforward関数で引数をkwargsとして渡してください。カスタムアテンション関数は、このシグネチャと返却形式に従う必要があります。

```python
import torch
from transformers import AutoModelForCausalLM, AttentionInterface, AttentionMaskInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import sdpa_mask

def custom_attention(
    module: torch.nn.Module,  # required arg
    query: torch.Tensor,  # required arg
    key: torch.Tensor,  # required arg
    value: torch.Tensor,  # required arg
    attention_mask: Optional[torch.Tensor],  # required arg
    a_new_kwargs = None,  # You can now add as many kwargs as you need
    another_new_kwargs = None,  # You can now add as many kwargs as you need
    **kwargs,  # You need to accept **kwargs as models will pass other args
) -> tuple[torch.Tensor, Optional[torch.Tensor]]
    ...  # do your magic!
    return attn_output, attn_weights  # attn_weights are optional here

AttentionInterface.register("custom", custom_attention)
AttentionMaskInterface.register("custom", sdpa_mask)  # to leave the existing mask untouched

model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="custom")
model(torch.ones(1, 5, dtype=int), a_new_kwargs=..., another_new_kwargs=...)
```

モデルの[modeling code](https://github.com/huggingface/transformers/tree/main/src/transformers/models)を確認し、attention関数に渡される引数とkwargsを特定してください。

## AttentionMaskInterface

[AttentionMaskInterface](/docs/transformers/main/ja/internal/modeling_utils#transformers.AttentionMaskInterface) は、[`create_*_mask`](#build-an-attention-mask) 関数が参照して、マスクをアクティブなアテンションバックエンドが期待する形式に変換するためのレジストリです。FlexAttention には [BlockMask](https://docs.pytorch.org/docs/stable/nn.attention.flex_attention.html#torch.nn.attention.flex_attention.BlockMask) が必要で、SDPA には 4 次元テンソル、FlashAttention にはベースとなる 2 次元パディングマスクが必要です。[AttentionMaskInterface.register()](/docs/transformers/main/ja/internal/modeling_utils#transformers.AttentionInterface.register) を使用して、カスタムバックエンドを登録したり、既存のバックエンドのフォーマッターをオーバーライドしたりできます。

```python
import torch
from transformers import AttentionMaskInterface
from transformers.masking_utils import sdpa_mask

def my_new_sdpa_mask(*args, **kwargs):
    print("I just entered the attention mask computation")
    return sdpa_mask(*args, **kwargs)

AttentionMaskInterface.register("my_new_sdpa_mask", my_new_sdpa_mask)
```

アクティブな`attn_implementation`に対応するフォーマッタが登録されていない場合、マスクの作成はスキップされ、`attention_mask=None`がアテンション層に渡されます。

登録された関数は、このシグネチャと一致する必要があります。

```python
def custom_attention_mask(
    batch_size: int,  # required arg
    q_length: int,  # required arg
    kv_length: int,  # required arg
    q_offset: int = 0,  # required arg
    kv_offset: int = 0,  # required arg
    mask_function: Callable = causal_mask_function,  # required arg
    attention_mask: Optional[torch.Tensor] = None,  # required arg
    **kwargs,  # a few additional args may be passed as kwargs, especially the model's config is always passed
) -> Optional[torch.Tensor]:
```

`mask_function` 引数は、PyTorch の [mask_mod](https://pytorch.org/blog/flexattention/) 関数を模した `Callable` です。4 つのインデックス `(batch_idx, head_idx, q_idx, kv_idx)` を受け取り、その位置がアテンション計算に寄与するかどうかを示すブール値を返します。これは、[アテンションマスクの構築](#build-an-attention-mask) で `or_mask_function` および `and_mask_function` に使用されているのと同じ基本形です。

> [!TIP]
> `mask_function`がマスクの作成に失敗する場合、torch.exportにはこの[ワークアラウンド](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/executorch.py)を使用してください。

## アテンションマスクの構築

`create_*_mask` 関数群を使用してアテンションマスクを構築します。これらの関数は [transformers.masking_utils](https://github.com/huggingface/transformers/blob/main/src/transformers/masking_utils.py#L894) に含まれており、モデル設定からアクティブなアテンションバックエンドを読み取り、[AttentionMaskInterface](/docs/transformers/main/ja/internal/modeling_utils#transformers.AttentionMaskInterface) でそのバックエンドのマスクフォーマッタを検索し、バックエンドが期待する形式を返します。マスクの反転、拡張、キャストを自分で行う必要はありません。

アテンションパターンに一致する関数を選択します。

| 関数 | ユースケース |
|---|---|
| `create_causal_mask` | 各トークンが自分自身とそれ以前のトークンにアテンションを向けるdecoder-onlyモデル |
| `create_bidirectional_mask` | エンコーダーモデル、またはデコーダーからエンコーダー状態へのクロスアテンション |
| `create_sliding_window_causal_mask` | スライディングウィンドウアテンションパターンを持つデコーダーモデル |
| `create_chunked_causal_mask` | シーケンスを固定サイズのブロックに分割するデコーダーモデル |
| `create_bidirectional_sliding_window_mask` | スライディングウィンドウアテンションパターンを持つエンコーダーモデル |

> [!WARNING]
> レガシーの呼び出し可能なマスクヘルパー - `get_extended_attention_mask`、`create_extended_attention_mask_for_decoder`、`invert_attention_mask` - は非推奨の警告を出力し、将来のリリースで削除されます。代わりに`create_*_mask`関数を使用してください。

デコーダーのフォワードパス内で`create_causal_mask`を呼び出します。設定、入力埋め込み、ユーザーが指定した2次元の`attention_mask`、およびキャッシュを引数として渡します。この関数は、埋め込みからバッチサイズ、クエリ長、dtype、デバイスを読み取り、キャッシュからキー長を計算します。

```py
from transformers.masking_utils import create_causal_mask

attention_mask = create_causal_mask(
    config=self.config,
    inputs_embeds=inputs_embeds,
    attention_mask=attention_mask,
    past_key_values=past_key_values,
)
```

エンコーダーの自己注意機構には`create_bidirectional_mask`を呼び出します。エンコーダーはキャッシュを行わないため、`past_key_values`は省略します。

```py
from transformers.masking_utils import create_bidirectional_mask

attention_mask = create_bidirectional_mask(
    config=self.config,
    inputs_embeds=embedding_output,
    attention_mask=attention_mask,
)
```

クロスアテンションでは、`encoder_hidden_states` にエンコーダーの状態を渡すことで、マスクがデコーダーのクエリ長ではなく、エンコーダーのキーとバリューの長さを参照するようにします。

```py
encoder_attention_mask = create_bidirectional_mask(
    config=self.config,
    inputs_embeds=embedding_output,
    attention_mask=encoder_attention_mask,
    encoder_hidden_states=encoder_hidden_states,
)
```

`or_mask_function` と `and_mask_function` 引数を使用して、ベースマスクの上に追加の制約を付けます。`or_mask_function` を使用すると追加の位置がアテンションできるようになり、`and_mask_function` を使用するとベースパターンをさらに制限できます。どちらも、[AttentionMaskInterface](#attentionmaskinterface) で説明されている4インデックスの `mask_function` シグネチャに従います。`(batch_idx, head_idx, q_idx, kv_idx)` を受け取り、ブール値を返します。

> [!WARNING]
> `or_mask_function` と `and_mask_function` は任意の注意パターンを表現できますが、組み込みのパターンよりも遅く、ExecuTorch との互換性はありません。オーバーヘッドは、マスク生成がフォワードパスの時間に占める割合が大きくなる小規模モデル（約 2 億パラメータ）で特に顕著です。標準的な `create_*_mask` 関数では必要な表現ができない場合のみ、これらを使用してください。

例えば、どこでも`True`を返す関数をオーバーレイして因果マスクを完全に双方向なものに変換できます。因果パターンとの和集合により、すべてのトークンが他のすべてのトークンに注意を払うようになります。

```py
mask_kwargs = {
    "config": self.config,
    "inputs_embeds": inputs_embeds,
    "attention_mask": attention_mask,
    "past_key_values": past_key_values,
    "position_ids": position_ids,
    "or_mask_function": lambda *args: torch.tensor(True, dtype=torch.bool),
}

attention_mask = create_causal_mask(**mask_kwargs)
```

生成時、[generate()](/docs/transformers/main/ja/main_classes/text_generation#transformers.GenerationMixin.generate) は `create_masks_for_generate` を介してマスクを構築します。この関数は、モデルの構成に基づいて適切な `create_*_mask` にディスパッチします。モデルクラスでこれをオーバーライドすることで、生成用のカスタムなマスキング戦略を適用できます。

## カスタムの4Dアテンションマスクを渡す

Pass your own 4D mask when you need an attention pattern the `create_*_mask` functions can't express. A 4D mask has the shape `(batch_size, 1, query_length, kv_length)`, where `1` broadcasts the same mask across every attention head. Transformers detects and uses the mask as-is, skipping `create_*_mask`.

4Dマスクは、2つの値の慣用法のいずれかを使用します。

| dtype | attend | マスクアウト |
|---|---|---|
| boolean | `True` | `False` |
| float | `0.0` | `-inf` |

float 規約では、softmax の前にマスクをアテンションスコアに加算します。スコアに `0.0` を加えても値は変化しないため、その位置は貢献します。一方、スコアに `-inf` を加えると softmax 後にゼロになるため、その位置は除外されます。

> [!IMPORTANT]
> 受け入れられる規約はアテンションバックエンドによって異なります。`sdpa` はブール値または浮動小数点マスクを受け取ります。`eager` はマスクをスコアに加算するため、浮動小数点マスクのみを受け取ります。`flash_attention_2` と `flex_attention` は独自の形式（2次元のパディングマスクと [BlockMask](https://docs.pytorch.org/docs/stable/nn.attention.flex_attention.html#torch.nn.attention.flex_attention.BlockMask)）を使用し、生の4次元マスクは受け入れません。

よくある間違いとして、2次元のpaddingマスクで使われる`1`/`0`の慣習を、float型の4次元マスクにそのまま適用してしまうことです。マスクはスコアに足し算されるため、`0.0`は位置を保持し、`1.0`はわずかなバイアスを加えるだけです。

以下の例では、誤ったマスクと正しいマスクを比較しています。どちらも同じ`1`/`0`の因果パターンから始まります。

```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", attn_implementation="sdpa")
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")

input_ids = tokenizer("my favorite condiment on a", return_tensors="pt").input_ids
seq_len = input_ids.shape[1]

# 1 attends, 0 masks
causal = torch.tril(torch.ones(seq_len, seq_len))

# wrong: 1.0/0.0 floats are added to the scores, so 0.0 keeps a token and 1.0 barely changes it
wrong_mask = causal[None, None]

# correct: 0.0 attends, -inf masks
correct_mask = torch.where(causal.bool(), 0.0, float("-inf"))[None, None]
```

`0.0`がマスクに使用される値であるため、誤ったマスクでは何も除外されず、すべての位置が保持されます。

```text
        wrong_mask                          correct_mask
   (1 attends, 0 masks)              (0 attends, -inf masks)

      k0 k1 k2 k3 k4                     k0   k1   k2   k3   k4
   q0  1  0  0  0  0                  q0  0  -inf -inf -inf -inf
   q1  1  1  0  0  0                  q1  0   0   -inf -inf -inf
   q2  1  1  1  0  0                  q2  0   0    0   -inf -inf
   q3  1  1  1  1  0                  q3  0   0    0    0   -inf
   q4  1  1  1  1  1                  q4  0   0    0    0    0
```

## 双方向アテンション

Decoder-onlyモデルはデフォルトでcausal（単方向）アテンションを使用し、各トークンは自分自身とそれ以前のトークンのみにアテンションを向けます。`is_causal=False`を設定すると、すべてのトークンが他のすべてのトークンにアテンションを向ける双方向アテンションに切り替わります。これにより、Decoder-onlyモデルをテキストエンコーダーとして使用できるようになり、例えば埋め込みの生成が可能になります。

> [!NOTE]
> これは因果的（デコーダー）モデルでのみ機能します。エンコーダーモデルをデコーダーモデルに変換するものではありません。

モデル設定で`is_causal=False`を設定すると、すべてのフォワードパスで双方向アテンションがデフォルトになります。

```py
from transformers import AutoModel, AutoConfig

config = AutoConfig.from_pretrained("meta-llama/Llama-3.2-1B")
config.is_causal = False

model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B", config=config)

# all forward passes now use bidirectional attention
outputs = model(**inputs)
```

モデルを2回読み込むことなく、causal attentionとbidirectional attentionの切り替えを行うには、モデルのconfigではなく、forward呼び出しで`is_causal`を渡します。このkwargはconfigを一時的に上書きし、呼び出し後に元の状態に復元されます。

```py
from transformers import AutoModel

model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B")

# run with bidirectional attention
outputs = model(**inputs, is_causal=False)

# run with default causal attention
outputs = model(**inputs)
```