kempnerforge.model.modality

Modality-injection container for Transformer.forward.

ModalityContext groups all “inputs that flow into the existing residual stream” plus per-token routing tags consumed by the residual stream, so Transformer.forward stays narrow regardless of how many architectures are active. Each VLM arch fills the fields it needs:

  • Joint-Decoder fills prefix_embeds + output_slice (image tokens prepended to the text sequence; output_slice trims them off the hidden state before the LM head).

  • Cross-Attention fills image_features + image_mask (image K/V flowing into separate cross-attention blocks; the residual stream itself carries text only).

  • Mixture-of-Transformers fills prefix_embeds + output_slice + modality_ids. The residual stream carries (image, text) concatenated; modality_ids tags every position with its modality so each layer’s MoTBlock can route per-modality projections + global self-attention.

  • Pipeline-parallel middle stages fill inputs_embeds (pre-embedded hidden state passed across stage boundaries).

  • Packed Joint-Decoder batches fill inline_embeds + inline_rows + inline_positions: image tokens written over placeholder positions inside the token-embedded rows, so each document’s images sit before its own text and doc_ids keeps attention inside the document.

  • A strategy that chooses the RoPE positions or the attention visibility of the residual it lays out (an image’s tokens seeing each other, crops of one image sharing positions) fills position_ids and attention_group_ids. Transformer.forward reads them as its arguments of the same names, which may not be passed alongside.

  • A strategy whose positions are not one integer per token (a 2-D or multi-axis RoPE over image crops) fills rope_cos and rope_sin instead: the rotation itself, (batch, seq_len, head_dim // 2) or (1, seq_len, head_dim // 2) shared by every row, which the transformer applies in place of its own table. The strategy builds them from precompute_rope_frequencies with the model’s head_dim and rope_theta, so the same code positions every forward, training and inference alike. forward moves the rotation tensors to the residual’s device for each call; build them on that device to avoid repeated host-to-device transfers.

Cross-arg invariants involving kv_caches (a Transformer.forward argument, not a ModalityContext field) are enforced at the top of Transformer.forward, not in __post_init__.

Classes

ModalityContext

Modality-injection container.
class kempnerforge.model.modality.ModalityContext[source]

Bases: object

Modality-injection container.

Invariants enforced in __post_init__:

  • At most one of inputs_embeds, prefix_embeds, image_features, inline_embeds may be set; they are mutually exclusive composition routes into the residual stream.

  • inline_embeds (M, dim) requires inline_rows and inline_positions, both (M,) int64, and they require it: token m is written at [inline_rows[m], inline_positions[m]] of the embedded rows.

  • output_rows and output_positions come together, (M,) int64, and exclude output_slice: the head runs on h[output_rows, output_positions] only, giving (1, M, vocab) logits, which is how a packed row skips its image slots and padding tail.

  • image_mask requires image_features to be set (a free-standing image_mask with no features is a programming error).

  • modality_ids requires prefix_embeds, inputs_embeds or inline_embeds to be set (routing without a residual extension is meaningless).

  • key_padding_mask requires prefix_embeds, inputs_embeds or inline_embeds to be set (it is a key-validity mask over the residual sequence; the Cross-Attention arch masks image K/V via image_mask instead).

  • position_ids and attention_group_ids are (batch, seq_len) int64 over the residual stream, with the meaning of the Transformer.forward arguments of those names; the dtype and rank are checked here, the shape against the residual in forward.

  • rope_cos and rope_sin come together, as floating (batch or 1, seq_len, head_dim // 2) tensors of one shape, dtype and device, and exclude position_ids; forward checks the shape against the residual.

output_slice composes with the tokens path AND with the inputs_embeds path; it is not constrained intra-context. The cross-arg constraint (output_slice vs kv_caches) lives on Transformer.forward instead.

inputs_embeds: torch.Tensor | None = None
prefix_embeds: torch.Tensor | None = None
output_slice: slice | None = None
image_features: torch.Tensor | None = None
image_mask: torch.Tensor | None = None
modality_ids: torch.Tensor | None = None
key_padding_mask: torch.Tensor | None = None
inline_embeds: torch.Tensor | None = None
inline_rows: torch.Tensor | None = None
inline_positions: torch.Tensor | None = None
output_rows: torch.Tensor | None = None
output_positions: torch.Tensor | None = None
position_ids: torch.Tensor | None = None
attention_group_ids: torch.Tensor | None = None
rope_cos: torch.Tensor | None = None
rope_sin: torch.Tensor | None = None
__init__(inputs_embeds=None, prefix_embeds=None, output_slice=None, image_features=None, image_mask=None, modality_ids=None, key_padding_mask=None, inline_embeds=None, inline_rows=None, inline_positions=None, output_rows=None, output_positions=None, position_ids=None, attention_group_ids=None, rope_cos=None, rope_sin=None)
Parameters:
Return type:

None