generalize attention utils

Author: rdyro

llama3/llama3_jax/attention_cache_utils.py

@@ -18,6 +18,13 @@
 _pad_after = lambda x, l, axis: jnp.pad(x, [(0, 0) if i != axis else (0, l - x.shape[i]) for i in range(x.ndim)])
 
 
+def safe_zip(*args):
+  if len(args) == 0:
+    return []
+  assert all(len(arg) == len(args[0]) for arg in args)
+  return zip(*args)
+
+
 def _transpose_attention_tree(kv_list: list[PyTree], time_axis: int):
     "From a list of cache entries stacked along layer idx (in transit) to stacked along batch, layers split into list."
 
@@ -28,7 +35,7 @@ def _transpose_attention_tree(kv_list: list[PyTree], time_axis: int):
     for i, c in enumerate(kv_list[0]):
         els = [[_split(z) for z in jax.tree.leaves(kv[i])] for kv in kv_list]  # [B, R_flat, L]
         els = jax.tree.map(lambda *xs: jnp.concatenate(xs, axis=0), *els)  # [R_flat, L]
-        leaves_list = list(zip(*els))  # [L, R_flat]
+        leaves_list = list(safe_zip(*els))  # [L, R_flat]
         out[i] = [jax.tree.unflatten(jax.tree.structure(c), leaves) for leaves in leaves_list]  # [L, R]
     return tuple(out), max_seq_len
 
@@ -41,7 +48,7 @@ def _transpose_attention_tree(kv_list: list[PyTree], time_axis: int):
 @partial(jax.jit, donate_argnames=("cache",))
 def _kvcache_update_cache(
     cache: KVCache,
-    kvs: list[tuple[list[jax.Array | QuantArray], list[jax.Array | QuantArray]]],
+    kvs: list[tuple[jax.Array | QuantArray, ...]],
     batch_idxs: list[jax.Array],
     actual_lens: list[jax.Array],
     update_mask: list[bool] | None = None,
@@ -62,15 +69,17 @@ def _update_element(x, u):
         # update_permute = [batch_dim, time_dim] + update_permute
         return x.at[batch_idxs[:, None], :, time_indices, ...].set(u.transpose(update_permute), mode="drop")
 
-    cache_k, cache_v = jax.tree.map(_update_element, (cache.k, cache.v), kvs)
+    cache_kvs = jax.tree.map(_update_element, cache.buffers, kvs)
     cache_starts = cache.starts.at[batch_idxs].set(start_time, mode="drop")
     cache_iter = jnp.where(uninitialized_cache, jnp.max(actual_lens), cache.iter)
-    return dataclasses.replace(cache, k=cache_k, v=cache_v, iter=cache_iter, starts=cache_starts)
+
+    buffer_names = [field.name for field in dataclasses.fields(cache)][:len(cache_kvs)]
+    return dataclasses.replace(cache, **dict(safe_zip(buffer_names, cache_kvs)), iter=cache_iter, starts=cache_starts)
 
 
 def kvcache_update_cache(
     cache: KVCache,
-    kvs: list[tuple[list[jax.Array | QuantArray], list[jax.Array | QuantArray]]],
+    kvs: list[tuple[jax.Array | QuantArray, ...]],
     batch_idxs: list[jax.Array],
     actual_lens: list[jax.Array],
 ):
@@ -85,7 +94,7 @@ def kvcache_update_cache(
 def kvcache_get_entry(cache: KVCache, batch_idx: jax.Array):
     shift = -cache.starts[batch_idx]
     assert cache.time_axis > 0
-    kvs = jax.tree.map(lambda x: jnp.roll(x[batch_idx, ...], shift=shift, axis=cache.time_axis - 1), (cache.k, cache.v))
+    kvs = jax.tree.map(lambda x: jnp.roll(x[batch_idx, ...], shift=shift, axis=cache.time_axis - 1), cache.buffers)
     kvs = (jax.tree.map(lambda *xs: jnp.stack(xs, 0), kvs[0]), jax.tree.map(lambda *xs: jnp.stack(xs, 0), kvs[1]))
     true_len = cache.fill_len()[batch_idx]
     return kvs, true_len
@@ -109,13 +118,13 @@ def _find_empty_pages(free_pages: jax.Array, k: int, proposal_pages: jax.Array |
         return jax.lax.top_k(free_pages, k)[1]
 
 
-def _paged_update_slice(cache: PagedKVCache, k: jax.Array | QuantArray, v: jax.Array | QuantArray, *, layer_idx: int):
-    key_heads = cache.k[layer_idx].shape[0]
-    assert v.shape[:-1] == k.shape[:-1] == (cache.batch_size, key_heads, 1)
+def _paged_update_slice(cache: PagedKVCache, kv: tuple[jax.Array | QuantArray, ...], *, layer_idx: int):
+    #key_heads = cache.buffers[0][layer_idx].shape[0]
+    #assert v.shape[:-1] == k.shape[:-1] == (cache.batch_size, key_heads, 1)  # TODO write this generically
     needs_next_page = (cache.lengths % cache.page_size) == 0
     page_table_idx = cache.lengths // cache.page_size
     current_page_cursor = jnp.take_along_axis(cache.block_tables, page_table_idx[:, None], axis=-1)[..., 0]
-    avg_pages_per_batch_entry = round(cache.k[layer_idx].shape[0] / cache.batch_size)
+    avg_pages_per_batch_entry = round(cache.buffers[0][layer_idx].shape[0] / cache.batch_size)
     even_batch_spread = jnp.arange(cache.batch_size) * avg_pages_per_batch_entry
     proposal_pages = jnp.where(cache.lengths == 0, even_batch_spread, current_page_cursor + 1)
     free_pages = _find_empty_pages(cache.free_pages, cache.batch_size, proposal_pages=proposal_pages)
@@ -127,27 +136,28 @@ def _paged_update_slice(cache: PagedKVCache, k: jax.Array | QuantArray, v: jax.A
     # for batch index update the target slice is (heads, i, j, head_dim)
     # so transpose update (batch, heads, seq, head_dim) -> (batch, heads, head_dim) -> (heads, batch, head_dim)
     _update = lambda dest, src: dest.at[:, page_cursor, inpage_cursor, ...].set(src.squeeze(2).swapaxes(0, 1))
-    cache.k[layer_idx], cache.v[layer_idx] = jax.tree.map(_update, (cache.k[layer_idx], cache.v[layer_idx]), (k, v))
+    for buffer, new_buffer in safe_zip(cache.buffers, kv):
+        buffer[layer_idx] = jax.tree.map(_update, buffer[layer_idx], new_buffer)
 
     batch_idx = jnp.arange(cache.batch_size)
     new_block_tables = cache.block_tables.at[batch_idx, new_lengths // cache.page_size].set(page_cursor)
 
     new_free_pages = cache.free_pages.at[page_cursor].set(False, mode="drop")
     new_state = dict(lengths=new_lengths, block_tables=new_block_tables, free_pages=new_free_pages)
-    return cache.k[layer_idx], cache.v[layer_idx], new_state
+    return tuple(buffer[layer_idx] for buffer in cache.buffers), new_state
 
 
-def paged_update_slice(cache: PagedKVCache, k: jax.Array | QuantArray, v: jax.Array | QuantArray, *, layer_idx: int):
+def paged_update_slice(cache: PagedKVCache, kv: tuple[jax.Array | QuantArray, ...], *, layer_idx: int):
     repl_sharding = jax.typeof(cache.lengths).sharding
-    kv_sharding = jax.tree.map(lambda x: jax.typeof(x).sharding, (cache.k[layer_idx], cache.v[layer_idx]))
-    sharding = (*kv_sharding, dict(lengths=repl_sharding, block_tables=repl_sharding, free_pages=repl_sharding))
-    return auto_axes(partial(_paged_update_slice, layer_idx=layer_idx), out_sharding=sharding)(cache, k, v)
+    kv_sharding = jax.tree.map(lambda x: jax.typeof(x).sharding, tuple(buffer[layer_idx] for buffer in cache.buffers))
+    sharding = (kv_sharding, dict(lengths=repl_sharding, block_tables=repl_sharding, free_pages=repl_sharding))
+    return auto_axes(partial(_paged_update_slice, layer_idx=layer_idx), out_sharding=sharding)(cache, kv)
 
 
 @partial(jax.jit, donate_argnames=("cache",))
 def _batch_paged_update_sequences(
     cache: PagedKVCache,
-    kvs: list[tuple[list[jax.Array | QuantArray], list[jax.Array | QuantArray]]],
+    kvs: list[tuple[jax.Array | QuantArray, ...]],
     batch_idxs: list[jax.Array],
     actual_lens: list[jax.Array],
     update_mask: list[bool] | None = None,
@@ -156,9 +166,7 @@ def _batch_paged_update_sequences(
     batch_idxs = jnp.where(update_mask, jnp.array(batch_idxs), 2**30)  # send masked to nowhere
     actual_lens = jnp.minimum(jnp.array(actual_lens), jnp.array([jax.tree.leaves(kv)[0].shape[2] for kv in kvs]))
 
-    kvs, max_seq_len = _transpose_attention_tree(
-        kvs, time_axis=2
-    )  # undo stacking along the layer dimension for transit
+    kvs, max_seq_len = _transpose_attention_tree(kvs, time_axis=2)  # undo stack along layer dimension in transit
 
     # clear existing pages
     actual_page_num = jnp.rint(jnp.ceil(cache.lengths[batch_idxs] / cache.page_size)).astype(jnp.int32)
@@ -186,21 +194,23 @@ def _update_element(x, u):
         update_permute = [1, 0, 2] + [i for i in range(u.ndim) if i not in (0, 1, 2)]
         return x.at[:, pages_idx, ...].set(u.transpose(update_permute), mode="drop")
 
-    cache_k, cache_v = jax.tree.map(_update_element, (cache.k, cache.v), kvs)
+    new_buffers = jax.tree.map(_update_element, cache.buffers, kvs)
     block_tables_idx = jnp.where(
         update_mask[:, None] & (pages_arange[None, :] < actual_page_num[:, None]), pages_arange[None, :], 2**30
     )
     new_block_tables = cache.block_tables.at[batch_idxs[:, None], block_tables_idx].set(pages_idx, mode="drop")
     new_free_pages = new_free_pages.at[pages_idx.reshape(-1)].set(False, mode="drop")
     new_lengths = cache.lengths.at[batch_idxs].set(actual_lens, mode="drop")
+
+    named_buffers = dict(zip([field.name for field in dataclasses.fields(cache)][:len(new_buffers)], new_buffers))
     return dataclasses.replace(
-        cache, k=cache_k, v=cache_v, lengths=new_lengths, block_tables=new_block_tables, free_pages=new_free_pages
+        cache, **named_buffers, lengths=new_lengths, block_tables=new_block_tables, free_pages=new_free_pages
     )
 
 
 def batch_paged_update_sequences(
     cache: KVCache,
-    kvs: list[tuple[list[jax.Array | QuantArray], list[jax.Array | QuantArray]]],
+    kvs: list[tuple[jax.Array | QuantArray, ...]],
     batch_idxs: list[jax.Array],
     actual_lens: list[jax.Array],
 ):
@@ -222,5 +232,5 @@ def batch_paged_get_entry(cache: PagedKVCache, batch_idx: jax.Array, max_seq_len
     _get = lambda x: jnp.where(mask[None, :, *([None] * (x.ndim - 3))], _reshape_out(x[:, page_indices, ...]), 0)
 
     # stack along layer dimensions for transit
-    kvs = tuple(jax.tree.map(lambda *xs: jnp.stack(xs, 0), *z) for z in jax.tree.map(_get, (cache.k, cache.v)))
+    kvs = tuple(jax.tree.map(lambda *xs: jnp.stack(xs, 0), *z) for z in jax.tree.map(_get, cache.buffers))
     return kvs, true_len

llama3/llama3_jax/model.py

@@ -37,6 +37,7 @@
 except ModuleNotFoundError:
     from jax.sharding import auto_axes as _auto_axes, reshard
 from jax.experimental.pallas.ops.gpu import paged_attention
+from etils import epath
 
 from . import ragged_attention
 from . import attention_cache_utils
@@ -216,7 +217,7 @@ class ArrayInfo:
 _count_left_padding = lambda ids, pad_id=0: auto_axes(
     lambda ids: jnp.sum(jnp.cumsum(ids != pad_id, axis=-1) == 0, axis=-1), out_sharding=P(None)
 )(ids)
-_length_minus_padding = lambda segment_ids: auto_axes(
+_length_minus_right_padding = lambda segment_ids: auto_axes(
     lambda segment_ids: jnp.sum(jnp.cumsum(jnp.flip(segment_ids != 0, -1), axis=-1) > 0, -1), out_sharding=P(None)
 )(segment_ids)
 
@@ -411,7 +412,7 @@ class KVCache(_Init):
     iter: jax.Array  # []  # sequences are right-aligned for slice update performance
     starts: jax.Array  # [batch_size]  # sequences are right-aligned, we need start indices
     batch_size: int = 0
-    size: int = 0
+    size: int = 2 ** 30
     time_axis: int = 2
 
     @classmethod
@@ -428,6 +429,7 @@ def abstract(cls, cfg: Config, batch_size: int):
             # -1 means unintialized since iter (cursor) must be 0 <= iter < len - 1
             iter=ArrayInfo((), jnp.int32, (), jax.nn.initializers.constant(-1)),
             starts=ArrayInfo((batch_size,), jnp.int32, ("batch",), jax.nn.initializers.zeros),
+            size=cfg.max_seq_len,
         )
         if cfg.quant_cache:
             _quantize = partial(quantize, axis=-1, scale_dtype=cfg.quant_scale_dtype, zero_init=True)
@@ -447,8 +449,11 @@ def abstract(cls, cfg: Config, batch_size: int):
         return cache
 
     def fill_len(self) -> jax.Array:
-        length = jnp.where(self.iter > self.starts, self.iter - self.starts, self.size + self.iter - self.starts)
-        return jnp.where(self.iter >= 0, length, 0)
+        return jnp.where(self.iter >= 0, (self.iter - self.starts) % self.size, 0)
+
+    @property
+    def buffers(self) -> tuple[jax.Array, ...]:
+        return (self.k, self.v)
 
     update_slice = None
     insert_sequences = staticmethod(attention_cache_utils.kvcache_update_cache)
@@ -463,7 +468,7 @@ class PagedKVCache(_Init):
     block_tables: jax.Array  # [batch_size, pages_per_seq]
     free_pages: jax.Array  # [total_num_pages]
     batch_size: int = 0
-    size: int = 2**31 - 1
+    size: int = 2**30
     page_size: int = 0
 
     @classmethod
@@ -501,6 +506,10 @@ def abstract(cls, cfg: "Config", batch_size: int, total_num_pages: int, page_siz
     def fill_len(self) -> jax.Array:
         return self.lengths
 
+    @property
+    def buffers(self) -> tuple[jax.Array, ...]:
+        return (self.k, self.v)
+
     update_slice = staticmethod(attention_cache_utils.paged_update_slice)
     insert_sequences = staticmethod(attention_cache_utils.batch_paged_update_sequences)
     get_sequence = staticmethod(attention_cache_utils.batch_paged_get_entry)
@@ -807,12 +816,9 @@ def attention_block(
         q, k = apply_rotary_embedding(q, sin, cos), apply_rotary_embedding(k, sin, cos)
 
     if cfg.quant_cache:
-        k = QuantArray(
-            *quantize(k, -1, scale_dtype=cfg.quant_scale_dtype), out_scaling=True, scale_expand_dims=(-2, -3)
-        )
-        v = QuantArray(
-            *quantize(v, -1, scale_dtype=cfg.quant_scale_dtype), out_scaling=False, scale_expand_dims=(-2, -3)
-        )
+        _quantize = partial(quantize, axis=-1, scale_dtype=cfg.quant_scale_dtype)
+        k = QuantArray(*_quantize(k), out_scaling=True, scale_expand_dims=(-2, -3))
+        v = QuantArray(*_quantize(v), out_scaling=False, scale_expand_dims=(-2, -3))
 
     with jax.named_scope("cache_update"):
         paged_state, starts = None, None
@@ -825,23 +831,21 @@ def attention_block(
             ) % cache.size  # [B, T]
 
             q_segment_ids = jnp.where(segment_ids != 0, 1, 0)
-            incremental_position = jnp.max(_length_minus_padding(segment_ids))
+            incremental_position = jnp.max(_length_minus_right_padding(segment_ids))
             # i.e. valid below where we've written things [B, T]
-            kv_segment_ids = (
-                (time_indices >= 0) & (time_indices < cache.fill_len()[:, None] + incremental_position)
-            ).astype(jnp.int32)
-            q_offset = cache.fill_len() - _count_left_padding(segment_ids)
+            kv_segment_ids = (time_indices >= 0) & (time_indices < cache.fill_len()[:, None] + incremental_position)
+            q_offset = cache.fill_len() - _count_left_padding(segment_ids, 0)  # 0 is the pad "token" for segment_ids
             starts, lengths = cache.starts, cache.fill_len()
             cache_updates = (k, v)
         elif is_type(cache, PagedKVCache):
             cache: PagedKVCache
-            k, v, paged_state = PagedKVCache.update_slice(cache, k=k, v=v, layer_idx=idx)
+            (k, v), paged_state = PagedKVCache.update_slice(cache, (k, v), layer_idx=idx)
             cache_updates = (k, v, paged_state)
         else:
             # this supports prefill only; no support for a ring cache buffer here
             q_segment_ids, kv_segment_ids = segment_ids, segment_ids
             q_offset = jnp.zeros(x.shape[0], dtype=jnp.int32)
-            starts, lengths = _count_left_padding(segment_ids, 0), _length_minus_padding(kv_segment_ids)
+            starts, lengths = _count_left_padding(segment_ids, 0), _length_minus_right_padding(kv_segment_ids)
             cache_updates = (k, v)
 
     # Compute attention
@@ -931,15 +935,12 @@ def forward(
         x, cache_updates = forward_layer(x, segment_ids, layer, sin, cos, idx, cfg, cache)
         all_cache_updates.append(cache_updates)
 
-    # Final layer norm.
-    x = rms_norm(x, weights.gamma_final)
-
-    # Project to vocabulary size
-    logits = einsum("btd,dv->btv", x, weights.lm_head)
+    x = rms_norm(x, weights.gamma_final)  # Final layer norm.
+    logits = einsum("btd,dv->btv", x, weights.lm_head)  # Project to vocabulary size
 
     if is_type(cache, KVCache):
         cache.k, cache.v = [z[0] for z in all_cache_updates], [z[1] for z in all_cache_updates]
-        new_iter = (jnp.maximum(0, cache.iter) + jnp.max(_length_minus_padding(segment_ids))) % cache.size
+        new_iter = (jnp.maximum(0, cache.iter) + jnp.max(_length_minus_right_padding(segment_ids))) % cache.size
         cache = dataclasses.replace(cache, iter=new_iter)
         return logits, cache
     elif is_type(cache, PagedKVCache):

llama3/pyproject.toml

@@ -19,6 +19,8 @@ dependencies = [
     #"datasets",
     "gcsfs",
     "etils",
+    "importlib_resources",
+    "absl-py",
 ]
 
 # we don't need CUDA torch

llama3/tests/test_model.py

@@ -54,15 +54,15 @@ def test_model_init(self, quant):
     @parameterized.product(quant=[False, True])
     def test_cache_init(self, quant):
         cfg = dataclasses.replace(self.small_cfg, quant_cache=quant)
-        cache = l3jax.KVCache.init(random.key(0), cfg, 2, cfg.max_seq_len)
+        cache = l3jax.KVCache.init(random.key(0), cfg, 2)
         del cache
 
     @parameterized.product(quant_weights=[False, True], quant_cache=[True, False])
     def test_prefill_decode(self, quant_weights, quant_cache):
         cfg = dataclasses.replace(self.small_cfg, quant_layer=quant_weights, quant_cache=quant_cache)
         tokens = jnp.ones((1, 32), dtype=jnp.int32)
         weights = l3jax.Weights.init(random.key(0), cfg)
-        cache = l3jax.KVCache.init(random.key(0), cfg, tokens.shape[0], cfg.max_seq_len)
+        cache = l3jax.KVCache.init(random.key(0), cfg, tokens.shape[0])
         with use_mesh(cfg.mesh):
             max_tokens, _, cache = l3jax.prefill(tokens, weights, cache, cfg)
         next_tokens = max_tokens[:, :-1]