base

base

Base class for vision-language models used by VLA decoders.

GenerativeVLM

GenerativeVLM(input_keys, pretrained, frozen, model_dtype=None, max_text_length=None)

Bases: LanguageEncoderMixin, GenerativeLanguageModel, ABC

Base for VLM components that fuse vision and language in one language-model pass.

Initialize common VLM input wiring and runtime options.

Parameters:

Name	Type	Description	Default
input_keys	str | list[str]	RGB camera keys consumed by the VLM. Tokenized text and padding-mask keys are added by the base class.	required
pretrained	bool	Whether the concrete subclass loads pretrained weights.	required
frozen	bool	Whether all model weights are frozen.	required
model_dtype	str | None	Optional precision string for model parameter dtype.	None
max_text_length	int | None	Optional text sequence length used when running the language model. None lets the concrete subclass set a model-specific default.	None

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def __init__(
    self,
    input_keys: str | list[str],
    pretrained: bool,
    frozen: bool,
    model_dtype: str | None = None,
    max_text_length: int | None = None,
):
    """Initialize common VLM input wiring and runtime options.

    Args:
        input_keys: RGB camera keys consumed by the VLM. Tokenized text and
            padding-mask keys are added by the base class.
        pretrained: Whether the concrete subclass loads pretrained weights.
        frozen: Whether all model weights are frozen.
        model_dtype: Optional precision string for model parameter dtype.
        max_text_length: Optional text sequence length used when running the
            language model. ``None`` lets the concrete subclass set a
            model-specific default.
    """
    if isinstance(input_keys, str):
        input_keys = [input_keys]
    token_keys = {
        SampleKey.TOKENIZED_OBSERVATIONS.value,
        SampleKey.IS_PAD_OBSERVATION.value,
    }
    camera_keys = [key for key in input_keys if key not in token_keys]
    all_keys = camera_keys + [
        SampleKey.TOKENIZED_OBSERVATIONS.value,
        SampleKey.IS_PAD_OBSERVATION.value,
    ]
    specification = InputSpecification(
        keys=all_keys,
        required_camera_modalities=[CameraModality.RGB],
        required=[SampleKey.TOKENIZED_OBSERVATIONS.value],
        requires_tokenized=True,
    )
    super().__init__(
        input_specification=specification,
        pretrained=pretrained,
        frozen=frozen,
        model_dtype=model_dtype,
    )
    self.camera_keys = camera_keys
    self.is_multi_camera = len(camera_keys) > 1
    self._setup_language_keys(
        output_modality=EncoderOutputKeys.FUSED_RGB_LANGUAGE.value
    )
    if max_text_length is not None:
        self.max_text_length = max_text_length

total_image_tokens property

total_image_tokens

Total image tokens across all cameras.

layers property

layers

Language-model transformer layers used by VLA backbones.

rotary_embedding property

rotary_embedding

Rotary embedding module used by the language-model transformer layers.

text_config property

text_config

Text-model config for constructing VLA expert layers.

embed_input_ids

embed_input_ids(token_ids)

Embed language-vocabulary token IDs with the VLM language tower.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def embed_input_ids(self, token_ids: torch.Tensor) -> torch.Tensor:
    """Embed language-vocabulary token IDs with the VLM language tower."""
    return self._embed_language(token_ids=token_ids)

forward_language_model

forward_language_model(input_ids=None, inputs_embeds=None, attention_mask=None, past_key_values=None, use_cache=False, cache_position=None, position_ids=None, output_hidden_states=True)

Run the VLM language tower over caller-provided token inputs.

Parameters:

Name	Type	Description	Default
input_ids	Tensor | None	Optional token IDs with shape (B, S).	None
inputs_embeds	Tensor | None	Optional token embeddings with shape (B, S, D).	None
attention_mask	Tensor | None	Optional language-model attention mask.	None
past_key_values	Cache | tuple[tuple[Tensor, ...], ...] | None	Optional cached key/value tensors.	None
use_cache	bool	Whether to return/update cached key/value tensors.	False
cache_position	Tensor | None	Optional HuggingFace KV-cache slots for the tokens in this call. During cached decoding, if the prefix has length P, the next token uses cache slot P so its key/value is appended after the prefix.	None
position_ids	Tensor | None	Optional positions for the language model positional encoding, with shape (B, S). These should count real tokens, not padding: [PAD, PAD, t0, t1] should pass [0, 0, 0, 1] so t0 and t1 get positions 0 and 1.	None
output_hidden_states	bool	Whether to return hidden states.	True

Returns:

Type	Description
CausalLanguageModelOutput	Causal language-model output with logits shape (B, S, V).

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def forward_language_model(
    self,
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    past_key_values: Cache | tuple[tuple[torch.Tensor, ...], ...] | None = None,
    use_cache: bool = False,
    cache_position: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    output_hidden_states: bool = True,
) -> CausalLanguageModelOutput:
    """Run the VLM language tower over caller-provided token inputs.

    Args:
        input_ids: Optional token IDs with shape ``(B, S)``.
        inputs_embeds: Optional token embeddings with shape ``(B, S, D)``.
        attention_mask: Optional language-model attention mask.
        past_key_values: Optional cached key/value tensors.
        use_cache: Whether to return/update cached key/value tensors.
        cache_position: Optional HuggingFace KV-cache slots for the tokens
            in this call. During cached decoding, if the prefix has length
            ``P``, the next token uses cache slot ``P`` so its key/value is
            appended after the prefix.
        position_ids: Optional positions for the language model positional
            encoding, with shape ``(B, S)``. These should count real tokens,
            not padding: ``[PAD, PAD, t0, t1]`` should pass
            ``[0, 0, 0, 1]`` so ``t0`` and ``t1`` get positions ``0`` and
            ``1``.
        output_hidden_states: Whether to return hidden states.

    Returns:
        Causal language-model output with logits shape ``(B, S, V)``.
    """
    language_model_inputs = {
        "input_ids": input_ids,
        "inputs_embeds": inputs_embeds,
        "attention_mask": attention_mask,
        "past_key_values": past_key_values,
        "use_cache": use_cache,
        "output_hidden_states": output_hidden_states,
        "return_dict": True,
    }
    if cache_position is not None:
        language_model_inputs["cache_position"] = cache_position
    if position_ids is not None:
        language_model_inputs["position_ids"] = position_ids
    return self._get_language_model()(**language_model_inputs)

encode

encode(inputs)

Embed images and text, concatenate, and run the language model.

Parameters:

Name	Type	Description	Default
inputs	dict[str, Tensor]	Camera images, tokenized text, and optional padding mask.	required

Returns:

Type	Description
dict[str, Tensor]	Dict with fused sequential features and padding mask.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def encode(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Embed images and text, concatenate, and run the language model.

    Args:
        inputs: Camera images, tokenized text, and optional padding mask.

    Returns:
        Dict with fused sequential features and padding mask.
    """
    inputs_embeds, full_padding_mask = self._assemble_multimodal_embeddings(
        inputs=inputs
    )
    full_attention_mask = (~full_padding_mask).to(torch.long)
    lm_outputs = self._get_language_model()(
        inputs_embeds=inputs_embeds,
        attention_mask=full_attention_mask,
    )
    fused_features = lm_outputs.last_hidden_state
    return {
        EncoderOutputKeys.FUSED_RGB_LANGUAGE.value: fused_features,
        self.padding_mask_name: full_padding_mask,
    }

forward

forward(inputs)

Run the VLM on raw inputs with or without an observation horizon.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def forward(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Run the VLM on raw inputs with or without an observation horizon."""
    temporal = self._flatten_temporal(inputs=inputs)
    if temporal is None:
        return self.encode(inputs=inputs)
    flattened, batch_size, temporal_length = temporal
    outputs = self.encode(inputs=flattened)
    return self._unflatten_temporal(
        outputs=outputs,
        batch_size=batch_size,
        temporal_length=temporal_length,
    )

build_prefix

build_prefix(inputs)

Build VLM prefix input embeddings for VLA decoder execution.

The prefix is the pre-transformer multimodal input sequence. The consuming decoder runs the language-model layers over it exactly once itself (interleaved expert layers or causal-LM prefill), so the prefix must not already be language-model output states.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def build_prefix(
    self, inputs: dict[str, torch.Tensor]
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Build VLM prefix input embeddings for VLA decoder execution.

    The prefix is the pre-transformer multimodal input sequence. The
    consuming decoder runs the language-model layers over it exactly once
    itself (interleaved expert layers or causal-LM prefill), so the prefix
    must not already be language-model output states.
    """
    temporal = self._flatten_temporal(inputs=inputs)
    if temporal is None:
        return self._assemble_multimodal_embeddings(inputs=inputs)
    flattened, batch_size, temporal_length = temporal
    prefix, padding_mask = self._assemble_multimodal_embeddings(inputs=flattened)
    sequence_length, hidden_dimension = prefix.shape[1], prefix.shape[2]
    # (B*T, S, D) -> (B, T*S, D)
    prefix = prefix.reshape(
        batch_size, temporal_length * sequence_length, hidden_dimension
    )
    padding_mask = padding_mask.reshape(
        batch_size, temporal_length * sequence_length
    )
    return prefix, padding_mask

validate_input_metadata

validate_input_metadata(key, metadata)

Validate metadata for VLM camera and tokenized-text inputs.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def validate_input_metadata(self, key: str, metadata: BaseMetadata) -> str | None:
    """Validate metadata for VLM camera and tokenized-text inputs."""
    if key in self.camera_keys:
        if not isinstance(metadata, CameraMetadata):
            return f"Expected CameraMetadata for '{key}', got {type(metadata).__name__}"
    else:
        if isinstance(metadata, CameraMetadata):
            return (
                f"{type(self).__name__} cannot process image data for '{key}'. "
                f"Got CameraMetadata, expected tokenized text input."
            )
    return None

get_vocab_size

get_vocab_size()

Return the VLM language model vocabulary size.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def get_vocab_size(self) -> int:
    """Return the VLM language model vocabulary size."""
    return self._get_language_model().config.vocab_size

get_backbone_layers

get_backbone_layers()

Return the language-model transformer layers for interleaved decoding.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def get_backbone_layers(self) -> nn.ModuleList:
    """Return the language-model transformer layers for interleaved decoding."""
    return self._get_language_model().layers

get_rotary_embedding

get_rotary_embedding()

Return the language-model rotary positional encoding module.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def get_rotary_embedding(self) -> nn.Module:
    """Return the language-model rotary positional encoding module."""
    return self._get_language_model().rotary_emb

get_backbone_hidden_dim

get_backbone_hidden_dim()

Return the VLM language model hidden dimension.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def get_backbone_hidden_dim(self) -> int:
    """Return the VLM language model hidden dimension."""
    return self.hidden_dimension

get_text_config

get_text_config()

Return the VLM language model configuration.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

def get_text_config(self) -> PretrainedConfig:
    """Return the VLM language model configuration."""
    return self._get_language_model().config

extract_key_value staticmethod

extract_key_value(vlm_layer, hidden_states)

Project hidden states to key and value through a VLM layer's projections.

Does not apply RoPE — use extract_query_key_value when positional information is needed.

Parameters:

Name	Type	Description	Default
vlm_layer	Module	Pretrained VLM transformer layer.	required
hidden_states	Tensor	(B, S, D).	required

Returns:

Type	Description
tuple[Tensor, Tensor]	(keys, values) each (B, S, key_value_dimension).

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

@staticmethod
def extract_key_value(
    vlm_layer: nn.Module,
    hidden_states: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Project hidden states to key and value through a VLM layer's projections.

    Does not apply RoPE — use ``extract_query_key_value`` when positional
    information is needed.

    Args:
        vlm_layer: Pretrained VLM transformer layer.
        hidden_states: (B, S, D).

    Returns:
        (keys, values) each (B, S, key_value_dimension).
    """
    normalized = vlm_layer.input_layernorm(hidden_states)
    return (
        vlm_layer.self_attn.k_proj(normalized),
        vlm_layer.self_attn.v_proj(normalized),
    )

compute_rope staticmethod

compute_rope(rotary_embedding, hidden_states, position_ids)

Compute (cos, sin) RoPE components for given positions.

Parameters:

Name	Type	Description	Default
rotary_embedding	Module	The language-model rotary embedding module.	required
hidden_states	Tensor	Tensor whose dtype/device to match.	required
position_ids	Tensor	Position indices (B, S).	required

Returns:

Type	Description
tuple[Tensor, Tensor]	(cos, sin) each (B, 1, S, head_dim) for head broadcast.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

@staticmethod
def compute_rope(
    rotary_embedding: nn.Module,
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute (cos, sin) RoPE components for given positions.

    Args:
        rotary_embedding: The language-model rotary embedding module.
        hidden_states: Tensor whose dtype/device to match.
        position_ids: Position indices (B, S).

    Returns:
        (cos, sin) each (B, 1, S, head_dim) for head broadcast.
    """
    cos, sin = rotary_embedding(hidden_states, position_ids)
    # (B, S, head_dim) → (B, 1, S, head_dim) for head broadcast
    return cos.unsqueeze(1), sin.unsqueeze(1)

build_additive_attention_mask staticmethod

build_additive_attention_mask(attention_mask, dtype)

Convert a bool mask to the additive mask expected by HF decoder layers.

Parameters:

Name	Type	Description	Default
attention_mask	Tensor | None	Boolean attention mask where True means masked. Accepts any broadcast-compatible shape used by HF decoder layers, typically (B, 1, Q, K).	required
dtype	dtype	Floating dtype for the returned additive mask.	required

Returns:

Type	Description
Tensor | None	Additive attention mask with 0 for valid locations and
Tensor | None	torch.finfo(dtype).min for masked locations. Returns None
Tensor | None	only when no mask is provided. An all-visible mask is returned as
Tensor | None	explicit zeros: HF decoder layers fall back to causal attention
Tensor | None	when given None, which would silently override the intended
Tensor | None	bidirectional structure.

Raises:

Type	Description
ValueError	If dtype is not a floating-point dtype.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

@staticmethod
def build_additive_attention_mask(
    attention_mask: torch.Tensor | None,
    dtype: torch.dtype,
) -> torch.Tensor | None:
    """Convert a bool mask to the additive mask expected by HF decoder layers.

    Args:
        attention_mask: Boolean attention mask where ``True`` means masked.
            Accepts any broadcast-compatible shape used by HF decoder layers,
            typically ``(B, 1, Q, K)``.
        dtype: Floating dtype for the returned additive mask.

    Returns:
        Additive attention mask with ``0`` for valid locations and
        ``torch.finfo(dtype).min`` for masked locations. Returns ``None``
        only when no mask is provided. An all-visible mask is returned as
        explicit zeros: HF decoder layers fall back to causal attention
        when given ``None``, which would silently override the intended
        bidirectional structure.

    Raises:
        ValueError: If ``dtype`` is not a floating-point dtype.
    """
    if attention_mask is None:
        return None
    if not torch.empty((), dtype=dtype).is_floating_point():
        raise ValueError(f"dtype must be floating point, got {dtype}.")
    boolean_mask = attention_mask.to(dtype=torch.bool)
    additive_mask = torch.zeros(
        boolean_mask.shape,
        dtype=dtype,
        device=boolean_mask.device,
    )
    return additive_mask.masked_fill(boolean_mask, torch.finfo(dtype).min)

extract_key_value_with_rope staticmethod

extract_key_value_with_rope(vlm_layer, hidden_states, rotary_embedding, position_ids)

Project hidden states to K/V with RoPE applied to keys.

Parameters:

Name	Type	Description	Default
vlm_layer	Module	Pretrained VLM transformer layer.	required
hidden_states	Tensor	(B, S, D).	required
rotary_embedding	Module	The language-model rotary embedding module.	required
position_ids	Tensor	(B, S) position indices.	required

Returns:

Type	Description
Tensor	(keys, values) — keys have RoPE applied and are flattened to
Tensor	(B, S, key_value_dimension), values are flat (B, S, key_value_dimension).

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

@staticmethod
def extract_key_value_with_rope(
    vlm_layer: nn.Module,
    hidden_states: torch.Tensor,
    rotary_embedding: nn.Module,
    position_ids: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Project hidden states to K/V with RoPE applied to keys.

    Args:
        vlm_layer: Pretrained VLM transformer layer.
        hidden_states: (B, S, D).
        rotary_embedding: The language-model rotary embedding module.
        position_ids: (B, S) position indices.

    Returns:
        (keys, values) — keys have RoPE applied and are flattened to
        (B, S, key_value_dimension), values are flat (B, S, key_value_dimension).
    """
    normalized = vlm_layer.input_layernorm(hidden_states)
    attention = vlm_layer.self_attn
    batch_size, sequence_length, _ = normalized.shape
    head_dimension = attention.head_dim
    number_of_key_value_heads = attention.config.num_key_value_heads
    keys_flat = attention.k_proj(normalized)
    values_flat = attention.v_proj(normalized)
    keys_headed = keys_flat.view(
        batch_size, sequence_length, number_of_key_value_heads, head_dimension
    ).transpose(1, 2)
    cos, sin = GenerativeVLM.compute_rope(
        rotary_embedding=rotary_embedding,
        hidden_states=keys_headed,
        position_ids=position_ids[:, :sequence_length],
    )
    keys_headed = RotaryPositionalEncoding.apply_rotation_half(
        tensor=keys_headed, sine=sin, cosine=cos
    )
    keys_with_rope = (
        keys_headed.transpose(1, 2)
        .contiguous()
        .view(batch_size, sequence_length, -1)
    )
    return keys_with_rope, values_flat

extract_query_key_value staticmethod

extract_query_key_value(vlm_layer, hidden_states, rotary_embedding, position_ids)

Project hidden states to Q/K/V and apply rotary positional encoding.

Works for any HF decoder layer that exposes input_layernorm, self_attn.{q,k,v}_proj, head_dim, and config.num_{attention,key_value}_heads.

Parameters:

Name	Type	Description	Default
vlm_layer	Module	Pretrained VLM transformer layer.	required
hidden_states	Tensor	(B, S, D).	required
rotary_embedding	Module	The language-model rotary embedding module.	required
position_ids	Tensor	(B, S) position indices.	required

Returns:

Name	Type	Description
	Tensor	(query, key, value) with RoPE on Q and K.
Shapes	Tensor	Q (B, H, S, D_head), K (B, KV_H, S, D_head), V same as K.

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

@staticmethod
def extract_query_key_value(
    vlm_layer: nn.Module,
    hidden_states: torch.Tensor,
    rotary_embedding: nn.Module,
    position_ids: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Project hidden states to Q/K/V and apply rotary positional encoding.

    Works for any HF decoder layer that exposes ``input_layernorm``,
    ``self_attn.{q,k,v}_proj``, ``head_dim``, and
    ``config.num_{attention,key_value}_heads``.

    Args:
        vlm_layer: Pretrained VLM transformer layer.
        hidden_states: (B, S, D).
        rotary_embedding: The language-model rotary embedding module.
        position_ids: (B, S) position indices.

    Returns:
        (query, key, value) with RoPE on Q and K.
        Shapes: Q (B, H, S, D_head), K (B, KV_H, S, D_head), V same as K.
    """
    normalized = vlm_layer.input_layernorm(hidden_states)
    attention = vlm_layer.self_attn
    batch_size, sequence_length, _ = normalized.shape
    head_dimension = attention.head_dim
    number_of_heads = attention.config.num_attention_heads
    number_of_key_value_heads = attention.config.num_key_value_heads
    query = (
        attention.q_proj(normalized)
        .view(batch_size, sequence_length, number_of_heads, head_dimension)
        .transpose(1, 2)
    )
    key = (
        attention.k_proj(normalized)
        .view(
            batch_size, sequence_length, number_of_key_value_heads, head_dimension
        )
        .transpose(1, 2)
    )
    value = (
        attention.v_proj(normalized)
        .view(
            batch_size, sequence_length, number_of_key_value_heads, head_dimension
        )
        .transpose(1, 2)
    )
    cos, sin = GenerativeVLM.compute_rope(
        rotary_embedding=rotary_embedding,
        hidden_states=value,
        position_ids=position_ids[:, :sequence_length],
    )
    query = RotaryPositionalEncoding.apply_rotation_half(
        tensor=query, sine=sin, cosine=cos
    )
    key = RotaryPositionalEncoding.apply_rotation_half(
        tensor=key, sine=sin, cosine=cos
    )
    return query, key, value

apply_residual_feedforward staticmethod

apply_residual_feedforward(vlm_layer, vlm_residual, vlm_attention_output)

Complete a transformer layer after externally computed attention.

Called by interleaved decoders that run joint attention outside the VLM layer and need it to finish: O-projection → normalization → residual add → feedforward → residual add.

Handles both Llama-style layers (single post-attention norm) and Gemma2-style layers (sandwich norms around sublayer outputs).

Parameters:

Name	Type	Description	Default
vlm_layer	Module	Pretrained VLM transformer layer.	required
vlm_residual	Tensor	Hidden states before attention (B, S, D).	required
vlm_attention_output	Tensor	Raw attention output before O-projection (B, S, H*D_head).	required

Returns:

Type	Description
Tensor	Updated hidden states (B, S, D).

Source code in src/versatil/models/decoding/generative_language_models/vision_language/base.py

@staticmethod
def apply_residual_feedforward(
    vlm_layer: nn.Module,
    vlm_residual: torch.Tensor,
    vlm_attention_output: torch.Tensor,
) -> torch.Tensor:
    """Complete a transformer layer after externally computed attention.

    Called by interleaved decoders that run joint attention outside the VLM
    layer and need it to finish: O-projection → normalization → residual
    add → feedforward → residual add.

    Handles both Llama-style layers (single post-attention norm) and
    Gemma2-style layers (sandwich norms around sublayer outputs).

    Args:
        vlm_layer: Pretrained VLM transformer layer.
        vlm_residual: Hidden states before attention (B, S, D).
        vlm_attention_output: Raw attention output before O-projection (B, S, H*D_head).

    Returns:
        Updated hidden states (B, S, D).
    """
    hidden_states = vlm_layer.self_attn.o_proj(vlm_attention_output)
    if hasattr(vlm_layer, "pre_feedforward_layernorm"):
        hidden_states = vlm_layer.post_attention_layernorm(hidden_states)
        hidden_states = vlm_residual + hidden_states
        residual = hidden_states
        hidden_states = vlm_layer.pre_feedforward_layernorm(hidden_states)
        hidden_states = vlm_layer.mlp(hidden_states)
        hidden_states = vlm_layer.post_feedforward_layernorm(hidden_states)
    else:
        hidden_states = vlm_residual + hidden_states
        residual = hidden_states
        hidden_states = vlm_layer.post_attention_layernorm(hidden_states)
        hidden_states = vlm_layer.mlp(hidden_states)
    return residual + hidden_states