openvino.opset17.qsa_prepare

openvino.opset17.qsa_prepare(raw_index_keys: Node | int | float | ndarray, past_block_index_keys: Node | int | float | ndarray, key_length: Node | int | float | ndarray, rope_cos: Node | int | float | ndarray, rope_sin: Node | int | float | ndarray, k_norm_weight: Node | int | float | ndarray, past_incomplete_raw_keys: Node | int | float | ndarray, past_incomplete_rope: Node | int | float | ndarray, key_block_size: int, epsilon: float = 1e-06, name: str | None = None) → Node

Return a node which builds one index key per complete key block for sparse attention.

A block key is the mean of the raw index keys of the key_block_size consecutive key tokens of the block, RMS-normalized with k_norm_weight and rotated with the RoPE of the first token of the block. Only the blocks completed by the new tokens are computed. All inputs are positional; a feature is disabled by connecting an empty constant.

Parameters:

• raw_index_keys – Raw index key of every new token, shape [N, L, Di].

• past_block_index_keys – Block index keys computed by the previous calls, shape [N, Bp, Di].

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

• rope_cos – RoPE cos of every new token, shape [N, L, Dr]; empty constant disables the rotation.

• rope_sin – RoPE sin of every new token, shape [N, L, Dr]; empty constant disables the rotation.

• k_norm_weight – Scale of the block key normalization, shape [Di]; empty constant disables it.

• past_incomplete_raw_keys – Raw keys of the past tokens of the incomplete block, shape [N, key_block_size - 1, Di].

• past_incomplete_rope – RoPE cos and sin of the first token of the incomplete block, shape [N, 2 * Dr]; empty constant if and only if rope_cos and rope_sin are.

• key_block_size – Number of consecutive key tokens forming one key block, greater than 1.

• epsilon – Value added to the mean square in the RMS normalization.

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

Returns:

The new node performing QSAPrepare operation.