class

MultiQueryAttention

extendsMultiheadAttention

MultiQueryAttention(embed_dim: int, num_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)

source edit

Multi-query attention — multi-head attention with a single key/value head.

A thin MultiheadAttention that fixes num_kv_heads = 1 (so it takes no num_kv_heads argument): all num_heads query heads share one key/value head — the extreme of grouped-query attention. This minimises the key/value projection and, decisively, the K/V cache (one head's worth of K/V per layer instead of num_heads), trading some quality for the smallest decode-memory footprint. Exactly equivalent to MultiheadAttention(..., num_kv_heads=1) (forward, shapes, and the KV-cache contract are inherited verbatim); this subclass only reshapes the constructor for discoverability.

Parameters

embed_dimint

Total model dimension. Must be divisible by num_heads.

num_headsint

Number of query heads (the single K/V head is shared across all of them).

dropoutfloat= 0.0

Dropout probability applied to the attention weights.

biasbool= True

If True, add a learnable bias to the input and output projections.

add_bias_kvbool= False

If True, prepend a learnable bias row to the keys and values.

add_zero_attnbool= False

If True, append a zero row to the keys and values so a query can attend to "nothing".

kdimint or None= None

Feature dimension of the key input; defaults to embed_dim.

vdimint or None= None

Feature dimension of the value input; defaults to embed_dim.

batch_firstbool= False

If True, inputs/outputs are (batch, seq, feature); otherwise (seq, batch, feature).

deviceDeviceLike= None

Device on which the parameters are allocated.

dtypeDTypeLike= None

Data type of the parameters.

Notes

The single key/value head is projected once (k_proj_weight / v_proj_weight of shape (head_dim, …)) and repeated num_heads times via lucid.nn.functional.repeat_kv before the attention scores, so the output is bit-identical to standard attention whose K/V heads are all tied.

Reference: Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need", 2019 (arXiv:1911.02150).

Examples

>>> import lucid, lucid.nn as nn
>>> mqa = nn.MultiQueryAttention(embed_dim=512, num_heads=8, batch_first=True)
>>> mqa.num_kv_heads
1
>>> tuple(mqa.k_proj_weight.shape)   # 1 K/V head × head_dim 64
(64, 512)

Used by 1

lucid.nn.modules

Constructors

dunder

init

→None

__init__(embed_dim: int, num_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)

source edit

Initialise multi-query attention (num_kv_heads = 1).

MultiQueryAttention(embed_dim: int, num_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)

>>> import lucid, lucid.nn as nn >>> mqa = nn.MultiQueryAttention(embed_dim=512, num_heads=8, batch_first=True) >>> mqa.num_kv_heads 1 >>> tuple(mqa.k_proj_weight.shape) # 1 K/V head × head_dim 64 (64, 512)

__init__(embed_dim: int, num_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)