electron-rare/mlx-lm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Make MambaCache compatible with batch generation for nemotron-h (#690)

Browse Source 
* Make MambaCache compatible with batch generation

* fix: Support right-padding masking in ArraysCache, add tests

* almost working

* test pass

* update models + gated delta

* rebase + fix

* fix

* allow batching in server

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Nikhil Mitra
2026-01-08 16:42:36 -05:00
committed by GitHub
parent 5cce1495e0
commit 43082feafa
16 changed files with 561 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx_lm/generate.py
+11 -7
View File
@@ -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
)
mlx_lm/models/cache.py
+36 -9
View File
@@ -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
mlx_lm/models/falcon_h1.py
+43 -25
View File
@@ -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)
mlx_lm/models/gated_delta.py
+2 -4
View File
@@ -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
mlx_lm/models/granitemoehybrid.py
+42 -25
View File
@@ -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)
mlx_lm/models/kimi_linear.py
+26 -11
View File
@@ -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
mlx_lm/models/lfm2.py
+20 -9
View File
@@ -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
mlx_lm/models/lfm2_moe.py
+20 -9
View File
@@ -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
mlx_lm/models/mamba2.py
+30 -12
View File
@@ -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)
mlx_lm/models/nemotron_h.py
+30 -14
View File
@@ -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)
mlx_lm/models/plamo2.py
+58 -35
View File
@@ -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)
mlx_lm/models/qwen3_next.py
+11 -2
View File
@@ -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))
mlx_lm/models/ssm.py
+22 -3
View File
@@ -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(
mlx_lm/server.py
+3 -1
View File
@@ -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
tests/test_generate.py
+121
View File
@@ -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()
tests/test_models.py
+86 -15
View File
@@ -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__":