[OPSD] Request last-position logits only in HybridEngine rollout prefill

[jinyouzhi]

Related to #8197.

Summary

Rollout prefill only needs the last position's logits to pick the first token. Two prefill call sites in HybridEngineRollout called the model without logits_to_keep, so the lm_head produced full [batch, prompt_len, vocab] logits:

• _generate_graph (graph-capture path)
• _continuous_prefill_cache, explicit static-cache branch (models that still expose _supports_cache_class)

This PR passes logits_to_keep=1 when the model's forward declares the argument, mirroring how Transformers' generate() decides. Models without it are unchanged: their **kwargs may be forwarded to submodules, so the argument is not passed blindly. The module.generate() branch already does this internally and is untouched.

Scope

• Graph-capture mode: always benefits (prefill calls the model directly).
• Continuous batching: benefits on Transformers versions that still expose _supports_cache_class and support logits_to_keep (>= 4.50). On Transformers 5.x this path goes through module.generate(), which already keeps only the last logits, so behavior is unchanged there.

Measurements

Prefill forward with StaticCache. Random-weight Llama with vocab 151936, hidden 1024, 4 layers, bf16. 1× Intel Arc Pro B70.

B × T	full logits	logits_to_keep=1	peak memory
8 × 1024	141.7 ms	43.3 ms	2.46 → 0.32 GiB
8 × 2048	279.0 ms	79.6 ms	4.92 → 0.64 GiB
16 × 2048	544.7 ms	149.5 ms	9.84 → 1.27 GiB

The selected first token is identical in all cases, and it matches the argmax of a plain HF forward. Small models gain the most because the lm_head dominates; for larger models the time saving shrinks, but the full-vocab activation is still avoided.

Tests

• New test_continuous_prefill_requests_last_logits_only: a model following the Transformers logits_to_keep contract gets logits_to_keep=1 at prefill, and the generated tokens are unchanged. The test fails without the change.
• tests/unit/runtime/rollout/: 63 passed, 1 skipped (CPU).