ttnn.transformer.scaled_dot_product_attention

ttnn.transformer.scaled_dot_product_attention(input_tensor_q: ttnn.Tensor, input_tensor_k: ttnn.Tensor, input_tensor_v: ttnn.Tensor, *, attn_mask: ttnn.Tensor | None = None, is_casual: bool = true, memory_config: ttnn.MemoryConfig | None = None, queue_id: int | None = 0, scale: float | None = None, program_config: SDPAProgramConfig | None = None, compute_kernel_config: ttnn.DeviceComputeKernelConfig | None = None) → ttnn.Tensor

Causal scaled dot product attention. This API mimicks the PyTorch API of the same name. The implementation is FlashAttention-2.”

Accepts a SDPAProgramConfig which specifies the grid size and chunk tiles in the Q and K sequence lengths. The op parallelizes over b, nqh, and Q’s s dimension.

Parameters:

  • input_tensor_q (ttnn.Tensor) – the input tensor. [b x nqh x s x dh]

  • input_tensor_k (ttnn.Tensor) – the input tensor. [b x nkv x s x dh]

  • input_tensor_v (ttnn.Tensor) – the input tensor. [b x nkv x s x dh]

Keyword Arguments:

  • attn_mask (ttnn.Tensor, optional) – Defaults to None. [b x 1 x s x s]. Head broadcasting is implied.

  • is_casual (bool) – Defaults to true.

  • memory_config (ttnn.MemoryConfig, optional) – Memory configuration for the operation. Defaults to None.

  • queue_id (int, optional) – command queue id. Defaults to 0.

  • scale (float, optional) – Defaults to None.

  • program_config (SDPAProgramConfig, optional) – Defaults to None.

  • compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) – Defaults to None.

Returns:

ttnn.Tensor – the output tensor [b x nqh x s x dh].