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, with an optional tensor KV cache and split-half (rotary_interleaved=0) 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).

  • key (T): Key tensor with shape (batch, sequence, kv_num_heads * head_size).

  • value (T): Value tensor with shape (batch, sequence, kv_num_heads * value_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, head_size / 2); required when do_rotary is set.

  • sin_cache (T): Optional rotary sine cache with shape (max_sequence_length, head_size / 2); 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).

Outputs

  • output (T): Output tensor with the same shape and type as query.

  • 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)

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

  • rotary_interleaved (INT): Rotary layout; only 0 (split-half) is supported. (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, and rotary cache 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).