GroupQueryAttention (CPU)
=========================

* Domain: ``com.microsoft``
* Device: ``CPU``
* Kernel: ``onnx_light_cpu::GroupQueryAttention``
* Supported types: ``FLOAT``, ``FLOAT16``, ``BFLOAT16``
* Opset bounds: since opset 1

Operator support
----------------

.. list-table::
   :header-rows: 1
   :widths: 30 70

   * - Capability
     - Implementation
   * - Shape inference
     - :cpp:func:`onnx_light_cpu::ComputeShapeGroupQueryAttention`
   * - Peak memory
     - :cpp:func:`onnx_light_cpu::ComputePeakMemoryGroupQueryAttention`
   * - Fusion patterns
     - :cpp:class:`onnx_light_cpu::GroupQueryAttentionFusionPattern`
   * - Gradient
     - :cpp:func:`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).