zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 01:12:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add state property to PlamoCache

Browse Source
This commit is contained in:
Shunta Saito 2025-02-23 16:06:38 +09:00
parent 21c0abaf23
commit e2d9d619c4
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
llms/mlx_lm/models/plamo2.py
View File

@ -194,6 +194,10 @@ class PlamoAttentionCache(nn.Module):
assert value.shape[2] == L assert value.shape[2] == L
self.key = key self.key = key
self.value = value self.value = value
@property
def state(self) -> tuple[mx.array, mx.array]:
return self.key, self.value
class PlamoMambaCache(nn.Module): class PlamoMambaCache(nn.Module):
@ -207,6 +211,9 @@ class PlamoMambaCache(nn.Module):
self.conv_state = conv_state self.conv_state = conv_state
self.ssm_state = ssm_state self.ssm_state = ssm_state
@property
def state(self) -> tuple[mx.array, mx.array]:
return self.conv_state, self.ssm_state
PlamoLayerCache = PlamoAttentionCache | PlamoMambaCache PlamoLayerCache = PlamoAttentionCache | PlamoMambaCache