46da74fea2
* Unify attention mask creation in LLMs.
Currently, each model implementation in `mlx-examples/llms/models` has ad-hoc
code to create a mask for the attention mechanism. This usually takes the form:
```
mask = None
if h.shape[1] > 1:
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
mask = mask.astype(h.dtype)
```
This correctly creates a mask only if the input consists of more than one token.
But this code assumes the multi-token input is at the beginning of inference.
If, for example, we are evaluating multiple tokens because of speculative
decoding or prompt cache reuse, this mask will not have the correct shape and
and will cause the raising of an exception in the attention computation.
Some of the models correctly implement the mask creation with code like this:
```
mask = None
if h.shape[1] > 1:
mask = create_additive_causal_mask(
h.shape[1], cache[0].offset if cache is not None else 0
)
mask = mask.astype(h.dtype)
```
This commit unifies the attention mask creation for all models with a new
function `create_attention_mask`, reducing code duplication and helping all
models support inference performance enhancements like those mentioned above.
* Allow batches in LLM key-value cache
The current implementation of the LLM key-value cache assumes that
the input batch is of size 1. Input batching (evaluating multiple
alterative inputs at the same time) can be a valuable tool for
speculative sampling and other techniques.
This change removes the hard-coded batch size from the code that
resizes the key-value cache.
* Simplify causal mask creation
Use the same codepath regardless of whether there's an offset or
not. Addresses [this comment](https://github.com/ml-explore/mlx-examples/pull/911#discussion_r1691459717).
* Use old-style type annotation to avoid linter error
80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
import inspect
|
|
from dataclasses import dataclass
|
|
from typing import List, Optional
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
|
|
class KVCache:
|
|
|
|
def __init__(self, head_dim, n_kv_heads):
|
|
self.n_kv_heads = n_kv_heads
|
|
if isinstance(head_dim, int):
|
|
self.k_head_dim = self.v_head_dim = head_dim
|
|
elif isinstance(head_dim, tuple) and len(head_dim) == 2:
|
|
self.k_head_dim, self.v_head_dim = head_dim
|
|
else:
|
|
raise ValueError("head_dim must be an int or a tuple of two ints")
|
|
self.keys = None
|
|
self.values = None
|
|
self.offset = 0
|
|
self.step = 256
|
|
|
|
def update_and_fetch(self, keys, values):
|
|
prev = self.offset
|
|
if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
|
|
B = keys.shape[0]
|
|
n_steps = (self.step + keys.shape[2] - 1) // self.step
|
|
k_shape = (B, self.n_kv_heads, n_steps * self.step, self.k_head_dim)
|
|
v_shape = (B, self.n_kv_heads, n_steps * self.step, self.v_head_dim)
|
|
new_k = mx.zeros(k_shape, keys.dtype)
|
|
new_v = mx.zeros(v_shape, values.dtype)
|
|
if self.keys is not None:
|
|
if prev % self.step != 0:
|
|
self.keys = self.keys[..., :prev, :]
|
|
self.values = self.values[..., :prev, :]
|
|
self.keys = mx.concatenate([self.keys, new_k], axis=2)
|
|
self.values = mx.concatenate([self.values, new_v], axis=2)
|
|
else:
|
|
self.keys, self.values = new_k, new_v
|
|
|
|
self.offset += keys.shape[2]
|
|
self.keys[..., prev : self.offset, :] = keys
|
|
self.values[..., prev : self.offset, :] = values
|
|
return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
|
|
|
|
|
|
@dataclass
|
|
class BaseModelArgs:
|
|
@classmethod
|
|
def from_dict(cls, params):
|
|
return cls(
|
|
**{
|
|
k: v
|
|
for k, v in params.items()
|
|
if k in inspect.signature(cls).parameters
|
|
}
|
|
)
|
|
|
|
|
|
def create_additive_causal_mask(N: int, offset: int = 0):
|
|
rinds = mx.arange(offset + N)
|
|
linds = mx.arange(offset, offset + N) if offset else rinds
|
|
mask = linds[:, None] < rinds[None]
|
|
return mask * -1e9
|
|
|
|
|
|
def create_attention_mask(h: mx.array, cache: Optional[List[KVCache]] = None):
|
|
T = h.shape[1]
|
|
if T > 1:
|
|
# Input consists of multiple tokens, create a causal mask so that prior
|
|
# tokens do not give attention to later tokens. If a cache is in place
|
|
# (because e.g. prompt reuse), offset the mask accordingly.
|
|
offset = cache[0].offset if cache is not None and cache[0] is not None else 0
|
|
mask = create_additive_causal_mask(T, offset)
|
|
mask = mask.astype(h.dtype)
|
|
else:
|
|
mask = None
|
|
return mask
|