Porting LLaMA 2 to JAX, issues with dynamic slices 😢

[Artur-Galstyan]

Hi everyone!

I'm porting LLaMA2 to Jax using Equinox, but I'm having some trouble with the dynamic slices and the fact that it keeps reJITting.

Consider this:

logits = jitted_model(
    tokens,
    start_pos, # <<< this is an integer
    mask=mask,
    inference_mode=True,
    freqs_cis=freqs_cis,
)

If I keep start_pos the same, everything is blazingly fast (takes 0.009 seconds per run). Nice.

But as soon as it changes, it takes close to 4 seconds, because it recompiles all the attention layers.

Here's what I mean:

(For brevity, here's only the relevant snippet)

    # =========== This is what makes it slow ===========
    cache_k = jax.lax.dynamic_update_slice(
        cache_k, xk, (0, bsz, start_pos, start_pos + seqlen)
    )
    cache_v = jax.lax.dynamic_update_slice(
        cache_v, xv, (0, bsz, start_pos, start_pos + seqlen)
    )
    start_indices = (0, 0, 0, 0)
    slice_sizes = (bsz, start_pos + seqlen, self.n_local_kv_heads, self.head_dim)
    keys = jax.lax.dynamic_slice(cache_k, start_indices, slice_sizes)
    values = jax.lax.dynamic_slice(cache_v, start_indices, slice_sizes)
    # print(f"{keys.shape=}") # 1, 2, 32, 128 | (bs, start_pos + seqlen, n_local_kv_heads, head_dim)
    # print(f"{values.shape=}") # 1, 2, 32, 128
    # ===============================================

The problem is the keys array (and values too, but since they're the same shapes, I'll only refer to keys). As start_pos becomes larger, so does the keys array, e.g.:

keys.shape=(1, 1, 32, 128)
keys.shape=(1, 2, 32, 128)
keys.shape=(1, 3, 32, 128)

etc.

I do however know, the maximum size it can have, but I'm not sure how I can use that information to make it all be fast. Does anyone have an idea?

[aniquetahir]

You might find this reference helpful:
https://github.com/ayaka14732/llama-2-jax