braindecode.functional.scaled_dot_product_attention#

braindecode.functional.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, scale=None, fused_in_lazy=False, exact=False)[source]#

F.scaled_dot_product_attention; Gaudi’s fused kernel on HPU.

Parameters:

  • query (Tensor) – Shape (..., seq, head_dim); key/value may have another sequence length.

  • key (Tensor) – Shape (..., seq, head_dim); key/value may have another sequence length.

  • value (Tensor) – Shape (..., seq, head_dim); key/value may have another sequence length.

  • attn_mask (Optional[Tensor]) – Boolean (True takes part) or additive mask, as in torch.

  • dropout_p (float) – Dropout on the attention weights; pass 0.0 in eval mode.

  • scale (Optional[float]) – Score scale; None is 1 / sqrt(head_dim).

  • fused_in_lazy (bool) – Use Gaudi’s fused kernel in HPU lazy mode too.

  • exact (bool) – Run the textbook op order instead of torch’s kernels: softmax((q @ k^T) / sqrt(d)) @ v, or softmax((q @ k^T) * scale) @ v when scale is given, with dropout on the weights. Bit-identical to a hand-written attention; torch’s kernels round differently (its math path scales q and k by sqrt(scale) each before the matmul, its fused CPU kernel runs a blocked softmax), about 1e-6 apart.

Return type:

Tensor

Returns:

  • torch.Tensor – Shape (..., seq, head_dim).

  • In HPU eager mode, torch’s SDPA runs op by op (the full

  • (batch, heads, seq, seq) scores in device memory), 2-3.5x slower than

  • the bridge’s fused kernel on long sequences. The fused kernel is used

  • without a mask or dropout (its dropout would draw other masks). In lazy

  • mode only with ``fused_in_lazy=True`` (the graph compiler fuses torch’s)

  • SDPA too, faster on long sequences, slower on short ones.