GroupQueryAttention (CPU)#

  • Domain: com.microsoft

  • Device: CPU

  • Kernel: onnx_light_cpu::GroupQueryAttention

  • Supported types: FLOAT, FLOAT16, BFLOAT16

  • Opset bounds: since opset 1

Operator support#

Capability

Implementation

Shape inference

onnx_light_cpu::ComputeShapeGroupQueryAttention()

Peak memory

onnx_light_cpu::ComputePeakMemoryGroupQueryAttention()

Fusion patterns

onnx_light_cpu::GroupQueryAttentionFusionPattern

Gradient

onnx_light_cpu::RegisterCustomOperatorGradients()

LightOpSchema#

Version: 1

Grouped-query attention over rank-3 Q/K/V tensors (or packed QKV in query with empty key/value inputs), with an optional tensor KV cache and split-half or interleaved rotary position embedding applied to Q/K at the position derived from seqlens_k (or an explicit position_ids).

Inputs

  • query (T): Query tensor with shape (batch, sequence, num_heads * head_size), or packed QKV (batch, sequence, (num_heads + 2 * kv_num_heads) * head_size) when key/value are empty. Each packed row contains three contiguous segments in [Q | K | V] order, with widths num_heads * head_size, kv_num_heads * head_size, and kv_num_heads * head_size.

  • key (T): Key tensor with shape (batch, sequence, kv_num_heads * head_size), or empty for packed QKV.

  • value (T): Value tensor with shape (batch, sequence, kv_num_heads * value_head_size), or empty for packed QKV (where value_head_size equals head_size).

  • past_key (T_CACHE): Optional KV cache key tensor with shape (batch, kv_num_heads, past_sequence_length, head_size).

  • past_value (T_CACHE): Optional KV cache value tensor with shape (batch, kv_num_heads, past_sequence_length, value_head_size).

  • seqlens_k (M): INT32 tensor containing one total sequence length minus one per batch.

  • total_sequence_length (M): INT32 scalar total sequence length.

  • cos_cache (T): Optional rotary cosine cache with shape (max_sequence_length, rotary_dim / 2), where rotary_dim is positive, even, and <= head_size; required when do_rotary is set.

  • sin_cache (T): Optional rotary sine cache matching cos_cache; only the first rotary_dim coordinates are rotated, leaving the head suffix unchanged; required when do_rotary is set.

  • position_ids (I): Optional INT64 tensor with shape (batch, sequence) giving the rotary position of the first token of each batch; only valid when do_rotary is set.

  • attention_bias (T): Optional additive/boolean mask broadcastable to (batch, num_heads, sequence, total_sequence_length).

  • head_sink (T): Optional tensor with shape (num_heads,) and the query type. Each value is a virtual softmax logit contributing to normalization without contributing to the value reduction.

Outputs

  • output (T): Output tensor with shape (batch, sequence, num_heads * value_head_size) and query type.

  • present_key (T_CACHE): Optional present KV cache key tensor with shape (batch, kv_num_heads, past_sequence_length + sequence, head_size); the concatenation of past_key with the (rotated) current key.

  • present_value (T_CACHE): Optional present KV cache value tensor with shape (batch, kv_num_heads, past_sequence_length + sequence, value_head_size); the concatenation of past_value with the current value.

Attributes

  • causal (INT): Whether to apply causal masking. (optional, default: 1)

  • do_rotary (INT): Whether to apply rotary position embedding. (optional, default: 0)

  • kv_num_heads (INT): Number of key/value heads. (required)

  • local_window_size (INT): Causal window token count including the current token, in [1, INT32_MAX], or -1 for full attention. Requires causal=1 when set; does not truncate past/present caches. (optional, default: -1)

  • num_heads (INT): Number of query heads. (required)

  • rotary_interleaved (INT): Rotary layout: 0 split-half, 1 interleaved. (optional, default: 0)

  • scale (FLOAT): Optional attention-score scale. (optional)

  • softcap (FLOAT): Optional attention-score softcap. (optional, default: 0)

Type constraints

  • T: Constrain query, key, value, output, rotary cache, and head_sink tensors to floating-point tensors. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).

  • T_CACHE: Constrain the supported non-quantized KV cache tensors. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).

  • M: Constrain sequence-length inputs to INT32. Allowed types: tensor(int32).

  • I: Constrain position_ids to INT64. Allowed types: tensor(int64).