lalamo 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

lalamo/modules/attention.py

@@ -9,9 +9,10 @@ from jax import numpy as jnp
 from jax import vmap
 from jaxtyping import Array, Bool, DTypeLike, Float, Int, PRNGKeyArray
 
+from lalamo.common import dummy_array
 from lalamo.modules.normalization import RMSNorm, RMSNormConfig
 
-from .common import AttentionType, LalamoModule, ParameterTree, WeightLayout
+from .common import AttentionType, LalamoModule, ParameterTree
 from .kv_cache import DynamicKVCacheLayer, KVCacheLayer, StaticKVCacheLayer
 from .linear import LinearBase, LinearConfig
 from .rope import PositionalEmbeddings
@@ -44,8 +45,6 @@ def _soft_capped_attention_kernel(
 ) -> Float[Array, "dst_tokens heads head_channels"]:
     _, num_heads, head_dim = queries.shape
     _, num_groups, _ = keys.shape
-    if scale is None:
-        scale = head_dim**-0.5
     group_size = num_heads // num_groups
     keys = _repeat_kv(keys, group_size)
     values = _repeat_kv(values, group_size)
@@ -59,7 +58,11 @@ def _soft_capped_attention_kernel(
     if mask is not None:
         attention_logits = jnp.where(mask, attention_logits, jnp.array(float("-inf"), dtype=attention_logits.dtype))
 
-    attention_logits = attention_logits * scale
+    if scale is None:
+        scale_val = head_dim**-0.5
+    else:
+        scale_val = float(scale)
+    attention_logits = attention_logits * scale_val
     attention_logits = apply_soft_capping(attention_logits, logit_soft_cap)
     attention_weights = jax.nn.softmax(attention_logits, axis=-1)
     return einsum(
@@ -70,7 +73,7 @@ def _soft_capped_attention_kernel(
 
 
 class AttentionResult(NamedTuple):
-    outputs: Float[Array, "suffix_tokens channels"]
+    outputs: Float[Array, "*batch suffix_tokens channels"]
     kv_cache: KVCacheLayer | None = None
 
 
@@ -83,6 +86,7 @@ class AttentionConfig:
     key_norm_config: RMSNormConfig | None
 
     logit_soft_cap: float | None
+    has_sinks: bool
     has_qkv_biases: bool
     has_out_biases: bool
 
@@ -130,12 +134,18 @@ class AttentionConfig:
         else:
             key_norm = None
 
+        if self.has_sinks:
+            sinks = jnp.zeros((num_heads,), dtype=qkv_projection.activation_precision)
+        else:
+            sinks = None
+
         return Attention(
             self,
             qkv_projection=qkv_projection,
             out_projection=out_projection,
             query_norm=query_norm,
             key_norm=key_norm,
+            sinks=sinks,
             num_heads=num_heads,
             num_groups=num_groups,
             head_dim=head_dim,
@@ -183,12 +193,18 @@ class AttentionConfig:
         else:
             key_norm = None
 
+        if self.has_sinks:
+            sinks = dummy_array(num_heads, qkv_projection.activation_precision)
+        else:
+            sinks = None
+
         return Attention(
             self,
             qkv_projection=qkv_projection,
             out_projection=out_projection,
             query_norm=query_norm,
             key_norm=key_norm,
+            sinks=sinks,
             num_heads=num_heads,
             num_groups=num_groups,
             head_dim=head_dim,
@@ -205,6 +221,8 @@ class Attention(LalamoModule[AttentionConfig]):
     query_norm: RMSNorm | None
     key_norm: RMSNorm | None
 
+    sinks: Float[Array, " heads"] | None
+
     num_heads: int = eqx.field(static=True)
     num_groups: int = eqx.field(static=True)
     head_dim: int = eqx.field(static=True)
@@ -234,6 +252,10 @@ class Attention(LalamoModule[AttentionConfig]):
     def attention_type(self) -> AttentionType:
         return AttentionType.SLIDING_WINDOW if self.sliding_window_size is not None else AttentionType.GLOBAL
 
+    @property
+    def has_sinks(self) -> bool:
+        return self.sinks is not None
+
     def __post_init__(self) -> None:
         if self.qkv_projection.has_biases != self.config.has_qkv_biases:
             raise ValueError(
@@ -285,6 +307,12 @@ class Attention(LalamoModule[AttentionConfig]):
                 f" ({self.num_groups} * {self.head_dim} = {self.num_groups * self.head_dim}),"
                 f" got {v_output_dim}",
             )
+        if self.sinks is not None:
+            (num_sink_heads,) = self.sinks.shape
+            if num_sink_heads != self.num_heads:
+                raise ValueError(
+                    f"Number of sink heads must be equal to number of heads ({self.num_heads}), got {num_sink_heads}",
+                )
 
     @eqx.filter_jit
     def __call__(
@@ -325,12 +353,22 @@ class Attention(LalamoModule[AttentionConfig]):
         keys = apply_positional_embeddings(keys)
 
         if kv_cache is None:
-            updated_kv_cache = DynamicKVCacheLayer.init(keys, values, length=length_without_padding)
+            updated_kv_cache = DynamicKVCacheLayer.init(self.has_sinks, keys, values, length=length_without_padding)
         else:
             updated_kv_cache = kv_cache.extend(keys, values, added_length=length_without_padding)
 
         num_suffix_tokens, _, _ = queries.shape
-        mask = updated_kv_cache.attention_mask(num_suffix_tokens, self.is_causal, self.sliding_window_size)
+        mask = updated_kv_cache.attention_mask(
+            num_suffix_tokens,
+            self.is_causal,
+            length_without_padding,
+            self.sliding_window_size,
+        )
+        if self.sinks is not None:
+            sink_bias = jnp.zeros((self.num_heads, *mask.shape), dtype=queries.dtype)
+            sink_bias = sink_bias.at[:, :, 0].set(self.sinks[:, None])
+        else:
+            sink_bias = None
 
         if self.config.logit_soft_cap is not None:
             attention_output = _soft_capped_attention_kernel(
@@ -346,6 +384,7 @@ class Attention(LalamoModule[AttentionConfig]):
                 queries,
                 updated_kv_cache.keys,
                 updated_kv_cache.values,
+                bias=sink_bias,
                 mask=mask,
                 scale=self.scale,
             )
@@ -366,41 +405,55 @@ class Attention(LalamoModule[AttentionConfig]):
         )
 
     def init_static_kv_cache(self, capacity: int) -> StaticKVCacheLayer:
-        return StaticKVCacheLayer.empty(capacity, self.num_groups, self.head_dim, self.activation_precision)
-
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
-        result = dict(
-            qkv_projection=self.qkv_projection.export_weights(weight_layout),
-            out_projection=self.out_projection.export_weights(weight_layout),
+        return StaticKVCacheLayer.empty(
+            self.has_sinks,
+            capacity,
+            self.num_groups,
+            self.head_dim,
+            self.activation_precision,
         )
+
+    def export_weights(self) -> ParameterTree:
+        result: dict[str, ParameterTree | Array] = {
+            "qkv_projection": self.qkv_projection.export_weights(),
+            "out_projection": self.out_projection.export_weights(),
+        }
         if self.query_norm is not None:
-            result["query_norm"] = self.query_norm.export_weights(weight_layout)
+            result["query_norm"] = self.query_norm.export_weights()
         if self.key_norm is not None:
-            result["key_norm"] = self.key_norm.export_weights(weight_layout)
+            result["key_norm"] = self.key_norm.export_weights()
+        if self.sinks is not None:
+            assert isinstance(self.sinks, Array)
+            result["sinks"] = self.sinks
         return result
 
     def import_weights(
         self,
         weights: ParameterTree[Array],
-        weight_layout: WeightLayout = WeightLayout.AUTO,
     ) -> Self:
         assert isinstance(weights, Mapping)
         assert isinstance(weights["qkv_projection"], Mapping)
         assert isinstance(weights["out_projection"], Mapping)
         if self.query_norm is not None:
             assert isinstance(weights["query_norm"], Mapping)
-            query_norm = self.query_norm.import_weights(weights["query_norm"], weight_layout)
+            query_norm = self.query_norm.import_weights(weights["query_norm"])
         else:
             query_norm = None
         if self.key_norm is not None:
             assert isinstance(weights["key_norm"], Mapping)
-            key_norm = self.key_norm.import_weights(weights["key_norm"], weight_layout)
+            key_norm = self.key_norm.import_weights(weights["key_norm"])
         else:
             key_norm = None
+        if self.sinks is not None:
+            assert isinstance(weights["sinks"], Array)
+            sinks = weights["sinks"]
+        else:
+            sinks = None
         return replace(
             self,
-            qkv_projection=self.qkv_projection.import_weights(weights["qkv_projection"], weight_layout),
-            out_projection=self.out_projection.import_weights(weights["out_projection"], weight_layout),
+            qkv_projection=self.qkv_projection.import_weights(weights["qkv_projection"]),
+            out_projection=self.out_projection.import_weights(weights["out_projection"]),
             query_norm=query_norm,
             key_norm=key_norm,
+            sinks=sinks,
         )

lalamo/modules/common.py

@@ -6,79 +6,31 @@ from typing import Self
 
 import equinox as eqx
 from cattrs import Converter
-from einops import rearrange
 from jax import numpy as jnp
-from jaxtyping import Array, DTypeLike, Float
+from jaxtyping import Array, DTypeLike
 
 from lalamo.common import ParameterTree
 
 __all__ = [
     "AttentionType",
     "DummyUnionMember",
+    "ForwardPassMode",
     "LalamoModule",
     "config_converter",
-    "from_layout",
-    "into_layout",
     "register_config_union",
 ]
 
 
-class WeightLayout(Enum):
-    AUTO = "auto"
-    INPUT_OUTPUT = "input_output"
-    OUTPUT_INPUT = "output_input"
-
-    def __str__(self) -> str:
-        match self:
-            case WeightLayout.AUTO:
-                return "auto"
-            case WeightLayout.INPUT_OUTPUT:
-                return "(input, output)"
-            case WeightLayout.OUTPUT_INPUT:
-                return "(output, input)"
-
-
-_DEFAULT_WEIGHT_LAYOUT = WeightLayout.INPUT_OUTPUT
-
-
-def into_layout(
-    weights: Float[Array, "in_channels out_channels"],
-    layout: WeightLayout,
-) -> Float[Array, "in_channels out_channels"] | Float[Array, "out_channels in_channels"]:
-    if layout == WeightLayout.AUTO:
-        layout = _DEFAULT_WEIGHT_LAYOUT
-    match layout:
-        case WeightLayout.OUTPUT_INPUT:
-            return weights
-        case WeightLayout.INPUT_OUTPUT:
-            return rearrange(
-                weights,
-                "total_out_channels in_channels -> in_channels total_out_channels",
-            )
-
-
-def from_layout(
-    weights: ParameterTree | Array,
-    layout: WeightLayout,
-) -> Array:
-    assert isinstance(weights, Array)
-    if layout == WeightLayout.AUTO:
-        layout = _DEFAULT_WEIGHT_LAYOUT
-    match layout:
-        case WeightLayout.OUTPUT_INPUT:
-            return weights
-        case WeightLayout.INPUT_OUTPUT:
-            return rearrange(
-                weights,
-                "in_channels total_out_channels -> total_out_channels in_channels",
-            )
-
-
 class AttentionType(Enum):
     GLOBAL = "global"
     SLIDING_WINDOW = "sliding_window"
 
 
+class ForwardPassMode(Enum):
+    MULTI_TOKEN = "multi_token"
+    SINGLE_TOKEN = "single_token"
+
+
 class LalamoModule[ConfigT](eqx.Module):
     config: ConfigT = eqx.field(static=True)
 
@@ -87,13 +39,12 @@ class LalamoModule[ConfigT](eqx.Module):
     def activation_precision(self) -> DTypeLike: ...
 
     @abstractmethod
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree[Array]: ...
+    def export_weights(self) -> ParameterTree[Array]: ...
 
     @abstractmethod
     def import_weights(
         self,
         weights: ParameterTree[Array],
-        weight_layout: WeightLayout = WeightLayout.AUTO,
     ) -> Self: ...
 
 

lalamo/modules/decoder.py

@@ -8,9 +8,10 @@ from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
 from lalamo.common import ParameterTree
+from lalamo.modules.utils import vmap_twice
 
-from .common import AttentionType, LalamoModule, WeightLayout
-from .decoder_layer import DecoderLayer, DecoderLayerConfig, DecoderLayerResult
+from .common import AttentionType, ForwardPassMode, LalamoModule
+from .decoder_layer import DecoderLayer, DecoderLayerConfig, DecoderLayerForwardPassConfig, DecoderLayerResult
 from .embedding import EmbeddingBase, EmbeddingConfig
 from .kv_cache import KVCache
 from .normalization import RMSNorm, RMSNormConfig
@@ -20,13 +21,17 @@ __all__ = [
     "Decoder",
     "DecoderActivationTrace",
     "DecoderConfig",
+    "DecoderForwardPassConfig",
     "DecoderResult",
 ]
 
 
+type DecoderForwardPassConfig = DecoderLayerForwardPassConfig
+
+
 class DecoderActivationTrace(eqx.Module):
-    token_ids: Int[Array, " suffix_tokens"]
-    token_positions: Int[Array, " suffix_tokens"]
+    token_ids: Int[Array, "batch suffix_tokens"]
+    token_positions: Int[Array, "batch suffix_tokens"]
     kv_cache: KVCache | None
 
     local_positional_embeddings: PositionalEmbeddings
@@ -34,7 +39,7 @@ class DecoderActivationTrace(eqx.Module):
 
     layer_results: tuple[DecoderLayerResult, ...]
 
-    output_norm: Float[Array, "suffix_tokens channels"]
+    output_norm: Float[Array, "batch suffix_tokens channels"]
 
     def export(self) -> ParameterTree:
         result = dict(
@@ -51,7 +56,7 @@ class DecoderActivationTrace(eqx.Module):
 
 
 class DecoderResult(eqx.Module):
-    logits: Float[Array, "suffix_tokens channels"]
+    logits: Float[Array, "batch suffix_tokens channels"]
     updated_kv_cache: KVCache | None = None
     activation_trace: DecoderActivationTrace | None = None
 
@@ -167,13 +172,9 @@ class DecoderConfig:
         )
 
         if self.local_rope_config:
-            assert self.sliding_window_sizes is not None
-            max_sliding_window_size = max(
-                window_size for window_size in self.sliding_window_sizes if window_size is not None
-            )
             local_rope = self.local_rope_config.init(
                 head_dim=self.head_dim,
-                num_timesteps=max(max_sliding_window_size, self.context_length),
+                num_timesteps=self.context_length,
             )
         else:
             local_rope = None
@@ -219,19 +220,31 @@ class Decoder(LalamoModule[DecoderConfig]):
     @eqx.filter_jit
     def __call__(
         self,
-        token_ids: Int[Array, " suffix_tokens"],
-        token_positions: Int[Array, " suffix_tokens"],
+        token_ids: Int[Array, "batch suffix_tokens"],
+        token_positions: Int[Array, "batch suffix_tokens"],
         kv_cache: KVCache | None = None,
         return_updated_kv_cache: bool = False,
         return_activation_trace: bool = False,
-        length_without_padding: Int[Array, ""] | int | None = None,
+        lengths_without_padding: Int[Array, " batch"] | None = None,
+        forward_pass_mode: ForwardPassMode = ForwardPassMode.MULTI_TOKEN,
+        forward_pass_config: DecoderForwardPassConfig | None = None,
     ) -> DecoderResult:
+        if token_ids.ndim != 2:
+            raise ValueError(
+                f"token_ids must be a 2D arrays of size (batch_size, sequence_length), got {token_ids.shape}",
+            )
+        if token_positions.ndim != 2:
+            raise ValueError(
+                "token_positions must be a 2D arrays of size (batch_size, sequence_length),"
+                f" got {token_positions.shape}",
+            )
+
         maybe_kv_cache = kv_cache or ([None] * len(self.layers))
-        inner_features = self.embedding.embed(token_ids)
+        inner_features = vmap(self.embedding.embed)(token_ids)
 
-        global_positional_embeddings = self.global_rope(token_positions)
+        global_positional_embeddings = vmap(self.global_rope)(token_positions)
         if self.local_rope is not None:
-            local_positional_embeddings = self.local_rope(token_positions)
+            local_positional_embeddings = vmap(self.local_rope)(token_positions)
         else:
             local_positional_embeddings = global_positional_embeddings
 
@@ -249,14 +262,16 @@ class Decoder(LalamoModule[DecoderConfig]):
                 kv_cache=kv_cache_slice,
                 return_updated_kv_cache=return_updated_kv_cache,
                 return_activation_trace=return_activation_trace,
-                length_without_padding=length_without_padding,
+                lengths_without_padding=lengths_without_padding,
+                forward_pass_mode=forward_pass_mode,
+                forward_pass_config=forward_pass_config,
             )
             inner_features = layer_result.outputs
             layer_results.append(layer_result)
             updated_kv_cache_layers.append(layer_result.updated_kv_cache)
 
-        normalized_outputs = vmap(self.output_norm, in_axes=0)(inner_features)
-        logits = vmap(self.embedding.readout, in_axes=0)(normalized_outputs)
+        normalized_outputs = vmap_twice(self.output_norm)(inner_features)
+        logits = vmap_twice(self.embedding.readout)(normalized_outputs)
 
         if return_activation_trace:
             activation_trace = DecoderActivationTrace(
@@ -282,24 +297,23 @@ class Decoder(LalamoModule[DecoderConfig]):
             activation_trace=activation_trace,
         )
 
-    def init_static_kv_cache(self, capacity: int) -> KVCache:
-        return KVCache(layer.init_static_kv_cache(capacity) for layer in self.layers)
+    def init_static_kv_cache(self, batch_size: int, capacity: int) -> KVCache:
+        return KVCache(layer.init_static_kv_cache(batch_size, capacity) for layer in self.layers)
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+    def export_weights(self) -> ParameterTree:
         result = dict(
-            embedding=self.embedding.export_weights(weight_layout),
-            global_rope=self.global_rope.export_weights(weight_layout),
-            layers=[layer.export_weights(weight_layout) for layer in self.layers],
-            output_norm=self.output_norm.export_weights(weight_layout),
+            embedding=self.embedding.export_weights(),
+            global_rope=self.global_rope.export_weights(),
+            layers=[layer.export_weights() for layer in self.layers],
+            output_norm=self.output_norm.export_weights(),
         )
         if self.local_rope:
-            result["local_rope"] = self.local_rope.export_weights(weight_layout)
+            result["local_rope"] = self.local_rope.export_weights()
         return result
 
     def import_weights(
         self,
         weights: ParameterTree[Array],
-        weight_layout: WeightLayout = WeightLayout.AUTO,
     ) -> Self:
         assert isinstance(weights, Mapping)
         assert isinstance(weights["embedding"], Mapping)
@@ -308,19 +322,19 @@ class Decoder(LalamoModule[DecoderConfig]):
         assert isinstance(weights["output_norm"], Mapping)
         if self.local_rope:
             assert isinstance(weights["local_rope"], Mapping)
-            local_rope = self.local_rope.import_weights(weights["local_rope"], weight_layout)
+            local_rope = self.local_rope.import_weights(weights["local_rope"])
         else:
             local_rope = None
 
         layers = []
         for layer, layer_weights in zip(self.layers, weights["layers"], strict=True):
             assert isinstance(layer_weights, Mapping)
-            layers.append(layer.import_weights(layer_weights, weight_layout))
+            layers.append(layer.import_weights(layer_weights))
         return replace(
             self,
-            embedding=self.embedding.import_weights(weights["embedding"], weight_layout),
-            global_rope=self.global_rope.import_weights(weights["global_rope"], weight_layout),
+            embedding=self.embedding.import_weights(weights["embedding"]),
+            global_rope=self.global_rope.import_weights(weights["global_rope"]),
             layers=tuple(layers),
-            output_norm=self.output_norm.import_weights(weights["output_norm"], weight_layout),
+            output_norm=self.output_norm.import_weights(weights["output_norm"]),
             local_rope=local_rope,
         )