Mlra Attention

Multi-Head Low-Rank Attention (MLRA).

Implements Multi-Head Low-Rank Attention, arXiv:2603.02188 (Liu et al. 2026). MLRA extends Multi-Head Latent Attention (MLA) by splitting the single shared latent head into mlra_num_latent_heads independent latent sub-spaces, each of rank r = latent_rank / num_latent_heads. The key property is partitionability: each latent sub-head can be assigned to a different tensor-parallel device, so each device loads only 1 / num_latent_heads of the KV cache instead of the whole cache (which MLA forces). This enables efficient 4-way TP decoding and delivers a 2.8x decoding speedup over MLA in the paper’s experiments, while reaching state-of-the-art perplexity and downstream task scores.

Formulation (per token x): the latent cache is split into L = num_latent_heads disjoint sub-vectors c_1, ..., c_L (concatenated into a single c_KV of rank r); keys and values are reconstructed per sub-head via per-block up-projections and concatenated into the full num_heads * head_dim dimension before standard softmax attention.

Reference:

Liu, S., Peng, H., Zhang, Z., Chen, Z., & Guo, Y. (2026). “Multi-Head Low-Rank Attention”. arXiv:2603.02188.

class src.model.attention.latent.mlra_attn.MLRAAttention(*args: Any, **kwargs: Any)

Bases: Module

Multi-Head Low-Rank Attention with partitionable latent heads.

Parameters:

config – Configuration object. Relevant attributes: hidden_size, num_heads, dropout, use_bitnet, mode, and optional mlra_latent_rank (default hidden_size // 2) and mlra_num_latent_heads (default 4, must divide latent_rank evenly).

hidden_size

Input dimensionality.

num_heads

Number of query heads.

head_dim

Per-head dimensionality.

latent_rank

Total latent rank.

num_latent_heads

Number of disjoint latent sub-spaces (L).

sub_rank

latent_rank // num_latent_heads.

dkv_proj

Latent down-projection hidden_size -> r.

uk_projs, uv_projs

ModuleList of per-sub-head up-projections.

q_proj

Query projection.

out_proj

Output projection.

dropout, mode

As in the rest of the family.

Raises:

ValueError – If hidden_size not divisible by num_heads; if mlra_num_latent_heads does not divide latent_rank.

forward(x: torch.Tensor, logical_layer_idx: int | None = None) → torch.Tensor

Compute MLRA attention with partitioned latent sub-heads.

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).