Length masking for batch inputs

[barronalex]

Allow passing an array of lengths to make_prompt_cache.

create_attention_mask then does the correct masking when you call an mlx-lm model with a batch of padded variable-length inputs and the cache you created above.

Very open to other ways of doing this, but I went with this because it allows us to leave the model code untouched.

[awni]

Very open to other ways of doing this, but I went with this because it allows us to leave the model code untouched.

I do wonder if we should revisit our design of model interface to include an optional mask. For cases that require flexible masks it will be a lot more .. well .. flexible. The other added benefit is it is more functional and easier to compile / export if we decide to go that route.

Other than the tedious aspect of changing the model interface.. how would that work for your use case?

[barronalex]

That would work great too and would definitely be easier to follow in the code.

The only gotcha will be making sure that the mask you pass has the same dtype as the rest of the model so you don’t accidentally upcast.

[awni]

The only gotcha will be making sure that the mask you pass has the same dtype as the rest of the model so you don’t accidentally upcast.

It's a good point and is easy to do. If we do go that route (and maybe either way) we should require the mask type to be the same type as keys/queries/values in the scaled_dot_product fast implementation. Or we could require that result_type(mask.dtype(), queries.dtype()) == queries.dtype() or something like that.

[chimezie]

+1 on "revisit our design of model interface to include an optional mask."

[barronalex]

I'll redraft this to add a mask input everywhere. I'll make a PR in the core repo too for the scaled_dot_product_attention change.

[barronalex]

OK all done.

I'm a little worried we're going to break anyone using model(x, cache) but model(x, mask, cache) is more consistent with all of the other places we use the mask.

[awni]

Nice, thanks a ton for making that change!!

anyone using model(x, cache) but model(x, mask, cache) is more consistent with all of the other places we use the mask.

It's an easy fix for them to make it a kwarg. Let's land it and see what happens.

[awni]

Looks great, thanks!!

[chimezie]

So, just a question to help my understanding of the implications of this: you no longer need to 'remove' any prefix that was cached beforehand (as said here)?

[awni]

So, just a question to help my understanding of the implications of this: you no longer need to 'remove' any prefix that was cached beforehand

Right now this doesn't change any behavior in mlx lm for training or fine-tuning. It's just some additional functionality that let's you specify a mask and some lengths to the mask creation function.

For cases like:

• prefill prefx
• generate with question A
• trim cache to prefix length
• generate with question B

You could use a mask instead but I would in general not advise it since the cache will grow and you will be doing a whole bunch of wasted computation.

[llllvvuu]

This masks out the most recent tokens (right-padding), right?

print(create_causal_mask(3, 1, lengths=mx.array([1, 2, 3, 1])))