Gta Attention

Grouped-head laTenT Attention (GTA).

Implements Grouped-Head laTenT Attention, arXiv:2506.17286 (Sun et al. 2025). GTA attacks the redundancy that attention maps across heads exhibit high similarity (much of the per-head computation is unnecessary) and that the value cache can be heavily compressed. It combines two components:

  1. Shared attention map mechanism: a single attention score tensor is computed per group of heads and reused across all heads in the group, shrinking the key cache.

  2. Nonlinear value decoder with learned projections: the value cache is compressed into a low-rank latent space by a down-projection, and reconstructed by a non-linear (silu) decoder before the output projection, further cutting memory.

The paper reports GTA cuts attention FLOPs by up to 62.5% versus GQA and shrinks the KV cache by up to 70%, while avoiding the extra overhead of Multi-Head Latent Attention, achieving a 2x end-to-end inference speedup.

Reference:

Sun, L., Deng, C., Jiang, J., Wu, X., Zhang, H., Chen, L., Ni, L., & Wang, J. (2025). “GTA: Grouped-head latenT Attention”. arXiv:2506.17286.

class src.model.attention.latent.gta_attn.GTAAttention(*args: Any, **kwargs: Any)[source]

Bases: Module

Grouped-Head laTenT Attention with shared maps + latent values.

Parameters:

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

hidden_size

Input dimensionality.

num_heads

Number of query heads H.

head_dim

Per-head dimensionality.

num_groups

Number of head groups sharing an attention map.

group_size

num_heads // num_groups.

value_latent_rank

Latent rank of the value cache.

q_proj, k_proj

Query/key projections.

dv_proj

Value down-projection hidden_size -> value_latent_rank.

uv_proj

Non-linear value decoder value_latent_rank -> num_heads*head_dim (silu activated).

out_proj

Output projection.

dropout, mode

As in the rest of the family.

Raises:

ValueError – If hidden_size not divisible by num_heads or gta_num_shared_groups not dividing num_heads.

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

Compute GTA attention with shared group maps and latent values.

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