fix: replace manual decode loop with pipelined generation in SpecPrefill Phase 4

[Vigilans]

Summary

Fixes #247. Try to make my first reviewed contribution here :)

SpecPrefill Phase 4 called model() + mx.eval() in a plain Python loop
without mx.async_eval pipelining, causing the GPU to idle between steps
(~0.3 tok/s vs 30+ tok/s normal decode).

• Extend MLXLanguageModel.stream_generate() to accept a prompt_cache
  parameter and mx.array/list[int] prompt types, so SpecPrefill can
  hand off the pre-populated sparse-prefill cache to the engine's standard
  pipelined decode path
• Replace the manual decode loop in both SpecPrefill paths:
  • Non-MLLM: hand off to self._model.stream_generate(prompt_cache=cache)
  • MLLM text-only: hand off to mlx_lm.stream_generate(prompt_cache=bc)
    (matching its normal path)
• First-token EOS check is preserved since stream_generate does not
  validate whether the prompt itself is EOS

Design decisions

Handoff target: Model layer self._model.stream_generate

There are three layers that can drive decode: Engine layer
(SimpleEngine.stream_generate), Model layer
(MLXLanguageModel.stream_generate), and mlx_lm.stream_generate directly.

The Engine layer is not suitable because _stream_generate_specprefill already
holds self._generation_lock (asyncio.Lock, non-reentrant) — calling back
into SimpleEngine.stream_generate would deadlock. It would also re-run
SpecPrefill eligibility checks and re-tokenize the prompt unnecessarily
(though no infinite recursion: the single-token prompt falls below the
threshold).

The Model layer (self._model.stream_generate) is the right fit. The
non-MLLM normal decode path already calls it (line 362), so routing
SpecPrefill Phase 4 through the same method maintains symmetry and reuses
its stop sequence handling. For the MLLM text-only path, self._text_model
is a raw mlx_lm TextModel (nn.Module built by build_text_model) with no
stream_generate method, so it calls mlx_lm.stream_generate directly —
again symmetric with the MLLM normal path (line 1090).

Extending MLXLanguageModel.stream_generate interface

The Model layer originally only accepted prompt: str and had no way to
pass a pre-populated cache. Two changes were needed:

1. prompt_cache parameter: forwarded via kwargs to
   mlx_lm.stream_generate, which passes it through to generate_step.
   When prompt_cache is not None, generate_step skips creating a new
   cache and uses the provided one directly.
2. prompt type broadened to str | mx.array | list[int]: Phase 4
   passes mx.array([first_token_id]) as prompt. mlx_lm.stream_generate
   already accepts all three types natively; the only adaptation needed was
   gating tokenizer.encode() behind an isinstance(prompt, str) check
   for the prompt token count.

Test plan

• Existing test_model_stream_generate still passes
• New tests: test_model_stream_generate_with_prompt_cache,
  test_model_stream_generate_with_list_prompt
• End-to-end: SpecPrefill activated on 12k-token prompt, Phase 4
  decode runs at normal speed

[Vigilans]

btw, I noticed SpecPrefill currently only triggers on the streaming path (stream_chat → stream_generate). The non-streaming chat() method calls self._model.chat() directly and bypasses stream_generate entirely, so SpecPrefill never activates for ·stream=False` requests. Is this intentional, or it will be implemented in the future?

[Thump604]

@Vigilans, @waybarrios: independent technical review of this PR.

Verification of the fix

Confirmed against current upstream main (b4fa030). The diff modifies two _run_specprefill call sites in vllm_mlx/engine/simple.py (the SimpleEngine streaming path and the BatchedEngine SpecPrefill path) to:

1. Sample the first token from SpecPrefill`s final logits (preserved behavior)
2. Hand off subsequent decode to mlx_lm.stream_generate with the SpecPrefill-prepared prompt_cache

The new prompt_cache parameter is added to MLXLanguageModel.stream_generate in vllm_mlx/models/llm.py:177 and forwarded into mlx_lms stream_generatevia the existingmtp_kwargs` dict.

Technical correctness

The root cause analysis in #247 holds. The manual model() + mx.eval(y) decode loop bypasses the pipelined kernel scheduling that mlx_lm.stream_generate provides, which is why the GPU stays idle (<1W vs 90-93W) and decode crawls at 0.3 tok/s vs 30+. The fix routes through the standard stream_generate path which gets full kernel pipelining and proper async eval batching.

Both _run_specprefill call sites get the same treatment, which is correct. Both BatchedEngine and SimpleEngine had the same bug.

Backward compatibility

prompt_cache=None is the default, existing callers behave identically. The prompt: Union[str, mx.array, list[int]] signature widening is safe. Existing string callers still work, new token-list and array callers are needed for the SpecPrefill handoff.

Minor observations (not blocking)

1. mtp_kwargs is now a slight misnomer since it also carries prompt_cache for non-MTP cases. Renaming to extra_kwargs would be clearer in a follow-up.
2. First-token decode is duplicated between the manual sample-and-decode block and what stream_generate would do internally. This is necessary because the first token must come from SpecPrefill`s final logits, not from a fresh forward pass. A code comment noting this would help future readers, but the logic is correct.
3. The two new tests in tests/test_llm.py exercise the new prompt_cache and list[int] prompt paths. An end-to-end SpecPrefill+stream_generate test that asserts the decode speed regression does not return would close the loop, but the production evidence in SpecPrefill Phase 4 in SimpleEngine: very slow decode using model() #247 is hard.

On the related observation in your comment

The non-streaming chat() method bypassing SpecPrefill is a separate, related gap. It is on a different code path (the bypass happens before the streaming pipeline is reached) and would need its own change. The two issues are complementary and the fixes would compose well if both land.

Recommendation

Merge candidate. Real fix to a real performance regression, hard production evidence, sound implementation, reasonable tests.

[Thump604]

This is a correct fix and a big one. The manual model() + mx.eval(y) per-token loop foregoes MLX's kernel pipelining, which is exactly what gets SpecPrefill Phase 4 stuck at ~0.3 tok/s while the same model at the same context runs ~30 tok/s through a normal decode. Routing through mlx_lm.stream_generate(prompt=mx.array([first_token]), prompt_cache=cache) is the clean fix.

On your note about SpecPrefill currently only triggering on the streaming path: that's a separate accumulator gap I'm working on. Once both this PR and that refactor land, non-streaming /v1/chat/completions and /v1/completions will also reach Phase 4, and this fix becomes load-bearing there too. Worth mentioning in the PR description if you want to foreshadow it.

Post-merge I'd like to run this against a 16K/64K/128K needle harness on Qwen 3.5 122B + 2B draft as an independent regression signal. If that's useful I can post the run output here as a follow-up comment.

Approving.

[Thump604]

Hey @Vigilans - the pipelined decode approach is correct and we've validated the same pattern independently on our local staging. Big speedup. Currently conflicting with main - can you rebase? Good to merge after that.

[Vigilans]

Hey @Vigilans - the pipelined decode approach is correct and we've validated the same pattern independently on our local staging. Big speedup. Currently conflicting with main - can you rebase? Good to merge after that.

Rebased on latest main and fixed the ruff lint errors from the previous CI run.

[janhilgard]

+1, agreeing with Thump604's approval. The fix is correct and well-scoped.

The manual model() + mx.eval(y) per-token loop was indeed the bottleneck -- it serializes every decode step and prevents MLX's kernel pipelining from overlapping compute and memory operations. Routing through mlx_lm.stream_generate(prompt=mx.array([first_token_id]), prompt_cache=cache) is the right fix.

A few minor observations (none blocking):

1. stream_generate signature broadening in llm.py -- accepting Union[str, mx.array, list[int]] is a nice generalization. The TYPE_CHECKING guard for import mlx.core as mx is correct since mx is only needed for the type annotation.

2. First token is decoded standalone (self._text_tokenizer.decode([first_token_id])) before handing off to stream_generate. This means the first token's text may differ slightly from what stream_generate's internal detokenizer would produce (e.g., for tokens that need surrounding context for correct decode, like byte-fallback tokens). In practice this is unlikely to matter for the first token after SpecPrefill, but worth being aware of.

3. The MLLM path (second _run_specprefill at ~line 1204) uses mlx_stream_generate (the raw mlx_lm function) directly, while the LLM path uses self._model.stream_generate (the wrapper). This asymmetry is fine since the MLLM path needs the raw model + tokenizer, but it means MTP is not available in the MLLM SpecPrefill path (the wrapper handles MTP kwargs). Not a regression since MTP was not available there before either.

CI has no checks running -- would be good to trigger a run to confirm green before merge. Branch is mergeable with no conflicts.

[Thump604]

This is approved by both reviewers and ready to go. Could you push a rebase onto current main to trigger CI? The branch hasn't had checks run.