positional_encoding

positional_encoding

Positional encoding for GPT transformer.

create_positional_encoding

create_positional_encoding(encoding_type, embedding_dimension, maximum_sequence_length, number_of_heads=None, base_frequency=10000.0, learnable_frequencies=False)

Factory function to create positional encoding.

Parameters:

Name	Type	Description	Default
encoding_type	str	Type of encoding (use PositionalEncodingType enum values)	required
embedding_dimension	int	Model embedding dimension	required
maximum_sequence_length	int	Maximum sequence length	required
number_of_heads	int | None	Number of attention heads (required for RoPE)	None
base_frequency	float	Base frequency for RoPE	10000.0
learnable_frequencies	bool	Whether to make RoPE frequencies learnable	False

Returns:

Type	Description
PositionalEncoding1D | RotaryPositionalEncoding1D	Positional encoding module

Raises:

Type	Description
ValueError	If encoding_type is not supported or required args missing

Source code in src/versatil/models/layers/transformer/positional_encoding.py

def create_positional_encoding(
    encoding_type: str,
    embedding_dimension: int,
    maximum_sequence_length: int,
    number_of_heads: int | None = None,
    base_frequency: float = 10000.0,
    learnable_frequencies: bool = False,
) -> PositionalEncoding1D | RotaryPositionalEncoding1D:
    """Factory function to create positional encoding.

    Args:
        encoding_type: Type of encoding (use PositionalEncodingType enum values)
        embedding_dimension: Model embedding dimension
        maximum_sequence_length: Maximum sequence length
        number_of_heads: Number of attention heads (required for RoPE)
        base_frequency: Base frequency for RoPE
        learnable_frequencies: Whether to make RoPE frequencies learnable

    Returns:
        Positional encoding module

    Raises:
        ValueError: If encoding_type is not supported or required args missing
    """
    if encoding_type == PositionalEncodingType.SINUSOIDAL.value:
        return SinusoidalPositionalEncoding1D(
            embedding_dimension=embedding_dimension,
            maximum_sequence_length=maximum_sequence_length,
        )
    elif encoding_type == PositionalEncodingType.LEARNED.value:
        return LearnedPositionalEncoding1D(
            embedding_dimension=embedding_dimension,
            maximum_sequence_length=maximum_sequence_length,
        )
    elif encoding_type == PositionalEncodingType.ROPE.value:
        if number_of_heads is None:
            raise ValueError("number_of_heads is required for RoPE positional encoding")
        return RotaryPositionalEncoding1D(
            embedding_dimension=embedding_dimension,
            number_of_heads=number_of_heads,
            base_frequency=base_frequency,
            learnable_frequencies=learnable_frequencies,
        )
    else:
        raise ValueError(
            f"Unsupported positional encoding type: {encoding_type}. "
            f"Must be one of {[e.value for e in PositionalEncodingType]}."
        )

apply_rope_positional_encoding

apply_rope_positional_encoding(queries, keys, positional_encoding, cache_position=0)

Apply positional encoding to queries and keys.

Handles both Sinusoidal (added to embeddings) and RoPE (applied via rotation).

Parameters:

Name	Type	Description	Default
queries	Tensor	Query tensor (B, number_of_heads, query_len, head_dim)	required
keys	Tensor	Key tensor (B, number_of_heads, key_len, head_dim) including cached keys	required
positional_encoding	Module	Positional encoding module	required
cache_position	int	Starting position for queries (0 for initial forward, cache_len for generation)	0

Returns:

Type	Description
tuple[Tensor, Tensor]	Tuple of (queries_with_pos, keys_with_pos)

Source code in src/versatil/models/layers/transformer/positional_encoding.py

def apply_rope_positional_encoding(
    queries: torch.Tensor,
    keys: torch.Tensor,
    positional_encoding: nn.Module,
    cache_position: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply positional encoding to queries and keys.

    Handles both Sinusoidal (added to embeddings) and RoPE (applied via rotation).

    Args:
        queries: Query tensor (B, number_of_heads, query_len, head_dim)
        keys: Key tensor (B, number_of_heads, key_len, head_dim) including cached keys
        positional_encoding: Positional encoding module
        cache_position: Starting position for queries (0 for initial forward, cache_len for generation)

    Returns:
        Tuple of (queries_with_pos, keys_with_pos)
    """
    if isinstance(positional_encoding, RotaryPositionalEncoding1D):
        # RoPE: apply rotation to Q and K
        # Both queries and keys are new segments at the same positions
        sequence_length = queries.shape[2]

        # Compute rotation components for positions [cache_position, ..., cache_position + seq_len - 1]
        sine, cosine = positional_encoding.compute_rotation_components(
            cache_position + sequence_length
        )
        sine = sine[cache_position : cache_position + sequence_length]
        cosine = cosine[cache_position : cache_position + sequence_length]

        # Expand for batch and heads: (seq_len, head_dim) -> (1, 1, seq_len, head_dim)
        sine = sine.unsqueeze(0).unsqueeze(0)
        cosine = cosine.unsqueeze(0).unsqueeze(0)

        # Apply rotation to both queries and keys
        queries = positional_encoding.apply_rotation(queries, sine, cosine)
        keys = positional_encoding.apply_rotation(keys, sine, cosine)

        return queries, keys

    else:
        # Unknown type - return unchanged
        logging.warning(
            "Positional encoding module is not an instance of RotaryPositionalEncoding. Skipping."
        )
        return queries, keys