Repository navigation
Add cache to process inputs #1111
Description
We currently apply apply-chat-template and process inputs in two places:
-
Data level
Line 58 in ac5328b
-
Rollout level
slime/slime/rollout/sglang_rollout.py
Line 99 in ac5328b
This causes duplicated computation and unnecessary overhead.
Why Two Stages Are Necessary
We cannot merge these into a single stage(right now, open to optimize):
- Data level: required to check context length early (with tokenization).
- Rollout level: performs tokenization and handles multimodal inputs.
Proposed Solution
Introduce caching to avoid duplicate work.
- Add a cache at the data level for processed chat templates.
- At the rollout level, reuse the cached result when available.
Activity
liuhao-labs commentedon Aug 23, 2026
Hi, I rechecked this against current
main. The duplicate work is still present:filter_long_prompt()tokenizes retained text prompts only to test their length, discards those IDs, and_prepare_prompt_ids()encodes the same rendered prompt again during the default rollout. #1113 improved tokenization correctness but did not retain or reuse those IDs.I can submit a focused CPU-tested fix that caches only retained text prompt IDs when max-length filtering is enabled and consumes that cache in the default rollout path. I would intentionally leave multimodal processor outputs uncached: retaining pixel tensors for the whole
Datasetand deepcopying them for every sampled sibling would create a much larger host-memory regression.Would this text-only, default-rollout scope be the preferred first PR, or would you prefer a compact cache representation/API that custom rollout functions can explicitly opt into? I will keep
max_length=None, custom rollout, and multimodal behavior unchanged.nanjiangwill commentedon Aug 23, 2026
replied in #2309 (comment)