feat(api): per-request SpecPrefill overrides on /v1/completions

[Thump604]

Summary

/v1/chat/completions already accepts per-request specprefill and specprefill_keep_pct overrides and threads them into engine.chat(). /v1/completions does not. This PR wires the same two fields through to /v1/completions on both the non-streaming and streaming paths.

Why

A SpecPrefill user with mixed /v1/chat/completions and /v1/completions traffic currently cannot opt individual /v1/completions requests into or out of SpecPrefill, and cannot tune specprefill_keep_pct per request on that endpoint. The server config is the only control. This closes that gap symmetrically against the existing chat endpoint.

Changes

• vllm_mlx/api/models.py: add specprefill: bool | None and specprefill_keep_pct: float | None to CompletionRequest, matching the fields already on ChatCompletionRequest.
• vllm_mlx/server.py::create_completion: extract both and pass as gen_kwargs into engine.generate(), mirroring the pattern at server.py:1421 in create_chat_completion.
• vllm_mlx/server.py::stream_completion: apply the same extraction so streaming clients have the same control.

Both new fields default to None, so existing behavior is unchanged for clients that don't set them. No schema changes to ChatCompletionRequest. No engine-side changes needed: SimpleEngine.stream_generate already pops specprefill and specprefill_keep_pct from kwargs (simple.py:307-308).

Total delta: +20 lines across 2 files.

Test plan

• POST /v1/completions with prompt + max_tokens but no specprefill field still returns 200 with content (default behavior unchanged).
• POST /v1/completions with {"specprefill": true} engages SpecPrefill on an eligible SimpleEngine model (non-MLLM with --specprefill-draft-model).
• POST /v1/completions with {"specprefill": false} bypasses SpecPrefill on the same eligible model.
• POST /v1/completions with {"specprefill_keep_pct": 0.2} uses 20% keep instead of the server default.
• POST /v1/completions with stream=true respects all three per-request controls above.

I've verified the default-behavior case locally on a runtime whose CompletionRequest schema is a strict superset of this PR's changes. The minimal version in this branch is structurally a subset and should behave identically for the default path.

[Thump604]

Quick merge signal: this is the API half of the /v1/completions SpecPrefill override pair with #266. The branch is mergeable, CI is green, and the local runtime validation on our side was clean once both halves were in place.

If you're doing a small sweep, #265 + #266 are the pair to take together.

[waybarrios]

Minor style note: in create_chat_completion the specprefill kwargs are added directly to the existing chat_kwargs dict, but here a separate gen_kwargs dict is created and spread with **gen_kwargs. Works the same, just a slightly different pattern — not a blocker, just a stylistic inconsistency.

[waybarrios]

To clarify my earlier comment — here's what I mean by the style difference.

In create_chat_completion, all kwargs go into a single dict:

chat_kwargs = {
    "max_tokens": request.max_tokens or _default_max_tokens,
    "temperature": _resolve_temperature(request.temperature),
    "top_p": _resolve_top_p(request.top_p),
}
if request.specprefill is not None:
    chat_kwargs["specprefill"] = request.specprefill
if request.specprefill_keep_pct is not None:
    chat_kwargs["specprefill_keep_pct"] = request.specprefill_keep_pct

await engine.chat(..., **chat_kwargs)

In this PR, the base args are passed inline and the specprefill ones go into a separate dict that gets spread:

gen_kwargs = {}
if request.specprefill is not None:
    gen_kwargs["specprefill"] = request.specprefill
if request.specprefill_keep_pct is not None:
    gen_kwargs["specprefill_keep_pct"] = request.specprefill_keep_pct

await engine.generate(
    prompt=prompt,
    max_tokens=...,
    temperature=...,
    top_p=...,
    stop=...,
    **gen_kwargs,
)

For consistency, you could unify into a single dict like chat does:

gen_kwargs = {
    "max_tokens": request.max_tokens or _default_max_tokens,
    "temperature": _resolve_temperature(request.temperature),
    "top_p": _resolve_top_p(request.top_p),
    "stop": request.stop,
}
if request.specprefill is not None:
    gen_kwargs["specprefill"] = request.specprefill
if request.specprefill_keep_pct is not None:
    gen_kwargs["specprefill_keep_pct"] = request.specprefill_keep_pct

await engine.generate(prompt=prompt, **gen_kwargs)

Just a nit — the PR works fine as-is.

[Thump604]

I aligned the /v1/completions SpecPrefill override handling style in 77066d9 so it now builds and mutates a single generation kwargs dict, matching the chat-completions path. Behavior is unchanged; this just removes the stylistic inconsistency you called out.

[Thump604]

Restacked on current main.

I kept the single-dict kwargs style you called out, but preserved the current completions-path repetition_penalty handling while adding the per-request specprefill / specprefill_keep_pct overrides.

Local verification on the rebased branch:

• python -m py_compile vllm_mlx/server.py vllm_mlx/api/models.py
• pytest -q tests/test_api_models.py tests/test_server.py

[janhilgard]

LGTM — clean API-side complement to #266 (already merged).

• specprefill and specprefill_keep_pct fields on CompletionRequest mirror the existing ChatCompletionRequest schema symmetrically
• Both create_completion and stream_completion forward overrides via a single generate_kwargs dict, matching the chat endpoint style
• Default None preserves existing behavior for clients that don't set these fields
• Rebase on current main is clean