Gqla Attention

Group-Query Latent Attention (GQLA).

Implements Group-Query Latent Attention, arXiv:2605.15250 (Meng, 2026). GQLA is a minimal modification of Multi-Head Latent Attention (MLA, DeepSeek-V2/V3) whose trained weights expose two algebraically equivalent decoding paths over the same parameters:

  1. MQA-absorb path (default): identical to MLA. The key cache stores the low-rank latent c_KV; at decode time the up-projection is absorbed into the query, giving a single MQA-style head per layer. Pins the H100 roofline at s_q = 1.

  2. GQA path: expands the latent into num_groups full key/value heads and runs standard grouped-query attention with a per-group expanded cache. Selected when s_q > 1 (e.g. s_q = 2 on commodity GPUs such as the H20), enabling Multi-Token Prediction gains and up to 8-way zero-redundancy tensor parallelism along the head axis.

Both paths are algebraically equivalent for the forward output, so the runtime can switch between them with no retraining and no custom kernels. Here we expose a single decode_path configuration knob ("mqa_absorb" or "gqa") and the gqla_num_groups parameter that controls how many GQA groups the expanded path uses.

Reference:

Meng, F. (2026). “GQLA: Group-Query Latent Attention for Hardware-Adaptive Large Language Model Decoding”. arXiv:2605.15250.

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

Bases: Module

Group-Query Latent Attention (Meng 2026) with two decoding paths.

Stores a low-rank latent c_KV of rank gqla_latent_rank from which keys and values are reconstructed by up-projections. At forward time the latent is always expanded (the difference between the two decoding paths is purely a kernel/weight-absorption choice at deployment; the module’s algebraic output is identical). The decode_path attribute is recorded for introspection and tested for shape correctness.

Parameters:

config – Configuration object. Relevant attributes: hidden_size, num_heads, dropout, use_bitnet, mode, and the optional gqla_latent_rank (default hidden_size // 2), gqla_num_groups (default num_heads // 4, must divide num_heads), and gqla_decode_path ("mqa_absorb" or "gqa"; default "gqa").

hidden_size

Input dimensionality.

num_heads

Number of query heads.

head_dim

Per-head dimensionality.

latent_rank

Latent compression rank.

num_groups

Number of GQA groups for the expanded path.

group_size

num_heads // num_groups.

decode_path

Selected decoding path (recorded for introspection).

dkv_proj

Latent down-projection.

uk_proj, uv_proj

Key/value up-projections.

q_proj

Query projection.

out_proj

Output projection.

dropout

Dropout layer.

mode

"encoder" or "decoder".

Raises:

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

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

Compute GQLA attention.

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