Mla Attention

Multi-Head Latent Attention (MLA + RoPE) for small models.

Implements the latent attention variant studied by Mehta et al. (2025), arXiv:2506.09342, for small language models. MLA jointly compresses keys and values into a single low-rank latent vector c_KV whose rank r_kv is strictly smaller than hidden_size; only the latent is cached and stored during inference, so the KV-cache footprint drops from 2 * num_heads * head_dim to r_kv per token. The paper’s Pareto- optimal configuration is r_kv = hidden_size // 2 plus rotary positional embeddings (RoPE) applied to the decompressed queries and keys, which recovers (and slightly exceeds) standard MHA quality while halving the cache.

Formulation (per token x):

c_KV = W_DKV x # latent, shape r_kv k = W_UK c_KV # keys, shape num_heads*head_dim v = W_UV c_KV # values, shape num_heads*head_dim q = W_Q x # queries, shape num_heads*head_dim [q, k] = RoPE([q, k]) # rotary on decompressed Q/K out = SDPA(q, k, v) * W_O # standard softmax attention

Reference:: Mehta, S., Dandekar, R., Dandekar, R., & Panat, S. (2025). “Latent Multi-Head Attention for Small Language Models”, arXiv:2506.09342.

class src.model.attention.latent.mla_attn.MLAAttention(*args: Any, **kwargs: Any)[source]

Bases: Module

Multi-Head Latent Attention with RoPE (Mehta et al. 2025).

Compresses the key-value cache into a low-rank latent vector via a shared down-projection W_DKV and two up-projections W_UK, W_UV. RoPE is applied to the decompressed query and key tensors, which the paper shows is essential to recover MHA quality on small models.

Parameters:: config – Model configuration object with attributes hidden_size, num_heads, dropout, use_bitnet, mode and the optional mla_latent_rank (default hidden_size // 2) and rope_base (default 10000.0).

hidden_size: Input embedding dimensionality.

num_heads: Number of attention heads.

head_dim: Dimensionality per head (hidden_size // num_heads).

latent_rank: Rank r_kv of the joint key-value latent.

rope_base: RoPE base frequency.

dkv_proj: Down-projection hidden_size -> r_kv (latent).

uk_proj: Key up-projection r_kv -> num_heads*head_dim.

uv_proj: Value up-projection r_kv -> num_heads*head_dim.

q_proj: Query projection hidden_size -> num_heads*head_dim.

out_proj: Output projection.

dropout: Dropout layer.

mode: "encoder" or "decoder".

Raises:: ValueError – If hidden_size is not divisible by num_heads.

forward(x: torch.Tensor, logical_layer_idx: int | None = None) → torch.Tensor[source]

Compute latent attention with RoPE on decompressed Q/K.

Parameters:

x – Input tensor of shape (batch_size, seq_len, hidden_size).
logical_layer_idx – Unused; accepted for interface compatibility.

Returns:

Output tensor of shape (batch_size, seq_len, hidden_size).