zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-26 02:33:21 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add theta cache for Rope and mask cache for ALiBi (#375)

Browse Source
This commit is contained in:
Anchen 2024-01-07 00:22:58 -08:00 committed by GitHub
parent c6d2878c1a
commit 48f6ca8c3a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 47 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
python/mlx/nn/layers/positional_encoding.py
View File

@ -25,8 +25,15 @@ class RoPE(Module):
base (float, optional): The base used to compute angular frequency for base (float, optional): The base used to compute angular frequency for
each dimension in the positional encodings. Default: ``10000``. each dimension in the positional encodings. Default: ``10000``.
scale (float, optional): The scale used to scale the positions. Default: ``1.0``. scale (float, optional): The scale used to scale the positions. Default: ``1.0``.
Attributes:
_cos_sin_theta_key (tuple): Cached key for the precomputed cosine and sine values.
_cos_sin_theta_value (tuple): Cached cosine and sine values.
""" """
_cos_sin_theta_key = None
_cos_sin_theta_value = None
def __init__( def __init__(
self, self,
dims: int, dims: int,
@ -86,8 +93,9 @@ class RoPE(Module):
return mx.reshape(rx, shape) return mx.reshape(rx, shape)
@staticmethod @classmethod
def create_cos_sin_theta( def create_cos_sin_theta(
cls,
N: int, N: int,
D: int, D: int,
offset: int = 0, offset: int = 0,
@ -95,11 +103,14 @@ class RoPE(Module):
scale: float = 1.0, scale: float = 1.0,
dtype=mx.float32, dtype=mx.float32,
): ):
if (N, D, offset, base, scale, dtype) != cls._cos_sin_theta_key:
D = D // 2 D = D // 2
positions = mx.arange(offset, N, dtype=dtype) * scale positions = mx.arange(offset, N, dtype=dtype) * scale
freqs = mx.exp(-mx.arange(0.0, D, dtype=dtype) * (math.log(base) / D)) freqs = mx.exp(-mx.arange(0.0, D, dtype=dtype) * (math.log(base) / D))
theta = mx.reshape(positions, (-1, 1)) * mx.reshape(freqs, (1, -1)) theta = mx.reshape(positions, (-1, 1)) * mx.reshape(freqs, (1, -1))
return mx.cos(theta), mx.sin(theta) cls._cos_sin_theta_key = (N, D, offset, base, scale, dtype)
cls._cos_sin_theta_value = (mx.cos(theta), mx.sin(theta))
return cls._cos_sin_theta_value
class SinusoidalPositionalEncoding(Module): class SinusoidalPositionalEncoding(Module):
@ -163,14 +174,25 @@ class SinusoidalPositionalEncoding(Module):
class ALiBi(Module): class ALiBi(Module):
@staticmethod _alibi_mask_key = None
_alibi_mask = None
@classmethod
def create_alibi_matrix( def create_alibi_matrix(
cls,
q_sequence_length: int, q_sequence_length: int,
k_sequence_length: int, k_sequence_length: int,
num_heads: int, num_heads: int,
offset: int, offset: int,
dtype=mx.float32, 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) x1 = mx.arange(offset, q_sequence_length)
x2 = mx.arange(0, k_sequence_length) x2 = mx.arange(0, k_sequence_length)
distance_matrix = -mx.abs( distance_matrix = -mx.abs(
@ -178,7 +200,16 @@ class ALiBi(Module):
) )
alibi_slope = ALiBi.create_alibi_slope(num_heads=num_heads) alibi_slope = ALiBi.create_alibi_slope(num_heads=num_heads)
alibi_mask = (distance_matrix * alibi_slope).astype(dtype) alibi_mask = (distance_matrix * alibi_slope).astype(dtype)
return alibi_mask cls._alibi_mask_key = (
q_sequence_length,
k_sequence_length,
num_heads,
offset,
dtype,
)
cls._alibi_mask = alibi_mask
return cls._alibi_mask
@staticmethod @staticmethod
def create_alibi_slope(num_heads): def create_alibi_slope(num_heads):