tensorplay.nn.functional.scaled_dot_product_attention

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

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

Computes scaled dot product attention on query, key and value. Routes to the fused native kernel when possible, otherwise follows torch’s math 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 (impl=1), 'math' forces the composed reference path; 'mem_efficient', matching torch’s memory-efficient backend, is not available in this build.