openvino.opset17.sparse_attention#

openvino.opset17.sparse_attention(query: Node | int | float | ndarray, key: Node | int | float | ndarray, value: Node | int | float | ndarray, block_indices: Node | int | float | ndarray, block_count: Node | int | float | ndarray, key_length: Node | int | float | ndarray | None = None, scale: Node | int | float | ndarray | None = None, sinks: Node | int | float | ndarray | None = None, *, key_block_size: int, query_block_size: int = 1, causal: bool = True, name: str | None = None) → Node#

Return a node which computes attention over the key blocks selected for each query.

Every query token attends only to the tokens of its selected key blocks and, in the causal mode, to the tokens of the incomplete block that contains it.

Parameters:

  • query – Attention query of shape [N, H, L, E].

  • key – Attention key of shape [N, Hkv, S, E].

  • value – Attention value of shape [N, Hkv, S, Ev].

  • block_indices – Selected key blocks of shape [N, Hs, Lq, KB].

  • block_count – Number of valid entries of every row of block_indices, shape [N, Hs, Lq].

  • key_length – Number of valid key tokens of every sequence, shape [N] or [1].

  • scale – Attention scale factor used instead of the default 1 / sqrt(E).

  • sinks – Attention sink logit of every query head, shape [H]. Requires scale.

  • key_block_size – Number of consecutive key tokens forming one key block.

  • query_block_size – Number of consecutive query tokens sharing one selection.

  • causal – Whether the query tokens are the last key tokens and attend to their incomplete block.

  • name – The optional name for the new output node.

Returns:

The new node performing SparseAttention operation.