Class ov::op::v17::SparseAttention#

class SparseAttention : public ov::op::Op#

SparseAttention - Scaled dot-product attention computed only over the key blocks selected for each query.

Computes grouped-query scaled dot-product attention in which every query token attends only to the tokens of its selected key blocks, instead of to the whole key sequence. Key blocks are groups of key_block_size consecutive key tokens; the ones a query attends to are listed in block_indices and, in the causal mode, the incomplete block containing the query is added. Query tokens share a selection in groups of query_block_size tokens. The selected blocks differ per query token and are normally produced by SparseAttentionIndexer. The optional sinks hold one learnable logit per query head.

Public Functions

SparseAttention(const OutputVector &inputs, int64_t key_block_size, int64_t query_block_size = 1, bool causal = true)#

Constructs a SparseAttention operation from an ordered list of inputs.

Parameters:

  • inputs – query, key, value, block_indices, block_count and optionally key_length and 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

SparseAttention(const Output<Node> &query, const Output<Node> &key, const Output<Node> &value, const Output<Node> &block_indices, const Output<Node> &block_count, int64_t key_block_size, int64_t query_block_size = 1, bool causal = true)#

Constructs a SparseAttention operation without padding and with the default scale.

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]

SparseAttention(const Output<Node> &query, const Output<Node> &key, const Output<Node> &value, const Output<Node> &block_indices, const Output<Node> &block_count, const Output<Node> &key_length, int64_t key_block_size, int64_t query_block_size = 1, bool causal = true)#

Constructs a SparseAttention operation for a batch with padding.

Parameters:

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

SparseAttention(const Output<Node> &query, const Output<Node> &key, const Output<Node> &value, const Output<Node> &block_indices, const Output<Node> &block_count, const Output<Node> &key_length, const Output<Node> &scale, int64_t key_block_size, int64_t query_block_size = 1, bool causal = true)#

Constructs a SparseAttention operation with an explicit attention scale.

Parameters:

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

SparseAttention(const Output<Node> &query, const Output<Node> &key, const Output<Node> &value, const Output<Node> &block_indices, const Output<Node> &block_count, const Output<Node> &key_length, const Output<Node> &scale, const Output<Node> &sinks, int64_t key_block_size, int64_t query_block_size = 1, bool causal = true)#

Constructs a SparseAttention operation with attention sinks.

Parameters:

sinks – Attention sink logit of every query head, shape [H]

virtual void validate_and_infer_types() override#

Verifies that attributes and inputs are consistent and computes output shapes and element types. Must be implemented by concrete child classes so that it can be run any number of times.

Throws if the node is invalid.