MultiHeadAttention (CPU)#

  • Domain: com.microsoft

  • Device: CPU

  • Kernel: onnx_light_cpu::MultiHeadAttention

  • Supported types: FLOAT, FLOAT16, BFLOAT16

  • Opset bounds: since opset 1

Operator support#

Capability

Implementation

Shape inference

onnx_light_cpu::ComputeShapeMultiHeadAttention()

Peak memory

onnx_light_cpu::ComputePeakMemoryMultiHeadAttention()

Fusion patterns

None

Gradient

Not provided

LightOpSchema#

Version: 1

ONNX Runtime multi-head self/cross attention. Rank-5 packed QKV, shared cache buffers and qk output are not implemented by this CPU provider.

Inputs

  • query (T): (B,S,N*H).

  • key (T): (B,L,N*H) or cross-attention (B,N,L,H).

  • value (T): (B,L,N*Hv) or cross-attention (B,N,L,Hv).

  • bias (T): Optional projection bias (2*N*H+N*Hv); cross-attention uses query bias.

  • key_padding_mask (M): Optional int32 lengths (B), or binary mask (B,T)/(B,S,T).

  • attention_bias (T): Optional additive (1 or B,1 or N,S,T).

  • past_key (T): Optional paired cache (B,N,P,H).

  • past_value (T): Optional paired cache (B,N,P,Hv).

  • past_sequence_length (M): Unsupported shared-buffer input.

  • cache_indirection (M): Unsupported shared-buffer input.

Outputs

  • output (T): (B,S,N*Hv).

  • present_key (T): Optional (B,N,P+L,H).

  • present_value (T): Optional (B,N,P+L,Hv).

Attributes

  • mask_filter_value (FLOAT): Mask additive penalty. (optional, default: -10000)

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

  • scale (FLOAT): Defaults to 1/sqrt(H). (optional)

  • unidirectional (INT): Causal mask flag. (optional, default: 0)

Type constraints

  • T: Matching floating-point tensors. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).

  • M: Int32 padding mask. Allowed types: tensor(int32).