kempnerforge.model.attention

Multi-head attention with Grouped-Query Attention (GQA) support.

GQA is the general case:

• n_kv_heads == n_heads → standard Multi-Head Attention (MHA)

• n_kv_heads == 1 → Multi-Query Attention (MQA)

• 1 < n_kv_heads < n_heads → Grouped-Query Attention (GQA)

Classes

Attention	Grouped-Query Attention with RoPE and SDPA.
KVCache	Pre-allocated KV cache for autoregressive generation.

class kempnerforge.model.attention.KVCache

Bases: object

Pre-allocated KV cache for autoregressive generation.

Stores key and value tensors for all previous positions, enabling incremental decoding without recomputing attention over the full sequence. Keys are stored after RoPE application but before GQA expansion.

__init__(batch_size, max_seq_len, n_kv_heads, head_dim, dtype, device)

Parameters:

• batch_size (int)

• max_seq_len (int)

• n_kv_heads (int)

• head_dim (int)

• dtype (torch.dtype)

• device (torch.device)

Return type:

None

update(k_new, v_new)

Append new key/value entries and return full cached tensors.

Parameters:

• k_new (torch.Tensor) – New keys, shape (batch, n_kv_heads, new_seq_len, head_dim).

• v_new (torch.Tensor) – New values, shape (batch, n_kv_heads, new_seq_len, head_dim).

Returns:

Tuple of (all_keys, all_values), each (batch, n_kv_heads, total_seq_len, head_dim).

Return type:

tuple[torch.Tensor, torch.Tensor]

class kempnerforge.model.attention.Attention

Bases: Module

Grouped-Query Attention with RoPE and SDPA.

__init__(dim, n_heads, n_kv_heads, head_dim=None, qk_norm=False, sdpa_backend='auto', norm_eps=1e-05)

Parameters:

• dim (int)

• n_heads (int)

• n_kv_heads (int)

• head_dim (int | None)

• qk_norm (bool)

• sdpa_backend (str)

• norm_eps (float)

Return type:

None

forward(x, rope_cos, rope_sin, *, kv_cache=None, attn_mask=None)

Forward pass.

Attention consumes the visibility mask; it never derives one. Every visibility rule (document isolation, bidirectional groups, key validity, and their combination) is expressed by kempnerforge.model.masking, which builds this one object once per forward and shares it across layers. Without a mask, attention is causal, aligned to the end of the key axis when tokens are appended to a filled kv_cache.

Parameters:

• x (torch.Tensor) – Input tensor of shape (batch, seq_len, dim).

• rope_cos (torch.Tensor) – RoPE cosine frequencies, shape (seq_len, head_dim // 2) shared across the batch or (batch, seq_len, head_dim // 2) per row.

• rope_sin (torch.Tensor) – RoPE sine frequencies, same shape as rope_cos.

• kv_cache (KVCache | None) – Optional KV cache for incremental generation. The tokens of this forward are appended to it, and each attends to every cached key and to the appended keys up to its own position. Mutually exclusive with attn_mask: a cache grows the key axis every step, so the mask would have to be rebuilt each step rather than shared, and a decode step needs no mask at all.

• attn_mask (torch.Tensor | BlockMask | None) – The mask from kempnerforge.model.masking.build_attention_mask: a FlexAttention BlockMask, a dense bool tensor broadcastable to (batch, n_heads, seq_len, seq_len) where True = attend, or None for plain causal attention.

Returns:

Output tensor of shape (batch, seq_len, dim).

Raises:

• ValueError – If attn_mask is combined with kv_cache.

• NotImplementedError – If capture_attention_weights is set and attn_mask is a BlockMask.

Return type:

torch.Tensor