feat: add --prefill-step-size CLI flag

[kol22]

Summary

MLLMSchedulerConfig.prefill_step_size defaults to 1024 but isn't exposed as a
CLI argument. The MLLM batch generator enforces
total_prompt_tokens <= prefill_step_size * batch_count, so any single vision
request exceeding 1024 tokens fails with:

ValueError: Total prompt tokens (2289) exceeds safe limit (1024)

Since vision tokens alone typically exceed 1024 (images contribute ~1400+
tokens), this effectively blocks all MLLM inference under the default config.

Fix

Add --prefill-step-size to both serve and bench commands. Default of 0
means "use engine default" — 2048 for LLM, 1024 for MLLM — preserving existing
behavior. When set, the value flows through SchedulerConfig to both the LLM
scheduler and the MLLM scheduler.

Example usage:

vllm-mlx serve model-name --continuous-batching --prefill-step-size 16384

Test

Tested with Qwen3-VL-32B-Instruct-8bit. Before this fix, vision requests fail
with the ValueError above. After passing --prefill-step-size 16384, they
complete successfully.

[waybarrios]

Looked through this PR. The core plumbing for serve works correctly, the flag flows through SchedulerConfig into MLLMSchedulerConfig as expected.

A few things I noticed:

1. The --prefill-step-size flag is also added to bench_parser, but bench_command uses AsyncEngineCore which only runs the LLM Scheduler. The MLLM scheduler path lives in BatchedEngine._start_mllm(), which bench never calls. So the flag gets accepted on vllm-mlx bench but does nothing. Might be worth removing it from bench or adding a note.

2. The flag name --prefill-step-size is pretty generic, but it only controls the MLLM path (mllm_prefill_step_size). SchedulerConfig already has a separate prefill_step_size field (default 2048) for the LLM BatchGenerator. Someone running a plain LLM model with --prefill-step-size 512 would see the flag accepted with no effect. Something like --mllm-prefill-step-size would make the scope clearer and match the internal field name.

3. The 0 means use default contract lives only in the CLI layer. The dataclass field is Optional[int] with no validation, so if anything constructs SchedulerConfig(mllm_prefill_step_size=0) directly, that 0 gets forwarded straight to MLLMSchedulerConfig and you'd end up with max_batch_tokens = 0.

None of these are blockers, the fix for the original MLLM prefill guard issue works.

[waybarrios]

Would you mind addressing point 1 (removing the flag from bench_parser or wiring it to the MLLM bench path) and point 2 (renaming to --mllm-prefill-step-size) before merging? Those two would avoid confusion for users. Point 3 is minor and can be addressed later if you prefer.

[waybarrios]

btw @kol22 great work!

[kol22]

ee69b94

Addressed all three points :

1. Removed the prefill override flag from bench_parser since bench_command runs AsyncEngineCore (LLM scheduler path)
2. Renamed the serve flag to --mllm-prefill-step-size to make scope explicit and avoid confusion with the LLM prefill_step_size.
3. Added validation in SchedulerConfig.__post_init__ so mllm_prefill_step_size must be > 0 when provided (prevents direct non-CLI construction with 0 from propagating).

Appreciate the feedback and quick review!

[janhilgard]

@waybarrios — all three review items were addressed back on Feb 25 (removed from bench, renamed to --mllm-prefill-step-size, added > 0 validation). This has been waiting for a follow-up review for ~7 weeks now.

It does have merge conflicts from recent main changes — @kol22 would you mind rebasing? After that it should be ready to merge.

[kol22]

@janhilgard - resolved & rebased. Should be good to go.

[janhilgard]

LGTM — all three review items from @waybarrios have been addressed:

1. Removed from bench_parser — the flag is now serve-only, matching the actual MLLM scheduler path.
2. Renamed to --mllm-prefill-step-size — clear scope, matches the internal field name.
3. __post_init__ validation — mllm_prefill_step_size must be > 0 when provided, preventing the 0 propagation edge case.

Rebase on current main is clean. Code is minimal and consistent with existing patterns. Ready to merge.