tensorplay.nn.functional.scaled_dot_product_attention

tensorplay.nn.functional.scaled_dot_product_attention(query: Tensor, key: Tensor, value: Tensor, attn_mask=None, dropout_p: float = 0.0, is_causal: bool = False, scale=None, enable_gqa: bool = False, backend: str | None = None) → Tensor[source]

scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, backend=None) -> Tensor

Computes scaled dot product attention on query, key and value. Routes to reference:

$$\text{Attention}(Q, K, V) = \text{softmax}(\frac{Q K^T}{\sqrt{E}}) V$$

Parameters:

backend (str, optional) – 'flash' | 'mem_efficient', 'math', or None to pick automatically. 'flash' selects the fused flash-attention kernel, 'math' forces the composed reference path. When None, the routing candidate order is governed by tensorplay.nn.attention.sdpa_kernel().