lalamo 0.6.4__py3-none-any.whl → 0.6.6__py3-none-any.whl

lalamo/modules/token_mixers/mamba.py

@@ -25,9 +25,65 @@ __all__ = [
     "Mamba2Result",
     "SeparableCausalConv",
     "SeparableCausalConvConfig",
+    "exp_segsum",
+    "fused_ssd_intra_chunk",
 ]
 
 
+def exp_segsum(x: Float[Array, "... T"]) -> Float[Array, "... T T"]:
+    """Compute exp(segsum(x)) as lower-triangular matrix using cumsum difference."""
+    seq_len = x.shape[-1]
+    cs = jnp.cumsum(x, axis=-1)
+    diff = cs[..., :, None] - cs[..., None, :]
+    mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=bool))
+    return jnp.where(mask, jnp.exp(diff), 0.0)
+
+
+def fused_ssd_intra_chunk(
+    a_cumsum: Float[Array, "groups heads_per_group chunks chunk_size"],
+    cb: Float[Array, "chunks chunk_size chunk_size groups"],
+    x: Float[Array, "chunks chunk_size groups heads_per_group head_dim"],
+) -> Float[Array, "chunks chunk_size groups heads_per_group head_dim"]:
+    """Compute intra-chunk diagonal block outputs for SSD.
+
+    Avoids materializing the full global L matrix by computing decay locally per (chunk, group, head).
+    """
+    groups, heads_per_group, chunks, chunk_size = a_cumsum.shape
+
+    def compute_one(
+        a_cs: Float[Array, " chunk_size"],
+        cb_slice: Float[Array, "chunk_size chunk_size"],
+        x_slice: Float[Array, "chunk_size head_dim"],
+    ) -> Float[Array, "chunk_size head_dim"]:
+        diff = a_cs[:, None] - a_cs[None, :]
+        mask = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.bool_))
+        decay_local = jnp.where(mask, jnp.exp(diff), 0.0)
+        weighted = decay_local * cb_slice
+        return weighted @ x_slice
+
+    def compute_chunk_group_head(chunk_idx: int, group_idx: int, head_idx: int) -> Float[Array, "chunk_size head_dim"]:
+        return compute_one(
+            a_cumsum[group_idx, head_idx, chunk_idx, :],
+            cb[chunk_idx, :, :, group_idx],
+            x[chunk_idx, :, group_idx, head_idx, :],
+        )
+
+    def over_heads(chunk_idx: int, group_idx: int) -> Float[Array, "heads_per_group chunk_size head_dim"]:
+        return jax.vmap(lambda head_idx: compute_chunk_group_head(chunk_idx, group_idx, head_idx))(
+            jnp.arange(heads_per_group),
+        )
+
+    def over_groups(chunk_idx: int) -> Float[Array, "groups heads_per_group chunk_size head_dim"]:
+        return jax.vmap(lambda group_idx: over_heads(chunk_idx, group_idx))(jnp.arange(groups))
+
+    result = jax.vmap(over_groups)(jnp.arange(chunks))
+
+    return rearrange(
+        result,
+        "chunks groups heads_per_group chunk_size head_dim -> chunks chunk_size groups heads_per_group head_dim",
+    )
+
+
 Mamba2Result = TokenMixerResult[Mamba2StateLayer]
 
 
@@ -156,6 +212,19 @@ class SeparableCausalConv(LalamoModule[SeparableCausalConvConfig]):
             updated_state,
         )
 
+    def step(
+        self,
+        token: Float[Array, " channels"],
+        state: Float[Array, "kernel_minus_1 channels"],
+    ) -> tuple[Float[Array, " channels"], Float[Array, "kernel_minus_1 channels"]]:
+        """Single-token conv update without full convolution overhead."""
+        full_input = jnp.concatenate([state, token[None, :]], axis=0)
+        output = einsum(full_input, self.weights, "kernel channels, channels kernel -> channels")
+        if self.biases is not None:
+            output = output + self.biases
+        new_state = jnp.concatenate([state[1:], token[None, :]], axis=0)
+        return output, new_state
+
     def export_weights(self) -> ParameterTree:
         result: dict[str, Array] = {"weights": self.weights}
         if self.biases is not None:
@@ -188,6 +257,8 @@ class Mamba2Config(TokenMixerConfigBase):
     has_in_biases: bool
     has_out_biases: bool
 
+    chunk_size: int = 256
+
     @property
     def inner_dim(self) -> int:
         return self.num_heads * self.head_dim
@@ -330,99 +401,258 @@ class Mamba2(TokenMixerBase[Mamba2Config, Mamba2StateLayer]):
                 f"Number of value heads ({self.num_heads}) must be divisible by number of groups ({self.num_groups})",
             )
 
-    def _scan(
+    def _step(
+        self,
+        values: Float[Array, "heads head_dim"],
+        keys: Float[Array, "groups state_dim"],
+        queries: Float[Array, "groups state_dim"],
+        dt_log: Float[Array, " heads"],
+        state: Float[Array, "heads head_dim state_dim"],
+    ) -> tuple[Float[Array, "heads head_dim"], Float[Array, "heads head_dim state_dim"]]:
+        """Single-token SSM state update without scan overhead."""
+        heads_per_group = self.num_heads // self.num_groups
+
+        dt = jax.nn.softplus(dt_log)
+        decay = jnp.exp(-dt)[:, None, None]
+        mix = dt[:, None, None]
+
+        keys_expanded = jnp.repeat(keys, heads_per_group, axis=0)
+        queries_expanded = jnp.repeat(queries, heads_per_group, axis=0)
+        values_norm = values / (dt[:, None] + 1e-8)
+
+        input_contribution = mix * values_norm[:, :, None] * keys_expanded[:, None, :]
+        new_state = decay * state + input_contribution
+        output = einsum(new_state, queries_expanded, "heads head_dim state_dim, heads state_dim -> heads head_dim")
+
+        return output, new_state
+
+    def _decode_step(
         self,
-        hidden_states: Float[Array, "suffix_tokens heads head_channels"],
-        input_projection: Float[Array, "suffix_tokens groups state_channels"],
-        output_projection: Float[Array, "suffix_tokens groups state_channels"],
-        time_delta_log: Float[Array, "suffix_tokens heads"],
-        initial_state: Float[Array, "heads head_channels state_channels"],
+        inputs: Float[Array, "1 channels"],
+        state: Mamba2StateLayer,
+    ) -> Mamba2Result:
+        """Optimized path for single-token decode without scan machinery."""
+        token = inputs[0]
+
+        conv_in, gate, dt_log = self.in_projection(token)
+        conv_out, new_conv_state = self.conv.step(conv_in, state.conv_state)
+        conv_activated = self.config.activation(conv_out)
+
+        values_flat, input_proj_flat, output_proj_flat = jnp.split(
+            conv_activated,
+            [self.inner_dim, self.inner_dim + self.num_groups * self.state_dim],
+        )
+        values = rearrange(values_flat, "(heads head_dim) -> heads head_dim", heads=self.num_heads)
+        keys = rearrange(input_proj_flat, "(groups state_dim) -> groups state_dim", groups=self.num_groups)
+        queries = rearrange(output_proj_flat, "(groups state_dim) -> groups state_dim", groups=self.num_groups)
+
+        y, new_ssm_state = self._step(values, keys, queries, dt_log, state.ssm_state)
+
+        y = y + self.skip_connection_weight[:, None] * values
+        y = rearrange(y, "heads head_dim -> (heads head_dim)")
+        gated = y * jax.nn.silu(gate + self.gate_bias)
+        (output,) = self.out_projection(gated)
+
+        return Mamba2Result(
+            outputs=output[None, :],
+            state=Mamba2StateLayer(new_conv_state, new_ssm_state),
+        )
+
+    def _chunked_scan(
+        self,
+        values: Float[Array, "suffix_tokens heads head_dim"],
+        keys: Float[Array, "suffix_tokens groups state_dim"],
+        queries: Float[Array, "suffix_tokens groups state_dim"],
+        dt: Float[Array, "suffix_tokens heads"],
+        initial_state: Float[Array, "heads head_dim state_dim"],
+        chunk_size: int,
         num_steps: Int[Array, ""] | int,
-    ) -> tuple[
-        Float[Array, "suffix_tokens heads head_channels"],
-        Float[Array, "heads head_channels state_channels"],
-    ]:
-        def scan_fn(
-            index_and_carry_state: tuple[Int[Array, ""], Float[Array, "heads head_channels state_channels"]],
-            step_inputs: tuple[
-                Float[Array, "heads head_channels"],
-                Float[Array, "groups state_channels"],
-                Float[Array, "groups state_channels"],
-                Float[Array, " heads"],
-            ],
-        ) -> tuple[
-            tuple[Int[Array, ""], Float[Array, "heads head_channels state_channels"]],
-            Float[Array, "heads head_channels"],
-        ]:
-            index, carry_state = index_and_carry_state
-            hidden_state_t, input_proj_t, output_proj_t, time_delta_log_t = step_inputs
-            dt = jax.nn.softplus(time_delta_log_t)[:, None]
-            heads_per_group = self.num_heads // self.num_groups
-
-            hidden_grouped = rearrange(
-                hidden_state_t,
-                "(groups heads) head_channels -> groups heads head_channels",
-                groups=self.num_groups,
-                heads=heads_per_group,
-            )
-            x_norm_grouped = hidden_grouped / (
-                dt.reshape(self.num_heads)[
-                    rearrange(
-                        jnp.arange(self.num_heads),
-                        "(groups heads)-> groups heads",
-                        groups=self.num_groups,
-                        heads=heads_per_group,
-                    )
-                ][:, :, None]
-                + 1e-8
-            )
+        d: Float[Array, " heads"] | None = None,
+        z: Float[Array, "suffix_tokens heads head_dim"] | None = None,
+        z_bias: Float[Array, "heads head_dim"] | None = None,
+    ) -> tuple[Float[Array, "suffix_tokens heads head_dim"], Float[Array, "heads head_dim state_dim"]]:
+        """Chunked parallel scan implementing the SSD algorithm."""
+        seq_len = values.shape[0]
+        num_steps = jnp.asarray(num_steps, dtype=jnp.int32)
+
+        pad_len = (chunk_size - seq_len % chunk_size) % chunk_size
+        if pad_len > 0:
+            values = jnp.pad(values, ((0, pad_len), (0, 0), (0, 0)))
+            keys = jnp.pad(keys, ((0, pad_len), (0, 0), (0, 0)))
+            queries = jnp.pad(queries, ((0, pad_len), (0, 0), (0, 0)))
+            dt = jnp.pad(dt, ((0, pad_len), (0, 0)))
+            if z is not None:
+                z = jnp.pad(z, ((0, pad_len), (0, 0), (0, 0)))
+
+        values_orig = values
+        keys_orig = keys
+        dt_orig = dt
+
+        padded_len = values.shape[0]
+        position_indices = jnp.arange(padded_len)
+        valid_mask = (position_indices < num_steps).astype(values.dtype)
+        values = values * valid_mask[:, None, None]
+        keys = keys * valid_mask[:, None, None]
+
+        values = rearrange(
+            values,
+            "(chunks chunk_size) (groups heads_per_group) head_dim"
+            " -> chunks chunk_size groups heads_per_group head_dim",
+            chunk_size=chunk_size,
+            groups=self.num_groups,
+        )
+        log_decay = rearrange(
+            -dt,
+            "(chunks chunk_size) (groups heads_per_group) -> groups heads_per_group chunks chunk_size",
+            chunk_size=chunk_size,
+            groups=self.num_groups,
+        )
+        keys_chunked = rearrange(
+            keys,
+            "(chunks chunk_size) groups state_dim -> chunks chunk_size groups state_dim",
+            chunk_size=chunk_size,
+        )
+        queries_chunked = rearrange(
+            queries,
+            "(chunks chunk_size) groups state_dim -> chunks chunk_size groups state_dim",
+            chunk_size=chunk_size,
+        )
+        log_decay_cumsum = jnp.cumsum(log_decay, axis=-1)
 
-            decay = jnp.exp(-dt)[:, :, None]
-            mix = dt[:, :, None]
-            decay_group = rearrange(
-                decay,
-                "(groups heads) 1 1 -> groups heads 1 1",
-                groups=self.num_groups,
-                heads=heads_per_group,
-            )
-            mix_group = rearrange(
-                mix,
-                "(groups heads) 1 1 -> groups heads 1 1",
-                groups=self.num_groups,
-                heads=heads_per_group,
-            )
+        queries_keys_prod = einsum(
+            queries_chunked,
+            keys_chunked,
+            "chunks query_pos groups state_dim, chunks key_pos groups state_dim -> chunks query_pos key_pos groups",
+        )
+        y_diag = fused_ssd_intra_chunk(log_decay_cumsum, queries_keys_prod, values)
+
+        decay_states = jnp.exp(log_decay_cumsum[:, :, :, -1:] - log_decay_cumsum)
+        states = einsum(
+            keys_chunked,
+            decay_states,
+            values,
+            "chunks chunk_size groups state_dim, groups heads_per_group chunks chunk_size,"
+            " chunks chunk_size groups heads_per_group head_dim"
+            " -> chunks groups heads_per_group head_dim state_dim",
+        )
 
-            input_contribution_group = mix_group * x_norm_grouped[:, :, :, None] * input_proj_t[:, None, None, :]
-            carry_state_group = rearrange(
-                carry_state,
-                "(groups heads) head_channels state_channels -> groups heads head_channels state_channels",
-                groups=self.num_groups,
-                heads=heads_per_group,
-            )
-            updated_state_group = decay_group * carry_state_group + input_contribution_group
+        initial_state_grouped = rearrange(
+            initial_state,
+            "(groups heads_per_group) head_dim state_dim -> groups heads_per_group head_dim state_dim",
+            groups=self.num_groups,
+        )
+        states = jnp.concatenate([initial_state_grouped[None, ...], states], axis=0)
+        log_decay_chunk_ends = jnp.pad(log_decay_cumsum[:, :, :, -1], ((0, 0), (0, 0), (1, 0)))
+        decay_chunk = exp_segsum(log_decay_chunk_ends)
+        new_states = einsum(
+            decay_chunk,
+            states,
+            "groups heads_per_group out_idx chunks,"
+            " chunks groups heads_per_group head_dim state_dim"
+            " -> out_idx groups heads_per_group head_dim state_dim",
+        )
+        states = new_states[:-1]
+
+        state_decay_out = jnp.exp(log_decay_cumsum)
+        y_off = einsum(
+            queries_chunked,
+            states,
+            state_decay_out,
+            "chunks chunk_size groups state_dim,"
+            " chunks groups heads_per_group head_dim state_dim,"
+            " groups heads_per_group chunks chunk_size"
+            " -> chunks chunk_size groups heads_per_group head_dim",
+        )
+
+        y = y_diag + y_off
+        if d is not None:
+            d_grouped = rearrange(d, "(groups heads_per_group) -> groups heads_per_group", groups=self.num_groups)
+            y = y + d_grouped[None, None, :, :, None] * values
+        y = rearrange(
+            y,
+            "chunks chunk_size groups heads_per_group head_dim"
+            " -> (chunks chunk_size) (groups heads_per_group) head_dim",
+        )
+
+        if z is not None:
+            gate = z + z_bias[None, :, :] if z_bias is not None else z
+            y = y * jax.nn.silu(gate)
 
-            output_group = einsum(
-                updated_state_group,
-                output_proj_t,
-                "groups heads head_channels state_channels, groups state_channels -> groups heads head_channels",
+        y = y[:seq_len]
+
+        new_states_flat = rearrange(
+            new_states,
+            "chunks groups heads_per_group head_dim state_dim -> chunks (groups heads_per_group) head_dim state_dim",
+        )
+        final_state = self._compute_final_state(
+            values_orig,
+            keys_orig,
+            dt_orig,
+            new_states_flat,
+            num_steps,
+            chunk_size,
+        )
+
+        return y, final_state
+
+    def _compute_final_state(
+        self,
+        values: Float[Array, "suffix_tokens heads head_dim"],
+        keys: Float[Array, "suffix_tokens groups state_dim"],
+        dt: Float[Array, "suffix_tokens heads"],
+        chunk_states: Float[Array, "chunks_plus_1 heads head_dim state_dim"],
+        num_steps: Int[Array, ""],
+        chunk_size: int,
+    ) -> Float[Array, "heads head_dim state_dim"]:
+        """Compute the exact final state at position num_steps using precomputed chunk_states."""
+        heads_per_group = self.num_heads // self.num_groups
+
+        chunk_idx = num_steps // chunk_size
+        pos_in_chunk = num_steps % chunk_size
+        chunk_start_state = jax.lax.dynamic_index_in_dim(chunk_states, chunk_idx, axis=0, keepdims=False)
+
+        def at_boundary() -> Float[Array, "heads head_dim state_dim"]:
+            return chunk_start_state
+
+        def within_chunk() -> Float[Array, "heads head_dim state_dim"]:
+            chunk_start_pos = chunk_idx * chunk_size
+            values_chunk = jax.lax.dynamic_slice(
+                values,
+                (chunk_start_pos, 0, 0),
+                (chunk_size, values.shape[1], values.shape[2]),
             )
-            updated_state = rearrange(
-                updated_state_group,
-                "groups heads head_channels state_channels -> (groups heads) head_channels state_channels",
+            keys_chunk = jax.lax.dynamic_slice(
+                keys,
+                (chunk_start_pos, 0, 0),
+                (chunk_size, keys.shape[1], keys.shape[2]),
             )
-            output_t = rearrange(output_group, "groups heads head_channels -> (groups heads) head_channels")
+            dt_chunk = jax.lax.dynamic_slice(dt, (chunk_start_pos, 0), (chunk_size, dt.shape[1]))
 
-            propagated_state = jax.lax.cond(index < num_steps, lambda: updated_state, lambda: carry_state)
+            log_decay_cumsum = jnp.cumsum(-dt_chunk, axis=0)
+            last_pos_idx = pos_in_chunk - 1
+            log_decay_cumsum_at_last = jax.lax.dynamic_index_in_dim(
+                log_decay_cumsum,
+                last_pos_idx,
+                axis=0,
+                keepdims=False,
+            )
 
-            return (index + 1, propagated_state), output_t
+            decayed_start = jnp.exp(log_decay_cumsum_at_last)[:, None, None] * chunk_start_state
+
+            decay_to_last = jnp.exp(log_decay_cumsum_at_last[None, :] - log_decay_cumsum)
+            mask = jnp.arange(chunk_size) <= last_pos_idx
+            masked_decay = jnp.where(mask[:, None], decay_to_last, 0.0)
+            keys_expanded = jnp.repeat(keys_chunk, heads_per_group, axis=1)
+            input_contrib = einsum(
+                masked_decay,
+                keys_expanded,
+                values_chunk,
+                "chunk_size heads, chunk_size heads state_dim, chunk_size heads head_dim -> heads head_dim state_dim",
+            )
 
-        (_, final_state), outputs = jax.lax.scan(
-            scan_fn,
-            (jnp.zeros((), dtype=jnp.int32), initial_state),
-            (hidden_states, input_projection, output_projection, time_delta_log),
-        )
+            return decayed_start + input_contrib
 
-        return outputs, final_state
+        return jax.lax.cond(pos_in_chunk == 0, at_boundary, within_chunk)
 
     @eqx.filter_jit
     def __call__(
@@ -436,8 +666,6 @@ class Mamba2(TokenMixerBase[Mamba2Config, Mamba2StateLayer]):
         if positional_embeddings is not None:
             raise ValueError("Positional embeddings are not supported for Mamba2.")
 
-        conv_inputs, gate_values, time_delta_log = vmap(self.in_projection)(inputs)
-
         if state is None:
             state = Mamba2StateLayer.init(
                 self.config.kernel_size,
@@ -449,6 +677,13 @@ class Mamba2(TokenMixerBase[Mamba2Config, Mamba2StateLayer]):
                 self.activation_precision,
             )
 
+        seq_len, _ = inputs.shape
+
+        if seq_len == 1 and return_updated_state:
+            return self._decode_step(inputs, state)
+
+        conv_inputs, gate_values, time_delta_log = vmap(self.in_projection)(inputs)
+
         conv_output, updated_conv_state = self.conv(
             conv_inputs,
             length_without_padding,
@@ -466,50 +701,55 @@ class Mamba2(TokenMixerBase[Mamba2Config, Mamba2StateLayer]):
             axis=-1,
         )
 
-        hidden_states = rearrange(
+        values = rearrange(
             x_channels,
             "suffix_tokens (heads head_channels) -> suffix_tokens heads head_channels",
             heads=self.num_heads,
         )
-        input_projection = rearrange(
+        keys = rearrange(
             input_proj_channels,
             "suffix_tokens (groups state_channels) -> suffix_tokens groups state_channels",
             groups=self.num_groups,
         )
-        output_projection = rearrange(
+        queries = rearrange(
             output_proj_channels,
             "suffix_tokens (groups state_channels) -> suffix_tokens groups state_channels",
             groups=self.num_groups,
         )
-        time_delta_log = rearrange(
-            time_delta_log,
-            "suffix_tokens heads -> suffix_tokens heads",
-            heads=self.num_heads,
-        )
 
         if length_without_padding is None:
             length_without_padding, _ = inputs.shape
 
-        ssm_outputs, final_ssm_state = self._scan(
-            hidden_states,
-            input_projection,
-            output_projection,
-            time_delta_log,
+        gate_values_reshaped = rearrange(
+            gate_values,
+            "suffix_tokens (heads head_channels) -> suffix_tokens heads head_channels",
+            heads=self.num_heads,
+        )
+        gate_bias_reshaped = rearrange(
+            self.gate_bias,
+            "(heads head_channels) -> heads head_channels",
+            heads=self.num_heads,
+        )
+
+        dt = jax.nn.softplus(time_delta_log)
+        ssm_outputs, final_ssm_state = self._chunked_scan(
+            values,
+            keys,
+            queries,
+            dt,
             state.ssm_state,
+            self.config.chunk_size,
             length_without_padding,
+            d=self.skip_connection_weight,
+            z=gate_values_reshaped,
+            z_bias=gate_bias_reshaped,
         )
 
-        skip_contribution = self.skip_connection_weight[None, :, None] * hidden_states
-        ssm_outputs = ssm_outputs + skip_contribution
-
-        ssm_outputs = rearrange(
+        ssm_outputs_flat = rearrange(
             ssm_outputs,
             "suffix_tokens heads head_channels -> suffix_tokens (heads head_channels)",
         )
-
-        gated_outputs = ssm_outputs * jax.nn.silu(gate_values + self.gate_bias)
-
-        (outputs,) = vmap(self.out_projection)(gated_outputs)
+        (outputs,) = vmap(self.out_projection)(ssm_outputs_flat)
 
         if return_updated_state:
             assert updated_conv_state is not None

lalamo/speculator/__init__.py

@@ -1,5 +1,4 @@
 from .common import Speculator
-from .estimator import estimate_batchsize_from_memory
 from .inference import CollectTracesEvent, inference_collect_traces
 from .ngram import NGramSpeculator
 from .utils import SpeculatorTrainingEvent, train_speculator
@@ -9,7 +8,6 @@ __all__ = [
     "NGramSpeculator",
     "Speculator",
     "SpeculatorTrainingEvent",
-    "estimate_batchsize_from_memory",
     "inference_collect_traces",
     "train_speculator",
 ]