Add prefill_step_size as load param

[will-lms]

1. Add prefill_step_size as a parameter to load_model
2. Increase the default step size to 2048 default across the board. This was already the default for BatchedModelKit. Users (or the app) can set a lower value if desired.
3. Add tests for new parameter.

[will-lms]

Codex Review

The new prefill-step-size override is wired through most generation paths, but it is still ignored for
VisionModelKit-backed image requests. That leaves the advertised escape hatch ineffective for a real subset of
multimodal users.

Review comment:

• [P2] Honor prefill_step_size on VisionModelKit image requests — mlx-engine/mlx_engine/
  vision_model_kit/vision_model_kit.py:141-144
  On VisionModelKit-backed multimodal models (for example Qwen2/2.5-VL), this path still returns a fake one-token
  prompt and does the real prompt prefill inside VisionModelWrapper.call() as a single full-sequence pass.
  That makes the new prefill_step_size escape hatch a silent no-op whenever images_b64 is non-empty, so users
  lowering it to work around the new 2048 default will still see the same peak-memory behavior on image requests.

Will's response

This is true. The old prefill step size default was also not applied in the mlx-vlm image prompt path. It is out of scope to add for now. mlx-vlm does have a prefill_step_size but does not have a prompt_progress_callback equivalent.

[mattjcly]

nit: could call this resolve_prefill_step_size since not a pure validation function (resolves default if None)

Personal preference, totally non-blocking

[will-lms]

I had this first, but codex was unhappy b/c "it is also resolving, not just validating." It clearly does both, but the resolve_and_validate_... name is a mouth full.

[mattjcly]

For your consideration - from what I can tell prefill_step_size doesn't really have to be a load parameter, but instead could be a create_generator/inference-time parameter.

I think this would more directly model the domain of capabilities of this underlying API at this time and enable more flexibility without needing to reload the model.

I acknowledge llama.cpp treats batch size as a load parameter, so there would be a conceptual divergence there.

Also non-blocking IMO

[will-lms]

Yes I made this choice b/c we already treat prefill step size as a Load-time parameter for llama.cpp. I opted to keep consistency but agree that for this engine we could treat it is Prediction-time.

Going to defer for now as we don't have a use-case that would benefit from making it Prediction-time.