[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).

Contribution checklist

Related work. This relates to #8197 (rollout profiling). A search for logits_to_keep and rollout prefill found no other PR.

Sibling paths checked. HybridEngineRollout has three prefill entry points:

• graph capture: changed;
• the explicit static-cache branch of continuous batching: changed;
• module.generate(): already keeps only the last logits, so it is untouched.

Commands and results.

pytest -p no:randomly tests/unit/runtime/rollout/

At head fba8f9184: 63 passed, 1 skipped, both with DS_ACCELERATOR=cpu and on 1× Intel Arc Pro B70 (XPU), Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11), torch 2.14.0.dev20260804+xpu, transformers 5.17.0. The measurement table above was taken on the same B70 setup.

Numerics, performance, memory. The selected first token is unchanged. Time and peak-memory results are in the Measurements table above.

Compatibility matrix.

	Validated	Not validated
Transformers	5.17.0 (static-cache prefill on device)	The _supports_cache_class branch on <5.x is covered only by the unit test
Rollout path	Prefill forward on device; continuous prefill in the unit test	End-to-end graph-capture rollout on device
ZeRO		Hybrid engine with ZeRO-3
Precision	bf16	fp16
Accelerator	XPU, CPU	CUDA

AI assistance. This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-17T14:54:43.432657Z	dd42f0e	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[nathon-lee]

Hi, Hi, @jinyouzhi Thanks for the optimization! The changes look good to me, and I didn't find any blocking issues.

One small, non-blocking suggestion: could the new test use position-dependent logits? Since every position currently predicts token 7, it would also pass if the wrong position were selected.

[jinyouzhi]

Hi, Hi, @jinyouzhi Thanks for the optimization! The changes look good to me, and I didn't find any blocking issues.

One small, non-blocking suggestion: could the new test use position-dependent logits? Since every position currently predicts token 7, it would also pass if the wrong position were selected.

Thank you for your remindering, I have updated the cases.