diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 1fc2775..8fb6b79 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -913,6 +913,8 @@ def _merge_caches(caches): cache = BatchKVCache.merge([c[i] for c in caches]) elif isinstance(caches[0][i], RotatingKVCache): cache = BatchRotatingKVCache.merge([c[i] for c in caches]) + elif isinstance(caches[0][i], ArraysCache): + cache = ArraysCache.merge([c[i] for c in caches]) else: raise ValueError( f"{type(caches[0][i])} does not yet support batching with history" @@ -1027,9 +1029,11 @@ class BatchGenerator: def _process_prompts(self, prompts): uids, inputs, max_tokens, caches, samplers, logits_processors = zip(*prompts) + if hasattr(caches[0][0], "keys"): + cache_is_empty = all(c[0].keys is None for c in caches) + else: + cache_is_empty = all(c[0][0] is None for c in caches) - cache_lengths = [cache.cache_length(c) for c in caches] - max_cache_length = max(cache_lengths) lengths = [len(p) for p in inputs] max_length = max(lengths) padding = [max_length - l for l in lengths] @@ -1042,7 +1046,7 @@ class BatchGenerator: # New prompts so # 1. Left-pad the inputs # 2. Process - if max_cache_length == 0: + if cache_is_empty: inputs = _left_pad_prompts(inputs, max_length=max_length) prompt_cache = _make_cache(self.model, padding) @@ -1058,7 +1062,6 @@ class BatchGenerator: for uid, length in zip(uids, lengths) ] ) - mx.clear_cache() # Further prompt processing so we need to # 1. Merge the KV caches and prepare for right padded prompts @@ -1088,12 +1091,13 @@ class BatchGenerator: ) mx.clear_cache() - for c in prompt_cache: - c.finalize() mx.eval([c.state for c in prompt_cache]) - mx.clear_cache() inputs = last_inputs + for c in prompt_cache: + c.finalize() + mx.clear_cache() + y, logprobs = self._step( inputs, prompt_cache, samplers, logits_processors, tokens ) diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py index 841490d..93a38f4 100644 --- a/mlx_lm/models/cache.py +++ b/mlx_lm/models/cache.py @@ -551,6 +551,7 @@ class ArraysCache(_BaseCache): def __init__(self, size, left_padding: Optional[List[int]] = None): self.cache = [None] * size self.left_padding = mx.array(left_padding) if left_padding else None + self.lengths = None def __setitem__(self, idx, value): self.cache[idx] = value @@ -571,27 +572,53 @@ class ArraysCache(_BaseCache): In-place filter to keep just the given indices in the cache. """ self.cache = [c[batch_indices] for c in self.cache] - self.left_padding = None def extend(self, other): """ In-place extend this cache with the other cache. """ self.cache = [mx.concatenate([c, o]) for c, o in zip(self.cache, other.cache)] + + def extract(self, idx): + cache = ArraysCache(len(self.cache)) + cache.cache = [c[idx : idx + 1] for c in self.cache] + return cache + + def prepare(self, lengths=None, **kwargs): + self.lengths = mx.array(lengths) + + def finalize(self): + self.lengths = None self.left_padding = None + def advance(self, N): + if self.lengths is not None: + self.lengths -= N + if self.left_padding is not None: + self.left_padding -= N + def make_mask(self, N: int): - if self.cache[0] is None and self.left_padding is not None: - return mx.arange(N) >= self.left_padding[:, None] + if self.left_padding is not None: + pos = mx.arange(N) + return pos >= self.left_padding[:, None] + elif self.lengths is not None: + pos = mx.arange(N) + return pos < self.lengths[:, None] else: return None - def extract(self, idx): - """ - Extract a single item from the batched cache. - """ - cache = ArraysCache(len(self.cache)) - cache.cache = [c[idx : idx + 1] if c is not None else None for c in self.cache] + @classmethod + def merge(cls, caches): + n_state = len(caches[0].cache) + B = len(caches) + cache = cls(n_state) + for e in range(n_state): + c0 = caches[0][e] + shape = list(c0.shape) + shape[0] = B + cache[e] = mx.zeros(shape, c0.dtype) + for i in range(B): + cache[e][i : i + 1] = caches[i][e] return cache diff --git a/mlx_lm/models/falcon_h1.py b/mlx_lm/models/falcon_h1.py index 2683949..8e1b1b1 100644 --- a/mlx_lm/models/falcon_h1.py +++ b/mlx_lm/models/falcon_h1.py @@ -231,21 +231,36 @@ class FalconH1Mixer(nn.Module): self.intermediate_size, self.hidden_size, bias=args.projectors_bias ) - def _apply_conv( - self, conv_input: mx.array, cache: Optional[MambaCache] = None + def _conv( + self, + conv_input: mx.array, + cache: Optional[MambaCache], + mask: Optional[mx.array], ) -> mx.array: - if cache is None or cache[0] is None: - conv_state = mx.zeros( - (conv_input.shape[0], self.conv_kernel_size - 1, self.conv_dim), - dtype=conv_input.dtype, - ) - else: - conv_state = cache[0] - - padded_input = mx.concatenate([conv_state, conv_input], axis=1) + if mask is not None: + conv_input = mx.where(mask[..., None], conv_input, 0) if cache is not None: - cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :] + if cache[0] is None: + conv_state = mx.zeros( + (conv_input.shape[0], self.conv_kernel_size - 1, self.conv_dim), + dtype=conv_input.dtype, + ) + else: + conv_state = cache[0] + padded_input = mx.concatenate([conv_state, conv_input], axis=1) + n_keep = self.conv_kernel_size - 1 + if cache.lengths is not None: + t = padded_input.shape[1] + ends = mx.clip(cache.lengths, 0, t - n_keep) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(padded_input, positions, axis=1) + else: + cache[0] = padded_input[:, -n_keep:, :] + else: + padded_input = mx.pad( + conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)] + ) conv_output = self.conv1d(padded_input) return nn.silu(conv_output) @@ -256,17 +271,20 @@ class FalconH1Mixer(nn.Module): B: mx.array, C: mx.array, dt: mx.array, - state: Optional[mx.array] = None, - mask: Optional[mx.array] = None, + cache: Optional[MambaCache], + mask: Optional[mx.array], ) -> mx.array: batch_size, seq_len, _ = hidden_states.shape - hidden_states = hidden_states.reshape( batch_size, seq_len, self.num_heads, self.head_dim ) B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size) C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size) - + if cache: + state = cache[1] + lengths = cache.lengths + else: + state, lengths = None, None y, state = ssm_update( hidden_states, self.A_log, @@ -278,9 +296,11 @@ class FalconH1Mixer(nn.Module): state, self.time_step_limit, mask, + lengths, ) - - return y.reshape(batch_size, seq_len, self.intermediate_size), state + if cache: + cache[1] = state + return y.reshape(batch_size, seq_len, self.intermediate_size) def __call__(self, input_states, cache=None, mask: Optional[mx.array] = None): projected_states = self.in_proj(input_states) @@ -291,11 +311,9 @@ class FalconH1Mixer(nn.Module): axis=-1, ) - if mask is not None: - conv_input = mx.where(mask[..., None], conv_input, 0) - conv_output = self._apply_conv(conv_input, cache) + conv_output = self._conv(conv_input, cache, mask) - hidden_states_ssm, B, C = mx.split( + hidden_states, B, C = mx.split( conv_output, [ self.intermediate_size, @@ -303,10 +321,10 @@ class FalconH1Mixer(nn.Module): ], axis=-1, ) - state = cache[1] if cache else None - y, state = self._ssm(hidden_states_ssm, B, C, dt, state, mask) + + y = self._ssm(hidden_states, B, C, dt, cache, mask=mask) if cache: - cache[1] = state + cache.advance(y.shape[1]) if self.mamba_rms_norm: y = self.norm(y, gate) diff --git a/mlx_lm/models/gated_delta.py b/mlx_lm/models/gated_delta.py index 25aab21..3bb45ac 100644 --- a/mlx_lm/models/gated_delta.py +++ b/mlx_lm/models/gated_delta.py @@ -161,11 +161,9 @@ def _gated_delta_step_ops( state = state + k[..., None, :] * delta[..., None] # Output projection along key dim with q y = (state * q[..., None, :]).sum(axis=-1) # [B, H, Dv] + if mask is not None: - if mask.ndim == 2: - mask = mx.expand_dims(mask, axes=(2, 3)) - elif mask.ndim == 3: - mask = mx.expand_dims(mask, axis=-1) + mask = mx.expand_dims(mask, axis=(1, 2, 3)) state = mx.where(mask, state, old_state) return y, state diff --git a/mlx_lm/models/granitemoehybrid.py b/mlx_lm/models/granitemoehybrid.py index b239189..b2a8d2a 100644 --- a/mlx_lm/models/granitemoehybrid.py +++ b/mlx_lm/models/granitemoehybrid.py @@ -119,21 +119,36 @@ class GraniteMoeHybridMamba2Mixer(nn.Module): self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias ) - def _apply_conv( - self, conv_input: mx.array, cache: Optional[MambaCache] = None + def _conv( + self, + conv_input: mx.array, + cache: Optional[MambaCache], + mask: Optional[mx.array], ) -> mx.array: - if cache is None or cache[0] is None: - conv_state = mx.zeros( - (conv_input.shape[0], self.conv_kernel_size - 1, self.conv_dim), - dtype=conv_input.dtype, - ) - else: - conv_state = cache[0] - - padded_input = mx.concatenate([conv_state, conv_input], axis=1) + if mask is not None: + conv_input = mx.where(mask[..., None], conv_input, 0) if cache is not None: - cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :] + if cache[0] is None: + conv_state = mx.zeros( + (conv_input.shape[0], self.conv_kernel_size - 1, self.conv_dim), + dtype=conv_input.dtype, + ) + else: + conv_state = cache[0] + padded_input = mx.concatenate([conv_state, conv_input], axis=1) + n_keep = self.conv_kernel_size - 1 + if cache.lengths is not None: + t = padded_input.shape[1] + ends = mx.clip(cache.lengths, 0, t - n_keep) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(padded_input, positions, axis=1) + else: + cache[0] = padded_input[:, -n_keep:, :] + else: + padded_input = mx.pad( + conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)] + ) conv_output = self.conv1d(padded_input) return nn.silu(conv_output) @@ -144,8 +159,8 @@ class GraniteMoeHybridMamba2Mixer(nn.Module): B: mx.array, C: mx.array, dt: mx.array, - state: Optional[mx.array] = None, - mask: Optional[mx.array] = None, + cache: Optional[MambaCache], + mask: Optional[mx.array], ) -> mx.array: batch_size, seq_len, _ = hidden_states.shape @@ -154,26 +169,33 @@ class GraniteMoeHybridMamba2Mixer(nn.Module): ) B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size) C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size) + if cache: + state = cache[1] + lengths = cache.lengths + else: + state, lengths = None, None y, state = ssm_update( hidden_states, self.A_log, B, C, - self.D, + self.D.astype(hidden_states.dtype), dt, self.dt_bias, state, self.time_step_limit, mask, ) + if cache: + cache[1] = state - return y.reshape(batch_size, seq_len, self.intermediate_size), state + return y.reshape(batch_size, seq_len, self.intermediate_size) def __call__( self, hidden_states: mx.array, - mask: Optional[mx.array] = None, + mask: Optional[mx.array], cache: Optional[MambaCache] = None, ) -> mx.array: @@ -184,11 +206,7 @@ class GraniteMoeHybridMamba2Mixer(nn.Module): [self.intermediate_size, self.intermediate_size + self.conv_dim], axis=-1, ) - - if mask is not None: - conv_input = mx.where(mask[..., None], conv_input, 0) - conv_output = self._apply_conv(conv_input, cache) - + conv_output = self._conv(conv_input, cache, mask) hidden_states_ssm, B, C = mx.split( conv_output, [ @@ -197,10 +215,9 @@ class GraniteMoeHybridMamba2Mixer(nn.Module): ], axis=-1, ) - state = cache[1] if cache else None - y, state = self._ssm(hidden_states_ssm, B, C, dt, state, mask) + y = self._ssm(hidden_states_ssm, B, C, dt, cache, mask) if cache: - cache[1] = state + cache.advance(y.shape[1]) y = self.norm(y, gate) return self.out_proj(y) diff --git a/mlx_lm/models/kimi_linear.py b/mlx_lm/models/kimi_linear.py index 219906e..fa73d12 100644 --- a/mlx_lm/models/kimi_linear.py +++ b/mlx_lm/models/kimi_linear.py @@ -259,18 +259,30 @@ class ShortConv1d(nn.Module): ) def __call__( - self, x: mx.array, cache: Optional[mx.array] + self, + x: mx.array, + state: Optional[mx.array], + mask: Optional[mx.array], + lengths: Optional[mx.array], ) -> Tuple[mx.array, mx.array]: - if cache is None: - pad = mx.zeros( + if mask is not None: + x = mx.where(mask[..., None], x, 0) + + if state is None: + state = mx.zeros( (x.shape[0], self.kernel_size - 1, x.shape[-1]), dtype=x.dtype ) - else: - pad = cache - conv_input = mx.concatenate([pad, x], axis=1) + conv_input = mx.concatenate([state, x], axis=1) out = nn.silu(self.conv(conv_input)) - new_cache = conv_input[:, -self.kernel_size + 1 :, :] - return out, new_cache + n_keep = self.kernel_size - 1 + if lengths is not None: + ends = mx.clip(cache.lengths, 0, x.shape[1]) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + new_state = mx.take_along_axis(conv_input, positions, axis=1) + else: + new_state = conv_input[:, -n_keep:, :] + + return out, new_state class KimiDeltaAttention(nn.Module): @@ -323,9 +335,11 @@ class KimiDeltaAttention(nn.Module): if cache is not None: conv_state, ssm_state = cache + lengths = cache.lengths else: conv_state = None ssm_state = None + lengths = None if conv_state is None: s = mx.zeros((B, self.conv_kernel - 1, self.projection_dim), dtype=dtype) @@ -335,9 +349,9 @@ class KimiDeltaAttention(nn.Module): else: q_state, k_state, v_state = conv_state - q_conv, q_state = self.q_conv(self.q_proj(x), q_state) - k_conv, k_state = self.k_conv(self.k_proj(x), k_state) - v_conv, v_state = self.v_conv(self.v_proj(x), v_state) + q_conv, q_state = self.q_conv(self.q_proj(x), q_state, mask, lengths) + k_conv, k_state = self.k_conv(self.k_proj(x), k_state, mask, lengths) + v_conv, v_state = self.v_conv(self.v_proj(x), v_state, mask, lengths) if cache is not None: cache[0] = (q_state, k_state, v_state) @@ -374,6 +388,7 @@ class KimiDeltaAttention(nn.Module): if cache is not None: cache[1] = ssm_state + cache.advance(T) gate = self.g_b_proj(self.g_a_proj(x)).reshape( B, T, self.num_heads, self.head_dim diff --git a/mlx_lm/models/lfm2.py b/mlx_lm/models/lfm2.py index 62504ac..a46ad38 100644 --- a/mlx_lm/models/lfm2.py +++ b/mlx_lm/models/lfm2.py @@ -138,17 +138,28 @@ class ShortConv(nn.Module): Bx = B * x if mask is not None: Bx = mx.where(mask[..., None], Bx, 0) - state = None - if cache is not None: - state = cache[0] - if state is None: - state = mx.zeros( - (Bx.shape[0], self.L_cache - 1, self.args.hidden_size), dtype=Bx.dtype - ) - Bx = mx.concatenate([state, Bx], axis=-2) if cache is not None: - cache[0] = Bx[:, -(self.L_cache - 1) :] + if cache[0] is None: + state = mx.zeros( + (Bx.shape[0], self.L_cache - 1, self.args.hidden_size), + dtype=Bx.dtype, + ) + else: + state = cache[0] + Bx = mx.concatenate([state, Bx], axis=1) + n_keep = self.L_cache - 1 + t = x.shape[1] + if cache.lengths is not None: + ends = mx.clip(cache.lengths, 0, t) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(Bx, positions, axis=1) + else: + cache[0] = Bx[:, -n_keep:, :] + cache.advance(t) + else: + Bx = mx.pad(Bx, [(0, 0), (self.L_cache - 1, 0), (0, 0)]) + conv_out = self.conv(Bx) y = C * conv_out diff --git a/mlx_lm/models/lfm2_moe.py b/mlx_lm/models/lfm2_moe.py index 2869d9a..64b42eb 100644 --- a/mlx_lm/models/lfm2_moe.py +++ b/mlx_lm/models/lfm2_moe.py @@ -139,17 +139,28 @@ class ShortConv(nn.Module): Bx = B * x if mask is not None: Bx = mx.where(mask[..., None], Bx, 0) - state = None - if cache is not None: - state = cache[0] - if state is None: - state = mx.zeros( - (Bx.shape[0], self.L_cache - 1, self.args.hidden_size), dtype=Bx.dtype - ) - Bx = mx.concatenate([state, Bx], axis=-2) if cache is not None: - cache[0] = Bx[:, -(self.L_cache - 1) :] + if cache[0] is None: + state = mx.zeros( + (Bx.shape[0], self.L_cache - 1, self.args.hidden_size), + dtype=Bx.dtype, + ) + else: + state = cache[0] + Bx = mx.concatenate([state, Bx], axis=1) + n_keep = self.L_cache - 1 + t = x.shape[1] + if cache.lengths is not None: + ends = mx.clip(cache.lengths, 0, t) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(Bx, positions, axis=1) + else: + cache[0] = Bx[:, -n_keep:, :] + cache.advance(t) + else: + Bx = mx.pad(Bx, [(0, 0), (self.L_cache - 1, 0), (0, 0)]) + conv_out = self.conv(Bx) y = C * conv_out diff --git a/mlx_lm/models/mamba2.py b/mlx_lm/models/mamba2.py index 82a09b3..a026f6e 100644 --- a/mlx_lm/models/mamba2.py +++ b/mlx_lm/models/mamba2.py @@ -93,9 +93,15 @@ class Mamba2Block(nn.Module): self.intermediate_size, self.hidden_size, bias=args.use_bias ) - def _apply_conv( - self, conv_input: mx.array, cache: Optional[MambaCache] = None + def _conv( + self, + conv_input: mx.array, + cache: Optional[MambaCache], + mask: Optional[mx.array], ) -> mx.array: + if mask is not None: + conv_input = mx.where(mask[..., None], conv_input, 0) + if cache is not None: if cache[0] is None: conv_state = mx.zeros( @@ -105,7 +111,14 @@ class Mamba2Block(nn.Module): else: conv_state = cache[0] padded_input = mx.concatenate([conv_state, conv_input], axis=1) - cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :, :] + n_keep = self.conv_kernel_size - 1 + if cache.lengths is not None: + t = padded_input.shape[1] + ends = mx.clip(cache.lengths, 0, t - n_keep) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(padded_input, positions, axis=1) + else: + cache[0] = padded_input[:, -n_keep:, :] else: padded_input = mx.pad( conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)] @@ -120,8 +133,8 @@ class Mamba2Block(nn.Module): B: mx.array, C: mx.array, dt: mx.array, - state: Optional[mx.array] = None, - mask: Optional[mx.array] = None, + cache: Optional[MambaCache], + mask: Optional[mx.array], ) -> mx.array: batch_size, seq_len, _ = hidden_states.shape hidden_states = hidden_states.reshape( @@ -129,6 +142,11 @@ class Mamba2Block(nn.Module): ) B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size) C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size) + if cache: + state = cache[1] + lengths = cache.lengths + else: + state, lengths = None, None y, state = ssm_update( hidden_states, self.A_log, @@ -140,8 +158,11 @@ class Mamba2Block(nn.Module): state, self.time_step_limit, mask, + lengths, ) - return y.reshape(batch_size, seq_len, self.intermediate_size), state + if cache: + cache[1] = state + return y.reshape(batch_size, seq_len, self.intermediate_size) def __call__( self, @@ -155,9 +176,7 @@ class Mamba2Block(nn.Module): [self.intermediate_size, self.intermediate_size + self.conv_dim], axis=-1, ) - if mask is not None: - conv_input = mx.where(mask[..., None], conv_input, 0) - conv_output = self._apply_conv(conv_input, cache) + conv_output = self._conv(conv_input, cache, mask) hidden_states, B, C = mx.split( conv_output, [ @@ -166,10 +185,9 @@ class Mamba2Block(nn.Module): ], axis=-1, ) - state = cache[1] if cache else None - y, state = self._ssm(hidden_states, B, C, dt, state, mask=mask) + y = self._ssm(hidden_states, B, C, dt, cache, mask=mask) if cache: - cache[1] = state + cache.advance(y.shape[1]) y = self.norm(y, gate) return self.out_proj(y) diff --git a/mlx_lm/models/nemotron_h.py b/mlx_lm/models/nemotron_h.py index f821ae3..5bc2419 100644 --- a/mlx_lm/models/nemotron_h.py +++ b/mlx_lm/models/nemotron_h.py @@ -111,9 +111,15 @@ class NemotronHMamba2Mixer(nn.Module): self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias ) - def _apply_conv( - self, conv_input: mx.array, cache: Optional[MambaCache] = None + def _conv( + self, + conv_input: mx.array, + cache: Optional[MambaCache], + mask: Optional[mx.array], ) -> mx.array: + if mask is not None: + conv_input = mx.where(mask[..., None], conv_input, 0) + if cache is not None: if cache[0] is None: conv_state = mx.zeros( @@ -123,11 +129,19 @@ class NemotronHMamba2Mixer(nn.Module): else: conv_state = cache[0] padded_input = mx.concatenate([conv_state, conv_input], axis=1) - cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :, :] + n_keep = self.conv_kernel_size - 1 + if cache.lengths is not None: + t = padded_input.shape[1] + ends = mx.clip(cache.lengths, 0, t - n_keep) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(padded_input, positions, axis=1) + else: + cache[0] = padded_input[:, -n_keep:, :] else: padded_input = mx.pad( conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)] ) + conv_output = self.conv1d(padded_input) return nn.silu(conv_output) @@ -137,8 +151,8 @@ class NemotronHMamba2Mixer(nn.Module): B: mx.array, C: mx.array, dt: mx.array, - state: Optional[mx.array], - mask: Optional[mx.array] = None, + cache: Optional[MambaCache], + mask: Optional[mx.array], ) -> mx.array: batch_size, seq_len, _ = hidden_states.shape @@ -147,6 +161,11 @@ class NemotronHMamba2Mixer(nn.Module): ) B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size) C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size) + if cache: + state = cache[1] + lengths = cache.lengths + else: + state, lengths = None, None y, state = ssm_update( hidden_states, @@ -160,8 +179,10 @@ class NemotronHMamba2Mixer(nn.Module): self.time_step_limit, mask, ) + if cache: + cache[1] = state - return y.reshape(batch_size, seq_len, self.intermediate_size), state + return y.reshape(batch_size, seq_len, self.intermediate_size) def __call__( self, @@ -177,11 +198,7 @@ class NemotronHMamba2Mixer(nn.Module): [self.intermediate_size, self.intermediate_size + self.conv_dim], axis=-1, ) - if mask is not None: - conv_input = mx.where(mask[..., None], conv_input, 0) - - conv_output = self._apply_conv(conv_input, cache) - + conv_output = self._conv(conv_input, cache, mask) hidden_states_ssm, B, C = mx.split( conv_output, [ @@ -190,10 +207,9 @@ class NemotronHMamba2Mixer(nn.Module): ], axis=-1, ) - state = cache[1] if cache else None - y, state = self._ssm(hidden_states_ssm, B, C, dt, state, mask) + y = self._ssm(hidden_states_ssm, B, C, dt, cache, mask) if cache: - cache[1] = state + cache.advance(y.shape[1]) y = self.norm(y, gate) return self.out_proj(y) diff --git a/mlx_lm/models/plamo2.py b/mlx_lm/models/plamo2.py index cf96331..ec16136 100644 --- a/mlx_lm/models/plamo2.py +++ b/mlx_lm/models/plamo2.py @@ -54,27 +54,13 @@ class RMSNorm(nn.Module): ) -def causal_conv1d_update(conv_state, x, weight) -> tuple[mx.array, mx.array]: - dim = x.shape[-1] - state_len = conv_state.shape[-2] - x = mx.concatenate([conv_state, x], axis=-2) - conv_state = x[:, -state_len:] - out = mx.conv1d( - x, - weight, - padding=0, - groups=dim, - ) - return nn.silu(out), conv_state - - class Mamba(nn.Module): def __init__(self, config: ModelArgs) -> None: super().__init__() self.config = config self.hidden_size = config.hidden_size self.d_state = config.mamba_d_state - self.d_conv = config.mamba_d_conv + self.conv_kernel_size = config.mamba_d_conv self.chunk_size = config.mamba_chunk_size self.num_heads = config.mamba_num_heads self.hidden_size_per_head = config.hidden_size_per_head @@ -88,7 +74,7 @@ class Mamba(nn.Module): in_channels=self.intermediate_size, out_channels=self.intermediate_size, bias=False, - kernel_size=self.d_conv, + kernel_size=self.conv_kernel_size, groups=self.intermediate_size, padding=0, ) @@ -111,20 +97,63 @@ class Mamba(nn.Module): self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + def _conv( + self, + conv_input: mx.array, + cache: Optional[MambaCache], + mask: Optional[mx.array], + ) -> mx.array: + if mask is not None: + conv_input = mx.where(mask[..., None], conv_input, 0) + + if cache is not None: + if cache[0] is None: + conv_state = mx.zeros( + ( + conv_input.shape[0], + self.conv_kernel_size - 1, + self.intermediate_size, + ), + dtype=conv_input.dtype, + ) + else: + conv_state = cache[0] + padded_input = mx.concatenate([conv_state, conv_input], axis=1) + n_keep = self.conv_kernel_size - 1 + if cache.lengths is not None: + t = padded_input.shape[1] + ends = mx.clip(cache.lengths, 0, t - n_keep) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(padded_input, positions, axis=1) + else: + cache[0] = padded_input[:, -n_keep:, :] + else: + padded_input = mx.pad( + conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)] + ) + + conv_output = self.conv1d(padded_input) + return nn.silu(conv_output) + def _ssm( self, x: mx.array, B: mx.array, C: mx.array, dt: mx.array, - state: Optional[mx.array] = None, - mask: Optional[mx.array] = None, + cache: Optional[Any], + mask: Optional[mx.array], ) -> mx.array: batch_size, seq_len, _ = x.shape x = x.reshape(batch_size, seq_len, self.num_heads, self.hidden_size_per_head) B = B.reshape(batch_size, seq_len, 1, self.d_state) C = C.reshape(batch_size, seq_len, 1, self.d_state) + if cache: + state = cache[1] + lengths = cache.lengths + else: + state, lengths = None, None y, state = ssm_update( x, @@ -136,8 +165,11 @@ class Mamba(nn.Module): self.dt_bias, state, mask=mask, + lengths=lengths, ) - return y.reshape(batch_size, seq_len, self.intermediate_size), state + if cache: + cache[1] = state + return y.reshape(batch_size, seq_len, self.intermediate_size) def __call__( self, @@ -147,14 +179,6 @@ class Mamba(nn.Module): ): bsize, length, _ = hidden_states.shape - if cache is not None and cache[0] is not None: - conv_state = cache[0] - else: - conv_state = mx.zeros( - (bsize, self.d_conv - 1, self.intermediate_size), - dtype=hidden_states.dtype, - ) - zx = self.in_proj(hidden_states) zx = zx.reshape(bsize, length, self.num_heads, -1) # z: (bsize, length, num_heads, hidden_size_per_head) @@ -168,9 +192,8 @@ class Mamba(nn.Module): ) x = x.reshape(bsize, -1, self.num_heads * self.hidden_size_per_head) - if mask is not None: - x = mx.where(mask[..., None], x, 0) - x, conv_state = causal_conv1d_update(conv_state, x, self.conv1d.weight) + x = self._conv(x, cache, mask) + BCdt = self.bcdt_proj(x) B, C, dt = mx.split(BCdt, [self.d_state, self.d_state * 2], axis=-1) @@ -181,18 +204,18 @@ class Mamba(nn.Module): # (bsize, length, num_heads) dt = self.dt_proj(dt) - out, ssm_state = self._ssm( + out = self._ssm( x, B, C, dt, - cache[1] if cache else None, + cache, mask, ) + if cache: + cache.advance(out.shape[1]) + out = out * nn.silu(z.flatten(-2)) - if cache is not None: - cache[0] = conv_state - cache[1] = ssm_state return self.out_proj(out) diff --git a/mlx_lm/models/qwen3_next.py b/mlx_lm/models/qwen3_next.py index b5194ad..d3ca775 100644 --- a/mlx_lm/models/qwen3_next.py +++ b/mlx_lm/models/qwen3_next.py @@ -44,10 +44,10 @@ class ModelArgs(BaseModelArgs): rope_theta: float partial_rotary_factor: float max_position_embeddings: int + head_dim: int norm_topk_prob: bool = False tie_word_embeddings: bool = False attention_bias: bool = False - head_dim: Optional[int] = None rope_scaling: Optional[Dict[str, Union[float, str]]] = None full_attention_interval: int = 4 @@ -247,8 +247,16 @@ class Qwen3NextGatedDeltaNet(nn.Module): if mask is not None: mixed_qkv = mx.where(mask[..., None], mixed_qkv, 0) conv_input = mx.concatenate([conv_state, mixed_qkv], axis=1) + if cache is not None: - cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :] + n_keep = self.conv_kernel_size - 1 + if cache.lengths is not None: + ends = mx.clip(cache.lengths, 0, S) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(conv_input, positions, axis=1) + else: + cache[0] = conv_input[:, -n_keep:, :] + conv_out = nn.silu(self.conv1d(conv_input)) q, k, v = [ @@ -280,6 +288,7 @@ class Qwen3NextGatedDeltaNet(nn.Module): if cache is not None: cache[1] = state + cache.advance(S) out = self.norm(out, z) return self.out_proj(out.reshape(B, S, -1)) diff --git a/mlx_lm/models/ssm.py b/mlx_lm/models/ssm.py index ff28876..aeb16bf 100644 --- a/mlx_lm/models/ssm.py +++ b/mlx_lm/models/ssm.py @@ -114,6 +114,7 @@ def ssm_attn( state: Optional[mx.array] = None, time_step_limit: Tuple[float, float] = (0.001, 100.0), mask: Optional[mx.array] = None, + lengths: Optional[mx.array] = None, step: int = 256, ) -> Tuple[mx.array, mx.array]: """SSD-SSM forward pass. @@ -128,6 +129,7 @@ def ssm_attn( dt_bias: Bias for time deltas of shape (num_heads,). time_step_limit: Minimum and maximum value for time deltas. mask: Optional multiplicative mask. + lengths: Optional lenghts of sequences, assumed to be the full length if unspecified. step: Step size for processing x. Code modified from @@ -157,7 +159,14 @@ def ssm_attn( y = surrogate_attention_matrix @ dtx.swapaxes(1, 2) y = mx.swapaxes(y, 1, 2) - decay = decay[:, :, -1:, :].transpose(0, 3, 1, 2) + if lengths is not None: + pos = mx.maximum(mx.minimum(lengths, step) - 1, 0) + pos = mx.expand_dims(pos, (1, 2, 3)) + decay = mx.take_along_axis(decay, pos, axis=2) + else: + decay = decay[:, :, -1:, :] + + decay = decay.transpose(0, 3, 1, 2) B = mx.repeat(B, h // g, axis=1).swapaxes(2, 3) dtxdecay = dtx * decay dtxdecay = dtxdecay.swapaxes(1, 2).swapaxes(2, 3) @@ -167,10 +176,16 @@ def ssm_attn( if state is not None: exp_dtA_cumsum = mx.exp(mx.cumsum(dtA, axis=-2)) next_state += exp_dtA_cumsum[:, -1, :, None, None] * state - state = state.reshape((b, 1, g, repeats, dh, d)) C = C.reshape(b, s, g, 1, d, 1) - y_prev = (state @ C).squeeze(-1).flatten(2, 3) + y_prev = ( + (state.reshape((b, 1, g, repeats, dh, d)) @ C).squeeze(-1).flatten(2, 3) + ) y += exp_dtA_cumsum[..., None] * y_prev + if lengths is not None and state is not None: + next_state = mx.where( + mx.expand_dims(lengths < 0, (1, 2, 3)), state, next_state + ) + return y, next_state ys = [] @@ -183,6 +198,8 @@ def ssm_attn( state, None if mask is None else mask[..., i : i + step], ) + if lengths is not None: + lengths = lengths - step ys.append(y) y = mx.concatenate(ys, axis=1) + x * D.reshape(1, 1, h, 1) return y, state @@ -199,6 +216,7 @@ def ssm_update( state: Optional[mx.array] = None, time_step_limit: Tuple[float, float] = (0.001, 100.0), mask: Optional[mx.array] = None, + lengths: Optional[mx.array] = None, ): seq_len = hidden_states.shape[1] if ( @@ -218,6 +236,7 @@ def ssm_update( state, time_step_limit, mask=mask, + lengths=lengths, ) else: return ssm_update_kernel( diff --git a/mlx_lm/server.py b/mlx_lm/server.py index 425cd38..061fefc 100644 --- a/mlx_lm/server.py +++ b/mlx_lm/server.py @@ -35,7 +35,9 @@ from huggingface_hub import scan_cache_dir from ._version import __version__ from .generate import BatchGenerator, stream_generate from .models.cache import ( + ArraysCache, KVCache, + MambaCache, RotatingKVCache, can_trim_prompt_cache, make_prompt_cache, @@ -541,7 +543,7 @@ class ResponseGenerator: ): return False for c in self.model_provider.cache_types: - if c not in (KVCache, RotatingKVCache): + if c not in (KVCache, RotatingKVCache, ArraysCache, MambaCache): return False if args.seed is not None: return False diff --git a/tests/test_generate.py b/tests/test_generate.py index 6bd33ee..fee5801 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -1,5 +1,6 @@ # Copyright © 2024 Apple Inc. +import random import unittest from typing import List @@ -10,6 +11,7 @@ from mlx_lm.generate import ( GenerationResponse, batch_generate, generate, + generate_step, stream_generate, ) from mlx_lm.models.cache import RotatingKVCache @@ -511,6 +513,125 @@ class TestGenerate(unittest.TestCase): if rotating: del self.model.make_cache + def _continued_generation_test_helper(self, model): + def rand_prompt(n): + return [random.randint(0, 1000) for _ in range(n)] + + # Make the prompts + prompts_a = [ + rand_prompt(5), + rand_prompt(3), + rand_prompt(8), + rand_prompt(1), + ] + prompts_b = [ + rand_prompt(2), + rand_prompt(7), + rand_prompt(4), + rand_prompt(6), + ] + + # Generate once + batch_gen = BatchGenerator( + model, + stop_tokens={}, + max_tokens=10, + prefill_batch_size=4, + prefill_step_size=32, + completion_batch_size=2, + ) + + uids = batch_gen.insert(prompts_a) + caches = {uid: None for uid in uids} + while responses := batch_gen.next(): + for r in responses: + if r.finish_reason is not None: + caches[r.uid] = r.prompt_cache + + caches = [caches[uid] for uid in uids] + + # Generate the 2nd time + uids = batch_gen.insert(prompts_b, caches=caches) + batch_responses = {uid: [] for uid in uids} + while responses := batch_gen.next(): + for r in responses: + batch_responses[r.uid].append(r.logprobs) + + for e, uid in enumerate(uids): + for i, (_, logprobs) in enumerate( + generate_step( + mx.array(prompts_b[e]), + model, + max_tokens=10, + prompt_cache=caches[e], + ) + ): + batch_logprobs = batch_responses[uid][i] + self.assertTrue( + mx.allclose(batch_logprobs, logprobs, rtol=1e-4, atol=1e-4) + ) + + def test_batch_continued_generation_ssm(self): + from mlx_lm.models import mamba2 + + random.seed(0) + mx.random.seed(4) + + # Make a small SSM model + args = mamba2.ModelArgs( + model_type="mamba2", + num_heads=8, + head_dim=16, + vocab_size=1000, + hidden_size=128, + intermediate_size=128, + state_size=32, + num_hidden_layers=4, + layer_norm_epsilon=1e-4, + conv_kernel=3, + n_groups=4, + use_bias=False, + use_conv_bias=False, + tie_word_embeddings=True, + time_step_limit=(0.01, 10), + time_step_rank="auto", + ) + model = mamba2.Model(args) + self._continued_generation_test_helper(model) + + def test_batch_continued_generation_gated_delta(self): + from mlx_lm.models import qwen3_next + + random.seed(0) + mx.random.seed(4) + args = qwen3_next.ModelArgs( + model_type="qwen3_next", + hidden_size=128, + num_hidden_layers=4, + intermediate_size=128, + num_attention_heads=8, + num_key_value_heads=4, + vocab_size=1000, + linear_num_value_heads=4, + linear_num_key_heads=4, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_conv_kernel_dim=3, + num_experts=4, + num_experts_per_tok=2, + decoder_sparse_step=1, + shared_expert_intermediate_size=128, + mlp_only_layers=[0], + moe_intermediate_size=128, + rms_norm_eps=1e-5, + head_dim=64, + rope_theta=1000.0, + partial_rotary_factor=0.5, + max_position_embeddings=1000, + ) + model = qwen3_next.Model(args) + self._continued_generation_test_helper(model) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_models.py b/tests/test_models.py index a329c0e..5bdcf16 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1915,7 +1915,6 @@ class TestModels(unittest.TestCase): "n_groups": 4, "use_bias": False, "use_conv_bias": False, - "chunk_size": 32, "tie_word_embeddings": True, "time_step_limit": (0.01, 10), "time_step_rank": "auto", @@ -1998,6 +1997,31 @@ class TestModels(unittest.TestCase): "group_norm_size": 1, "max_position_embeddings": 1000, }, + { + "model_type": "qwen3_next", + "hidden_size": 128, + "num_hidden_layers": 4, + "intermediate_size": 128, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "vocab_size": 1000, + "linear_num_value_heads": 4, + "linear_num_key_heads": 4, + "linear_key_head_dim": 32, + "linear_value_head_dim": 32, + "linear_conv_kernel_dim": 3, + "num_experts": 4, + "num_experts_per_tok": 2, + "decoder_sparse_step": 1, + "shared_expert_intermediate_size": 128, + "mlp_only_layers": [0], + "moe_intermediate_size": 128, + "rms_norm_eps": 1e-5, + "head_dim": 64, + "rope_theta": 1000.0, + "partial_rotary_factor": 0.5, + "max_position_embeddings": 1000, + }, { "model_type": "kimi_linear", "vocab_size": 1000, @@ -2235,6 +2259,50 @@ class TestModels(unittest.TestCase): self.assertTrue(mx.allclose(out, out_m, atol=1e-4, rtol=1e-4)) self.assertTrue(mx.allclose(out_state, out_state_m, atol=1e-4, rtol=1e-4)) + def test_ssm_right_pad(self): + batch_size = 1 + n_group = 1 + num_heads = 48 + head_dim = 64 + state_dim = 128 + seq_len = 4 + pad = 2 + + hidden_states = mx.random.normal( + shape=(batch_size, seq_len + pad, num_heads, head_dim) + ) + B = mx.random.normal(shape=(batch_size, seq_len + pad, n_group, state_dim)) + C = mx.random.normal(shape=(batch_size, seq_len + pad, n_group, state_dim)) + dt = mx.random.normal(shape=(batch_size, seq_len + pad, num_heads)) + dt_bias = mx.random.normal(shape=(num_heads,)) + A_log = mx.random.normal(shape=(num_heads,)) + D = mx.random.normal(shape=(num_heads,)) + out, out_state = ssm_attn( + hidden_states[:, :-pad], + A_log, + B[:, :-pad], + C[:, :-pad], + D, + dt[:, :-pad], + dt_bias, + ) + mask = mx.array([[True] * seq_len + [False] * pad]) + lengths = mx.array([seq_len]) + out_m, out_state_m = ssm_attn( + hidden_states, + A_log, + B, + C, + D, + dt, + dt_bias, + mask=mask, + lengths=lengths, + ) + out_m = out_m[:, :-pad] + self.assertTrue(mx.allclose(out, out_m, atol=1e-4, rtol=1e-4)) + self.assertTrue(mx.allclose(out_state, out_state_m, atol=1e-4, rtol=1e-4)) + def test_gated_delta(self): mx.random.seed(0) for B in [1, 2]: @@ -2269,23 +2337,26 @@ class TestModels(unittest.TestCase): k = mx.random.normal(shape=(B, T, Hk, Dk)) v = mx.random.normal(shape=(B, T, Hv, Dv)) g = mx.random.normal(shape=(B, T, Hv)) - mask = mx.array([[False, True, True]]) beta = mx.random.normal(shape=(B, T, Hv)) state = mx.random.normal(shape=(B, Hv, Dk, Dv)) - y_gt, st_gt = gated_delta_ops( - q[:, 1:], - k[:, 1:], - v[:, 1:], - g[:, 1:], - beta[:, 1:], - state, - ) - for fn in [gated_delta_ops, gated_delta_kernel]: - y, st = fn(q, k, v, g, beta, state, mask) - y = y[:, 1:] - self.assertTrue(mx.allclose(y, y_gt, rtol=1e-4, atol=1e-4)) - self.assertTrue(mx.allclose(st, st_gt, rtol=1e-4, atol=1e-3)) + for s, e, mask in [ + (1, 3, mx.array([[False, True, True]])), + (0, 2, mx.array([[True, True, False]])), + ]: + y_gt, st_gt = gated_delta_ops( + q[:, s:e], + k[:, s:e], + v[:, s:e], + g[:, s:e], + beta[:, s:e], + state, + ) + for fn in [gated_delta_ops, gated_delta_kernel]: + y, st = fn(q, k, v, g, beta, state, mask) + y = y[:, s:e] + self.assertTrue(mx.allclose(y, y_gt, rtol=1e-4, atol=1e-4)) + self.assertTrue(mx.allclose(st, st_gt, rtol=1e-4, atol=1e-3)) if __name__ == "__main__":