zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-21 20:46:50 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fused Operations in delta, B, C = ... :. Before: 57.822 tokens-per-sec, after: 83.890 tokens-per-sec

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2025-01-20 18:39:22 +01:00
parent e43ac7c90e
commit 9494a275ac
1 changed files with 4 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
llms/mlx_lm/models/mamba.py
View File

@ -127,14 +127,10 @@ class MambaBlock(nn.Module):
A = -mx.exp(self.A_log)
D = self.D
deltaBC = self.x_proj(x)
delta, B, C = mx.split(
deltaBC,
indices_or_sections=[
self.time_step_rank,
self.time_step_rank + self.ssm_state_size,
],
axis=-1,
)
delta, B, C = map(self.mixer_norm if self.use_bcdt_rms else lambda x: x,
mx.split(deltaBC, [self.time_step_rank,
self.time_step_rank + self.ssm_state_size],
axis=-1))
if self.use_bcdt_rms:
delta, B, C = map(self.mixer_norm, (delta, B, C))
delta = nn.softplus(self.dt_proj(delta))