Slice and repeat

[AakashKumarNain]

The question I am asking here is a part of the workflow that is very common nowadays, especially for the KV cache used in LLMs and MLLMs, and I am unable to think of any good solution to it. I will try to keep the example simple.

Let's say you have an array x with shape (batch_size, seqlen, dim2, dim3). At every call, we want to select all elements in this array up to a given position i.e. x[:, :pos, :, :] and then repeat the selected elements along a certain axis.

Normal workflow (without jit)

key = jax.random.PRNGKey(1234)

bs = 2
seqlen = 5
dim2 = 3
dim3 = 4
pos = 2
num_repeats = 3

x = jax.random.randint(key=key, shape=(bs, seqlen, dim2, dim3), minval=2, maxval=100)
x_slice_v1 = x[:, :pos, :, :]
x_slice_v1 = jnp.repeat(x_slice_v1, num_repeats, axis=2)

# output
Array([[[[55, 49, 88, 61],
         [55, 49, 88, 61],
         [55, 49, 88, 61],
         [33, 10, 33, 10],
         [33, 10, 33, 10],
         [33, 10, 33, 10],
         [51, 68, 79, 23],
         [51, 68, 79, 23],
         [51, 68, 79, 23]],

        [[84, 45, 62, 59],
         [84, 45, 62, 59],
         [84, 45, 62, 59],
         [56,  6, 43, 85],
         [56,  6, 43, 85],
         [56,  6, 43, 85],
         [31, 96, 60, 92],
         [31, 96, 60, 92],
         [31, 96, 60, 92]]],


       [[[32, 87, 51, 83],
         [32, 87, 51, 83],
         [32, 87, 51, 83],
         [90, 10, 98, 85],
         [90, 10, 98, 85],
         [90, 10, 98, 85],
         [92, 55, 48, 97],
         [92, 55, 48, 97],
         [92, 55, 48, 97]],

        [[31, 16, 70, 26],
         [31, 16, 70, 26],
         [31, 16, 70, 26],
         [ 7, 16, 43, 50],
         [ 7, 16, 43, 50],
         [ 7, 16, 43, 50],
         [47, 35, 69, 48],
         [47, 35, 69, 48],
         [47, 35, 69, 48]]]], dtype=int32)

JIT-based workflow

For jit based workflow, there are two options:

1. Use jax.lax.dynamic_slice(...) to get the relevant elements and then apply repeat -> This is not ideal. Why? Too many recompilations leads to extreme poor performance as this operation is repeated in many many layers.
2. Always return an array of the same shape with the relevant positions filled with the relevant values, and other values are populated with some random number like zero. This is better than the first option, but then the values obtained after repeat will be different. Here is an example

@partial(jax.jit, static_argnums=(1,))
def get_slice(x, pos):
    seqlen = x.shape[1]
    idx = (jnp.arange(seqlen) < pos)[None, :, None, None]
    x_slice = jnp.where(idx, x, 0)
    return x_slice

x_slice_v2 = get_slice(x, pos)

The output of the above operation is this:

Array([[[[55, 49, 88, 61],
         [33, 10, 33, 10],
         [51, 68, 79, 23]],

        [[84, 45, 62, 59],
         [56,  6, 43, 85],
         [31, 96, 60, 92]],

        [[ 0,  0,  0,  0],
         [ 0,  0,  0,  0],
         [ 0,  0,  0,  0]],

        [[ 0,  0,  0,  0],
         [ 0,  0,  0,  0],
         [ 0,  0,  0,  0]],

        [[ 0,  0,  0,  0],
         [ 0,  0,  0,  0],
         [ 0,  0,  0,  0]]],


       [[[32, 87, 51, 83],
         [90, 10, 98, 85],
         [92, 55, 48, 97]],

        [[31, 16, 70, 26],
         [ 7, 16, 43, 50],
         [47, 35, 69, 48]],

        [[ 0,  0,  0,  0],
         [ 0,  0,  0,  0],
         [ 0,  0,  0,  0]],

        [[ 0,  0,  0,  0],
         [ 0,  0,  0,  0],
         [ 0,  0,  0,  0]],

        [[ 0,  0,  0,  0],
         [ 0,  0,  0,  0],
         [ 0,  0,  0,  0]]]], dtype=int32)

The slice extracted are correct but If I apply repeat on top of this like jnp.repeat(x_slice_v2, num_repeats, axis=2), it will result in a different array than the one we obtained with the normal workflow. I know that this is the expected output of this operation.

My question is "How can I obtain the value for x_slice_v2 where the output is equal to x_slice_v1?"

[AakashKumarNain]

Any suggestions? @jakevdp @mattjj

[jakevdp]

I'm confused as to why you can't use your x_slice_v1 code under JIT. It seems like it should work, so long as all variables affecting array size are marked as static.

[AakashKumarNain]

Hmm, on a second thought the cache grows in the same for both the function. I will close this. Thanks