feat(mlx_lm)!: batch_generate

[llllvvuu]

This is based on @willccbb's implementation at https://github.com/willccbb/mlx_parallm.

BREAKING CHANGE: generate_step takes (bs, seq_len) instead of (seq_len,). In particular, sampler and logits_processors will need to handle logits of shape (bs, vocab_size) instead of (vocab_size,).

[llllvvuu]

Kind of interesting: for quantized models, the throughput is doesn't go up a lot between small bs (bs=1,2,3,4), but then it starts to go up a lot at higher bs, which is the opposite of what I expected intuitively. For unquantized models the throughput does goes up between small bs. I observe the same on @willccbb's original repo.

[awni]

I think it makes sense to minimize the complexity to the generate function (which is becoming a bit spaghetti) to split out the batched generation into a separate function called batch_generate. I would simplify that function to have fewer arguments (like no formatter, no printing during generation, verbose only prints the timings (e.g. as you have it now).

Also maybe more tricky is the fact that I think for this to be correct, the causal masks need to consider the left padding in the input (please correct me if I'm wrong about that). This has two implications:

1. Probably we'd need to add a mask parameter to the model __call__ functions and provide an appropriately constructed mask for the batch case.
2. The Rotating KV cache will be broken in this case (it keeps the initial tokens which would be the padded tokens) and when rotates the mask would need to be updated to consider the padding (which is a bit complicated/tedious). In this case I may suggest disabling this option entirely..

Let me know what you think about the above.

[llllvvuu]

I think it makes sense to minimize the complexity to the generate function (which is becoming a bit spaghetti) to split out the batched generation into a separate function called batch_generate. I would simplify that function to have fewer arguments (like no formatter, no printing during generation, verbose only prints the timings (e.g. as you have it now).

Makes sense to me, will implement.

Also maybe more tricky is the fact that I think for this to be correct, the causal masks need to consider the left padding in the input (please correct me if I'm wrong about that). This has two implications:

1. Probably we'd need to add a mask parameter to the model __call__ functions and provide an appropriately constructed mask for the batch case.

Yes, this sounds straightforward enough.

2. The Rotating KV cache will be broken in this case (it keeps the initial tokens which would be the padded tokens) and when rotates the mask would need to be updated to consider the padding (which is a bit complicated/tedious). In this case I may suggest disabling this option entirely..

I'll do a bit of thinking if there's an easy way to handle this, otherwise I'll remove that parameter in batch_generate.

Will update when these changes are ready!

[awni]

@llllvvuu are you coming back to this?

[llllvvuu]

hey @awni , sorry for the delay, I'd been job hunting this month. I should be able to get back to this in ~a week

[awni]

No worries, just checking. I'll follow up in a week or so.

[nath1295]

Just realised the attention mask has been mentioned in this PR, which is the reason I raised this issue #1044

[llllvvuu]

Sorry for the delay @awni . I took advantage of #1173 to update this PR. It is pending versioned release of ml-explore/mlx#1726 for the mask dtype.

I noticed one other potential issue: For absolute/rotary positional encodings, the position IDs of padded prompts won't start from 0 (this becomes more tricky if a padded prompt cache is added as then the position IDs should become non-contiguous IIUC). I'm not sure what the priority of this is or if it requires any change to mx.fast.rope.

[qinxuye]

Any update? we do need this to support parallel generation.

[awni]

Will get to this soon. Sorry for the delay.