Run body-only forwards for intermediate prefill chunks (#592)
On a chunked prefill, every chunk except the last samples nothing: the step still runs the full forward including the lm_head projection over the whole chunk, then throws the logits away. This PR runs the transformer body only on those steps. The hidden states still go through the forward-output submit, so the KV and GDN cache writes evaluate exactly as before; there is just no vocab projection and no sampling sync. The step qualifies only when every prefill entry is an intermediate chunk, there are no decode rows, and no drafter is installed (`propose()` runs unconditional per-step bookkeeping the skip would bypass). Body resolution lives on the model adapter next to `text_model`; a model whose structure cannot be resolved logs one warning and keeps the full forward. Reaching the sampling step without logits while any row must sample raises instead of guessing. This is flagless on purpose: on any step that produces output the path is structurally unchanged, and the fallback is the existing full forward. Measured on my M1 Ultra, branch vs base (`VLLM_METAL_MEMORY_FRACTION=0.5 vllm serve <model> --no-enable-prefix-caching`, greedy, prefill tok/s): | model | prompt | base |this pr | |---|---:|---:|---:| | Qwen3.6-35B-A3B-4bit | 16.2k tokens | 686.7 | 742.6 | | gemma-4-26B-A4B-it-OptiQ-4bit | 6k tokens | 389.9 | 413.8 | | gemma-4-e2b-it-4bit | 6k tokens | 1363.9 | 1556.8 | | Qwen3-0.6B | 6k tokens | 3654.5 | 4018.8 | --------- Signed-off-by: RickyChen / 陳昭儒 <ricky.chen@infinirc.com>
C
Chao-Ju Chen committed
05687887c12ae2e9962affa5bf868f2c0810e60c
Parent: e35b529
Committed by GitHub <noreply@github.com>
on 8/13/2026, 12:18:37 PM