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

Fix ALiBi slopes for non-power-of-2 num_heads (#3071)

Browse Source
This commit is contained in:
atharva
2026-01-29 07:23:11 -08:00
committed by GitHub
parent 0c6a895ed7
commit 590b4f1c16
1 changed files with 23 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mlx/nn/layers/positional_encoding.py
+23 -32
View File
@@ -115,47 +115,38 @@ class SinusoidalPositionalEncoding(Module):
class ALiBi(Module):
_alibi_mask_key = None
_alibi_mask = None
@classmethod
@staticmethod
def create_alibi_matrix(
cls,
q_sequence_length: int,
k_sequence_length: int,
num_heads: int,
offset: int,
dtype=mx.float32,
):
if (
q_sequence_length,
k_sequence_length,
num_heads,
offset,
dtype,
) != cls._alibi_mask_key:
x1 = mx.arange(offset, q_sequence_length)
x2 = mx.arange(0, k_sequence_length)
distance_matrix = -mx.abs(
mx.expand_dims(x1[:, None] - x2[None, :], axis=(0, 1))
)
alibi_slope = ALiBi.create_alibi_slope(num_heads=num_heads)
alibi_mask = (distance_matrix * alibi_slope).astype(dtype)
cls._alibi_mask_key = (
q_sequence_length,
k_sequence_length,
num_heads,
offset,
dtype,
)
cls._alibi_mask = alibi_mask
return cls._alibi_mask
x1 = mx.arange(offset, q_sequence_length)
x2 = mx.arange(0, k_sequence_length)
distance_matrix = -mx.abs(
mx.expand_dims(x1[:, None] - x2[None, :], axis=(0, 1))
)
alibi_slope = ALiBi.create_alibi_slope(num_heads=num_heads, dtype=dtype)
alibi_mask = (distance_matrix * alibi_slope).astype(dtype)
return alibi_mask
@staticmethod
def create_alibi_slope(num_heads):
x = (2**8) ** (1 / num_heads)
out = mx.power(x, -mx.arange(1, num_heads + 1))
def create_alibi_slope(num_heads, dtype):
def get_slopes(n: int):
if math.log2(n).is_integer():
start = 2 ** (-(2 ** -(math.log2(n) - 3)))
return [start * start**i for i in range(n)]
else:
closest_power_of_2 = 2 ** math.floor(math.log2(n))
return (
get_slopes(closest_power_of_2)
+ get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
)
slopes = get_slopes(num_heads)
out = mx.array(slopes, dtype=dtype)
return mx.expand_dims(out, axis=(-1, -2))
def __call__(self, attention_scores, offset=0, mask=None):