Cca Attention

Compressed Convolutional Attention (CCA) and CCGQA.

Implements Compressed Convolutional Attention (CCA) and Compressed Convolutional Grouped Query Attention (CCGQA) from arXiv:2510.04476 (Figliola, Alonso, Iyer, Anthony & Millidge, 2025, Zyphra).

Both variants down-project queries, keys, and values into a shared compressed latent space and perform the entire attention operation inside that latent – there are no Q/K/V up-projection matrices (unlike MLA, which up-projects back to full width before attending). Only a single output up-projection W̃_O maps the latent output back to the residual stream. This simultaneously reduces parameters, KV-cache size, and attention FLOPs by the compression factor C (MLA only shrinks the cache).

To make attention in the fully compressed latent space viable, CCA introduces three innovations, all toggleable via config:

Two convolutions on the packed q/k tensor: a depth-wise causal sequence convolution (mixes across positions) followed by a head-wise grouped channel convolution (mixes across channels within each head). The paper’s ablation shows two conv layers is optimal.
q-k-mean: adds the pre-convolution mean of q and k to the post-convolution values, increasing attention-diagonal sparsity when combined with QK-norm.
Value-shift: each attention head receives half its values from the current token and half from the previous token (a token-shift inductive bias borrowed from RWKV), implemented via two independent value projections.

After down-projection + convolutions + qk-mean + value-shift, QK L2-normalisation and a learnable key temperature β are applied, RoPE is applied directly in the latent (no separate RoPE head/cache needed, unlike MLA), and standard softmax attention is computed.

CCGQA extends CCA with GQA-style key/value head sharing applied inside the compressed latent, and decouples the query and KV compression rates: C₁ (query) and C₂ (KV) with C₂ ≥ C₁. The per-head latent dimension d_h must match between query and key heads, which enforces C₂ / C₁ = num_heads / num_kv_heads.

Reference:: Figliola, T., Alonso, N., Iyer, R., Anthony, Q., & Millidge, B. (2025). “Compressed Convolutional Attention: Efficient Attention in a Compressed Latent Space”. arXiv:2510.04476.

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

Bases: Module

Compressed Convolutional Attention (Figliola et al. 2025).

Down-projects q, k, v into a shared latent of dimension cca_latent_rank (ẽ = E / C), performs attention entirely in the latent with optional convolutions, qk-mean, and value-shift, then up-projects the output back to the residual stream via a single W̃_O.

Parameters:

config –

Model configuration object with attributes hidden_size, num_heads, dropout, use_bitnet, mode, and the optional CCA-specific fields:

cca_latent_rank – latent width ẽ (default hidden_size // 4, i.e. compression C = 4). Must be divisible by num_heads.
cca_num_conv_layers – 0, 1, or 2 convolution layers (default 2; the paper’s recommended setting).
cca_conv_kernel_seq – kernel size k_seq of the depth-wise causal sequence convolution (default 4).
cca_conv_kernel_ch – kernel size k_ch of the head-wise grouped channel convolution (default 3).
cca_qk_mean – enable the q-k-mean bias (default True).
cca_value_shift – enable value-shift with two value projections (default True; requires num_heads even and cca_latent_rank even).
rope_base – RoPE base frequency (default 10000.0).

hidden_size: Input embedding dimensionality.

num_heads: Number of attention heads.

latent_dim: Latent width ẽ.

latent_head_dim: Per-head latent dimensionality d_h = ẽ / num_heads.

num_conv_layers: Number of convolution layers (0, 1, or 2).

conv_kernel_seq: Sequence-conv kernel size.

conv_kernel_ch: Channel-conv kernel size.

qk_mean: Whether q-k-mean is enabled.

value_shift: Whether value-shift is enabled.

rope_base: RoPE base frequency.

linear_qk: Packed q/k down-projection E -> 2ẽ.

val_proj1, val_proj2: Value projections for value-shift (each E -> ẽ/2). Present only when value_shift is True.

val_proj: Single value projection E -> ẽ. Present only when value_shift is False.

out_proj: Output up-projection ẽ -> E (W̃_O).

conv_qk0: Depth-wise causal sequence Conv1d (or None).

conv_qk1: Head-wise grouped channel Conv1d (or None).

temp: Learnable key temperature β (scalar, init 0).

dropout: Dropout layer.

mode: "encoder" or "decoder".

scale: Attention softmax scale 1/sqrt(d_h).

Raises:: ValueError – If hidden_size not divisible by num_heads; if cca_latent_rank not divisible by num_heads; if cca_num_conv_layers not in {0, 1, 2}; if cca_value_shift is True but num_heads is odd or cca_latent_rank is odd; if latent_head_dim is odd (RoPE requires even).

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

Compute Compressed Convolutional 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).

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

Bases: Module

Compressed Convolutional Grouped Query Attention (Figliola et al. 2025).

Extends CCA with GQA-style key/value head sharing applied inside the compressed latent, and decouples the query and KV compression rates. The query latent has width E / C₁ and the KV latent has width E / C₂ with C₂ ≥ C₁. The per-head latent dimension d_h must be the same for query and key heads, which enforces the constraint:

ccgqa_query_latent_rank / num_heads
    == ccgqa_kv_latent_rank / ccgqa_num_kv_heads

i.e. C₂ / C₁ == num_heads / ccgqa_num_kv_heads (the GQA group size).

Parameters:

config –

Model configuration object with attributes hidden_size, num_heads, dropout, use_bitnet, mode, and the optional CCGQA-specific fields:

ccgqa_query_latent_rank – query latent width E / C₁ (default hidden_size // 2, i.e. C₁ = 2). Must be divisible by num_heads.
ccgqa_kv_latent_rank – KV latent width E / C₂ (default hidden_size // 8, i.e. C₂ = 8). Must be divisible by ccgqa_num_kv_heads and ≤ ccgqa_query_latent_rank.
ccgqa_num_kv_heads – number of KV (group) heads (default num_heads // 4; must divide num_heads).
ccgqa_num_conv_layers – 0, 1, or 2 (default 2).
ccgqa_conv_kernel_seq – sequence-conv kernel (default 4).
ccgqa_conv_kernel_ch – channel-conv kernel (default 3).
ccgqa_qk_mean – enable q-k-mean with B_group/E_group (default True).
ccgqa_value_shift – enable value-shift (default True; requires ccgqa_num_kv_heads even and ccgqa_kv_latent_rank even).
rope_base – RoPE base frequency (default 10000.0).

hidden_size: Input embedding dimensionality.

num_heads: Number of query heads.

num_kv_heads: Number of KV (group) heads.

group_size: num_heads // num_kv_heads.

latent_head_dim: Per-head latent dimensionality d_h.

query_latent_dim: Query latent width E / C₁.

kv_latent_dim: KV latent width E / C₂.

num_conv_layers: Number of convolution layers (0, 1, or 2).

qk_mean: Whether q-k-mean is enabled.

value_shift: Whether value-shift is enabled.

rope_base: RoPE base frequency.

linear_qk: Packed q/k down-projection E -> (query_latent + kv_latent).

val_proj1, val_proj2: Value projections for value-shift.

val_proj: Single value projection (when value_shift is False).

out_proj: Output up-projection query_latent -> E (W̃_O).

conv_qk0, conv_qk1: Convolution layers (or None).

temp: Learnable key temperature β.

dropout, mode, scale: As in CCA.

Raises:: ValueError – If any divisibility or constraint check fails (see Args above for the full list).

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

Compute Compressed Convolutional Grouped Query 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).