GroupedQueryAttention

Dot-product attention sharing keys and values across heads.

Abstract Signature:

GroupedQueryAttention(embed_dim: int, num_heads: int, num_kv_heads: int)

PyTorch

API: —
Strategy: Custom / Partial

Keras

API: keras.layers.MultiHeadAttention
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: flax.nnx.MultiHeadAttention
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: paxml.layers.GroupedQueryAttention
Strategy: Direct Mapping
Official Docs ↗