mlx_lm.server: how to set context size?

[chigkim]

It looks like I can set context size on mlx_lm.generate with --max-kv-size.
However, mlx_lm.server doesn't seem to have that option from command line nor api via http request.
Does it just allocate the maximum context size for the loaded model?
Or, max-kv-size is not what I'm thinking?
Thanks!

[awni]

Right now MLX LM server doesn't support that. No reason for that other than reactive feature support for MLX server..

To make it work basically we would need to get it from the request and set it here where the prompt cache is made. And if the max_kv_size is changed even if the model remains the same, we'd need to reset / make a new cache.

[chigkim]

Thanks for the response.
How does mlx_lm.server set the context size currently then?
Does it use the maximum context size of the loaded model, or does it allocate dynamically based on prompt size and max_tokens?
Thanks!

[awni]

The context size not set in the server. So it can grow arbitrarily large with the prompt / tokens being generated.

[chigkim]

Ah ok. If it just dynamically accommodates prompt + output, I guess you don't need this feature. Closing it.
Thanks!

[mattharrison]

You need to specify the size with the server and --max-tokens, the default is a pitiful 512.

Here's my writeup for using Aider with the server: Aider-AI/aider#4526

[digiperfect-tech]

@mattharrison Can we control context size and max tokens to generate separately? The docs suggest max-tokens controls how many tokens to generate (I assume it might just be setting the context size internally to that - but in general one would like context size to be much larger than max tokens to generate).

[awni]

max-tokens only controls the generated tokens. There is no context size limit in mlx_lm.server. In general you shouldn't exceed the model's context limit but right now we don't even check for that (maybe we should.. though it's somewhat rare to encounter it)

If you wanted to define a context limit less than the model's inherent limit, you'd need to define how that would behave.

• For example one could use a sliding window (which the RotatingKVCache implements).

• Another possibility is to simply truncate the prefix of the input.

If you have thoughts on what you are looking for there would be interested to hear.