SIGN IN SIGN UP

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