GroupedQueryAttention
MultiheadAttentionGroupedQueryAttention(embed_dim: int, num_heads: int, num_kv_heads: int, dropout: float = 0.0, bias: bool = True, add_bias_kv: bool = False, add_zero_attn: bool = False, kdim: int | None = None, vdim: int | None = None, batch_first: bool = False, device: DeviceLike = None, dtype: DTypeLike = None)Grouped-query attention — multi-head attention with fewer key/value heads.
A thin MultiheadAttention with num_kv_heads promoted to a
required argument (the parameter that defines GQA; on the base class it is
optional and defaults to num_heads). The num_kv_heads key/value heads
are split into groups, each shared by num_heads // num_kv_heads query heads,
so the model projects a smaller key/value space and — during incremental
decoding — carries a smaller K/V cache. This is the attention used by Llama 2 /
3, Mistral, Qwen, and Gemma. Exactly equivalent to
MultiheadAttention(..., num_kv_heads=num_kv_heads) (forward, shapes, and the
KV-cache contract are inherited verbatim); this subclass only reshapes the
constructor for discoverability.
Parameters
embed_dimintnum_heads.num_headsintnum_kv_headsintnum_heads. 1 is multi-query attention
(MultiQueryAttention); passing num_heads degenerates to
standard multi-head attention.dropoutfloat= 0.0biasbool= TrueTrue, add a learnable bias to the input and output projections.add_bias_kvbool= FalseTrue, prepend a learnable bias row to the keys and values.add_zero_attnbool= FalseTrue, append a zero row to the keys and values so a query can attend
to "nothing".kdimint or None= Noneembed_dim.vdimint or None= Noneembed_dim.batch_firstbool= FalseTrue, inputs/outputs are (batch, seq, feature); otherwise
(seq, batch, feature).deviceDeviceLike= NonedtypeDTypeLike= NoneNotes
The key/value projections are sized (num_kv_heads * head_dim, …) instead of
(embed_dim, …) (so k_proj_weight / v_proj_weight and the K/V cache
are smaller — the cache shrink is the main inference win: less memory bandwidth
per decode step). Just before the scores are formed, each K/V head is repeated
num_heads // num_kv_heads times via
lucid.nn.functional.repeat_kv so the head counts match; the output is
bit-identical to standard attention with K/V heads tied within each group.
Reference: Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebrón, and Sanghai, "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints", EMNLP 2023 (arXiv:2305.13245).
Examples
>>> import lucid, lucid.nn as nn
>>> gqa = nn.GroupedQueryAttention(embed_dim=512, num_heads=8, num_kv_heads=2,
... batch_first=True)
>>> x = lucid.randn(1, 10, 512)
>>> out, _ = gqa(x, x, x, need_weights=False)
>>> out.shape
(1, 10, 512)
>>> tuple(gqa.k_proj_weight.shape) # 2 K/V heads × head_dim 64 = 128
(128, 512)Used by 1
Constructors
1__init__
→None__init__(embed_dim: int, num_heads: int, num_kv_heads: int, dropout: float = 0.0, bias: bool = True, add_bias_kv: bool = False, add_zero_attn: bool = False, kdim: int | None = None, vdim: int | None = None, batch_first: bool = False, device: DeviceLike = None, dtype: DTypeLike = None)Initialise grouped-query attention; see the class docstring.