kempnerforge.model

Model architectures for KempnerForge.

class kempnerforge.model.Transformer

Bases: Module

Full transformer model built from ModelConfig.

Embedding → TransformerBlocks → Norm → Output Head

include_token_embedding=False / include_output_head=False build the model without that module: the attribute is None and registers no state-dict keys. Without the embedding, pass pre-embedded inputs via modality.inputs_embeds; without the head, forward returns the final hidden states. tie_embeddings=True requires both modules. generate() rejects a model built without either module. Tensor parallelism is not supported for such a model.

__init__(config, *, vlm_config=None, num_image_tokens=0, include_token_embedding=True, include_output_head=True)

Parameters:

• config (ModelConfig)

• vlm_config (VLMConfig | None)

• num_image_tokens (int)

• include_token_embedding (bool)

• include_output_head (bool)

Return type:

None

init_weights_and_freqs()

Initialize weights and RoPE frequencies after meta-device materialization.

Called after model.to_empty(device=...) to fill in parameter values and compute RoPE frequency table. Safe to call on already-initialized models (skips if freqs are already computed).

Return type:

None

set_moe_step(step, max_steps)

Set training step on all MoE routers for adaptive bias scheduling.

Reaches every sigmoid router, including each MoT modality’s.

Parameters:

• step (int)

• max_steps (int)

Return type:

None

get_moe_aux_loss()

Sum the auxiliary losses of every MoE router. Returns 0 if dense.

A MoT layer contributes one term per modality, so an image + text MoT model sums twice as many terms as a text model with the same MoE layers.

Return type:

torch.Tensor

get_moe_router_z_loss()

Sum the router z-losses (ST-MoE) of every MoE router. Returns 0 if dense.

Like get_moe_aux_loss, a MoT layer contributes one term per modality.

Return type:

torch.Tensor

get_expert_counts()

Collect per-layer expert utilization for flat MoE layers.

Returns {layer_idx: (num_experts,) tensor} for layers whose MLP is a MoEMLP (the standard, single-pool MoE). Returns {} for dense models, for MoT and for MoMa, whose per-modality routers surface through get_mot_expert_counts and get_moma_expert_counts instead.

Return type:

dict[int, torch.Tensor]

get_mot_expert_counts()

Collect per-layer, per-modality expert utilization for MoT MoE layers.

Returns {layer_idx: {modality: (num_experts,) tensor}} for every MoT layer whose per-modality FFNs are MoEMLP``s; returns ``{} otherwise. Each tensor holds the token count per expert from that modality’s router in the most recent forward, so the image and text counts of a layer sum to different totals.

Return type:

dict[int, dict[str, torch.Tensor]]

get_moma_expert_counts()

Collect per-layer, per-modality expert utilization for MoMa layers.

Returns {layer_idx: {modality: (num_experts_for_modality,) tensor}} for every layer whose MLP is a MoMaFFN; returns {} otherwise. Each modality’s expert count tensor reflects the most recent forward through that layer’s expert-choice router (paper Figure 5-style utilization). Counts on a fresh model (no forward yet) are the router’s init zeros.

Return type:

dict[int, dict[str, torch.Tensor]]

forward(tokens=None, *, modality=None, kv_caches=None, doc_ids=None, position_ids=None, attention_group_ids=None)

Forward pass.

Exactly one of tokens or modality.inputs_embeds must be provided. Modality-injection routes (prefix_embeds, output_slice, image_features, image_mask, modality_ids, inline_embeds) are grouped on the optional ModalityContext arg; see kempnerforge/model/modality.py for the full intra-context invariant table. inline_embeds are written over the embedded tokens at (inline_rows, inline_positions) before the first layer, which is how a packed row carries each document’s image tokens before its own text, and output_rows / output_positions pick the positions the head runs on, giving (1, M, vocab) logits.

Parameters:

• tokens (torch.Tensor | None) – Integer token ids, shape (batch, seq_len).

• modality (ModalityContext | None) – Optional ModalityContext bundling pre-embedded inputs, prefix embeds, output slicing, image features, and modality routing tags for VLM arches. None is the pure text-only forward. It may also carry position_ids and attention_group_ids (see those arguments) or the rotation itself as rope_cos and rope_sin, (batch, seq_len, head_dim // 2) over the residual, or (1, seq_len, head_dim // 2) shared by every row, for position schemes that are not one integer per token; the rotation excludes position_ids, covers the appended tokens under kv_caches as position_ids does, is applied without the model’s own tables, and is unsupported for MoT.

• kv_caches (list[KVCache] | None) – Optional list of KVCache (one per layer) for generation. The tokens of this forward are appended after the cached ones; without caller positions or a rotation their RoPE positions continue from the cache fill level. Cross-arg invariant: kv_caches forbids modality.prefix_embeds, modality.output_slice, modality.image_features, modality.modality_ids, modality.key_padding_mask, doc_ids and attention_group_ids, by argument or as modality.attention_group_ids (all training-only).

• doc_ids (torch.Tensor | None) – Optional per-token document IDs for packed sequences, shape (batch, seq_len) over the residual stream. Enables block-diagonal causal attention that isolates documents within packed sequences.

• position_ids (torch.Tensor | None) – Optional explicit per-token RoPE positions, shape (batch, seq_len) matching the residual reaching attention (including any modality.prefix_embeds), dtype int64, values in [0, model.max_seq_len). Rows may differ, so positions can e.g. restart at zero at each document of a packed sequence. None (default) assigns contiguous positions starting at the KV-cache fill level, exactly as before. With kv_caches the tensor covers the tokens of this forward alone, the ones being appended, so a caller that positions its own tokens does so at prefill and at every decode step alike. Unsupported for MoT, which derives per-modality positions internally. A ModalityContext may carry the same tensor as modality.position_ids instead, which is how a strategy chooses the positions of the residual it lays out; passing both raises.

• attention_group_ids (torch.Tensor | None) – Optional per-token visibility group ids over the residual stream, shape (batch, seq_len), int64; negative means no group. Tokens sharing a nonnegative id see each other bidirectionally (e.g. prefix-LM spans): visible(q, k) = same_document AND (k <= q OR same_nonnegative_group). Without doc_ids the row is one document. None leaves attention unchanged; all-negative ids give the same visibility as None but still build a mask, so None is the unmasked fast path. On MoT, doc_ids and group ids need vlm.mot_modalities in residual order, ("image", "text"). May likewise come from modality.attention_group_ids, not both.

Returns:

Logits tensor of shape (batch, out_seq_len, vocab_size) where out_seq_len == seq_len normally or the sliced length when modality.output_slice is set, or (1, M, vocab_size) over the M positions modality.output_rows and output_positions pick. Built with include_output_head=False, returns the final hidden states (batch, out_seq_len, dim) instead.

Return type:

torch.Tensor

class kempnerforge.model.TransformerBlock

Bases: Module

Single transformer block with pre-norm architecture.

Structure: norm → attention → residual, norm → mlp → residual

__init__(config, layer_idx)

Parameters:

• config (ModelConfig)

• layer_idx (int)

Return type:

None

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

Run attention and the MLP over the residual stream.

Parameters:

• x (torch.Tensor) – Residual stream, shape (batch, seq_len, dim).

• rope_cos (torch.Tensor) – RoPE cosine frequencies for this block’s positions.

• rope_sin (torch.Tensor) – RoPE sine frequencies for this block’s positions.

• kv_cache (KVCache | None) – Optional KV cache for incremental generation.

• attn_mask (torch.Tensor | BlockMask | None) – The one mask from kempnerforge.model.masking.build_attention_mask, built once per forward and shared by every block.

Returns:

The updated residual stream, shape (batch, seq_len, dim).

Return type:

torch.Tensor

Modules

adapter	Vision-to-LLM adapter modules (the "connector").
attention	Multi-head attention with Grouped-Query Attention (GQA) support.
cross_attention	Cross-attention block for VLM Cross-Attention architecture.
embedding	Token embedding and output head for KempnerForge models.
generate	Autoregressive text generation with KV-cache.
hooks	Activation extraction hooks for mechanistic interpretability.
init	Weight initialization strategies for KempnerForge models.
masking	Attention mask builders.
mlp	Feed-forward network implementations for KempnerForge models.
modality	Modality-injection container for Transformer.forward.
moe	Mixture-of-Experts feed-forward layer for KempnerForge models.
moma	Mixture of Modality-Aware Experts (MoMa) operator, FFN, and block.
mot	Mixture-of-Transformers (MoT) operator and block.
norm	Normalization layers for KempnerForge models.
position	Rotary Position Embedding (RoPE) for KempnerForge models.
router	MoE router implementations for KempnerForge models.
transformer	Transformer model for KempnerForge.
vision	Vision encoders for VLM training.
vlm	Vision-language model wrapper.